fix(antigravity): bypass quota cooldowns and credit hints when cooling is disabled

- Bypass short cooldown checks and recording in execution flows when cooling is disabled globally or per auth.
- Skip marking credits permanently disabled and refreshing credit hints when cooling is disabled.
- Export quota cooldown status helper functions for auth and configuration evaluations.

Closes: #4793
This commit is contained in:
Luis Pater
2026-09-02 22:01:55 +08:00
parent d0fb44ca95
commit 272c1cff4e
5 changed files with 314 additions and 26 deletions

View File

@@ -282,7 +282,7 @@ func clearAntigravityCreditsFailureState(auth *cliproxyauth.Auth) {
antigravityCreditsFailureByAuth.Delete(strings.TrimSpace(auth.ID))
}
func markAntigravityCreditsPermanentlyDisabled(auth *cliproxyauth.Auth) {
if auth == nil || strings.TrimSpace(auth.ID) == "" {
if auth == nil || strings.TrimSpace(auth.ID) == "" || antigravityCoolingDisabled(auth, nil) {
return
}
authID := strings.TrimSpace(auth.ID)
@@ -343,7 +343,7 @@ func newAntigravityStatusErr(statusCode int, body []byte) statusErr {
return err
}
func (e *AntigravityExecutor) maybeRefreshAntigravityCreditsHint(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) {
if e == nil || auth == nil || !antigravityCreditsRetryEnabled(e.cfg) {
if e == nil || auth == nil || !antigravityCreditsRetryEnabled(e.cfg) || antigravityCoolingDisabled(auth, e.cfg) {
return
}
if ctx != nil && ctx.Err() != nil {
@@ -561,6 +561,10 @@ func antigravityShortCooldownKVKey(auth *cliproxyauth.Auth, modelName string) st
return "cpa:antigravity:short-cooldown:" + authID + ":" + homekv.HashKeyPart(modelName)
}
func antigravityCoolingDisabled(auth *cliproxyauth.Auth, cfg *config.Config) bool {
return cliproxyauth.QuotaCooldownDisabledForAuthWithConfig(auth, cfg)
}
func antigravityIsInShortCooldown(auth *cliproxyauth.Auth, modelName string, now time.Time) (bool, time.Duration) {
inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(context.Background(), auth, modelName, now)
if errCooldown != nil {
@@ -571,6 +575,9 @@ func antigravityIsInShortCooldown(auth *cliproxyauth.Auth, modelName string, now
}
func antigravityIsInShortCooldownRequired(ctx context.Context, auth *cliproxyauth.Auth, modelName string, now time.Time) (bool, time.Duration, error) {
if antigravityCoolingDisabled(auth, nil) {
return false, 0, nil
}
kvKey := antigravityShortCooldownKVKey(auth, modelName)
client, homeMode, errClient := currentAntigravityKVClient()
if homeMode {
@@ -626,6 +633,9 @@ func markAntigravityShortCooldown(auth *cliproxyauth.Auth, modelName string, now
}
func markAntigravityShortCooldownRequired(ctx context.Context, auth *cliproxyauth.Auth, modelName string, now time.Time, duration time.Duration) error {
if antigravityCoolingDisabled(auth, nil) {
return nil
}
kvKey := antigravityShortCooldownKVKey(auth, modelName)
client, homeMode, errClient := currentAntigravityKVClient()
if homeMode {

View File

@@ -0,0 +1,262 @@
package executor
import (
"context"
"io"
"net/http"
"strings"
"sync"
"testing"
"time"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
)
func TestAntigravityDisableCooling_ShortCooldownBypassedInHomeMode(t *testing.T) {
resetAntigravityCreditsRetryState()
t.Cleanup(resetAntigravityCreditsRetryState)
cliproxyauth.SetQuotaCooldownDisabled(true)
t.Cleanup(func() { cliproxyauth.SetQuotaCooldownDisabled(false) })
client := newFakeAntigravityKVClient()
useFakeAntigravityKVClient(t, client, true, nil)
cfg := &config.Config{
DisableCooling: true,
Home: config.HomeConfig{
Enabled: true,
},
}
exec := NewAntigravityExecutor(cfg)
auth := &cliproxyauth.Auth{
ID: "home-cooling-disabled-auth",
Metadata: map[string]any{
"access_token": "token",
"project_id": "test-project",
},
}
modelName := "claude-sonnet-4-5"
now := time.Now()
duration := 30 * time.Second
// 1. In Home mode with DisableCooling, marking short cooldown should be a no-op
if errMark := markAntigravityShortCooldownRequired(context.Background(), auth, modelName, now, duration); errMark != nil {
t.Fatalf("markAntigravityShortCooldownRequired() error = %v", errMark)
}
if client.setCount != 0 {
t.Fatalf("KVSet count = %d, want 0 when DisableCooling is true", client.setCount)
}
// 2. Pre-populate KV key manually to simulate existing key; read should still return false when cooling disabled
antigravityShortCooldownByAuth = sync.Map{}
client.values[antigravityShortCooldownKVKey(auth, modelName)] = []byte("9999999999999999999")
inCooldown, remaining, errRead := antigravityIsInShortCooldownRequired(context.Background(), auth, modelName, now)
if errRead != nil {
t.Fatalf("antigravityIsInShortCooldownRequired() error = %v", errRead)
}
if inCooldown || remaining > 0 {
t.Fatalf("inCooldown = %v, remaining = %v, want false/0 when DisableCooling is true", inCooldown, remaining)
}
// 3. In Execute, upstream 429 RATE_LIMIT_EXCEEDED should not record short cooldown to KV
client.setCount = 0
upstreamResp := `{"error":{"code":429,"message":"Rate limit exceeded","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"RATE_LIMIT_EXCEEDED"}]}}`
execCtx := context.WithValue(context.Background(), "cliproxy.roundtripper", roundTripperFunc(func(req *http.Request) (*http.Response, error) {
h := make(http.Header)
h.Set("Retry-After", "10")
h.Set("Content-Type", "application/json")
return &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: h,
Body: io.NopCloser(strings.NewReader(upstreamResp)),
}, nil
}))
req := cliproxyexecutor.Request{
Model: modelName,
Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`),
}
opts := cliproxyexecutor.Options{
SourceFormat: sdktranslator.FormatClaude,
}
_, _ = exec.Execute(execCtx, auth, req, opts)
if client.setCount != 0 {
t.Fatalf("Execute recorded short cooldown to KV store (setCount = %d), want 0 when cooling disabled", client.setCount)
}
}
func TestAntigravityDisableCooling_ExecuteStreamBypassesShortCooldown(t *testing.T) {
resetAntigravityCreditsRetryState()
t.Cleanup(resetAntigravityCreditsRetryState)
client := newFakeAntigravityKVClient()
useFakeAntigravityKVClient(t, client, true, nil)
cfg := &config.Config{
DisableCooling: true,
Home: config.HomeConfig{
Enabled: true,
},
}
exec := NewAntigravityExecutor(cfg)
auth := &cliproxyauth.Auth{
ID: "home-cooling-disabled-auth-stream",
Metadata: map[string]any{
"access_token": "token",
"project_id": "test-project",
},
}
modelName := "claude-sonnet-4-5"
upstreamResp := `{"error":{"code":429,"message":"Rate limit exceeded","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"RATE_LIMIT_EXCEEDED"}]}}`
execCtx := context.WithValue(context.Background(), "cliproxy.roundtripper", roundTripperFunc(func(req *http.Request) (*http.Response, error) {
h := make(http.Header)
h.Set("Retry-After", "10")
h.Set("Content-Type", "application/json")
return &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: h,
Body: io.NopCloser(strings.NewReader(upstreamResp)),
}, nil
}))
req := cliproxyexecutor.Request{
Model: modelName,
Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`),
}
opts := cliproxyexecutor.Options{
SourceFormat: sdktranslator.FormatClaude,
}
_, _ = exec.ExecuteStream(execCtx, auth, req, opts)
if client.setCount != 0 {
t.Fatalf("ExecuteStream recorded short cooldown to KV store (setCount = %d), want 0 when cooling disabled", client.setCount)
}
}
func TestAntigravityDisableCooling_AuthOverrideBypassesShortCooldown(t *testing.T) {
resetAntigravityCreditsRetryState()
t.Cleanup(resetAntigravityCreditsRetryState)
client := newFakeAntigravityKVClient()
useFakeAntigravityKVClient(t, client, false, nil)
cfg := &config.Config{
DisableCooling: false,
}
exec := NewAntigravityExecutor(cfg)
auth := &cliproxyauth.Auth{
ID: "override-cooling-disabled-auth",
Metadata: map[string]any{
"disable_cooling": true,
"access_token": "token",
"project_id": "test-project",
},
}
modelName := "claude-sonnet-4-5"
now := time.Now()
duration := 30 * time.Second
// In memory map: marking short cooldown should be skipped for auth with disable_cooling override
if errMark := markAntigravityShortCooldownRequired(context.Background(), auth, modelName, now, duration); errMark != nil {
t.Fatalf("markAntigravityShortCooldownRequired() error = %v", errMark)
}
if _, loaded := antigravityShortCooldownByAuth.Load(antigravityShortCooldownKey(auth, modelName)); loaded {
t.Fatalf("antigravityShortCooldownByAuth stored key, want skipped for auth with disable_cooling override")
}
// Pre-populate in-memory map; read should return false
antigravityShortCooldownByAuth.Store(antigravityShortCooldownKey(auth, modelName), now.Add(time.Hour))
inCooldown, remaining, errRead := antigravityIsInShortCooldownRequired(context.Background(), auth, modelName, now)
if errRead != nil {
t.Fatalf("antigravityIsInShortCooldownRequired() error = %v", errRead)
}
if inCooldown || remaining > 0 {
t.Fatalf("inCooldown = %v, remaining = %v, want false/0 when auth has disable_cooling override", inCooldown, remaining)
}
// In Execute, upstream 429 RATE_LIMIT_EXCEEDED should not record short cooldown to memory map
antigravityShortCooldownByAuth = sync.Map{}
upstreamResp := `{"error":{"code":429,"message":"Rate limit exceeded","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"RATE_LIMIT_EXCEEDED"}]}}`
execCtx := context.WithValue(context.Background(), "cliproxy.roundtripper", roundTripperFunc(func(req *http.Request) (*http.Response, error) {
h := make(http.Header)
h.Set("Retry-After", "10")
h.Set("Content-Type", "application/json")
return &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: h,
Body: io.NopCloser(strings.NewReader(upstreamResp)),
}, nil
}))
req := cliproxyexecutor.Request{
Model: modelName,
Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`),
}
opts := cliproxyexecutor.Options{
SourceFormat: sdktranslator.FormatClaude,
}
_, _ = exec.Execute(execCtx, auth, req, opts)
if _, loaded := antigravityShortCooldownByAuth.Load(antigravityShortCooldownKey(auth, modelName)); loaded {
t.Fatalf("Execute recorded short cooldown to memory map, want skipped when auth has disable_cooling override")
}
}
func TestAntigravityDisableCooling_CreditsHintRefreshBypassedInHomeMode(t *testing.T) {
resetAntigravityCreditsRetryState()
t.Cleanup(resetAntigravityCreditsRetryState)
client := newFakeAntigravityKVClient()
useFakeAntigravityKVClient(t, client, true, nil)
cfg := &config.Config{
DisableCooling: true,
Home: config.HomeConfig{
Enabled: true,
},
QuotaExceeded: config.QuotaExceeded{
AntigravityCredits: true,
},
}
exec := NewAntigravityExecutor(cfg)
auth := &cliproxyauth.Auth{
ID: "home-refresh-cooling-disabled-auth",
Metadata: map[string]any{
"access_token": "token",
"project_id": "test-project",
},
}
exec.maybeRefreshAntigravityCreditsHint(context.Background(), auth, "token")
if client.setNXCount != 0 {
t.Fatalf("KVSetNX count = %d, want 0 when DisableCooling is true", client.setNXCount)
}
}
func TestAntigravityDisableCooling_CreditsPermanentlyDisabledBypassed(t *testing.T) {
resetAntigravityCreditsRetryState()
t.Cleanup(resetAntigravityCreditsRetryState)
client := newFakeAntigravityKVClient()
useFakeAntigravityKVClient(t, client, true, nil)
auth := &cliproxyauth.Auth{
ID: "home-permanently-disabled-auth",
Metadata: map[string]any{
"disable_cooling": true,
},
}
markAntigravityCreditsPermanentlyDisabled(auth)
if client.setCount != 0 {
t.Fatalf("KVSet count = %d, want 0 when DisableCooling is true", client.setCount)
}
if cliproxyauth.HasKnownAntigravityCreditsHint(auth.ID) {
t.Fatalf("credits hint was stored for auth with DisableCooling true")
}
}

View File

@@ -28,12 +28,14 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
return resp, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"}
}
baseModel := thinking.ParseSuffix(req.Model).ModelName
if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil {
return resp, homeKVUnavailableStatusErr(errCooldown)
} else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) {
log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining)
d := remaining
return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
if !antigravityCoolingDisabled(auth, e.cfg) {
if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil {
return resp, homeKVUnavailableStatusErr(errCooldown)
} else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) {
log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining)
d := remaining
return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
}
}
isClaude := strings.Contains(strings.ToLower(baseModel), "claude")
@@ -137,7 +139,7 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
decision := decideAntigravity429(bodyBytes)
switch decision.kind {
case antigravity429DecisionShortCooldownSwitchAuth:
if decision.retryAfter != nil && *decision.retryAfter > 0 {
if decision.retryAfter != nil && *decision.retryAfter > 0 && !antigravityCoolingDisabled(auth, e.cfg) {
if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil {
err = homeKVUnavailableStatusErr(errMarkCooldown)
return resp, err
@@ -145,7 +147,7 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel)
}
case antigravity429DecisionFullQuotaExhausted:
if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) {
if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) && !antigravityCoolingDisabled(auth, e.cfg) {
markAntigravityCreditsPermanentlyDisabled(auth)
}
// No credits logic - just fall through to error return below
@@ -182,12 +184,14 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
// executeClaudeNonStream performs a claude non-streaming request to the Antigravity API.
func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {
baseModel := thinking.ParseSuffix(req.Model).ModelName
if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil {
return resp, homeKVUnavailableStatusErr(errCooldown)
} else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) {
log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining)
d := remaining
return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
if !antigravityCoolingDisabled(auth, e.cfg) {
if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil {
return resp, homeKVUnavailableStatusErr(errCooldown)
} else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) {
log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining)
d := remaining
return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
}
}
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
@@ -294,7 +298,7 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *
switch decision.kind {
case antigravity429DecisionShortCooldownSwitchAuth:
if decision.retryAfter != nil && *decision.retryAfter > 0 {
if decision.retryAfter != nil && *decision.retryAfter > 0 && !antigravityCoolingDisabled(auth, e.cfg) {
if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil {
err = homeKVUnavailableStatusErr(errMarkCooldown)
return resp, err
@@ -302,7 +306,7 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *
log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel)
}
case antigravity429DecisionFullQuotaExhausted:
if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) {
if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) && !antigravityCoolingDisabled(auth, e.cfg) {
markAntigravityCreditsPermanentlyDisabled(auth)
}
// No credits logic - just fall through to error return below

View File

@@ -27,12 +27,14 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
baseModel := thinking.ParseSuffix(req.Model).ModelName
ctx = context.WithValue(ctx, "alt", "")
if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil {
return nil, homeKVUnavailableStatusErr(errCooldown)
} else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) {
log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining)
d := remaining
return nil, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
if !antigravityCoolingDisabled(auth, e.cfg) {
if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil {
return nil, homeKVUnavailableStatusErr(errCooldown)
} else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) {
log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining)
d := remaining
return nil, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
}
}
reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
@@ -140,7 +142,7 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
switch decision.kind {
case antigravity429DecisionShortCooldownSwitchAuth:
if decision.retryAfter != nil && *decision.retryAfter > 0 {
if decision.retryAfter != nil && *decision.retryAfter > 0 && !antigravityCoolingDisabled(auth, e.cfg) {
if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil {
err = homeKVUnavailableStatusErr(errMarkCooldown)
return nil, err
@@ -148,7 +150,7 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
log.Debugf("antigravity executor: short quota cooldown (%s) for model %s recorded", *decision.retryAfter, baseModel)
}
case antigravity429DecisionFullQuotaExhausted:
if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) {
if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) && !antigravityCoolingDisabled(auth, e.cfg) {
markAntigravityCreditsPermanentlyDisabled(auth)
}
// No credits logic - just fall through to error return below

View File

@@ -36,6 +36,16 @@ func SetTransientErrorCooldownSeconds(seconds int) {
transientErrorCooldownSeconds.Store(int64(seconds))
}
// QuotaCooldownDisabledForAuth returns whether cooling is disabled for the auth under global settings.
func QuotaCooldownDisabledForAuth(auth *Auth) bool {
return quotaCooldownDisabledForAuth(auth)
}
// QuotaCooldownDisabledForAuthWithConfig returns whether cooling is disabled for the auth with the given config.
func QuotaCooldownDisabledForAuthWithConfig(auth *Auth, cfg *internalconfig.Config) bool {
return quotaCooldownDisabledForAuthWithConfig(auth, cfg)
}
func quotaCooldownDisabledForAuth(auth *Auth) bool {
return quotaCooldownDisabledForAuthWithConfig(auth, nil)
}