github-mirrors/CLIProxyAPI
1
0
Fork 0
mirror of https://mirror.skon.top/github.com/router-for-me/CLIProxyAPI synced 2026-04-30 16:20:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #3089 from XYenon/feat/session-affinity

Browse Source 
feat: support Codex/PI session headers for session affinity
This commit is contained in:
Luis Pater
2026-04-28 23:54:16 +08:00
committed by GitHub
parent c4965befe7 3ac39dcc7d
commit f95ab4cdf1
4 changed files with 91 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
config.example.yaml
View File

@@ -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

4
internal/config/config.go
View File

@@ -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"`

43
sdk/cliproxy/auth/selector.go
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)
}

54
sdk/cliproxy/auth/selector_test.go
View File

@@ -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) {