diff --git a/config.example.yaml b/config.example.yaml index d8fbed3c6..2db55c4d4 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -209,6 +209,9 @@ routing: # followed by prompt_cache_key, Responses conversation IDs, legacy body IDs, # execution or derived session identity, and the existing first-message hash fallback. # Automatic failover is always enabled when bound auth becomes unavailable. + # An established binding outranks credential priority: once a session is bound, that + # credential is kept even if a higher-priority credential recovers. Credential priority + # still decides cold bindings, requests without a session, and post-failover rebinding. session-affinity: false # default: false # How long session-to-auth bindings are retained. Default: 1h session-affinity-ttl: "1h" diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index 6a8562d12..43c763975 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -287,6 +287,14 @@ func (m *Manager) SetRoundTripperProvider(p RoundTripperProvider) { } func (m *Manager) availableAuthsForRouteModel(auths []*Auth, provider, routeModel string, now time.Time) ([]*Auth, error) { + return m.availableAuthsForRouteModelWithPriorityMode(auths, provider, routeModel, now, false) +} + +func (m *Manager) availableAuthsForRouteModelAcrossPriorities(auths []*Auth, provider, routeModel string, now time.Time) ([]*Auth, error) { + return m.availableAuthsForRouteModelWithPriorityMode(auths, provider, routeModel, now, true) +} + +func (m *Manager) availableAuthsForRouteModelWithPriorityMode(auths []*Auth, provider, routeModel string, now time.Time, allPriorities bool) ([]*Auth, error) { if len(auths) == 0 { return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"} } @@ -325,20 +333,32 @@ func (m *Manager) availableAuthsForRouteModel(auths []*Auth, provider, routeMode return nil, &Error{Code: "auth_unavailable", Message: "no auth available"} } - bestPriority := 0 - found := false - for priority := range availableByPriority { - if !found || priority > bestPriority { - bestPriority = priority - found = true + return availableAuthsFromPriorityBuckets(availableByPriority, allPriorities), nil +} + +// availableAuthsForSelector reports the candidates handed to priority-scoped consumers such as +// the plugin scheduler, plus the candidates handed to the configured selector. Both are equal +// unless session affinity is active, in which case the selector additionally receives lower +// priority tiers so an established binding can be validated instead of being preempted by a +// recovered higher-priority credential. +func (m *Manager) availableAuthsForSelector(selector Selector, auths []*Auth, provider, routeModel string, now time.Time) (priorityAuths, selectorAuths []*Auth, err error) { + if _, sessionAffinity := selector.(*SessionAffinitySelector); !sessionAffinity { + priorityAuths, err = m.availableAuthsForRouteModel(auths, provider, routeModel, now) + if err != nil { + return nil, nil, err } + priorityAuths = cloneAuthSlice(priorityAuths) + return priorityAuths, priorityAuths, nil } - available := availableByPriority[bestPriority] - if len(available) > 1 { - sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID }) + // One availability pass and one clone pass serve both lists: the highest priority tier is a + // subset of the across-priority candidates, so it is narrowed from the same cloned auths. + selectorAuths, err = m.availableAuthsForRouteModelAcrossPriorities(auths, provider, routeModel, now) + if err != nil { + return nil, nil, err } - return available, nil + selectorAuths = cloneAuthSlice(selectorAuths) + return highestPriorityAuths(selectorAuths), selectorAuths, nil } func selectionArgForSelector(selector Selector, routeModel string) string { @@ -1011,12 +1031,11 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op m.mu.RUnlock() return nil, nil, &Error{Code: "auth_not_found", Message: "no auth available"} } - available, errAvailable := m.availableAuthsForRouteModel(candidates, provider, model, time.Now()) + available, selectorAuths, errAvailable := m.availableAuthsForSelector(selector, candidates, provider, model, time.Now()) if errAvailable != nil { m.mu.RUnlock() return nil, nil, errAvailable } - available = cloneAuthSlice(available) m.mu.RUnlock() selected, handled, errPick := m.pickViaPluginScheduler(ctx, pluginScheduler, provider, []string{provider}, model, opts, tried, available) @@ -1025,7 +1044,7 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op } if !handled { selectorCtx := withWeightedSelectorStateModel(ctx, selector, model) - selected, errPick = selector.Pick(selectorCtx, provider, selectionArgForSelector(selector, model), opts, available) + selected, errPick = selector.Pick(selectorCtx, provider, selectionArgForSelector(selector, model), opts, selectorAuths) if errPick != nil { return nil, nil, errPick } @@ -1329,12 +1348,11 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m m.mu.RUnlock() return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} } - available, errAvailable := m.availableAuthsForRouteModel(candidates, "mixed", model, time.Now()) + available, selectorAuths, errAvailable := m.availableAuthsForSelector(selector, candidates, "mixed", model, time.Now()) if errAvailable != nil { m.mu.RUnlock() return nil, nil, "", errAvailable } - available = cloneAuthSlice(available) m.mu.RUnlock() selected, handled, errPick := m.pickViaPluginScheduler(ctx, pluginScheduler, "mixed", providers, model, opts, tried, available) @@ -1343,7 +1361,7 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m } if !handled { selectorCtx := withWeightedSelectorStateModel(ctx, selector, model) - selected, errPick = selector.Pick(selectorCtx, "mixed", selectionArgForSelector(selector, model), opts, available) + selected, errPick = selector.Pick(selectorCtx, "mixed", selectionArgForSelector(selector, model), opts, selectorAuths) if errPick != nil { return nil, nil, "", errPick } diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 00656f42d..bf6fca761 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -269,6 +269,14 @@ func collectAvailableByPriority(auths []*Auth, model string, now time.Time) (ava } func getAvailableAuths(auths []*Auth, provider, model string, now time.Time) ([]*Auth, error) { + return getAvailableAuthsWithPriorityMode(auths, provider, model, now, false) +} + +func getAvailableAuthsAcrossPriorities(auths []*Auth, provider, model string, now time.Time) ([]*Auth, error) { + return getAvailableAuthsWithPriorityMode(auths, provider, model, now, true) +} + +func getAvailableAuthsWithPriorityMode(auths []*Auth, provider, model string, now time.Time, allPriorities bool) ([]*Auth, error) { if len(auths) == 0 { return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"} } @@ -289,20 +297,73 @@ func getAvailableAuths(auths []*Auth, provider, model string, now time.Time) ([] return nil, &Error{Code: "auth_unavailable", Message: "no auth available"} } + return availableAuthsFromPriorityBuckets(availableByPriority, allPriorities), nil +} + +// availableAuthsFromPriorityBuckets flattens availability buckets into a stable, ID-sorted slice. +// When allPriorities is false only the highest available priority tier is returned. +// When allPriorities is true every tier is merged, so the result carries no priority ordering: +// use it for membership checks or feed it to highestPriorityAuths, never as a priority-ordered +// selection order. +func availableAuthsFromPriorityBuckets(availableByPriority map[int][]*Auth, allPriorities bool) []*Auth { + var candidates []*Auth + if allPriorities { + total := 0 + for _, bucket := range availableByPriority { + total += len(bucket) + } + candidates = make([]*Auth, 0, total) + for _, bucket := range availableByPriority { + candidates = append(candidates, bucket...) + } + } else { + bestPriority := 0 + found := false + for priority := range availableByPriority { + if !found || priority > bestPriority { + bestPriority = priority + found = true + } + } + bucket := availableByPriority[bestPriority] + candidates = make([]*Auth, 0, len(bucket)) + candidates = append(candidates, bucket...) + } + if len(candidates) > 1 { + sort.Slice(candidates, func(i, j int) bool { return candidates[i].ID < candidates[j].ID }) + } + return candidates +} + +// highestPriorityAuths narrows an availability slice to its highest priority tier while +// preserving the input order. The input slice is returned unchanged when every candidate +// already shares the highest priority, so the common single-tier case allocates nothing. +func highestPriorityAuths(auths []*Auth) []*Auth { + if len(auths) <= 1 { + return auths + } bestPriority := 0 - found := false - for priority := range availableByPriority { - if !found || priority > bestPriority { + bestCount := 0 + for _, auth := range auths { + priority := authPriority(auth) + switch { + case bestCount == 0 || priority > bestPriority: bestPriority = priority - found = true + bestCount = 1 + case priority == bestPriority: + bestCount++ } } - - available := availableByPriority[bestPriority] - if len(available) > 1 { - sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID }) + if bestCount == len(auths) { + return auths } - return available, nil + highest := make([]*Auth, 0, bestCount) + for _, auth := range auths { + if authPriority(auth) == bestPriority { + highest = append(highest, auth) + } + } + return highest } // Pick selects the next available auth for the provider in a round-robin manner. @@ -567,26 +628,38 @@ func NewSessionAffinitySelectorWithConfig(cfg SessionAffinityConfig) *SessionAff // Explicit Claude Code, Codex, OpenCode, pi, and request-body session signals // precede execution metadata, stable derived identity, and the legacy hash fallback. // +// An established binding outranks credential priority: a bound credential that is still +// available is reused even when a higher-priority credential recovers. Credential priority +// applies to cold bindings, requests without a session, and genuine bound-credential +// failover, so the fallback selector only ever receives the highest available priority tier. +// // Note: The cache key includes provider, session ID, and model to handle cases where // a session uses multiple models (e.g., gemini-2.5-pro and gemini-3-flash-preview) // that may be supported by different auth credentials, and to avoid cross-provider conflicts. func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { entry := selectorLogEntry(ctx) primaryID, fallbackID := extractSessionIDs(opts.Headers, opts.OriginalRequest, opts.Metadata) - if primaryID == "" { - entry.Debugf("session-affinity: no session ID extracted, falling back to default selector | provider=%s model=%s", provider, model) - return s.fallback.Pick(ctx, provider, model, opts, auths) - } - now := time.Now() availabilityCandidates := auths if _, weighted := s.fallback.(*WeightedRoundRobinSelector); weighted { availabilityCandidates = positiveWeightAuths(auths) } - available, err := getAvailableAuths(availabilityCandidates, provider, model, now) + if primaryID == "" { + fallbackAuths, errAvailable := getAvailableAuths(availabilityCandidates, provider, model, now) + if errAvailable != nil { + return nil, errAvailable + } + entry.Debugf("session-affinity: no session ID extracted, falling back to default selector | provider=%s model=%s", provider, model) + return s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) + } + + // A single availability pass serves both lookups: the bound credential is validated against + // every priority tier, while the fallback selector keeps seeing only the highest tier. + available, err := getAvailableAuthsAcrossPriorities(availabilityCandidates, provider, model, now) if err != nil { return nil, err } + fallbackAuths := highestPriorityAuths(available) cacheKey := provider + "::" + primaryID + "::" + model fallbackKey := "" @@ -610,7 +683,7 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri } } // Cached auth not available, reselect via fallback selector for even distribution - auth, err := s.fallback.Pick(ctx, provider, model, opts, auths) + auth, err := s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) if err != nil { return nil, err } @@ -631,7 +704,7 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri } } - auth, err := s.fallback.Pick(ctx, provider, model, opts, auths) + auth, err := s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) if err != nil { return nil, err } diff --git a/sdk/cliproxy/auth/session_affinity_priority_test.go b/sdk/cliproxy/auth/session_affinity_priority_test.go new file mode 100644 index 000000000..adb1c67bf --- /dev/null +++ b/sdk/cliproxy/auth/session_affinity_priority_test.go @@ -0,0 +1,178 @@ +package auth + +import ( + "context" + "net/http" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +func TestManagerSessionAffinityPreservesBindingAcrossHigherPriorityRecovery(t *testing.T) { + for _, testCase := range []struct { + name string + providerSuffix string + pick func(*Manager, context.Context, string, string, cliproxyexecutor.Options) (*Auth, error) + }{ + { + name: "single provider", + providerSuffix: "single", + pick: func(manager *Manager, ctx context.Context, provider, model string, opts cliproxyexecutor.Options) (*Auth, error) { + auth, _, errPick := manager.pickNext(ctx, provider, model, opts, nil) + return auth, errPick + }, + }, + { + name: "mixed provider", + providerSuffix: "mixed", + pick: func(manager *Manager, ctx context.Context, provider, model string, opts cliproxyexecutor.Options) (*Auth, error) { + auth, _, _, errPick := manager.pickNextMixed(ctx, []string{provider}, model, opts, nil) + return auth, errPick + }, + }, + } { + t.Run(testCase.name, func(t *testing.T) { + ctx := context.Background() + provider := "affinity-priority-" + testCase.providerSuffix + model := "affinity-priority-model" + highID := provider + "-high" + lowID := provider + "-low" + + manager := NewManager(nil, nil, nil) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Hour, + }) + defer affinity.Stop() + manager.SetSelector(affinity) + manager.RegisterExecutor(schedulerTestExecutor{provider: provider}) + + for _, auth := range []*Auth{ + {ID: highID, Provider: provider, Status: StatusActive, Attributes: map[string]string{"priority": "1"}}, + {ID: lowID, Provider: provider, Status: StatusActive, Attributes: map[string]string{"priority": "0"}}, + } { + if _, errRegister := manager.Register(WithSkipPersist(ctx), auth); errRegister != nil { + t.Fatalf("Register(%s): %v", auth.ID, errRegister) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + } + + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.DerivedSessionIDMetadataKey: "stable-session", + }} + pick := func(pickOpts cliproxyexecutor.Options) *Auth { + t.Helper() + auth, errPick := testCase.pick(manager, ctx, provider, model, pickOpts) + if errPick != nil { + t.Fatalf("pick: %v", errPick) + } + if auth == nil { + t.Fatal("pick returned nil auth") + } + return auth + } + + if got := pick(opts); got.ID != highID { + t.Fatalf("cold binding = %q, want high priority %q", got.ID, highID) + } + + manager.MarkResult(ctx, Result{ + AuthID: highID, + Provider: provider, + Model: model, + Success: false, + Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "quota"}, + }) + if got := pick(opts); got.ID != lowID { + t.Fatalf("failover binding = %q, want %q", got.ID, lowID) + } + + expireSessionAffinityPriorityModelCooldown(t, manager, highID, model) + if got := pick(opts); got.ID != lowID { + t.Fatalf("binding after higher-priority recovery = %q, want sticky %q", got.ID, lowID) + } + + newSessionOpts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.DerivedSessionIDMetadataKey: "new-session", + }} + if got := pick(newSessionOpts); got.ID != highID { + t.Fatalf("cold binding for new session = %q, want high priority %q", got.ID, highID) + } + + manager.MarkResult(ctx, Result{ + AuthID: lowID, + Provider: provider, + Model: model, + Success: false, + Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "quota"}, + }) + if got := pick(opts); got.ID != highID { + t.Fatalf("binding after bound auth became unavailable = %q, want %q", got.ID, highID) + } + }) + } +} + +func TestSessionAffinityFallbackOnlyReceivesHighestAvailablePriority(t *testing.T) { + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: lastAuthSelector{}, + TTL: time.Hour, + }) + defer selector.Stop() + + high := &Auth{ID: "a-high", Provider: "test", Status: StatusActive, Attributes: map[string]string{"priority": "1"}} + low := &Auth{ID: "z-low", Provider: "test", Status: StatusActive, Attributes: map[string]string{"priority": "0"}} + auths := []*Auth{high, low} + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.DerivedSessionIDMetadataKey: "stable-session", + }} + + assertPick := func(label string, pickOpts cliproxyexecutor.Options, wantID string) { + t.Helper() + got, errPick := selector.Pick(context.Background(), "test", "model", pickOpts, auths) + if errPick != nil { + t.Fatalf("%s: %v", label, errPick) + } + if got == nil { + t.Fatalf("%s = nil, want %q", label, wantID) + } + if got.ID != wantID { + t.Fatalf("%s = %q, want %q", label, got.ID, wantID) + } + } + + assertPick("cold binding", opts, high.ID) + assertPick("no-session fallback", cliproxyexecutor.Options{}, high.ID) + + high.Unavailable = true + assertPick("fallback after bound auth became unavailable", opts, low.ID) +} + +type lastAuthSelector struct{} + +func (lastAuthSelector) Pick(_ context.Context, _, _ string, _ cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { + if len(auths) == 0 { + return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"} + } + return auths[len(auths)-1], nil +} + +func expireSessionAffinityPriorityModelCooldown(t *testing.T, manager *Manager, authID, model string) { + t.Helper() + manager.mu.Lock() + defer manager.mu.Unlock() + auth := manager.auths[authID] + if auth == nil { + t.Fatalf("auth %q not found", authID) + } + state := auth.ModelStates[model] + if state == nil { + t.Fatalf("model state %q not found for auth %q", model, authID) + } + expired := time.Now().Add(-time.Second) + state.NextRetryAfter = expired + state.Quota.NextRecoverAt = expired +}