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:
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user