947 lines · 35.0 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 (	"bytes"	"context"	"crypto/sha256"	"crypto/tls"	"encoding/hex"	"encoding/json"	"fmt"	"net/http"	"strings"	"time" 	"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/runtime/executor/helps"	internalsignature "github.com/router-for-me/CLIProxyAPI/v7/internal/signature"	antigravityclaude "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/antigravity/claude"	"github.com/router-for-me/CLIProxyAPI/v7/internal/util"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"	"github.com/router-for-me/CLIProxyAPI/v7/sdk/proxyutil"	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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"	antigravityAuthType                    = "antigravity"	antigravityRequestTokenSafetyWindow    = 5 * time.Minute	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**") // 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}} func (e *AntigravityExecutor) obfuscateSensitiveWords(payload []byte) []byte {	if e == nil || e.cfg == nil || len(e.cfg.Antigravity.SensitiveWords) == 0 {		return payload	}	matcher := helps.BuildSensitiveWordMatcher(e.cfg.Antigravity.SensitiveWords)	return helps.ObfuscateSensitiveWordsInSystemInstruction(payload, matcher)} // Each Antigravity credential gets its own HTTP/1.1 connection pool. Sessions routed// to the same auth reuse that pool, while different OAuth identities never share a// TCP/TLS connection, matching the native client's one-credential process model.// The cache is bounded so pools cannot accumulate when keys churn.var (	antigravityBaseTransport = defaultAntigravityBaseTransport()	antigravityTransports    = helps.NewTransportCache[antigravityTransportKey](antigravityTransportCacheCapacity)) const (	// antigravityTransportCacheCapacity caps how many Antigravity connection pools stay	// alive. The bound exists only to stop entries from accumulating when keys churn, for	// example when a credential's proxy is rotated through the management API or when an	// SDK embedder supplies a freshly built base transport per request.	//	// It is sized for large deployments on purpose. An unused cache entry costs under 1 KB	// and no goroutines, so capacity is close to free, whereas evicting a pool that is	// still in active use forces the next request on that credential to redo the TCP + TLS	// handshake and defeats the point of caching. Credential counts in the low thousands	// are expected once Home-managed pools are included.	//	// Capacity is therefore NOT the lever for bounding memory: an idle pooled connection	// costs roughly 38 KB plus three goroutines, and that total is driven by live traffic	// and reclaimed by IdleConnTimeout. Shrinking this number does not save that memory,	// it only causes pool thrashing.	antigravityTransportCacheCapacity = 8192 	// antigravityDefaultMaxIdleConnsPerHost sets the default number of idle connections	// to retain per host per credential when connection pooling is enabled.	// Matches Go's DefaultMaxIdleConnsPerHost (2) and the native Antigravity binary.	antigravityDefaultMaxIdleConnsPerHost = 2 	// antigravityMaxAllowedMaxIdleConnsPerHost is the hard upper bound on MaxIdleConnsPerHost (100).	// Prevents unbounded connection pool expansion in multi-credential environments.	antigravityMaxAllowedMaxIdleConnsPerHost = 100 	// antigravityDefaultIdleConnTimeout is the default idle connection timeout (30 seconds).	// Kept strictly far below Google Frontend (GFE / ESF) 240-second cutoff to prevent	// client-side reuse of half-closed connections that cause connection resets.	antigravityDefaultIdleConnTimeout = 30 * time.Second 	// antigravityMaxAllowedIdleConnTimeout is the hard upper bound on IdleConnTimeout (210 seconds).	// Kept strictly below Google Frontend (GFE / ESF) 240.0-second HTTP/1.1 idle keep-alive cutoff	// with a 30-second safety margin to eliminate timer race conditions.	antigravityMaxAllowedIdleConnTimeout = 210 * time.Second 	// antigravityAnonymousTransportScope is the pool scope for auth objects that carry	// no identity at all. Reaching it means the auth has no ID, no source path and no	// token of any kind, so there is no credential to keep isolated and a single shared	// pool is safe. Allocating a private pool per request instead would leak a	// connection pool, and the goroutines managing it, on every call.	antigravityAnonymousTransportScope = "anonymous") type antigravityPoolSettings struct {	shortMode           bool	idleConnTimeout     time.Duration	maxIdleConnsPerHost int} func resolveAntigravityPoolSettings(cfg *config.Config) antigravityPoolSettings {	// By default, upstream connection pooling is disabled (shortMode = true, maxIdleConnsPerHost = -1)	// to prevent socket buildup and stale connection errors across rotating credentials.	settings := antigravityPoolSettings{		shortMode:           true,		maxIdleConnsPerHost: -1,		idleConnTimeout:     0,	}	if cfg == nil {		return settings	} 	// Pooling is active ONLY when explicitly enabled: true	if cfg.Antigravity.ConnectionPool.Enabled == nil || !*cfg.Antigravity.ConnectionPool.Enabled {		return settings	} 	// Enabled is true: initialize default pool settings	settings.shortMode = false	settings.idleConnTimeout = antigravityDefaultIdleConnTimeout	settings.maxIdleConnsPerHost = antigravityDefaultMaxIdleConnsPerHost 	rawTimeout := strings.TrimSpace(cfg.Antigravity.ConnectionPool.IdleConnTimeout)	if rawTimeout != "" {		d, err := time.ParseDuration(rawTimeout)		if err != nil {			log.Warnf("antigravity executor: invalid idle-conn-timeout %q: %v, using default %v", rawTimeout, err, antigravityDefaultIdleConnTimeout)		} else {			if d <= 0 {				settings.shortMode = true				settings.maxIdleConnsPerHost = -1				return settings			}			if d > antigravityMaxAllowedIdleConnTimeout {				d = antigravityMaxAllowedIdleConnTimeout			}			settings.idleConnTimeout = d		}	} 	if cfg.Antigravity.ConnectionPool.MaxIdleConnsPerHost != nil {		val := *cfg.Antigravity.ConnectionPool.MaxIdleConnsPerHost		if val < 0 {			settings.shortMode = true			settings.maxIdleConnsPerHost = -1			return settings		}		if val > antigravityMaxAllowedMaxIdleConnsPerHost {			val = antigravityMaxAllowedMaxIdleConnsPerHost		}		settings.maxIdleConnsPerHost = val	} 	if settings.maxIdleConnsPerHost < 0 {		settings.shortMode = true	} 	return settings} // ResetAntigravityTransports purges all cached Antigravity connection pools and closes their idle connections.// Used during configuration hot-reloads to ensure updated pool parameters apply immediately.func ResetAntigravityTransports() {	antigravityTransports.Purge()} // AntigravityTransportsLen reports the current number of cached Antigravity transports.func AntigravityTransportsLen() int {	return antigravityTransports.Len()} // closeAntigravityAuthIdleTransports closes and removes all idle connections for the given auth.func closeAntigravityAuthIdleTransports(auth *cliproxyauth.Auth) {	if auth == nil {		return	}	scope := antigravityTransportScope(auth)	if scope == "" || scope == antigravityAnonymousTransportScope {		return	}	antigravityTransports.CloseMatching(func(key antigravityTransportKey) bool {		return key.credential == scope	})} // antigravityTransportKey identifies one connection pool. At most one of proxy and// base is set: proxy for a credential-scoped proxy pool, base for a transport handed// in through the request context, and neither for a direct pool.// Resolved pool settings (shortMode, idleConnTimeout, maxIdleConnsPerHost) are included// in the key to prevent stale transport reuse across configuration hot-reloads under load.type antigravityTransportKey struct {	credential          string	proxy               string	base                *http.Transport	shortMode           bool	idleConnTimeout     time.Duration	maxIdleConnsPerHost int} func defaultAntigravityBaseTransport() *http.Transport {	if transport, ok := http.DefaultTransport.(*http.Transport); ok && transport != nil {		return transport	}	return &http.Transport{}} func cloneTransportWithHTTP11(base *http.Transport, cfgs ...*config.Config) *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()	}	// Native Antigravity sends no ALPN extension. With HTTP/2 disabled above,	// an empty NextProtos keeps the wire shape aligned while using HTTP/1.1.	clone.TLSClientConfig.NextProtos = nil	applyAntigravityPoolLimits(clone, cfgs...)	return clone} // applyAntigravityPoolLimits configures connection pool parameters for Antigravity.// Default: IdleConnTimeout=30s, MaxIdleConnsPerHost=2.// IdleConnTimeout is strictly capped at 240s (antigravityMaxAllowedIdleConnTimeout)// matching the Google Frontend (GFE / ESF) HTTP/1.1 idle cutoff.// If short-lived connection mode is configured, MaxIdleConnsPerHost is set to -1// so idle connections are closed immediately upon request completion.func applyAntigravityPoolLimits(transport *http.Transport, cfgs ...*config.Config) {	if transport == nil {		return	}	var cfg *config.Config	if len(cfgs) > 0 {		cfg = cfgs[0]	}	settings := resolveAntigravityPoolSettings(cfg)	if settings.shortMode {		transport.MaxIdleConnsPerHost = -1		transport.DisableKeepAlives = false		transport.IdleConnTimeout = 0		return	} 	// If the operator base transport already disabled pooling (< 0), honor it.	if transport.MaxIdleConnsPerHost < 0 {		return	} 	// Go treats 0 as DefaultMaxIdleConnsPerHost (2). Ensure at least configured/default limit.	if transport.MaxIdleConnsPerHost < settings.maxIdleConnsPerHost {		transport.MaxIdleConnsPerHost = settings.maxIdleConnsPerHost	} 	// MaxIdleConns caps the pool across all hosts.	if transport.MaxIdleConns > 0 && transport.MaxIdleConns < transport.MaxIdleConnsPerHost {		transport.MaxIdleConns = transport.MaxIdleConnsPerHost	} 	// Apply IdleConnTimeout:	// If the config explicitly specified a timeout, apply settings.idleConnTimeout.	// If config did not specify a timeout:	// - if base transport already has a timeout > 0, preserve it (capped at 240s).	// - otherwise, apply default 30s.	rawTimeout := ""	if cfg != nil {		rawTimeout = strings.TrimSpace(cfg.Antigravity.ConnectionPool.IdleConnTimeout)	}	if rawTimeout != "" {		transport.IdleConnTimeout = settings.idleConnTimeout	} else if transport.IdleConnTimeout == 0 || transport.IdleConnTimeout == 90*time.Second {		transport.IdleConnTimeout = settings.idleConnTimeout	} else if transport.IdleConnTimeout > antigravityMaxAllowedIdleConnTimeout {		transport.IdleConnTimeout = antigravityMaxAllowedIdleConnTimeout	}} // antigravityHTTP11Transport returns the HTTP/1.1 pool shared by every request that// uses the same credential and the same base transport. The base is either the// process default or a transport provided through the request context.func antigravityHTTP11Transport(auth *cliproxyauth.Auth, base *http.Transport, cfgs ...*config.Config) *http.Transport {	if base == nil {		return nil	}	var cfg *config.Config	if len(cfgs) > 0 {		cfg = cfgs[0]	}	settings := resolveAntigravityPoolSettings(cfg)	key := antigravityTransportKey{		credential:          antigravityTransportScope(auth),		base:                base,		shortMode:           settings.shortMode,		idleConnTimeout:     settings.idleConnTimeout,		maxIdleConnsPerHost: settings.maxIdleConnsPerHost,	}	transport, errGet := antigravityTransports.Get(key, func() (*http.Transport, error) {		return cloneTransportWithHTTP11(base, cfgs...), nil	})	if errGet != nil {		// Defensive only: the builder above cannot fail. Never return nil here, because a		// nil Transport makes http.Client fall back to http.DefaultTransport, which		// advertises h2 over ALPN and would break the Antigravity wire fingerprint.		log.Debugf("antigravity executor: cache HTTP/1.1 transport failed: %v", errGet)		return cloneTransportWithHTTP11(base, cfgs...)	}	return transport} // antigravityProxiedHTTP11Transport returns the credential-scoped HTTP/1.1 pool for// one proxy setting, or nil when the proxy setting cannot be turned into a// transport. Keying on the normalized proxy string rather than on a prebuilt// transport keeps one pool per credential and proxy instead of one per request.func antigravityProxiedHTTP11Transport(auth *cliproxyauth.Auth, proxyURL string, cfgs ...*config.Config) *http.Transport {	proxyURL = strings.TrimSpace(proxyURL)	if proxyURL == "" {		return nil	}	var cfg *config.Config	if len(cfgs) > 0 {		cfg = cfgs[0]	}	settings := resolveAntigravityPoolSettings(cfg)	key := antigravityTransportKey{		credential:          antigravityTransportScope(auth),		proxy:               proxyURL,		shortMode:           settings.shortMode,		idleConnTimeout:     settings.idleConnTimeout,		maxIdleConnsPerHost: settings.maxIdleConnsPerHost,	}	transport, errGet := antigravityTransports.Get(key, func() (*http.Transport, error) {		base, _, errBuild := proxyutil.BuildHTTPTransport(proxyURL)		if errBuild != nil {			return nil, errBuild		}		if base == nil {			return nil, fmt.Errorf("antigravity executor: proxy setting produced no transport")		}		return cloneTransportWithHTTP11(base, cfgs...), nil	})	if errGet != nil {		// The caller falls back to NewProxyAwareHTTPClient, which reports the failure		// and applies the context transport fallback.		return nil	}	return transport} // antigravityTransportScope returns the connection-pool scope for one credential.// Runtime auths always carry an ID. Incomplete auth objects, such as those built by// tests, plugins or SDK embedders, fall back to another stable credential marker so// they neither share a pool with an unrelated OAuth identity nor allocate a fresh// pool, and with it a fresh set of pool goroutines, on every single request.func antigravityTransportScope(auth *cliproxyauth.Auth) string {	if auth == nil {		return antigravityAnonymousTransportScope	}	if id := strings.TrimSpace(auth.ID); id != "" {		return "id:" + id	}	if auth.Attributes != nil {		if path := strings.TrimSpace(auth.Attributes[cliproxyauth.AttributePath]); path != "" {			return "path:" + path		}		if source := strings.TrimSpace(auth.Attributes[cliproxyauth.AttributeSource]); source != "" {			return "source:" + source		}	}	// Fall back to the credential material itself. Auth.Label is deliberately not used:	// it is documented as an optional human readable label for logging and carries no	// uniqueness guarantee, so two different OAuth identities sharing one label would	// wrongly share a TCP/TLS pool.	//	// The refresh token is preferred over the access token because it stays stable	// across token rotation. Keying on the access token would move a credential to a new	// pool on every refresh, and would also strand refresh requests themselves, which	// run before any access token exists.	if refresh := strings.TrimSpace(metaStringValue(auth.Metadata, "refresh_token")); refresh != "" {		return antigravityCredentialScope("refresh:", refresh)	}	if access := strings.TrimSpace(metaStringValue(auth.Metadata, "access_token")); access != "" {		return antigravityCredentialScope("token:", access)	}	return antigravityAnonymousTransportScope} // antigravityCredentialScope derives a pool scope from secret credential material.// Only a short digest is retained, and it is never logged, so a pool key cannot be// used to recover the credential it came from.func antigravityCredentialScope(prefix, secret string) string {	digest := sha256.Sum256([]byte(secret))	return prefix + hex.EncodeToString(digest[:8])} // newAntigravityHTTPClient creates an HTTP client specifically for Antigravity,// enforcing HTTP/1.1 by disabling HTTP/2 to match the native Antigravity client, which// negotiates TLS 1.3 without advertising an ALPN protocol and therefore never uses h2.// The underlying Transport is always shared so keep-alive connections survive across// requests instead of forcing a fresh TCP + TLS handshake every time.func newAntigravityHTTPClient(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) *http.Client {	// Native Antigravity reuses one transport across requests. Opt into a	// credential-scoped proxy transport only here so other providers keep their	// existing lifecycle and different OAuth identities remain isolated.	if proxyURL := antigravityProxyURL(cfg, auth); proxyURL != "" {		if transport := antigravityProxiedHTTP11Transport(auth, proxyURL, cfg); transport != nil {			return &http.Client{Transport: transport, Timeout: timeout}		}		// Fall through so NewProxyAwareHTTPClient reports the failure and applies the		// context transport fallback, preserving the previous behavior.	} 	client := helps.NewProxyAwareHTTPClient(ctx, cfg, auth, timeout)	// Direct requests share an HTTP/1.1 pool only within the selected credential.	if client.Transport == nil {		client.Transport = antigravityHTTP11Transport(auth, antigravityBaseTransport, cfg)		return client	} 	// Preserve a context-provided transport while forcing HTTP/1.1. The cache key	// includes credential identity, so sharing the base does not share TLS pools.	transport, ok := client.Transport.(*http.Transport)	if !ok {		// A RoundTripper that is not an *http.Transport owns its own protocol behavior.		return client	}	if transport == nil {		// A typed-nil *http.Transport still satisfies the interface nil check in		// NewProxyAwareHTTPClient. Leaving it in place would make http.Client fall back		// to http.DefaultTransport, which advertises h2 over ALPN and breaks the		// Antigravity fingerprint, so substitute the process base transport.		transport = antigravityBaseTransport	}	client.Transport = antigravityHTTP11Transport(auth, transport, cfg)	return client} func antigravityProxyURL(cfg *config.Config, auth *cliproxyauth.Auth) string {	if auth != nil {		if proxyURL := strings.TrimSpace(auth.ProxyURL); proxyURL != "" {			return proxyURL		}	}	if cfg != nil {		return strings.TrimSpace(cfg.ProxyURL)	}	return ""} func sanitizeAntigravityGeminiRequestSignatures(modelName string, rawJSON []byte) []byte {	if !antigravityUsesReasoningReplayCache(modelName) {		return rawJSON	}	rawJSON = internalsignature.SanitizeGeminiRequestThoughtSignatures(rawJSON, "request.contents")	return normalizeAntigravityGeminiFunctionResponseRoles(rawJSON)} // ensureAntigravityGeminiLeadingUserContent prepends a synthetic empty user turn// after every contents rewrite, including reasoning replay. Claude targets are// left unchanged because the adapter rejects empty text parts.func ensureAntigravityGeminiLeadingUserContent(modelName string, payload []byte) []byte {	if strings.Contains(strings.ToLower(modelName), "claude") {		return payload

Lines 1–500 of 947.