diff --git a/config.example.yaml b/config.example.yaml index 1b365d87..a394f979 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -91,6 +91,7 @@ max-retry-interval: 30 quota-exceeded: switch-project: true # Whether to automatically switch to another project when a quota is exceeded switch-preview-model: true # Whether to automatically switch to a preview model when a quota is exceeded + antigravity-credits: true # Whether to retry Antigravity quota_exhausted 429s once with enabledCreditTypes=["GOOGLE_ONE_AI"] # Routing strategy for selecting credentials when multiple match. routing: diff --git a/internal/config/config.go b/internal/config/config.go index c4156e97..ceb2e7bd 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -194,6 +194,10 @@ type QuotaExceeded struct { // SwitchPreviewModel indicates whether to automatically switch to a preview model when a quota is exceeded. SwitchPreviewModel bool `yaml:"switch-preview-model" json:"switch-preview-model"` + + // AntigravityCredits indicates whether to retry Antigravity quota_exhausted 429s once + // on the same credential with enabledCreditTypes=["GOOGLE_ONE_AI"]. + AntigravityCredits bool `yaml:"antigravity-credits" json:"antigravity-credits"` } // RoutingConfig configures how credentials are selected for requests. diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index 18079a43..6ee972a7 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -47,12 +47,41 @@ const ( defaultAntigravityAgent = "antigravity/1.19.6 darwin/arm64" antigravityAuthType = "antigravity" refreshSkew = 3000 * time.Second + antigravityCreditsRetryTTL = 5 * time.Hour // systemInstruction = "You are Antigravity, a powerful agentic AI coding assistant designed by the Google Deepmind team working on Advanced Agentic Coding.You are pair programming with a USER to solve their coding task. The task may require creating a new codebase, modifying or debugging an existing codebase, or simply answering a question.**Absolute paths only****Proactiveness**" ) +type antigravity429Category string + +const ( + antigravity429Unknown antigravity429Category = "unknown" + antigravity429RateLimited antigravity429Category = "rate_limited" + antigravity429QuotaExhausted antigravity429Category = "quota_exhausted" +) + var ( - randSource = rand.New(rand.NewSource(time.Now().UnixNano())) - randSourceMutex sync.Mutex + randSource = rand.New(rand.NewSource(time.Now().UnixNano())) + randSourceMutex sync.Mutex + antigravityCreditsExhaustedByAuth sync.Map + antigravityPreferCreditsByModel sync.Map + antigravityQuotaExhaustedKeywords = []string{ + "quota_exhausted", + "quota exhausted", + } + antigravityCreditsExhaustedKeywords = []string{ + "google_one_ai", + "insufficient credit", + "insufficient credits", + "not enough credit", + "not enough credits", + "credit exhausted", + "credits exhausted", + "credit balance", + "minimumcreditamountforusage", + "minimum credit amount for usage", + "minimum credit", + "resource has been exhausted", + } ) // AntigravityExecutor proxies requests to the antigravity upstream. @@ -183,6 +212,231 @@ func (e *AntigravityExecutor) HttpRequest(ctx context.Context, auth *cliproxyaut return httpClient.Do(httpReq) } +func injectEnabledCreditTypes(payload []byte) []byte { + if len(payload) == 0 { + return nil + } + if !gjson.ValidBytes(payload) { + return nil + } + updated, err := sjson.SetRawBytes(payload, "enabledCreditTypes", []byte(`["GOOGLE_ONE_AI"]`)) + if err != nil { + return nil + } + return updated +} + +func classifyAntigravity429(body []byte) antigravity429Category { + if len(body) == 0 { + return antigravity429Unknown + } + lowerBody := strings.ToLower(string(body)) + for _, keyword := range antigravityQuotaExhaustedKeywords { + if strings.Contains(lowerBody, keyword) { + return antigravity429QuotaExhausted + } + } + status := strings.TrimSpace(gjson.GetBytes(body, "error.status").String()) + if !strings.EqualFold(status, "RESOURCE_EXHAUSTED") { + return antigravity429Unknown + } + details := gjson.GetBytes(body, "error.details") + if !details.Exists() || !details.IsArray() { + return antigravity429Unknown + } + for _, detail := range details.Array() { + if detail.Get("@type").String() != "type.googleapis.com/google.rpc.ErrorInfo" { + continue + } + reason := strings.TrimSpace(detail.Get("reason").String()) + if strings.EqualFold(reason, "QUOTA_EXHAUSTED") { + return antigravity429QuotaExhausted + } + if strings.EqualFold(reason, "RATE_LIMIT_EXCEEDED") { + return antigravity429RateLimited + } + } + return antigravity429Unknown +} + +func antigravityCreditsRetryEnabled(cfg *config.Config) bool { + return cfg != nil && cfg.QuotaExceeded.AntigravityCredits +} + +func antigravityCreditsExhausted(auth *cliproxyauth.Auth, now time.Time) bool { + if auth == nil || strings.TrimSpace(auth.ID) == "" { + return false + } + value, ok := antigravityCreditsExhaustedByAuth.Load(auth.ID) + if !ok { + return false + } + until, ok := value.(time.Time) + if !ok || until.IsZero() { + antigravityCreditsExhaustedByAuth.Delete(auth.ID) + return false + } + if !until.After(now) { + antigravityCreditsExhaustedByAuth.Delete(auth.ID) + return false + } + return true +} + +func markAntigravityCreditsExhausted(auth *cliproxyauth.Auth, now time.Time) { + if auth == nil || strings.TrimSpace(auth.ID) == "" { + return + } + antigravityCreditsExhaustedByAuth.Store(auth.ID, now.Add(antigravityCreditsRetryTTL)) +} + +func clearAntigravityCreditsExhausted(auth *cliproxyauth.Auth) { + if auth == nil || strings.TrimSpace(auth.ID) == "" { + return + } + antigravityCreditsExhaustedByAuth.Delete(auth.ID) +} + +func antigravityPreferCreditsKey(auth *cliproxyauth.Auth, modelName string) string { + if auth == nil { + return "" + } + authID := strings.TrimSpace(auth.ID) + modelName = strings.TrimSpace(modelName) + if authID == "" || modelName == "" { + return "" + } + return authID + "|" + modelName +} + +func antigravityShouldPreferCredits(auth *cliproxyauth.Auth, modelName string, now time.Time) bool { + key := antigravityPreferCreditsKey(auth, modelName) + if key == "" { + return false + } + value, ok := antigravityPreferCreditsByModel.Load(key) + if !ok { + return false + } + until, ok := value.(time.Time) + if !ok || until.IsZero() { + antigravityPreferCreditsByModel.Delete(key) + return false + } + if !until.After(now) { + antigravityPreferCreditsByModel.Delete(key) + return false + } + return true +} + +func markAntigravityPreferCredits(auth *cliproxyauth.Auth, modelName string, now time.Time, retryAfter *time.Duration) { + key := antigravityPreferCreditsKey(auth, modelName) + if key == "" { + return + } + until := now.Add(antigravityCreditsRetryTTL) + if retryAfter != nil && *retryAfter > 0 { + until = now.Add(*retryAfter) + } + antigravityPreferCreditsByModel.Store(key, until) +} + +func clearAntigravityPreferCredits(auth *cliproxyauth.Auth, modelName string) { + key := antigravityPreferCreditsKey(auth, modelName) + if key == "" { + return + } + antigravityPreferCreditsByModel.Delete(key) +} + +func shouldMarkAntigravityCreditsExhausted(statusCode int, body []byte, reqErr error) bool { + if reqErr != nil || statusCode == 0 { + return false + } + if statusCode >= http.StatusInternalServerError || statusCode == http.StatusRequestTimeout { + return false + } + lowerBody := strings.ToLower(string(body)) + for _, keyword := range antigravityCreditsExhaustedKeywords { + if strings.Contains(lowerBody, keyword) { + return true + } + } + return false +} + +func newAntigravityStatusErr(statusCode int, body []byte) statusErr { + err := statusErr{code: statusCode, msg: string(body)} + if statusCode == http.StatusTooManyRequests { + if retryAfter, parseErr := parseRetryDelay(body); parseErr == nil && retryAfter != nil { + err.retryAfter = retryAfter + } + } + return err +} + +func (e *AntigravityExecutor) attemptCreditsFallback( + ctx context.Context, + auth *cliproxyauth.Auth, + httpClient *http.Client, + token string, + modelName string, + payload []byte, + stream bool, + alt string, + baseURL string, + originalBody []byte, +) (*http.Response, bool) { + if !antigravityCreditsRetryEnabled(e.cfg) { + return nil, false + } + if classifyAntigravity429(originalBody) != antigravity429QuotaExhausted { + return nil, false + } + now := time.Now() + if antigravityCreditsExhausted(auth, now) { + return nil, false + } + creditsPayload := injectEnabledCreditTypes(payload) + if len(creditsPayload) == 0 { + return nil, false + } + + httpReq, errReq := e.buildRequest(ctx, auth, token, modelName, creditsPayload, stream, alt, baseURL) + if errReq != nil { + recordAPIResponseError(ctx, e.cfg, errReq) + return nil, true + } + httpResp, errDo := httpClient.Do(httpReq) + if errDo != nil { + recordAPIResponseError(ctx, e.cfg, errDo) + return nil, true + } + if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices { + retryAfter, _ := parseRetryDelay(originalBody) + markAntigravityPreferCredits(auth, modelName, now, retryAfter) + clearAntigravityCreditsExhausted(auth) + return httpResp, true + } + + recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + bodyBytes, errRead := io.ReadAll(httpResp.Body) + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close credits fallback response body error: %v", errClose) + } + if errRead != nil { + recordAPIResponseError(ctx, e.cfg, errRead) + return nil, true + } + appendAPIResponseChunk(ctx, e.cfg, bodyBytes) + if shouldMarkAntigravityCreditsExhausted(httpResp.StatusCode, bodyBytes, nil) { + clearAntigravityPreferCredits(auth, modelName) + markAntigravityCreditsExhausted(auth, now) + } + return nil, true +} + // Execute performs a non-streaming request to the Antigravity API. func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { if opts.Alt == "responses/compact" { @@ -237,7 +491,15 @@ attemptLoop: var lastErr error for idx, baseURL := range baseURLs { - httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, translated, false, opts.Alt, baseURL) + requestPayload := translated + usedCreditsDirect := false + if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) { + if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 { + requestPayload = creditsPayload + usedCreditsDirect = true + } + } + httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, false, opts.Alt, baseURL) if errReq != nil { err = errReq return resp, err @@ -272,6 +534,36 @@ attemptLoop: } appendAPIResponseChunk(ctx, e.cfg, bodyBytes) + if httpResp.StatusCode == http.StatusTooManyRequests { + if usedCreditsDirect { + if shouldMarkAntigravityCreditsExhausted(httpResp.StatusCode, bodyBytes, nil) { + clearAntigravityPreferCredits(auth, baseModel) + markAntigravityCreditsExhausted(auth, time.Now()) + } + } else { + creditsResp, _ := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, false, opts.Alt, baseURL, bodyBytes) + if creditsResp != nil { + recordAPIResponseMetadata(ctx, e.cfg, creditsResp.StatusCode, creditsResp.Header.Clone()) + creditsBody, errCreditsRead := io.ReadAll(creditsResp.Body) + if errClose := creditsResp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close credits success response body error: %v", errClose) + } + if errCreditsRead != nil { + recordAPIResponseError(ctx, e.cfg, errCreditsRead) + err = errCreditsRead + return resp, err + } + appendAPIResponseChunk(ctx, e.cfg, creditsBody) + reporter.publish(ctx, parseAntigravityUsage(creditsBody)) + var param any + converted := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, creditsBody, ¶m) + resp = cliproxyexecutor.Response{Payload: converted, Headers: creditsResp.Header.Clone()} + reporter.ensurePublished(ctx) + return resp, nil + } + } + } + if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { log.Debugf("antigravity executor: upstream error status: %d, body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), bodyBytes)) lastStatus = httpResp.StatusCode @@ -295,13 +587,7 @@ attemptLoop: continue attemptLoop } } - sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)} - if httpResp.StatusCode == http.StatusTooManyRequests { - if retryAfter, parseErr := parseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil { - sErr.retryAfter = retryAfter - } - } - err = sErr + err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) return resp, err } @@ -315,13 +601,7 @@ attemptLoop: switch { case lastStatus != 0: - sErr := statusErr{code: lastStatus, msg: string(lastBody)} - if lastStatus == http.StatusTooManyRequests { - if retryAfter, parseErr := parseRetryDelay(lastBody); parseErr == nil && retryAfter != nil { - sErr.retryAfter = retryAfter - } - } - err = sErr + err = newAntigravityStatusErr(lastStatus, lastBody) case lastErr != nil: err = lastErr default: @@ -379,7 +659,15 @@ attemptLoop: var lastErr error for idx, baseURL := range baseURLs { - httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, translated, true, opts.Alt, baseURL) + requestPayload := translated + usedCreditsDirect := false + if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) { + if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 { + requestPayload = creditsPayload + usedCreditsDirect = true + } + } + httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL) if errReq != nil { err = errReq return resp, err @@ -428,6 +716,23 @@ attemptLoop: return resp, err } appendAPIResponseChunk(ctx, e.cfg, bodyBytes) + if httpResp.StatusCode == http.StatusTooManyRequests { + if usedCreditsDirect { + if shouldMarkAntigravityCreditsExhausted(httpResp.StatusCode, bodyBytes, nil) { + clearAntigravityPreferCredits(auth, baseModel) + markAntigravityCreditsExhausted(auth, time.Now()) + } + } else { + creditsResp, _ := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, true, opts.Alt, baseURL, bodyBytes) + if creditsResp != nil { + httpResp = creditsResp + recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + } + } + } + if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices { + goto streamSuccessClaudeNonStream + } lastStatus = httpResp.StatusCode lastBody = append([]byte(nil), bodyBytes...) lastErr = nil @@ -449,16 +754,11 @@ attemptLoop: continue attemptLoop } } - sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)} - if httpResp.StatusCode == http.StatusTooManyRequests { - if retryAfter, parseErr := parseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil { - sErr.retryAfter = retryAfter - } - } - err = sErr + err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) return resp, err } + streamSuccessClaudeNonStream: out := make(chan cliproxyexecutor.StreamChunk) go func(resp *http.Response) { defer close(out) @@ -520,13 +820,7 @@ attemptLoop: switch { case lastStatus != 0: - sErr := statusErr{code: lastStatus, msg: string(lastBody)} - if lastStatus == http.StatusTooManyRequests { - if retryAfter, parseErr := parseRetryDelay(lastBody); parseErr == nil && retryAfter != nil { - sErr.retryAfter = retryAfter - } - } - err = sErr + err = newAntigravityStatusErr(lastStatus, lastBody) case lastErr != nil: err = lastErr default: @@ -782,7 +1076,15 @@ attemptLoop: var lastErr error for idx, baseURL := range baseURLs { - httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, translated, true, opts.Alt, baseURL) + requestPayload := translated + usedCreditsDirect := false + if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) { + if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 { + requestPayload = creditsPayload + usedCreditsDirect = true + } + } + httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL) if errReq != nil { err = errReq return nil, err @@ -830,6 +1132,23 @@ attemptLoop: return nil, err } appendAPIResponseChunk(ctx, e.cfg, bodyBytes) + if httpResp.StatusCode == http.StatusTooManyRequests { + if usedCreditsDirect { + if shouldMarkAntigravityCreditsExhausted(httpResp.StatusCode, bodyBytes, nil) { + clearAntigravityPreferCredits(auth, baseModel) + markAntigravityCreditsExhausted(auth, time.Now()) + } + } else { + creditsResp, _ := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, true, opts.Alt, baseURL, bodyBytes) + if creditsResp != nil { + httpResp = creditsResp + recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + } + } + } + if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices { + goto streamSuccessExecuteStream + } lastStatus = httpResp.StatusCode lastBody = append([]byte(nil), bodyBytes...) lastErr = nil @@ -851,16 +1170,11 @@ attemptLoop: continue attemptLoop } } - sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)} - if httpResp.StatusCode == http.StatusTooManyRequests { - if retryAfter, parseErr := parseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil { - sErr.retryAfter = retryAfter - } - } - err = sErr + err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) return nil, err } + streamSuccessExecuteStream: out := make(chan cliproxyexecutor.StreamChunk) go func(resp *http.Response) { defer close(out) @@ -911,13 +1225,7 @@ attemptLoop: switch { case lastStatus != 0: - sErr := statusErr{code: lastStatus, msg: string(lastBody)} - if lastStatus == http.StatusTooManyRequests { - if retryAfter, parseErr := parseRetryDelay(lastBody); parseErr == nil && retryAfter != nil { - sErr.retryAfter = retryAfter - } - } - err = sErr + err = newAntigravityStatusErr(lastStatus, lastBody) case lastErr != nil: err = lastErr default: @@ -1479,7 +1787,7 @@ func antigravityWait(ctx context.Context, wait time.Duration) error { } } -func antigravityBaseURLFallbackOrder(auth *cliproxyauth.Auth) []string { +var antigravityBaseURLFallbackOrder = func(auth *cliproxyauth.Auth) []string { if base := resolveCustomAntigravityBaseURL(auth); base != "" { return []string{base} } diff --git a/internal/runtime/executor/antigravity_executor_credits_test.go b/internal/runtime/executor/antigravity_executor_credits_test.go new file mode 100644 index 00000000..13ab662b --- /dev/null +++ b/internal/runtime/executor/antigravity_executor_credits_test.go @@ -0,0 +1,423 @@ +package executor + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v6/internal/config" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator" +) + +func resetAntigravityCreditsRetryState() { + antigravityCreditsExhaustedByAuth = sync.Map{} + antigravityPreferCreditsByModel = sync.Map{} +} + +func TestClassifyAntigravity429(t *testing.T) { + t.Run("quota exhausted", func(t *testing.T) { + body := []byte(`{"error":{"status":"RESOURCE_EXHAUSTED","message":"QUOTA_EXHAUSTED"}}`) + if got := classifyAntigravity429(body); got != antigravity429QuotaExhausted { + t.Fatalf("classifyAntigravity429() = %q, want %q", got, antigravity429QuotaExhausted) + } + }) + + t.Run("structured rate limit", func(t *testing.T) { + body := []byte(`{ + "error": { + "status": "RESOURCE_EXHAUSTED", + "details": [ + {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "reason": "RATE_LIMIT_EXCEEDED"}, + {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.5s"} + ] + } + }`) + if got := classifyAntigravity429(body); got != antigravity429RateLimited { + t.Fatalf("classifyAntigravity429() = %q, want %q", got, antigravity429RateLimited) + } + }) + + t.Run("structured quota exhausted", func(t *testing.T) { + body := []byte(`{ + "error": { + "status": "RESOURCE_EXHAUSTED", + "details": [ + {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "reason": "QUOTA_EXHAUSTED"} + ] + } + }`) + if got := classifyAntigravity429(body); got != antigravity429QuotaExhausted { + t.Fatalf("classifyAntigravity429() = %q, want %q", got, antigravity429QuotaExhausted) + } + }) + + t.Run("unknown", func(t *testing.T) { + body := []byte(`{"error":{"message":"too many requests"}}`) + if got := classifyAntigravity429(body); got != antigravity429Unknown { + t.Fatalf("classifyAntigravity429() = %q, want %q", got, antigravity429Unknown) + } + }) +} + +func TestInjectEnabledCreditTypes(t *testing.T) { + body := []byte(`{"model":"gemini-2.5-flash","request":{}}`) + got := injectEnabledCreditTypes(body) + if got == nil { + t.Fatal("injectEnabledCreditTypes() returned nil") + } + if !strings.Contains(string(got), `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) { + t.Fatalf("injectEnabledCreditTypes() = %s, want enabledCreditTypes", string(got)) + } + + if got := injectEnabledCreditTypes([]byte(`not json`)); got != nil { + t.Fatalf("injectEnabledCreditTypes() for invalid json = %s, want nil", string(got)) + } +} + +func TestShouldMarkAntigravityCreditsExhausted(t *testing.T) { + for _, body := range [][]byte{ + []byte(`{"error":{"message":"Insufficient GOOGLE_ONE_AI credits"}}`), + []byte(`{"error":{"message":"minimumCreditAmountForUsage requirement not met"}}`), + []byte(`{"error":{"message":"Resource has been exhausted"}}`), + } { + if !shouldMarkAntigravityCreditsExhausted(http.StatusForbidden, body, nil) { + t.Fatalf("shouldMarkAntigravityCreditsExhausted(%s) = false, want true", string(body)) + } + } + if shouldMarkAntigravityCreditsExhausted(http.StatusServiceUnavailable, []byte(`{"error":{"message":"credits exhausted"}}`), nil) { + t.Fatal("shouldMarkAntigravityCreditsExhausted() = true for 5xx, want false") + } +} + +func TestAntigravityExecute_RetriesQuotaExhaustedWithCredits(t *testing.T) { + resetAntigravityCreditsRetryState() + t.Cleanup(resetAntigravityCreditsRetryState) + + var ( + mu sync.Mutex + requestBodies []string + ) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + _ = r.Body.Close() + + mu.Lock() + requestBodies = append(requestBodies, string(body)) + reqNum := len(requestBodies) + mu.Unlock() + + if reqNum == 1 { + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":{"status":"RESOURCE_EXHAUSTED","message":"QUOTA_EXHAUSTED"}}`)) + return + } + + if !strings.Contains(string(body), `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) { + t.Fatalf("second request body missing enabledCreditTypes: %s", string(body)) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]}}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}}`)) + })) + defer server.Close() + + exec := NewAntigravityExecutor(&config.Config{ + QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true}, + }) + auth := &cliproxyauth.Auth{ + ID: "auth-credits-ok", + Attributes: map[string]string{ + "base_url": server.URL, + }, + Metadata: map[string]any{ + "access_token": "token", + "project_id": "project-1", + "expired": time.Now().Add(1 * time.Hour).Format(time.RFC3339), + }, + } + + resp, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gemini-2.5-flash", + Payload: []byte(`{"request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatAntigravity, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + if len(resp.Payload) == 0 { + t.Fatal("Execute() returned empty payload") + } + + mu.Lock() + defer mu.Unlock() + if len(requestBodies) != 2 { + t.Fatalf("request count = %d, want 2", len(requestBodies)) + } +} + +func TestAntigravityExecute_SkipsCreditsRetryWhenAlreadyExhausted(t *testing.T) { + resetAntigravityCreditsRetryState() + t.Cleanup(resetAntigravityCreditsRetryState) + + var requestCount int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount++ + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":{"status":"RESOURCE_EXHAUSTED","message":"QUOTA_EXHAUSTED"}}`)) + })) + defer server.Close() + + exec := NewAntigravityExecutor(&config.Config{ + QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true}, + }) + auth := &cliproxyauth.Auth{ + ID: "auth-credits-exhausted", + Attributes: map[string]string{ + "base_url": server.URL, + }, + Metadata: map[string]any{ + "access_token": "token", + "project_id": "project-1", + "expired": time.Now().Add(1 * time.Hour).Format(time.RFC3339), + }, + } + markAntigravityCreditsExhausted(auth, time.Now()) + + _, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gemini-2.5-flash", + Payload: []byte(`{"request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatAntigravity, + }) + if err == nil { + t.Fatal("Execute() error = nil, want 429") + } + sErr, ok := err.(statusErr) + if !ok { + t.Fatalf("Execute() error type = %T, want statusErr", err) + } + if got := sErr.StatusCode(); got != http.StatusTooManyRequests { + t.Fatalf("Execute() status code = %d, want %d", got, http.StatusTooManyRequests) + } + if requestCount != 1 { + t.Fatalf("request count = %d, want 1", requestCount) + } +} + +func TestAntigravityExecute_PrefersCreditsAfterSuccessfulFallback(t *testing.T) { + resetAntigravityCreditsRetryState() + t.Cleanup(resetAntigravityCreditsRetryState) + + var ( + mu sync.Mutex + requestBodies []string + ) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + _ = r.Body.Close() + + mu.Lock() + requestBodies = append(requestBodies, string(body)) + reqNum := len(requestBodies) + mu.Unlock() + + switch reqNum { + case 1: + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":{"status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"QUOTA_EXHAUSTED"},{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"10s"}]}}`)) + case 2, 3: + if !strings.Contains(string(body), `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) { + t.Fatalf("request %d body missing enabledCreditTypes: %s", reqNum, string(body)) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"OK"}]}}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}}`)) + default: + t.Fatalf("unexpected request count %d", reqNum) + } + })) + defer server.Close() + + exec := NewAntigravityExecutor(&config.Config{ + QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true}, + }) + auth := &cliproxyauth.Auth{ + ID: "auth-prefer-credits", + Attributes: map[string]string{ + "base_url": server.URL, + }, + Metadata: map[string]any{ + "access_token": "token", + "project_id": "project-1", + "expired": time.Now().Add(1 * time.Hour).Format(time.RFC3339), + }, + } + + request := cliproxyexecutor.Request{ + Model: "gemini-2.5-flash", + Payload: []byte(`{"request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`), + } + opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatAntigravity} + + if _, err := exec.Execute(context.Background(), auth, request, opts); err != nil { + t.Fatalf("first Execute() error = %v", err) + } + if _, err := exec.Execute(context.Background(), auth, request, opts); err != nil { + t.Fatalf("second Execute() error = %v", err) + } + + mu.Lock() + defer mu.Unlock() + if len(requestBodies) != 3 { + t.Fatalf("request count = %d, want 3", len(requestBodies)) + } + if strings.Contains(requestBodies[0], `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) { + t.Fatalf("first request unexpectedly used credits: %s", requestBodies[0]) + } + if !strings.Contains(requestBodies[1], `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) { + t.Fatalf("fallback request missing credits: %s", requestBodies[1]) + } + if !strings.Contains(requestBodies[2], `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) { + t.Fatalf("preferred request missing credits: %s", requestBodies[2]) + } +} + +func TestAntigravityExecute_PreservesBaseURLFallbackAfterCreditsRetryFailure(t *testing.T) { + resetAntigravityCreditsRetryState() + t.Cleanup(resetAntigravityCreditsRetryState) + + var ( + mu sync.Mutex + firstCount int + secondCount int + ) + + firstServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + _ = r.Body.Close() + + mu.Lock() + firstCount++ + reqNum := firstCount + mu.Unlock() + + switch reqNum { + case 1: + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":{"status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"QUOTA_EXHAUSTED"}]}}`)) + case 2: + if !strings.Contains(string(body), `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) { + t.Fatalf("credits retry missing enabledCreditTypes: %s", string(body)) + } + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error":{"message":"permission denied"}}`)) + default: + t.Fatalf("unexpected first server request count %d", reqNum) + } + })) + defer firstServer.Close() + + secondServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + secondCount++ + mu.Unlock() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]}}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}}`)) + })) + defer secondServer.Close() + + exec := NewAntigravityExecutor(&config.Config{ + QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true}, + }) + auth := &cliproxyauth.Auth{ + ID: "auth-baseurl-fallback", + Attributes: map[string]string{ + "base_url": firstServer.URL, + }, + Metadata: map[string]any{ + "access_token": "token", + "project_id": "project-1", + "expired": time.Now().Add(1 * time.Hour).Format(time.RFC3339), + }, + } + + originalOrder := antigravityBaseURLFallbackOrder + defer func() { antigravityBaseURLFallbackOrder = originalOrder }() + antigravityBaseURLFallbackOrder = func(auth *cliproxyauth.Auth) []string { + return []string{firstServer.URL, secondServer.URL} + } + + resp, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gemini-2.5-flash", + Payload: []byte(`{"request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatAntigravity, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + if len(resp.Payload) == 0 { + t.Fatal("Execute() returned empty payload") + } + if firstCount != 2 { + t.Fatalf("first server request count = %d, want 2", firstCount) + } + if secondCount != 1 { + t.Fatalf("second server request count = %d, want 1", secondCount) + } +} + +func TestAntigravityExecute_DoesNotDirectInjectCreditsWhenFlagDisabled(t *testing.T) { + resetAntigravityCreditsRetryState() + t.Cleanup(resetAntigravityCreditsRetryState) + + var requestBodies []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + _ = r.Body.Close() + requestBodies = append(requestBodies, string(body)) + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":{"status":"RESOURCE_EXHAUSTED","message":"QUOTA_EXHAUSTED"}}`)) + })) + defer server.Close() + + exec := NewAntigravityExecutor(&config.Config{ + QuotaExceeded: config.QuotaExceeded{AntigravityCredits: false}, + }) + auth := &cliproxyauth.Auth{ + ID: "auth-flag-disabled", + Attributes: map[string]string{ + "base_url": server.URL, + }, + Metadata: map[string]any{ + "access_token": "token", + "project_id": "project-1", + "expired": time.Now().Add(1 * time.Hour).Format(time.RFC3339), + }, + } + markAntigravityPreferCredits(auth, "gemini-2.5-flash", time.Now(), nil) + + _, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gemini-2.5-flash", + Payload: []byte(`{"request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatAntigravity, + }) + if err == nil { + t.Fatal("Execute() error = nil, want 429") + } + if len(requestBodies) != 1 { + t.Fatalf("request count = %d, want 1", len(requestBodies)) + } + if strings.Contains(requestBodies[0], `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) { + t.Fatalf("request unexpectedly used enabledCreditTypes with flag disabled: %s", requestBodies[0]) + } +} diff --git a/internal/watcher/diff/config_diff.go b/internal/watcher/diff/config_diff.go index fccdaf8d..11f9093e 100644 --- a/internal/watcher/diff/config_diff.go +++ b/internal/watcher/diff/config_diff.go @@ -80,6 +80,9 @@ func BuildConfigChangeDetails(oldCfg, newCfg *config.Config) []string { if oldCfg.QuotaExceeded.SwitchPreviewModel != newCfg.QuotaExceeded.SwitchPreviewModel { changes = append(changes, fmt.Sprintf("quota-exceeded.switch-preview-model: %t -> %t", oldCfg.QuotaExceeded.SwitchPreviewModel, newCfg.QuotaExceeded.SwitchPreviewModel)) } + if oldCfg.QuotaExceeded.AntigravityCredits != newCfg.QuotaExceeded.AntigravityCredits { + changes = append(changes, fmt.Sprintf("quota-exceeded.antigravity-credits: %t -> %t", oldCfg.QuotaExceeded.AntigravityCredits, newCfg.QuotaExceeded.AntigravityCredits)) + } if oldCfg.Routing.Strategy != newCfg.Routing.Strategy { changes = append(changes, fmt.Sprintf("routing.strategy: %s -> %s", oldCfg.Routing.Strategy, newCfg.Routing.Strategy)) diff --git a/internal/watcher/diff/config_diff_test.go b/internal/watcher/diff/config_diff_test.go index c7b73f11..2d45aa57 100644 --- a/internal/watcher/diff/config_diff_test.go +++ b/internal/watcher/diff/config_diff_test.go @@ -229,7 +229,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) { MaxRetryCredentials: 1, MaxRetryInterval: 1, WebsocketAuth: false, - QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false}, + QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false, AntigravityCredits: false}, ClaudeKey: []config.ClaudeKey{{APIKey: "c1"}}, CodexKey: []config.CodexKey{{APIKey: "x1"}}, AmpCode: config.AmpCode{UpstreamAPIKey: "keep", RestrictManagementToLocalhost: false}, @@ -253,7 +253,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) { MaxRetryCredentials: 3, MaxRetryInterval: 3, WebsocketAuth: true, - QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true}, + QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true, AntigravityCredits: true}, ClaudeKey: []config.ClaudeKey{ {APIKey: "c1", BaseURL: "http://new", ProxyURL: "http://p", Headers: map[string]string{"H": "1"}, ExcludedModels: []string{"a"}}, {APIKey: "c2"}, @@ -297,6 +297,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) { expectContains(t, details, "nonstream-keepalive-interval: 0 -> 5") expectContains(t, details, "quota-exceeded.switch-project: false -> true") expectContains(t, details, "quota-exceeded.switch-preview-model: false -> true") + expectContains(t, details, "quota-exceeded.antigravity-credits: false -> true") expectContains(t, details, "api-keys count: 1 -> 2") expectContains(t, details, "claude-api-key count: 1 -> 2") expectContains(t, details, "codex-api-key count: 1 -> 2") @@ -320,7 +321,7 @@ func TestBuildConfigChangeDetails_AllBranches(t *testing.T) { MaxRetryCredentials: 1, MaxRetryInterval: 1, WebsocketAuth: false, - QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false}, + QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false, AntigravityCredits: false}, GeminiKey: []config.GeminiKey{ {APIKey: "g-old", BaseURL: "http://g-old", ProxyURL: "http://gp-old", Headers: map[string]string{"A": "1"}}, }, @@ -374,7 +375,7 @@ func TestBuildConfigChangeDetails_AllBranches(t *testing.T) { MaxRetryCredentials: 3, MaxRetryInterval: 3, WebsocketAuth: true, - QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true}, + QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true, AntigravityCredits: true}, GeminiKey: []config.GeminiKey{ {APIKey: "g-new", BaseURL: "http://g-new", ProxyURL: "http://gp-new", Headers: map[string]string{"A": "2"}, ExcludedModels: []string{"x", "y"}}, }, @@ -437,6 +438,7 @@ func TestBuildConfigChangeDetails_AllBranches(t *testing.T) { expectContains(t, changes, "ws-auth: false -> true") expectContains(t, changes, "quota-exceeded.switch-project: false -> true") expectContains(t, changes, "quota-exceeded.switch-preview-model: false -> true") + expectContains(t, changes, "quota-exceeded.antigravity-credits: false -> true") expectContains(t, changes, "api-keys: values updated (count unchanged, redacted)") expectContains(t, changes, "gemini[0].base-url: http://g-old -> http://g-new") expectContains(t, changes, "gemini[0].proxy-url: http://gp-old -> http://gp-new")