| @@ -5,6 +5,7 @@ import ( · 4 unmodified lines |
| 5 | 5 | "errors" |
| 6 | 6 | "fmt" |
| 7 | 7 | "net/http" |
| 8 | + "slices" |
| 8 | 9 | "sync" |
| 9 | 10 | "testing" |
| 10 | 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 | 1450 | // TestManager_UnknownUpstreamErrorRotatesAndPenalizesModelOnly pins the upstream |
| 1357 | 1451 | // 500 "status":"UNKNOWN" contract. It is an upstream internal failure, not a |
| 1358 | 1452 | // request fault, so the request must fall through to the next credential. The |