Skip to content

Invalidate session affinity on auth-wide failures and force fresh upstream sessions #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
clcc2019 merged 1 commit into from
Jun 3, 2026
Merged

Invalidate session affinity on auth-wide failures and force fresh upstream sessions #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 63 additions & 0 deletions sdk/cliproxy/auth/conductor.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
}
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
}
Expand All @@ -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
Expand Down
58 changes: 58 additions & 0 deletions sdk/cliproxy/auth/conductor_overrides_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
15 changes: 15 additions & 0 deletions sdk/cliproxy/auth/selector.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand All @@ -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
Expand Down
46 changes: 39 additions & 7 deletions sdk/cliproxy/auth/session_cache.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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()
}
Loading