diff --git a/sdk/cliproxy/auth/scheduler.go b/sdk/cliproxy/auth/scheduler.go index 8bec6123b..18471444f 100644 --- a/sdk/cliproxy/auth/scheduler.go +++ b/sdk/cliproxy/auth/scheduler.go @@ -127,10 +127,18 @@ func restoreReadyViewCursors(view *readyView, state readyViewCursorState) { view.cursor = normalizeCursor(state.cursor, len(view.flat)) } weights := scheduledWeightVector(view.flat) - if len(state.weightedState.current) == 0 || !weightVectorsEqual(state.weightedState.weights, weights) { + if len(state.weightedState.current) == 0 || weightsConfigChanged(state.weightedState.weights, weights) { return } - view.weightedState.current = state.weightedState.current + current := make(map[string]int64, len(view.flat)) + for _, entry := range view.flat { + if entry != nil && entry.auth != nil { + if val, ok := state.weightedState.current[entry.auth.ID]; ok { + current[entry.auth.ID] = val + } + } + } + view.weightedState.current = current view.weightedState.weights = weights } @@ -1066,22 +1074,6 @@ func scheduledWeightVectorMatching(entries []*scheduledAuth, predicate func(*sch } func pickSmoothWeightedScheduled(entries []*scheduledAuth, current map[string]int64, predicate func(*scheduledAuth) bool) *scheduledAuth { - active := make(map[string]struct{}, len(entries)) - for _, entry := range entries { - if entry == nil || entry.auth == nil || entry.meta == nil || entry.meta.weight <= 0 { - continue - } - if predicate != nil && !predicate(entry) { - continue - } - active[entry.auth.ID] = struct{}{} - } - for authID := range current { - if _, ok := active[authID]; !ok { - delete(current, authID) - } - } - var picked *scheduledAuth var pickedCurrent int64 var totalWeight int64 diff --git a/sdk/cliproxy/auth/scheduler_test.go b/sdk/cliproxy/auth/scheduler_test.go index 55d93dc1d..5438a376b 100644 --- a/sdk/cliproxy/auth/scheduler_test.go +++ b/sdk/cliproxy/auth/scheduler_test.go @@ -554,6 +554,37 @@ func TestSchedulerPick_MixedProvidersResetsCreditsWhenWeightsChange(t *testing.T } } +func TestSchedulerPickMixed_RetryTriedFilterPreservesSmoothWeightedDistribution(t *testing.T) { + t.Parallel() + + authA := &Auth{ID: "auth-a", Provider: "provider-a"} + authB := &Auth{ID: "auth-b", Provider: "provider-b"} + authC := &Auth{ID: "auth-c", Provider: "provider-c"} + authD := &Auth{ID: "auth-d", Provider: "provider-d"} + auths := []*Auth{authA, authB, authC, authD} + + scheduler := newSchedulerForTest(&WeightedRoundRobinSelector{}, auths...) + providers := []string{"provider-a", "provider-b", "provider-c", "provider-d"} + + // Simulate retries where auth-a failed and is in the tried filter: + // Verify that retry picks rotate smoothly across auth-b, auth-c, auth-d without alphabetical bias towards auth-b. + retryCounts := make(map[string]int) + tried := map[string]struct{}{"auth-a": {}} + for index := 0; index < 30; index++ { + picked, _, errPick := scheduler.pickMixed(context.Background(), providers, "", cliproxyexecutor.Options{}, tried) + if errPick != nil { + t.Fatalf("pickMixed(tried) error = %v", errPick) + } + retryCounts[picked.ID]++ + } + + for _, authID := range []string{"auth-b", "auth-c", "auth-d"} { + if retryCounts[authID] != 10 { + t.Fatalf("auth %q retry picks = %d, want 10 (even distribution without alphabetical bias, counts=%#v)", authID, retryCounts[authID], retryCounts) + } + } +} + func TestSchedulerPick_MixedProvidersPrefersHighestPriorityTier(t *testing.T) { t.Parallel() diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 7a053b78a..657be2379 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -24,10 +24,15 @@ import ( ) // RoundRobinSelector provides a simple provider scoped round-robin selection strategy. +// +// Rotation continues from the identity of the previous pick rather than from a numeric +// index. Candidate slices shrink whenever a retry excludes already tried credentials or a +// credential enters cooldown, and indexing a monotonic counter into a shrinking slice +// silently re-seats the rotation, which starves some credentials and hammers others. type RoundRobinSelector struct { - mu sync.Mutex - cursors map[string]int - maxKeys int + mu sync.Mutex + lastPicked map[string]string + maxKeys int } // WeightedRoundRobinSelector provides smooth weighted round-robin selection. @@ -377,29 +382,41 @@ func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, o available = preferCodexWebsocketAuths(ctx, provider, available) key := provider + ":" + canonicalModelKey(model) s.mu.Lock() - if s.cursors == nil { - s.cursors = make(map[string]int) + defer s.mu.Unlock() + if s.lastPicked == nil { + s.lastPicked = make(map[string]string) } limit := s.maxKeys if limit <= 0 { limit = 4096 } - s.ensureCursorKey(key, limit) - index := s.cursors[key] - if index >= 2_147_483_640 { - index = 0 - } - s.cursors[key] = index + 1 - s.mu.Unlock() - return available[index%len(available)], nil + s.ensureRotationKey(key, limit) + picked := available[successorIndex(available, s.lastPicked[key])] + s.lastPicked[key] = picked.ID + return picked, nil } -// ensureCursorKey ensures the cursor map has capacity for the given key. +// successorIndex returns the index of the first candidate ordered after lastID, wrapping to +// the start of the ring. Candidates arrive sorted by ID, so this resumes the rotation at the +// credential that follows the previous pick even when candidates were filtered out in +// between. An empty lastID starts at the head. +func successorIndex(available []*Auth, lastID string) int { + if lastID == "" { + return 0 + } + index := sort.Search(len(available), func(i int) bool { return available[i].ID > lastID }) + if index >= len(available) { + return 0 + } + return index +} + +// ensureRotationKey ensures the rotation map has capacity for the given key. // Must be called with s.mu held. -func (s *RoundRobinSelector) ensureCursorKey(key string, limit int) { - if _, ok := s.cursors[key]; !ok && len(s.cursors) >= limit { - s.cursors = make(map[string]int) +func (s *RoundRobinSelector) ensureRotationKey(key string, limit int) { + if _, ok := s.lastPicked[key]; !ok && len(s.lastPicked) >= limit { + s.lastPicked = make(map[string]string) } } @@ -450,23 +467,60 @@ func (s *WeightedRoundRobinSelector) Pick(ctx context.Context, provider, model s return picked, nil } +// maxSmoothWeightedStateEntries bounds a single accumulator map so credentials that are +// removed permanently cannot leak entries. Real pools stay far below this bound, so the +// transient subsets produced by retry exclusions and cooldowns are never pruned. +const maxSmoothWeightedStateEntries = 1024 + +// prepare syncs the configured weights into the state without discarding accumulated +// credits. Credits are reset only when a credential's configured weight actually changes, +// never when the candidate set shrinks temporarily (retry exclusions, cooldowns, session +// affinity), because discarding credits there would collapse selection onto the first +// candidate in slice order. func (s *smoothWeightedState) prepare(weights map[string]int64) { - if s.current == nil || !weightVectorsEqual(s.weights, weights) { - s.current = make(map[string]int64) + if s.current == nil || weightsConfigChanged(s.weights, weights) { + s.current = make(map[string]int64, len(weights)) } - s.weights = weights + if s.weights == nil { + s.weights = make(map[string]int64, len(weights)) + } + for authID, weight := range weights { + s.weights[authID] = weight + } + s.pruneStale(weights) } -func weightVectorsEqual(left, right map[string]int64) bool { - if len(left) != len(right) { - return false +// pruneStale drops entries for credentials outside the current candidate set, but only +// once a map exceeds the safety bound, so ordinary transient exclusions keep their credits. +func (s *smoothWeightedState) pruneStale(weights map[string]int64) { + if len(s.current) <= maxSmoothWeightedStateEntries && len(s.weights) <= maxSmoothWeightedStateEntries { + return } - for authID, weight := range left { - if right[authID] != weight { - return false + for authID := range s.current { + if _, ok := weights[authID]; !ok { + delete(s.current, authID) } } - return true + for authID := range s.weights { + if _, ok := weights[authID]; !ok { + delete(s.weights, authID) + } + } +} + +// weightsConfigChanged reports whether any credential present in both vectors has a +// different configured weight. Credentials that are merely missing from one side are +// ignored, since a candidate subset is not a configuration change. +func weightsConfigChanged(left, right map[string]int64) bool { + if len(left) == 0 { + return false + } + for authID, weight := range right { + if previous, ok := left[authID]; ok && previous != weight { + return true + } + } + return false } func authWeightVector(auths []*Auth) map[string]int64 { @@ -483,7 +537,6 @@ func authWeightVector(auths []*Auth) map[string]int64 { } func pickSmoothWeightedAuth(auths []*Auth, current map[string]int64) *Auth { - active := make(map[string]struct{}, len(auths)) var picked *Auth var pickedCurrent int64 var totalWeight int64 @@ -492,7 +545,6 @@ func pickSmoothWeightedAuth(auths []*Auth, current map[string]int64) *Auth { if auth == nil || weight <= 0 { continue } - active[auth.ID] = struct{}{} current[auth.ID] = saturatingAddInt64(current[auth.ID], weight) totalWeight = saturatingAddInt64(totalWeight, weight) if picked == nil || current[auth.ID] > pickedCurrent { @@ -500,11 +552,6 @@ func pickSmoothWeightedAuth(auths []*Auth, current map[string]int64) *Auth { pickedCurrent = current[auth.ID] } } - for authID := range current { - if _, ok := active[authID]; !ok { - delete(current, authID) - } - } if picked == nil { return nil } diff --git a/sdk/cliproxy/auth/selector_test.go b/sdk/cliproxy/auth/selector_test.go index 8df0520a7..546e628ee 100644 --- a/sdk/cliproxy/auth/selector_test.go +++ b/sdk/cliproxy/auth/selector_test.go @@ -275,6 +275,62 @@ func TestWeightedRoundRobinSelectorPick_DefaultWeightIsOne(t *testing.T) { } } +func TestWeightedRoundRobinSelectorPick_SubsetFilteringDoesNotResetAccumulatorOrFavorFirstAlphabetical(t *testing.T) { + t.Parallel() + + selector := &WeightedRoundRobinSelector{} + authB := &Auth{ID: "auth-b"} + authC := &Auth{ID: "auth-c"} + authD := &Auth{ID: "auth-d"} + subsetPool := []*Auth{authB, authC, authD} // auth-a excluded (e.g. tried or cooling) + + // Simulate repeated failover calls where auth-a is excluded: + // Verify that auth-b, auth-c, auth-d are picked evenly (10 each) rather than auth-b taking 100% of picks. + retryCounts := make(map[string]int) + for index := 0; index < 30; index++ { + got, errPick := selector.Pick(context.Background(), "provider", "model", cliproxyexecutor.Options{}, subsetPool) + if errPick != nil { + t.Fatalf("Pick(subset) error = %v", errPick) + } + retryCounts[got.ID]++ + } + + for _, authID := range []string{"auth-b", "auth-c", "auth-d"} { + if retryCounts[authID] != 10 { + t.Fatalf("auth %q retry picks = %d, want 10 (even distribution without alphabetical bias, counts=%#v)", authID, retryCounts[authID], retryCounts) + } + } +} + +func TestSmoothWeightedStatePrepare_KeepsCreditsForTransientSubsetsAndBoundsGrowth(t *testing.T) { + t.Parallel() + + state := &smoothWeightedState{} + state.prepare(map[string]int64{"a": 1, "b": 1}) + state.current["a"] = -2 + state.current["b"] = 1 + + // A shrinking candidate set must not discard credits. + state.prepare(map[string]int64{"b": 1}) + if state.current["a"] != -2 || state.current["b"] != 1 { + t.Fatalf("credits after subset prepare = %#v, want a:-2 b:1", state.current) + } + + // A real weight change resets credits. + state.prepare(map[string]int64{"b": 5}) + if len(state.current) != 0 { + t.Fatalf("credits after weight change = %#v, want empty", state.current) + } + + // Long-lived churn must stay bounded instead of leaking one entry per removed credential. + for index := 0; index < maxSmoothWeightedStateEntries*3; index++ { + state.prepare(map[string]int64{fmt.Sprintf("churn-%d", index): 5}) + } + if len(state.current) > maxSmoothWeightedStateEntries || len(state.weights) > maxSmoothWeightedStateEntries { + t.Fatalf("state grew unbounded: current=%d weights=%d, want <= %d", len(state.current), len(state.weights), maxSmoothWeightedStateEntries) + } +} + func TestRoundRobinSelectorPick_PriorityBuckets(t *testing.T) { t.Parallel() @@ -663,6 +719,108 @@ func TestRoundRobinSelectorPick_ThinkingSuffixSharesCursor(t *testing.T) { } } +func TestRoundRobinSelectorPick_ResumesRotationAcrossRetryExclusions(t *testing.T) { + t.Parallel() + + ids := []string{"aaa", "bbb", "ccc", "ddd", "eee"} + auths := make([]*Auth, 0, len(ids)) + for _, id := range ids { + auths = append(auths, &Auth{ID: id}) + } + + // Every request burns three attempts, so each attempt must consume the next slot of one + // shared rotation. Re-seating the rotation on the shrunken candidate slice would starve + // the head of the tier and hammer its tail. + selector := &RoundRobinSelector{} + const requests = 50 + firstAttempt := make(map[string]int) + allAttempts := make(map[string]int) + for index := 0; index < requests; index++ { + tried := make(map[string]struct{}) + for attempt := 0; attempt < 3; attempt++ { + candidates := make([]*Auth, 0, len(auths)) + for _, auth := range auths { + if _, used := tried[auth.ID]; !used { + candidates = append(candidates, auth) + } + } + got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, candidates) + if errPick != nil { + t.Fatalf("Pick() request %d attempt %d error = %v", index, attempt, errPick) + } + if attempt == 0 { + firstAttempt[got.ID]++ + } + allAttempts[got.ID]++ + tried[got.ID] = struct{}{} + } + } + + for _, id := range ids { + if firstAttempt[id] != requests/len(ids) { + t.Fatalf("auth %q first attempts = %d, want %d (counts=%#v)", id, firstAttempt[id], requests/len(ids), firstAttempt) + } + if allAttempts[id] != requests*3/len(ids) { + t.Fatalf("auth %q total attempts = %d, want %d (counts=%#v)", id, allAttempts[id], requests*3/len(ids), allAttempts) + } + } +} + +func TestWeightedRoundRobinSelectorPick_KeepsWeightRatiosWhenCandidatesAreExcluded(t *testing.T) { + t.Parallel() + + // auth-a is excluded exactly as a retry or cooldown would exclude it. The survivors must + // keep their configured 3:1 ratio instead of collapsing onto the first candidate. + selector := &WeightedRoundRobinSelector{} + authA := &Auth{ID: "auth-a", Attributes: map[string]string{AttributeWeight: "5"}} + authB := &Auth{ID: "auth-b", Attributes: map[string]string{AttributeWeight: "3"}} + authC := &Auth{ID: "auth-c", Attributes: map[string]string{AttributeWeight: "1"}} + + fullPool := []*Auth{authA, authB, authC} + for index := 0; index < 9; index++ { + if _, errPick := selector.Pick(context.Background(), "codex", "model", cliproxyexecutor.Options{}, fullPool); errPick != nil { + t.Fatalf("Pick(full) #%d error = %v", index, errPick) + } + } + + survivors := []*Auth{authB, authC} + counts := make(map[string]int) + for index := 0; index < 400; index++ { + got, errPick := selector.Pick(context.Background(), "codex", "model", cliproxyexecutor.Options{}, survivors) + if errPick != nil { + t.Fatalf("Pick(survivors) #%d error = %v", index, errPick) + } + counts[got.ID]++ + } + if counts["auth-b"] != 300 || counts["auth-c"] != 100 { + t.Fatalf("survivor picks = %#v, want auth-b:300 auth-c:100 (3:1)", counts) + } +} + +func TestSuccessorIndex_WrapsAndSkipsFilteredCandidates(t *testing.T) { + t.Parallel() + + available := []*Auth{{ID: "aaa"}, {ID: "ccc"}, {ID: "eee"}} + tests := []struct { + name string + lastID string + want int + }{ + {name: "no previous pick starts at head", lastID: "", want: 0}, + {name: "resumes after previous pick", lastID: "aaa", want: 1}, + {name: "resumes after filtered-out pick", lastID: "bbb", want: 1}, + {name: "wraps at the end of the ring", lastID: "eee", want: 0}, + {name: "wraps for removed trailing pick", lastID: "zzz", want: 0}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := successorIndex(available, tt.lastID); got != tt.want { + t.Fatalf("successorIndex(%q) = %d, want %d", tt.lastID, got, tt.want) + } + }) + } +} + func TestRoundRobinSelectorPick_CursorKeyCap(t *testing.T) { t.Parallel() @@ -676,14 +834,14 @@ func TestRoundRobinSelectorPick_CursorKeyCap(t *testing.T) { selector.mu.Lock() defer selector.mu.Unlock() - if selector.cursors == nil { - t.Fatalf("selector.cursors = nil") + if selector.lastPicked == nil { + t.Fatalf("selector.lastPicked = nil") } - if len(selector.cursors) != 1 { - t.Fatalf("len(selector.cursors) = %d, want %d", len(selector.cursors), 1) + if len(selector.lastPicked) != 1 { + t.Fatalf("len(selector.lastPicked) = %d, want %d", len(selector.lastPicked), 1) } - if _, ok := selector.cursors["gemini:m3"]; !ok { - t.Fatalf("selector.cursors missing key %q", "gemini:m3") + if _, ok := selector.lastPicked["gemini:m3"]; !ok { + t.Fatalf("selector.lastPicked missing key %q", "gemini:m3") } }