1495 lines · 50.1 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"	"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/registry"	"github.com/router-for-me/CLIProxyAPI/v6/internal/util"	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"	antigravityModelsPath          = "/v1internal:fetchAvailableModels"	antigravityClientID            = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"	antigravityClientSecret        = "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"	defaultAntigravityAgent        = "antigravity/1.104.0 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}} // Identifier returns the executor identifier.func (e *AntigravityExecutor) Identifier() string { return antigravityAuthType } // PrepareRequest prepares the HTTP request for execution (no-op for Antigravity).func (e *AntigravityExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil } // 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) {	isClaude := strings.Contains(strings.ToLower(req.Model), "claude")	if isClaude || strings.Contains(req.Model, "gemini-3-pro") {		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(), req.Model, auth)	defer reporter.trackFailure(ctx, &err) 	from := opts.SourceFormat	to := sdktranslator.FromString("antigravity")	originalPayload := bytes.Clone(req.Payload)	if len(opts.OriginalRequest) > 0 {		originalPayload = bytes.Clone(opts.OriginalRequest)	}	originalTranslated := sdktranslator.TranslateRequest(from, to, req.Model, originalPayload, false)	translated := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false) 	translated = ApplyThinkingMetadataCLI(translated, req.Metadata, req.Model)	translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated)	translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, translated)	translated = normalizeAntigravityThinking(req.Model, translated, isClaude)	translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated, originalTranslated) 	baseURLs := antigravityBaseURLFallbackOrder(auth)	httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0) 	var lastStatus int	var lastBody []byte	var lastErr error 	for idx, baseURL := range baseURLs {		httpReq, errReq := e.buildRequest(ctx, auth, token, req.Model, 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			}			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, bytes.Clone(opts.OriginalRequest), translated, bodyBytes, &param)		resp = cliproxyexecutor.Response{Payload: []byte(converted)}		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} // 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) {	token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)	if errToken != nil {		return resp, errToken	}	if updatedAuth != nil {		auth = updatedAuth	} 	reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)	defer reporter.trackFailure(ctx, &err) 	from := opts.SourceFormat	to := sdktranslator.FromString("antigravity")	originalPayload := bytes.Clone(req.Payload)	if len(opts.OriginalRequest) > 0 {		originalPayload = bytes.Clone(opts.OriginalRequest)	}	originalTranslated := sdktranslator.TranslateRequest(from, to, req.Model, originalPayload, true)	translated := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true) 	translated = ApplyThinkingMetadataCLI(translated, req.Metadata, req.Model)	translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated)	translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, translated)	translated = normalizeAntigravityThinking(req.Model, translated, true)	translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated, originalTranslated) 	baseURLs := antigravityBaseURLFallbackOrder(auth)	httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0) 	var lastStatus int	var lastBody []byte	var lastErr error 	for idx, baseURL := range baseURLs {		httpReq, errReq := e.buildRequest(ctx, auth, token, req.Model, 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			}			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		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, parseAntigravityUsage(resp.Payload))		var param any		converted := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), translated, resp.Payload, &param)		resp = cliproxyexecutor.Response{Payload: []byte(converted)}		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} 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 usageResult := root.Get("usageMetadata"); usageResult.Exists() {			usageRaw = usageResult.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				}

Lines 1–500 of 1495.