mirror of
https://github.com/router-for-me/CLIProxyAPIPlus.git
synced 2026-05-08 04:45:43 +08:00
278a89824c
Thinking blocks with empty signatures come from proxy-generated responses (Antigravity/Gemini routed as Claude). These should be silently dropped from the request payload before forwarding, not rejected with 400. Fixes 10 "missing thinking signature" errors.
2419 lines
81 KiB
Go
2419 lines
81 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/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
|
|
antigravityCreditsRetryTTL = 5 * time.Hour
|
|
antigravityCreditsAutoDisableDuration = 5 * time.Hour
|
|
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 {
|
|
Count int
|
|
DisabledUntil time.Time
|
|
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
|
|
antigravityPreferCreditsByModel sync.Map
|
|
antigravityShortCooldownByAuth sync.Map
|
|
antigravityQuotaExhaustedKeywords = []string{
|
|
"quota_exhausted",
|
|
"quota exhausted",
|
|
}
|
|
antigravityCreditsExhaustedKeywords = []string{
|
|
"google_one_ai",
|
|
"insufficient credit",
|
|
"insufficient credits",
|
|
"not enough credit",
|
|
"not enough credits",
|
|
"credit exhausted",
|
|
"credits exhausted",
|
|
"credit balance",
|
|
"minimumcreditamountforusage",
|
|
"minimum credit amount for usage",
|
|
"minimum credit",
|
|
"resource has been exhausted",
|
|
}
|
|
)
|
|
|
|
// 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 empty signatures (proxy-generated).
|
|
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
|
|
}
|
|
|
|
lowerBody := strings.ToLower(string(body))
|
|
for _, keyword := range antigravityQuotaExhaustedKeywords {
|
|
if strings.Contains(lowerBody, keyword) {
|
|
decision.kind = antigravity429DecisionFullQuotaExhausted
|
|
decision.reason = "quota_exhausted"
|
|
return decision
|
|
}
|
|
}
|
|
|
|
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() {
|
|
decision.kind = antigravity429DecisionSoftRetry
|
|
return decision
|
|
}
|
|
|
|
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
|
|
}
|
|
}
|
|
|
|
decision.kind = antigravity429DecisionSoftRetry
|
|
return decision
|
|
}
|
|
|
|
func antigravityHasQuotaResetDelayOrModelInfo(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
|
|
}
|
|
if strings.TrimSpace(detail.Get("metadata.quotaResetDelay").String()) != "" {
|
|
return true
|
|
}
|
|
if strings.TrimSpace(detail.Get("metadata.model").String()) != "" {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func antigravityCreditsRetryEnabled(cfg *config.Config) bool {
|
|
return cfg != nil && cfg.QuotaExceeded.AntigravityCredits
|
|
}
|
|
|
|
func antigravityCreditsFailureStateForAuth(auth *cliproxyauth.Auth) (string, antigravityCreditsFailureState, bool) {
|
|
if auth == nil || strings.TrimSpace(auth.ID) == "" {
|
|
return "", antigravityCreditsFailureState{}, false
|
|
}
|
|
authID := strings.TrimSpace(auth.ID)
|
|
value, ok := antigravityCreditsFailureByAuth.Load(authID)
|
|
if !ok {
|
|
return authID, antigravityCreditsFailureState{}, true
|
|
}
|
|
state, ok := value.(antigravityCreditsFailureState)
|
|
if !ok {
|
|
antigravityCreditsFailureByAuth.Delete(authID)
|
|
return authID, antigravityCreditsFailureState{}, true
|
|
}
|
|
return authID, state, true
|
|
}
|
|
|
|
func antigravityCreditsDisabled(auth *cliproxyauth.Auth, now time.Time) bool {
|
|
authID, state, ok := antigravityCreditsFailureStateForAuth(auth)
|
|
if !ok {
|
|
return false
|
|
}
|
|
if state.PermanentlyDisabled {
|
|
return true
|
|
}
|
|
if state.DisabledUntil.IsZero() {
|
|
return false
|
|
}
|
|
if state.DisabledUntil.After(now) {
|
|
return true
|
|
}
|
|
antigravityCreditsFailureByAuth.Delete(authID)
|
|
return false
|
|
}
|
|
|
|
func recordAntigravityCreditsFailure(auth *cliproxyauth.Auth, now time.Time) {
|
|
authID, state, ok := antigravityCreditsFailureStateForAuth(auth)
|
|
if !ok {
|
|
return
|
|
}
|
|
if state.PermanentlyDisabled {
|
|
antigravityCreditsFailureByAuth.Store(authID, state)
|
|
return
|
|
}
|
|
state.Count++
|
|
state.DisabledUntil = now.Add(antigravityCreditsAutoDisableDuration)
|
|
antigravityCreditsFailureByAuth.Store(authID, state)
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
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 antigravityPreferCreditsKey(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
|
|
}
|
|
|
|
func antigravityShouldPreferCredits(auth *cliproxyauth.Auth, modelName string, now time.Time) bool {
|
|
key := antigravityPreferCreditsKey(auth, modelName)
|
|
if key == "" {
|
|
return false
|
|
}
|
|
value, ok := antigravityPreferCreditsByModel.Load(key)
|
|
if !ok {
|
|
return false
|
|
}
|
|
until, ok := value.(time.Time)
|
|
if !ok || until.IsZero() {
|
|
antigravityPreferCreditsByModel.Delete(key)
|
|
return false
|
|
}
|
|
if !until.After(now) {
|
|
antigravityPreferCreditsByModel.Delete(key)
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
|
|
func markAntigravityPreferCredits(auth *cliproxyauth.Auth, modelName string, now time.Time, retryAfter *time.Duration) {
|
|
key := antigravityPreferCreditsKey(auth, modelName)
|
|
if key == "" {
|
|
return
|
|
}
|
|
until := now.Add(antigravityCreditsRetryTTL)
|
|
if retryAfter != nil && *retryAfter > 0 {
|
|
until = now.Add(*retryAfter)
|
|
}
|
|
antigravityPreferCreditsByModel.Store(key, until)
|
|
}
|
|
|
|
func clearAntigravityPreferCredits(auth *cliproxyauth.Auth, modelName string) {
|
|
key := antigravityPreferCreditsKey(auth, modelName)
|
|
if key == "" {
|
|
return
|
|
}
|
|
antigravityPreferCreditsByModel.Delete(key)
|
|
}
|
|
|
|
func shouldMarkAntigravityCreditsExhausted(statusCode int, body []byte, reqErr error) bool {
|
|
if reqErr != nil || statusCode == 0 {
|
|
return false
|
|
}
|
|
if statusCode >= http.StatusInternalServerError || statusCode == http.StatusRequestTimeout {
|
|
return false
|
|
}
|
|
lowerBody := strings.ToLower(string(body))
|
|
for _, keyword := range antigravityCreditsExhaustedKeywords {
|
|
if strings.Contains(lowerBody, keyword) {
|
|
if keyword == "resource has been exhausted" &&
|
|
statusCode == http.StatusTooManyRequests &&
|
|
decideAntigravity429(body).kind == antigravity429DecisionSoftRetry &&
|
|
!antigravityHasQuotaResetDelayOrModelInfo(body) {
|
|
return false
|
|
}
|
|
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
|
|
}
|
|
|
|
func (e *AntigravityExecutor) attemptCreditsFallback(
|
|
ctx context.Context,
|
|
auth *cliproxyauth.Auth,
|
|
httpClient *http.Client,
|
|
token string,
|
|
modelName string,
|
|
payload []byte,
|
|
stream bool,
|
|
alt string,
|
|
baseURL string,
|
|
originalBody []byte,
|
|
) (*http.Response, bool) {
|
|
if !antigravityCreditsRetryEnabled(e.cfg) {
|
|
return nil, false
|
|
}
|
|
if decideAntigravity429(originalBody).kind != antigravity429DecisionFullQuotaExhausted {
|
|
return nil, false
|
|
}
|
|
now := time.Now()
|
|
if shouldForcePermanentDisableCredits(originalBody) {
|
|
clearAntigravityPreferCredits(auth, modelName)
|
|
markAntigravityCreditsPermanentlyDisabled(auth)
|
|
return nil, false
|
|
}
|
|
|
|
if antigravityHasExplicitCreditsBalanceExhaustedReason(originalBody) {
|
|
clearAntigravityPreferCredits(auth, modelName)
|
|
markAntigravityCreditsPermanentlyDisabled(auth)
|
|
return nil, false
|
|
}
|
|
|
|
if antigravityCreditsDisabled(auth, now) {
|
|
return nil, false
|
|
}
|
|
creditsPayload := injectEnabledCreditTypes(payload)
|
|
if len(creditsPayload) == 0 {
|
|
return nil, false
|
|
}
|
|
|
|
httpReq, errReq := e.buildRequest(ctx, auth, token, modelName, creditsPayload, stream, alt, baseURL)
|
|
if errReq != nil {
|
|
helps.RecordAPIResponseError(ctx, e.cfg, errReq)
|
|
clearAntigravityPreferCredits(auth, modelName)
|
|
recordAntigravityCreditsFailure(auth, now)
|
|
return nil, true
|
|
}
|
|
httpResp, errDo := httpClient.Do(httpReq)
|
|
if errDo != nil {
|
|
helps.RecordAPIResponseError(ctx, e.cfg, errDo)
|
|
clearAntigravityPreferCredits(auth, modelName)
|
|
recordAntigravityCreditsFailure(auth, now)
|
|
return nil, true
|
|
}
|
|
if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices {
|
|
retryAfter, _ := parseRetryDelay(originalBody)
|
|
markAntigravityPreferCredits(auth, modelName, now, retryAfter)
|
|
clearAntigravityCreditsFailureState(auth)
|
|
return httpResp, true
|
|
}
|
|
|
|
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 credits fallback response body error: %v", errClose)
|
|
}
|
|
if errRead != nil {
|
|
helps.RecordAPIResponseError(ctx, e.cfg, errRead)
|
|
clearAntigravityPreferCredits(auth, modelName)
|
|
recordAntigravityCreditsFailure(auth, now)
|
|
return nil, true
|
|
}
|
|
helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes)
|
|
if shouldForcePermanentDisableCredits(bodyBytes) {
|
|
clearAntigravityPreferCredits(auth, modelName)
|
|
markAntigravityCreditsPermanentlyDisabled(auth)
|
|
return nil, true
|
|
}
|
|
|
|
if antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) {
|
|
clearAntigravityPreferCredits(auth, modelName)
|
|
markAntigravityCreditsPermanentlyDisabled(auth)
|
|
return nil, true
|
|
}
|
|
|
|
clearAntigravityPreferCredits(auth, modelName)
|
|
recordAntigravityCreditsFailure(auth, now)
|
|
return nil, true
|
|
}
|
|
|
|
func (e *AntigravityExecutor) handleDirectCreditsFailure(ctx context.Context, auth *cliproxyauth.Auth, modelName string, reqErr error) {
|
|
if reqErr != nil {
|
|
if shouldForcePermanentDisableCredits(reqErrBody(reqErr)) {
|
|
clearAntigravityPreferCredits(auth, modelName)
|
|
markAntigravityCreditsPermanentlyDisabled(auth)
|
|
return
|
|
}
|
|
|
|
if antigravityHasExplicitCreditsBalanceExhaustedReason(reqErrBody(reqErr)) {
|
|
clearAntigravityPreferCredits(auth, modelName)
|
|
markAntigravityCreditsPermanentlyDisabled(auth)
|
|
return
|
|
}
|
|
|
|
helps.RecordAPIResponseError(ctx, e.cfg, reqErr)
|
|
}
|
|
clearAntigravityPreferCredits(auth, modelName)
|
|
recordAntigravityCreditsFailure(auth, time.Now())
|
|
}
|
|
func reqErrBody(reqErr error) []byte {
|
|
if reqErr == nil {
|
|
return nil
|
|
}
|
|
msg := reqErr.Error()
|
|
if strings.TrimSpace(msg) == "" {
|
|
return nil
|
|
}
|
|
return []byte(msg)
|
|
}
|
|
|
|
func shouldForcePermanentDisableCredits(body []byte) bool {
|
|
return antigravityHasExplicitCreditsBalanceExhaustedReason(body)
|
|
}
|
|
|
|
// 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 {
|
|
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)
|
|
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel)
|
|
|
|
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
|
|
usedCreditsDirect := false
|
|
if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) {
|
|
if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 {
|
|
requestPayload = creditsPayload
|
|
usedCreditsDirect = true
|
|
}
|
|
}
|
|
|
|
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 and skipping credits fallback", *decision.retryAfter, baseModel)
|
|
}
|
|
case antigravity429DecisionFullQuotaExhausted:
|
|
if usedCreditsDirect {
|
|
clearAntigravityPreferCredits(auth, baseModel)
|
|
recordAntigravityCreditsFailure(auth, time.Now())
|
|
} else {
|
|
creditsResp, _ := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, false, opts.Alt, baseURL, bodyBytes)
|
|
if creditsResp != nil {
|
|
helps.RecordAPIResponseMetadata(ctx, e.cfg, creditsResp.StatusCode, creditsResp.Header.Clone())
|
|
creditsBody, errCreditsRead := io.ReadAll(creditsResp.Body)
|
|
if errClose := creditsResp.Body.Close(); errClose != nil {
|
|
log.Errorf("antigravity executor: close credits success response body error: %v", errClose)
|
|
}
|
|
if errCreditsRead != nil {
|
|
helps.RecordAPIResponseError(ctx, e.cfg, errCreditsRead)
|
|
err = errCreditsRead
|
|
return resp, err
|
|
}
|
|
helps.AppendAPIResponseChunk(ctx, e.cfg, creditsBody)
|
|
reporter.Publish(ctx, helps.ParseAntigravityUsage(creditsBody))
|
|
var param any
|
|
converted := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, creditsBody, ¶m)
|
|
resp = cliproxyexecutor.Response{Payload: converted, Headers: creditsResp.Header.Clone()}
|
|
reporter.EnsurePublished(ctx)
|
|
return resp, nil
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
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 {
|
|
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)
|
|
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel)
|
|
|
|
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
|
|
usedCreditsDirect := false
|
|
if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) {
|
|
if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 {
|
|
requestPayload = creditsPayload
|
|
usedCreditsDirect = true
|
|
}
|
|
}
|
|
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 and skipping credits fallback", *decision.retryAfter, baseModel)
|
|
}
|
|
case antigravity429DecisionFullQuotaExhausted:
|
|
if usedCreditsDirect {
|
|
clearAntigravityPreferCredits(auth, baseModel)
|
|
recordAntigravityCreditsFailure(auth, time.Now())
|
|
} else {
|
|
creditsResp, _ := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, true, opts.Alt, baseURL, bodyBytes)
|
|
if creditsResp != nil {
|
|
httpResp = creditsResp
|
|
helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices {
|
|
goto streamSuccessClaudeNonStream
|
|
}
|
|
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
|
|
}
|
|
|
|
streamSuccessClaudeNonStream:
|
|
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)
|
|
}
|
|
|
|
out <- cliproxyexecutor.StreamChunk{Payload: payload}
|
|
}
|
|
if errScan := scanner.Err(); errScan != nil {
|
|
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
|
|
reporter.PublishFailure(ctx)
|
|
out <- cliproxyexecutor.StreamChunk{Err: errScan}
|
|
} 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 {
|
|
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)
|
|
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel)
|
|
|
|
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
|
|
usedCreditsDirect := false
|
|
if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) {
|
|
if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 {
|
|
requestPayload = creditsPayload
|
|
usedCreditsDirect = true
|
|
}
|
|
}
|
|
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 cooldown and skipping credits fallback", *decision.retryAfter, baseModel)
|
|
}
|
|
case antigravity429DecisionFullQuotaExhausted:
|
|
if usedCreditsDirect {
|
|
clearAntigravityPreferCredits(auth, baseModel)
|
|
recordAntigravityCreditsFailure(auth, time.Now())
|
|
} else {
|
|
creditsResp, _ := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, true, opts.Alt, baseURL, bodyBytes)
|
|
if creditsResp != nil {
|
|
httpResp = creditsResp
|
|
helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices {
|
|
goto streamSuccessExecuteStream
|
|
}
|
|
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
|
|
}
|
|
|
|
streamSuccessExecuteStream:
|
|
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 {
|
|
out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}
|
|
}
|
|
}
|
|
tail := sdktranslator.TranslateStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, []byte("[DONE]"), ¶m)
|
|
for i := range tail {
|
|
out <- cliproxyexecutor.StreamChunk{Payload: tail[i]}
|
|
}
|
|
if errScan := scanner.Err(); errScan != nil {
|
|
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
|
|
reporter.PublishFailure(ctx)
|
|
out <- cliproxyexecutor.StreamChunk{Err: errScan}
|
|
} 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)) {
|
|
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) 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)
|
|
}
|
|
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) 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)
|
|
|
|
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 {
|
|
if auth != nil {
|
|
if auth.Attributes != nil {
|
|
if ua := strings.TrimSpace(auth.Attributes["user_agent"]); ua != "" {
|
|
return ua
|
|
}
|
|
}
|
|
if auth.Metadata != nil {
|
|
if ua, ok := auth.Metadata["user_agent"].(string); ok && strings.TrimSpace(ua) != "" {
|
|
return strings.TrimSpace(ua)
|
|
}
|
|
}
|
|
}
|
|
return misc.AntigravityUserAgent()
|
|
}
|
|
|
|
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 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{
|
|
antigravityBaseURLProd,
|
|
antigravityBaseURLDaily,
|
|
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
|
|
}
|