1591 lines · 51.9 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/v6/internal/config"	"github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"	"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.19.6 darwin/arm64"	antigravityAuthType            = "antigravity"	refreshSkew                    = 3000 * 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**") var (	randSource      = rand.New(rand.NewSource(time.Now().UnixNano()))	randSourceMutex sync.Mutex) // 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 := 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} // 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)} // 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	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)	} 	token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)	if errToken != nil {		return resp, errToken	}	if updatedAuth != nil {		auth = updatedAuth	} 	reporter := 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	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 := payloadRequestedModel(opts, req.Model)	translated = 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 {			httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, translated, false, opts.Alt, baseURL)			if errReq != nil {				err = errReq				return resp, err			} 			httpResp, errDo := httpClient.Do(httpReq)			if errDo != nil {				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			} 			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 {				recordAPIResponseError(ctx, e.cfg, errRead)				err = errRead				return resp, err			}			appendAPIResponseChunk(ctx, e.cfg, bodyBytes) 			if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices {				log.Debugf("antigravity executor: upstream error status: %d, body: %s", httpResp.StatusCode, 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 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					}				}				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					}				}				err = sErr				return resp, err			} 			reporter.publish(ctx, parseAntigravityUsage(bodyBytes))			var param any			converted := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, bodyBytes, &param)			resp = cliproxyexecutor.Response{Payload: converted, Headers: httpResp.Header.Clone()}			reporter.ensurePublished(ctx)			return resp, nil		} 		switch {		case lastStatus != 0:			sErr := statusErr{code: lastStatus, msg: string(lastBody)}			if lastStatus == http.StatusTooManyRequests {				if retryAfter, parseErr := parseRetryDelay(lastBody); parseErr == nil && retryAfter != nil {					sErr.retryAfter = retryAfter				}			}			err = sErr		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 	token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)	if errToken != nil {		return resp, errToken	}	if updatedAuth != nil {		auth = updatedAuth	} 	reporter := 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	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 := payloadRequestedModel(opts, req.Model)	translated = 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 {			httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, translated, true, opts.Alt, baseURL)			if errReq != nil {				err = errReq				return resp, err			} 			httpResp, errDo := httpClient.Do(httpReq)			if errDo != nil {				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			}			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 {					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				}				appendAPIResponseChunk(ctx, e.cfg, 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 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					}				}				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					}				}				err = sErr				return resp, err			} 			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()					appendAPIResponseChunk(ctx, e.cfg, line) 					// Filter usage metadata for all models					// Only retain usage statistics in the terminal chunk					line = FilterSSEUsageMetadata(line) 					payload := jsonPayload(line)					if payload == nil {						continue					} 					if detail, ok := parseAntigravityStreamUsage(payload); ok {						reporter.publish(ctx, detail)					} 					out <- cliproxyexecutor.StreamChunk{Payload: payload}				}				if errScan := scanner.Err(); errScan != nil {					recordAPIResponseError(ctx, e.cfg, errScan)					reporter.publishFailure(ctx)					out <- cliproxyexecutor.StreamChunk{Err: errScan}				} else {					reporter.ensurePublished(ctx)				}			}(httpResp) 			var buffer bytes.Buffer

Lines 1–500 of 1591.