| @@ -0,0 +1,459 @@ |
| 1 | +package auth |
| 2 | + |
| 3 | +import ( |
| 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 | + |
| 15 | +func 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 | + |
| 72 | +func 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 | + |
| 119 | +func 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 | + |
| 144 | +func 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 | + |
| 200 | +func 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 | + |
| 251 | +type slowPoolSimulatedExecutor struct { |
| 252 | + m *Manager |
| 253 | + model string |
| 254 | + authIDs []string |
| 255 | + calls []string |
| 256 | + executeErr error |
| 257 | + mu sync.Mutex |
| 258 | +} |
| 259 | + |
| 260 | +func (e *slowPoolSimulatedExecutor) Identifier() string { return "google" } |
| 261 | + |
| 262 | +func (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 | + |
| 282 | +func (*slowPoolSimulatedExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { |
| 283 | + return nil, nil |
| 284 | +} |
| 285 | + |
| 286 | +func (*slowPoolSimulatedExecutor) Refresh(context.Context, *Auth) (*Auth, error) { |
| 287 | + return nil, nil |
| 288 | +} |
| 289 | + |
| 290 | +func (*slowPoolSimulatedExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { |
| 291 | + return cliproxyexecutor.Response{}, nil |
| 292 | +} |
| 293 | + |
| 294 | +func (*slowPoolSimulatedExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { |
| 295 | + return nil, nil |
| 296 | +} |
| 297 | + |
| 298 | +func 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 | + |
| 360 | +func 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 | + |
| 404 | +func 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 | +} |