mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-09-03 06:35:00 +08:00
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:
@@ -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 {
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user