From e92014a27b9aa06ae907508d0c2987136b01a18d Mon Sep 17 00:00:00 2001 From: warelik Date: Fri, 21 Aug 2026 21:49:30 +0300 Subject: [PATCH] fix(auth): retain session affinity on transient errors Port upstream #5109: primary session binding is now retained across 5xx/429/408 and cloudflare challenges, while terminal auth rejections (401/402/403/404, invalid_grant, unsupported model) still release it. A sticky temporary fallback cache keeps failover from flapping while the primary credential cools down, then returns to the warm cache on recovery. This flips the Plus-pinned release-on-failure behavior and removes the session quarantine mechanism that is no longer needed. --- sdk/cliproxy/auth/selector.go | 303 +++--- sdk/cliproxy/auth/selector_review_p2_test.go | 72 -- sdk/cliproxy/auth/selector_test.go | 91 +- .../auth/session_affinity_fix_test.go | 288 ++---- .../auth/session_affinity_metadata_test.go | 6 +- .../auth/session_affinity_priority_test.go | 16 +- .../auth/session_affinity_retention_test.go | 862 ++++++++++++++++++ sdk/cliproxy/auth/session_cache.go | 17 + 8 files changed, 1194 insertions(+), 461 deletions(-) create mode 100644 sdk/cliproxy/auth/session_affinity_retention_test.go diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index b8a4ebe14..a254c0a2b 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -627,10 +627,9 @@ func availabilityBlock(unavailable, quotaExceeded bool, nextRetryAfter, nextReco // It extracts session ID from multiple sources and maintains session-to-auth // mappings with automatic failover when the bound auth becomes unavailable. type SessionAffinitySelector struct { - fallback Selector - cache *SessionCache - quarantine *SessionCache - bindMu sync.Mutex + fallback Selector + cache *SessionCache + fallbackCache *SessionCache } // SessionAffinityConfig configures the session affinity selector. @@ -656,9 +655,9 @@ func NewSessionAffinitySelectorWithConfig(cfg SessionAffinityConfig) *SessionAff cfg.TTL = time.Hour } return &SessionAffinitySelector{ - fallback: cfg.Fallback, - cache: NewSessionCache(cfg.TTL), - quarantine: NewSessionCache(cfg.TTL), + fallback: cfg.Fallback, + cache: NewSessionCache(cfg.TTL), + fallbackCache: NewSessionCache(cfg.TTL), } } @@ -681,7 +680,6 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri } opts.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey] = provider opts.Metadata[cliproxyexecutor.SessionAffinityModelMetadataKey] = model - primaryID, fallbackID := extractSessionIDs(opts.Headers, opts.OriginalRequest, opts.Metadata) now := time.Now() excluded := extractExcludedAuthIDs(opts.Metadata) @@ -711,7 +709,7 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if fallbackID != "" && fallbackID != primaryID { fallbackKey = provider + "::" + fallbackID + "::" + modelKey } - available = s.excludeSessionQuarantine(cacheKey, fallbackKey, available) + fallbackAuths := highestPriorityAuths(available) bind := func(authID string) { if fallbackKey != "" { @@ -721,101 +719,112 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri s.cache.Set(cacheKey, authID) } - // Fast path outside bindMu: reuse valid cached binding without holding bindMu. - 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 + collectTempFallbackKeys := func() []string { + keys := []string{cacheKey} + if fallbackKey != "" { + keys = append(keys, fallbackKey) + } + if aliases := s.cache.Aliases(cacheKey); len(aliases) > 0 { + for _, alias := range aliases { + if alias != "" { + keys = append(keys, alias) + } } } - } else if fallbackKey != "" { - if cachedAuthID, ok := s.cache.Get(fallbackKey); 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 + if fallbackKey != "" { + if aliases := s.cache.Aliases(fallbackKey); len(aliases) > 0 { + for _, alias := range aliases { + if alias != "" { + keys = append(keys, alias) + } } } } + return keys + } + bindTempFallback := func(authID string) { + if s.fallbackCache != nil { + s.fallbackCache.SetAliases(authID, collectTempFallbackKeys()...) + } + } + invalidateTempFallback := func() { + if s.fallbackCache != nil { + for _, key := range collectTempFallbackKeys() { + s.fallbackCache.Invalidate(key) + } + } + } + getTempFallbackAuth := func() (*Auth, bool) { + if s.fallbackCache == nil { + return nil, false + } + for _, key := range collectTempFallbackKeys() { + if tempAuthID, ok := s.fallbackCache.GetAndRefresh(key); ok { + for _, auth := range available { + if auth.ID == tempAuthID { + return auth, true + } + } + } + } + return nil, false } - s.bindMu.Lock() - defer s.bindMu.Unlock() - - // Under bindMu, re-check if a concurrent request refreshed or rebound the session. if cachedAuthID, ok := s.cache.GetAndRefresh(cacheKey); ok { for _, auth := range available { if auth.ID == cachedAuthID { - entry.Infof("session-affinity: concurrent cache hit | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) + invalidateTempFallback() 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 } } - } else if fallbackKey != "" { + // Primary cached auth is unavailable (cooling down). + // Check for an active sticky temporary fallback binding: + if fallbackAuth, ok := getTempFallbackAuth(); ok { + entry.Infof("session-affinity: sticky fallback cache hit | session=%s primary_cooling=%s fallback_auth=%s provider=%s model=%s", truncateSessionID(primaryID), cachedAuthID, fallbackAuth.ID, provider, model) + return fallbackAuth, nil + } + // Reselect fallback auth and record it as the sticky temporary fallback across aliases + auth, err := s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) + if err != nil { + return nil, err + } + bindTempFallback(auth.ID) + entry.Infof("session-affinity: cache hit but auth unavailable, reselected sticky fallback | session=%s primary_cooling=%s fallback_auth=%s provider=%s model=%s", truncateSessionID(primaryID), cachedAuthID, auth.ID, provider, model) + return auth, nil + } + + if fallbackKey != "" { if cachedAuthID, ok := s.cache.Get(fallbackKey); ok { for _, auth := range available { if auth.ID == cachedAuthID { - entry.Infof("session-affinity: concurrent cache hit | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) + invalidateTempFallback() 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 } } + if fallbackAuth, ok := getTempFallbackAuth(); ok { + entry.Infof("session-affinity: sticky secondary fallback cache hit | session=%s fallback=%s temp_auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), fallbackAuth.ID, provider, model) + return fallbackAuth, nil + } + auth, err := s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) + if err != nil { + return nil, err + } + bindTempFallback(auth.ID) + entry.Infof("session-affinity: fallback cache hit but auth unavailable, reselected sticky fallback | session=%s fallback=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), auth.ID, provider, model) + return auth, nil } } - // Authoritative stale observation conducted under bindMu using non-refreshing token read. - // Observe both alias groups: they may be split across different auths, in - // which case failover must reconcile both, not just the first one found. - staleKey := cacheKey - staleAuthID, staleGen, staleAliases, hasStale := s.cache.GetWithGeneration(cacheKey) - splitAuthID := "" - var splitGen uint64 - var splitAliases []string - hasSplit := false - if fallbackKey != "" { - splitAuthID, splitGen, splitAliases, hasSplit = s.cache.GetWithGeneration(fallbackKey) - } - splitGroups := hasStale && hasSplit && staleAuthID != splitAuthID - if !hasStale && hasSplit { - staleKey = fallbackKey - staleAuthID, staleGen, staleAliases, hasStale = splitAuthID, splitGen, splitAliases, true - } - auth, err := s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) if err != nil { return nil, err } - - if hasStale { - if splitGroups { - // Split alias groups (prompt-cache and conversation aliases bound to - // different auths): merge BOTH alias sets into a single group bound - // to the selected auth. Rebinding the groups separately would leave - // two groups on the same auth, and later housekeeping (OnResult) - // processes only the group holding the request's primary key — the - // surviving split group would keep selecting a failed auth. - if !s.mergeSplitAliasGroupsCAS(cacheKey, fallbackKey, auth.ID) { - entry.Infof("session-affinity: split-group merge lost to concurrent writer after retries | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) - } - } else { - additional := []string{cacheKey} - if fallbackKey != "" { - additional = append(additional, fallbackKey) - } - if s.rebindAliasGroupCAS(staleKey, staleAuthID, staleGen, staleAliases, auth.ID, additional) { - entry.Infof("session-affinity: rebound stale alias group | session=%s oldAuth=%s newAuth=%s gen=%d", truncateSessionID(primaryID), staleAuthID, auth.ID, staleGen) - } else { - entry.Infof("session-affinity: CAS rebind aborted due to concurrent mutation, serving selected auth statelessly | session=%s auth=%s", truncateSessionID(primaryID), auth.ID) - } - } - return auth, nil - } - bind(auth.ID) - entry.Infof("session-affinity: cache miss, bound candidate | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) + 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 } @@ -916,79 +925,97 @@ func (s *SessionAffinitySelector) OnResult(res Result) { if fallbackID != "" && fallbackID != primaryID { fallbackKey = ns + "::" + fallbackID + "::" + nsModel } - - if res.Success { + collectResultTempFallbackKeys := func() []string { + keys := []string{cacheKey} if fallbackKey != "" { - s.cache.SetAliases(res.AuthID, cacheKey, fallbackKey) - } else if current, ok := s.cache.Get(cacheKey); !ok || current == res.AuthID { - // Create or refresh in place; a delayed success from a stale auth - // must not steal back a binding that already rebound to another. - s.cache.Set(cacheKey, res.AuthID) + keys = append(keys, fallbackKey) } - return - } - - if res.Error != nil && shouldSkipCredentialCooldown(res.Error) { - return - } - - var aliases []string - if authID, _, groupAliases, ok := s.cache.GetWithGeneration(cacheKey); ok && authID == res.AuthID { - aliases = groupAliases - s.cache.Invalidate(cacheKey) - } else if fallbackKey != "" { - if authID, _, groupAliases, ok := s.cache.GetWithGeneration(fallbackKey); ok && authID == res.AuthID { - aliases = groupAliases - s.cache.Invalidate(fallbackKey) + if aliases := s.cache.Aliases(cacheKey); len(aliases) > 0 { + for _, alias := range aliases { + if alias != "" { + keys = append(keys, alias) + } + } + } + if fallbackKey != "" { + if aliases := s.cache.Aliases(fallbackKey); len(aliases) > 0 { + for _, alias := range aliases { + if alias != "" { + keys = append(keys, alias) + } + } + } } + return keys } - if len(aliases) == 0 { - aliases = []string{cacheKey, fallbackKey} + if res.Success { + // Refresh an existing binding, but also create one when this result + // comes from a path (e.g. stream wrapper) that did not call Pick first. + // If the primary cache is already bound to a different auth, leave it + // untouched: a fallback success must not displace the primary. + if current, ok := s.cache.Get(cacheKey); !ok || current == res.AuthID { + if fallbackKey != "" { + s.cache.SetAliases(res.AuthID, cacheKey, fallbackKey) + } else { + s.cache.Set(cacheKey, res.AuthID) + } + } else { + s.cache.Touch(cacheKey, res.AuthID) + if fallbackKey != "" { + s.cache.Touch(fallbackKey, res.AuthID) + } + } + if s.fallbackCache != nil { + for _, tk := range collectResultTempFallbackKeys() { + if current, ok := s.fallbackCache.Get(tk); ok && current == res.AuthID { + s.fallbackCache.Touch(tk, res.AuthID) + } + } + } + return } - s.quarantineSessionAuth(aliases, res.AuthID, res.RetryAfter) -} -func (s *SessionAffinitySelector) excludeSessionQuarantine(cacheKey, fallbackKey string, auths []*Auth) []*Auth { - if s == nil || s.quarantine == nil || len(auths) == 0 { - return auths - } - filtered := make([]*Auth, 0, len(auths)) - for _, auth := range auths { - if auth == nil { - continue + if res.Error != nil && isTerminalSessionAffinityError(res.Error) { + s.cache.CompareAndDelete(cacheKey, res.AuthID) + if fallbackKey != "" { + s.cache.CompareAndDelete(fallbackKey, res.AuthID) } - blocked := false - for _, key := range []string{cacheKey, fallbackKey} { - if key == "" { - continue - } - if _, ok := s.quarantine.Get(key + "::failed::" + auth.ID); ok { - blocked = true - break + if s.fallbackCache != nil { + for _, tk := range collectResultTempFallbackKeys() { + s.fallbackCache.CompareAndDelete(tk, res.AuthID) } } - if !blocked { - filtered = append(filtered, auth) - } } - return filtered } -func (s *SessionAffinitySelector) quarantineSessionAuth(cacheKeys []string, authID string, retryAfter *time.Duration) { - if s == nil || s.quarantine == nil || authID == "" { - return +// isTerminalSessionAffinityError reports whether a failure represents a permanent +// credential or authorization rejection (such as an invalid API key, revoked grant, +// depleted balance, or unsupported model on the account) that warrants purging +// the long-lived session binding from cache. Transient errors (5xx, 429, timeouts, +// cloudflare challenge) retain affinity so the session returns to its warm prompt +// cache once the cooldown or rate limit clears. +func isTerminalSessionAffinityError(err *Error) bool { + if err == nil { + return false } - delay := 5 * time.Second - if retryAfter != nil && *retryAfter > 0 { - delay = *retryAfter + if shouldSkipCredentialCooldown(err) { + return false } - expiresAt := time.Now().Add(delay) - for _, key := range cacheKeys { - if key == "" { - continue - } - quarantineKey := key + "::failed::" + authID - s.quarantine.setAliasesUntil(authID, expiresAt, quarantineKey) + if isInvalidGrantResultError(err) || isModelSupportResultError(err) { + return true + } + if isCloudflareChallengeResultError(err) { + return false + } + statusCode := statusCodeFromResult(err) + switch statusCode { + case http.StatusUnauthorized, // 401: invalid API key / unauthorized + http.StatusPaymentRequired, // 402: insufficient balance / credits depleted + http.StatusForbidden, // 403: account banned / forbidden + http.StatusNotFound: // 404: model not found / unsupported for account + return true + default: + return false } } @@ -1015,8 +1042,8 @@ func (s *SessionAffinitySelector) Stop() { if s.cache != nil { s.cache.Stop() } - if s.quarantine != nil { - s.quarantine.Stop() + if s.fallbackCache != nil { + s.fallbackCache.Stop() } } @@ -1026,8 +1053,8 @@ func (s *SessionAffinitySelector) InvalidateAuth(authID string) { if s.cache != nil { s.cache.InvalidateAuth(authID) } - if s.quarantine != nil { - s.quarantine.InvalidateAuth(authID) + if s.fallbackCache != nil { + s.fallbackCache.InvalidateAuth(authID) } } diff --git a/sdk/cliproxy/auth/selector_review_p2_test.go b/sdk/cliproxy/auth/selector_review_p2_test.go index 2dedce380..574371706 100644 --- a/sdk/cliproxy/auth/selector_review_p2_test.go +++ b/sdk/cliproxy/auth/selector_review_p2_test.go @@ -1,15 +1,12 @@ package auth import ( - "context" "errors" "fmt" "net/http" "slices" "testing" "time" - - cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) // Regression tests mirrored from CLIProxyAPI PR #4881 follow-up @@ -98,75 +95,6 @@ func TestGetAvailableAuthsSkipsNilCandidates(t *testing.T) { } } -// TestPickRebindsSplitAffinityGroupsOnFailover mirrors the CPA regression -// guard for the codex P2 finding on PR #4881. The CPAPlus binding design has -// no splitConflict skip: on a miss it rebinds the observed stale group via -// CompareAndReplaceAliases and absorbs both session keys into it, which -// converges the split groups onto the selected auth. This test locks that -// convergence in. -func TestPickRebindsSplitAffinityGroupsOnFailover(t *testing.T) { - t.Parallel() - - model := "test-model" - provider := "gemini" - primaryKey := provider + "::pck:pk1::" + model - fallbackKey := provider + "::conv:c1::" + model - - cooled := func(id string) *Auth { - return &Auth{ - ID: id, - ModelStates: map[string]*ModelState{ - model: { - Status: StatusActive, - Unavailable: true, - NextRetryAfter: time.Now().Add(60 * time.Second), - Quota: QuotaState{ - Exceeded: true, - NextRecoverAt: time.Now().Add(60 * time.Second), - }, - }, - }, - } - } - authA := cooled("auth-a") - authB := cooled("auth-b") - authC := &Auth{ - ID: "auth-c", - ModelStates: map[string]*ModelState{ - model: {Status: StatusActive}, - }, - } - - selector := NewSessionAffinitySelector(&FillFirstSelector{}) - selector.cache.SetAliases("auth-a", primaryKey) - selector.cache.SetAliases("auth-b", fallbackKey) - - payload := []byte(`{"prompt_cache_key":"pk1","conversation":{"id":"c1"}}`) - opts := cliproxyexecutor.Options{OriginalRequest: payload, Metadata: map[string]any{}} - auth, err := selector.Pick(context.Background(), provider, model, opts, []*Auth{authA, authB, authC}) - if err != nil { - t.Fatalf("Pick() error = %v, want nil", err) - } - if auth != authC { - t.Fatalf("Pick() = %v, want auth-c (only available auth)", auth.ID) - } - - gotPrimary, genP, aliasesPrimary, okPrimary := selector.cache.GetWithGeneration(primaryKey) - if !okPrimary || gotPrimary != "auth-c" { - t.Fatalf("primary group after failover = %q (ok=%v), want auth-c", gotPrimary, okPrimary) - } - gotFallback, genF, _, okFallback := selector.cache.GetWithGeneration(fallbackKey) - if !okFallback || gotFallback != "auth-c" { - t.Fatalf("fallback group after failover = %q (ok=%v), want auth-c", gotFallback, okFallback) - } - if genP == 0 || genP != genF { - t.Fatalf("split groups not merged into one: primary gen=%d, fallback gen=%d", genP, genF) - } - if !slices.Contains(aliasesPrimary, fallbackKey) { - t.Fatalf("primary group aliases %v missing fallback key %q", aliasesPrimary, fallbackKey) - } -} - // TestCompareAndDeleteGroupRejectsStaleObservation covers the codex P2 // follow-up on PR #4881: when a concurrent request refreshes or extends the // fallback group between the observation and the delete, a stale merge diff --git a/sdk/cliproxy/auth/selector_test.go b/sdk/cliproxy/auth/selector_test.go index ba2d96880..c029874cf 100644 --- a/sdk/cliproxy/auth/selector_test.go +++ b/sdk/cliproxy/auth/selector_test.go @@ -822,7 +822,7 @@ func TestSessionAffinitySelector_ThinkingSuffixVariantsPreserveBindingAndRelease t.Fatalf("third Pick() auth.ID = %q, want %q (thinking suffix variant should keep session stickiness)", third.ID, first.ID) } - // Failure on a thinking-suffix variant (with explicit metadata) should properly release the session binding + // Terminal failure on a thinking-suffix variant (with explicit metadata) should properly release the session binding optsWithMetadata := cliproxyexecutor.Options{ OriginalRequest: payload, Metadata: map[string]any{ @@ -835,7 +835,7 @@ func TestSessionAffinitySelector_ThinkingSuffixVariantsPreserveBindingAndRelease Model: "claude-sonnet-4-5(high)", AuthID: first.ID, Success: false, - Error: &Error{Code: "rate_limited", Message: "rate limited"}, + Error: &Error{Code: "unauthorized", HTTPStatus: http.StatusUnauthorized, Message: "invalid api key"}, Options: optsWithMetadata, }) @@ -864,7 +864,6 @@ func TestSessionAffinitySelector_WeightedBindingRebindsAfterWeightBecomesZero(t if errFirst != nil { t.Fatalf("first Pick() error = %v", errFirst) } - selector.OnResult(Result{AuthID: first.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) if first.ID != authA.ID { t.Fatalf("first Pick() auth.ID = %q, want %q", first.ID, authA.ID) } @@ -874,7 +873,6 @@ func TestSessionAffinitySelector_WeightedBindingRebindsAfterWeightBecomesZero(t if errSecond != nil { t.Fatalf("Pick() after weight update error = %v", errSecond) } - selector.OnResult(Result{AuthID: second.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) if second.ID != authB.ID { t.Fatalf("Pick() after weight update auth.ID = %q, want %q", second.ID, authB.ID) } @@ -882,10 +880,21 @@ func TestSessionAffinitySelector_WeightedBindingRebindsAfterWeightBecomesZero(t authA.Attributes[AttributeWeight] = "10" third, errThird := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) if errThird != nil { - t.Fatalf("Pick() after rebind error = %v", errThird) + t.Fatalf("Pick() after weight restored error = %v", errThird) + } + if third.ID != authA.ID { + t.Fatalf("Pick() after weight restored auth.ID = %q, want original bound auth %q", third.ID, authA.ID) + } + + // When authA is invalidated while remaining weight 0, session permanently rebinds to authB + authA.Attributes[AttributeWeight] = "0" + selector.InvalidateAuth(authA.ID) + fourth, errFourth := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) + if errFourth != nil { + t.Fatalf("Pick() after invalidation error = %v", errFourth) } - if third.ID != authB.ID { - t.Fatalf("Pick() after rebind auth.ID = %q, want sticky auth %q", third.ID, authB.ID) + if fourth.ID != authB.ID { + t.Fatalf("Pick() after invalidation auth.ID = %q, want %q", fourth.ID, authB.ID) } } @@ -1008,9 +1017,8 @@ func TestSessionAffinitySelector_FailoverWhenAuthUnavailable(t *testing.T) { if err != nil { t.Fatalf("Pick() error = %v", err) } - selector.OnResult(Result{AuthID: first.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) - // Remove the bound auth from available list (simulating rate limit) + // Remove the bound auth from available list (simulating rate limit / transient exclusion) availableWithoutFirst := make([]*Auth, 0, len(auths)-1) for _, a := range auths { if a.ID != first.ID { @@ -1018,7 +1026,7 @@ func TestSessionAffinitySelector_FailoverWhenAuthUnavailable(t *testing.T) { } } - // With failover enabled, should pick a new auth + // While original auth is temporarily unavailable, should pick a fallback auth second, err := selector.Pick(context.Background(), "claude", "claude-3", opts, availableWithoutFirst) if err != nil { t.Fatalf("Pick() after failover error = %v", err) @@ -1026,13 +1034,29 @@ func TestSessionAffinitySelector_FailoverWhenAuthUnavailable(t *testing.T) { if second.ID == first.ID { t.Fatalf("Pick() after failover returned same auth %q, expected different", first.ID) } - selector.OnResult(Result{AuthID: second.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) - // Subsequent picks should consistently return the new binding + // When original bound auth becomes available again, affinity is retained (returns to warm cache) + recovered, errRecovered := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) + if errRecovered != nil { + t.Fatalf("Pick() after recovery error = %v", errRecovered) + } + if recovered.ID != first.ID { + t.Fatalf("Pick() after recovery = %q, want original bound auth %q", recovered.ID, first.ID) + } + + // When original auth is explicitly invalidated (permanent failover), session rebinds + selector.InvalidateAuth(first.ID) + third, errThird := selector.Pick(context.Background(), "claude", "claude-3", opts, availableWithoutFirst) + if errThird != nil { + t.Fatalf("Pick() after invalidation error = %v", errThird) + } + if third.ID == first.ID { + t.Fatalf("Pick() after invalidation returned invalidated auth %q", first.ID) + } for i := 0; i < 5; i++ { got, _ := selector.Pick(context.Background(), "claude", "claude-3", opts, availableWithoutFirst) - if got.ID != second.ID { - t.Fatalf("Pick() #%d after failover inconsistent: got %q, want %q", i, got.ID, second.ID) + if got.ID != third.ID { + t.Fatalf("Pick() #%d after permanent rebind inconsistent: got %q, want %q", i, got.ID, third.ID) } } } @@ -1562,43 +1586,6 @@ func TestSessionAffinitySelectorCombinedIdentifiersBindConversationFallback(t *t } } -func TestSessionAffinitySelectorFailureQuarantinesAllAliases(t *testing.T) { - selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ - Fallback: &RoundRobinSelector{}, - TTL: time.Minute, - }) - defer selector.Stop() - auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} - provider := "responses-alias-group-failure" - model := "gpt-test" - - combined := cliproxyexecutor.Options{OriginalRequest: []byte(`{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`), Metadata: map[string]any{}} - first, err := selector.Pick(context.Background(), provider, model, combined, auths) - if err != nil { - t.Fatalf("combined-identifier Pick() error = %v", err) - } - selector.OnResult(Result{AuthID: first.ID, Provider: provider, Model: model, Options: combined, Success: true}) - - promptOnly := cliproxyexecutor.Options{OriginalRequest: []byte(`{"prompt_cache_key":"shared-cache-bucket"}`), Metadata: map[string]any{}} - failed, err := selector.Pick(context.Background(), provider, model, promptOnly, auths) - if err != nil { - t.Fatalf("prompt-only Pick() error = %v", err) - } - if failed.ID != first.ID { - t.Fatalf("prompt-only alias selected %q, want %q", failed.ID, first.ID) - } - selector.OnResult(Result{AuthID: failed.ID, Provider: provider, Model: model, Options: promptOnly, Error: &Error{Code: "upstream_failed", Message: "upstream failed", Retryable: true}}) - - conversationOnly := cliproxyexecutor.Options{OriginalRequest: []byte(`{"conversation":{"id":"conversation-session"}}`), Metadata: map[string]any{}} - next, err := selector.Pick(context.Background(), provider, model, conversationOnly, auths) - if err != nil { - t.Fatalf("conversation-only Pick() error = %v", err) - } - if next.ID == failed.ID { - t.Fatalf("conversation alias reused failed auth %q", failed.ID) - } -} - func TestSessionCacheCompareAndReplaceAliasesPreservesNewerBinding(t *testing.T) { cache := NewSessionCache(time.Minute) defer cache.Stop() @@ -2857,7 +2844,7 @@ func TestSessionAffinitySelector_FallbackReselectReceivesOnlyAvailable(t *testin // threads the request-scoped set of failed auth IDs through to the selector so a // failed auth is never re-picked for the remainder of that request, even though // it is still locally "available". -func TestSessionAffinitySelector_RequestScopedExclusionBreaksCarousel(t *testing.T) { +func TestSessionAffinitySelector_RequestScopedExclusionBreaksCarouselWithRecording(t *testing.T) { t.Parallel() rec := &recordingFallbackSelector{inner: &RoundRobinSelector{}} diff --git a/sdk/cliproxy/auth/session_affinity_fix_test.go b/sdk/cliproxy/auth/session_affinity_fix_test.go index d8be2cd7b..3e73f97e1 100644 --- a/sdk/cliproxy/auth/session_affinity_fix_test.go +++ b/sdk/cliproxy/auth/session_affinity_fix_test.go @@ -95,7 +95,7 @@ func TestSessionAffinity_RetryableFailureInvalidatesMatchingBinding(t *testing.T t.Fatalf("precondition failed: cache key should be bound") } - // Retryable failure (429 Rate Limit) + // Retryable/transient failure (429 Rate Limit) must RETAIN the binding. selector.OnResult(Result{ AuthID: authA.ID, Provider: "provider", @@ -105,8 +105,9 @@ func TestSessionAffinity_RetryableFailureInvalidatesMatchingBinding(t *testing.T Options: opts, }) - if bound, ok := selector.cache.Get(cacheKey); ok { - t.Fatalf("expected cacheKey %q to be invalidated on 429 failure, but still bound to %q", cacheKey, bound) + bound, ok := selector.cache.Get(cacheKey) + if !ok || bound != authA.ID { + t.Fatalf("expected cacheKey %q to be retained on 429 failure; got bound=%q ok=%v", cacheKey, bound, ok) } } @@ -159,8 +160,9 @@ func TestSessionAffinity_ExhaustedRequestDoesNotLeaveLastFailedAuthBound(t *test opts := cliproxyexecutor.Options{ Headers: http.Header{"X-Session-Id": []string{"sess-exhausted-123"}}, } + cacheKey := "provider::header:sess-exhausted-123::model" - // Attempt 1: picks Auth A, fails 429 + // Attempt 1: picks Auth A, fails 429 (transient - binding retained) picked1, _ := selector.Pick(context.Background(), "provider", "model", opts, auths) selector.OnResult(Result{ AuthID: picked1.ID, @@ -170,8 +172,11 @@ func TestSessionAffinity_ExhaustedRequestDoesNotLeaveLastFailedAuthBound(t *test Error: &Error{HTTPStatus: http.StatusTooManyRequests}, Options: opts, }) + if bound, ok := selector.cache.Get(cacheKey); !ok || bound != picked1.ID { + t.Fatalf("primary binding should be retained after transient 429; bound=%q ok=%v", bound, ok) + } - // Attempt 2 within request: excludes Auth A, picks Auth B, fails 429 + // Attempt 2 within request: excludes Auth A, picks Auth B as sticky fallback, fails 429 opts2 := cliproxyexecutor.Options{ Headers: http.Header{"X-Session-Id": []string{"sess-exhausted-123"}}, Metadata: map[string]any{ @@ -179,6 +184,9 @@ func TestSessionAffinity_ExhaustedRequestDoesNotLeaveLastFailedAuthBound(t *test }, } picked2, _ := selector.Pick(context.Background(), "provider", "model", opts2, auths) + if picked2.ID == picked1.ID { + t.Fatalf("fallback pick returned excluded auth %q", picked1.ID) + } selector.OnResult(Result{ AuthID: picked2.ID, Provider: "provider", @@ -188,10 +196,21 @@ func TestSessionAffinity_ExhaustedRequestDoesNotLeaveLastFailedAuthBound(t *test Options: opts2, }) - // Verify session cache is left clean (unbound) - cacheKey := "provider::header:sess-exhausted-123::model" - if bound, ok := selector.cache.Get(cacheKey); ok { - t.Fatalf("exhausted request left last failed auth bound=%q in cache", bound) + // Primary binding must still be the original auth; fallback is stored temporarily. + if bound, ok := selector.cache.Get(cacheKey); !ok || bound != picked1.ID { + t.Fatalf("primary binding should be retained after fallback transient 429; bound=%q ok=%v", bound, ok) + } + if fb, ok := selector.fallbackCache.Get(cacheKey); !ok || fb != picked2.ID { + t.Fatalf("fallback cache should hold %q after fallback failure; got %q ok=%v", picked2.ID, fb, ok) + } + + // Once the primary auth is available again, the session returns to it and clears the fallback. + second, _ := selector.Pick(context.Background(), "provider", "model", opts, auths) + if second.ID != picked1.ID { + t.Fatalf("exhausted request should return to original binding %q, got %q", picked1.ID, second.ID) + } + if _, ok := selector.fallbackCache.Get(cacheKey); ok { + t.Fatalf("temporary fallback should be cleared when primary recovers") } } @@ -291,9 +310,12 @@ func TestSessionAffinity_CachedAuthUnavailableRebindsFallback(t *testing.T) { t.Fatalf("Pick = %v/%v, want B", picked, err) } - bound, ok := selector.cache.Get(cacheKey) - if !ok || bound != authB.ID { - t.Fatalf("expected stale A to be replaced by pre-bound B; cache=%q ok=%v", bound, ok) + // Primary binding stays with the original auth; the fallback is sticky but temporary. + if bound, ok := selector.cache.Get(cacheKey); !ok || bound != authA.ID { + t.Fatalf("expected primary binding to remain %q after fallback; got %q ok=%v", authA.ID, bound, ok) + } + if fb, ok := selector.fallbackCache.Get(cacheKey); !ok || fb != authB.ID { + t.Fatalf("expected fallback cache to hold %q; got %q ok=%v", authB.ID, fb, ok) } } @@ -316,22 +338,27 @@ func TestSessionAffinity_FallbackBFailsLeavesCacheEmpty(t *testing.T) { if picked.ID != authB.ID { t.Fatalf("Pick = %q, want B", picked.ID) } - // B fails with retryable error. + // B fails with a transient 429. The primary binding is retained and B is stored temporarily. selector.OnResult(Result{AuthID: picked.ID, Provider: "provider", Model: "model", Success: false, Error: &Error{HTTPStatus: http.StatusTooManyRequests}, Options: opts}) - // Cache must be empty (no stale A, no B). - if bound, ok := selector.cache.Get(cacheKey); ok { - t.Fatalf("cache should be empty after B failure, got %q", bound) + if bound, ok := selector.cache.Get(cacheKey); !ok || bound != authA.ID { + t.Fatalf("primary binding should be retained after fallback failure; got %q ok=%v", bound, ok) + } + if fb, ok := selector.fallbackCache.Get(cacheKey); !ok || fb != authB.ID { + t.Fatalf("fallback cache should hold %q after transient failure; got %q ok=%v", authB.ID, fb, ok) } - // Immediate second request starts from normal fallback (A), not stale B affinity. + // Once A is available again, the session returns to it and clears the temporary fallback. second, _ := selector.Pick(context.Background(), "provider", "model", opts, []*Auth{authA, authB}) if second.ID != authA.ID { - t.Fatalf("second request should reselect from fallback, got %q", second.ID) + t.Fatalf("second request should return to primary binding %q, got %q", authA.ID, second.ID) + } + if _, ok := selector.fallbackCache.Get(cacheKey); ok { + t.Fatalf("temporary fallback should be cleared when primary recovers") } } -func TestSessionAffinity_FallbackBSucceedsBindsB(t *testing.T) { +func TestSessionAffinity_FallbackBSucceedsBindsTemporaryFallback(t *testing.T) { authA := &Auth{ID: "auth-a"} authB := &Auth{ID: "auth-b"} @@ -349,11 +376,23 @@ func TestSessionAffinity_FallbackBSucceedsBindsB(t *testing.T) { if picked.ID != authB.ID { t.Fatalf("Pick = %q, want B", picked.ID) } + // Fallback success must not overwrite the primary binding; it is stored as a sticky temporary fallback. selector.OnResult(Result{AuthID: picked.ID, Provider: "provider", Model: "model", Success: true, Options: opts}) - bound, ok := selector.cache.Get(cacheKey) - if !ok || bound != authB.ID { - t.Fatalf("B should be bound after success; bound=%q ok=%v", bound, ok) + if bound, ok := selector.cache.Get(cacheKey); !ok || bound != authA.ID { + t.Fatalf("primary binding should remain %q after fallback success; got %q ok=%v", authA.ID, bound, ok) + } + if fb, ok := selector.fallbackCache.Get(cacheKey); !ok || fb != authB.ID { + t.Fatalf("fallback cache should hold %q after success; got %q ok=%v", authB.ID, fb, ok) + } + + // When the primary auth is available again, the session returns to it. + second, _ := selector.Pick(context.Background(), "provider", "model", opts, []*Auth{authA, authB}) + if second.ID != authA.ID { + t.Fatalf("second request should return to primary binding %q, got %q", authA.ID, second.ID) + } + if _, ok := selector.fallbackCache.Get(cacheKey); ok { + t.Fatalf("temporary fallback should be cleared when primary recovers") } } @@ -404,14 +443,14 @@ func TestSessionAffinity_StreamFailureThroughWrapperInvalidates(t *testing.T) { t.Fatalf("precondition: bound") } - // Stream fails with retryable upstream error (503). + // Stream fails with a transient upstream error (503). The binding must be retained. errChunk := cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusServiceUnavailable}} res := manager.wrapStreamResult(ctx, auth, "stream-provider", "stream-model", opts, nil, []cliproxyexecutor.StreamChunk{errChunk}, closedStreamChunks(), OAuthModelAliasResult{}, false) for range res.Chunks { } - if bound, ok := affinity.cache.Get(cacheKey); ok { - t.Fatalf("stream failure should invalidate affinity; still bound=%q", bound) + if bound, ok := affinity.cache.Get(cacheKey); !ok || bound != auth.ID { + t.Fatalf("stream failure should retain affinity; bound=%q ok=%v", bound, ok) } } func optsWithMixedNamespace(opts cliproxyexecutor.Options) cliproxyexecutor.Options { @@ -460,7 +499,7 @@ func TestSessionAffinity_MixedNamespace_PickRecordsAndOnResultBindsCanonicalKey( } } -func TestSessionAffinity_MixedNamespace_FailureLeavesCacheEmpty(t *testing.T) { +func TestSessionAffinity_MixedNamespace_FailureRetainsBinding(t *testing.T) { gemini := &Auth{ID: "gemini-auth", Provider: "gemini"} fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { return available[0], nil @@ -471,10 +510,11 @@ func TestSessionAffinity_MixedNamespace_FailureLeavesCacheEmpty(t *testing.T) { cacheKey := "mixed::header:mixed-fail-12345::model" picked, _ := selector.Pick(context.Background(), "mixed", "model", opts, []*Auth{gemini}) + // A 429 is transient: the canonical binding must be retained. selector.OnResult(Result{AuthID: picked.ID, Provider: "gemini", Model: "model", Success: false, Error: &Error{HTTPStatus: http.StatusTooManyRequests}, Options: opts}) - if bound, ok := selector.cache.Get(cacheKey); ok { - t.Fatalf("mixed cache should be empty after failure; still bound=%q", bound) + if bound, ok := selector.cache.Get(cacheKey); !ok || bound != picked.ID { + t.Fatalf("mixed cache should retain binding after transient failure; bound=%q ok=%v", bound, ok) } } @@ -494,9 +534,13 @@ func TestSessionAffinity_MixedNamespace_StaleAuthRebindsFallback(t *testing.T) { if picked.ID != authB.ID { t.Fatalf("Pick = %q, want B", picked.ID) } - bound, ok := selector.cache.Get(cacheKey) - if !ok || bound != authB.ID { - t.Fatalf("expected stale A to be replaced by pre-bound B; cache=%q ok=%v", bound, ok) + + // The primary binding stays with A; the fallback is sticky but temporary. + if bound, ok := selector.cache.Get(cacheKey); !ok || bound != authA.ID { + t.Fatalf("expected primary binding to remain %q; got %q ok=%v", authA.ID, bound, ok) + } + if fb, ok := selector.fallbackCache.Get(cacheKey); !ok || fb != authB.ID { + t.Fatalf("expected fallback cache to hold %q; got %q ok=%v", authB.ID, fb, ok) } } @@ -522,7 +566,7 @@ func TestSessionAffinity_MixedNamespace_StaleFailureCannotDeleteNewerSuccess(t * } } -func TestSessionAffinity_MixedNamespace_StreamBindsAndInvalidatesCanonicalKey(t *testing.T) { +func TestSessionAffinity_MixedNamespace_StreamBindsAndRetainsCanonicalKey(t *testing.T) { ctx := context.Background() manager := NewManager(nil, nil, nil) affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{Fallback: &RoundRobinSelector{}, TTL: time.Hour}) @@ -547,13 +591,13 @@ func TestSessionAffinity_MixedNamespace_StreamBindsAndInvalidatesCanonicalKey(t t.Fatalf("stream mixed success should bind canonical key; bound=%q ok=%v", bound, ok) } - // Failure invalidates the same canonical key. + // A 503 stream failure is transient: the canonical binding must be retained. errChunk := cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusServiceUnavailable}} res2 := manager.wrapStreamResult(ctx, auth, "gemini", "stream-model", opts, nil, []cliproxyexecutor.StreamChunk{errChunk}, closedStreamChunks(), OAuthModelAliasResult{}, false) for range res2.Chunks { } - if bound, ok := affinity.cache.Get(cacheKey); ok { - t.Fatalf("stream mixed failure should invalidate canonical key; still bound=%q", bound) + if bound, ok := affinity.cache.Get(cacheKey); !ok || bound != auth.ID { + t.Fatalf("stream mixed failure should retain canonical key; bound=%q ok=%v", bound, ok) } } @@ -577,7 +621,7 @@ func TestSessionAffinity_SingleProviderStillUsesActualProviderKey(t *testing.T) } } -func TestSessionAffinity_MixedNamespace_SecondRequestSkipsUnavailable(t *testing.T) { +func TestSessionAffinity_MixedNamespace_SecondRequestReturnsToPrimary(t *testing.T) { authA := &Auth{ID: "auth-a", Provider: "gemini"} authB := &Auth{ID: "auth-b", Provider: "gemini"} fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { @@ -591,20 +635,27 @@ func TestSessionAffinity_MixedNamespace_SecondRequestSkipsUnavailable(t *testing // Bind A under the canonical key. selector.OnResult(Result{AuthID: authA.ID, Provider: "gemini", Model: "model", Success: true, Options: opts}) - // A deterministic-unavailable (only B in list) -> Pick skips A, gets B; B fails -> cache empty. + // A is unavailable (only B in list) -> Pick selects B as sticky fallback. picked, _ := selector.Pick(context.Background(), "mixed", "model", opts, []*Auth{authB}) if picked.ID != authB.ID { t.Fatalf("Pick = %q, want B", picked.ID) } + // B fails with a transient 503. The primary binding to A is retained. selector.OnResult(Result{AuthID: picked.ID, Provider: "gemini", Model: "model", Success: false, Error: &Error{HTTPStatus: http.StatusServiceUnavailable}, Options: opts}) - if _, ok := selector.cache.Get(cacheKey); ok { - t.Fatalf("cache should be empty after B failure") + if bound, ok := selector.cache.Get(cacheKey); !ok || bound != authA.ID { + t.Fatalf("primary binding should be retained after fallback failure; got %q ok=%v", bound, ok) + } + if fb, ok := selector.fallbackCache.Get(cacheKey); !ok || fb != authB.ID { + t.Fatalf("fallback cache should hold %q; got %q ok=%v", authB.ID, fb, ok) } - // Immediate second request with both available reselects from fallback (A), not stale B. + // Once A is available again, the session returns to it and clears the temporary fallback. second, _ := selector.Pick(context.Background(), "mixed", "model", opts, []*Auth{authA, authB}) if second.ID != authA.ID { - t.Fatalf("second request should reselect from fallback, got %q", second.ID) + t.Fatalf("second request should return to primary binding %q, got %q", authA.ID, second.ID) + } + if _, ok := selector.fallbackCache.Get(cacheKey); ok { + t.Fatalf("temporary fallback should be cleared when primary recovers") } } func optsWithAffinityNamespaces(opts cliproxyexecutor.Options, provider, model string) cliproxyexecutor.Options { @@ -675,7 +726,7 @@ func TestSessionAffinity_ModelNamespace_SingleProviderAliasRewrite(t *testing.T) } } -func TestSessionAffinity_ModelNamespace_FailureClearsRouteBinding(t *testing.T) { +func TestSessionAffinity_ModelNamespace_FailureRetainsRouteBinding(t *testing.T) { auth := &Auth{ID: "auth-a", Provider: "gemini"} fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { return available[0], nil @@ -691,10 +742,10 @@ func TestSessionAffinity_ModelNamespace_FailureClearsRouteBinding(t *testing.T) t.Fatalf("precondition: route binding should exist") } - // Failure with a rewritten Result model clears the canonical route binding. + // 503 is transient: the canonical route binding must be retained. selector.OnResult(Result{AuthID: auth.ID, Provider: "gemini", Model: "gemini-3.5-flash-lite", Success: false, Error: &Error{HTTPStatus: http.StatusServiceUnavailable}, Options: opts}) - if bound, ok := selector.cache.Get(routeKey); ok { - t.Fatalf("route binding not cleared after failure; bound=%q", bound) + if bound, ok := selector.cache.Get(routeKey); !ok || bound != auth.ID { + t.Fatalf("route binding should be retained after transient failure; bound=%q ok=%v", bound, ok) } } @@ -720,7 +771,7 @@ func TestSessionAffinity_ModelNamespace_StaleFailureCannotDeleteNewerSuccess(t * } } -func TestSessionAffinity_ModelNamespace_StreamRewriteBindsRouteKey(t *testing.T) { +func TestSessionAffinity_ModelNamespace_StreamRewriteBindsAndRetainsRouteKey(t *testing.T) { ctx := context.Background() manager := NewManager(nil, nil, nil) affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{Fallback: &RoundRobinSelector{}, TTL: time.Hour}) @@ -745,13 +796,13 @@ func TestSessionAffinity_ModelNamespace_StreamRewriteBindsRouteKey(t *testing.T) t.Fatalf("stream rewrite should bind route key; bound=%q ok=%v", bound, ok) } - // Stream failure with rewritten model clears the route key. + // A 503 stream failure is transient: the route key binding must be retained. errChunk := cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusServiceUnavailable}} res2 := manager.wrapStreamResult(ctx, auth, "gemini", "gemini-3.5-flash-lite", opts, nil, []cliproxyexecutor.StreamChunk{errChunk}, closedStreamChunks(), OAuthModelAliasResult{}, false) for range res2.Chunks { } - if bound, ok := affinity.cache.Get(routeKey); ok { - t.Fatalf("stream failure should clear route key; still bound=%q", bound) + if bound, ok := affinity.cache.Get(routeKey); !ok || bound != auth.ID { + t.Fatalf("stream failure should retain route key; bound=%q ok=%v", bound, ok) } } @@ -773,145 +824,6 @@ func TestSessionAffinity_ModelNamespace_MetadataAbsentUsesResultModel(t *testing } } -func TestSessionAffinity_QuarantinesRetryAfterForSameSessionOnly(t *testing.T) { - authA := &Auth{ID: "auth-a", Provider: "gemini"} - authB := &Auth{ID: "auth-b", Provider: "gemini"} - fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { - return available[0], nil - }) - selector := NewSessionAffinitySelector(fallback) - defer selector.Stop() - - opts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"cooldown-session-one"}}}, "mixed", ".gemini-flash") - first, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) - if err != nil || first.ID != authA.ID { - t.Fatalf("first Pick = %v/%v, want auth-a", first, err) - } - - retryAfter := 53 * time.Second - selector.OnResult(Result{ - AuthID: authA.ID, - Provider: "gemini", - Model: "gemini-3.6-flash", - Success: false, - Error: &Error{HTTPStatus: http.StatusTooManyRequests}, - RetryAfter: &retryAfter, - Options: opts, - }) - - second, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) - if err != nil || second.ID != authB.ID { - t.Fatalf("same-session retry Pick = %v/%v, want auth-b", second, err) - } - - otherOpts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"cooldown-session-two"}}}, "mixed", ".gemini-flash") - other, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", otherOpts, []*Auth{authA, authB}) - if err != nil || other.ID != authA.ID { - t.Fatalf("other-session Pick = %v/%v, want auth-a", other, err) - } -} - -func TestSessionAffinity_QuarantinesMultipleFailedAuths(t *testing.T) { - authA := &Auth{ID: "auth-a", Provider: "gemini"} - authB := &Auth{ID: "auth-b", Provider: "gemini"} - authC := &Auth{ID: "auth-c", Provider: "gemini"} - fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { - return available[0], nil - }) - selector := NewSessionAffinitySelector(fallback) - defer selector.Stop() - - opts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"cooldown-multiple"}}}, "mixed", ".gemini-flash") - retryAfter := 53 * time.Second - for _, auth := range []*Auth{authA, authB} { - picked, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB, authC}) - if err != nil || picked.ID != auth.ID { - t.Fatalf("Pick before failing %s = %v/%v", auth.ID, picked, err) - } - selector.OnResult(Result{ - AuthID: auth.ID, - Provider: "gemini", - Model: "gemini-3.6-flash", - Success: false, - Error: &Error{HTTPStatus: http.StatusTooManyRequests}, - RetryAfter: &retryAfter, - Options: opts, - }) - } - - third, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB, authC}) - if err != nil || third.ID != authC.ID { - t.Fatalf("third Pick = %v/%v, want auth-c", third, err) - } -} - -func TestSessionAffinity_QuarantineExpires(t *testing.T) { - authA := &Auth{ID: "auth-a", Provider: "gemini"} - authB := &Auth{ID: "auth-b", Provider: "gemini"} - fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { - return available[0], nil - }) - selector := NewSessionAffinitySelector(fallback) - defer selector.Stop() - - opts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"cooldown-expiry"}}}, "mixed", ".gemini-flash") - retryAfter := 20 * time.Millisecond - selector.OnResult(Result{ - AuthID: authA.ID, - Provider: "gemini", - Model: "gemini-3.6-flash", - Success: false, - Error: &Error{HTTPStatus: http.StatusTooManyRequests}, - RetryAfter: &retryAfter, - Options: opts, - }) - - before, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) - if err != nil || before.ID != authB.ID { - t.Fatalf("Pick before expiry = %v/%v, want auth-b", before, err) - } - selector.OnResult(Result{AuthID: authB.ID, Provider: "gemini", Model: "gemini-3.6-flash", Success: false, Error: &Error{HTTPStatus: http.StatusBadGateway}, Options: opts}) - time.Sleep(30 * time.Millisecond) - after, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) - if err != nil || after.ID != authA.ID { - t.Fatalf("Pick after expiry = %v/%v, want auth-a", after, err) - } -} - -func TestSessionAffinity_StaleSuccessDoesNotClearNewerQuarantine(t *testing.T) { - authA := &Auth{ID: "auth-a", Provider: "gemini"} - authB := &Auth{ID: "auth-b", Provider: "gemini"} - fallback := pickFuncSelector(func(_ context.Context, _, _ string, _ cliproxyexecutor.Options, available []*Auth) (*Auth, error) { - return available[0], nil - }) - selector := NewSessionAffinitySelector(fallback) - defer selector.Stop() - - opts := optsWithAffinityNamespaces(cliproxyexecutor.Options{Headers: http.Header{"X-Session-Id": []string{"cooldown-stale-success"}}}, "mixed", ".gemini-flash") - retryAfter := 53 * time.Second - selector.OnResult(Result{ - AuthID: authA.ID, - Provider: "gemini", - Model: "gemini-3.6-flash", - Success: false, - Error: &Error{HTTPStatus: http.StatusTooManyRequests}, - RetryAfter: &retryAfter, - Options: opts, - }) - selector.OnResult(Result{ - AuthID: authA.ID, - Provider: "gemini", - Model: "gemini-3.6-flash", - Success: true, - Options: opts, - }) - - got, err := selector.Pick(context.Background(), "mixed", ".gemini-flash", opts, []*Auth{authA, authB}) - if err != nil || got.ID != authB.ID { - t.Fatalf("Pick after stale success = %v/%v, want auth-b while auth-a remains quarantined", got, err) - } -} - func TestSessionAffinity_RequestScoped400DoesNotQuarantine(t *testing.T) { authA := &Auth{ID: "auth-a", Provider: "gemini"} authB := &Auth{ID: "auth-b", Provider: "gemini"} diff --git a/sdk/cliproxy/auth/session_affinity_metadata_test.go b/sdk/cliproxy/auth/session_affinity_metadata_test.go index 9103ba78f..1a2cc44b5 100644 --- a/sdk/cliproxy/auth/session_affinity_metadata_test.go +++ b/sdk/cliproxy/auth/session_affinity_metadata_test.go @@ -19,11 +19,11 @@ type failExecutor struct { func (e *failExecutor) Identifier() string { return e.provider } func (e *failExecutor) Execute(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { e.calls.Add(1) - return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusInternalServerError, Message: "upstream failure"} + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusUnauthorized, Message: "invalid api key"} } func (e *failExecutor) ExecuteStream(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { e.calls.Add(1) - return nil, &Error{HTTPStatus: http.StatusInternalServerError, Message: "upstream failure"} + return nil, &Error{HTTPStatus: http.StatusUnauthorized, Message: "invalid api key"} } func (e *failExecutor) Refresh(ctx context.Context, auth *Auth) (*Auth, error) { return auth, nil } func (e *failExecutor) CountTokens(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { @@ -252,7 +252,7 @@ func TestSessionAffinityOnResultWithMismatchedNamespaceFailsToUnbind(t *testing. Provider: "gemini", // actual provider Model: model, Success: false, - Error: &Error{HTTPStatus: http.StatusInternalServerError}, + Error: &Error{HTTPStatus: http.StatusUnauthorized, Message: "invalid api key"}, Options: cliproxyexecutor.Options{ Headers: http.Header{"X-Session-Id": []string{"sess-ns-1"}}, Metadata: map[string]any{ diff --git a/sdk/cliproxy/auth/session_affinity_priority_test.go b/sdk/cliproxy/auth/session_affinity_priority_test.go index 7426cf270..6f7f26a62 100644 --- a/sdk/cliproxy/auth/session_affinity_priority_test.go +++ b/sdk/cliproxy/auth/session_affinity_priority_test.go @@ -108,11 +108,10 @@ func TestManagerSessionAffinityPreservesBindingAcrossHigherPriorityRecovery(t *t }) expireSessionAffinityPriorityModelCooldown(t, manager, highID, model) - // The affinity namespace fix makes the mixed selection path bind and read - // under the canonical pool key, so the lowID binding is retained across - // higher-priority recovery in both the single- and mixed-provider subtests. - if got := pick(opts); got.ID != lowID { - t.Fatalf("binding after higher-priority recovery = %q, want sticky %q", got.ID, lowID) + // Session affinity is retained on transient failures and returns to the + // original bound credential once its cooldown clears. + if got := pick(opts); got.ID != highID { + t.Fatalf("binding after higher-priority recovery = %q, want recovered original %q", got.ID, highID) } newSessionOpts := cliproxyexecutor.Options{Metadata: map[string]any{ @@ -130,9 +129,10 @@ func TestManagerSessionAffinityPreservesBindingAcrossHigherPriorityRecovery(t *t Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "quota"}, Options: opts, }) - got, errPick := testCase.pick(manager, ctx, provider, model, opts) - if errPick == nil || got != nil { - t.Fatalf("binding after all session auths failed = %v/%v, want no candidate until quarantine expires", got, errPick) + // Session affinity is retained: even when the low-priority fallback + // cools down, the original high-priority binding is still in cache. + if got := pick(opts); got.ID != highID { + t.Fatalf("binding after bound auth became unavailable = %q, want %q", got.ID, highID) } }) } diff --git a/sdk/cliproxy/auth/session_affinity_retention_test.go b/sdk/cliproxy/auth/session_affinity_retention_test.go new file mode 100644 index 000000000..7bc9e013e --- /dev/null +++ b/sdk/cliproxy/auth/session_affinity_retention_test.go @@ -0,0 +1,862 @@ +package auth + +import ( + "context" + "net/http" + "sync/atomic" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type configurableTestExecutor struct { + provider string + calls atomic.Int32 + handler func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) +} + +func (e *configurableTestExecutor) Identifier() string { return e.provider } +func (e *configurableTestExecutor) Execute(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.calls.Add(1) + if e.handler != nil { + return e.handler(ctx, auth, req, opts) + } + return cliproxyexecutor.Response{Payload: []byte(`{"ok":true}`)}, nil +} +func (e *configurableTestExecutor) ExecuteStream(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.calls.Add(1) + return nil, nil +} +func (e *configurableTestExecutor) Refresh(ctx context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} +func (e *configurableTestExecutor) CountTokens(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (e *configurableTestExecutor) HttpRequest(ctx context.Context, auth *Auth, req *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestSessionAffinity_Transient503RetainsBindingAcrossRecovery(t *testing.T) { + ctx := context.Background() + p1 := "affinity-503-p1" + p2 := "affinity-503-p2" + model := "test-model-503" + auth1ID := "auth-1-503" + auth2ID := "auth-2-503" + + manager := NewManager(nil, nil, nil) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Hour, + }) + defer affinity.Stop() + manager.SetSelector(affinity) + + var auth1ShouldFail atomic.Bool + auth1ShouldFail.Store(true) + + exec1 := &configurableTestExecutor{ + provider: p1, + handler: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if auth1ShouldFail.Load() { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusServiceUnavailable, Message: "service unavailable 503"} + } + return cliproxyexecutor.Response{Payload: []byte(`{"served_by":"auth-1"}`)}, nil + }, + } + exec2 := &configurableTestExecutor{ + provider: p2, + handler: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{Payload: []byte(`{"served_by":"auth-2"}`)}, nil + }, + } + manager.RegisterExecutor(exec1) + manager.RegisterExecutor(exec2) + + for _, auth := range []*Auth{ + {ID: auth1ID, Provider: p1, Status: StatusActive}, + {ID: auth2ID, Provider: p2, Status: StatusActive}, + } { + if _, errRegister := manager.Register(WithSkipPersist(ctx), auth); errRegister != nil { + t.Fatalf("Register(%s): %v", auth.ID, errRegister) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + } + + req := cliproxyexecutor.Request{Model: model} + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{"sess-retention-503"}}, + } + + // 1. First Execute: auth-1 fails with transient 503, retry falls over to auth-2 which succeeds. + resp, errExec := manager.Execute(ctx, []string{p1, p2}, req, opts) + if errExec != nil { + t.Fatalf("first Execute failed: %v", errExec) + } + if string(resp.Payload) != `{"served_by":"auth-2"}` { + t.Fatalf("first Execute payload = %s, want auth-2", string(resp.Payload)) + } + + // Session affinity binding must STILL point to auth-1 (retained despite transient failure) + sessionKey := "mixed::header:sess-retention-503::" + model + boundAuthID, ok := affinity.cache.Get(sessionKey) + if !ok { + t.Fatalf("expected sessionKey %q to remain in cache after transient 503", sessionKey) + } + if boundAuthID != auth1ID { + t.Fatalf("sessionKey bound to %q, want %q (transient 503 must not purge affinity)", boundAuthID, auth1ID) + } + + // 2. Cooldown for auth-1 clears + auth1ShouldFail.Store(false) + expireSessionAffinityPriorityModelCooldown(t, manager, auth1ID, model) + + // 3. Second Execute for the SAME session must return to auth-1 where prompt cache lives + resp2, errExec2 := manager.Execute(ctx, []string{p1, p2}, req, opts) + if errExec2 != nil { + t.Fatalf("second Execute failed: %v", errExec2) + } + if string(resp2.Payload) != `{"served_by":"auth-1"}` { + t.Fatalf("second Execute payload = %s, want auth-1 (session should return to original warm cache auth)", string(resp2.Payload)) + } +} + +func TestSessionAffinity_Transient429RetryAfterRetainsBindingAcrossRecovery(t *testing.T) { + ctx := context.Background() + p1 := "affinity-429-p1" + p2 := "affinity-429-p2" + model := "test-model-429" + auth1ID := "auth-1-429" + auth2ID := "auth-2-429" + + manager := NewManager(nil, nil, nil) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Hour, + }) + defer affinity.Stop() + manager.SetSelector(affinity) + + var auth1ShouldFail atomic.Bool + auth1ShouldFail.Store(true) + + retryAfterDuration := 100 * time.Millisecond + exec1 := &configurableTestExecutor{ + provider: p1, + handler: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if auth1ShouldFail.Load() { + return cliproxyexecutor.Response{}, &retryAfterStatusError{ + status: http.StatusTooManyRequests, + retryAfter: retryAfterDuration, + message: "rate limited 429", + } + } + return cliproxyexecutor.Response{Payload: []byte(`{"served_by":"auth-1"}`)}, nil + }, + } + exec2 := &configurableTestExecutor{ + provider: p2, + handler: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{Payload: []byte(`{"served_by":"auth-2"}`)}, nil + }, + } + manager.RegisterExecutor(exec1) + manager.RegisterExecutor(exec2) + + for _, auth := range []*Auth{ + {ID: auth1ID, Provider: p1, Status: StatusActive}, + {ID: auth2ID, Provider: p2, Status: StatusActive}, + } { + if _, errRegister := manager.Register(WithSkipPersist(ctx), auth); errRegister != nil { + t.Fatalf("Register(%s): %v", auth.ID, errRegister) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + } + + req := cliproxyexecutor.Request{Model: model} + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{"sess-retention-429"}}, + } + + // 1. First Execute: auth-1 returns 429, request falls over to auth-2 which succeeds. + resp, errExec := manager.Execute(ctx, []string{p1, p2}, req, opts) + if errExec != nil { + t.Fatalf("first Execute failed: %v", errExec) + } + if string(resp.Payload) != `{"served_by":"auth-2"}` { + t.Fatalf("first Execute payload = %s, want auth-2", string(resp.Payload)) + } + + // Binding must STILL be auth-1 + sessionKey := "mixed::header:sess-retention-429::" + model + boundAuthID, ok := affinity.cache.Get(sessionKey) + if !ok { + t.Fatalf("expected sessionKey %q to remain in cache after 429 rate limit", sessionKey) + } + if boundAuthID != auth1ID { + t.Fatalf("sessionKey bound to %q, want %q", boundAuthID, auth1ID) + } + + // 2. Cooldown for auth-1 clears + auth1ShouldFail.Store(false) + expireSessionAffinityPriorityModelCooldown(t, manager, auth1ID, model) + + // 3. Next Execute returns to auth-1 + resp2, errExec2 := manager.Execute(ctx, []string{p1, p2}, req, opts) + if errExec2 != nil { + t.Fatalf("second Execute failed: %v", errExec2) + } + if string(resp2.Payload) != `{"served_by":"auth-1"}` { + t.Fatalf("second Execute payload = %s, want auth-1", string(resp2.Payload)) + } +} + +func TestSessionAffinity_Terminal401InvalidAPIKeyUnbindsSession(t *testing.T) { + ctx := context.Background() + p1 := "affinity-401-p1" + p2 := "affinity-401-p2" + model := "test-model-401" + auth1ID := "auth-1-401" + auth2ID := "auth-2-401" + + manager := NewManager(nil, nil, nil) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Hour, + }) + defer affinity.Stop() + manager.SetSelector(affinity) + + exec1 := &configurableTestExecutor{ + provider: p1, + handler: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, &Error{ + Code: "unauthorized", + HTTPStatus: http.StatusUnauthorized, + Message: "invalid_api_key", + } + }, + } + exec2 := &configurableTestExecutor{ + provider: p2, + handler: func(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{Payload: []byte(`{"served_by":"auth-2"}`)}, nil + }, + } + manager.RegisterExecutor(exec1) + manager.RegisterExecutor(exec2) + + for _, auth := range []*Auth{ + {ID: auth1ID, Provider: p1, Status: StatusActive}, + {ID: auth2ID, Provider: p2, Status: StatusActive}, + } { + if _, errRegister := manager.Register(WithSkipPersist(ctx), auth); errRegister != nil { + t.Fatalf("Register(%s): %v", auth.ID, errRegister) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + } + + req := cliproxyexecutor.Request{Model: model} + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{"sess-retention-401"}}, + } + + // 1. First Execute: auth-1 fails with 401 terminal error, unbinds, falls over to auth-2 which succeeds and binds. + resp, errExec := manager.Execute(ctx, []string{p1, p2}, req, opts) + if errExec != nil { + t.Fatalf("first Execute failed: %v", errExec) + } + if string(resp.Payload) != `{"served_by":"auth-2"}` { + t.Fatalf("first Execute payload = %s, want auth-2", string(resp.Payload)) + } + + // Session is now permanently rebound to auth-2 + sessionKey := "mixed::header:sess-retention-401::" + model + boundAuthID, ok := affinity.cache.Get(sessionKey) + if !ok { + t.Fatalf("expected sessionKey %q to be bound to auth-2", sessionKey) + } + if boundAuthID != auth2ID { + t.Fatalf("sessionKey bound to %q, want %q (terminal 401 should rebind to next auth)", boundAuthID, auth2ID) + } + + // 2. Subsequent requests for the same session stay on auth-2 + resp2, errExec2 := manager.Execute(ctx, []string{p1, p2}, req, opts) + if errExec2 != nil { + t.Fatalf("second Execute failed: %v", errExec2) + } + if string(resp2.Payload) != `{"served_by":"auth-2"}` { + t.Fatalf("second Execute payload = %s, want auth-2", string(resp2.Payload)) + } +} + +func TestSessionAffinitySelector_RequestScopedExclusionBreaksCarousel(t *testing.T) { + t.Parallel() + + fallback := &RoundRobinSelector{} + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: fallback, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{ + {ID: "auth-a"}, + {ID: "auth-b"}, + {ID: "auth-c"}, + } + + payload := []byte(`{"metadata":{"user_id":"user_xxx_account__session_carousel-break"}}`) + opts := cliproxyexecutor.Options{OriginalRequest: payload} + + // 1. First pick establishes affinity binding to auth-a + first, err := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) + if err != nil { + t.Fatalf("Pick() error = %v", err) + } + if first.ID != "auth-a" { + t.Fatalf("initial pick = %q, want auth-a", first.ID) + } + selector.OnResult(Result{AuthID: first.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) + + // 2. Simulated retries within one request: auth-a failed and is excluded from available candidates + availableWithoutFirst := make([]*Auth, 0, len(auths)-1) + for _, a := range auths { + if a.ID != first.ID { + availableWithoutFirst = append(availableWithoutFirst, a) + } + } + + // 20 successive retry attempts within the request must NEVER return auth-a + for attempt := 0; attempt < 20; attempt++ { + got, errPick := selector.Pick(context.Background(), "claude", "claude-3", opts, availableWithoutFirst) + if errPick != nil { + t.Fatalf("attempt %d Pick() error = %v", attempt, errPick) + } + if got.ID == first.ID { + t.Fatalf("attempt %d returned excluded auth %q, expected different", attempt, first.ID) + } + } + + // 3. New request after recovery with full candidates returns to auth-a (warm cache retained) + recovered, errRecovered := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) + if errRecovered != nil { + t.Fatalf("Pick() after recovery error = %v", errRecovered) + } + if recovered.ID != first.ID { + t.Fatalf("Pick() after recovery = %q, want original bound auth %q", recovered.ID, first.ID) + } +} + +func TestSessionAffinitySelector_OnResult_TransientVsTerminalClassification(t *testing.T) { + t.Parallel() + + type testCase struct { + name string + err *Error + wantRetain bool + } + + cases := []testCase{ + { + name: "500 Internal Server Error (transient)", + err: &Error{HTTPStatus: http.StatusInternalServerError, Message: "internal error"}, + wantRetain: true, + }, + { + name: "502 Bad Gateway (transient)", + err: &Error{HTTPStatus: http.StatusBadGateway, Message: "bad gateway"}, + wantRetain: true, + }, + { + name: "503 Service Unavailable (transient)", + err: &Error{HTTPStatus: http.StatusServiceUnavailable, Message: "service unavailable"}, + wantRetain: true, + }, + { + name: "504 Gateway Timeout (transient)", + err: &Error{HTTPStatus: http.StatusGatewayTimeout, Message: "gateway timeout"}, + wantRetain: true, + }, + { + name: "429 Too Many Requests / Quota (transient)", + err: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "rate limit exceeded"}, + wantRetain: true, + }, + { + name: "408 Request Timeout (transient)", + err: &Error{HTTPStatus: http.StatusRequestTimeout, Message: "request timeout"}, + wantRetain: true, + }, + { + name: "Cloudflare challenge (transient)", + err: &Error{HTTPStatus: http.StatusForbidden, Message: "just a moment... cloudflare challenge"}, + wantRetain: true, + }, + { + name: "400 Bad Request / client fault (skip cooldown, retain)", + err: &Error{HTTPStatus: http.StatusBadRequest, Message: `{"error":{"type":"invalid_request_error"}}`}, + wantRetain: true, + }, + { + name: "401 Unauthorized / invalid API key (terminal)", + err: &Error{HTTPStatus: http.StatusUnauthorized, Message: "invalid_api_key"}, + wantRetain: false, + }, + { + name: "402 Payment Required / out of credits (terminal)", + err: &Error{HTTPStatus: http.StatusPaymentRequired, Message: "insufficient balance"}, + wantRetain: false, + }, + { + name: "403 Forbidden / account disabled (terminal)", + err: &Error{HTTPStatus: http.StatusForbidden, Message: "account suspended"}, + wantRetain: false, + }, + { + name: "404 Not Found / unsupported model (terminal)", + err: &Error{HTTPStatus: http.StatusNotFound, Message: "model not found for plan"}, + wantRetain: false, + }, + { + name: "invalid_grant OAuth token revoked (terminal)", + err: &Error{HTTPStatus: http.StatusBadRequest, Message: `{"error":"invalid_grant"}`}, + wantRetain: false, + }, + { + name: "model_not_supported error (terminal)", + err: &Error{HTTPStatus: http.StatusBadRequest, Message: "requested model is not supported for your plan"}, + wantRetain: false, + }, + } + + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Hour, + }) + defer selector.Stop() + + sessionID := "sess-classify-" + tc.name + model := "test-model" + authID := "auth-target" + cacheKey := "mixed::header:" + sessionID + "::" + model + + selector.cache.Set(cacheKey, authID) + + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{sessionID}}, + Metadata: map[string]any{ + cliproxyexecutor.SessionAffinityProviderMetadataKey: "mixed", + cliproxyexecutor.SessionAffinityModelMetadataKey: model, + }, + } + + selector.OnResult(Result{ + AuthID: authID, + Provider: "claude", + Model: model, + Success: false, + Error: tc.err, + Options: opts, + }) + + got, ok := selector.cache.Get(cacheKey) + if tc.wantRetain { + if !ok || got != authID { + t.Fatalf("cache binding purged, want retained auth %q (ok=%v, got=%q)", authID, ok, got) + } + } else { + if ok { + t.Fatalf("cache binding unexpectedly retained %q, want purged", got) + } + } + }) + } +} + +func TestSessionAffinity_StickyTemporaryFallbackDuringPrimaryCooldown(t *testing.T) { + t.Parallel() + + fallback := &RoundRobinSelector{} + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: fallback, + TTL: time.Hour, + }) + defer selector.Stop() + + authA := &Auth{ID: "auth-a"} + authB := &Auth{ID: "auth-b"} + authC := &Auth{ID: "auth-c"} + allAuths := []*Auth{authA, authB, authC} + + sessionID := "sess-sticky-fallback-test" + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{sessionID}}, + } + + // 1. Initial Pick: auth-a is picked and bound as primary + first, err := selector.Pick(context.Background(), "claude", "claude-3", opts, allAuths) + if err != nil { + t.Fatalf("initial Pick() error = %v", err) + } + if first.ID != "auth-a" { + t.Fatalf("initial Pick() = %q, want auth-a", first.ID) + } + selector.OnResult(Result{AuthID: first.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) + + // 2. Primary auth-a is cooling down (only auth-b and auth-c available) + coolingAuths := []*Auth{authB, authC} + + // First pick during cooldown chooses a fallback auth (e.g. auth-b) + firstFallback, err := selector.Pick(context.Background(), "claude", "claude-3", opts, coolingAuths) + if err != nil { + t.Fatalf("first fallback Pick() error = %v", err) + } + firstFallbackID := firstFallback.ID + + // 5 consecutive picks during cooldown MUST all return the exact same fallback auth (sticky, no round-robin wandering) + for i := 1; i <= 5; i++ { + got, errPick := selector.Pick(context.Background(), "claude", "claude-3", opts, coolingAuths) + if errPick != nil { + t.Fatalf("consecutive fallback Pick %d error = %v", i, errPick) + } + if got.ID != firstFallbackID { + t.Fatalf("consecutive fallback Pick %d = %q, want sticky %q (prevent round-robin wandering during cooldown)", i, got.ID, firstFallbackID) + } + selector.OnResult(Result{AuthID: got.ID, Provider: "claude", Model: "claude-3", Options: opts, Success: true}) + } + + // 3. Primary auth-a recovers (allAuths available again) + recovered, errRecover := selector.Pick(context.Background(), "claude", "claude-3", opts, allAuths) + if errRecover != nil { + t.Fatalf("recovered Pick() error = %v", errRecover) + } + if recovered.ID != "auth-a" { + t.Fatalf("recovered Pick() = %q, want primary %q", recovered.ID, "auth-a") + } + + // 4. If primary auth-a cools down again, and first fallback auth (auth-b) ALSO fails/cools down: + onlyC := []*Auth{authC} + nextFallback, errNext := selector.Pick(context.Background(), "claude", "claude-3", opts, onlyC) + if errNext != nil { + t.Fatalf("fallback when B down Pick() error = %v", errNext) + } + if nextFallback.ID != "auth-c" { + t.Fatalf("fallback when B down = %q, want auth-c", nextFallback.ID) + } + + // Subsequent picks stick to auth-c + for i := 1; i <= 3; i++ { + got, errPick := selector.Pick(context.Background(), "claude", "claude-3", opts, onlyC) + if errPick != nil { + t.Fatalf("subsequent fallback C Pick %d error = %v", i, errPick) + } + if got.ID != "auth-c" { + t.Fatalf("subsequent fallback C Pick %d = %q, want auth-c", i, got.ID) + } + } +} + +func TestSessionAffinity_StickyTemporaryFallbackTTLExpiry(t *testing.T) { + t.Parallel() + + shortTTL := 50 * time.Millisecond + fallback := &RoundRobinSelector{} + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: fallback, + TTL: shortTTL, + }) + defer selector.Stop() + + authA := &Auth{ID: "auth-a"} + authB := &Auth{ID: "auth-b"} + authC := &Auth{ID: "auth-c"} + allAuths := []*Auth{authA, authB, authC} + + sessionID := "sess-sticky-fallback-ttl" + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{sessionID}}, + } + + // 1. Initial binding to auth-a + _, err := selector.Pick(context.Background(), "claude", "claude-3", opts, allAuths) + if err != nil { + t.Fatalf("initial Pick() error = %v", err) + } + + // 2. Primary cooling down, pick fallback + coolingAuths := []*Auth{authB, authC} + fb1, err := selector.Pick(context.Background(), "claude", "claude-3", opts, coolingAuths) + if err != nil { + t.Fatalf("fallback Pick() error = %v", err) + } + + // 3. Before TTL expires: sticks to fb1 + fb2, err := selector.Pick(context.Background(), "claude", "claude-3", opts, coolingAuths) + if err != nil { + t.Fatalf("consecutive fallback Pick() error = %v", err) + } + if fb2.ID != fb1.ID { + t.Fatalf("fallback Pick() before TTL = %q, want sticky %q", fb2.ID, fb1.ID) + } + + // 4. Wait for TTL to expire + time.Sleep(shortTTL * 2) + + // After TTL expires, temporary key is evicted and a new pick can be made + fb3, err := selector.Pick(context.Background(), "claude", "claude-3", opts, coolingAuths) + if err != nil { + t.Fatalf("fallback Pick() after TTL error = %v", err) + } + if fb3 == nil { + t.Fatal("fallback Pick() after TTL returned nil") + } +} + +func TestSessionAffinity_StickyTemporaryFallbackSharedAcrossAliases(t *testing.T) { + t.Parallel() + + fallback := &RoundRobinSelector{} + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: fallback, + TTL: time.Hour, + }) + defer selector.Stop() + + authA := &Auth{ID: "auth-a"} + authB := &Auth{ID: "auth-b"} + authC := &Auth{ID: "auth-c"} + allAuths := []*Auth{authA, authB, authC} + + // Payload with both prompt_cache_key (primary) and conversation.id (fallback alias) + payloadBoth := []byte(`{"prompt_cache_key":"pck-shared-test","conversation":{"id":"conv-shared-test"}}`) + optsBoth := cliproxyexecutor.Options{OriginalRequest: payloadBoth} + + // Payload with ONLY conversation.id (the alias) + payloadConvOnly := []byte(`{"conversation":{"id":"conv-shared-test"}}`) + optsConvOnly := cliproxyexecutor.Options{OriginalRequest: payloadConvOnly} + + // Payload with ONLY prompt_cache_key (the primary) + payloadPckOnly := []byte(`{"prompt_cache_key":"pck-shared-test"}`) + optsPckOnly := cliproxyexecutor.Options{OriginalRequest: payloadPckOnly} + + // 1. Initial request with both aliases establishes binding to auth-a on both + first, err := selector.Pick(context.Background(), "claude", "claude-3", optsBoth, allAuths) + if err != nil { + t.Fatalf("initial Pick() error = %v", err) + } + if first.ID != "auth-a" { + t.Fatalf("initial Pick() = %q, want auth-a", first.ID) + } + selector.OnResult(Result{AuthID: first.ID, Provider: "claude", Model: "claude-3", Options: optsBoth, Success: true}) + + // 2. Primary auth-a cools down (only auth-b and auth-c available) + coolingAuths := []*Auth{authB, authC} + + // Request under optsBoth selects fallback (e.g. auth-c) and must bind temp fallback to BOTH aliases + fallback1, err := selector.Pick(context.Background(), "claude", "claude-3", optsBoth, coolingAuths) + if err != nil { + t.Fatalf("first fallback Pick() error = %v", err) + } + expectedFallbackID := fallback1.ID + + // 3. Subsequent request identifying the SAME session via conversation alias ONLY + // Must hit the sticky temporary fallback (expectedFallbackID) rather than alternating via round-robin to another auth! + gotConv, errConv := selector.Pick(context.Background(), "claude", "claude-3", optsConvOnly, coolingAuths) + if errConv != nil { + t.Fatalf("Pick() with conversation alias error = %v", errConv) + } + if gotConv.ID != expectedFallbackID { + t.Fatalf("Pick() with conversation alias = %q, want sticky fallback %q (temporary fallback must be shared across aliases)", gotConv.ID, expectedFallbackID) + } + + // 4. Subsequent request identifying the session via prompt_cache_key ONLY + gotPck, errPck := selector.Pick(context.Background(), "claude", "claude-3", optsPckOnly, coolingAuths) + if errPck != nil { + t.Fatalf("Pick() with pck alias error = %v", errPck) + } + if gotPck.ID != expectedFallbackID { + t.Fatalf("Pick() with pck alias = %q, want sticky fallback %q", gotPck.ID, expectedFallbackID) + } + + // 5. Primary auth-a recovers: request under either alias returns auth-a and clears temporary fallbacks + recoveredConv, errRecConv := selector.Pick(context.Background(), "claude", "claude-3", optsConvOnly, allAuths) + if errRecConv != nil { + t.Fatalf("recovered Pick(conv) error = %v", errRecConv) + } + if recoveredConv.ID != "auth-a" { + t.Fatalf("recovered Pick(conv) = %q, want primary %q", recoveredConv.ID, "auth-a") + } + + recoveredPck, errRecPck := selector.Pick(context.Background(), "claude", "claude-3", optsPckOnly, allAuths) + if errRecPck != nil { + t.Fatalf("recovered Pick(pck) error = %v", errRecPck) + } + if recoveredPck.ID != "auth-a" { + t.Fatalf("recovered Pick(pck) = %q, want primary %q", recoveredPck.ID, "auth-a") + } +} + +func TestSessionAffinity_StickyTemporaryFallbackSecondaryBranch(t *testing.T) { + t.Parallel() + + fallback := &RoundRobinSelector{} + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: fallback, + TTL: time.Hour, + }) + defer selector.Stop() + + authA := &Auth{ID: "auth-a"} + authB := &Auth{ID: "auth-b"} + authC := &Auth{ID: "auth-c"} + allAuths := []*Auth{authA, authB, authC} + + payloadBoth := []byte(`{"prompt_cache_key":"pck-sec-branch","conversation":{"id":"conv-sec-branch"}}`) + optsBoth := cliproxyexecutor.Options{OriginalRequest: payloadBoth} + + payloadConvOnly := []byte(`{"conversation":{"id":"conv-sec-branch"}}`) + optsConvOnly := cliproxyexecutor.Options{OriginalRequest: payloadConvOnly} + + payloadPckOnly := []byte(`{"prompt_cache_key":"pck-sec-branch"}`) + optsPckOnly := cliproxyexecutor.Options{OriginalRequest: payloadPckOnly} + + // 1. Initial request binds both aliases to auth-a + first, err := selector.Pick(context.Background(), "claude", "claude-3", optsBoth, allAuths) + if err != nil { + t.Fatalf("initial Pick() error = %v", err) + } + if first.ID != "auth-a" { + t.Fatalf("initial Pick() = %q, want auth-a", first.ID) + } + selector.OnResult(Result{AuthID: first.ID, Provider: "claude", Model: "claude-3", Options: optsBoth, Success: true}) + + // 2. Primary auth-a cools down + coolingAuths := []*Auth{authB, authC} + + // Request under conv alias ONLY initiates fallback pick (secondary fallback branch) + fallbackFromConv, err := selector.Pick(context.Background(), "claude", "claude-3", optsConvOnly, coolingAuths) + if err != nil { + t.Fatalf("fallback Pick from conv alias error = %v", err) + } + expectedFallbackID := fallbackFromConv.ID + + // 3. Subsequent request under pck alias MUST return the same fallback + gotPck, errPck := selector.Pick(context.Background(), "claude", "claude-3", optsPckOnly, coolingAuths) + if errPck != nil { + t.Fatalf("Pick() with pck alias error = %v", errPck) + } + if gotPck.ID != expectedFallbackID { + t.Fatalf("Pick() with pck alias = %q, want %q", gotPck.ID, expectedFallbackID) + } + + // 4. Subsequent request under both aliases MUST return the same fallback + gotBoth, errBoth := selector.Pick(context.Background(), "claude", "claude-3", optsBoth, coolingAuths) + if errBoth != nil { + t.Fatalf("Pick() with both aliases error = %v", errBoth) + } + if gotBoth.ID != expectedFallbackID { + t.Fatalf("Pick() with both aliases = %q, want %q", gotBoth.ID, expectedFallbackID) + } + + // 5. Recovery of auth-a returns to auth-a under any alias + recovered, errRec := selector.Pick(context.Background(), "claude", "claude-3", optsBoth, allAuths) + if errRec != nil { + t.Fatalf("recovered Pick() error = %v", errRec) + } + if recovered.ID != "auth-a" { + t.Fatalf("recovered Pick() = %q, want %q", recovered.ID, "auth-a") + } +} + +func TestSessionAffinity_ModelFallbackSuffixDoesNotCollideWithTemporaryFallback(t *testing.T) { + t.Parallel() + + fallback := &RoundRobinSelector{} + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: fallback, + TTL: time.Hour, + }) + defer selector.Stop() + + authA := &Auth{ID: "auth-a"} + authB := &Auth{ID: "auth-b"} + authC := &Auth{ID: "auth-c"} + allAuths := []*Auth{authA, authB, authC} + + sessionID := "sess-fallback-suffix-collision" + opts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Session-Id": []string{sessionID}}, + } + + baseModel := "gpt-4" + suffixedModel := "gpt-4::fallback" + + // 1. Establish binding for suffixedModel ("gpt-4::fallback") to auth-a + pickSuffixed1, err := selector.Pick(context.Background(), "openai", suffixedModel, opts, allAuths) + if err != nil { + t.Fatalf("initial pick for suffixed model error = %v", err) + } + if pickSuffixed1.ID != "auth-a" { + t.Fatalf("initial pick for suffixed model = %q, want auth-a", pickSuffixed1.ID) + } + selector.OnResult(Result{AuthID: pickSuffixed1.ID, Provider: "openai", Model: suffixedModel, Options: opts, Success: true}) + + // 2. Establish binding for baseModel ("gpt-4") to auth-b + pickBase1, err := selector.Pick(context.Background(), "openai", baseModel, opts, []*Auth{authB, authC}) + if err != nil { + t.Fatalf("initial pick for base model error = %v", err) + } + if pickBase1.ID != "auth-b" { + t.Fatalf("initial pick for base model = %q, want auth-b", pickBase1.ID) + } + selector.OnResult(Result{AuthID: pickBase1.ID, Provider: "openai", Model: baseModel, Options: opts, Success: true}) + + // 3. Primary auth-b for baseModel cools down; temporary fallback for baseModel is selected as auth-c + coolingBaseAuths := []*Auth{authC} + pickBaseFallback, err := selector.Pick(context.Background(), "openai", baseModel, opts, coolingBaseAuths) + if err != nil { + t.Fatalf("fallback pick for base model error = %v", err) + } + if pickBaseFallback.ID != "auth-c" { + t.Fatalf("fallback pick for base model = %q, want auth-c", pickBaseFallback.ID) + } + selector.OnResult(Result{AuthID: pickBaseFallback.ID, Provider: "openai", Model: baseModel, Options: opts, Success: true}) + + // 4. Request for suffixedModel ("gpt-4::fallback") MUST NOT be overwritten or corrupted by baseModel's temporary fallback + pickSuffixed2, err := selector.Pick(context.Background(), "openai", suffixedModel, opts, allAuths) + if err != nil { + t.Fatalf("subsequent pick for suffixed model error = %v", err) + } + if pickSuffixed2.ID != "auth-a" { + t.Fatalf("subsequent pick for suffixed model = %q, want bound %q (must not collide with base model's temporary fallback)", pickSuffixed2.ID, "auth-a") + } + + // 5. When baseModel's primary auth-b recovers, invalidating baseModel's temporary fallback MUST NOT purge suffixedModel's binding + recoveredBase, err := selector.Pick(context.Background(), "openai", baseModel, opts, allAuths) + if err != nil { + t.Fatalf("recovered pick for base model error = %v", err) + } + if recoveredBase.ID != "auth-b" { + t.Fatalf("recovered pick for base model = %q, want auth-b", recoveredBase.ID) + } + + pickSuffixed3, err := selector.Pick(context.Background(), "openai", suffixedModel, opts, allAuths) + if err != nil { + t.Fatalf("pick for suffixed model after base recovery error = %v", err) + } + if pickSuffixed3.ID != "auth-a" { + t.Fatalf("pick for suffixed model after base recovery = %q, want bound %q (must not be purged when base model clears temporary fallback)", pickSuffixed3.ID, "auth-a") + } +} diff --git a/sdk/cliproxy/auth/session_cache.go b/sdk/cliproxy/auth/session_cache.go index b220990ab..570c1ca84 100644 --- a/sdk/cliproxy/auth/session_cache.go +++ b/sdk/cliproxy/auth/session_cache.go @@ -72,6 +72,23 @@ func (c *SessionCache) Get(sessionID string) (string, bool) { return "", false } +// Aliases returns all currently known alias identifiers for the session entry. +func (c *SessionCache) Aliases(sessionID string) []string { + if sessionID == "" { + return nil + } + now := time.Now() + c.mu.RLock() + defer c.mu.RUnlock() + entry, ok := c.entries[sessionID] + if !ok || !now.Before(entry.expiresAt) { + return nil + } + res := make([]string, len(entry.aliases)) + copy(res, entry.aliases) + return res +} + // 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) {