fix(antigravity): preserve fallback and honor config gate

2026-04-03 03:06:36 +00:00 · 2026-03-31 00:14:05 +08:00
parent 88dd9c715d
commit a0bf33eca6
2 changed files with 139 additions and 17 deletions
--- 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}
 	}
--- 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])
+	}
+}