mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-05-12 09:05:57 +08:00
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)
2397 lines
80 KiB
Go
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, ¶m)
|
|
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, ¶m)
|
|
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), ¶m)
|
|
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]"), ¶m)
|
|
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
|
|
}
|