diff --git a/internal/registry/model_registry.go b/internal/registry/model_registry.go index 15e495e7a..b1c3bdee4 100644 --- a/internal/registry/model_registry.go +++ b/internal/registry/model_registry.go @@ -1305,6 +1305,16 @@ func (r *ModelRegistry) GetAvailableModels(handlerType string) []map[string]any return models } +// isQuotaSuspension reports whether a suspension is a quota cooldown. A model in a +// quota cooldown stays listed: the limit is temporary, and a request meanwhile gets a +// 429 with its retry time instead of an unknown model. A credential-wide quota +// ("credential_quota") suspends every model of that credential, so it is a quota +// cooldown too; counting it as another failure hides all of them when no other +// credential serves them. +func isQuotaSuspension(reason string) bool { + return strings.EqualFold(reason, "quota") || strings.EqualFold(reason, "credential_quota") +} + func modelRegistrationAvailability(registration *ModelRegistration, now time.Time) (bool, time.Time) { if registration == nil { return false, time.Time{} @@ -1328,23 +1338,23 @@ func modelRegistrationAvailability(registration *ModelRegistration, now time.Tim cooldownSuspended := 0 otherSuspended := 0 - quotaAndOtherSuspended := 0 + quotaAndSuspended := 0 if registration.SuspendedClients != nil { for clientID, reason := range registration.SuspendedClients { - if strings.EqualFold(reason, "quota") { + if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { + quotaAndSuspended++ + } + if isQuotaSuspension(reason) { cooldownSuspended++ continue } otherSuspended++ - if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { - quotaAndOtherSuspended++ - } } } - // A credential-wide quota can mark the same client both quota-exceeded and + // A quota cooldown can mark the same client both quota-exceeded and // suspended. Count that unavailable client only once. - effectiveClients := availableClients - expiredClients - otherSuspended + quotaAndOtherSuspended + effectiveClients := availableClients - expiredClients - cooldownSuspended - otherSuspended + quotaAndSuspended if effectiveClients < 0 { effectiveClients = 0 } @@ -1505,7 +1515,7 @@ func (r *ModelRegistry) GetAvailableModelsByProvider(provider string) []*ModelIn expiredClients := 0 cooldownSuspended := 0 otherSuspended := 0 - quotaAndOtherSuspended := 0 + quotaAndSuspended := 0 if ok && registration != nil { if registration.QuotaExceededClients != nil { for clientID, quotaTime := range registration.QuotaExceededClients { @@ -1528,20 +1538,20 @@ func (r *ModelRegistry) GetAvailableModelsByProvider(provider string) []*ModelIn if p, okProvider := r.clientProviders[clientID]; !okProvider || p != provider { continue } - if strings.EqualFold(reason, "quota") { + if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { + quotaAndSuspended++ + } + if isQuotaSuspension(reason) { cooldownSuspended++ continue } otherSuspended++ - if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { - quotaAndOtherSuspended++ - } } } } availableClients := entry.count - effectiveClients := availableClients - expiredClients - otherSuspended + quotaAndOtherSuspended + effectiveClients := availableClients - expiredClients - cooldownSuspended - otherSuspended + quotaAndSuspended if effectiveClients < 0 { effectiveClients = 0 } diff --git a/internal/registry/model_registry_sole_credential_quota_test.go b/internal/registry/model_registry_sole_credential_quota_test.go new file mode 100644 index 000000000..4a5a141cb --- /dev/null +++ b/internal/registry/model_registry_sole_credential_quota_test.go @@ -0,0 +1,84 @@ +package registry + +import ( + "testing" + "time" +) + +// A credential-wide 429 marks the model that received it "quota" and the credential's +// other models "credential_quota". With no other credential to serve them, all of them +// must stay listed like any quota cooldown instead of leaving only the model that +// received the 429 until a later request re-projects the credential. +func TestCredentialQuotaKeepsSoleCredentialModelsListed(t *testing.T) { + r := newTestModelRegistry() + models := []*ModelInfo{ + {ID: "claude-sonnet-5", OwnedBy: "anthropic", Type: "claude"}, + {ID: "claude-opus-5-5", OwnedBy: "anthropic", Type: "claude"}, + {ID: "claude-haiku-4-5", OwnedBy: "anthropic", Type: "claude"}, + } + r.RegisterClient("claude-oauth", "claude", models) + _, epoch := r.GetModelsAndEpochForClient("claude-oauth") + if !r.ApplyClientModelProjections("claude-oauth", epoch, 1, []ClientModelProjection{ + {ModelID: "claude-sonnet-5", Suspended: true, SuspendReason: "quota", QuotaExceeded: true}, + {ModelID: "claude-opus-5-5", Suspended: true, SuspendReason: "credential_quota", QuotaExceeded: true}, + // A sibling without a recorded state of its own inherits the credential's suspension. + {ModelID: "claude-haiku-4-5", Suspended: true, SuspendReason: "credential_quota"}, + }) { + t.Fatal("failed to apply the credential-wide quota state") + } + if got := len(r.GetAvailableModels("claude")); got != len(models) { + t.Errorf("/v1/models lists %d models during a credential-wide quota, want %d", got, len(models)) + } + if got := len(r.GetAvailableModelsByProvider("claude")); got != len(models) { + t.Errorf("provider catalog lists %d models during a credential-wide quota, want %d", got, len(models)) + } + if got := len(r.GetAvailableModelInfos()); got != len(models) { + t.Errorf("metadata catalog lists %d models during a credential-wide quota, want %d", got, len(models)) + } + + // After the quota window, the suspension alone still keeps each model listed. + registration := r.models["claude-opus-5-5"] + past := time.Now().Add(-2 * modelQuotaExceededWindow) + registration.QuotaExceededClients["claude-oauth"] = &past + if available, _ := modelRegistrationAvailability(registration, time.Now()); !available { + t.Error("a credential-wide quota hid the model after the quota window") + } + if got := len(r.GetAvailableModelsByProvider("claude")); got != len(models) { + t.Errorf("provider catalog lists %d models after the quota window, want %d", got, len(models)) + } + + // A suspension that is not a quota cooldown still hides the model. + _, epoch = r.GetModelsAndEpochForClient("claude-oauth") + if !r.ApplyClientModelProjections("claude-oauth", epoch, 2, []ClientModelProjection{ + {ModelID: "claude-opus-5-5", Suspended: true, SuspendReason: "unauthorized"}, + }) { + t.Fatal("failed to apply the unauthorized state") + } + for _, model := range r.GetAvailableModelsByProvider("claude") { + if model.ID == "claude-opus-5-5" { + t.Error("a model whose only credential is unauthorized is still listed") + } + } + + // When a second credential is suspended for a non-quota reason, a sibling model + // with only credential_quota (no active 5-minute quota timestamp) must also be hidden. + r.RegisterClient("claude-api-backup", "claude", []*ModelInfo{ + {ID: "claude-haiku-4-5", OwnedBy: "anthropic", Type: "claude"}, + }) + _, backupEpoch := r.GetModelsAndEpochForClient("claude-api-backup") + if !r.ApplyClientModelProjections("claude-api-backup", backupEpoch, 1, []ClientModelProjection{ + {ModelID: "claude-haiku-4-5", Suspended: true, SuspendReason: "manual"}, + }) { + t.Fatal("failed to apply the manual suspension to backup client") + } + for _, model := range r.GetAvailableModelsByProvider("claude") { + if model.ID == "claude-haiku-4-5" { + t.Error("a model with one credential_quota client and one manual-suspended client is still listed in provider catalog") + } + } + for _, model := range r.GetAvailableModels("claude") { + if id, _ := model["id"].(string); id == "claude-haiku-4-5" { + t.Error("a model with one credential_quota client and one manual-suspended client is still listed in /v1/models") + } + } +}