2404 lines · 80.5 KB · Go
RawDownload
// 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/v7/internal/cache"	"github.com/router-for-me/CLIProxyAPI/v7/internal/config"	"github.com/router-for-me/CLIProxyAPI/v7/internal/misc"	"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"	"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"	"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"	antigravityclaude "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/antigravity/claude"	"github.com/router-for-me/CLIProxyAPI/v7/internal/util"	sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"	sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/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 instancefunc 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

Lines 1–500 of 2404.