History

fix(auth): enforce minimum cooldown floor and track attempted credentials on 429

- Enforce a 10-second minimum cooldown floor for quota errors to prevent retry storms from sub-second `Retry-After` values. - Track attempted authentication credentials per round across execution flows. - Ensure credentials attempted during a failed 429 round respect the cooldown floor instead of triggering zero-wait retries.

Luis Pater committed 2 months ago

Commit 18e01a71 parent f4161755 files · +538 −13Browse files

Unified · Split

5 files
  1. sdk/cliproxy/auth/conductor_cooldown.go +10 −2
  2. sdk/cliproxy/auth/conductor_execution.go +30 −6
  3. sdk/cliproxy/auth/conductor_refresh.go +1 −0
  4. sdk/cliproxy/auth/conductor_selection.go +38 −5
  5. sdk/cliproxy/auth/conductor_subsecond_cooldown_test.go +459 −0
sdk/cliproxy/auth/conductor_cooldown.go+10−2
@@ -845,7 +845,11 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { · 844 unmodified lines
845 backoffLevel := state.Quota.BackoffLevel845 backoffLevel := state.Quota.BackoffLevel
846 if !disableCooling {846 if !disableCooling {
847 if result.RetryAfter != nil {847 if result.RetryAfter != nil {
848 next = now.Add(*result.RetryAfter)848 cooldown := *result.RetryAfter
849 if cooldown < minQuotaCooldownFloor {
850 cooldown = minQuotaCooldownFloor
851 }
852 next = now.Add(cooldown)
849 } else {853 } else {
850 next, backoffLevel = quotaCooldownAfterFailure(state.Quota, now)854 next, backoffLevel = quotaCooldownAfterFailure(state.Quota, now)
851 }855 }
@@ -2050,7 +2054,11 @@ func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Durati · 1198 unmodified lines
2050 var next time.Time2054 var next time.Time
2051 if !disableCooling {2055 if !disableCooling {
2052 if retryAfter != nil {2056 if retryAfter != nil {
2053 next = now.Add(*retryAfter)2057 cooldown := *retryAfter
2058 if cooldown < minQuotaCooldownFloor {
2059 cooldown = minQuotaCooldownFloor
2060 }
2061 next = now.Add(cooldown)
2054 } else {2062 } else {
2055 next, auth.Quota.BackoffLevel = quotaCooldownAfterFailure(auth.Quota, now)2063 next, auth.Quota.BackoffLevel = quotaCooldownAfterFailure(auth.Quota, now)
2056 }2064 }
sdk/cliproxy/auth/conductor_execution.go+30−6
@@ -135,7 +135,9 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye · 134 unmodified lines
135 var preferredUpstreamErr error135 var preferredUpstreamErr error
136 retryModel := authSelectionModelFromOptions(opts, req.Model)136 retryModel := authSelectionModelFromOptions(opts, req.Model)
137 for attempt := 0; ; attempt++ {137 for attempt := 0; ; attempt++ {
138 resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, attempt, defaultRequestRetry)138 roundAttempted := make(map[string]struct{})
139 roundOpts := withAttemptedAuthTracker(opts, roundAttempted)
140 resp, errExec := m.executeMixedOnce(ctx, normalized, req, roundOpts, maxRetryCredentials, attempt, defaultRequestRetry)
139 if errExec == nil {141 if errExec == nil {
140 return resp, nil142 return resp, nil
141 }143 }
@@ -146,7 +148,7 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye · 4 unmodified lines
146 preferredUpstreamErr = errExec148 preferredUpstreamErr = errExec
147 }149 }
148 lastErr = errExec150 lastErr = errExec
149 wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry)151 wait, shouldRetry := m.shouldRetryAfterErrorWithAttempted(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry, roundAttempted)
150 if !shouldRetry {152 if !shouldRetry {
151 break153 break
152 }154 }
@@ -192,7 +194,9 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip · 39 unmodified lines
192 var preferredUpstreamErr error194 var preferredUpstreamErr error
193 retryModel := authSelectionModelFromOptions(opts, req.Model)195 retryModel := authSelectionModelFromOptions(opts, req.Model)
194 for attempt := 0; ; attempt++ {196 for attempt := 0; ; attempt++ {
195 resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, attempt, defaultRequestRetry)197 roundAttempted := make(map[string]struct{})
198 roundOpts := withAttemptedAuthTracker(opts, roundAttempted)
199 resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, roundOpts, maxRetryCredentials, attempt, defaultRequestRetry)
196 if errExec == nil {200 if errExec == nil {
197 return resp, nil201 return resp, nil
198 }202 }
@@ -203,7 +207,7 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip · 4 unmodified lines
203 preferredUpstreamErr = errExec207 preferredUpstreamErr = errExec
204 }208 }
205 lastErr = errExec209 lastErr = errExec
206 wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry)210 wait, shouldRetry := m.shouldRetryAfterErrorWithAttempted(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry, roundAttempted)
207 if !shouldRetry {211 if !shouldRetry {
208 break212 break
209 }213 }
@@ -247,7 +251,9 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli · 37 unmodified lines
247 retryRoundPending := false251 retryRoundPending := false
248 retryRoundWaited := false252 retryRoundWaited := false
249 for {253 for {
250 result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, &homeRetryLimit, attempt, defaultRequestRetry)254 roundAttempted := make(map[string]struct{})
255 roundOpts := withAttemptedAuthTracker(opts, roundAttempted)
256 result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, roundOpts, maxRetryCredentials, &homeRetryLimit, attempt, defaultRequestRetry)
251 if errStream == nil {257 if errStream == nil {
252 return result, nil258 return result, nil
253 }259 }
@@ -272,7 +278,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli · 18 unmodified lines
272 return nil, unwrapExecutionBoundaryError(errStream)278 return nil, unwrapExecutionBoundaryError(errStream)
273 }279 }
274 lastErr = errStream280 lastErr = errStream
275 wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errStream, attempt, normalized, retryModel, maxWait, homeRetryLimit, defaultRequestRetry)281 wait, shouldRetry := m.shouldRetryAfterErrorWithAttempted(ctx, opts, errStream, attempt, normalized, retryModel, maxWait, homeRetryLimit, defaultRequestRetry, roundAttempted)
276 if !shouldRetry {282 if !shouldRetry {
277 break283 break
278 }284 }
@@ -1098,6 +1104,24 @@ func shouldExcludeHomeAuthAfterStreamError(ctx context.Context, _ *Auth, err err · 819 unmodified lines
1098 return true1104 return true
1099}1105}
1100 1106
1107func withAttemptedAuthTracker(opts cliproxyexecutor.Options, attempted map[string]struct{}) cliproxyexecutor.Options {
1108 if attempted == nil {
1109 return opts
1110 }
1111 meta := cloneRequestMetadata(opts.Metadata)
1112 prevCallback, _ := meta[cliproxyexecutor.SelectedAuthCallbackMetadataKey].(func(string))
1113 meta[cliproxyexecutor.SelectedAuthCallbackMetadataKey] = func(authID string) {
1114 if strings.TrimSpace(authID) != "" {
1115 attempted[authID] = struct{}{}
1116 }
1117 if prevCallback != nil {
1118 prevCallback(authID)
1119 }
1120 }
1121 opts.Metadata = meta
1122 return opts
1123}
1124
1101func cloneRequestMetadata(src map[string]any) map[string]any {1125func cloneRequestMetadata(src map[string]any) map[string]any {
1102 if len(src) == 0 {1126 if len(src) == 0 {
1103 return make(map[string]any, 4)1127 return make(map[string]any, 4)
sdk/cliproxy/auth/conductor_refresh.go+1−0
@@ -32,6 +32,7 @@ const ( · 31 unmodified lines
32 refreshIneffectiveBackoff = 30 * time.Second32 refreshIneffectiveBackoff = 30 * time.Second
33 quotaBackoffBase = time.Second33 quotaBackoffBase = time.Second
34 quotaBackoffMax = 30 * time.Minute34 quotaBackoffMax = 30 * time.Minute
35 minQuotaCooldownFloor = 10 * time.Second
35 transientErrorCooldown = time.Minute36 transientErrorCooldown = time.Minute
36)37)
37 38
sdk/cliproxy/auth/conductor_selection.go+38−5
@@ -999,6 +999,10 @@ func credentialRetryRoundStateEligible(lastErr *Error, quotaExceeded bool) bool · 998 unmodified lines
999}999}
1000 1000
1001func (m *Manager) closestCooldownWait(providers []string, model string, attempt int, eligibility authSelectionEligibility, pinnedAuthID string, defaultRequestRetry int) (time.Duration, bool) {1001func (m *Manager) closestCooldownWait(providers []string, model string, attempt int, eligibility authSelectionEligibility, pinnedAuthID string, defaultRequestRetry int) (time.Duration, bool) {
1002 return m.closestCooldownWaitWithAttempted(providers, model, attempt, eligibility, pinnedAuthID, defaultRequestRetry, 0, nil)
1003}
1004
1005func (m *Manager) closestCooldownWaitWithAttempted(providers []string, model string, attempt int, eligibility authSelectionEligibility, pinnedAuthID string, defaultRequestRetry int, status int, attempted map[string]struct{}) (time.Duration, bool) {
1002 if m == nil || len(providers) == 0 {1006 if m == nil || len(providers) == 0 {
1003 return 0, false1007 return 0, false
1004 }1008 }
@@ -1050,13 +1054,38 @@ func (m *Manager) closestCooldownWait(providers []string, model string, attempt · 45 unmodified lines
1050 if !retryEligible {1054 if !retryEligible {
1051 continue1055 continue
1052 }1056 }
1053 if next.IsZero() {1057
1054 return 0, true1058 wasAttempted := false
1059 if len(attempted) > 0 {
1060 _, wasAttempted = attempted[auth.ID]
1055 }1061 }
1056 wait := next.Sub(now)1062 coolingDisabled := m.cooldownDisabledForAuth(auth)
1057 if wait < 0 {1063 if !wasAttempted || coolingDisabled || status != http.StatusTooManyRequests {
1064 if next.IsZero() {
1065 return 0, true
1066 }
1067 wait := next.Sub(now)
1068 if wait < 0 {
1069 continue
1070 }
1071 if !found || wait < minWait {
1072 minWait = wait
1073 found = true
1074 }
1058 continue1075 continue
1059 }1076 }
1077
1078 // auth was already attempted in the round that just failed with 429,
1079 // and cooling is enabled. It must not trigger an immediate zero-wait retry round.
1080 var wait time.Duration
1081 if next.IsZero() {
1082 wait = minQuotaCooldownFloor
1083 } else {
1084 wait = next.Sub(now)
1085 if wait < minQuotaCooldownFloor {
1086 wait = minQuotaCooldownFloor
1087 }
1088 }
1060 if !found || wait < minWait {1089 if !found || wait < minWait {
1061 minWait = wait1090 minWait = wait
1062 found = true1091 found = true
@@ -1131,6 +1160,10 @@ func (m *Manager) shouldRetryAfterError(err error, attempt int, providers []stri · 68 unmodified lines
1131// immediately. If every eligible credential still needs a positive cooldown,1160// immediately. If every eligible credential still needs a positive cooldown,
1132// retry stops without waiting.1161// retry stops without waiting.
1133func (m *Manager) shouldRetryAfterErrorWithHomeRetryLimit(ctx context.Context, opts cliproxyexecutor.Options, err error, attempt int, providers []string, model string, maxWait time.Duration, homeRetryLimit int, defaultRequestRetry int) (time.Duration, bool) {1162func (m *Manager) shouldRetryAfterErrorWithHomeRetryLimit(ctx context.Context, opts cliproxyexecutor.Options, err error, attempt int, providers []string, model string, maxWait time.Duration, homeRetryLimit int, defaultRequestRetry int) (time.Duration, bool) {
1163 return m.shouldRetryAfterErrorWithAttempted(ctx, opts, err, attempt, providers, model, maxWait, homeRetryLimit, defaultRequestRetry, nil)
1164}
1165
1166func (m *Manager) shouldRetryAfterErrorWithAttempted(ctx context.Context, opts cliproxyexecutor.Options, err error, attempt int, providers []string, model string, maxWait time.Duration, homeRetryLimit int, defaultRequestRetry int, attempted map[string]struct{}) (time.Duration, bool) {
1134 if err == nil {1167 if err == nil {
1135 return 0, false1168 return 0, false
1136 }1169 }
@@ -1184,7 +1217,7 @@ func (m *Manager) shouldRetryAfterErrorWithHomeRetryLimit(ctx context.Context, o · 47 unmodified lines
1184 if !isCredentialRetryRoundStatus(status) || !m.retryAllowed(attempt, providers, model, eligibility, pinnedAuthID, defaultRequestRetry) {1217 if !isCredentialRetryRoundStatus(status) || !m.retryAllowed(attempt, providers, model, eligibility, pinnedAuthID, defaultRequestRetry) {
1185 return 0, false1218 return 0, false
1186 }1219 }
1187 wait, found := m.closestCooldownWait(providers, model, attempt, eligibility, pinnedAuthID, defaultRequestRetry)1220 wait, found := m.closestCooldownWaitWithAttempted(providers, model, attempt, eligibility, pinnedAuthID, defaultRequestRetry, status, attempted)
1188 if found {1221 if found {
1189 if wait > 0 && (maxWait <= 0 || wait > maxWait) {1222 if wait > 0 && (maxWait <= 0 || wait > maxWait) {
1190 return 0, false1223 return 0, false
sdk/cliproxy/auth/conductor_subsecond_cooldown_test.go+459−0
@@ -0,0 +1,459 @@
1package auth
2
3import (
4 "context"
5 "net/http"
6 "sync"
7 "testing"
8 "time"
9
10 internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
11 "github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
12 cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
13)
14
15func TestMarkResult_429SubSecondRetryAfter_EnforcesMinimumCooldownFloor(t *testing.T) {
16 withQuotaCooldownEnabled(t)
17
18 manager := NewManager(nil, nil, nil)
19 auth := &Auth{
20 ID: "auth-gemini-subsecond",
21 Provider: "google",
22 }
23
24 model := "gemini-3.1-flash-image"
25 reg := registry.GetGlobalRegistry()
26 reg.RegisterClient(auth.ID, "google", []*registry.ModelInfo{{ID: model}})
27 t.Cleanup(func() { reg.UnregisterClient(auth.ID) })
28
29 if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil {
30 t.Fatalf("Register returned error: %v", errRegister)
31 }
32
33 now := time.Now()
34 subSecond := 708 * time.Millisecond
35 manager.MarkResult(context.Background(), Result{
36 AuthID: auth.ID,
37 Provider: "google",
38 Model: model,
39 Success: false,
40 RetryAfter: &subSecond,
41 Error: &Error{
42 HTTPStatus: http.StatusTooManyRequests,
43 Message: "RESOURCE_EXHAUSTED",
44 },
45 })
46
47 updated, ok := manager.GetByID(auth.ID)
48 if !ok || updated == nil {
49 t.Fatal("expected auth to be present")
50 }
51 state := updated.ModelStates[model]
52 if state == nil {
53 t.Fatal("expected model state to be present")
54 }
55 if state.NextRetryAfter.IsZero() {
56 t.Fatal("expected NextRetryAfter to be set")
57 }
58
59 // Sub-second RetryAfter (e.g. 708ms) must be clamped to at least minQuotaCooldownFloor (10s)
60 // to prevent instant unblocking and multi-round retry storms across credentials.
61 minExpected := now.Add(10 * time.Second)
62 if state.NextRetryAfter.Before(minExpected) {
63 t.Fatalf("expected NextRetryAfter >= %v (at least 10s cooldown floor), got %v (diff %v)",
64 minExpected, state.NextRetryAfter, state.NextRetryAfter.Sub(now))
65 }
66 if state.Quota.NextRecoverAt.Before(minExpected) {
67 t.Fatalf("expected Quota.NextRecoverAt >= %v (at least 10s cooldown floor), got %v (diff %v)",
68 minExpected, state.Quota.NextRecoverAt, state.Quota.NextRecoverAt.Sub(now))
69 }
70}
71
72func TestMarkResult_429LongerRetryAfter_PreservedAboveFloor(t *testing.T) {
73 withQuotaCooldownEnabled(t)
74
75 manager := NewManager(nil, nil, nil)
76 auth := &Auth{
77 ID: "auth-gemini-longer",
78 Provider: "google",
79 }
80
81 model := "gemini-3.1-flash-image"
82 reg := registry.GetGlobalRegistry()
83 reg.RegisterClient(auth.ID, "google", []*registry.ModelInfo{{ID: model}})
84 t.Cleanup(func() { reg.UnregisterClient(auth.ID) })
85
86 if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil {
87 t.Fatalf("Register returned error: %v", errRegister)
88 }
89
90 now := time.Now()
91 longDuration := 5 * time.Minute
92 manager.MarkResult(context.Background(), Result{
93 AuthID: auth.ID,
94 Provider: "google",
95 Model: model,
96 Success: false,
97 RetryAfter: &longDuration,
98 Error: &Error{
99 HTTPStatus: http.StatusTooManyRequests,
100 Message: "RESOURCE_EXHAUSTED",
101 },
102 })
103
104 updated, ok := manager.GetByID(auth.ID)
105 if !ok || updated == nil {
106 t.Fatal("expected auth to be present")
107 }
108 state := updated.ModelStates[model]
109 if state == nil {
110 t.Fatal("expected model state to be present")
111 }
112
113 expected := now.Add(longDuration)
114 if state.NextRetryAfter.Before(expected.Add(-time.Second)) || state.NextRetryAfter.After(expected.Add(time.Second)) {
115 t.Fatalf("expected NextRetryAfter ~ %v, got %v", expected, state.NextRetryAfter)
116 }
117}
118
119func TestApplyAuthFailureState_429SubSecondRetryAfter_EnforcesMinimumCooldownFloor(t *testing.T) {
120 now := time.Now()
121 subSecond := 708 * time.Millisecond
122 auth := &Auth{ID: "auth-level-subsecond"}
123 quotaErr := &Error{
124 HTTPStatus: http.StatusTooManyRequests,
125 Message: "RESOURCE_EXHAUSTED",
126 }
127
128 applyAuthFailureState(auth, quotaErr, &subSecond, now, false)
129
130 if auth.NextRetryAfter.IsZero() {
131 t.Fatal("expected NextRetryAfter to be set")
132 }
133 minExpected := now.Add(10 * time.Second)
134 if auth.NextRetryAfter.Before(minExpected) {
135 t.Fatalf("expected auth NextRetryAfter >= %v (at least 10s cooldown floor), got %v (diff %v)",
136 minExpected, auth.NextRetryAfter, auth.NextRetryAfter.Sub(now))
137 }
138 if auth.Quota.NextRecoverAt.Before(minExpected) {
139 t.Fatalf("expected auth Quota.NextRecoverAt >= %v (at least 10s cooldown floor), got %v (diff %v)",
140 minExpected, auth.Quota.NextRecoverAt, auth.Quota.NextRecoverAt.Sub(now))
141 }
142}
143
144func TestManager_Execute_429SubSecondRetryAfter_PreventsRetryStorm(t *testing.T) {
145 withQuotaCooldownEnabled(t)
146
147 manager := NewManager(nil, nil, nil)
148 manager.SetRetryConfig(5, 5*time.Second, 6)
149
150 subSecond := 708 * time.Millisecond
151 executor := &authFallbackExecutor{
152 id: "google",
153 executeErrors: map[string]error{
154 "auth-storm-1": &retryAfterStatusError{
155 status: http.StatusTooManyRequests,
156 message: "quota exhausted",
157 retryAfter: subSecond,
158 },
159 "auth-storm-2": &retryAfterStatusError{
160 status: http.StatusTooManyRequests,
161 message: "quota exhausted",
162 retryAfter: subSecond,
163 },
164 },
165 }
166 manager.RegisterExecutor(executor)
167
168 reg := registry.GetGlobalRegistry()
169 model := "gemini-3.1-flash-image"
170
171 for _, id := range []string{"auth-storm-1", "auth-storm-2"} {
172 auth := &Auth{
173 ID: id,
174 Provider: "google",
175 }
176 if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil {
177 t.Fatalf("Register returned error: %v", errRegister)
178 }
179 reg.RegisterClient(id, "google", []*registry.ModelInfo{{ID: model}})
180 defer reg.UnregisterClient(id)
181 }
182
183 req := cliproxyexecutor.Request{Model: model}
184 _, errExecute := manager.Execute(context.Background(), []string{"google"}, req, cliproxyexecutor.Options{})
185 if errExecute == nil {
186 t.Fatal("expected execute error")
187 }
188 if statusCodeFromError(errExecute) != http.StatusTooManyRequests {
189 t.Fatalf("execute status = %d, want %d", statusCodeFromError(errExecute), http.StatusTooManyRequests)
190 }
191
192 // Each account should be attempted once within round 0.
193 // It should NOT enter a tight loop retrying rounds because cooldown (>= 10s) exceeds maxWait (5s).
194 calls := executor.ExecuteCalls()
195 if len(calls) != 2 {
196 t.Fatalf("execute calls = %d, want exactly 2 (initial round only, no retry storm)", len(calls))
197 }
198}
199
200func TestClosestCooldownWaitWithAttempted_ExpiredCooldownDoesNotTriggerZeroWaitForAttemptedAuth(t *testing.T) {
201 withQuotaCooldownEnabled(t)
202
203 manager := NewManager(nil, nil, nil)
204 now := time.Now()
205 expired := now.Add(-5 * time.Second)
206
207 model := "gemini-3.1-flash-image"
208 auth := &Auth{
209 ID: "auth-attempted-expired",
210 Provider: "google",
211 ModelStates: map[string]*ModelState{
212 model: {
213 Status: StatusError,
214 Unavailable: true,
215 NextRetryAfter: expired,
216 Quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: expired},
217 },
218 },
219 }
220
221 reg := registry.GetGlobalRegistry()
222 reg.RegisterClient(auth.ID, "google", []*registry.ModelInfo{{ID: model}})
223 t.Cleanup(func() { reg.UnregisterClient(auth.ID) })
224
225 if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil {
226 t.Fatalf("Register returned error: %v", errRegister)
227 }
228
229 eligibility := authSelectionEligibilityForRequest(context.Background(), cliproxyexecutor.Options{})
230
231 // Case 1: auth was NOT attempted in the failed round (untried credential).
232 // Because its cooldown is expired, it should be immediately available (wait = 0, found = true).
233 waitUntried, foundUntried := manager.closestCooldownWaitWithAttempted([]string{"google"}, model, 0, eligibility, "", 5, http.StatusTooManyRequests, nil)
234 if !foundUntried || waitUntried != 0 {
235 t.Fatalf("untried credential: wait = %v, found = %t; want (0, true)", waitUntried, foundUntried)
236 }
237
238 // Case 2: auth WAS attempted in the failed round that returned 429.
239 // Even though its individual cooldown has expired, it must NOT trigger an immediate zero-wait retry round.
240 // It must enforce a cooldown floor of at least minQuotaCooldownFloor (10s).
241 attempted := map[string]struct{}{auth.ID: {}}
242 waitAttempted, foundAttempted := manager.closestCooldownWaitWithAttempted([]string{"google"}, model, 0, eligibility, "", 5, http.StatusTooManyRequests, attempted)
243 if !foundAttempted {
244 t.Fatal("expected candidate to be found for retry after cooldown")
245 }
246 if waitAttempted < minQuotaCooldownFloor {
247 t.Fatalf("attempted credential after 429: wait = %v, want >= %v (must not be zero-wait)", waitAttempted, minQuotaCooldownFloor)
248 }
249}
250
251type slowPoolSimulatedExecutor struct {
252 m *Manager
253 model string
254 authIDs []string
255 calls []string
256 executeErr error
257 mu sync.Mutex
258}
259
260func (e *slowPoolSimulatedExecutor) Identifier() string { return "google" }
261
262func (e *slowPoolSimulatedExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
263 e.mu.Lock()
264 e.calls = append(e.calls, auth.ID)
265 e.mu.Unlock()
266
267 // When auth-slow-2 executes, auth-slow-1 has already failed and been marked.
268 // We mutate auth-slow-1 directly in manager.auths to simulate that the pool iteration
269 // was slow enough that auth-slow-1's cooldown expired in the past before round 0 finished.
270 if auth.ID == e.authIDs[1] && e.m != nil {
271 e.m.mu.Lock()
272 if a1 := e.m.auths[e.authIDs[0]]; a1 != nil && a1.ModelStates[e.model] != nil {
273 expired := time.Now().Add(-5 * time.Second)
274 a1.ModelStates[e.model].NextRetryAfter = expired
275 a1.ModelStates[e.model].Quota.NextRecoverAt = expired
276 }
277 e.m.mu.Unlock()
278 }
279 return cliproxyexecutor.Response{}, e.executeErr
280}
281
282func (*slowPoolSimulatedExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
283 return nil, nil
284}
285
286func (*slowPoolSimulatedExecutor) Refresh(context.Context, *Auth) (*Auth, error) {
287 return nil, nil
288}
289
290func (*slowPoolSimulatedExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
291 return cliproxyexecutor.Response{}, nil
292}
293
294func (*slowPoolSimulatedExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
295 return nil, nil
296}
297
298func TestManager_Execute_SlowPoolExpiredCooldown_PreventsRetryStorm(t *testing.T) {
299 withQuotaCooldownEnabled(t)
300
301 manager := NewManager(nil, nil, nil)
302 manager.SetRetryConfig(5, 5*time.Second, 6)
303
304 model := "gemini-3.1-flash-image"
305 authIDs := []string{"auth-slow-1", "auth-slow-2"}
306
307 subSecond := 708 * time.Millisecond
308 executor := &slowPoolSimulatedExecutor{
309 m: manager,
310 model: model,
311 authIDs: authIDs,
312 executeErr: &retryAfterStatusError{
313 status: http.StatusTooManyRequests,
314 message: "quota exhausted",
315 retryAfter: subSecond,
316 },
317 }
318 manager.RegisterExecutor(executor)
319
320 reg := registry.GetGlobalRegistry()
321 for _, id := range authIDs {
322 auth := &Auth{ID: id, Provider: "google"}
323 if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil {
324 t.Fatalf("Register error: %v", errRegister)
325 }
326 reg.RegisterClient(id, "google", []*registry.ModelInfo{{ID: model}})
327 defer reg.UnregisterClient(id)
328 }
329
330 req := cliproxyexecutor.Request{Model: model}
331 _, errExecute := manager.Execute(context.Background(), []string{"google"}, req, cliproxyexecutor.Options{})
332 if errExecute == nil {
333 t.Fatal("expected execute error")
334 }
335 if statusCodeFromError(errExecute) != http.StatusTooManyRequests {
336 t.Fatalf("execute status = %d, want %d", statusCodeFromError(errExecute), http.StatusTooManyRequests)
337 }
338
339 // Confirm that auth-slow-1's cooldown genuinely expired in the past before round 0 completed.
340 updatedA1, ok := manager.GetByID(authIDs[0])
341 if !ok || updatedA1 == nil || updatedA1.ModelStates[model] == nil {
342 t.Fatal("expected auth-slow-1 state to exist")
343 }
344 if !updatedA1.ModelStates[model].NextRetryAfter.Before(time.Now()) {
345 t.Fatalf("expected auth-slow-1 NextRetryAfter to be expired in past, got %v", updatedA1.ModelStates[model].NextRetryAfter)
346 }
347
348 // Even though auth-slow-1's cooldown expired before round 0 finished,
349 // because both auths were attempted in round 0 and failed with 429,
350 // they require a minimum cooldown floor (>= 10s) which exceeds maxWait (5s).
351 // Exactly 2 calls (round 0 only) should occur. No retry storm!
352 executor.mu.Lock()
353 callCount := len(executor.calls)
354 executor.mu.Unlock()
355 if callCount != 2 {
356 t.Fatalf("executor calls = %d, want exactly 2 (round 0 only, calls=%v)", callCount, executor.calls)
357 }
358}
359
360func TestManager_ShouldRetryAfterError_429EnforcesCooldownWaitEvenWithLargeMaxWait(t *testing.T) {
361 withQuotaCooldownEnabled(t)
362
363 manager := NewManager(nil, nil, nil)
364 now := time.Now()
365 expired := now.Add(-5 * time.Second)
366
367 model := "gemini-3.1-flash-image"
368 auth := &Auth{
369 ID: "auth-large-maxwait",
370 Provider: "google",
371 ModelStates: map[string]*ModelState{
372 model: {
373 Status: StatusError,
374 Unavailable: true,
375 NextRetryAfter: expired,
376 Quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: expired},
377 },
378 },
379 }
380
381 reg := registry.GetGlobalRegistry()
382 reg.RegisterClient(auth.ID, "google", []*registry.ModelInfo{{ID: model}})
383 t.Cleanup(func() { reg.UnregisterClient(auth.ID) })
384
385 if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil {
386 t.Fatalf("Register returned error: %v", errRegister)
387 }
388
389 opts := cliproxyexecutor.Options{}
390 err429 := &Error{HTTPStatus: http.StatusTooManyRequests, Message: "RESOURCE_EXHAUSTED"}
391 attempted := map[string]struct{}{auth.ID: {}}
392
393 // With maxWait = 30s (which allows waiting up to 30s):
394 // Because auth was attempted and failed with 429, it must return wait >= 10s (NOT wait = 0).
395 wait, shouldRetry := manager.shouldRetryAfterErrorWithAttempted(context.Background(), opts, err429, 0, []string{"google"}, model, 30*time.Second, -1, 5, attempted)
396 if !shouldRetry {
397 t.Fatal("expected shouldRetry = true when maxWait (30s) >= cooldown floor (10s)")
398 }
399 if wait < minQuotaCooldownFloor {
400 t.Fatalf("wait = %v, want >= %v (must enforce positive cooldown wait, not zero-wait loop)", wait, minQuotaCooldownFloor)
401 }
402}
403
404func TestClosestCooldownWaitWithAttempted_RespectsManagerAndProviderCoolingOverrides(t *testing.T) {
405 withQuotaCooldownEnabled(t)
406
407 manager := NewManager(nil, nil, nil)
408 now := time.Now()
409 expired := now.Add(-5 * time.Second)
410
411 model := "compat-model-override"
412 auth := &Auth{
413 ID: "auth-compat-override",
414 Provider: "openai-compatibility",
415 Attributes: map[string]string{
416 "provider_key": "custom-llm",
417 },
418 ModelStates: map[string]*ModelState{
419 model: {
420 Status: StatusError,
421 Unavailable: true,
422 NextRetryAfter: expired,
423 Quota: QuotaState{Exceeded: true, Reason: "quota", NextRecoverAt: expired},
424 },
425 },
426 }
427
428 reg := registry.GetGlobalRegistry()
429 reg.RegisterClient(auth.ID, "openai-compatibility", []*registry.ModelInfo{{ID: model}})
430 t.Cleanup(func() { reg.UnregisterClient(auth.ID) })
431
432 if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil {
433 t.Fatalf("Register returned error: %v", errRegister)
434 }
435
436 attempted := map[string]struct{}{auth.ID: {}}
437 eligibility := authSelectionEligibilityForRequest(context.Background(), cliproxyexecutor.Options{})
438
439 // 1. Without provider override (cooling enabled): attempted auth after 429 enforces wait >= 10s.
440 wait, found := manager.closestCooldownWaitWithAttempted([]string{"openai-compatibility"}, model, 0, eligibility, "", 5, http.StatusTooManyRequests, attempted)
441 if !found || wait < minQuotaCooldownFloor {
442 t.Fatalf("cooling enabled: wait = %v, want >= %v", wait, minQuotaCooldownFloor)
443 }
444
445 // 2. With provider override setting disable_cooling = true:
446 disableCooling := true
447 manager.SetConfig(&internalconfig.Config{
448 OpenAICompatibility: []internalconfig.OpenAICompatibility{{
449 Name: "custom-llm",
450 DisableCooling: &disableCooling,
451 }},
452 })
453
454 // Now cooling is disabled via provider override: attempted auth should allow immediate retry (wait = 0, found = true).
455 waitDisabled, foundDisabled := manager.closestCooldownWaitWithAttempted([]string{"openai-compatibility"}, model, 0, eligibility, "", 5, http.StatusTooManyRequests, attempted)
456 if !foundDisabled || waitDisabled != 0 {
457 t.Fatalf("provider cooling disabled: wait = %v, found = %t; want (0, true)", waitDisabled, foundDisabled)
458 }
459}

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t