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.
This commit is contained in:
VooDisss
2026-03-27 18:52:58 +02:00
parent 4c4cbd44da
commit 6962e09dd9
3 changed files with 50 additions and 16 deletions
+30 -10
View File
@@ -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)
}