From 6962e09dd945120760b3ae3166d9e9480a1afd7a Mon Sep 17 00:00:00 2001 From: VooDisss Date: Fri, 27 Mar 2026 18:52:58 +0200 Subject: [PATCH] fix(auth): scope affinity by provider Keep sticky auth affinity limited to matching providers and stop persisting execution-session IDs as long-lived affinity keys so provider switching and normal streaming traffic do not create incorrect pins or stale affinity state. --- sdk/api/handlers/handlers.go | 1 - sdk/cliproxy/auth/conductor.go | 40 +++++++++++++++----- sdk/cliproxy/auth/conductor_affinity_test.go | 25 +++++++++--- 3 files changed, 50 insertions(+), 16 deletions(-) diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go index 420d1fcc..69dd6007 100644 --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -210,7 +210,6 @@ func requestExecutionMetadata(ctx context.Context) map[string]any { } if executionSessionID := executionSessionIDFromContext(ctx); executionSessionID != "" { meta[coreexecutor.ExecutionSessionMetadataKey] = executionSessionID - meta[authAffinityMetadataKey] = executionSessionID } else if ctx != nil { if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil { if apiKey, exists := ginCtx.Get("apiKey"); exists && apiKey != nil { diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 147b0ece..5c38654b 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -2246,8 +2246,17 @@ func authAffinityKeyFromMetadata(meta map[string]any) string { } } -func (m *Manager) AuthAffinity(key string) string { +func scopedAuthAffinityKey(provider, key string) string { + provider = strings.TrimSpace(strings.ToLower(provider)) key = strings.TrimSpace(key) + if provider == "" || key == "" { + return "" + } + return provider + "|" + key +} + +func (m *Manager) AuthAffinity(provider, key string) string { + key = scopedAuthAffinityKey(provider, key) if m == nil || key == "" { return "" } @@ -2256,12 +2265,12 @@ func (m *Manager) AuthAffinity(key string) string { return strings.TrimSpace(m.affinity[key]) } -func (m *Manager) applyAuthAffinity(opts *cliproxyexecutor.Options) { +func (m *Manager) applyAuthAffinity(provider string, opts *cliproxyexecutor.Options) { if m == nil || opts == nil || pinnedAuthIDFromMetadata(opts.Metadata) != "" { return } if affinityKey := authAffinityKeyFromMetadata(opts.Metadata); affinityKey != "" { - if affinityAuthID := m.AuthAffinity(affinityKey); affinityAuthID != "" { + if affinityAuthID := m.AuthAffinity(provider, affinityKey); affinityAuthID != "" { if opts.Metadata == nil { opts.Metadata = make(map[string]any) } @@ -2275,15 +2284,15 @@ func (m *Manager) persistAuthAffinity(entry *log.Entry, opts cliproxyexecutor.Op return } if affinityKey := authAffinityKeyFromMetadata(opts.Metadata); affinityKey != "" { - m.SetAuthAffinity(affinityKey, authID) + m.SetAuthAffinity(provider, affinityKey, authID) if entry != nil && log.IsLevelEnabled(log.DebugLevel) { entry.Debugf("auth affinity pinned auth_id=%s provider=%s model=%s", authID, provider, model) } } } -func (m *Manager) SetAuthAffinity(key, authID string) { - key = strings.TrimSpace(key) +func (m *Manager) SetAuthAffinity(provider, key, authID string) { + key = scopedAuthAffinityKey(provider, key) authID = strings.TrimSpace(authID) if m == nil || key == "" || authID == "" { return @@ -2296,8 +2305,8 @@ func (m *Manager) SetAuthAffinity(key, authID string) { m.affinityMu.Unlock() } -func (m *Manager) ClearAuthAffinity(key string) { - key = strings.TrimSpace(key) +func (m *Manager) ClearAuthAffinity(provider, key string) { + key = scopedAuthAffinityKey(provider, key) if m == nil || key == "" { return } @@ -2389,7 +2398,7 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op } func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) { - m.applyAuthAffinity(&opts) + m.applyAuthAffinity(provider, &opts) if !m.useSchedulerFastPath() { return m.pickNextLegacy(ctx, provider, model, opts, tried) } @@ -2504,7 +2513,18 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m } func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) { - m.applyAuthAffinity(&opts) + if pinnedAuthIDFromMetadata(opts.Metadata) == "" { + for _, provider := range providers { + providerKey := strings.TrimSpace(strings.ToLower(provider)) + if providerKey == "" { + continue + } + m.applyAuthAffinity(providerKey, &opts) + if pinnedAuthIDFromMetadata(opts.Metadata) != "" { + break + } + } + } if !m.useSchedulerFastPath() { return m.pickNextMixedLegacy(ctx, providers, model, opts, tried) } diff --git a/sdk/cliproxy/auth/conductor_affinity_test.go b/sdk/cliproxy/auth/conductor_affinity_test.go index e84f7c96..363e2367 100644 --- a/sdk/cliproxy/auth/conductor_affinity_test.go +++ b/sdk/cliproxy/auth/conductor_affinity_test.go @@ -52,7 +52,7 @@ func TestManagerPickNextMixedUsesAuthAffinity(t *testing.T) { t.Fatalf("Register(codex-b) error = %v", errRegister) } - manager.SetAuthAffinity("idem-1", "codex-b") + manager.SetAuthAffinity("codex", "idem-1", "codex-b") opts := cliproxyexecutor.Options{Metadata: map[string]any{"auth_affinity_key": "idem-1"}} got, _, provider, errPick := manager.pickNextMixed(context.Background(), []string{"codex"}, "gpt-5.4", opts, map[string]struct{}{}) @@ -74,12 +74,27 @@ func TestManagerAuthAffinityRoundTrip(t *testing.T) { t.Parallel() manager := NewManager(nil, nil, nil) - manager.SetAuthAffinity("idem-2", "auth-1") - if got := manager.AuthAffinity("idem-2"); got != "auth-1" { + manager.SetAuthAffinity("codex", "idem-2", "auth-1") + if got := manager.AuthAffinity("codex", "idem-2"); got != "auth-1" { t.Fatalf("AuthAffinity = %q, want %q", got, "auth-1") } - manager.ClearAuthAffinity("idem-2") - if got := manager.AuthAffinity("idem-2"); got != "" { + manager.ClearAuthAffinity("codex", "idem-2") + if got := manager.AuthAffinity("codex", "idem-2"); got != "" { t.Fatalf("AuthAffinity after clear = %q, want empty", got) } } + +func TestManagerAuthAffinityScopedByProvider(t *testing.T) { + t.Parallel() + + manager := NewManager(nil, nil, nil) + manager.SetAuthAffinity("codex", "shared-key", "codex-auth") + manager.SetAuthAffinity("gemini", "shared-key", "gemini-auth") + + if got := manager.AuthAffinity("codex", "shared-key"); got != "codex-auth" { + t.Fatalf("codex affinity = %q, want %q", got, "codex-auth") + } + if got := manager.AuthAffinity("gemini", "shared-key"); got != "gemini-auth" { + t.Fatalf("gemini affinity = %q, want %q", got, "gemini-auth") + } +}