diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index ce3a1341df8..65c709dac5f 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -3514,6 +3514,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { schedulerDirty := false persistAuthID := "" var schedulerSnapshot *Auth + invalidateAuthAffinity := false m.mu.Lock() if auth, ok := m.auths[result.AuthID]; ok && auth != nil { @@ -3648,6 +3649,9 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { if result.AuthScoped { auth.Quota.AuthScope = true } + if auth.Unavailable && !auth.NextRetryAfter.IsZero() && auth.NextRetryAfter.After(now) { + invalidateAuthAffinity = true + } if result.Model != "" { state := ensureModelState(auth, result.Model) state.Unavailable = true @@ -3683,6 +3687,9 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { if m.scheduler != nil && schedulerSnapshot != nil { m.scheduler.upsertAuth(schedulerSnapshot) } + if invalidateAuthAffinity { + m.invalidateSessionAffinityForAuth(result.AuthID) + } if clearModelQuota && result.Model != "" { registry.GetGlobalRegistry().ClearModelQuotaExceeded(result.AuthID, result.Model) @@ -3703,6 +3710,32 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { m.hook.OnResult(ctx, result) } +func (m *Manager) invalidateSessionAffinityForAuth(authID string) { + if m == nil || strings.TrimSpace(authID) == "" { + return + } + seen := make(map[*SessionAffinitySelector]struct{}, 2) + invalidate := func(selector Selector) { + affinity, ok := selector.(*SessionAffinitySelector) + if !ok || affinity == nil { + return + } + if _, okSeen := seen[affinity]; okSeen { + return + } + seen[affinity] = struct{}{} + affinity.InvalidateAuth(authID) + } + + m.mu.RLock() + selector := m.selector + kiroSelector := m.kiroSelector + m.mu.RUnlock() + + invalidate(selector) + invalidate(kiroSelector) +} + func (m *Manager) markCleanModelSuccessResult(ctx context.Context, result Result) bool { if m == nil || !result.Success || result.AuthID == "" || result.Model == "" { return false @@ -4887,6 +4920,7 @@ func (m *Manager) pickNextSingleWithSchedulerAffinity(ctx context.Context, affin } cacheKey := sessionAffinityCacheKey(provider, primaryID, opts.Metadata) + forceFreshUpstream := false if cachedAuthID, ok := affinity.cache.GetAndRefresh(cacheKey); ok { if auth, executor, okCached, errPick := m.pickCachedSingleWithScheduler(ctx, provider, model, opts, tried, cachedAuthID); errPick != nil || okCached { return auth, executor, errPick @@ -4895,14 +4929,24 @@ func (m *Manager) pickNextSingleWithSchedulerAffinity(ctx context.Context, affin forceNewUpstreamSessionForNextCredential(&opts) } + if affinity.cache.ConsumeForceNew(cacheKey) { + forceFreshUpstream = true + } + if fallbackID != "" && fallbackID != primaryID { fallbackKey := sessionAffinityCacheKey(provider, fallbackID, opts.Metadata) + if affinity.cache.ConsumeForceNew(fallbackKey) { + forceFreshUpstream = true + } if cachedAuthID, ok := affinity.cache.Get(fallbackKey); ok { auth, executor, okCached, errPick := m.pickCachedSingleWithScheduler(ctx, provider, model, opts, tried, cachedAuthID) if errPick != nil { return auth, executor, errPick } if okCached { + if forceFreshUpstream { + forceNewUpstreamSessionForNextCredential(&opts) + } affinity.cache.Set(cacheKey, auth.ID) return auth, executor, nil } @@ -4911,6 +4955,10 @@ func (m *Manager) pickNextSingleWithSchedulerAffinity(ctx context.Context, affin } } + if forceFreshUpstream { + forceNewUpstreamSessionForNextCredential(&opts) + } + auth, executor, errPick := m.pickNextSingleStableWithScheduler(ctx, provider, model, opts, tried, cacheKey) if errPick != nil { return nil, nil, errPick @@ -5141,6 +5189,7 @@ func (m *Manager) pickNextMixedWithSchedulerAffinity(ctx context.Context, affini } cacheKey := sessionAffinityCacheKey("mixed", primaryID, opts.Metadata) + forceFreshUpstream := false if cachedAuthID, ok := affinity.cache.GetAndRefresh(cacheKey); ok { if auth, executor, provider, okCached, errPick := m.pickCachedMixedWithScheduler(ctx, providers, model, opts, tried, cachedAuthID); errPick != nil || okCached { return auth, executor, provider, errPick @@ -5149,14 +5198,24 @@ func (m *Manager) pickNextMixedWithSchedulerAffinity(ctx context.Context, affini forceNewUpstreamSessionForNextCredential(&opts) } + if affinity.cache.ConsumeForceNew(cacheKey) { + forceFreshUpstream = true + } + if fallbackID != "" && fallbackID != primaryID { fallbackKey := sessionAffinityCacheKey("mixed", fallbackID, opts.Metadata) + if affinity.cache.ConsumeForceNew(fallbackKey) { + forceFreshUpstream = true + } if cachedAuthID, ok := affinity.cache.Get(fallbackKey); ok { auth, executor, provider, okCached, errPick := m.pickCachedMixedWithScheduler(ctx, providers, model, opts, tried, cachedAuthID) if errPick != nil { return auth, executor, provider, errPick } if okCached { + if forceFreshUpstream { + forceNewUpstreamSessionForNextCredential(&opts) + } affinity.cache.Set(cacheKey, auth.ID) return auth, executor, provider, nil } @@ -5165,6 +5224,10 @@ func (m *Manager) pickNextMixedWithSchedulerAffinity(ctx context.Context, affini } } + if forceFreshUpstream { + forceNewUpstreamSessionForNextCredential(&opts) + } + auth, executor, provider, errPick := m.pickNextMixedStableWithScheduler(ctx, providers, model, opts, tried, cacheKey) if errPick != nil { return nil, nil, "", errPick diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go index 85d633be202..1b09d348dce 100644 --- a/sdk/cliproxy/auth/conductor_overrides_test.go +++ b/sdk/cliproxy/auth/conductor_overrides_test.go @@ -564,6 +564,64 @@ func TestManager_CachedUnavailableAuthReselectForcesNewUpstreamSession(t *testin } } +func TestManager_AuthWideFailureInvalidatesAffinityAndForcesFreshSession(t *testing.T) { + affinity := NewSessionAffinitySelector(&RoundRobinSelector{}) + m := NewManager(nil, affinity, nil) + m.SetRetryConfig(0, 0, 0) + executor := &unauthorizedFailoverSessionExecutor{} + m.RegisterExecutor(executor) + + baseID := uuid.NewString() + auth1 := &Auth{ID: baseID + "-auth-1", Provider: "codex"} + auth2 := &Auth{ID: baseID + "-auth-2", Provider: "codex"} + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth1.ID, "codex", []*registry.ModelInfo{{ID: "test-model"}}) + reg.RegisterClient(auth2.ID, "codex", []*registry.ModelInfo{{ID: "test-model"}}) + t.Cleanup(func() { + reg.UnregisterClient(auth1.ID) + reg.UnregisterClient(auth2.ID) + affinity.Stop() + }) + + if _, errRegister := m.Register(context.Background(), auth1); errRegister != nil { + t.Fatalf("register auth1: %v", errRegister) + } + if _, errRegister := m.Register(context.Background(), auth2); errRegister != nil { + t.Fatalf("register auth2: %v", errRegister) + } + + metadata := map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "session-1", + } + cacheKey := "providers:codex::exec:session-1" + affinity.cache.Set(cacheKey, auth1.ID) + m.MarkResult(context.Background(), Result{ + AuthID: auth1.ID, + Provider: "codex", + Model: "test-model", + Success: false, + Error: &Error{HTTPStatus: http.StatusUnauthorized, Message: "unauthorized"}, + }) + if cachedAuthID, ok := affinity.cache.Get(cacheKey); ok { + t.Fatalf("expected auth-wide failure to invalidate affinity cache, got %q", cachedAuthID) + } + + _, errExecute := m.Execute(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{Metadata: metadata}) + if errExecute != nil { + t.Fatalf("Execute error: %v", errExecute) + } + forced := executor.ForcedSessions() + if len(forced) != 1 { + t.Fatalf("forced sessions = %#v, want one call", forced) + } + if strings.TrimSpace(forced[0]) == "" { + t.Fatalf("first call after auth-wide affinity invalidation should force a fresh upstream session: %#v", forced) + } + if cachedAuthID, ok := affinity.cache.Get(cacheKey); !ok || cachedAuthID != auth2.ID { + t.Fatalf("expected session to rebind to auth2, got auth=%q ok=%v", cachedAuthID, ok) + } +} + func TestManager_OuterRetryPreservesForcedUpstreamSession(t *testing.T) { m := NewManager(nil, nil, nil) m.SetRetryConfig(1, 100*time.Millisecond, 1) diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 5740a4a8bcf..43d775eb175 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -585,6 +585,7 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri } cacheKey := sessionAffinityCacheKey(provider, primaryID, opts.Metadata) + forceFreshUpstream := false if cachedAuthID, ok := s.cache.GetAndRefresh(cacheKey); ok { for _, auth := range available { @@ -604,11 +605,21 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri return auth, nil } + if s.cache.ConsumeForceNew(cacheKey) { + forceFreshUpstream = true + } + if fallbackID != "" && fallbackID != primaryID { fallbackKey := sessionAffinityCacheKey(provider, fallbackID, opts.Metadata) + if s.cache.ConsumeForceNew(fallbackKey) { + forceFreshUpstream = true + } if cachedAuthID, ok := s.cache.Get(fallbackKey); ok { for _, auth := range available { if auth.ID == cachedAuthID { + if forceFreshUpstream { + forceNewUpstreamSessionForNextCredential(&opts) + } s.cache.Set(cacheKey, auth.ID) infoLog("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 @@ -618,6 +629,10 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri } } + if forceFreshUpstream { + forceNewUpstreamSessionForNextCredential(&opts) + } + auth, err := s.fallback.Pick(ctx, provider, model, opts, auths) if err != nil { return nil, err diff --git a/sdk/cliproxy/auth/session_cache.go b/sdk/cliproxy/auth/session_cache.go index ba44d93608c..1434587db39 100644 --- a/sdk/cliproxy/auth/session_cache.go +++ b/sdk/cliproxy/auth/session_cache.go @@ -13,10 +13,11 @@ type sessionEntry struct { // SessionCache provides TTL-based session to auth mapping with automatic cleanup. type SessionCache struct { - mu sync.RWMutex - entries map[string]sessionEntry - ttl time.Duration - stopCh chan struct{} + mu sync.RWMutex + entries map[string]sessionEntry + forceNew map[string]time.Time + ttl time.Duration + stopCh chan struct{} } // NewSessionCache creates a cache with the specified TTL. @@ -26,9 +27,10 @@ func NewSessionCache(ttl time.Duration) *SessionCache { ttl = 30 * time.Minute } c := &SessionCache{ - entries: make(map[string]sessionEntry), - ttl: ttl, - stopCh: make(chan struct{}), + entries: make(map[string]sessionEntry), + forceNew: make(map[string]time.Time), + ttl: ttl, + stopCh: make(chan struct{}), } go c.cleanupLoop() return c @@ -107,15 +109,40 @@ func (c *SessionCache) InvalidateAuth(authID string) { if authID == "" { return } + now := time.Now() c.mu.Lock() + if c.forceNew == nil { + c.forceNew = make(map[string]time.Time) + } for sid, entry := range c.entries { if entry.authID == authID { delete(c.entries, sid) + if entry.expiresAt.After(now) { + c.forceNew[sid] = entry.expiresAt + } } } c.mu.Unlock() } +// ConsumeForceNew reports whether the session was recently unbound from an +// unavailable auth and clears the marker. Callers use this to start a fresh +// upstream provider session when the same downstream session is rebound to a +// different credential. +func (c *SessionCache) ConsumeForceNew(sessionID string) bool { + if sessionID == "" { + return false + } + now := time.Now() + c.mu.Lock() + expiresAt, ok := c.forceNew[sessionID] + if ok { + delete(c.forceNew, sessionID) + } + c.mu.Unlock() + return ok && expiresAt.After(now) +} + // Stop terminates the background cleanup goroutine. func (c *SessionCache) Stop() { select { @@ -146,5 +173,10 @@ func (c *SessionCache) cleanup() { delete(c.entries, sid) } } + for sid, expiresAt := range c.forceNew { + if now.After(expiresAt) { + delete(c.forceNew, sid) + } + } c.mu.Unlock() }