diff --git a/config.example.yaml b/config.example.yaml index 653b78421..0db308f97 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -197,9 +197,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: metadata.user_id (Claude Code session format), - # X-Session-ID, Session_id (Codex), X-Client-Request-Id (PI), conversation_id, - # or first few messages hash. + # Explicit Claude Code, Codex, OpenCode, and pi session headers are preferred, + # followed by prompt_cache_key, Responses conversation IDs, legacy body IDs, + # execution or derived session identity, and the existing first-message hash fallback. # 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_types.go b/internal/config/config_types.go index 76eb6ef30..c0eaa1d6a 100644 --- a/internal/config/config_types.go +++ b/internal/config/config_types.go @@ -207,9 +207,9 @@ type RoutingConfig struct { Strategy string `yaml:"strategy,omitempty" json:"strategy,omitempty"` // SessionAffinity enables universal session-sticky routing for all clients. - // Session IDs are extracted from multiple sources: - // metadata.user_id (Claude Code session format), X-Session-ID, Session_id (Codex), - // X-Client-Request-Id (PI), metadata.user_id, conversation_id, or message hash. + // Explicit Claude Code, Codex, OpenCode, and pi session headers are preferred, + // followed by prompt_cache_key, Responses conversation IDs, legacy body IDs, + // execution or derived session identity, and the existing message-content hash fallback. // 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/conductor.go b/sdk/cliproxy/auth/conductor.go index cf537b6df..d25524a6c 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -121,6 +121,7 @@ type Manager struct { // homeSessionSelections owns retained Home selections for websocket sessions. homeSessionSelections map[string]map[homeSessionSelectionKey]*HomeDispatchSelection homeSessionLocks sync.Map + homeSessionAliases homeSessionAliasCache // providerOffsets tracks per-model provider rotation state for multi-provider routing. providerOffsets map[string]int homeDispatchBundle atomic.Pointer[HomeDispatchBundle] diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index d21ab818a..6d24ce468 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -117,6 +117,10 @@ func (m *Manager) setConfigSnapshotLocked(cfg *internalconfig.Config) bool { m.mu.RLock() oldCooldownStore := m.cooldownStore m.mu.RUnlock() + previousCfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) + if homeSessionAliasTTL(previousCfg) != homeSessionAliasTTL(cfg) { + m.homeSessionAliases.clear() + } m.runtimeConfig.Store(cfg) clearedCooldowns := m.clearDisabledCooldownStates(cfg) if clearedCooldowns && oldCooldownStore != nil { diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go index 0c85f17a3..92ef35084 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -509,6 +509,7 @@ func (m *Manager) clearHomeRuntimeAuths() { m.clearHomeRuntimeAuthsLocked() selections := m.takeAllHomeSessionSelectionsLocked() m.mu.Unlock() + m.homeSessionAliases.clear() for _, selection := range selections { selection.End("home_disabled") } @@ -716,7 +717,7 @@ func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, o return nil, &Error{Code: "home_unavailable", Message: "home execution registry unavailable", Retryable: true, HTTPStatus: http.StatusServiceUnavailable} } - sessionID := ExtractSessionID(opts.Headers, opts.OriginalRequest, opts.Metadata) + sessionID := m.homeDispatchSessionID(opts) dispatchHeaders := homeDispatchHeaders(ctx, opts.Headers) raw, errRPop := client.RPopAuth(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata)) if errRPop != nil { diff --git a/sdk/cliproxy/auth/home_session_alias.go b/sdk/cliproxy/auth/home_session_alias.go new file mode 100644 index 000000000..f422441a4 --- /dev/null +++ b/sdk/cliproxy/auth/home_session_alias.go @@ -0,0 +1,243 @@ +package auth + +import ( + "container/list" + "strings" + "sync" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +const ( + defaultHomeSessionAliasTTL = time.Hour + homeSessionAliasCleanupOps = 256 + homeSessionAliasSoftLimit = 4096 +) + +type homeSessionAliasEntry struct { + canonical string + expiresAt time.Time + aliases []string +} + +// homeSessionAliasCache reconciles multiple client identifiers for one Home +// session without changing Home's single-session-ID protocol. +type homeSessionAliasCache struct { + mu sync.Mutex + entries map[string]homeSessionAliasEntry + groups map[string]homeSessionAliasEntry + evictionOrder *list.List + evictionElements map[string]*list.Element + ops uint64 +} + +func (c *homeSessionAliasCache) canonical(primary, fallback string, ttl time.Duration, now time.Time) string { + primary = strings.TrimSpace(primary) + fallback = strings.TrimSpace(fallback) + if primary == "" { + return "" + } + if ttl <= 0 { + ttl = defaultHomeSessionAliasTTL + } + + c.mu.Lock() + defer c.mu.Unlock() + c.ensureInitializedLocked() + c.ops++ + if c.ops%homeSessionAliasCleanupOps == 0 { + c.cleanupLocked(now) + } + + canonical := primary + aliases := mergeSessionAliases(nil, primary, fallback) + previousGroups := make(map[string]homeSessionAliasEntry, 2) + remember := func(entry homeSessionAliasEntry) { + previousGroups[entry.canonical] = entry + } + + primaryFound := false + canonicalFromLiveAlias := false + if existing, ok := c.entryLocked(primary, now); ok { + primaryFound = true + canonicalFromLiveAlias = true + canonical = existing.canonical + remember(existing) + aliases = mergeSessionAliases(aliases, existing.aliases...) + } + if fallback != "" && fallback != primary { + if existing, ok := c.entryLocked(fallback, now); ok { + canonicalFromLiveAlias = true + if !primaryFound { + canonical = existing.canonical + } + remember(existing) + aliases = mergeSessionAliases(aliases, existing.aliases...) + } + } + if canonicalFromLiveAlias { + if existing, ok := c.groupLocked(canonical, now); ok { + remember(existing) + aliases = mergeSessionAliases(aliases, existing.aliases...) + } + } + if !canonicalFromLiveAlias { + if _, ok := c.groupLocked(canonical, now); ok { + return canonical + } + } + aliases = compactHomeSessionAliases(mergeSessionAliases(aliases, canonical)) + for _, previous := range previousGroups { + c.removeGroupLocked(previous) + } + + c.setGroupLocked(homeSessionAliasEntry{ + canonical: canonical, + expiresAt: now.Add(ttl), + aliases: aliases, + }) + c.enforceLimitLocked(homeSessionAliasSoftLimit) + return canonical +} + +func (c *homeSessionAliasCache) ensureInitializedLocked() { + if c.entries == nil { + c.entries = make(map[string]homeSessionAliasEntry) + } + if c.groups == nil { + c.groups = make(map[string]homeSessionAliasEntry) + } + if c.evictionOrder == nil { + c.evictionOrder = list.New() + } + if c.evictionElements == nil { + c.evictionElements = make(map[string]*list.Element) + } +} + +func (c *homeSessionAliasCache) entryLocked(alias string, now time.Time) (homeSessionAliasEntry, bool) { + entry, ok := c.entries[alias] + if !ok { + return homeSessionAliasEntry{}, false + } + if now.Before(entry.expiresAt) { + return entry, true + } + if group, exists := c.groups[entry.canonical]; exists && sameHomeSessionAliasGroup(group, entry) { + c.removeGroupLocked(group) + } else { + delete(c.entries, alias) + } + return homeSessionAliasEntry{}, false +} + +func (c *homeSessionAliasCache) groupLocked(canonical string, now time.Time) (homeSessionAliasEntry, bool) { + entry, ok := c.groups[canonical] + if !ok { + return homeSessionAliasEntry{}, false + } + if now.Before(entry.expiresAt) { + return entry, true + } + c.removeGroupLocked(entry) + return homeSessionAliasEntry{}, false +} + +func (c *homeSessionAliasCache) setGroupLocked(entry homeSessionAliasEntry) { + if existing, ok := c.groups[entry.canonical]; ok { + c.removeGroupLocked(existing) + } + entry.aliases = append([]string(nil), entry.aliases...) + c.groups[entry.canonical] = entry + for _, alias := range entry.aliases { + c.entries[alias] = entry + } + c.evictionElements[entry.canonical] = c.evictionOrder.PushBack(entry.canonical) +} + +func (c *homeSessionAliasCache) removeGroupLocked(entry homeSessionAliasEntry) { + current, ok := c.groups[entry.canonical] + if !ok || !sameHomeSessionAliasGroup(current, entry) { + return + } + for _, alias := range current.aliases { + mapped, exists := c.entries[alias] + if exists && sameHomeSessionAliasGroup(mapped, current) { + delete(c.entries, alias) + } + } + delete(c.groups, current.canonical) + if element, exists := c.evictionElements[current.canonical]; exists { + c.evictionOrder.Remove(element) + delete(c.evictionElements, current.canonical) + } +} + +func sameHomeSessionAliasGroup(left, right homeSessionAliasEntry) bool { + return left.canonical == right.canonical && left.expiresAt.Equal(right.expiresAt) && + equalSessionAliases(left.aliases, right.aliases) +} + +func (c *homeSessionAliasCache) enforceLimitLocked(limit int) { + if limit <= 0 { + return + } + for len(c.entries) > limit { + oldest := c.evictionOrder.Front() + if oldest == nil { + return + } + canonical, _ := oldest.Value.(string) + entry, ok := c.groups[canonical] + if !ok { + c.evictionOrder.Remove(oldest) + delete(c.evictionElements, canonical) + continue + } + c.removeGroupLocked(entry) + } +} + +func (c *homeSessionAliasCache) cleanupLocked(now time.Time) { + for _, entry := range c.groups { + if !now.Before(entry.expiresAt) { + c.removeGroupLocked(entry) + } + } +} + +func (c *homeSessionAliasCache) clear() { + c.mu.Lock() + c.entries = nil + c.groups = nil + c.evictionOrder = nil + c.evictionElements = nil + c.ops = 0 + c.mu.Unlock() +} + +func homeSessionAliasTTL(cfg *internalconfig.Config) time.Duration { + if cfg == nil { + return defaultHomeSessionAliasTTL + } + raw := strings.TrimSpace(cfg.Routing.SessionAffinityTTL) + if raw == "" { + return defaultHomeSessionAliasTTL + } + parsed, errParse := time.ParseDuration(raw) + if errParse != nil || parsed <= 0 { + return defaultHomeSessionAliasTTL + } + return parsed +} + +func (m *Manager) homeDispatchSessionID(opts cliproxyexecutor.Options) string { + primary, fallback := extractSessionIDs(opts.Headers, opts.OriginalRequest, opts.Metadata) + if primary == "" || m == nil { + return primary + } + cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) + return m.homeSessionAliases.canonical(primary, fallback, homeSessionAliasTTL(cfg), time.Now()) +} diff --git a/sdk/cliproxy/auth/home_session_alias_test.go b/sdk/cliproxy/auth/home_session_alias_test.go new file mode 100644 index 000000000..d271aab11 --- /dev/null +++ b/sdk/cliproxy/auth/home_session_alias_test.go @@ -0,0 +1,329 @@ +package auth + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "sync" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type sessionAliasCaptureDispatcher struct { + mu sync.Mutex + sessions []string +} + +func (*sessionAliasCaptureDispatcher) HeartbeatOK() bool { return true } + +func (d *sessionAliasCaptureDispatcher) RPopAuth(_ context.Context, _ string, sessionID string, _ http.Header, _ int) ([]byte, error) { + d.mu.Lock() + d.sessions = append(d.sessions, sessionID) + d.mu.Unlock() + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ + ID: "home-session-alias-auth", + Provider: "home-session-alias", + Status: StatusActive, + }}) +} + +func (*sessionAliasCaptureDispatcher) AbortAmbiguousDispatch() {} + +func (d *sessionAliasCaptureDispatcher) sessionIDs() []string { + d.mu.Lock() + defer d.mu.Unlock() + return append([]string(nil), d.sessions...) +} + +func TestHomeSessionAliasCacheClearsWhenConfiguredTTLChanges(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{ + Home: internalconfig.HomeConfig{Enabled: true}, + Routing: internalconfig.RoutingConfig{SessionAffinityTTL: "1h"}, + }) + combined := cliproxyexecutor.Options{OriginalRequest: []byte( + `{"conversation":{"id":"ttl-conversation"},"prompt_cache_key":"ttl-prompt"}`, + )} + conversationOnly := cliproxyexecutor.Options{OriginalRequest: []byte( + `{"conversation":{"id":"ttl-conversation"}}`, + )} + if got := manager.homeDispatchSessionID(combined); got != "pck:ttl-prompt" { + t.Fatalf("combined canonical = %q, want pck:ttl-prompt", got) + } + if got := manager.homeDispatchSessionID(conversationOnly); got != "pck:ttl-prompt" { + t.Fatalf("conversation canonical before reload = %q, want existing prompt canonical", got) + } + + manager.SetConfig(&internalconfig.Config{ + Home: internalconfig.HomeConfig{Enabled: true}, + Routing: internalconfig.RoutingConfig{SessionAffinityTTL: "1m"}, + }) + if got := manager.homeDispatchSessionID(conversationOnly); got != "conv:ttl-conversation" { + t.Fatalf("conversation canonical after TTL change = %q, want cleared alias cache", got) + } +} + +func TestHomeDispatchCanonicalizesPromptCacheAndConversationAliases(t *testing.T) { + tests := []struct { + name string + payloads []string + want string + }{ + { + name: "conversation then combined then prompt cache", + payloads: []string{ + `{"conversation":{"id":"conversation-session"}}`, + `{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`, + `{"prompt_cache_key":"shared-cache-bucket"}`, + }, + want: "conv:conversation-session", + }, + { + name: "prompt cache then combined then conversation", + payloads: []string{ + `{"prompt_cache_key":"shared-cache-bucket"}`, + `{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`, + `{"conversation":{"id":"conversation-session"}}`, + }, + want: "pck:shared-cache-bucket", + }, + { + name: "combined request establishes prompt cache primary", + payloads: []string{ + `{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`, + `{"conversation":{"id":"conversation-session"}}`, + `{"prompt_cache_key":"shared-cache-bucket"}`, + }, + want: "pck:shared-cache-bucket", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dispatcher := &sessionAliasCaptureDispatcher{} + manager := newHomeSelectionTestManager(t, dispatcher) + manager.RegisterExecutor(schedulerTestExecutor{provider: "home-session-alias"}) + + for _, payload := range tt.payloads { + selection, errSelection := manager.pickHomeDispatchSelection(context.Background(), "gpt-test", cliproxyexecutor.Options{ + OriginalRequest: []byte(payload), + }) + if errSelection != nil { + t.Fatalf("pickHomeDispatchSelection() error = %v", errSelection) + } + selection.End("test_complete") + } + + got := dispatcher.sessionIDs() + if len(got) != len(tt.payloads) { + t.Fatalf("Home session IDs = %#v, want %d entries", got, len(tt.payloads)) + } + for index, sessionID := range got { + if sessionID != tt.want { + t.Fatalf("Home session ID[%d] = %q, want %q; all=%#v", index, sessionID, tt.want, got) + } + } + }) + } +} + +func TestHomeSessionAliasCachePrimaryAccessRefreshesWholeAliasGroup(t *testing.T) { + var cache homeSessionAliasCache + now := time.Now() + const primary = "pck:shared-cache-bucket" + const fallback = "conv:conversation-session" + + if got := cache.canonical(primary, fallback, time.Minute, now); got != primary { + t.Fatalf("initial canonical = %q, want %q", got, primary) + } + cache.mu.Lock() + fallbackEntry := cache.entries[fallback] + fallbackEntry.expiresAt = now.Add(-time.Second) + cache.entries[fallback] = fallbackEntry + cache.mu.Unlock() + + if got := cache.canonical(primary, "", time.Minute, now.Add(10*time.Second)); got != primary { + t.Fatalf("primary-only canonical = %q, want %q", got, primary) + } + if got := cache.canonical(fallback, "", time.Minute, now.Add(20*time.Second)); got != primary { + t.Fatalf("fallback canonical after active primary traffic = %q, want %q", got, primary) + } +} + +func TestHomeSessionAliasCacheSharedPromptKeyPreservesConversationAliases(t *testing.T) { + var cache homeSessionAliasCache + now := time.Now() + const promptKey = "pck:shared-cache-bucket" + const conversationA = "conv:conversation-a" + const conversationB = "conv:conversation-b" + + if got := cache.canonical(promptKey, conversationA, time.Minute, now); got != promptKey { + t.Fatalf("conversation A canonical = %q, want %q", got, promptKey) + } + if got := cache.canonical(promptKey, conversationB, time.Minute, now.Add(time.Second)); got != promptKey { + t.Fatalf("conversation B canonical = %q, want %q", got, promptKey) + } + if got := cache.canonical(conversationA, "", time.Minute, now.Add(2*time.Second)); got != promptKey { + t.Fatalf("conversation A alias canonical = %q, want %q", got, promptKey) + } + if got := cache.canonical(conversationB, "", time.Minute, now.Add(3*time.Second)); got != promptKey { + t.Fatalf("conversation B alias canonical = %q, want %q", got, promptKey) + } +} + +func TestHomeSessionAliasCacheConversationIDContainingPromptMarkerRemainsStable(t *testing.T) { + var cache homeSessionAliasCache + now := time.Now() + const promptKey = "pck:shared-cache-bucket" + const conversation = "conv:a::pck:b" + if got := cache.canonical(promptKey, conversation, time.Minute, now); got != promptKey { + t.Fatalf("combined canonical = %q, want %q", got, promptKey) + } + if got := cache.canonical(conversation, "", time.Minute, now.Add(time.Second)); got != promptKey { + t.Fatalf("conversation-only canonical = %q, want %q", got, promptKey) + } +} + +func TestHomeSessionAliasCacheSharedPromptKeyCapsStableAliasesByRecency(t *testing.T) { + var cache homeSessionAliasCache + now := time.Now() + const promptKey = "pck:shared-cache-bucket" + for index := 0; index < 128; index++ { + conversation := fmt.Sprintf("conv:conversation-%03d", index) + cache.canonical(promptKey, conversation, time.Minute, now.Add(time.Duration(index)*time.Second)) + } + + cache.mu.Lock() + defer cache.mu.Unlock() + if len(cache.entries) > 65 { + t.Fatalf("home alias entries = %d, want one prompt key plus at most 64 stable aliases", len(cache.entries)) + } + if _, ok := cache.entries["conv:conversation-127"]; !ok { + t.Fatal("newest Home conversation alias was not retained") + } + if _, ok := cache.entries["conv:conversation-000"]; ok { + t.Fatal("oldest Home conversation alias was retained after stable-alias cap") + } +} + +func TestHomeSessionAliasCacheRotatingPrimaryEvictsObsoleteAliases(t *testing.T) { + var cache homeSessionAliasCache + now := time.Now() + const fallback = "conv:conversation-session" + wantCanonical := "pck:cache-00" + for index := 0; index < 16; index++ { + primary := fmt.Sprintf("pck:cache-%02d", index) + if got := cache.canonical(primary, fallback, time.Minute, now.Add(time.Duration(index)*time.Second)); got != wantCanonical { + t.Fatalf("canonical at index %d = %q, want %q", index, got, wantCanonical) + } + } + latest := "pck:cache-15" + + cache.mu.Lock() + defer cache.mu.Unlock() + if len(cache.entries) != 2 { + t.Fatalf("home alias entries = %d, want only latest primary and fallback", len(cache.entries)) + } + if _, ok := cache.entries[latest]; !ok { + t.Fatalf("latest primary %q was not retained", latest) + } + if _, ok := cache.entries[fallback]; !ok { + t.Fatalf("fallback %q was not retained", fallback) + } + if _, ok := cache.entries[wantCanonical]; ok { + t.Fatalf("obsolete canonical alias %q was retained as a lookup key", wantCanonical) + } + if aliases := cache.entries[fallback].aliases; len(aliases) != 2 { + t.Fatalf("home fallback alias group = %#v, want exactly two active identifiers", aliases) + } +} + +func TestHomeSessionAliasCacheDoesNotReconnectCompactedCanonicalAlias(t *testing.T) { + var cache homeSessionAliasCache + now := time.Now() + const obsoletePrompt = "pck:cache-a" + const currentPrompt = "pck:cache-b" + const conversation = "conv:conversation-session" + + if got := cache.canonical(obsoletePrompt, conversation, time.Minute, now); got != obsoletePrompt { + t.Fatalf("initial canonical = %q, want %q", got, obsoletePrompt) + } + if got := cache.canonical(currentPrompt, conversation, time.Minute, now.Add(time.Second)); got != obsoletePrompt { + t.Fatalf("rotated canonical = %q, want stable %q", got, obsoletePrompt) + } + + cache.mu.Lock() + if _, ok := cache.entries[obsoletePrompt]; ok { + cache.mu.Unlock() + t.Fatalf("obsolete prompt alias %q remained live after compaction", obsoletePrompt) + } + cache.mu.Unlock() + + if got := cache.canonical(obsoletePrompt, "", time.Minute, now.Add(2*time.Second)); got != obsoletePrompt { + t.Fatalf("obsolete prompt canonical = %q, want standalone %q", got, obsoletePrompt) + } + + cache.mu.Lock() + conversationEntry, conversationOK := cache.entries[conversation] + currentEntry, currentOK := cache.entries[currentPrompt] + _, obsoleteOK := cache.entries[obsoletePrompt] + cache.mu.Unlock() + if obsoleteOK { + t.Fatalf("stale canonical %q replaced the live group", obsoletePrompt) + } + if !conversationOK || !currentOK || !sameHomeSessionAliasGroup(conversationEntry, currentEntry) { + t.Fatalf("live aliases were disconnected: conversation=%#v current=%#v", conversationEntry, currentEntry) + } + if got := cache.canonical(conversation, "", time.Minute, now.Add(3*time.Second)); got != obsoletePrompt { + t.Fatalf("live conversation canonical = %q, want %q", got, obsoletePrompt) + } +} + +func TestHomeSessionAliasCacheSoftLimitEvictsOldestTouchedGroup(t *testing.T) { + var cache homeSessionAliasCache + now := time.Now() + const oldest = "session:zzzz-oldest" + cache.canonical(oldest, "", time.Hour, now) + for index := 0; index < homeSessionAliasSoftLimit; index++ { + cache.canonical(fmt.Sprintf("session:%05d", index), "", time.Hour, now) + } + + cache.mu.Lock() + defer cache.mu.Unlock() + if len(cache.entries) > homeSessionAliasSoftLimit { + t.Fatalf("alias entries = %d, want at most %d", len(cache.entries), homeSessionAliasSoftLimit) + } + if _, ok := cache.entries[oldest]; ok { + t.Fatalf("oldest insertion %q remained after incremental eviction", oldest) + } + if _, ok := cache.entries["session:00000"]; !ok { + t.Fatal("newer insertion was evicted instead of the oldest group") + } +} + +func TestHomeSessionAliasCacheEnforcesSoftLimit(t *testing.T) { + var cache homeSessionAliasCache + now := time.Now() + for i := 0; i < homeSessionAliasSoftLimit+32; i++ { + cache.canonical(fmt.Sprintf("session:%05d", i), "", time.Hour, now.Add(time.Duration(i)*time.Nanosecond)) + } + + cache.mu.Lock() + entryCount := len(cache.entries) + _, oldestPresent := cache.entries["session:00000"] + _, newestPresent := cache.entries[fmt.Sprintf("session:%05d", homeSessionAliasSoftLimit+31)] + cache.mu.Unlock() + if entryCount > homeSessionAliasSoftLimit { + t.Fatalf("alias entries = %d, want at most %d", entryCount, homeSessionAliasSoftLimit) + } + if oldestPresent { + t.Fatal("oldest alias remained after enforcing soft limit") + } + if !newestPresent { + t.Fatal("newest alias was evicted while enforcing soft limit") + } +} diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 4230bddc0..4be062570 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -7,7 +7,6 @@ import ( "hash/fnv" "math" "net/http" - "regexp" "sort" "strconv" "strings" @@ -362,10 +361,6 @@ func isAuthBlockedForModel(auth *Auth, model string, now time.Time) (bool, block return false, blockReasonNone, time.Time{} } -// sessionPattern matches Claude Code user_id format: -// user_{hash}_account__session_{uuid} -var sessionPattern = regexp.MustCompile(`_session_([a-f0-9-]+)$`) - // SessionAffinitySelector wraps another selector with session-sticky behavior. // It extracts session ID from multiple sources and maintains session-to-auth // mappings with automatic failover when the bound auth becomes unavailable. @@ -403,14 +398,8 @@ func NewSessionAffinitySelectorWithConfig(cfg SessionAffinityConfig) *SessionAff } // Pick selects an auth with session affinity when possible. -// Priority for session ID extraction: -// 1. metadata.user_id containing a Claude Code session -// 2. Explicit session headers -// 3. X-Client-Request-Id header -// 4. Explicit request-body session and user fields -// 5. Explicit execution session metadata -// 6. Stable context-derived session identity -// 7. Legacy message hash fallback +// Explicit Claude Code, Codex, OpenCode, pi, and request-body session signals +// precede execution metadata, stable derived identity, and the legacy hash 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) @@ -430,10 +419,22 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri } cacheKey := provider + "::" + primaryID + "::" + model + fallbackKey := "" + if fallbackID != "" && fallbackID != primaryID { + fallbackKey = provider + "::" + fallbackID + "::" + model + } + bind := func(authID string) { + if fallbackKey != "" { + s.cache.SetAliases(authID, cacheKey, fallbackKey) + return + } + s.cache.Set(cacheKey, authID) + } if cachedAuthID, ok := s.cache.GetAndRefresh(cacheKey); ok { for _, auth := range available { if auth.ID == cachedAuthID { + bind(auth.ID) entry.Infof("session-affinity: cache hit | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) return auth, nil } @@ -443,17 +444,16 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if err != nil { return nil, err } - s.cache.Set(cacheKey, auth.ID) + bind(auth.ID) entry.Infof("session-affinity: cache hit but auth unavailable, reselected | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) return auth, nil } - if fallbackID != "" && fallbackID != primaryID { - fallbackKey := provider + "::" + fallbackID + "::" + model + if fallbackKey != "" { if cachedAuthID, ok := s.cache.Get(fallbackKey); ok { for _, auth := range available { if auth.ID == cachedAuthID { - s.cache.Set(cacheKey, auth.ID) + bind(auth.ID) entry.Infof("session-affinity: fallback cache hit | session=%s fallback=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), auth.ID, provider, model) return auth, nil } @@ -465,7 +465,7 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if err != nil { return nil, err } - s.cache.Set(cacheKey, auth.ID) + bind(auth.ID) entry.Infof("session-affinity: cache miss, new binding | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) return auth, nil } @@ -503,113 +503,124 @@ func (s *SessionAffinitySelector) InvalidateAuth(authID string) { } } -// ExtractSessionID extracts session identifier from multiple sources. +// normalizedSessionCandidate validates an explicit client-provided session signal. +// It keeps opaque printable IDs intact while rejecting values that are unsafe or +// implausibly large for routing keys and logs. +func normalizedSessionCandidate(raw string) string { + return cliproxysession.NormalizeExplicitID(raw) +} + +func sessionHeaderValue(headers http.Header, name string) string { + if headers == nil { + return "" + } + if value := normalizedSessionCandidate(headers.Get(name)); value != "" { + return value + } + for key, values := range headers { + if !strings.EqualFold(key, name) { + continue + } + for _, raw := range values { + if value := normalizedSessionCandidate(raw); value != "" { + return value + } + } + } + return "" +} + +// ExtractSessionID extracts a session identifier from explicit client signals, +// then falls back to execution metadata, derived identity, and message history. // Priority order: -// 1. metadata.user_id containing a Claude Code session -// 2. Explicit session headers -// 3. X-Client-Request-Id header -// 4. Explicit request-body session and user fields -// 5. Explicit execution session metadata -// 6. Stable context-derived session identity -// 7. Legacy message hash fallback +// 1. X-Claude-Code-Session-Id +// 2. Claude Code metadata.user_id session +// 3. Session-Id / Session_id (Codex and compatible clients) +// 4. X-Session-ID +// 5. X-Session-Affinity (OpenCode) +// 6. X-Client-Request-Id (pi Responses) +// 7. session_id / sessionId +// 8. prompt_cache_key, with conversation / conversation.id as an alias +// 9. metadata.user_id and conversation_id legacy body fields +// 10. explicit execution session metadata +// 11. stable context-derived session identity +// 12. stable hash from initial message content func ExtractSessionID(headers http.Header, payload []byte, metadata map[string]any) string { primary, _ := extractSessionIDs(headers, payload, metadata) return primary } // extractSessionIDs returns (primaryID, fallbackID) for session affinity. -// primaryID: full hash including assistant response (stable after first turn) -// fallbackID: short hash without assistant (used to inherit binding from first turn) +// fallbackID preserves an earlier binding when a stronger body identifier appears +// later, and lets callers bind both identifiers when both are present. func extractSessionIDs(headers http.Header, payload []byte, metadata map[string]any) (string, string) { - // 1. metadata.user_id with Claude Code session format (highest priority) - if len(payload) > 0 { - userID := gjson.GetBytes(payload, "metadata.user_id").String() - if userID != "" { - // Old format: user_{hash}_account__session_{uuid} - if matches := sessionPattern.FindStringSubmatch(userID); len(matches) >= 2 { - id := "claude:" + matches[1] - return id, "" - } - // New format: JSON object with session_id field - // e.g. {"device_id":"...","account_uuid":"...","session_id":"uuid"} - if len(userID) > 0 && userID[0] == '{' { - if sid := gjson.Get(userID, "session_id").String(); sid != "" { - return "claude:" + sid, "" - } - } - } + if sid := sessionHeaderValue(headers, "X-Claude-Code-Session-Id"); sid != "" { + return "claude:" + sid, "" } - - // 2. X-Session-ID header - if sessionID := sessionHeaderValue(headers, "X-Session-ID"); sessionID != "" { - return "header:" + sessionID, "" + if sid := cliproxysession.ClaudeMetadataSessionID(payload); sid != "" { + return "claude:" + sid, "" } - - // 3. Session_id header (Codex) - if sessionID := sessionHeaderValue(headers, "Session-Id"); sessionID != "" { - return "codex:" + sessionID, "" + if sid := sessionHeaderValue(headers, "Session-Id"); sid != "" { + return "codex:" + sid, "" } - if sessionID := sessionHeaderValue(headers, "Session_id"); sessionID != "" { - return "codex:" + sessionID, "" + if sid := sessionHeaderValue(headers, "Session_id"); sid != "" { + return "codex:" + sid, "" } - - // 4. X-Client-Request-Id header (PI) - if requestID := sessionHeaderValue(headers, "X-Client-Request-Id"); requestID != "" { - return "clientreq:" + requestID, "" + if sid := sessionHeaderValue(headers, "X-Session-ID"); sid != "" { + return "header:" + sid, "" + } + if sid := sessionHeaderValue(headers, "X-Session-Affinity"); sid != "" { + return "affinity:" + sid, "" + } + if sid := sessionHeaderValue(headers, "X-Client-Request-Id"); sid != "" { + return "clientreq:" + sid, "" } if len(payload) > 0 { - // 5. Explicit request-body session fields. for _, path := range []string{"session_id", "sessionId"} { - if sessionID := strings.TrimSpace(gjson.GetBytes(payload, path).String()); sessionID != "" { - return "session:" + sessionID, "" + if sid := normalizedSessionCandidate(gjson.GetBytes(payload, path).String()); sid != "" { + return "session:" + sid, "" } } - if userID := strings.TrimSpace(gjson.GetBytes(payload, "metadata.user_id").String()); userID != "" { + + conversationID := "" + conversation := gjson.GetBytes(payload, "conversation") + if sid := normalizedSessionCandidate(conversation.Get("id").String()); sid != "" { + conversationID = "conv:" + sid + } else if conversation.Type == gjson.String { + if sid := normalizedSessionCandidate(conversation.String()); sid != "" { + conversationID = "conv:" + sid + } + } + if sid := normalizedSessionCandidate(gjson.GetBytes(payload, "prompt_cache_key").String()); sid != "" { + return "pck:" + sid, conversationID + } + if conversationID != "" { + return conversationID, "" + } + + if userID := normalizedSessionCandidate(gjson.GetBytes(payload, "metadata.user_id").String()); userID != "" { return "user:" + userID, "" } - if conversationID := strings.TrimSpace(gjson.GetBytes(payload, "conversation_id").String()); conversationID != "" { + if conversationID := normalizedSessionCandidate(gjson.GetBytes(payload, "conversation_id").String()); conversationID != "" { return "conv:" + conversationID, "" } - if promptCacheKey := strings.TrimSpace(gjson.GetBytes(payload, "prompt_cache_key").String()); promptCacheKey != "" { - return "prompt:" + promptCacheKey, "" - } } - // 6. Explicit long-lived execution session. if executionID, ok := metadata[cliproxyexecutor.ExecutionSessionMetadataKey].(string); ok { - if executionID = strings.TrimSpace(executionID); executionID != "" { + if executionID = normalizedSessionCandidate(executionID); executionID != "" { return "execution:" + executionID, "" } } - - // 7. Stable context-derived session identity. - if derivedID := cliproxysession.DerivedID(metadata); derivedID != "" { + if derivedID := normalizedSessionCandidate(cliproxysession.DerivedID(metadata)); derivedID != "" { return "derived:" + derivedID, "" } - if len(payload) == 0 { return "", "" } - - // 8. Legacy hash-based fallback from message content. return extractMessageHashIDs(payload) } -func sessionHeaderValue(headers http.Header, name string) string { - for key, values := range headers { - if !strings.EqualFold(key, name) { - continue - } - for _, value := range values { - if value = strings.TrimSpace(value); value != "" { - return value - } - } - } - return "" -} - func extractMessageHashIDs(payload []byte) (primaryID, fallbackID string) { var systemPrompt, firstUserMsg, firstAssistantMsg string diff --git a/sdk/cliproxy/auth/selector_test.go b/sdk/cliproxy/auth/selector_test.go index 51e87206d..cb1bfb40f 100644 --- a/sdk/cliproxy/auth/selector_test.go +++ b/sdk/cliproxy/auth/selector_test.go @@ -12,6 +12,8 @@ import ( "time" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + cliproxysession "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/session" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" ) func TestFillFirstSelectorPick_Deterministic(t *testing.T) { @@ -612,7 +614,7 @@ func TestSessionAffinitySelector_FailoverWhenAuthUnavailable(t *testing.T) { func TestExtractSessionID_ClaudeCodePriorityOverHeader(t *testing.T) { t.Parallel() - // Claude Code metadata.user_id should have highest priority, even when X-Session-ID header is present + // Claude Code metadata.user_id remains higher priority than a generic X-Session-ID header. headers := make(http.Header) headers.Set("X-Session-ID", "header-session-id") @@ -1028,6 +1030,227 @@ func TestSessionAffinitySelector_ThreeScenarios(t *testing.T) { }) } +func TestSessionAffinitySelectorBodyIdentifierTransitionsPreserveBinding(t *testing.T) { + t.Parallel() + + bothPayload := []byte(`{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`) + primaryID, fallbackID := extractSessionIDs(nil, bothPayload, nil) + if primaryID != "pck:shared-cache-bucket" || fallbackID != "conv:conversation-session" { + t.Fatalf("extractSessionIDs() = (%q, %q), want prompt-cache primary with conversation fallback", primaryID, fallbackID) + } + + for _, tt := range []struct { + name string + firstPayload []byte + }{ + {name: "prompt cache first", firstPayload: []byte(`{"prompt_cache_key":"shared-cache-bucket"}`)}, + {name: "conversation first", firstPayload: []byte(`{"conversation":{"id":"conversation-session"}}`)}, + } { + t.Run(tt.name, func(t *testing.T) { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + provider := "responses-transition-" + tt.name + + first, err := selector.Pick(context.Background(), provider, "gpt-test", cliproxyexecutor.Options{OriginalRequest: tt.firstPayload}, auths) + if err != nil { + t.Fatalf("first Pick() error = %v", err) + } + second, err := selector.Pick(context.Background(), provider, "gpt-test", cliproxyexecutor.Options{OriginalRequest: bothPayload}, auths) + if err != nil { + t.Fatalf("combined-identifier Pick() error = %v", err) + } + if second.ID != first.ID { + t.Fatalf("combined identifiers changed auth from %q to %q", first.ID, second.ID) + } + }) + } +} + +func TestSessionAffinitySelectorCombinedIdentifiersBindConversationFallback(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + provider := "responses-combined-to-conversation" + + combined := []byte(`{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`) + conversationOnly := []byte(`{"conversation":{"id":"conversation-session"}}`) + first, err := selector.Pick(context.Background(), provider, "gpt-test", cliproxyexecutor.Options{OriginalRequest: combined}, auths) + if err != nil { + t.Fatalf("combined-identifier Pick() error = %v", err) + } + second, err := selector.Pick(context.Background(), provider, "gpt-test", cliproxyexecutor.Options{OriginalRequest: conversationOnly}, auths) + if err != nil { + t.Fatalf("conversation-only Pick() error = %v", err) + } + if second.ID != first.ID { + t.Fatalf("dropping prompt_cache_key changed auth from %q to %q", first.ID, second.ID) + } +} + +func TestSessionAffinitySelectorPrimaryTrafficKeepsConversationAliasAlive(t *testing.T) { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + provider := "responses-active-primary-alias" + model := "gpt-test" + combined := []byte(`{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`) + promptOnly := []byte(`{"prompt_cache_key":"shared-cache-bucket"}`) + conversationOnly := []byte(`{"conversation":{"id":"conversation-session"}}`) + + first, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: combined}, auths) + if err != nil { + t.Fatalf("combined Pick() error = %v", err) + } + conversationKey := provider + "::conv:conversation-session::" + model + selector.cache.mu.Lock() + conversationEntry := selector.cache.entries[conversationKey] + conversationEntry.expiresAt = time.Now().Add(-time.Second) + selector.cache.entries[conversationKey] = conversationEntry + selector.cache.mu.Unlock() + + primary, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: promptOnly}, auths) + if err != nil { + t.Fatalf("prompt-only Pick() error = %v", err) + } + if primary.ID != first.ID { + t.Fatalf("prompt-only auth = %q, want %q", primary.ID, first.ID) + } + fallback, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: conversationOnly}, auths) + if err != nil { + t.Fatalf("conversation-only Pick() error = %v", err) + } + if fallback.ID != first.ID { + t.Fatalf("conversation alias expired during active primary traffic: got %q, want %q", fallback.ID, first.ID) + } +} + +func TestSessionAffinitySelectorSharedPromptKeyPreservesConversationAliases(t *testing.T) { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + provider := "responses-shared-prompt-key" + model := "gpt-test" + + combinedA := []byte(`{"conversation":{"id":"conversation-a"},"prompt_cache_key":"shared-cache-bucket"}`) + combinedB := []byte(`{"conversation":{"id":"conversation-b"},"prompt_cache_key":"shared-cache-bucket"}`) + conversationA := []byte(`{"conversation":{"id":"conversation-a"}}`) + conversationB := []byte(`{"conversation":{"id":"conversation-b"}}`) + + first, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: combinedA}, auths) + if err != nil { + t.Fatalf("conversation A combined Pick() error = %v", err) + } + second, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: combinedB}, auths) + if err != nil { + t.Fatalf("conversation B combined Pick() error = %v", err) + } + if second.ID != first.ID { + t.Fatalf("shared prompt key changed auth from %q to %q", first.ID, second.ID) + } + for name, payload := range map[string][]byte{"conversation A": conversationA, "conversation B": conversationB} { + picked, errPick := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: payload}, auths) + if errPick != nil { + t.Fatalf("%s Pick() error = %v", name, errPick) + } + if picked.ID != first.ID { + t.Fatalf("%s alias selected %q, want %q", name, picked.ID, first.ID) + } + } +} + +func TestSessionAffinitySelectorConversationIDContainingPromptMarkerRemainsStable(t *testing.T) { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + provider := "responses-opaque-conversation" + model := "gpt-test" + combined := []byte(`{"conversation":{"id":"a::pck:b"},"prompt_cache_key":"shared-cache-bucket"}`) + conversationOnly := []byte(`{"conversation":{"id":"a::pck:b"}}`) + + first, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: combined}, auths) + if err != nil { + t.Fatalf("combined Pick() error = %v", err) + } + second, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: conversationOnly}, auths) + if err != nil { + t.Fatalf("conversation-only Pick() error = %v", err) + } + if second.ID != first.ID { + t.Fatalf("opaque conversation alias selected %q, want %q", second.ID, first.ID) + } +} + +func TestSessionCacheSharedPromptKeyCapsStableAliasesByRecency(t *testing.T) { + cache := NewSessionCache(time.Minute) + defer cache.Stop() + const promptKey = "openai::pck:shared-cache-bucket::gpt-test" + for index := 0; index < 128; index++ { + conversation := fmt.Sprintf("openai::conv:conversation-%03d::gpt-test", index) + cache.SetAliases("auth-a", promptKey, conversation) + } + + cache.mu.RLock() + defer cache.mu.RUnlock() + if len(cache.entries) > 65 { + t.Fatalf("cache entries = %d, want one prompt key plus at most 64 stable aliases", len(cache.entries)) + } + if _, ok := cache.entries["openai::conv:conversation-127::gpt-test"]; !ok { + t.Fatal("newest conversation alias was not retained") + } + if _, ok := cache.entries["openai::conv:conversation-000::gpt-test"]; ok { + t.Fatal("oldest conversation alias was retained after stable-alias cap") + } +} + +func TestSessionCacheRotatingPrimaryEvictsObsoleteAliases(t *testing.T) { + cache := NewSessionCache(time.Minute) + defer cache.Stop() + + const fallback = "openai::conv:conversation-session::gpt-test" + for index := 0; index < 16; index++ { + primary := fmt.Sprintf("openai::pck:cache-%02d::gpt-test", index) + cache.SetAliases("auth-a", primary, fallback) + } + latest := "openai::pck:cache-15::gpt-test" + oldest := "openai::pck:cache-00::gpt-test" + + cache.mu.RLock() + defer cache.mu.RUnlock() + if len(cache.entries) != 2 { + t.Fatalf("cache entries = %d, want only latest primary and fallback", len(cache.entries)) + } + if _, ok := cache.entries[latest]; !ok { + t.Fatalf("latest primary %q was not retained", latest) + } + if _, ok := cache.entries[fallback]; !ok { + t.Fatalf("fallback %q was not retained", fallback) + } + if _, ok := cache.entries[oldest]; ok { + t.Fatalf("obsolete primary %q was retained", oldest) + } + if aliases := cache.entries[fallback].aliases; len(aliases) != 2 { + t.Fatalf("fallback alias group = %#v, want exactly two active identifiers", aliases) + } +} + func TestSessionAffinitySelector_MultiModelSession(t *testing.T) { t.Parallel() @@ -1314,3 +1537,248 @@ func TestSessionAffinitySelector_Concurrent(t *testing.T) { default: } } + +func TestExtractSessionIDNativeSignals(t *testing.T) { + t.Parallel() + tests := []struct { + name string + headers http.Header + payload string + want string + }{ + { + name: "claude code header", + headers: http.Header{"X-Claude-Code-Session-Id": []string{"claude-session"}}, + want: "claude:claude-session", + }, + { + name: "lowercase claude code header", + headers: http.Header{"x-claude-code-session-id": []string{"lowercase-session"}}, + want: "claude:lowercase-session", + }, + { + name: "codex hyphen header", + headers: http.Header{"Session-Id": []string{"codex-session"}}, + want: "codex:codex-session", + }, + { + name: "codex underscore header", + headers: http.Header{"Session_id": []string{"legacy-codex-session"}}, + want: "codex:legacy-codex-session", + }, + { + name: "open code session affinity", + headers: http.Header{"X-Session-Affinity": []string{"ses_opencode"}}, + want: "affinity:ses_opencode", + }, + { + name: "prompt cache key", + payload: `{"prompt_cache_key":"prompt-session"}`, + want: "pck:prompt-session", + }, + { + name: "responses conversation object", + payload: `{"conversation":{"id":"conv-object"}}`, + want: "conv:conv-object", + }, + { + name: "responses conversation string", + payload: `{"conversation":"conv-string"}`, + want: "conv:conv-string", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := ExtractSessionID(tt.headers, []byte(tt.payload), nil); got != tt.want { + t.Fatalf("ExtractSessionID() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestExtractSessionIDNativeSignalPriority(t *testing.T) { + t.Parallel() + tests := []struct { + name string + headers http.Header + payload string + want string + }{ + { + name: "claude header beats metadata", + headers: http.Header{ + "X-Claude-Code-Session-Id": []string{"header-session"}, + }, + payload: `{"metadata":{"user_id":"user_hash_account__session_22222222-2222-4222-8222-222222222222"}}`, + want: "claude:header-session", + }, + { + name: "claude metadata beats codex header", + headers: http.Header{ + "Session-Id": []string{"codex-session"}, + }, + payload: `{"metadata":{"user_id":"user_hash_account__session_22222222-2222-4222-8222-222222222222"}}`, + want: "claude:22222222-2222-4222-8222-222222222222", + }, + { + name: "codex header beats x session id and prompt key", + headers: http.Header{ + "Session-Id": []string{"codex-session"}, + "X-Session-Id": []string{"generic-session"}, + }, + payload: `{"prompt_cache_key":"prompt-session"}`, + want: "codex:codex-session", + }, + { + name: "x session id beats affinity", + headers: http.Header{ + "X-Session-Id": []string{"generic-session"}, + "X-Session-Affinity": []string{"affinity-session"}, + }, + want: "header:generic-session", + }, + { + name: "prompt cache key beats conversation id", + payload: `{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`, + want: "pck:shared-cache-bucket", + }, + { + name: "client request id beats body fallbacks", + headers: http.Header{ + "X-Client-Request-Id": []string{"client-session"}, + }, + payload: `{"prompt_cache_key":"prompt-session","conversation":{"id":"conversation-session"}}`, + want: "clientreq:client-session", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := ExtractSessionID(tt.headers, []byte(tt.payload), nil); got != tt.want { + t.Fatalf("ExtractSessionID() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestExtractSessionIDRejectsInvalidExplicitSignals(t *testing.T) { + t.Parallel() + tooLong := strings.Repeat("a", 257) + tests := []struct { + name string + headers http.Header + payload string + want string + }{ + { + name: "whitespace", + headers: http.Header{"X-Claude-Code-Session-Id": []string{" "}}, + want: "", + }, + { + name: "newline", + headers: http.Header{"X-Session-Id": []string{"bad\nsession"}}, + want: "", + }, + { + name: "control character", + headers: http.Header{"Session-Id": []string{"bad\x00session"}}, + want: "", + }, + { + name: "too long", + headers: http.Header{"X-Client-Request-Id": []string{tooLong}}, + want: "", + }, + { + name: "invalid stronger signal falls through", + headers: http.Header{ + "X-Claude-Code-Session-Id": []string{"bad\nsession"}, + "Session-Id": []string{"valid-codex"}, + }, + want: "codex:valid-codex", + }, + { + name: "invalid prompt key falls through to conversation", + payload: `{"prompt_cache_key":" ","conversation":{"id":"valid-conversation"}}`, + want: "conv:valid-conversation", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := ExtractSessionID(tt.headers, []byte(tt.payload), nil); got != tt.want { + t.Fatalf("ExtractSessionID() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestExtractSessionIDClaudeMetadataParsesBeforeBoundingSessionID(t *testing.T) { + t.Parallel() + const sessionID = "11111111-1111-4111-8111-111111111111" + metadata := map[string]string{ + "device_id": strings.Repeat("d", 64), + "account_uuid": "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa", + "session_id": sessionID, + "organization_uuid": "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb", + "email": "user@example.com", + } + + for _, tt := range []struct { + name string + encode func(any) ([]byte, error) + }{ + {name: "rich compact json", encode: json.Marshal}, + {name: "pretty printed json", encode: func(v any) ([]byte, error) { return json.MarshalIndent(v, "", " ") }}, + } { + t.Run(tt.name, func(t *testing.T) { + userID, errMarshal := tt.encode(metadata) + if errMarshal != nil { + t.Fatalf("marshal metadata: %v", errMarshal) + } + payload, errPayload := json.Marshal(map[string]any{ + "metadata": map[string]string{"user_id": string(userID)}, + }) + if errPayload != nil { + t.Fatalf("marshal payload: %v", errPayload) + } + if got := ExtractSessionID(nil, payload, nil); got != "claude:"+sessionID { + t.Fatalf("ExtractSessionID() = %q, want %q", got, "claude:"+sessionID) + } + }) + } +} + +func TestSessionAffinitySelectorUsesRequestPayloadWhenOriginalRequestMissing(t *testing.T) { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + request := cliproxyexecutor.Request{ + Model: "gpt-test", + Payload: []byte(`{"conversation":{"id":"request-only-conversation"},"input":"hello"}`), + } + _, opts := cliproxysession.Enrich(request, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + }) + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + + first, errFirst := selector.Pick(context.Background(), "openai", request.Model, opts, auths) + if errFirst != nil { + t.Fatalf("first Pick() error = %v", errFirst) + } + second, errSecond := selector.Pick(context.Background(), "openai", request.Model, opts, auths) + if errSecond != nil { + t.Fatalf("second Pick() error = %v", errSecond) + } + if second.ID != first.ID { + t.Fatalf("request-only conversation changed auth from %q to %q", first.ID, second.ID) + } +} diff --git a/sdk/cliproxy/auth/session_cache.go b/sdk/cliproxy/auth/session_cache.go index a812e581b..54bf3867b 100644 --- a/sdk/cliproxy/auth/session_cache.go +++ b/sdk/cliproxy/auth/session_cache.go @@ -1,14 +1,18 @@ package auth import ( + "strings" "sync" "time" ) -// sessionEntry stores auth binding with expiration. +const maxStableSessionAliases = 64 + +// sessionEntry stores an auth binding, its identifier aliases, and expiration. type sessionEntry struct { authID string expiresAt time.Time + aliases []string } // SessionCache provides TTL-based session to auth mapping with automatic cleanup. @@ -40,66 +44,212 @@ func (c *SessionCache) Get(sessionID string) (string, bool) { if sessionID == "" { return "", false } + now := time.Now() c.mu.RLock() entry, ok := c.entries[sessionID] + if ok && now.Before(entry.expiresAt) { + c.mu.RUnlock() + return entry.authID, true + } c.mu.RUnlock() if !ok { return "", false } - if time.Now().After(entry.expiresAt) { - c.mu.Lock() - delete(c.entries, sessionID) - c.mu.Unlock() + + c.mu.Lock() + defer c.mu.Unlock() + entry, ok = c.entries[sessionID] + if !ok { return "", false } - return entry.authID, true + if time.Now().Before(entry.expiresAt) { + return entry.authID, true + } + c.removeAliasGroupLocked(entry) + return "", false } -// GetAndRefresh retrieves the auth ID bound to a session and refreshes TTL on hit. -// This extends the binding lifetime for active sessions. +// GetAndRefresh retrieves the auth ID bound to a session and refreshes the TTL +// for every identifier known to represent the same logical session. func (c *SessionCache) GetAndRefresh(sessionID string) (string, bool) { if sessionID == "" { return "", false } now := time.Now() c.mu.Lock() + defer c.mu.Unlock() entry, ok := c.entries[sessionID] if !ok { - c.mu.Unlock() return "", false } - if now.After(entry.expiresAt) { - delete(c.entries, sessionID) - c.mu.Unlock() + if !now.Before(entry.expiresAt) { + c.removeAliasGroupLocked(entry) return "", false } - // Refresh TTL on successful access - entry.expiresAt = now.Add(c.ttl) - c.entries[sessionID] = entry - c.mu.Unlock() + + aliases := compactSessionAliases(mergeSessionAliases([]string{sessionID}, entry.aliases...)) + c.replaceAliasGroupsLocked(entry.authID, now.Add(c.ttl), aliases, entry) return entry.authID, true } -// Set binds a session to an auth ID with TTL refresh. +// Set binds a session to an auth ID with TTL refresh. Existing aliases for the +// same logical session remain attached when the binding is refreshed or moved. func (c *SessionCache) Set(sessionID, authID string) { - if sessionID == "" || authID == "" { - return - } - c.mu.Lock() - c.entries[sessionID] = sessionEntry{ - authID: authID, - expiresAt: time.Now().Add(c.ttl), - } - c.mu.Unlock() + c.SetAliases(authID, sessionID) } -// Invalidate removes a specific session binding. +// SetAliases binds multiple identifiers for one logical session to an auth ID. +func (c *SessionCache) SetAliases(authID string, sessionIDs ...string) { + if authID == "" { + return + } + now := time.Now() + c.mu.Lock() + defer c.mu.Unlock() + + aliases := mergeSessionAliases(nil, sessionIDs...) + previousGroups := make([]sessionEntry, 0, len(sessionIDs)) + for _, sessionID := range sessionIDs { + entry, ok := c.entries[sessionID] + if !ok { + continue + } + if !now.Before(entry.expiresAt) { + c.removeAliasGroupLocked(entry) + continue + } + previousGroups = append(previousGroups, entry) + aliases = mergeSessionAliases(aliases, entry.aliases...) + } + aliases = compactSessionAliases(aliases) + if len(aliases) == 0 { + return + } + c.replaceAliasGroupsLocked(authID, now.Add(c.ttl), aliases, previousGroups...) +} + +func (c *SessionCache) replaceAliasGroupsLocked(authID string, expiresAt time.Time, aliases []string, previousGroups ...sessionEntry) { + for _, previous := range previousGroups { + c.removeAliasGroupLocked(previous) + } + entry := sessionEntry{authID: authID, expiresAt: expiresAt, aliases: aliases} + for _, alias := range aliases { + c.entries[alias] = entry + } +} + +func (c *SessionCache) removeAliasGroupLocked(entry sessionEntry) { + for _, alias := range entry.aliases { + current, ok := c.entries[alias] + if !ok || current.authID != entry.authID || !current.expiresAt.Equal(entry.expiresAt) || + !equalSessionAliases(current.aliases, entry.aliases) { + continue + } + delete(c.entries, alias) + } +} + +func compactSessionAliases(aliases []string) []string { + return compactSessionAliasesWith(aliases, isLocalPromptCacheSessionAlias) +} + +func compactHomeSessionAliases(aliases []string) []string { + return compactSessionAliasesWith(aliases, func(alias string) bool { + return strings.HasPrefix(alias, "pck:") + }) +} + +func compactSessionAliasesWith(aliases []string, isPromptCacheAlias func(string) bool) []string { + compacted := make([]string, 0, len(aliases)) + hasPromptCacheKey := false + stableAliases := 0 + for _, alias := range aliases { + if isPromptCacheAlias(alias) { + if hasPromptCacheKey { + continue + } + hasPromptCacheKey = true + } else { + if stableAliases >= maxStableSessionAliases { + continue + } + stableAliases++ + } + compacted = append(compacted, alias) + } + return compacted +} + +func isLocalPromptCacheSessionAlias(alias string) bool { + if strings.HasPrefix(alias, "pck:") { + return true + } + _, sessionAndModel, ok := strings.Cut(alias, "::") + return ok && strings.HasPrefix(sessionAndModel, "pck:") +} + +func equalSessionAliases(left, right []string) bool { + if len(left) != len(right) { + return false + } + for index := range left { + if left[index] != right[index] { + return false + } + } + return true +} + +func mergeSessionAliases(existing []string, candidates ...string) []string { + aliases := make([]string, 0, len(existing)+len(candidates)) + seen := make(map[string]struct{}, cap(aliases)) + add := func(alias string) { + if alias == "" { + return + } + if _, ok := seen[alias]; ok { + return + } + seen[alias] = struct{}{} + aliases = append(aliases, alias) + } + for _, alias := range existing { + add(alias) + } + for _, alias := range candidates { + add(alias) + } + return aliases +} + +// Invalidate removes a specific session binding without allowing another alias +// in the same group to recreate it on its next refresh. func (c *SessionCache) Invalidate(sessionID string) { if sessionID == "" { return } c.mu.Lock() + entry, ok := c.entries[sessionID] delete(c.entries, sessionID) + if ok { + for _, alias := range entry.aliases { + if alias == sessionID { + continue + } + current, exists := c.entries[alias] + if !exists || current.authID != entry.authID { + continue + } + filtered := make([]string, 0, len(current.aliases)) + for _, candidate := range current.aliases { + if candidate != sessionID { + filtered = append(filtered, candidate) + } + } + current.aliases = filtered + c.entries[alias] = current + } + } c.mu.Unlock() } @@ -144,7 +294,7 @@ func (c *SessionCache) cleanup() { now := time.Now() c.mu.Lock() for sid, entry := range c.entries { - if now.After(entry.expiresAt) { + if !now.Before(entry.expiresAt) { delete(c.entries, sid) } } diff --git a/sdk/cliproxy/session/identity.go b/sdk/cliproxy/session/identity.go index 3d6c9fffe..696e0372a 100644 --- a/sdk/cliproxy/session/identity.go +++ b/sdk/cliproxy/session/identity.go @@ -2,11 +2,14 @@ package session import ( + "bytes" "crypto/sha256" "encoding/hex" "encoding/json" "fmt" + "regexp" "strings" + "unicode" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" @@ -19,6 +22,8 @@ const ( instructionRuneLimit = 50 ) +var legacyClaudeSessionPattern = regexp.MustCompile(`_session_([a-f0-9-]+)$`) + type canonicalRoot struct { Version string `json:"version"` Format string `json:"format"` @@ -34,6 +39,41 @@ type canonicalPart struct { Value string `json:"value"` } +// NormalizeExplicitID validates an explicit client-provided session identifier. +// It preserves opaque printable values while rejecting oversized or control-bearing IDs. +func NormalizeExplicitID(raw string) string { + for _, r := range raw { + if unicode.IsControl(r) { + return "" + } + } + raw = strings.TrimSpace(raw) + if raw == "" || len(raw) > 256 { + return "" + } + return raw +} + +// ClaudeMetadataSessionID extracts the explicit Claude Code session from +// current JSON metadata or the legacy user_id suffix before bounding the +// surrounding metadata container. +func ClaudeMetadataSessionID(payload []byte) string { + if len(payload) == 0 { + return "" + } + userID := strings.TrimSpace(gjson.GetBytes(payload, "metadata.user_id").String()) + if userID == "" { + return "" + } + if strings.HasPrefix(userID, "{") { + return NormalizeExplicitID(gjson.Get(userID, "session_id").String()) + } + if matches := legacyClaudeSessionPattern.FindStringSubmatch(userID); len(matches) >= 2 { + return NormalizeExplicitID(matches[1]) + } + return "" +} + // CallerScope returns an irreversible namespace for a downstream caller credential. func CallerScope(value string) string { value = strings.TrimSpace(value) @@ -56,24 +96,26 @@ func DerivedID(metadata map[string]any) string { // Enrich derives a session identity once and places it in both request and option metadata. func Enrich(req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Request, cliproxyexecutor.Options) { payload := opts.OriginalRequest - if len(payload) == 0 { - payload = req.Payload + if len(payload) == 0 && len(req.Payload) > 0 { + opts.OriginalRequest = bytes.Clone(req.Payload) + payload = opts.OriginalRequest } - if executionID := firstMetadataString(cliproxyexecutor.ExecutionSessionMetadataKey, opts.Metadata, req.Metadata); executionID != "" { + if executionID := firstNormalizedMetadataID(cliproxyexecutor.ExecutionSessionMetadataKey, opts.Metadata, req.Metadata); executionID != "" { req.Metadata = metadataWithValue(metadataWithoutKey(req.Metadata, cliproxyexecutor.DerivedSessionIDMetadataKey), cliproxyexecutor.ExecutionSessionMetadataKey, executionID) opts.Metadata = metadataWithValue(metadataWithoutKey(opts.Metadata, cliproxyexecutor.DerivedSessionIDMetadataKey), cliproxyexecutor.ExecutionSessionMetadataKey, executionID) return req, opts } + req.Metadata = metadataWithoutKey(req.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey) + opts.Metadata = metadataWithoutKey(opts.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey) if hasExplicitSession(opts.Headers, payload) { req.Metadata = metadataWithoutKey(req.Metadata, cliproxyexecutor.DerivedSessionIDMetadataKey) opts.Metadata = metadataWithoutKey(opts.Metadata, cliproxyexecutor.DerivedSessionIDMetadataKey) return req, opts } - derivedID := DerivedID(opts.Metadata) - if derivedID == "" { - derivedID = DerivedID(req.Metadata) - } + derivedID := firstNormalizedMetadataID(cliproxyexecutor.DerivedSessionIDMetadataKey, opts.Metadata, req.Metadata) + req.Metadata = metadataWithoutKey(req.Metadata, cliproxyexecutor.DerivedSessionIDMetadataKey) + opts.Metadata = metadataWithoutKey(opts.Metadata, cliproxyexecutor.DerivedSessionIDMetadataKey) if derivedID == "" { callerScope := metadataString(opts.Metadata, cliproxyexecutor.CallerScopeMetadataKey) if callerScope == "" { @@ -90,8 +132,8 @@ func Enrich(req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (clipro } func hasExplicitSession(headers map[string][]string, payload []byte) bool { - for _, header := range []string{"X-Session-ID", "Session-Id", "Session_id", "X-Client-Request-Id"} { - if strings.TrimSpace(headerValue(headers, header)) != "" { + for _, header := range []string{"X-Claude-Code-Session-Id", "X-Session-ID", "Session-Id", "Session_id", "X-Session-Affinity", "X-Client-Request-Id"} { + if NormalizeExplicitID(headerValue(headers, header)) != "" { return true } } @@ -99,20 +141,35 @@ func hasExplicitSession(headers map[string][]string, payload []byte) bool { return false } root := gjson.ParseBytes(payload) - for _, path := range []string{"metadata.user_id", "session_id", "sessionId", "conversation_id", "prompt_cache_key"} { - if strings.TrimSpace(root.Get(path).String()) != "" { + for _, path := range []string{"session_id", "sessionId", "conversation_id", "prompt_cache_key"} { + if NormalizeExplicitID(root.Get(path).String()) != "" { return true } } - return false + if ClaudeMetadataSessionID(payload) != "" { + return true + } + userID := strings.TrimSpace(root.Get("metadata.user_id").String()) + if NormalizeExplicitID(userID) != "" { + return true + } + conversation := root.Get("conversation") + if NormalizeExplicitID(conversation.Get("id").String()) != "" { + return true + } + return conversation.Type == gjson.String && NormalizeExplicitID(conversation.String()) != "" } func headerValue(headers map[string][]string, name string) string { for key, values := range headers { - if !strings.EqualFold(key, name) || len(values) == 0 { + if !strings.EqualFold(key, name) { continue } - return values[0] + for _, value := range values { + if normalized := NormalizeExplicitID(value); normalized != "" { + return normalized + } + } } return "" } @@ -468,6 +525,22 @@ func metadataWithoutKey(metadata map[string]any, key string) map[string]any { return cloned } +func firstNormalizedMetadataID(key string, metadataSets ...map[string]any) string { + for _, metadata := range metadataSets { + if metadata == nil { + continue + } + raw, ok := metadata[key].(string) + if !ok { + continue + } + if normalized := NormalizeExplicitID(raw); normalized != "" { + return normalized + } + } + return "" +} + func firstMetadataString(key string, metadataSets ...map[string]any) string { for _, metadata := range metadataSets { if value := metadataString(metadata, key); value != "" { diff --git a/sdk/cliproxy/session/identity_test.go b/sdk/cliproxy/session/identity_test.go index 1fbc498cb..7211d3798 100644 --- a/sdk/cliproxy/session/identity_test.go +++ b/sdk/cliproxy/session/identity_test.go @@ -1,6 +1,7 @@ package session import ( + "bytes" "net/http" "strings" "testing" @@ -134,10 +135,42 @@ func TestEnrichSkipsDerivationForExplicitSessions(t *testing.T) { payload: []byte(`not-json`), headers: http.Header{"X-Session-ID": []string{"header-session"}}, }, + { + name: "Claude Code session header", + payload: []byte(`{"messages":[{"role":"user","content":"hello"}]}`), + headers: http.Header{"X-Claude-Code-Session-Id": []string{"claude-session"}}, + }, + { + name: "later valid multi-value session header", + payload: []byte(`{"messages":[{"role":"user","content":"hello"}]}`), + headers: http.Header{"X-Session-Affinity": []string{"", "later-valid-session"}}, + }, + { + name: "OpenCode affinity header", + payload: []byte(`{"messages":[{"role":"user","content":"hello"}]}`), + headers: http.Header{"X-Session-Affinity": []string{"opencode-session"}}, + }, + { + name: "Responses conversation object", + payload: []byte(`{"conversation":{"id":"conversation-session"},"messages":[{"role":"user","content":"hello"}]}`), + }, + { + name: "Responses conversation string", + payload: []byte(`{"conversation":"conversation-session","messages":[{"role":"user","content":"hello"}]}`), + }, { name: "metadata user id", payload: []byte(`{"metadata":{"user_id":"explicit-user"},"messages":[{"role":"user","content":"hello"}]}`), }, + { + name: "long legacy Claude metadata session", + payload: []byte(`{"metadata":{"user_id":"` + strings.Repeat("x", 300) + + `_session_ac980658-63bd-4fb3-97ba-8da64cb1e344"},"messages":[{"role":"user","content":"hello"}]}`), + }, + { + name: "JSON metadata user id without nested session", + payload: []byte(`{"metadata":{"user_id":"{\"device_id\":\"abc123\"}"},"messages":[{"role":"user","content":"hello"}]}`), + }, { name: "body session id", payload: []byte(`{"session_id":"body-session","messages":[{"role":"user","content":"hello"}]}`), @@ -196,6 +229,83 @@ func TestEnrichSkipsDerivationForExplicitSessions(t *testing.T) { } } +func TestEnrichDerivesAfterInvalidSessionIdentity(t *testing.T) { + t.Parallel() + + baseMessages := `"input":"hello"` + tests := []struct { + name string + payload []byte + headers http.Header + requestMetadata map[string]any + optionMetadata map[string]any + }{ + { + name: "oversized prompt cache key", + payload: []byte(`{"prompt_cache_key":"` + strings.Repeat("x", 257) + `",` + baseMessages + `}`), + }, + { + name: "trailing control character prompt cache key", + payload: []byte(`{"prompt_cache_key":"tenant\n",` + baseMessages + `}`), + }, + { + name: "leading control character prompt cache key", + payload: []byte(`{"prompt_cache_key":"\ttenant",` + baseMessages + `}`), + }, + { + name: "control character session header", + payload: []byte(`{` + baseMessages + `}`), + headers: http.Header{"X-Session-Affinity": []string{"bad\nsession"}}, + }, + { + name: "oversized execution session option metadata", + payload: []byte(`{"input":"hello"}`), + optionMetadata: map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: strings.Repeat("x", 257)}, + }, + { + name: "control character execution session request metadata", + payload: []byte(`{"input":"hello"}`), + requestMetadata: map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: "bad\nsession"}, + }, + { + name: "oversized retained derived session option metadata", + payload: []byte(`{"input":"hello"}`), + optionMetadata: map[string]any{cliproxyexecutor.DerivedSessionIDMetadataKey: strings.Repeat("x", 257)}, + }, + { + name: "control character retained derived session request metadata", + payload: []byte(`{"input":"hello"}`), + requestMetadata: map[string]any{cliproxyexecutor.DerivedSessionIDMetadataKey: "bad\nsession"}, + }, + } + for _, test := range tests { + test := test + t.Run(test.name, func(t *testing.T) { + t.Parallel() + req := cliproxyexecutor.Request{Payload: test.payload, Metadata: test.requestMetadata} + opts := cliproxyexecutor.Options{ + OriginalRequest: test.payload, + SourceFormat: sdktranslator.FormatOpenAIResponse, + Headers: test.headers, + Metadata: test.optionMetadata, + } + enrichedReq, enrichedOpts := Enrich(req, opts) + requestID := DerivedID(enrichedReq.Metadata) + optionsID := DerivedID(enrichedOpts.Metadata) + wantID := DeriveID(sdktranslator.FormatOpenAIResponse, test.payload, "") + if requestID != wantID || optionsID != wantID { + t.Fatalf("derived identities = request:%q options:%q, want %q", requestID, optionsID, wantID) + } + if got := metadataString(enrichedReq.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); got != "" { + t.Fatalf("request execution session = %q, want invalid value removed", got) + } + if got := metadataString(enrichedOpts.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); got != "" { + t.Fatalf("options execution session = %q, want invalid value removed", got) + } + }) + } +} + func TestEnrichCopiesDerivedIdentityToRequestAndOptions(t *testing.T) { t.Parallel() @@ -216,3 +326,23 @@ func TestEnrichCopiesDerivedIdentityToRequestAndOptions(t *testing.T) { t.Fatal("Enrich() mutated original request metadata") } } + +func TestEnrichCarriesRequestPayloadIntoSelectionOptions(t *testing.T) { + t.Parallel() + + payload := []byte(`{"conversation":{"id":"request-only-conversation"},"input":"hello"}`) + _, enrichedOpts := Enrich( + cliproxyexecutor.Request{Payload: payload}, + cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}, + ) + + if !bytes.Equal(enrichedOpts.OriginalRequest, payload) { + t.Fatalf("OriginalRequest = %q, want request payload %q", enrichedOpts.OriginalRequest, payload) + } + if len(enrichedOpts.OriginalRequest) > 0 && &enrichedOpts.OriginalRequest[0] == &payload[0] { + t.Fatal("OriginalRequest aliases Request.Payload instead of preserving a snapshot") + } + if got := DerivedID(enrichedOpts.Metadata); got != "" { + t.Fatalf("DerivedSessionID = %q, want explicit conversation to remain authoritative", got) + } +}