From adf052984f8b53bbc6ac8e4338297afee5b337f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ch=C3=A9n=20M=C3=B9?= <10558748+hkfires@users.noreply.github.com> Date: Tue, 25 Aug 2026 11:51:17 +0800 Subject: [PATCH] fix(antigravity): remove cross-endpoint fallback (#5209) (#5228) Fixes #5209 --- .../antigravity_executor_buildrequest_test.go | 28 + .../executor/antigravity_executor_credits.go | 90 --- .../antigravity_executor_credits_test.go | 23 - .../executor/antigravity_executor_execute.go | 643 +++++++----------- .../executor/antigravity_executor_request.go | 22 +- .../antigravity_executor_signature_test.go | 94 +++ .../executor/antigravity_executor_stream.go | 361 ++++------ .../executor/antigravity_executor_tokens.go | 174 ++--- 8 files changed, 569 insertions(+), 866 deletions(-) diff --git a/internal/runtime/executor/antigravity_executor_buildrequest_test.go b/internal/runtime/executor/antigravity_executor_buildrequest_test.go index 66390cba3..485db1978 100644 --- a/internal/runtime/executor/antigravity_executor_buildrequest_test.go +++ b/internal/runtime/executor/antigravity_executor_buildrequest_test.go @@ -13,6 +13,34 @@ import ( sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" ) +func TestResolveAntigravityRequestBaseURL(t *testing.T) { + t.Run("default uses daily endpoint", func(t *testing.T) { + if got := resolveAntigravityRequestBaseURL(&cliproxyauth.Auth{}); got != antigravityBaseURLDaily { + t.Fatalf("base URL = %q, want %q", got, antigravityBaseURLDaily) + } + }) + + t.Run("custom attribute endpoint remains supported", func(t *testing.T) { + auth := &cliproxyauth.Auth{Attributes: map[string]string{"base_url": "https://enterprise.example.com/"}} + if got := resolveAntigravityRequestBaseURL(auth); got != "https://enterprise.example.com" { + t.Fatalf("base URL = %q, want custom endpoint", got) + } + }) + + t.Run("custom auth file endpoint remains supported", func(t *testing.T) { + auth := &cliproxyauth.Auth{Metadata: map[string]any{"base_url": "https://enterprise.example.com/"}} + if got := resolveAntigravityRequestBaseURL(auth); got != "https://enterprise.example.com" { + t.Fatalf("base URL = %q, want custom auth file endpoint", got) + } + }) +} + +func TestAntigravityLoadCodeAssistBaseURLRemainsProdByDefault(t *testing.T) { + if got := antigravityLoadCodeAssistBaseURL(&cliproxyauth.Auth{}); got != antigravityBaseURLProd { + t.Fatalf("loadCodeAssist base URL = %q, want %q", got, antigravityBaseURLProd) + } +} + func TestAntigravityBuildRequest_SanitizesGeminiToolSchema(t *testing.T) { body := buildRequestBodyFromPayload(t, "gemini-2.5-pro") diff --git a/internal/runtime/executor/antigravity_executor_credits.go b/internal/runtime/executor/antigravity_executor_credits.go index bb049f04a..9cf15e7fc 100644 --- a/internal/runtime/executor/antigravity_executor_credits.go +++ b/internal/runtime/executor/antigravity_executor_credits.go @@ -525,57 +525,10 @@ func (e *AntigravityExecutor) updateAntigravityCreditsBalance(ctx context.Contex return } } -func antigravityShouldRetryNoCapacity(statusCode int, body []byte) bool { - if statusCode != http.StatusServiceUnavailable { - return false - } - if len(body) == 0 { - return false - } - msg := strings.ToLower(string(body)) - return strings.Contains(msg, "no capacity available") -} - -func antigravityShouldRetryTransientResourceExhausted429(statusCode int, body []byte) bool { - if statusCode != http.StatusTooManyRequests { - return false - } - if len(body) == 0 { - return false - } - if classifyAntigravity429(body) != antigravity429Unknown { - return false - } - status := strings.TrimSpace(gjson.GetBytes(body, "error.status").String()) - if !strings.EqualFold(status, "RESOURCE_EXHAUSTED") { - return false - } - msg := strings.ToLower(string(body)) - return strings.Contains(msg, "resource has been exhausted") -} - -func antigravityShouldRetrySoftRateLimit(statusCode int, body []byte) bool { - if statusCode != http.StatusTooManyRequests { - return false - } - return decideAntigravity429(body).kind == antigravity429DecisionSoftRetry -} - func antigravityShouldBypassShortCooldown(ctx context.Context, cfg *config.Config) bool { return cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(cfg) } -func antigravitySoftRateLimitDelay(attempt int) time.Duration { - if attempt < 0 { - attempt = 0 - } - base := time.Duration(attempt+1) * 500 * time.Millisecond - if base > 3*time.Second { - base = 3 * time.Second - } - return base -} - func antigravityShortCooldownKey(auth *cliproxyauth.Auth, modelName string) string { if auth == nil { return "" @@ -730,46 +683,3 @@ func homeKVUnavailableStatusErr(cause error) statusErr { } return statusErr{code: http.StatusServiceUnavailable, msg: fmt.Sprintf("home kv store unavailable: %v", cause)} } - -func antigravityNoCapacityRetryDelay(attempt int) time.Duration { - if attempt < 0 { - attempt = 0 - } - delay := time.Duration(attempt+1) * 250 * time.Millisecond - if delay > 2*time.Second { - delay = 2 * time.Second - } - return delay -} - -func antigravityTransient429RetryDelay(attempt int) time.Duration { - if attempt < 0 { - attempt = 0 - } - delay := time.Duration(attempt+1) * 100 * time.Millisecond - if delay > 500*time.Millisecond { - delay = 500 * time.Millisecond - } - return delay -} - -func antigravityInstantRetryDelay(wait time.Duration) time.Duration { - if wait <= 0 { - return 0 - } - return wait + 800*time.Millisecond -} - -func antigravityWait(ctx context.Context, wait time.Duration) error { - if wait <= 0 { - return nil - } - timer := time.NewTimer(wait) - defer timer.Stop() - select { - case <-ctx.Done(): - return ctx.Err() - case <-timer.C: - return nil - } -} diff --git a/internal/runtime/executor/antigravity_executor_credits_test.go b/internal/runtime/executor/antigravity_executor_credits_test.go index 3223c1bb4..bba63a64b 100644 --- a/internal/runtime/executor/antigravity_executor_credits_test.go +++ b/internal/runtime/executor/antigravity_executor_credits_test.go @@ -225,29 +225,6 @@ func TestClassifyAntigravity429(t *testing.T) { }) } -func TestAntigravityShouldRetryNoCapacity_Standard503(t *testing.T) { - body := []byte(`{ - "error": { - "code": 503, - "message": "No capacity available for model gemini-3.1-flash-image on the server", - "status": "UNAVAILABLE", - "details": [ - { - "@type": "type.googleapis.com/google.rpc.ErrorInfo", - "reason": "MODEL_CAPACITY_EXHAUSTED", - "domain": "cloudcode-pa.googleapis.com", - "metadata": { - "model": "gemini-3.1-flash-image" - } - } - ] - } - }`) - if !antigravityShouldRetryNoCapacity(http.StatusServiceUnavailable, body) { - t.Fatal("antigravityShouldRetryNoCapacity() = false, want true") - } -} - func TestInjectEnabledCreditTypes(t *testing.T) { body := []byte(`{"model":"claude-sonnet-4-6","request":{}}`) got := injectEnabledCreditTypes(body) diff --git a/internal/runtime/executor/antigravity_executor_execute.go b/internal/runtime/executor/antigravity_executor_execute.go index bc64a8544..192a701df 100644 --- a/internal/runtime/executor/antigravity_executor_execute.go +++ b/internal/runtime/executor/antigravity_executor_execute.go @@ -82,181 +82,101 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg) - baseURLs := antigravityBaseURLFallbackOrder(auth) + baseURL := resolveAntigravityRequestBaseURL(auth) httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) httpClient = reporter.TrackHTTPClient(httpClient) - // Credential retry rounds are owned by the conductor. Keep one upstream - // attempt per credential so request-retry is not consumed twice. - attempts := 1 - -attemptLoop: - for attempt := 0; attempt < attempts; attempt++ { - var lastStatus int - var lastBody []byte - var lastErr error - - for idx, baseURL := range baseURLs { - requestPayload := translated - if useCredits { - if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { - requestPayload = cp - helps.MarkCreditsUsed(ctx) - } - } - replayScope := antigravityReasoningReplayScope{} - if antigravityUsesReasoningReplayCache(baseModel) { - var errReplay error - requestPayload, replayScope, errReplay = prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, requestPayload) - if errReplay != nil { - err = errReplay - return resp, err - } - } - requestPayload = ensureAntigravityGeminiLeadingUserContent(baseModel, requestPayload) - - httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, false, opts.Alt, baseURL, helps.DerivedAntigravitySessionID(opts.Metadata, req.Metadata)) - if errReq != nil { - err = errReq - return resp, err - } - - httpResp, errDo := httpClient.Do(httpReq) - if errDo != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errDo) - if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { - return resp, errDo - } - lastStatus = 0 - lastBody = nil - lastErr = errDo - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - err = errDo - return resp, err - } - - helps.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 response body error: %v", errClose) - } - if errRead != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errRead) - err = errRead - return resp, err - } - helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) - - if httpResp.StatusCode == http.StatusTooManyRequests { - decision := decideAntigravity429(bodyBytes) - switch decision.kind { - case antigravity429DecisionInstantRetrySameAuth: - if attempt+1 < attempts { - if decision.retryAfter != nil && *decision.retryAfter > 0 { - wait := antigravityInstantRetryDelay(*decision.retryAfter) - log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait) - if errWait := antigravityWait(ctx, wait); errWait != nil { - return resp, errWait - } - } - continue attemptLoop - } - case antigravity429DecisionShortCooldownSwitchAuth: - if decision.retryAfter != nil && *decision.retryAfter > 0 { - if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { - err = homeKVUnavailableStatusErr(errMarkCooldown) - return resp, err - } - log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel) - } - case antigravity429DecisionFullQuotaExhausted: - if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { - markAntigravityCreditsPermanentlyDisabled(auth) - } - // No credits logic - just fall through to error return below - } - } - - if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { - log.Debugf("antigravity executor: upstream error status: %d, body: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), bodyBytes)) - lastStatus = httpResp.StatusCode - lastBody = append([]byte(nil), bodyBytes...) - lastErr = nil - if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - if antigravityShouldRetryTransientResourceExhausted429(httpResp.StatusCode, bodyBytes) && attempt+1 < attempts { - delay := antigravityTransient429RetryDelay(attempt) - log.Debugf("antigravity executor: transient 429 resource exhausted for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return resp, errWait - } - continue attemptLoop - } - if antigravityShouldRetryNoCapacity(httpResp.StatusCode, bodyBytes) { - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: no capacity on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - if attempt+1 < attempts { - delay := antigravityNoCapacityRetryDelay(attempt) - log.Debugf("antigravity executor: no capacity for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return resp, errWait - } - continue attemptLoop - } - } - if antigravityShouldRetrySoftRateLimit(httpResp.StatusCode, bodyBytes) { - if attempt+1 < attempts { - delay := antigravitySoftRateLimitDelay(attempt) - log.Debugf("antigravity executor: soft rate limit for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return resp, errWait - } - continue attemptLoop - } - } - if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, bodyBytes); errClear != nil { - // Report the upstream failure rather than the cleanup failure. - logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) - } - err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) - return resp, err - } - - // Success - if useCredits { - clearAntigravityCreditsFailureState(auth) - } - cacheAntigravityReasoningReplayFromResponse(ctx, replayScope, requestPayload, bodyBytes) - bodyBytes = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, bodyBytes) - reporter.Publish(ctx, helps.ParseAntigravityUsage(bodyBytes)) - var param any - converted := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bodyBytes, ¶m) - if responseFormat == sdktranslator.FormatOpenAIResponse { - converted = helps.EnsureResponsesUsageDetails(converted) - } - resp = cliproxyexecutor.Response{Payload: converted, Headers: httpResp.Header.Clone()} - reporter.EnsurePublished(ctx) - return resp, nil + // Credential retry rounds are owned by the conductor. Perform one upstream + // request per credential so request-retry is not consumed twice. + requestPayload := translated + if useCredits { + if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { + requestPayload = cp + helps.MarkCreditsUsed(ctx) } - - switch { - case lastStatus != 0: - err = newAntigravityStatusErr(lastStatus, lastBody) - case lastErr != nil: - err = lastErr - default: - err = statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} + } + replayScope := antigravityReasoningReplayScope{} + if antigravityUsesReasoningReplayCache(baseModel) { + var errReplay error + requestPayload, replayScope, errReplay = prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, requestPayload) + if errReplay != nil { + err = errReplay + return resp, err } + } + requestPayload = ensureAntigravityGeminiLeadingUserContent(baseModel, requestPayload) + + httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, false, opts.Alt, baseURL, helps.DerivedAntigravitySessionID(opts.Metadata, req.Metadata)) + if errReq != nil { + err = errReq return resp, err } - return resp, err + httpResp, errDo := httpClient.Do(httpReq) + if errDo != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errDo) + if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { + return resp, errDo + } + err = errDo + return resp, err + } + + helps.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 response body error: %v", errClose) + } + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + err = errRead + return resp, err + } + helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) + + if httpResp.StatusCode == http.StatusTooManyRequests { + decision := decideAntigravity429(bodyBytes) + switch decision.kind { + case antigravity429DecisionShortCooldownSwitchAuth: + if decision.retryAfter != nil && *decision.retryAfter > 0 { + if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { + err = homeKVUnavailableStatusErr(errMarkCooldown) + return resp, err + } + log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel) + } + case antigravity429DecisionFullQuotaExhausted: + if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { + markAntigravityCreditsPermanentlyDisabled(auth) + } + // No credits logic - just fall through to error return below + } + } + + if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { + log.Debugf("antigravity executor: upstream error status: %d, body: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), bodyBytes)) + if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, bodyBytes); errClear != nil { + // Report the upstream failure rather than the cleanup failure. + logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) + } + err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) + return resp, err + } + + // Success + if useCredits { + clearAntigravityCreditsFailureState(auth) + } + cacheAntigravityReasoningReplayFromResponse(ctx, replayScope, requestPayload, bodyBytes) + bodyBytes = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, bodyBytes) + reporter.Publish(ctx, helps.ParseAntigravityUsage(bodyBytes)) + var param any + converted := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bodyBytes, ¶m) + if responseFormat == sdktranslator.FormatOpenAIResponse { + converted = helps.EnsureResponsesUsageDetails(converted) + } + resp = cliproxyexecutor.Response{Payload: converted, Headers: httpResp.Header.Clone()} + reporter.EnsurePublished(ctx) + return resp, nil } // executeClaudeNonStream performs a claude non-streaming request to the Antigravity API. @@ -311,251 +231,164 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth * useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg) - baseURLs := antigravityBaseURLFallbackOrder(auth) + baseURL := resolveAntigravityRequestBaseURL(auth) httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) httpClient = reporter.TrackHTTPClient(httpClient) - // Credential retry rounds are owned by the conductor. Keep one upstream - // attempt per credential so request-retry is not consumed twice. - attempts := 1 - -attemptLoop: - for attempt := 0; attempt < attempts; attempt++ { - var lastStatus int - var lastBody []byte - var lastErr error - - for idx, baseURL := range baseURLs { - requestPayload := translated - if useCredits { - if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { - requestPayload = cp - helps.MarkCreditsUsed(ctx) - } - } - replayScope := antigravityReasoningReplayScope{} - if antigravityUsesReasoningReplayCache(baseModel) { - var errReplay error - requestPayload, replayScope, errReplay = prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, requestPayload) - if errReplay != nil { - err = errReplay - return resp, err - } - } - requestPayload = ensureAntigravityGeminiLeadingUserContent(baseModel, requestPayload) - httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL, helps.DerivedAntigravitySessionID(opts.Metadata, req.Metadata)) - if errReq != nil { - err = errReq - return resp, err - } - - httpResp, errDo := httpClient.Do(httpReq) - if errDo != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errDo) - if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { - return resp, errDo - } - lastStatus = 0 - lastBody = nil - lastErr = errDo - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - err = errDo - return resp, err - } - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { - bodyBytes, errRead := io.ReadAll(httpResp.Body) - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("antigravity executor: close response body error: %v", errClose) - } - if errRead != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errRead) - if errors.Is(errRead, context.Canceled) || errors.Is(errRead, context.DeadlineExceeded) { - err = errRead - return resp, err - } - if errCtx := ctx.Err(); errCtx != nil { - err = errCtx - return resp, err - } - lastStatus = 0 - lastBody = nil - lastErr = errRead - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: read error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - err = errRead - return resp, err - } - helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) - if httpResp.StatusCode == http.StatusTooManyRequests { - decision := decideAntigravity429(bodyBytes) - - switch decision.kind { - case antigravity429DecisionInstantRetrySameAuth: - if attempt+1 < attempts { - if decision.retryAfter != nil && *decision.retryAfter > 0 { - wait := antigravityInstantRetryDelay(*decision.retryAfter) - log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait) - if errWait := antigravityWait(ctx, wait); errWait != nil { - return resp, errWait - } - } - continue attemptLoop - } - case antigravity429DecisionShortCooldownSwitchAuth: - if decision.retryAfter != nil && *decision.retryAfter > 0 { - if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { - err = homeKVUnavailableStatusErr(errMarkCooldown) - return resp, err - } - log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel) - } - case antigravity429DecisionFullQuotaExhausted: - if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { - markAntigravityCreditsPermanentlyDisabled(auth) - } - // No credits logic - just fall through to error return below - } - } - - lastStatus = httpResp.StatusCode - lastBody = append([]byte(nil), bodyBytes...) - lastErr = nil - if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - if antigravityShouldRetryTransientResourceExhausted429(httpResp.StatusCode, bodyBytes) && attempt+1 < attempts { - delay := antigravityTransient429RetryDelay(attempt) - log.Debugf("antigravity executor: transient 429 resource exhausted for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return resp, errWait - } - continue attemptLoop - } - if antigravityShouldRetryNoCapacity(httpResp.StatusCode, bodyBytes) { - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: no capacity on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - if attempt+1 < attempts { - delay := antigravityNoCapacityRetryDelay(attempt) - log.Debugf("antigravity executor: no capacity for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return resp, errWait - } - continue attemptLoop - } - } - if antigravityShouldRetrySoftRateLimit(httpResp.StatusCode, bodyBytes) { - if attempt+1 < attempts { - delay := antigravitySoftRateLimitDelay(attempt) - log.Debugf("antigravity executor: soft rate limit for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return resp, errWait - } - continue attemptLoop - } - } - if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, bodyBytes); errClear != nil { - // Report the upstream failure rather than the cleanup failure. - logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) - } - err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) - return resp, err - } - - // Stream success - if useCredits { - clearAntigravityCreditsFailureState(auth) - } - replayAccumulator := newAntigravityReasoningReplayAccumulator(replayScope, requestPayload) - out := make(chan cliproxyexecutor.StreamChunk) - go func(resp *http.Response) { - defer close(out) - defer func() { - if errClose := resp.Body.Close(); errClose != nil { - log.Errorf("antigravity executor: close response body error: %v", errClose) - } - }() - scanner := bufio.NewScanner(resp.Body) - scanner.Buffer(nil, streamScannerBuffer) - for scanner.Scan() { - line := scanner.Bytes() - helps.AppendAPIResponseChunk(ctx, e.cfg, line) - if replayAccumulator != nil { - replayAccumulator.ObserveSSELine(line) - } - - // Filter usage metadata for all models - // Only retain usage statistics in the terminal chunk - line = helps.FilterSSEUsageMetadata(line) - - payload := helps.JSONPayload(line) - if payload == nil { - continue - } - - if detail, ok := helps.ParseAntigravityStreamUsage(payload); ok { - reporter.Publish(ctx, detail) - } - - out <- cliproxyexecutor.StreamChunk{Payload: payload} - } - if errScan := scanner.Err(); errScan != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errScan) - reporter.PublishFailure(ctx, errScan) - out <- cliproxyexecutor.StreamChunk{Err: errScan} - } else { - if replayAccumulator != nil { - replayAccumulator.Commit(ctx) - } - reporter.EnsurePublished(ctx) - } - }(httpResp) - - var buffer bytes.Buffer - for chunk := range out { - if chunk.Err != nil { - return resp, chunk.Err - } - if len(chunk.Payload) > 0 { - _, _ = buffer.Write(chunk.Payload) - _, _ = buffer.Write([]byte("\n")) - } - } - resp = cliproxyexecutor.Response{Payload: e.convertStreamToNonStream(buffer.Bytes())} - - resp.Payload = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, resp.Payload) - reporter.Publish(ctx, helps.ParseAntigravityUsage(resp.Payload)) - var param any - converted := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, resp.Payload, ¶m) - if responseFormat == sdktranslator.FormatOpenAIResponse { - converted = helps.EnsureResponsesUsageDetails(converted) - } - resp = cliproxyexecutor.Response{Payload: converted, Headers: httpResp.Header.Clone()} - reporter.EnsurePublished(ctx) - - return resp, nil + // Credential retry rounds are owned by the conductor. Perform one upstream + // request per credential so request-retry is not consumed twice. + requestPayload := translated + if useCredits { + if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { + requestPayload = cp + helps.MarkCreditsUsed(ctx) } - - switch { - case lastStatus != 0: - err = newAntigravityStatusErr(lastStatus, lastBody) - case lastErr != nil: - err = lastErr - default: - err = statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} + } + replayScope := antigravityReasoningReplayScope{} + if antigravityUsesReasoningReplayCache(baseModel) { + var errReplay error + requestPayload, replayScope, errReplay = prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, requestPayload) + if errReplay != nil { + err = errReplay + return resp, err } + } + requestPayload = ensureAntigravityGeminiLeadingUserContent(baseModel, requestPayload) + httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL, helps.DerivedAntigravitySessionID(opts.Metadata, req.Metadata)) + if errReq != nil { + err = errReq return resp, err } - return resp, err + httpResp, errDo := httpClient.Do(httpReq) + if errDo != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errDo) + if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { + return resp, errDo + } + err = errDo + return resp, err + } + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { + bodyBytes, errRead := io.ReadAll(httpResp.Body) + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close response body error: %v", errClose) + } + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + if errors.Is(errRead, context.Canceled) || errors.Is(errRead, context.DeadlineExceeded) { + err = errRead + return resp, err + } + if errCtx := ctx.Err(); errCtx != nil { + err = errCtx + return resp, err + } + err = errRead + return resp, err + } + helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) + if httpResp.StatusCode == http.StatusTooManyRequests { + decision := decideAntigravity429(bodyBytes) + + switch decision.kind { + case antigravity429DecisionShortCooldownSwitchAuth: + if decision.retryAfter != nil && *decision.retryAfter > 0 { + if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { + err = homeKVUnavailableStatusErr(errMarkCooldown) + return resp, err + } + log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel) + } + case antigravity429DecisionFullQuotaExhausted: + if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { + markAntigravityCreditsPermanentlyDisabled(auth) + } + // No credits logic - just fall through to error return below + } + } + + if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, bodyBytes); errClear != nil { + // Report the upstream failure rather than the cleanup failure. + logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) + } + err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) + return resp, err + } + + // Stream success + if useCredits { + clearAntigravityCreditsFailureState(auth) + } + replayAccumulator := newAntigravityReasoningReplayAccumulator(replayScope, requestPayload) + out := make(chan cliproxyexecutor.StreamChunk) + go func(resp *http.Response) { + defer close(out) + defer func() { + if errClose := resp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close response body error: %v", errClose) + } + }() + scanner := bufio.NewScanner(resp.Body) + scanner.Buffer(nil, streamScannerBuffer) + for scanner.Scan() { + line := scanner.Bytes() + helps.AppendAPIResponseChunk(ctx, e.cfg, line) + if replayAccumulator != nil { + replayAccumulator.ObserveSSELine(line) + } + + // Filter usage metadata for all models + // Only retain usage statistics in the terminal chunk + line = helps.FilterSSEUsageMetadata(line) + + payload := helps.JSONPayload(line) + if payload == nil { + continue + } + + if detail, ok := helps.ParseAntigravityStreamUsage(payload); ok { + reporter.Publish(ctx, detail) + } + + out <- cliproxyexecutor.StreamChunk{Payload: payload} + } + if errScan := scanner.Err(); errScan != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errScan) + reporter.PublishFailure(ctx, errScan) + out <- cliproxyexecutor.StreamChunk{Err: errScan} + } else { + if replayAccumulator != nil { + replayAccumulator.Commit(ctx) + } + reporter.EnsurePublished(ctx) + } + }(httpResp) + + var buffer bytes.Buffer + for chunk := range out { + if chunk.Err != nil { + return resp, chunk.Err + } + if len(chunk.Payload) > 0 { + _, _ = buffer.Write(chunk.Payload) + _, _ = buffer.Write([]byte("\n")) + } + } + resp = cliproxyexecutor.Response{Payload: e.convertStreamToNonStream(buffer.Bytes())} + + resp.Payload = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, resp.Payload) + reporter.Publish(ctx, helps.ParseAntigravityUsage(resp.Payload)) + var param any + converted := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, resp.Payload, ¶m) + if responseFormat == sdktranslator.FormatOpenAIResponse { + converted = helps.EnsureResponsesUsageDetails(converted) + } + resp = cliproxyexecutor.Response{Payload: converted, Headers: httpResp.Header.Clone()} + reporter.EnsurePublished(ctx) + + return resp, nil } func (e *AntigravityExecutor) convertStreamToNonStream(stream []byte) []byte { diff --git a/internal/runtime/executor/antigravity_executor_request.go b/internal/runtime/executor/antigravity_executor_request.go index d4c7a790c..5dfec83bc 100644 --- a/internal/runtime/executor/antigravity_executor_request.go +++ b/internal/runtime/executor/antigravity_executor_request.go @@ -31,7 +31,7 @@ func (e *AntigravityExecutor) buildRequest(ctx context.Context, auth *cliproxyau base := strings.TrimSuffix(baseURL, "/") if base == "" { - base = buildBaseURL(auth) + base = resolveAntigravityRequestBaseURL(auth) } path := antigravityGeneratePath if stream { @@ -383,9 +383,12 @@ func antigravityRequestNeedsSchemaSanitization(payload []byte) bool { } return false } -func buildBaseURL(auth *cliproxyauth.Auth) string { - if baseURLs := antigravityBaseURLFallbackOrder(auth); len(baseURLs) > 0 { - return baseURLs[0] + +// resolveAntigravityRequestBaseURL selects one request endpoint without cross-tier fallback. +// Consumer credentials default to daily; enterprise/GCP credentials can set base_url explicitly. +func resolveAntigravityRequestBaseURL(auth *cliproxyauth.Auth) string { + if base := resolveCustomAntigravityBaseURL(auth); base != "" { + return base } return antigravityBaseURLDaily } @@ -433,17 +436,6 @@ func antigravityConfiguredUserAgent(auth *cliproxyauth.Auth) string { return raw } -var antigravityBaseURLFallbackOrder = func(auth *cliproxyauth.Auth) []string { - if base := resolveCustomAntigravityBaseURL(auth); base != "" { - return []string{base} - } - return []string{ - antigravityBaseURLDaily, - antigravityBaseURLProd, - // antigravitySandboxBaseURLDaily, - } -} - func resolveCustomAntigravityBaseURL(auth *cliproxyauth.Auth) string { if auth == nil { return "" diff --git a/internal/runtime/executor/antigravity_executor_signature_test.go b/internal/runtime/executor/antigravity_executor_signature_test.go index b0d4d473e..98d9fa4ad 100644 --- a/internal/runtime/executor/antigravity_executor_signature_test.go +++ b/internal/runtime/executor/antigravity_executor_signature_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/base64" + "errors" "fmt" "io" "net/http" @@ -525,6 +526,99 @@ func TestAntigravityStreamDoesNotPrependLeadingUserForClaudeTarget(t *testing.T) } } +func TestAntigravityRequestPathsDoNotFallbackEndpoints(t *testing.T) { + type upstreamCall struct { + host string + path string + } + routes := []struct { + name string + kind string + model string + wantPath string + }{ + {name: "generate non-stream", kind: "execute", model: "gemini-3.6-flash-high", wantPath: antigravityGeneratePath}, + {name: "Claude non-stream", kind: "execute", model: "claude-sonnet-4-6", wantPath: antigravityStreamPath}, + {name: "generate stream", kind: "stream", model: "gemini-3.6-flash-high", wantPath: antigravityStreamPath}, + {name: "count tokens", kind: "count", model: "gemini-3.6-flash-high", wantPath: antigravityCountTokensPath}, + } + failures := []struct { + name string + transport bool + }{ + {name: "HTTP 429"}, + {name: "transport error", transport: true}, + } + + for _, route := range routes { + for _, failure := range failures { + t.Run(route.name+"/"+failure.name, func(t *testing.T) { + calls := make([]upstreamCall, 0, 1) + transportErr := errors.New("daily endpoint unavailable") + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", roundTripperFunc(func(req *http.Request) (*http.Response, error) { + calls = append(calls, upstreamCall{host: req.URL.Host, path: req.URL.Path}) + if failure.transport { + return nil, transportErr + } + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"code":429,"status":"RESOURCE_EXHAUSTED"}}`)), + }, nil + })) + auth := &cliproxyauth.Auth{Metadata: map[string]any{ + "access_token": "token", + "expired": time.Now().Add(2 * time.Hour).Format(time.RFC3339), + "project_id": "project-1", + }} + req := cliproxyexecutor.Request{ + Model: route.model, + Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatGemini, + ResponseFormat: sdktranslator.FormatGemini, + } + executor := NewAntigravityExecutor(&config.Config{}) + + var errRequest error + switch route.kind { + case "execute": + _, errRequest = executor.Execute(ctx, auth, req, opts) + case "stream": + _, errRequest = executor.ExecuteStream(ctx, auth, req, opts) + case "count": + _, errRequest = executor.CountTokens(ctx, auth, req, opts) + default: + t.Fatalf("unknown route kind %q", route.kind) + } + if errRequest == nil { + t.Fatal("request error = nil, want original upstream error") + } + if failure.transport { + if !errors.Is(errRequest, transportErr) { + t.Fatalf("request error = %v, want transport error", errRequest) + } + } else { + status, ok := errRequest.(interface{ StatusCode() int }) + if !ok || status.StatusCode() != http.StatusTooManyRequests { + t.Fatalf("request error = %v, want status 429", errRequest) + } + } + if len(calls) != 1 { + t.Fatalf("upstream calls = %v, want exactly one daily endpoint request", calls) + } + if got, want := calls[0].host, resolveHost(antigravityBaseURLDaily); got != want { + t.Fatalf("upstream host = %q, want %q", got, want) + } + if got := calls[0].path; got != route.wantPath { + t.Fatalf("upstream path = %q, want %q", got, route.wantPath) + } + }) + } + } +} + func TestAntigravityCountTokensMatchesTargetLeadingUserPolicy(t *testing.T) { tests := []struct { name string diff --git a/internal/runtime/executor/antigravity_executor_stream.go b/internal/runtime/executor/antigravity_executor_stream.go index 96e44d2f9..9c008f73a 100644 --- a/internal/runtime/executor/antigravity_executor_stream.go +++ b/internal/runtime/executor/antigravity_executor_stream.go @@ -78,246 +78,159 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg) - baseURLs := antigravityBaseURLFallbackOrder(auth) + baseURL := resolveAntigravityRequestBaseURL(auth) httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) httpClient = reporter.TrackHTTPClient(httpClient) - // Credential retry rounds are owned by the conductor. Keep one upstream - // attempt per credential so request-retry is not consumed twice. - attempts := 1 - -attemptLoop: - for attempt := 0; attempt < attempts; attempt++ { - var lastStatus int - var lastBody []byte - var lastErr error - - for idx, baseURL := range baseURLs { - requestPayload := translated - if useCredits { - if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { - requestPayload = cp - helps.MarkCreditsUsed(ctx) - } - } - replayScope := antigravityReasoningReplayScope{} - if antigravityUsesReasoningReplayCache(baseModel) { - var errReplay error - requestPayload, replayScope, errReplay = prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, requestPayload) - if errReplay != nil { - err = errReplay - return nil, err - } - } - requestPayload = ensureAntigravityGeminiLeadingUserContent(baseModel, requestPayload) - httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL, helps.DerivedAntigravitySessionID(opts.Metadata, req.Metadata)) - if errReq != nil { - err = errReq + // Credential retry rounds are owned by the conductor. Perform one upstream + // request per credential so request-retry is not consumed twice. + requestPayload := translated + if useCredits { + if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { + requestPayload = cp + helps.MarkCreditsUsed(ctx) + } + } + replayScope := antigravityReasoningReplayScope{} + if antigravityUsesReasoningReplayCache(baseModel) { + var errReplay error + requestPayload, replayScope, errReplay = prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, requestPayload) + if errReplay != nil { + err = errReplay + return nil, err + } + } + requestPayload = ensureAntigravityGeminiLeadingUserContent(baseModel, requestPayload) + httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL, helps.DerivedAntigravitySessionID(opts.Metadata, req.Metadata)) + if errReq != nil { + err = errReq + return nil, err + } + httpResp, errDo := httpClient.Do(httpReq) + if errDo != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errDo) + if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { + return nil, errDo + } + err = errDo + return nil, err + } + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { + bodyBytes, errRead := io.ReadAll(httpResp.Body) + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close response body error: %v", errClose) + } + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + if errors.Is(errRead, context.Canceled) || errors.Is(errRead, context.DeadlineExceeded) { + err = errRead return nil, err } - httpResp, errDo := httpClient.Do(httpReq) - if errDo != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errDo) - if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { - return nil, errDo - } - lastStatus = 0 - lastBody = nil - lastErr = errDo - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - err = errDo + if errCtx := ctx.Err(); errCtx != nil { + err = errCtx return nil, err } - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { - bodyBytes, errRead := io.ReadAll(httpResp.Body) - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("antigravity executor: close response body error: %v", errClose) - } - if errRead != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errRead) - if errors.Is(errRead, context.Canceled) || errors.Is(errRead, context.DeadlineExceeded) { - err = errRead + err = errRead + return nil, err + } + helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) + if httpResp.StatusCode == http.StatusTooManyRequests { + decision := decideAntigravity429(bodyBytes) + + switch decision.kind { + case antigravity429DecisionShortCooldownSwitchAuth: + if decision.retryAfter != nil && *decision.retryAfter > 0 { + if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { + err = homeKVUnavailableStatusErr(errMarkCooldown) return nil, err } - if errCtx := ctx.Err(); errCtx != nil { - err = errCtx - return nil, err - } - lastStatus = 0 - lastBody = nil - lastErr = errRead - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: read error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - err = errRead - return nil, err + log.Debugf("antigravity executor: short quota cooldown (%s) for model %s recorded", *decision.retryAfter, baseModel) } - helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) - if httpResp.StatusCode == http.StatusTooManyRequests { - decision := decideAntigravity429(bodyBytes) - - switch decision.kind { - case antigravity429DecisionInstantRetrySameAuth: - if attempt+1 < attempts { - if decision.retryAfter != nil && *decision.retryAfter > 0 { - wait := antigravityInstantRetryDelay(*decision.retryAfter) - log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait) - if errWait := antigravityWait(ctx, wait); errWait != nil { - return nil, errWait - } - } - continue attemptLoop - } - case antigravity429DecisionShortCooldownSwitchAuth: - if decision.retryAfter != nil && *decision.retryAfter > 0 { - if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { - err = homeKVUnavailableStatusErr(errMarkCooldown) - return nil, err - } - log.Debugf("antigravity executor: short quota cooldown (%s) for model %s recorded", *decision.retryAfter, baseModel) - } - case antigravity429DecisionFullQuotaExhausted: - if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { - markAntigravityCreditsPermanentlyDisabled(auth) - } - // No credits logic - just fall through to error return below - } + case antigravity429DecisionFullQuotaExhausted: + if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { + markAntigravityCreditsPermanentlyDisabled(auth) } - - lastStatus = httpResp.StatusCode - lastBody = append([]byte(nil), bodyBytes...) - lastErr = nil - if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - if antigravityShouldRetryTransientResourceExhausted429(httpResp.StatusCode, bodyBytes) && attempt+1 < attempts { - delay := antigravityTransient429RetryDelay(attempt) - log.Debugf("antigravity executor: transient 429 resource exhausted for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return nil, errWait - } - continue attemptLoop - } - if antigravityShouldRetryNoCapacity(httpResp.StatusCode, bodyBytes) { - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: no capacity on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - if attempt+1 < attempts { - delay := antigravityNoCapacityRetryDelay(attempt) - log.Debugf("antigravity executor: no capacity for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return nil, errWait - } - continue attemptLoop - } - } - if antigravityShouldRetrySoftRateLimit(httpResp.StatusCode, bodyBytes) { - if attempt+1 < attempts { - delay := antigravitySoftRateLimitDelay(attempt) - log.Debugf("antigravity executor: soft rate limit for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return nil, errWait - } - continue attemptLoop - } - } - if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, bodyBytes); errClear != nil { - // Report the upstream failure rather than the cleanup failure. - logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) - } - err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) - return nil, err + // No credits logic - just fall through to error return below } - - // Stream success - if useCredits { - clearAntigravityCreditsFailureState(auth) - } - replayAccumulator := newAntigravityReasoningReplayAccumulator(replayScope, requestPayload) - out := make(chan cliproxyexecutor.StreamChunk) - go func(resp *http.Response) { - defer close(out) - defer func() { - if errClose := resp.Body.Close(); errClose != nil { - log.Errorf("antigravity executor: close response line error: %v", errClose) - } - }() - scanner := bufio.NewScanner(resp.Body) - scanner.Buffer(nil, streamScannerBuffer) - claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) - var param any - for scanner.Scan() { - line := scanner.Bytes() - helps.AppendAPIResponseChunk(ctx, e.cfg, line) - if replayAccumulator != nil { - replayAccumulator.ObserveSSELine(line) - } - - // Filter usage metadata for all models - // Only retain usage statistics in the terminal chunk - line = helps.FilterSSEUsageMetadata(line) - - payload := helps.JSONPayload(line) - if payload == nil { - continue - } - - if detail, ok := helps.ParseAntigravityStreamUsage(payload); ok { - reporter.Publish(ctx, detail) - } - - payload = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, payload) - chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bytes.Clone(payload), ¶m, claudeInputTokens) - for i := range chunks { - select { - case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: - case <-ctx.Done(): - return - } - } - } - tail := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, []byte("[DONE]"), ¶m, claudeInputTokens) - for i := range tail { - select { - case out <- cliproxyexecutor.StreamChunk{Payload: tail[i]}: - case <-ctx.Done(): - return - } - } - if errScan := scanner.Err(); errScan != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errScan) - reporter.PublishFailure(ctx, errScan) - select { - case out <- cliproxyexecutor.StreamChunk{Err: errScan}: - case <-ctx.Done(): - } - } else { - if replayAccumulator != nil { - replayAccumulator.Commit(ctx) - } - reporter.EnsurePublished(ctx) - } - }(httpResp) - return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil } - switch { - case lastStatus != 0: - err = newAntigravityStatusErr(lastStatus, lastBody) - case lastErr != nil: - err = lastErr - default: - err = statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} + if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, bodyBytes); errClear != nil { + // Report the upstream failure rather than the cleanup failure. + logAntigravityReasoningReplayDegraded(replayScope, "invalidate", errClear) } + err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) return nil, err } - return nil, err + // Stream success + if useCredits { + clearAntigravityCreditsFailureState(auth) + } + replayAccumulator := newAntigravityReasoningReplayAccumulator(replayScope, requestPayload) + out := make(chan cliproxyexecutor.StreamChunk) + go func(resp *http.Response) { + defer close(out) + defer func() { + if errClose := resp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close response line error: %v", errClose) + } + }() + scanner := bufio.NewScanner(resp.Body) + scanner.Buffer(nil, streamScannerBuffer) + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) + var param any + for scanner.Scan() { + line := scanner.Bytes() + helps.AppendAPIResponseChunk(ctx, e.cfg, line) + if replayAccumulator != nil { + replayAccumulator.ObserveSSELine(line) + } + + // Filter usage metadata for all models + // Only retain usage statistics in the terminal chunk + line = helps.FilterSSEUsageMetadata(line) + + payload := helps.JSONPayload(line) + if payload == nil { + continue + } + + if detail, ok := helps.ParseAntigravityStreamUsage(payload); ok { + reporter.Publish(ctx, detail) + } + + payload = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, payload) + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bytes.Clone(payload), ¶m, claudeInputTokens) + for i := range chunks { + select { + case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: + case <-ctx.Done(): + return + } + } + } + tail := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, []byte("[DONE]"), ¶m, claudeInputTokens) + for i := range tail { + select { + case out <- cliproxyexecutor.StreamChunk{Payload: tail[i]}: + case <-ctx.Done(): + return + } + } + if errScan := scanner.Err(); errScan != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errScan) + reporter.PublishFailure(ctx, errScan) + select { + case out <- cliproxyexecutor.StreamChunk{Err: errScan}: + case <-ctx.Done(): + } + } else { + if replayAccumulator != nil { + replayAccumulator.Commit(ctx) + } + reporter.EnsurePublished(ctx) + } + }(httpResp) + return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil } diff --git a/internal/runtime/executor/antigravity_executor_tokens.go b/internal/runtime/executor/antigravity_executor_tokens.go index fb422131f..654124417 100644 --- a/internal/runtime/executor/antigravity_executor_tokens.go +++ b/internal/runtime/executor/antigravity_executor_tokens.go @@ -3,7 +3,6 @@ package executor import ( "bytes" "context" - "errors" "io" "net/http" "net/url" @@ -67,7 +66,7 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut payload = helps.DeleteJSONField(payload, "model") payload = helps.DeleteJSONField(payload, "request.safetySettings") - baseURLs := antigravityBaseURLFallbackOrder(auth) + base := resolveAntigravityRequestBaseURL(auth) httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) var authID, authLabel, authType, authValue string @@ -77,114 +76,71 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut authType, authValue = auth.AccountInfo() } - var lastStatus int - var lastBody []byte - var lastErr error - - for idx, baseURL := range baseURLs { - base := strings.TrimSuffix(baseURL, "/") - if base == "" { - base = buildBaseURL(auth) - } - - var requestURL strings.Builder - requestURL.WriteString(base) - requestURL.WriteString(antigravityCountTokensPath) - if opts.Alt != "" { - requestURL.WriteString("?$alt=") - requestURL.WriteString(url.QueryEscape(opts.Alt)) - } - - httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bytes.NewReader(payload)) - if errReq != nil { - return cliproxyexecutor.Response{}, errReq - } - // No httpReq.Close: keep the shared Antigravity connection pool usable. - httpReq.Header.Set("Content-Type", "application/json") - httpReq.Header.Set("Authorization", "Bearer "+token) - httpReq.Header.Set("User-Agent", resolveUserAgent(auth)) - if host := resolveHost(base); host != "" { - httpReq.Host = host - } - var attrs map[string]string - if auth != nil { - attrs = auth.Attributes - } - util.ApplyCustomHeadersFromAttrs(httpReq, attrs) - - helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ - URL: requestURL.String(), - Method: http.MethodPost, - Headers: httpReq.Header.Clone(), - Body: payload, - Provider: e.Identifier(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - }) - - httpResp, errDo := httpClient.Do(httpReq) - if errDo != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errDo) - if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { - return cliproxyexecutor.Response{}, errDo - } - lastStatus = 0 - lastBody = nil - lastErr = errDo - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - return cliproxyexecutor.Response{}, errDo - } - - helps.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 response body error: %v", errClose) - } - if errRead != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errRead) - return cliproxyexecutor.Response{}, errRead - } - helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) - - if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices { - count := gjson.GetBytes(bodyBytes, "totalTokens").Int() - translated := sdktranslator.TranslateTokenCount(respCtx, to, responseFormat, count, bodyBytes) - return cliproxyexecutor.Response{Payload: translated, Headers: httpResp.Header.Clone()}, nil - } - - lastStatus = httpResp.StatusCode - lastBody = append([]byte(nil), bodyBytes...) - lastErr = nil - if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)} - if httpResp.StatusCode == http.StatusTooManyRequests { - if retryAfter, parseErr := helps.ParseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil { - sErr.retryAfter = retryAfter - } - } - return cliproxyexecutor.Response{}, sErr + var requestURL strings.Builder + requestURL.WriteString(base) + requestURL.WriteString(antigravityCountTokensPath) + if opts.Alt != "" { + requestURL.WriteString("?$alt=") + requestURL.WriteString(url.QueryEscape(opts.Alt)) } - switch { - case lastStatus != 0: - sErr := statusErr{code: lastStatus, msg: string(lastBody)} - if lastStatus == http.StatusTooManyRequests { - if retryAfter, parseErr := helps.ParseRetryDelay(lastBody); parseErr == nil && retryAfter != nil { - sErr.retryAfter = retryAfter - } - } - return cliproxyexecutor.Response{}, sErr - case lastErr != nil: - return cliproxyexecutor.Response{}, lastErr - default: - return cliproxyexecutor.Response{}, statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} + httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bytes.NewReader(payload)) + if errReq != nil { + return cliproxyexecutor.Response{}, errReq } + // No httpReq.Close: keep the shared Antigravity connection pool usable. + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Authorization", "Bearer "+token) + httpReq.Header.Set("User-Agent", resolveUserAgent(auth)) + if host := resolveHost(base); host != "" { + httpReq.Host = host + } + var attrs map[string]string + if auth != nil { + attrs = auth.Attributes + } + util.ApplyCustomHeadersFromAttrs(httpReq, attrs) + + helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ + URL: requestURL.String(), + Method: http.MethodPost, + Headers: httpReq.Header.Clone(), + Body: payload, + Provider: e.Identifier(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + }) + + httpResp, errDo := httpClient.Do(httpReq) + if errDo != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errDo) + return cliproxyexecutor.Response{}, errDo + } + + helps.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 response body error: %v", errClose) + } + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + return cliproxyexecutor.Response{}, errRead + } + helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) + + if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices { + count := gjson.GetBytes(bodyBytes, "totalTokens").Int() + translated := sdktranslator.TranslateTokenCount(respCtx, to, responseFormat, count, bodyBytes) + return cliproxyexecutor.Response{Payload: translated, Headers: httpResp.Header.Clone()}, nil + } + + sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)} + if httpResp.StatusCode == http.StatusTooManyRequests { + if retryAfter, parseErr := helps.ParseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil { + sErr.retryAfter = retryAfter + } + } + return cliproxyexecutor.Response{}, sErr }