Files
CLIProxyAPI/internal/runtime/executor/antigravity_executor.go
1137043480 bf0e5c23f7 fix: prevent goroutine leaks in streaming executors via context-aware channel sends
All streaming executors use bare channel sends (out <- chunk) inside goroutines
that process upstream SSE responses. When the downstream consumer disconnects
(client timeout, network drop, etc.), these sends block indefinitely, causing
the goroutine and all associated resources (HTTP response body, scanner buffers,
translation state) to leak permanently.

Over time, leaked goroutines accumulate monotonically, leading to RSS growth
from ~30MB to 3.7GB+ and eventual OOM kills on resource-constrained VPS hosts.

Fix: Replace all bare 'out <- ...' sends with:
  select {
  case out <- ...:
  case <-ctx.Done():
    return
  }

This ensures goroutines terminate promptly when the request context is canceled,
allowing GC to reclaim all associated resources.

Affected executors (9 files, 36+ send sites):
- antigravity_executor.go (5 sites)
- gemini_cli_executor.go (6 sites)
- gemini_vertex_executor.go (6 sites)
- aistudio_executor.go (4 sites)
- gemini_executor.go (3 sites)
- openai_compat_executor.go (3 sites)
- claude_executor.go (4 sites)
- codex_executor.go (2 sites)
- kimi_executor.go (3 sites)
2026-05-03 11:25:04 -04:00

2397 lines
80 KiB
Go

// Package executor provides runtime execution capabilities for various AI service providers.
// This file implements the Antigravity executor that proxies requests to the antigravity
// upstream using OAuth credentials.
package executor
import (
"bufio"
"bytes"
"context"
"crypto/sha256"
"crypto/tls"
"encoding/binary"
"encoding/json"
"errors"
"fmt"
"io"
"math/rand"
"net/http"
"net/url"
"strconv"
"strings"
"sync"
"time"
"github.com/google/uuid"
"github.com/router-for-me/CLIProxyAPI/v6/internal/cache"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor/helps"
"github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
antigravityclaude "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/claude"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
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"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
const (
antigravityBaseURLDaily = "https://daily-cloudcode-pa.googleapis.com"
antigravitySandboxBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com"
antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com"
antigravityCountTokensPath = "/v1internal:countTokens"
antigravityStreamPath = "/v1internal:streamGenerateContent"
antigravityGeneratePath = "/v1internal:generateContent"
antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
antigravityClientSecret = "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"
defaultAntigravityAgent = "antigravity/1.21.9 darwin/arm64" // fallback only; overridden at runtime by misc.AntigravityUserAgent()
antigravityAuthType = "antigravity"
refreshSkew = 3000 * time.Second
antigravityCreditsHintRefreshInterval = 10 * time.Minute
antigravityCreditsHintRefreshTimeout = 5 * time.Second
antigravityShortQuotaCooldownThreshold = 5 * time.Minute
antigravityInstantRetryThreshold = 3 * time.Second
// 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
type antigravityCreditsFailureState struct {
PermanentlyDisabled bool
ExplicitBalanceExhausted bool
}
type antigravity429DecisionKind string
const (
antigravity429Unknown antigravity429Category = "unknown"
antigravity429RateLimited antigravity429Category = "rate_limited"
antigravity429QuotaExhausted antigravity429Category = "quota_exhausted"
antigravity429SoftRateLimit antigravity429Category = "soft_rate_limit"
antigravity429DecisionSoftRetry antigravity429DecisionKind = "soft_retry"
antigravity429DecisionInstantRetrySameAuth antigravity429DecisionKind = "instant_retry_same_auth"
antigravity429DecisionShortCooldownSwitchAuth antigravity429DecisionKind = "short_cooldown_switch_auth"
antigravity429DecisionFullQuotaExhausted antigravity429DecisionKind = "full_quota_exhausted"
)
type antigravity429Decision struct {
kind antigravity429DecisionKind
retryAfter *time.Duration
reason string
}
var (
randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
randSourceMutex sync.Mutex
antigravityCreditsFailureByAuth sync.Map
antigravityShortCooldownByAuth sync.Map
antigravityCreditsBalanceByAuth sync.Map // auth.ID → antigravityCreditsBalance
antigravityCreditsHintRefreshByID sync.Map // auth.ID → *antigravityCreditsHintRefreshState
antigravityQuotaExhaustedKeywords = []string{
"quota_exhausted",
"quota exhausted",
}
)
type antigravityCreditsBalance struct {
CreditAmount float64
MinCreditAmount float64
PaidTierID string
Known bool
}
type antigravityCreditsHintRefreshState struct {
mu sync.Mutex
lastAttempt time.Time
}
func antigravityAuthHasCredits(auth *cliproxyauth.Auth) bool {
if auth == nil || strings.TrimSpace(auth.ID) == "" {
return false
}
if hint, ok := cliproxyauth.GetAntigravityCreditsHint(auth.ID); ok && hint.Known {
return hint.Available
}
val, ok := antigravityCreditsBalanceByAuth.Load(strings.TrimSpace(auth.ID))
if !ok {
return true // optimistic: assume credits available when balance unknown
}
bal, valid := val.(antigravityCreditsBalance)
if !valid {
antigravityCreditsBalanceByAuth.Delete(strings.TrimSpace(auth.ID))
return false
}
if !bal.Known {
return false
}
available := bal.CreditAmount >= bal.MinCreditAmount
cliproxyauth.SetAntigravityCreditsHint(strings.TrimSpace(auth.ID), cliproxyauth.AntigravityCreditsHint{
Known: true,
Available: available,
CreditAmount: bal.CreditAmount,
MinCreditAmount: bal.MinCreditAmount,
PaidTierID: bal.PaidTierID,
UpdatedAt: time.Now(),
})
return available
}
// parseMetaFloat extracts a float64 from auth.Metadata (handles string and numeric types).
func parseMetaFloat(metadata map[string]any, key string) (float64, bool) {
v, ok := metadata[key]
if !ok {
return 0, false
}
switch typed := v.(type) {
case float64:
return typed, true
case int:
return float64(typed), true
case int64:
return float64(typed), true
case uint64:
return float64(typed), true
case json.Number:
if f, err := typed.Float64(); err == nil {
return f, true
}
case string:
if f, err := strconv.ParseFloat(strings.TrimSpace(typed), 64); err == nil {
return f, true
}
}
return 0, false
}
// AntigravityExecutor proxies requests to the antigravity upstream.
type AntigravityExecutor struct {
cfg *config.Config
}
// NewAntigravityExecutor creates a new Antigravity executor instance.
//
// Parameters:
// - cfg: The application configuration
//
// Returns:
// - *AntigravityExecutor: A new Antigravity executor instance
func NewAntigravityExecutor(cfg *config.Config) *AntigravityExecutor {
return &AntigravityExecutor{cfg: cfg}
}
// antigravityTransport is a singleton HTTP/1.1 transport shared by all Antigravity requests.
// It is initialized once via antigravityTransportOnce to avoid leaking a new connection pool
// (and the goroutines managing it) on every request.
var (
antigravityTransport *http.Transport
antigravityTransportOnce sync.Once
)
func cloneTransportWithHTTP11(base *http.Transport) *http.Transport {
if base == nil {
return nil
}
clone := base.Clone()
clone.ForceAttemptHTTP2 = false
// Wipe TLSNextProto to prevent implicit HTTP/2 upgrade.
clone.TLSNextProto = make(map[string]func(authority string, c *tls.Conn) http.RoundTripper)
if clone.TLSClientConfig == nil {
clone.TLSClientConfig = &tls.Config{}
} else {
clone.TLSClientConfig = clone.TLSClientConfig.Clone()
}
// Actively advertise only HTTP/1.1 in the ALPN handshake.
clone.TLSClientConfig.NextProtos = []string{"http/1.1"}
return clone
}
// initAntigravityTransport creates the shared HTTP/1.1 transport exactly once.
func initAntigravityTransport() {
base, ok := http.DefaultTransport.(*http.Transport)
if !ok {
base = &http.Transport{}
}
antigravityTransport = cloneTransportWithHTTP11(base)
}
// newAntigravityHTTPClient creates an HTTP client specifically for Antigravity,
// enforcing HTTP/1.1 by disabling HTTP/2 to perfectly mimic Node.js https defaults.
// The underlying Transport is a singleton to avoid leaking connection pools.
func newAntigravityHTTPClient(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) *http.Client {
antigravityTransportOnce.Do(initAntigravityTransport)
client := helps.NewProxyAwareHTTPClient(ctx, cfg, auth, timeout)
// If no transport is set, use the shared HTTP/1.1 transport.
if client.Transport == nil {
client.Transport = antigravityTransport
return client
}
// Preserve proxy settings from proxy-aware transports while forcing HTTP/1.1.
if transport, ok := client.Transport.(*http.Transport); ok {
client.Transport = cloneTransportWithHTTP11(transport)
}
return client
}
func validateAntigravityRequestSignatures(from sdktranslator.Format, rawJSON []byte) ([]byte, error) {
if from.String() != "claude" {
return rawJSON, nil
}
// Always strip thinking blocks with invalid signatures (empty or non-Claude-format).
rawJSON = antigravityclaude.StripEmptySignatureThinkingBlocks(rawJSON)
if cache.SignatureCacheEnabled() {
return rawJSON, nil
}
if !cache.SignatureBypassStrictMode() {
// Non-strict bypass: let the translator handle invalid signatures
// by dropping unsigned thinking blocks silently (no 400).
return rawJSON, nil
}
if err := antigravityclaude.ValidateClaudeBypassSignatures(rawJSON); err != nil {
return rawJSON, statusErr{code: http.StatusBadRequest, msg: err.Error()}
}
return rawJSON, nil
}
// Identifier returns the executor identifier.
func (e *AntigravityExecutor) Identifier() string { return antigravityAuthType }
// PrepareRequest injects Antigravity credentials into the outgoing HTTP request.
func (e *AntigravityExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error {
if req == nil {
return nil
}
token, _, errToken := e.ensureAccessToken(req.Context(), auth)
if errToken != nil {
return errToken
}
if strings.TrimSpace(token) == "" {
return statusErr{code: http.StatusUnauthorized, msg: "missing access token"}
}
req.Header.Set("Authorization", "Bearer "+token)
return nil
}
// HttpRequest injects Antigravity credentials into the request and executes it.
// It uses a whitelist approach: all incoming headers are stripped and only
// the minimum set required by the Antigravity protocol is explicitly set.
func (e *AntigravityExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) {
if req == nil {
return nil, fmt.Errorf("antigravity executor: request is nil")
}
if ctx == nil {
ctx = req.Context()
}
httpReq := req.WithContext(ctx)
// --- Whitelist: save only the headers we need from the original request ---
contentType := httpReq.Header.Get("Content-Type")
// Wipe ALL incoming headers
for k := range httpReq.Header {
delete(httpReq.Header, k)
}
// --- Set only the headers Antigravity actually sends ---
if contentType != "" {
httpReq.Header.Set("Content-Type", contentType)
}
// Content-Length is managed automatically by Go's http.Client from the Body
httpReq.Header.Set("User-Agent", resolveUserAgent(auth))
httpReq.Close = true // sends Connection: close
// Inject Authorization: Bearer <token>
if err := e.PrepareRequest(httpReq, auth); err != nil {
return nil, err
}
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
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 {
switch decideAntigravity429(body).kind {
case antigravity429DecisionInstantRetrySameAuth, antigravity429DecisionShortCooldownSwitchAuth:
return antigravity429RateLimited
case antigravity429DecisionFullQuotaExhausted:
return antigravity429QuotaExhausted
case antigravity429DecisionSoftRetry:
return antigravity429SoftRateLimit
default:
return antigravity429Unknown
}
}
func decideAntigravity429(body []byte) antigravity429Decision {
decision := antigravity429Decision{kind: antigravity429DecisionSoftRetry}
if len(body) == 0 {
return decision
}
if retryAfter, parseErr := parseRetryDelay(body); parseErr == nil && retryAfter != nil {
decision.retryAfter = retryAfter
}
status := strings.TrimSpace(gjson.GetBytes(body, "error.status").String())
if !strings.EqualFold(status, "RESOURCE_EXHAUSTED") {
return decision
}
details := gjson.GetBytes(body, "error.details")
if details.Exists() && details.IsArray() {
for _, detail := range details.Array() {
if detail.Get("@type").String() != "type.googleapis.com/google.rpc.ErrorInfo" {
continue
}
reason := strings.TrimSpace(detail.Get("reason").String())
decision.reason = reason
switch {
case strings.EqualFold(reason, "QUOTA_EXHAUSTED"):
decision.kind = antigravity429DecisionFullQuotaExhausted
return decision
case strings.EqualFold(reason, "RATE_LIMIT_EXCEEDED"):
if decision.retryAfter == nil {
decision.kind = antigravity429DecisionSoftRetry
return decision
}
switch {
case *decision.retryAfter < antigravityInstantRetryThreshold:
decision.kind = antigravity429DecisionInstantRetrySameAuth
case *decision.retryAfter < antigravityShortQuotaCooldownThreshold:
decision.kind = antigravity429DecisionShortCooldownSwitchAuth
default:
decision.kind = antigravity429DecisionFullQuotaExhausted
}
return decision
}
}
}
lowerBody := strings.ToLower(string(body))
for _, keyword := range antigravityQuotaExhaustedKeywords {
if strings.Contains(lowerBody, keyword) {
decision.kind = antigravity429DecisionFullQuotaExhausted
decision.reason = "quota_exhausted"
return decision
}
}
decision.kind = antigravity429DecisionSoftRetry
return decision
}
func antigravityCreditsRetryEnabled(cfg *config.Config) bool {
return cfg != nil && cfg.QuotaExceeded.AntigravityCredits
}
func clearAntigravityCreditsFailureState(auth *cliproxyauth.Auth) {
if auth == nil || strings.TrimSpace(auth.ID) == "" {
return
}
antigravityCreditsFailureByAuth.Delete(strings.TrimSpace(auth.ID))
}
func markAntigravityCreditsPermanentlyDisabled(auth *cliproxyauth.Auth) {
if auth == nil || strings.TrimSpace(auth.ID) == "" {
return
}
authID := strings.TrimSpace(auth.ID)
state := antigravityCreditsFailureState{
PermanentlyDisabled: true,
ExplicitBalanceExhausted: true,
}
antigravityCreditsFailureByAuth.Store(authID, state)
antigravityCreditsBalanceByAuth.Store(authID, antigravityCreditsBalance{
CreditAmount: 0,
MinCreditAmount: 1,
Known: true,
})
cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{
Known: true,
Available: false,
CreditAmount: 0,
MinCreditAmount: 1,
UpdatedAt: time.Now(),
})
}
func clearAntigravityCreditsPermanentlyDisabled(auth *cliproxyauth.Auth) {
if auth == nil || strings.TrimSpace(auth.ID) == "" {
return
}
antigravityCreditsFailureByAuth.Delete(strings.TrimSpace(auth.ID))
}
func antigravityHasExplicitCreditsBalanceExhaustedReason(body []byte) bool {
if len(body) == 0 {
return false
}
details := gjson.GetBytes(body, "error.details")
if !details.Exists() || !details.IsArray() {
return false
}
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, "INSUFFICIENT_G1_CREDITS_BALANCE") {
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
}
// 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" {
return resp, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"}
}
baseModel := thinking.ParseSuffix(req.Model).ModelName
if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) {
log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining)
d := remaining
return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
}
isClaude := strings.Contains(strings.ToLower(baseModel), "claude")
if isClaude || strings.Contains(baseModel, "gemini-3-pro") || strings.Contains(baseModel, "gemini-3.1-flash-image") {
return e.executeClaudeNonStream(ctx, auth, req, opts)
}
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
defer reporter.TrackFailure(ctx, &err)
from := opts.SourceFormat
to := sdktranslator.FromString("antigravity")
originalPayloadSource := req.Payload
if len(opts.OriginalRequest) > 0 {
originalPayloadSource = opts.OriginalRequest
}
originalPayload := originalPayloadSource
originalPayload, errValidate := validateAntigravityRequestSignatures(from, originalPayload)
if errValidate != nil {
return resp, errValidate
}
req.Payload = originalPayload
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
if errToken != nil {
return resp, errToken
}
if updatedAuth != nil {
auth = updatedAuth
}
originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, false)
translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false)
translated, err = thinking.ApplyThinking(translated, req.Model, from.String(), to.String(), e.Identifier())
if err != nil {
return resp, err
}
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
requestPath := helps.PayloadRequestPath(opts)
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel, requestPath)
useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg)
baseURLs := antigravityBaseURLFallbackOrder(auth)
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
attempts := antigravityRetryAttempts(auth, e.cfg)
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)
}
}
httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, false, opts.Alt, baseURL)
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 {
markAntigravityShortCooldown(auth, baseModel, time.Now(), *decision.retryAfter)
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
}
}
err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes)
return resp, err
}
// Success
if useCredits {
clearAntigravityCreditsFailureState(auth)
}
reporter.Publish(ctx, helps.ParseAntigravityUsage(bodyBytes))
var param any
converted := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, bodyBytes, &param)
resp = cliproxyexecutor.Response{Payload: converted, Headers: httpResp.Header.Clone()}
reporter.EnsurePublished(ctx)
return resp, 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"}
}
return resp, err
}
return resp, err
}
// executeClaudeNonStream performs a claude non-streaming request to the Antigravity API.
func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {
baseModel := thinking.ParseSuffix(req.Model).ModelName
if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) {
log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining)
d := remaining
return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
}
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
defer reporter.TrackFailure(ctx, &err)
from := opts.SourceFormat
to := sdktranslator.FromString("antigravity")
originalPayloadSource := req.Payload
if len(opts.OriginalRequest) > 0 {
originalPayloadSource = opts.OriginalRequest
}
originalPayload := originalPayloadSource
originalPayload, errValidate := validateAntigravityRequestSignatures(from, originalPayload)
if errValidate != nil {
return resp, errValidate
}
req.Payload = originalPayload
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
if errToken != nil {
return resp, errToken
}
if updatedAuth != nil {
auth = updatedAuth
}
originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true)
translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, true)
translated, err = thinking.ApplyThinking(translated, req.Model, from.String(), to.String(), e.Identifier())
if err != nil {
return resp, err
}
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
requestPath := helps.PayloadRequestPath(opts)
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel, requestPath)
useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg)
baseURLs := antigravityBaseURLFallbackOrder(auth)
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
attempts := antigravityRetryAttempts(auth, e.cfg)
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)
}
}
httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL)
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 {
markAntigravityShortCooldown(auth, baseModel, time.Now(), *decision.retryAfter)
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
}
}
err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes)
return resp, err
}
// Stream success
if useCredits {
clearAntigravityCreditsFailureState(auth)
}
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)
// 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)
}
select {
case out <- cliproxyexecutor.StreamChunk{Payload: payload}:
case <-ctx.Done():
return
}
}
if errScan := scanner.Err(); errScan != nil {
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done():
}
} else {
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())}
reporter.Publish(ctx, helps.ParseAntigravityUsage(resp.Payload))
var param any
converted := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, resp.Payload, &param)
resp = cliproxyexecutor.Response{Payload: converted, Headers: httpResp.Header.Clone()}
reporter.EnsurePublished(ctx)
return resp, 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"}
}
return resp, err
}
return resp, err
}
func (e *AntigravityExecutor) convertStreamToNonStream(stream []byte) []byte {
responseTemplate := ""
var traceID string
var finishReason string
var modelVersion string
var responseID string
var role string
var usageRaw string
parts := make([]map[string]interface{}, 0)
var pendingKind string
var pendingText strings.Builder
var pendingThoughtSig string
flushPending := func() {
if pendingKind == "" {
return
}
text := pendingText.String()
switch pendingKind {
case "text":
if strings.TrimSpace(text) == "" {
pendingKind = ""
pendingText.Reset()
pendingThoughtSig = ""
return
}
parts = append(parts, map[string]interface{}{"text": text})
case "thought":
if strings.TrimSpace(text) == "" && pendingThoughtSig == "" {
pendingKind = ""
pendingText.Reset()
pendingThoughtSig = ""
return
}
part := map[string]interface{}{"thought": true}
part["text"] = text
if pendingThoughtSig != "" {
part["thoughtSignature"] = pendingThoughtSig
}
parts = append(parts, part)
}
pendingKind = ""
pendingText.Reset()
pendingThoughtSig = ""
}
normalizePart := func(partResult gjson.Result) map[string]interface{} {
var m map[string]interface{}
_ = json.Unmarshal([]byte(partResult.Raw), &m)
if m == nil {
m = map[string]interface{}{}
}
sig := partResult.Get("thoughtSignature").String()
if sig == "" {
sig = partResult.Get("thought_signature").String()
}
if sig != "" {
m["thoughtSignature"] = sig
delete(m, "thought_signature")
}
if inlineData, ok := m["inline_data"]; ok {
m["inlineData"] = inlineData
delete(m, "inline_data")
}
return m
}
for _, line := range bytes.Split(stream, []byte("\n")) {
trimmed := bytes.TrimSpace(line)
if len(trimmed) == 0 || !gjson.ValidBytes(trimmed) {
continue
}
root := gjson.ParseBytes(trimmed)
responseNode := root.Get("response")
if !responseNode.Exists() {
if root.Get("candidates").Exists() {
responseNode = root
} else {
continue
}
}
responseTemplate = responseNode.Raw
if traceResult := root.Get("traceId"); traceResult.Exists() && traceResult.String() != "" {
traceID = traceResult.String()
}
if roleResult := responseNode.Get("candidates.0.content.role"); roleResult.Exists() {
role = roleResult.String()
}
if finishResult := responseNode.Get("candidates.0.finishReason"); finishResult.Exists() && finishResult.String() != "" {
finishReason = finishResult.String()
}
if modelResult := responseNode.Get("modelVersion"); modelResult.Exists() && modelResult.String() != "" {
modelVersion = modelResult.String()
}
if responseIDResult := responseNode.Get("responseId"); responseIDResult.Exists() && responseIDResult.String() != "" {
responseID = responseIDResult.String()
}
if usageResult := responseNode.Get("usageMetadata"); usageResult.Exists() {
usageRaw = usageResult.Raw
} else if usageMetadataResult := root.Get("usageMetadata"); usageMetadataResult.Exists() {
usageRaw = usageMetadataResult.Raw
}
if partsResult := responseNode.Get("candidates.0.content.parts"); partsResult.IsArray() {
for _, part := range partsResult.Array() {
hasFunctionCall := part.Get("functionCall").Exists()
hasInlineData := part.Get("inlineData").Exists() || part.Get("inline_data").Exists()
sig := part.Get("thoughtSignature").String()
if sig == "" {
sig = part.Get("thought_signature").String()
}
text := part.Get("text").String()
thought := part.Get("thought").Bool()
if hasFunctionCall || hasInlineData {
flushPending()
parts = append(parts, normalizePart(part))
continue
}
if thought || part.Get("text").Exists() {
kind := "text"
if thought {
kind = "thought"
}
if pendingKind != "" && pendingKind != kind {
flushPending()
}
pendingKind = kind
pendingText.WriteString(text)
if kind == "thought" && sig != "" {
pendingThoughtSig = sig
}
continue
}
flushPending()
parts = append(parts, normalizePart(part))
}
}
}
flushPending()
if responseTemplate == "" {
responseTemplate = `{"candidates":[{"content":{"role":"model","parts":[]}}]}`
}
partsJSON, _ := json.Marshal(parts)
updatedTemplate, _ := sjson.SetRawBytes([]byte(responseTemplate), "candidates.0.content.parts", partsJSON)
responseTemplate = string(updatedTemplate)
if role != "" {
updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "candidates.0.content.role", role)
responseTemplate = string(updatedTemplate)
}
if finishReason != "" {
updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "candidates.0.finishReason", finishReason)
responseTemplate = string(updatedTemplate)
}
if modelVersion != "" {
updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "modelVersion", modelVersion)
responseTemplate = string(updatedTemplate)
}
if responseID != "" {
updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "responseId", responseID)
responseTemplate = string(updatedTemplate)
}
if usageRaw != "" {
updatedTemplate, _ = sjson.SetRawBytes([]byte(responseTemplate), "usageMetadata", []byte(usageRaw))
responseTemplate = string(updatedTemplate)
} else if !gjson.Get(responseTemplate, "usageMetadata").Exists() {
updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.promptTokenCount", 0)
responseTemplate = string(updatedTemplate)
updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.candidatesTokenCount", 0)
responseTemplate = string(updatedTemplate)
updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.totalTokenCount", 0)
responseTemplate = string(updatedTemplate)
}
output := `{"response":{},"traceId":""}`
updatedOutput, _ := sjson.SetRawBytes([]byte(output), "response", []byte(responseTemplate))
output = string(updatedOutput)
if traceID != "" {
updatedOutput, _ = sjson.SetBytes([]byte(output), "traceId", traceID)
output = string(updatedOutput)
}
return []byte(output)
}
// ExecuteStream performs a streaming request to the Antigravity API.
func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) {
if opts.Alt == "responses/compact" {
return nil, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"}
}
baseModel := thinking.ParseSuffix(req.Model).ModelName
ctx = context.WithValue(ctx, "alt", "")
if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) {
log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining)
d := remaining
return nil, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
}
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
defer reporter.TrackFailure(ctx, &err)
from := opts.SourceFormat
to := sdktranslator.FromString("antigravity")
originalPayloadSource := req.Payload
if len(opts.OriginalRequest) > 0 {
originalPayloadSource = opts.OriginalRequest
}
originalPayload := originalPayloadSource
originalPayload, errValidate := validateAntigravityRequestSignatures(from, originalPayload)
if errValidate != nil {
return nil, errValidate
}
req.Payload = originalPayload
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
if errToken != nil {
return nil, errToken
}
if updatedAuth != nil {
auth = updatedAuth
}
originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true)
translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, true)
translated, err = thinking.ApplyThinking(translated, req.Model, from.String(), to.String(), e.Identifier())
if err != nil {
return nil, err
}
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
requestPath := helps.PayloadRequestPath(opts)
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel, requestPath)
useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg)
baseURLs := antigravityBaseURLFallbackOrder(auth)
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
attempts := antigravityRetryAttempts(auth, e.cfg)
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)
}
}
httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL)
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
}
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 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
}
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
}
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 {
markAntigravityShortCooldown(auth, baseModel, time.Now(), *decision.retryAfter)
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
}
}
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
}
}
err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes)
return nil, err
}
// Stream success
if useCredits {
clearAntigravityCreditsFailureState(auth)
}
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)
var param any
for scanner.Scan() {
line := scanner.Bytes()
helps.AppendAPIResponseChunk(ctx, e.cfg, 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)
}
chunks := sdktranslator.TranslateStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, bytes.Clone(payload), &param)
for i := range chunks {
select {
case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}:
case <-ctx.Done():
return
}
}
}
tail := sdktranslator.TranslateStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, []byte("[DONE]"), &param)
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)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errScan}:
case <-ctx.Done():
}
} else {
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"}
}
return nil, err
}
return nil, err
}
// Refresh refreshes the authentication credentials using the refresh token.
func (e *AntigravityExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
if auth == nil {
return auth, nil
}
updated, errRefresh := e.refreshToken(ctx, auth.Clone())
if errRefresh != nil {
return nil, errRefresh
}
return updated, nil
}
// CountTokens counts tokens for the given request using the Antigravity API.
func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
baseModel := thinking.ParseSuffix(req.Model).ModelName
from := opts.SourceFormat
to := sdktranslator.FromString("antigravity")
respCtx := context.WithValue(ctx, "alt", opts.Alt)
originalPayloadSource := req.Payload
if len(opts.OriginalRequest) > 0 {
originalPayloadSource = opts.OriginalRequest
}
originalPayloadSource, errValidate := validateAntigravityRequestSignatures(from, originalPayloadSource)
if errValidate != nil {
return cliproxyexecutor.Response{}, errValidate
}
req.Payload = originalPayloadSource
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
if errToken != nil {
return cliproxyexecutor.Response{}, errToken
}
if updatedAuth != nil {
auth = updatedAuth
}
if strings.TrimSpace(token) == "" {
return cliproxyexecutor.Response{}, statusErr{code: http.StatusUnauthorized, msg: "missing access token"}
}
// Prepare payload once (doesn't depend on baseURL)
payload := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false)
payload, err := thinking.ApplyThinking(payload, req.Model, from.String(), to.String(), e.Identifier())
if err != nil {
return cliproxyexecutor.Response{}, err
}
payload = deleteJSONField(payload, "project")
payload = deleteJSONField(payload, "model")
payload = deleteJSONField(payload, "request.safetySettings")
baseURLs := antigravityBaseURLFallbackOrder(auth)
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
var authID, authLabel, authType, authValue string
if auth != nil {
authID = auth.ID
authLabel = auth.Label
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
}
httpReq.Close = true
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, from, 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 := parseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil {
sErr.retryAfter = retryAfter
}
}
return cliproxyexecutor.Response{}, sErr
}
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
}
}
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"}
}
}
func (e *AntigravityExecutor) ensureAccessToken(ctx context.Context, auth *cliproxyauth.Auth) (string, *cliproxyauth.Auth, error) {
if auth == nil {
return "", nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"}
}
accessToken := metaStringValue(auth.Metadata, "access_token")
expiry := tokenExpiry(auth.Metadata)
if accessToken != "" && expiry.After(time.Now().Add(refreshSkew)) {
e.maybeRefreshAntigravityCreditsHint(ctx, auth, accessToken)
return accessToken, nil, nil
}
refreshCtx := context.Background()
if ctx != nil {
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt)
}
}
updated, errRefresh := e.refreshToken(refreshCtx, auth.Clone())
if errRefresh != nil {
return "", nil, errRefresh
}
return metaStringValue(updated.Metadata, "access_token"), updated, nil
}
func (e *AntigravityExecutor) maybeRefreshAntigravityCreditsHint(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) {
if e == nil || auth == nil || !antigravityCreditsRetryEnabled(e.cfg) {
return
}
if ctx != nil && ctx.Err() != nil {
return
}
authID := strings.TrimSpace(auth.ID)
if authID == "" {
return
}
if hint, ok := cliproxyauth.GetAntigravityCreditsHint(authID); ok && hint.Known {
return
}
if strings.TrimSpace(accessToken) == "" {
accessToken = metaStringValue(auth.Metadata, "access_token")
}
if strings.TrimSpace(accessToken) == "" {
return
}
state := &antigravityCreditsHintRefreshState{}
if existing, loaded := antigravityCreditsHintRefreshByID.LoadOrStore(authID, state); loaded {
if cast, ok := existing.(*antigravityCreditsHintRefreshState); ok && cast != nil {
state = cast
} else {
antigravityCreditsHintRefreshByID.Delete(authID)
antigravityCreditsHintRefreshByID.Store(authID, state)
}
}
now := time.Now()
if !state.mu.TryLock() {
return
}
if !state.lastAttempt.IsZero() && now.Sub(state.lastAttempt) < antigravityCreditsHintRefreshInterval {
state.mu.Unlock()
return
}
state.lastAttempt = now
refreshCtx := context.Background()
if ctx != nil {
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt)
}
}
refreshCtx, cancel := context.WithTimeout(refreshCtx, antigravityCreditsHintRefreshTimeout)
authCopy := auth.Clone()
go func(state *antigravityCreditsHintRefreshState, auth *cliproxyauth.Auth, token string) {
defer cancel()
defer state.mu.Unlock()
e.updateAntigravityCreditsBalance(refreshCtx, auth, token)
}(state, authCopy, accessToken)
}
func (e *AntigravityExecutor) refreshToken(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
if auth == nil {
return nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"}
}
refreshToken := metaStringValue(auth.Metadata, "refresh_token")
if refreshToken == "" {
return auth, statusErr{code: http.StatusUnauthorized, msg: "missing refresh token"}
}
form := url.Values{}
form.Set("client_id", antigravityClientID)
form.Set("client_secret", antigravityClientSecret)
form.Set("grant_type", "refresh_token")
form.Set("refresh_token", refreshToken)
httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, "https://oauth2.googleapis.com/token", strings.NewReader(form.Encode()))
if errReq != nil {
return auth, errReq
}
httpReq.Header.Set("Host", "oauth2.googleapis.com")
httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
// Real Antigravity uses Go's default User-Agent for OAuth token refresh
httpReq.Header.Set("User-Agent", "Go-http-client/2.0")
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
httpResp, errDo := httpClient.Do(httpReq)
if errDo != nil {
return auth, errDo
}
defer func() {
if errClose := httpResp.Body.Close(); errClose != nil {
log.Errorf("antigravity executor: close response body error: %v", errClose)
}
}()
bodyBytes, errRead := io.ReadAll(httpResp.Body)
if errRead != nil {
return auth, errRead
}
if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices {
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
}
}
return auth, sErr
}
var tokenResp struct {
AccessToken string `json:"access_token"`
RefreshToken string `json:"refresh_token"`
ExpiresIn int64 `json:"expires_in"`
TokenType string `json:"token_type"`
}
if errUnmarshal := json.Unmarshal(bodyBytes, &tokenResp); errUnmarshal != nil {
return auth, errUnmarshal
}
if auth.Metadata == nil {
auth.Metadata = make(map[string]any)
}
auth.Metadata["access_token"] = tokenResp.AccessToken
if tokenResp.RefreshToken != "" {
auth.Metadata["refresh_token"] = tokenResp.RefreshToken
}
auth.Metadata["expires_in"] = tokenResp.ExpiresIn
now := time.Now()
auth.Metadata["timestamp"] = now.UnixMilli()
auth.Metadata["expired"] = now.Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339)
auth.Metadata["type"] = antigravityAuthType
if errProject := e.ensureAntigravityProjectID(ctx, auth, tokenResp.AccessToken); errProject != nil {
log.Warnf("antigravity executor: ensure project id failed: %v", errProject)
}
e.updateAntigravityCreditsBalance(ctx, auth, tokenResp.AccessToken)
return auth, nil
}
func (e *AntigravityExecutor) ensureAntigravityProjectID(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) error {
if auth == nil {
return nil
}
if auth.Metadata["project_id"] != nil {
return nil
}
token := strings.TrimSpace(accessToken)
if token == "" {
token = metaStringValue(auth.Metadata, "access_token")
}
if token == "" {
return nil
}
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
projectID, errFetch := sdkAuth.FetchAntigravityProjectID(ctx, token, httpClient)
if errFetch != nil {
return errFetch
}
if strings.TrimSpace(projectID) == "" {
return nil
}
if auth.Metadata == nil {
auth.Metadata = make(map[string]any)
}
auth.Metadata["project_id"] = strings.TrimSpace(projectID)
return nil
}
func (e *AntigravityExecutor) updateAntigravityCreditsBalance(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) {
if auth == nil || strings.TrimSpace(auth.ID) == "" {
return
}
token := strings.TrimSpace(accessToken)
if token == "" {
token = metaStringValue(auth.Metadata, "access_token")
}
if token == "" {
return
}
userAgent := resolveLoadCodeAssistUserAgent(auth)
loadReqBody, errMarshal := json.Marshal(map[string]any{
"metadata": map[string]string{
"ide_type": "ANTIGRAVITY",
"ide_version": misc.AntigravityVersionFromUserAgent(userAgent),
"ide_name": "antigravity",
},
})
if errMarshal != nil {
log.Debugf("antigravity executor: marshal loadCodeAssist request error: %v", errMarshal)
return
}
baseURL := buildBaseURL(auth)
endpointURL := strings.TrimSuffix(baseURL, "/") + "/v1internal:loadCodeAssist"
httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, bytes.NewReader(loadReqBody))
if errReq != nil {
log.Debugf("antigravity executor: create loadCodeAssist request error: %v", errReq)
return
}
httpReq.Header.Set("Authorization", "Bearer "+token)
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("User-Agent", userAgent)
httpReq.Header.Set("X-Goog-Api-Client", misc.AntigravityGoogAPIClientUA)
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
httpResp, errDo := httpClient.Do(httpReq)
if errDo != nil {
log.Debugf("antigravity executor: loadCodeAssist request error: %v", errDo)
return
}
defer func() {
if errClose := httpResp.Body.Close(); errClose != nil {
log.Errorf("antigravity executor: close loadCodeAssist response body error: %v", errClose)
}
}()
bodyBytes, errRead := io.ReadAll(httpResp.Body)
if errRead != nil || httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices {
log.Debugf("antigravity executor: loadCodeAssist returned status %d, err=%v", httpResp.StatusCode, errRead)
return
}
authID := strings.TrimSpace(auth.ID)
paidTierID := strings.TrimSpace(gjson.GetBytes(bodyBytes, "paidTier.id").String())
credits := gjson.GetBytes(bodyBytes, "paidTier.availableCredits")
if !credits.IsArray() {
cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{
Known: true,
Available: false,
PaidTierID: paidTierID,
UpdatedAt: time.Now(),
})
return
}
for _, credit := range credits.Array() {
if !strings.EqualFold(credit.Get("creditType").String(), "GOOGLE_ONE_AI") {
continue
}
creditAmount, errCA := strconv.ParseFloat(strings.TrimSpace(credit.Get("creditAmount").String()), 64)
if errCA != nil {
continue
}
minAmount, errMA := strconv.ParseFloat(strings.TrimSpace(credit.Get("minimumCreditAmountForUsage").String()), 64)
if errMA != nil {
continue
}
bal := antigravityCreditsBalance{
CreditAmount: creditAmount,
MinCreditAmount: minAmount,
PaidTierID: paidTierID,
Known: true,
}
antigravityCreditsBalanceByAuth.Store(authID, bal)
cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{
Known: true,
Available: creditAmount >= minAmount,
CreditAmount: creditAmount,
MinCreditAmount: minAmount,
PaidTierID: paidTierID,
UpdatedAt: time.Now(),
})
if creditAmount >= minAmount {
clearAntigravityCreditsPermanentlyDisabled(auth)
}
return
}
}
func (e *AntigravityExecutor) buildRequest(ctx context.Context, auth *cliproxyauth.Auth, token, modelName string, payload []byte, stream bool, alt, baseURL string) (*http.Request, error) {
if token == "" {
return nil, statusErr{code: http.StatusUnauthorized, msg: "missing access token"}
}
base := strings.TrimSuffix(baseURL, "/")
if base == "" {
base = buildBaseURL(auth)
}
path := antigravityGeneratePath
if stream {
path = antigravityStreamPath
}
var requestURL strings.Builder
requestURL.WriteString(base)
requestURL.WriteString(path)
if stream {
if alt != "" {
requestURL.WriteString("?$alt=")
requestURL.WriteString(url.QueryEscape(alt))
} else {
requestURL.WriteString("?alt=sse")
}
} else if alt != "" {
requestURL.WriteString("?$alt=")
requestURL.WriteString(url.QueryEscape(alt))
}
// Extract project_id from auth metadata if available
projectID := ""
if auth != nil && auth.Metadata != nil {
if pid, ok := auth.Metadata["project_id"].(string); ok {
projectID = strings.TrimSpace(pid)
}
}
payload = geminiToAntigravity(modelName, payload, projectID)
payload, _ = sjson.SetBytes(payload, "model", modelName)
// Cap maxOutputTokens to model's max_completion_tokens from registry
if maxOut := gjson.GetBytes(payload, "request.generationConfig.maxOutputTokens"); maxOut.Exists() && maxOut.Type == gjson.Number {
if modelInfo := registry.LookupModelInfo(modelName, "antigravity"); modelInfo != nil && modelInfo.MaxCompletionTokens > 0 {
if int(maxOut.Int()) > modelInfo.MaxCompletionTokens {
payload, _ = sjson.SetBytes(payload, "request.generationConfig.maxOutputTokens", modelInfo.MaxCompletionTokens)
}
}
}
useAntigravitySchema := strings.Contains(modelName, "claude") || strings.Contains(modelName, "gemini-3-pro") || strings.Contains(modelName, "gemini-3.1-pro")
var (
bodyReader io.Reader
payloadLog []byte
)
if antigravityRequestNeedsSchemaSanitization(payload) {
payloadStr := string(payload)
paths := make([]string, 0)
util.Walk(gjson.Parse(payloadStr), "", "parametersJsonSchema", &paths)
for _, p := range paths {
payloadStr, _ = util.RenameKey(payloadStr, p, p[:len(p)-len("parametersJsonSchema")]+"parameters")
}
if useAntigravitySchema {
payloadStr = util.CleanJSONSchemaForAntigravity(payloadStr)
} else {
payloadStr = util.CleanJSONSchemaForGemini(payloadStr)
}
if strings.Contains(modelName, "claude") {
updated, _ := sjson.SetBytes([]byte(payloadStr), "request.toolConfig.functionCallingConfig.mode", "VALIDATED")
payloadStr = string(updated)
} else {
payloadStr, _ = sjson.Delete(payloadStr, "request.generationConfig.maxOutputTokens")
}
bodyReader = strings.NewReader(payloadStr)
if e.cfg != nil && e.cfg.RequestLog {
payloadLog = []byte(payloadStr)
}
} else {
if strings.Contains(modelName, "claude") {
payload, _ = sjson.SetBytes(payload, "request.toolConfig.functionCallingConfig.mode", "VALIDATED")
} else {
payload, _ = sjson.DeleteBytes(payload, "request.generationConfig.maxOutputTokens")
}
bodyReader = bytes.NewReader(payload)
if e.cfg != nil && e.cfg.RequestLog {
payloadLog = append([]byte(nil), payload...)
}
}
// if useAntigravitySchema {
// systemInstructionPartsResult := gjson.Get(payloadStr, "request.systemInstruction.parts")
// payloadStr, _ = sjson.SetBytes([]byte(payloadStr), "request.systemInstruction.role", "user")
// payloadStr, _ = sjson.SetBytes([]byte(payloadStr), "request.systemInstruction.parts.0.text", systemInstruction)
// payloadStr, _ = sjson.SetBytes([]byte(payloadStr), "request.systemInstruction.parts.1.text", fmt.Sprintf("Please ignore following [ignore]%s[/ignore]", systemInstruction))
// if systemInstructionPartsResult.Exists() && systemInstructionPartsResult.IsArray() {
// for _, partResult := range systemInstructionPartsResult.Array() {
// payloadStr, _ = sjson.SetRawBytes([]byte(payloadStr), "request.systemInstruction.parts.-1", []byte(partResult.Raw))
// }
// }
// }
httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bodyReader)
if errReq != nil {
return nil, errReq
}
httpReq.Close = true
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)
var authID, authLabel, authType, authValue string
if auth != nil {
authID = auth.ID
authLabel = auth.Label
authType, authValue = auth.AccountInfo()
}
helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{
URL: requestURL.String(),
Method: http.MethodPost,
Headers: httpReq.Header.Clone(),
Body: payloadLog,
Provider: e.Identifier(),
AuthID: authID,
AuthLabel: authLabel,
AuthType: authType,
AuthValue: authValue,
})
return httpReq, nil
}
func antigravityRequestNeedsSchemaSanitization(payload []byte) bool {
if gjson.GetBytes(payload, "request.tools.0").Exists() {
return true
}
if gjson.GetBytes(payload, "request.generationConfig.responseJsonSchema").Exists() {
return true
}
if gjson.GetBytes(payload, "request.generationConfig.responseSchema").Exists() {
return true
}
return false
}
func tokenExpiry(metadata map[string]any) time.Time {
if metadata == nil {
return time.Time{}
}
if expStr, ok := metadata["expired"].(string); ok {
expStr = strings.TrimSpace(expStr)
if expStr != "" {
if parsed, errParse := time.Parse(time.RFC3339, expStr); errParse == nil {
return parsed
}
}
}
expiresIn, hasExpires := int64Value(metadata["expires_in"])
tsMs, hasTimestamp := int64Value(metadata["timestamp"])
if hasExpires && hasTimestamp {
return time.Unix(0, tsMs*int64(time.Millisecond)).Add(time.Duration(expiresIn) * time.Second)
}
return time.Time{}
}
func metaStringValue(metadata map[string]any, key string) string {
if metadata == nil {
return ""
}
if v, ok := metadata[key]; ok {
switch typed := v.(type) {
case string:
return strings.TrimSpace(typed)
case []byte:
return strings.TrimSpace(string(typed))
}
}
return ""
}
func int64Value(value any) (int64, bool) {
switch typed := value.(type) {
case int:
return int64(typed), true
case int64:
return typed, true
case float64:
return int64(typed), true
case json.Number:
if i, errParse := typed.Int64(); errParse == nil {
return i, true
}
case string:
if strings.TrimSpace(typed) == "" {
return 0, false
}
if i, errParse := strconv.ParseInt(strings.TrimSpace(typed), 10, 64); errParse == nil {
return i, true
}
}
return 0, false
}
func buildBaseURL(auth *cliproxyauth.Auth) string {
if baseURLs := antigravityBaseURLFallbackOrder(auth); len(baseURLs) > 0 {
return baseURLs[0]
}
return antigravityBaseURLDaily
}
func resolveHost(base string) string {
parsed, errParse := url.Parse(base)
if errParse != nil {
return ""
}
if parsed.Host != "" {
return parsed.Host
}
return strings.TrimPrefix(strings.TrimPrefix(base, "https://"), "http://")
}
func resolveUserAgent(auth *cliproxyauth.Auth) string {
return misc.AntigravityRequestUserAgent(antigravityConfiguredUserAgent(auth))
}
func resolveLoadCodeAssistUserAgent(auth *cliproxyauth.Auth) string {
return misc.AntigravityLoadCodeAssistUserAgent(antigravityConfiguredUserAgent(auth))
}
func antigravityConfiguredUserAgent(auth *cliproxyauth.Auth) string {
raw := ""
if auth != nil {
if auth.Attributes != nil {
if ua := strings.TrimSpace(auth.Attributes["user_agent"]); ua != "" {
raw = ua
}
}
if raw == "" && auth.Metadata != nil {
if ua, ok := auth.Metadata["user_agent"].(string); ok && strings.TrimSpace(ua) != "" {
raw = strings.TrimSpace(ua)
}
}
}
return raw
}
func antigravityRetryAttempts(auth *cliproxyauth.Auth, cfg *config.Config) int {
retry := 0
if cfg != nil {
retry = cfg.RequestRetry
}
if auth != nil {
if override, ok := auth.RequestRetryOverride(); ok {
retry = override
}
}
if retry < 0 {
retry = 0
}
attempts := retry + 1
if attempts < 1 {
return 1
}
return attempts
}
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 ""
}
authID := strings.TrimSpace(auth.ID)
modelName = strings.TrimSpace(modelName)
if authID == "" || modelName == "" {
return ""
}
return authID + "|" + modelName + "|sc"
}
func antigravityIsInShortCooldown(auth *cliproxyauth.Auth, modelName string, now time.Time) (bool, time.Duration) {
key := antigravityShortCooldownKey(auth, modelName)
if key == "" {
return false, 0
}
value, ok := antigravityShortCooldownByAuth.Load(key)
if !ok {
return false, 0
}
until, ok := value.(time.Time)
if !ok || until.IsZero() {
antigravityShortCooldownByAuth.Delete(key)
return false, 0
}
remaining := until.Sub(now)
if remaining <= 0 {
antigravityShortCooldownByAuth.Delete(key)
return false, 0
}
return true, remaining
}
func markAntigravityShortCooldown(auth *cliproxyauth.Auth, modelName string, now time.Time, duration time.Duration) {
key := antigravityShortCooldownKey(auth, modelName)
if key == "" {
return
}
antigravityShortCooldownByAuth.Store(key, now.Add(duration))
}
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
}
}
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 ""
}
if auth.Attributes != nil {
if v := strings.TrimSpace(auth.Attributes["base_url"]); v != "" {
return strings.TrimSuffix(v, "/")
}
}
if auth.Metadata != nil {
if v, ok := auth.Metadata["base_url"].(string); ok {
v = strings.TrimSpace(v)
if v != "" {
return strings.TrimSuffix(v, "/")
}
}
}
return ""
}
func geminiToAntigravity(modelName string, payload []byte, projectID string) []byte {
template := payload
template, _ = sjson.SetBytes(template, "model", modelName)
template, _ = sjson.SetBytes(template, "userAgent", "antigravity")
isImageModel := strings.Contains(modelName, "image")
var reqType string
if isImageModel {
reqType = "image_gen"
} else {
reqType = "agent"
}
template, _ = sjson.SetBytes(template, "requestType", reqType)
// Use real project ID from auth if available, otherwise generate random (legacy fallback)
if projectID != "" {
template, _ = sjson.SetBytes(template, "project", projectID)
} else {
template, _ = sjson.SetBytes(template, "project", generateProjectID())
}
if isImageModel {
template, _ = sjson.SetBytes(template, "requestId", generateImageGenRequestID())
} else {
template, _ = sjson.SetBytes(template, "requestId", generateRequestID())
template, _ = sjson.SetBytes(template, "request.sessionId", generateStableSessionID(payload))
}
template, _ = sjson.DeleteBytes(template, "request.safetySettings")
if toolConfig := gjson.GetBytes(template, "toolConfig"); toolConfig.Exists() && !gjson.GetBytes(template, "request.toolConfig").Exists() {
template, _ = sjson.SetRawBytes(template, "request.toolConfig", []byte(toolConfig.Raw))
template, _ = sjson.DeleteBytes(template, "toolConfig")
}
return template
}
func generateRequestID() string {
return "agent-" + uuid.NewString()
}
func generateImageGenRequestID() string {
return fmt.Sprintf("image_gen/%d/%s/12", time.Now().UnixMilli(), uuid.NewString())
}
func generateSessionID() string {
randSourceMutex.Lock()
n := randSource.Int63n(9_000_000_000_000_000_000)
randSourceMutex.Unlock()
return "-" + strconv.FormatInt(n, 10)
}
func generateStableSessionID(payload []byte) string {
contents := gjson.GetBytes(payload, "request.contents")
if contents.IsArray() {
for _, content := range contents.Array() {
if content.Get("role").String() == "user" {
text := content.Get("parts.0.text").String()
if text != "" {
h := sha256.Sum256([]byte(text))
n := int64(binary.BigEndian.Uint64(h[:8])) & 0x7FFFFFFFFFFFFFFF
return "-" + strconv.FormatInt(n, 10)
}
}
}
}
return generateSessionID()
}
func generateProjectID() string {
adjectives := []string{"useful", "bright", "swift", "calm", "bold"}
nouns := []string{"fuze", "wave", "spark", "flow", "core"}
randSourceMutex.Lock()
adj := adjectives[randSource.Intn(len(adjectives))]
noun := nouns[randSource.Intn(len(nouns))]
randSourceMutex.Unlock()
randomPart := strings.ToLower(uuid.NewString())[:5]
return adj + "-" + noun + "-" + randomPart
}