diff --git a/config.example.yaml b/config.example.yaml index 22696069f..24e3d99c8 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -104,8 +104,9 @@ quota-exceeded: routing: strategy: "round-robin" # round-robin (default), fill-first # Enable universal session-sticky routing for all clients. - # Session IDs are extracted from: X-Session-ID header, Idempotency-Key, - # metadata.user_id, conversation_id, or first few messages hash. + # Session IDs are extracted from: metadata.user_id (Claude Code session format), + # X-Session-ID, Session_id (Codex), X-Amp-Thread-Id (Amp CLI), + # X-Client-Request-Id (PI), conversation_id, or first few messages hash. # Automatic failover is always enabled when bound auth becomes unavailable. session-affinity: false # default: false # How long session-to-auth bindings are retained. Default: 1h diff --git a/internal/config/config.go b/internal/config/config.go index 9817a8a71..1ee7aed53 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -226,7 +226,9 @@ type RoutingConfig struct { // SessionAffinity enables universal session-sticky routing for all clients. // Session IDs are extracted from multiple sources: - // X-Session-ID header, Idempotency-Key, metadata.user_id, conversation_id, or message hash. + // metadata.user_id (Claude Code session format), X-Session-ID, Session_id (Codex), + // X-Amp-Thread-Id (Amp CLI thread), X-Client-Request-Id (PI), metadata.user_id, + // conversation_id, or message hash. // Automatic failover is always enabled when bound auth becomes unavailable. SessionAffinity bool `yaml:"session-affinity,omitempty" json:"session-affinity,omitempty"` diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index f49979ce4..f0fe237c8 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -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) } diff --git a/sdk/cliproxy/auth/selector_test.go b/sdk/cliproxy/auth/selector_test.go index c3041b5ba..f6682c6fc 100644 --- a/sdk/cliproxy/auth/selector_test.go +++ b/sdk/cliproxy/auth/selector_test.go @@ -776,6 +776,46 @@ func TestExtractSessionID_Headers(t *testing.T) { } } +func TestExtractSessionID_CodexSessionIDHeader(t *testing.T) { + t.Parallel() + + headers := make(http.Header) + headers.Set("Session_id", "codex-session-123") + + got := ExtractSessionID(headers, nil, nil) + want := "codex:codex-session-123" + if got != want { + t.Errorf("ExtractSessionID() with Session_id = %q, want %q", got, want) + } +} + +func TestExtractSessionID_ClientRequestIDHeader(t *testing.T) { + t.Parallel() + + headers := make(http.Header) + headers.Set("X-Client-Request-Id", "pi-session-123") + + got := ExtractSessionID(headers, nil, nil) + want := "clientreq:pi-session-123" + if got != want { + t.Errorf("ExtractSessionID() with X-Client-Request-Id = %q, want %q", got, want) + } +} + +func TestExtractSessionID_CodexSessionIDPriorityOverClientRequestID(t *testing.T) { + t.Parallel() + + headers := make(http.Header) + headers.Set("X-Client-Request-Id", "pi-session-123") + headers.Set("Session_id", "codex-session-456") + + got := ExtractSessionID(headers, nil, nil) + want := "codex:codex-session-456" + if got != want { + t.Errorf("ExtractSessionID() = %q, want %q (Session_id should take priority over X-Client-Request-Id)", got, want) + } +} + func TestExtractSessionID_AmpThreadId(t *testing.T) { t.Parallel() @@ -789,6 +829,20 @@ func TestExtractSessionID_AmpThreadId(t *testing.T) { } } +func TestExtractSessionID_AmpThreadIdPriorityOverClientRequestID(t *testing.T) { + t.Parallel() + + headers := make(http.Header) + headers.Set("X-Amp-Thread-Id", "T-priority-test") + headers.Set("X-Client-Request-Id", "pi-session-123") + + got := ExtractSessionID(headers, nil, nil) + want := "amp:T-priority-test" + if got != want { + t.Errorf("ExtractSessionID() = %q, want %q (X-Amp-Thread-Id should take priority over X-Client-Request-Id)", got, want) + } +} + // TestExtractSessionID_AmpThreadIdLowerPriority verifies X-Amp-Thread-Id is lower // priority than Claude Code metadata.user_id but higher than conversation_id. func TestExtractSessionID_AmpThreadIdPriority(t *testing.T) {