| @@ -5,6 +5,7 @@ import ( · 4 unmodified lines |
| 5 | "errors" | 5 | "errors" |
| 6 | "fmt" | 6 | "fmt" |
| 7 | "net/http" | 7 | "net/http" |
| | 8 | "slices" |
| 8 | "sync" | 9 | "sync" |
| 9 | "testing" | 10 | "testing" |
| 10 | "time" | 11 | "time" |
| @@ -1353,6 +1354,99 @@ func TestManager_RequestScopedErrorStopsCredentialFallbackWithoutSuspendingAuth( · 1342 unmodified lines |
| 1353 | } | 1354 | } |
| 1354 | } | 1355 | } |
| 1355 | | 1356 | |
| | 1357 | func TestManager_DeepSeekInsufficientBalanceRotatesCredentialAndRebindsSession(t *testing.T) { |
| | 1358 | m := NewManager(nil, nil, nil) |
| | 1359 | m.SetRetryConfig(2, 30*time.Second, 0) |
| | 1360 | affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ |
| | 1361 | Fallback: &RoundRobinSelector{}, |
| | 1362 | TTL: time.Hour, |
| | 1363 | }) |
| | 1364 | defer affinity.Stop() |
| | 1365 | m.SetSelector(affinity) |
| | 1366 | |
| | 1367 | const provider = "openai-compatibility" |
| | 1368 | const model = "deepseek-v4-pro" |
| | 1369 | |
| | 1370 | executor := &authFallbackExecutor{ |
| | 1371 | id: provider, |
| | 1372 | executeErrors: map[string]error{ |
| | 1373 | "aa-empty-balance": &Error{ |
| | 1374 | HTTPStatus: http.StatusPaymentRequired, |
| | 1375 | Message: `{"error":{"message":"Insufficient Balance","type":"unknown_error","param":null,"code":"invalid_request_error"}}`, |
| | 1376 | }, |
| | 1377 | }, |
| | 1378 | } |
| | 1379 | m.RegisterExecutor(executor) |
| | 1380 | |
| | 1381 | depletedAuth := &Auth{ID: "aa-empty-balance", Provider: provider} |
| | 1382 | availableAuth := &Auth{ID: "bb-available-balance", Provider: provider} |
| | 1383 | |
| | 1384 | reg := registry.GetGlobalRegistry() |
| | 1385 | models := []*registry.ModelInfo{{ID: model}} |
| | 1386 | reg.RegisterClient(depletedAuth.ID, provider, models) |
| | 1387 | reg.RegisterClient(availableAuth.ID, provider, models) |
| | 1388 | t.Cleanup(func() { |
| | 1389 | reg.UnregisterClient(depletedAuth.ID) |
| | 1390 | reg.UnregisterClient(availableAuth.ID) |
| | 1391 | }) |
| | 1392 | |
| | 1393 | if _, errRegister := m.Register(context.Background(), depletedAuth); errRegister != nil { |
| | 1394 | t.Fatalf("register depleted auth: %v", errRegister) |
| | 1395 | } |
| | 1396 | if _, errRegister := m.Register(context.Background(), availableAuth); errRegister != nil { |
| | 1397 | t.Fatalf("register available auth: %v", errRegister) |
| | 1398 | } |
| | 1399 | |
| | 1400 | opts := cliproxyexecutor.Options{Metadata: map[string]any{ |
| | 1401 | cliproxyexecutor.DerivedSessionIDMetadataKey: "deepseek-insufficient-balance", |
| | 1402 | }} |
| | 1403 | beforeExecute := time.Now() |
| | 1404 | resp, errExecute := m.Execute( |
| | 1405 | context.Background(), |
| | 1406 | []string{provider}, |
| | 1407 | cliproxyexecutor.Request{Model: model}, |
| | 1408 | opts, |
| | 1409 | ) |
| | 1410 | if errExecute != nil { |
| | 1411 | t.Fatalf("expected fallback to the next credential, got error: %v", errExecute) |
| | 1412 | } |
| | 1413 | if got := string(resp.Payload); got != availableAuth.ID { |
| | 1414 | t.Fatalf("served by %q, want %q", got, availableAuth.ID) |
| | 1415 | } |
| | 1416 | |
| | 1417 | resp, errExecute = m.Execute( |
| | 1418 | context.Background(), |
| | 1419 | []string{provider}, |
| | 1420 | cliproxyexecutor.Request{Model: model}, |
| | 1421 | opts, |
| | 1422 | ) |
| | 1423 | if errExecute != nil { |
| | 1424 | t.Fatalf("expected rebound session to use the next credential, got error: %v", errExecute) |
| | 1425 | } |
| | 1426 | if got := string(resp.Payload); got != availableAuth.ID { |
| | 1427 | t.Fatalf("rebound session served by %q, want %q", got, availableAuth.ID) |
| | 1428 | } |
| | 1429 | wantCalls := []string{depletedAuth.ID, availableAuth.ID, availableAuth.ID} |
| | 1430 | if calls := executor.ExecuteCalls(); !slices.Equal(calls, wantCalls) { |
| | 1431 | t.Fatalf("credential calls = %v, want %v", calls, wantCalls) |
| | 1432 | } |
| | 1433 | |
| | 1434 | updatedDepleted, ok := m.GetByID(depletedAuth.ID) |
| | 1435 | if !ok || updatedDepleted == nil { |
| | 1436 | t.Fatal("expected depleted auth to remain registered") |
| | 1437 | } |
| | 1438 | state := updatedDepleted.ModelStates[model] |
| | 1439 | if state == nil { |
| | 1440 | t.Fatal("expected the depleted credential to be cooled down for the model") |
| | 1441 | } |
| | 1442 | if !state.Unavailable { |
| | 1443 | t.Fatal("expected the depleted credential to be unavailable for the model") |
| | 1444 | } |
| | 1445 | if state.NextRetryAfter.Before(beforeExecute.Add(29 * time.Minute)) { |
| | 1446 | t.Fatalf("cooldown expires at %v, want approximately 30 minutes", state.NextRetryAfter) |
| | 1447 | } |
| | 1448 | } |
| | 1449 | |
| 1356 | // TestManager_UnknownUpstreamErrorRotatesAndPenalizesModelOnly pins the upstream | 1450 | // TestManager_UnknownUpstreamErrorRotatesAndPenalizesModelOnly pins the upstream |
| 1357 | // 500 "status":"UNKNOWN" contract. It is an upstream internal failure, not a | 1451 | // 500 "status":"UNKNOWN" contract. It is an upstream internal failure, not a |
| 1358 | // request fault, so the request must fall through to the next credential. The | 1452 | // request fault, so the request must fall through to the next credential. The |