From a0bf33eca674cf3bff005a107d504359b23adea9 Mon Sep 17 00:00:00 2001 From: xixiwenxuanhe Date: Tue, 31 Mar 2026 00:14:05 +0800 Subject: [PATCH] fix(antigravity): preserve fallback and honor config gate --- .../runtime/executor/antigravity_executor.go | 24 +--- .../antigravity_executor_credits_test.go | 132 ++++++++++++++++++ 2 files changed, 139 insertions(+), 17 deletions(-) diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index 76ce9586..6ee972a7 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -493,7 +493,7 @@ attemptLoop: for idx, baseURL := range baseURLs { requestPayload := translated usedCreditsDirect := false - if antigravityShouldPreferCredits(auth, baseModel, time.Now()) { + if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) { if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 { requestPayload = creditsPayload usedCreditsDirect = true @@ -541,7 +541,7 @@ attemptLoop: markAntigravityCreditsExhausted(auth, time.Now()) } } else { - creditsResp, attemptedCredits := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, false, opts.Alt, baseURL, bodyBytes) + 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) @@ -561,10 +561,6 @@ attemptLoop: reporter.ensurePublished(ctx) return resp, nil } - if attemptedCredits { - err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) - return resp, err - } } } @@ -665,7 +661,7 @@ attemptLoop: for idx, baseURL := range baseURLs { requestPayload := translated usedCreditsDirect := false - if antigravityShouldPreferCredits(auth, baseModel, time.Now()) { + if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) { if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 { requestPayload = creditsPayload usedCreditsDirect = true @@ -727,13 +723,10 @@ attemptLoop: markAntigravityCreditsExhausted(auth, time.Now()) } } else { - creditsResp, attemptedCredits := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, true, opts.Alt, baseURL, bodyBytes) + 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()) - } else if attemptedCredits { - err = newAntigravityStatusErr(http.StatusTooManyRequests, bodyBytes) - return resp, err } } } @@ -1085,7 +1078,7 @@ attemptLoop: for idx, baseURL := range baseURLs { requestPayload := translated usedCreditsDirect := false - if antigravityShouldPreferCredits(auth, baseModel, time.Now()) { + if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) { if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 { requestPayload = creditsPayload usedCreditsDirect = true @@ -1146,13 +1139,10 @@ attemptLoop: markAntigravityCreditsExhausted(auth, time.Now()) } } else { - creditsResp, attemptedCredits := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, true, opts.Alt, baseURL, bodyBytes) + 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()) - } else if attemptedCredits { - err = newAntigravityStatusErr(http.StatusTooManyRequests, bodyBytes) - return nil, err } } } @@ -1797,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 index ecac0c83..13ab662b 100644 --- a/internal/runtime/executor/antigravity_executor_credits_test.go +++ b/internal/runtime/executor/antigravity_executor_credits_test.go @@ -289,3 +289,135 @@ func TestAntigravityExecute_PrefersCreditsAfterSuccessfulFallback(t *testing.T) 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]) + } +}