diff --git a/internal/registry/model_registry.go b/internal/registry/model_registry.go index 6ef90f75b..6a34f7065 100644 --- a/internal/registry/model_registry.go +++ b/internal/registry/model_registry.go @@ -1299,17 +1299,23 @@ func modelRegistrationAvailability(registration *ModelRegistration, now time.Tim cooldownSuspended := 0 otherSuspended := 0 + quotaAndOtherSuspended := 0 if registration.SuspendedClients != nil { - for _, reason := range registration.SuspendedClients { + for clientID, reason := range registration.SuspendedClients { if strings.EqualFold(reason, "quota") { cooldownSuspended++ continue } otherSuspended++ + if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { + quotaAndOtherSuspended++ + } } } - effectiveClients := availableClients - expiredClients - otherSuspended + // A credential-wide quota can mark the same client both quota-exceeded and + // suspended. Count that unavailable client only once. + effectiveClients := availableClients - expiredClients - otherSuspended + quotaAndOtherSuspended if effectiveClients < 0 { effectiveClients = 0 } @@ -1470,6 +1476,7 @@ func (r *ModelRegistry) GetAvailableModelsByProvider(provider string) []*ModelIn expiredClients := 0 cooldownSuspended := 0 otherSuspended := 0 + quotaAndOtherSuspended := 0 if ok && registration != nil { if registration.QuotaExceededClients != nil { for clientID, quotaTime := range registration.QuotaExceededClients { @@ -1497,12 +1504,15 @@ func (r *ModelRegistry) GetAvailableModelsByProvider(provider string) []*ModelIn continue } otherSuspended++ + if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { + quotaAndOtherSuspended++ + } } } } availableClients := entry.count - effectiveClients := availableClients - expiredClients - otherSuspended + effectiveClients := availableClients - expiredClients - otherSuspended + quotaAndOtherSuspended if effectiveClients < 0 { effectiveClients = 0 } @@ -1542,8 +1552,11 @@ func (r *ModelRegistry) GetModelCount(modelID string) int { } } suspendedClients := 0 - if registration.SuspendedClients != nil { - suspendedClients = len(registration.SuspendedClients) + for clientID := range registration.SuspendedClients { + if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { + continue + } + suspendedClients++ } result := registration.Count - expiredClients - suspendedClients if result < 0 { diff --git a/internal/registry/model_registry_credential_quota_regression_test.go b/internal/registry/model_registry_credential_quota_regression_test.go new file mode 100644 index 000000000..a20d65fc2 --- /dev/null +++ b/internal/registry/model_registry_credential_quota_regression_test.go @@ -0,0 +1,43 @@ +package registry + +import "testing" + +func TestAvailableModelsKeepsHealthyProviderDuringCredentialQuota(t *testing.T) { + r := newTestModelRegistry() + model := &ModelInfo{ID: "gpt-5.6-luna", OwnedBy: "openai", Type: "openai"} + r.RegisterClient("healthy-api-provider", "codex", []*ModelInfo{model}) + r.RegisterClient("exhausted-oauth-account", "codex", []*ModelInfo{model}) + _, epoch := r.GetModelsAndEpochForClient("exhausted-oauth-account") + if !r.ApplyClientModelProjections("exhausted-oauth-account", epoch, 1, []ClientModelProjection{{ + ModelID: model.ID, Suspended: true, SuspendReason: "credential_quota", QuotaExceeded: true, + }}) { + t.Fatal("failed to apply the observed OAuth quota state") + } + if got := r.GetAvailableModels("openai"); len(got) != 1 { + t.Errorf("/v1/models hid a model with a healthy API provider: got %v", got) + } + if got := r.GetAvailableModelsByProvider("codex"); len(got) != 1 { + t.Errorf("provider catalog hid a model with a healthy API provider: got %v", got) + } + if got := r.GetAvailableModelInfos(); len(got) != 1 { + t.Errorf("metadata catalog hid a model with a healthy API provider: got %v", got) + } + if got := r.GetModelCount(model.ID); got != 1 { + t.Errorf("healthy API provider count = %d, want 1", got) + } + _, epoch = r.GetModelsAndEpochForClient("healthy-api-provider") + if !r.ApplyClientModelProjections("healthy-api-provider", epoch, 1, []ClientModelProjection{{ + ModelID: model.ID, Suspended: true, SuspendReason: "manual", + }}) { + t.Fatal("failed to suspend the second provider") + } + if got := r.GetAvailableModels("openai"); len(got) != 0 { + t.Errorf("catalog should hide a model after both providers become unavailable: %v", got) + } + if got := r.GetAvailableModelsByProvider("codex"); len(got) != 0 { + t.Errorf("provider catalog should hide a model after both providers become unavailable: %v", got) + } + if got := r.GetModelCount(model.ID); got != 0 { + t.Errorf("unavailable provider count = %d, want 0", got) + } +} diff --git a/internal/registry/model_registry_quota_refresh_regression_test.go b/internal/registry/model_registry_quota_refresh_regression_test.go new file mode 100644 index 000000000..4d20c3029 --- /dev/null +++ b/internal/registry/model_registry_quota_refresh_regression_test.go @@ -0,0 +1,65 @@ +package registry + +import ( + "testing" + "time" +) + +// Healthy providers remain discoverable throughout quota windows and refreshes. +func TestCredentialQuotaKeepsHealthyCatalogAcrossRefresh(t *testing.T) { + base := time.Date(2026, 9, 22, 8, 53, 19, 0, time.UTC) + registration := &ModelRegistration{ + Count: 2, + QuotaExceededClients: map[string]*time.Time{"oauth": &base}, + SuspendedClients: map[string]string{"oauth": "credential_quota"}, + } + for _, tc := range []struct { + offset time.Duration + want bool + }{ + {0, true}, {5*time.Minute - time.Nanosecond, true}, + {5 * time.Minute, true}, {6 * time.Minute, true}, + } { + available, _ := modelRegistrationAvailability(registration, base.Add(tc.offset)) + t.Logf("since projection=%v listed=%v; OAuth suspension unchanged", tc.offset, available) + if available != tc.want { + t.Fatalf("at %v: listed=%v, want %v", tc.offset, available, tc.want) + } + } + // A registry refresh clears the old quota timestamp; reapplying an active + // cooldown starts the five-minute window again without another upstream 429. + r := newTestModelRegistry() + models := []*ModelInfo{{ID: "audit-luna", OwnedBy: "openai"}} + r.RegisterClient("healthy", "codex", models) + r.RegisterClient("oauth", "codex", models) + _, epoch := r.GetModelsAndEpochForClient("oauth") + projection := []ClientModelProjection{{ModelID: "audit-luna", Suspended: true, SuspendReason: "credential_quota", QuotaExceeded: true}} + if !r.ApplyClientModelProjections("oauth", epoch, 1, projection) { + t.Fatal("initial projection rejected") + } + if got := len(r.GetAvailableModels("openai")); got != 1 { + t.Fatalf("initial cached catalog size=%d, want 1", got) + } + cached := r.availableModelsCache["openai"] + if cached.expiresAt.IsZero() { + t.Fatal("quota-sensitive catalog has no automatic expiration") + } + old := time.Now().Add(-6 * time.Minute) + r.models["audit-luna"].QuotaExceededClients["oauth"] = &old + // Advance the observed timestamps instead of sleeping. Let the actual + // public listing path detect the expired cache without manual invalidation. + cached.expiresAt = time.Now().Add(-time.Second) + r.availableModelsCache["openai"] = cached + if got := len(r.GetAvailableModels("openai")); got != 1 { + t.Fatalf("after five-minute window: models=%d, want 1", got) + } + r.RegisterClient("oauth", "codex", models) + _, epoch = r.GetModelsAndEpochForClient("oauth") + if !r.ApplyClientModelProjections("oauth", epoch, 2, projection) { + t.Fatal("refreshed projection rejected") + } + if got := len(r.GetAvailableModels("openai")); got != 1 { + t.Fatalf("after refresh with unchanged OAuth cooldown: models=%d, want 1", got) + } + t.Log("healthy catalog stays visible before expiration, after expiration, and after re-registration") +} diff --git a/sdk/cliproxy/auth/catalog_credential_quota_regression_test.go b/sdk/cliproxy/auth/catalog_credential_quota_regression_test.go new file mode 100644 index 000000000..20bfe358d --- /dev/null +++ b/sdk/cliproxy/auth/catalog_credential_quota_regression_test.go @@ -0,0 +1,97 @@ +package auth + +import ( + "context" + "net/http" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" +) + +func TestCredentialQuotaKeepsHealthyCatalogAcrossRestart(t *testing.T) { + withQuotaCooldownEnabled(t) + ctx := context.Background() + m := NewManager(nil, nil, nil) + store := NewFileCooldownStateStore(t.TempDir()) + m.SetCooldownStateStore(store) + r := registry.GetGlobalRegistry() + models := []*registry.ModelInfo{{ID: "audit-luna"}, {ID: "audit-sol"}, {ID: "audit-terra"}} + for _, id := range []string{"audit-healthy-api", "audit-quota-oauth"} { + if _, err := m.Register(ctx, &Auth{ID: id, Provider: "codex", Status: StatusActive}); err != nil { + t.Fatal(err) + } + r.RegisterClient(id, "codex", models) + t.Cleanup(func() { r.UnregisterClient(id) }) + } + for _, model := range models { + m.MarkResult(ctx, Result{AuthID: "audit-quota-oauth", Provider: "codex", Model: model.ID, Success: true}) + } + retry := 3 * time.Hour + m.MarkResult(ctx, Result{ + AuthID: "audit-quota-oauth", Provider: "codex", Model: "audit-sol", + Error: &Error{HTTPStatus: http.StatusTooManyRequests, Code: "usage_limit_reached", Message: "quota exceeded"}, + RetryAfter: &retry, CredentialScope: true, + }) + oauth, _ := m.GetByID("audit-quota-oauth") + listed := map[string]bool{} + for _, model := range r.GetAvailableModels("openai") { + id, _ := model["id"].(string) + listed[id] = true + } + for _, model := range models { + state := oauth.ModelStates[model.ID] + if state == nil || !state.Quota.Exceeded { + t.Fatalf("expected model %s to have Quota.Exceeded = true", model.ID) + } + expectedReason := "credential_quota" + if model.ID == "audit-sol" { + expectedReason = "quota" + } + if state.Quota.Reason != expectedReason { + t.Fatalf("model %s: got Quota.Reason = %q, want %q", model.ID, state.Quota.Reason, expectedReason) + } + t.Logf("model=%s OAuth_reason=%s quota_exceeded=%v healthy_provider_blocked=%v listed=%v", + model.ID, state.Quota.Reason, state.Quota.Exceeded, false, listed[model.ID]) + wantListed := true + if listed[model.ID] != wantListed { + t.Errorf("healthy catalog lost %s: listed=%v, want %v", model.ID, listed[model.ID], wantListed) + } + } + healthy, _ := m.GetByID("audit-healthy-api") + for _, model := range models { + blocked, _, _ := isAuthBlockedForModel(healthy, model.ID, time.Now()) + if blocked { + t.Fatalf("healthy API provider unexpectedly blocked for %s", model.ID) + } + } + // Simulate the same service restart sequence using the real .cds store. + for _, id := range []string{"audit-healthy-api", "audit-quota-oauth"} { + r.UnregisterClient(id) + } + restarted := NewManager(nil, nil, nil) + for _, id := range []string{"audit-healthy-api", "audit-quota-oauth"} { + if _, err := restarted.Register(ctx, &Auth{ID: id, Provider: "codex", Status: StatusActive}); err != nil { + t.Fatal(err) + } + r.RegisterClient(id, "codex", models) + } + restarted.SetCooldownStateStore(store) + if err := restarted.RestoreCooldownStates(ctx); err != nil { + t.Fatal(err) + } + restarted.ReconcileRegistryModelStates(ctx, "audit-quota-oauth") + restored, _ := restarted.GetByID("audit-quota-oauth") + if restored.Quota.Reason != "credential_quota" || !restored.Quota.NextRecoverAt.After(time.Now()) { + t.Fatal("restart did not restore the active credential quota") + } + listedAfterRestart := map[string]bool{} + for _, model := range r.GetAvailableModels("openai") { + id, _ := model["id"].(string) + listedAfterRestart[id] = true + } + if !listedAfterRestart["audit-luna"] || !listedAfterRestart["audit-terra"] || !listedAfterRestart["audit-sol"] { + t.Fatalf("restart hid a healthy provider: %v", listedAfterRestart) + } + t.Log("real .cds persistence + restart + reconciliation keeps all models with a healthy provider visible") +} diff --git a/sdk/cliproxy/auth/scheduler_test.go b/sdk/cliproxy/auth/scheduler_test.go index e2dd7d908..129661653 100644 --- a/sdk/cliproxy/auth/scheduler_test.go +++ b/sdk/cliproxy/auth/scheduler_test.go @@ -2211,8 +2211,9 @@ func TestManager_SchedulerSharesThinkingSuffixCooldownAndRegistryState(t *testin if len(auth.ModelStates) != 1 || auth.ModelStates[baseModel] == nil { t.Fatalf("ModelStates = %+v, want only canonical key %q", auth.ModelStates, baseModel) } - if count := reg.GetModelCount(baseModel); count != 0 { - t.Fatalf("registry model count during cooldown = %d, want 0", count) + // Only thinking-auth-a is cooling down; thinking-auth-b remains available. + if count := reg.GetModelCount(baseModel); count != 1 { + t.Fatalf("registry model count during cooldown = %d, want 1", count) } for _, model := range []string{baseModel, baseModel + "(medium)", baseModel + "(low)"} { got, errPick := manager.scheduler.pickSingle(context.Background(), "gemini", model, cliproxyexecutor.Options{}, nil)