feat: support Codex/PI session headers for session affinity

Amp-Thread-ID: https://ampcode.com/threads/T-019dce25-c070-773a-ac52-11c541220b30
Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
XYenon
2026-04-27 17:08:49 +08:00
parent 34027da7f1
commit 3ac39dcc7d
4 changed files with 91 additions and 15 deletions
+31 -12
View File
@@ -469,11 +469,14 @@ func NewSessionAffinitySelectorWithConfig(cfg SessionAffinityConfig) *SessionAff
// Pick selects an auth with session affinity when possible.
// Priority for session ID extraction:
// 1. metadata.user_id (Claude Code format) - highest priority
// 1. metadata.user_id (Claude Code format with _session_{uuid}) - highest priority
// 2. X-Session-ID header
// 3. metadata.user_id (non-Claude Code format)
// 4. conversation_id field
// 5. Hash-based fallback from messages
// 3. Session_id header (Codex)
// 4. X-Amp-Thread-Id header (Amp CLI thread ID)
// 5. X-Client-Request-Id header (PI)
// 6. metadata.user_id (non-Claude Code format)
// 7. conversation_id field in request body
// 8. Stable hash from first few messages content (fallback)
//
// Note: The cache key includes provider, session ID, and model to handle cases where
// a session uses multiple models (e.g., gemini-2.5-pro and gemini-3-flash-preview)
@@ -570,10 +573,12 @@ func (s *SessionAffinitySelector) InvalidateAuth(authID string) {
// Priority order:
// 1. metadata.user_id (Claude Code format with _session_{uuid}) - highest priority for Claude Code clients
// 2. X-Session-ID header
// 3. X-Amp-Thread-Id header (Amp CLI thread ID)
// 4. metadata.user_id (non-Claude Code format)
// 5. conversation_id field in request body
// 6. Stable hash from first few messages content (fallback)
// 3. Session_id header (Codex)
// 4. X-Amp-Thread-Id header (Amp CLI thread ID)
// 5. X-Client-Request-Id header (PI)
// 6. metadata.user_id (non-Claude Code format)
// 7. conversation_id field in request body
// 8. Stable hash from first few messages content (fallback)
func ExtractSessionID(headers http.Header, payload []byte, metadata map[string]any) string {
primary, _ := extractSessionIDs(headers, payload, metadata)
return primary
@@ -609,29 +614,43 @@ func extractSessionIDs(headers http.Header, payload []byte, metadata map[string]
}
}
// 3. X-Amp-Thread-Id header (Amp CLI thread ID)
// 3. Session_id header (Codex)
if headers != nil {
if sid := headers.Get("Session_id"); sid != "" {
return "codex:" + sid, ""
}
}
// 4. X-Amp-Thread-Id header (Amp CLI thread ID)
if headers != nil {
if tid := headers.Get("X-Amp-Thread-Id"); tid != "" {
return "amp:" + tid, ""
}
}
// 5. X-Client-Request-Id header (PI)
if headers != nil {
if rid := headers.Get("X-Client-Request-Id"); rid != "" {
return "clientreq:" + rid, ""
}
}
if len(payload) == 0 {
return "", ""
}
// 4. metadata.user_id (non-Claude Code format)
// 6. metadata.user_id (non-Claude Code format)
userID := gjson.GetBytes(payload, "metadata.user_id").String()
if userID != "" {
return "user:" + userID, ""
}
// 5. conversation_id field
// 7. conversation_id field
if convID := gjson.GetBytes(payload, "conversation_id").String(); convID != "" {
return "conv:" + convID, ""
}
// 6. Hash-based fallback from message content
// 8. Hash-based fallback from message content
return extractMessageHashIDs(payload)
}