From 09471dd9daba6691e5810e18012e973407260358 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Thu, 3 Sep 2026 17:49:57 +0800 Subject: [PATCH] fix(auth): prevent individual model quota cooldowns from blocking credential - Ignore aggregated single-model quota cooldowns during credential-level availability checks. - Keep credentials eligible unless marked completely unavailable or subject to a credential-wide quota. Closes: #5371 --- .../auth/conductor_alias_cooldown_test.go | 242 ++++++++++++++++++ sdk/cliproxy/auth/selector.go | 10 +- 2 files changed, 251 insertions(+), 1 deletion(-) create mode 100644 sdk/cliproxy/auth/conductor_alias_cooldown_test.go diff --git a/sdk/cliproxy/auth/conductor_alias_cooldown_test.go b/sdk/cliproxy/auth/conductor_alias_cooldown_test.go new file mode 100644 index 00000000..8b0246d4 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_alias_cooldown_test.go @@ -0,0 +1,242 @@ +package auth + +import ( + "context" + "errors" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +func TestManagerExecute_ModelAliasRequestNotBlockedByOtherModelQuotaCooldown(t *testing.T) { + const ( + provider = "antigravity" + requestModel = "gemini-3.6-flash" + targetModel = "gemini-3.6-flash-high" + imageModel = "gemini-3.1-flash-image" + ) + + manager := NewManager(nil, &RoundRobinSelector{}, nil) + executor := &aliasRoutingExecutor{id: provider} + manager.RegisterExecutor(executor) + manager.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{ + provider: {{ + Name: targetModel, + Alias: requestModel, + Fork: true, + }}, + }) + + now := time.Now() + next := now.Add(1 * time.Hour) + + auth := &Auth{ + ID: "antigravity-auth-1", + Provider: provider, + Status: StatusActive, + ModelStates: map[string]*ModelState{ + targetModel: { + Status: StatusActive, + }, + imageModel: { + Status: StatusError, + Unavailable: true, + NextRetryAfter: next, + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: next, + }, + }, + }, + } + updateAggregatedAvailability(auth, now) + if !auth.Quota.Exceeded { + t.Fatalf("precondition failed: auth.Quota.Exceeded should be true after updateAggregatedAvailability") + } + if auth.Unavailable { + t.Fatalf("precondition failed: auth.Unavailable should be false since targetModel is active") + } + + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, provider, []*registry.ModelInfo{ + {ID: requestModel}, + {ID: targetModel}, + {ID: imageModel}, + }) + t.Cleanup(func() { + reg.UnregisterClient(auth.ID) + }) + manager.RefreshSchedulerEntry(auth.ID) + + resp, errExecute := manager.Execute( + context.Background(), + []string{provider}, + cliproxyexecutor.Request{Model: requestModel}, + cliproxyexecutor.Options{}, + ) + if errExecute != nil { + t.Fatalf("Execute() error = %v, want success", errExecute) + } + if string(resp.Payload) != targetModel { + t.Fatalf("Execute() payload = %q, want %q", string(resp.Payload), targetModel) + } +} + +func TestManagerSelectAuth_ModelAliasRequestNotBlockedByOtherModelQuotaCooldown(t *testing.T) { + const ( + provider = "antigravity" + requestModel = "gemini-3.6-flash" + targetModel = "gemini-3.6-flash-high" + imageModel = "gemini-3.1-flash-image" + ) + + manager := NewManager(nil, &RoundRobinSelector{}, nil) + executor := &aliasRoutingExecutor{id: provider} + manager.RegisterExecutor(executor) + manager.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{ + provider: {{ + Name: targetModel, + Alias: requestModel, + Fork: true, + }}, + }) + + now := time.Now() + next := now.Add(1 * time.Hour) + + auth := &Auth{ + ID: "antigravity-auth-2", + Provider: provider, + Status: StatusActive, + ModelStates: map[string]*ModelState{ + targetModel: { + Status: StatusActive, + }, + imageModel: { + Status: StatusError, + Unavailable: true, + NextRetryAfter: next, + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: next, + }, + }, + }, + } + updateAggregatedAvailability(auth, now) + if !auth.Quota.Exceeded { + t.Fatalf("precondition failed: auth.Quota.Exceeded should be true after updateAggregatedAvailability") + } + if auth.Unavailable { + t.Fatalf("precondition failed: auth.Unavailable should be false since targetModel is active") + } + + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, provider, []*registry.ModelInfo{ + {ID: requestModel}, + {ID: targetModel}, + {ID: imageModel}, + }) + t.Cleanup(func() { + reg.UnregisterClient(auth.ID) + }) + manager.RefreshSchedulerEntry(auth.ID) + + selected, errSelect := manager.SelectAuth( + context.Background(), + provider, + requestModel, + cliproxyexecutor.Options{}, + ) + if errSelect != nil { + t.Fatalf("SelectAuth() error = %v, want success", errSelect) + } + if selected == nil || selected.ID != auth.ID { + t.Fatalf("SelectAuth() selected = %#v, want %s", selected, auth.ID) + } +} + +func TestManagerExecute_ModelAliasRequestBlockedWhenTargetModelInQuotaCooldown(t *testing.T) { + const ( + provider = "antigravity" + requestModel = "gemini-3.6-flash" + targetModel = "gemini-3.6-flash-high" + ) + + manager := NewManager(nil, &RoundRobinSelector{}, nil) + executor := &aliasRoutingExecutor{id: provider} + manager.RegisterExecutor(executor) + manager.SetOAuthModelAlias(map[string][]internalconfig.OAuthModelAlias{ + provider: {{ + Name: targetModel, + Alias: requestModel, + Fork: true, + }}, + }) + + now := time.Now() + next := now.Add(1 * time.Hour) + + auth := &Auth{ + ID: "antigravity-auth-3", + Provider: provider, + Status: StatusActive, + ModelStates: map[string]*ModelState{ + targetModel: { + Status: StatusError, + Unavailable: true, + NextRetryAfter: next, + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: next, + }, + }, + }, + } + updateAggregatedAvailability(auth, now) + + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(auth.ID, provider, []*registry.ModelInfo{ + {ID: requestModel}, + {ID: targetModel}, + }) + t.Cleanup(func() { + reg.UnregisterClient(auth.ID) + }) + manager.RefreshSchedulerEntry(auth.ID) + + _, errExecute := manager.Execute( + context.Background(), + []string{provider}, + cliproxyexecutor.Request{Model: requestModel}, + cliproxyexecutor.Options{}, + ) + if errExecute == nil { + t.Fatal("Execute() error = nil, want cooldown error") + } + var cooldownErr *modelCooldownError + if !errors.As(errExecute, &cooldownErr) { + t.Fatalf("Execute() error = %T (%v), want *modelCooldownError", errExecute, errExecute) + } + if cooldownErr.model != requestModel { + t.Fatalf("cooldown model = %q, want %q", cooldownErr.model, requestModel) + } +} diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 510a8a52..c472c715 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -843,7 +843,15 @@ func isAuthBlockedForModel(auth *Auth, model string, now time.Time) (bool, block } return availabilityBlock(auth.Unavailable, auth.Quota.Exceeded, auth.NextRetryAfter, auth.Quota.NextRecoverAt, now) } - return availabilityBlock(auth.Unavailable, auth.Quota.Exceeded, auth.NextRetryAfter, auth.Quota.NextRecoverAt, now) + quotaExceeded := auth.Quota.Exceeded + // When model is empty and the credential has individual model states, auth.Quota.Exceeded + // is an aggregate of single-model quota cooldowns (reason "quota"). As long as the credential + // itself is not unavailable (not all models failed) and not under a credential-wide quota, + // do not treat individual model cooldowns as blocking the entire credential. + if len(auth.ModelStates) > 0 && auth.Quota.Reason != "credential_quota" && !auth.Unavailable { + quotaExceeded = false + } + return availabilityBlock(auth.Unavailable, quotaExceeded, auth.NextRetryAfter, auth.Quota.NextRecoverAt, now) } func availabilityBlock(unavailable, quotaExceeded bool, nextRetryAfter, nextRecoverAt, now time.Time) (bool, blockReason, time.Time) { -- 2.51.2