package executor import (	"bufio"	"bytes"	"context"	"encoding/json"	"fmt"	"io"	"math/rand"	"net/http"	"net/url"	"strconv"	"strings"	"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.sandbox.googleapis.com"	antigravityBaseURLAutopush     = "https://autopush-cloudcode-pa.sandbox.googleapis.com"	antigravityBaseURLProd         = "https://cloudcode-pa.googleapis.com"	antigravityStreamPath          = "/v1internal:streamGenerateContent"	antigravityGeneratePath        = "/v1internal:generateContent"	antigravityModelsPath          = "/v1internal:fetchAvailableModels"	antigravityClientID            = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"	antigravityClientSecret        = "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"	defaultAntigravityAgent        = "antigravity/1.11.5 windows/amd64"	antigravityAuthType            = "antigravity"	refreshSkew                    = 3000 * time.Second	streamScannerBuffer        int = 20_971_520) var randSource = rand.New(rand.NewSource(time.Now().UnixNano())) // AntigravityExecutor proxies requests to the antigravity upstream.type AntigravityExecutor struct {	cfg *config.Config} // NewAntigravityExecutor constructs a new executor instance.func NewAntigravityExecutor(cfg *config.Config) *AntigravityExecutor {	return &AntigravityExecutor{cfg: cfg}} // Identifier implements ProviderExecutor.func (e *AntigravityExecutor) Identifier() string { return antigravityAuthType } // PrepareRequest implements ProviderExecutor.func (e *AntigravityExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil } // applyThinkingMetadata applies thinking config from model suffix metadata (e.g., -reasoning, -thinking-N).// It trusts user intent when suffix is used, even if registry doesn't have Thinking metadata.func applyThinkingMetadata(translated []byte, metadata map[string]any, model string) []byte {	budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(metadata)	if !ok {		return translated	}	if budgetOverride != nil && util.ModelSupportsThinking(model) {		norm := util.NormalizeThinkingBudget(model, *budgetOverride)		budgetOverride = &norm	}	return util.ApplyGeminiCLIThinkingConfig(translated, budgetOverride, includeOverride)} // Execute handles non-streaming requests via the antigravity generate endpoint.func (e *AntigravityExecutor) Execute(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")	translated := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false) 	translated = applyThinkingMetadata(translated, req.Metadata, req.Model) 	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)			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			}			err = statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)}			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:		err = statusErr{code: lastStatus, msg: string(lastBody)}	case lastErr != nil:		err = lastErr	default:		err = statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"}	}	return resp, err} // ExecuteStream handles streaming requests via the antigravity upstream.func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (stream <-chan cliproxyexecutor.StreamChunk, err error) {	ctx = context.WithValue(ctx, "alt", "") 	token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)	if errToken != nil {		return nil, 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")	translated := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true) 	translated = applyThinkingMetadata(translated, req.Metadata, req.Model) 	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 nil, err		} 		httpResp, errDo := httpClient.Do(httpReq)		if errDo != nil {			recordAPIResponseError(ctx, e.cfg, 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		}		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)				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			}			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			}			err = statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)}			return nil, err		} 		out := make(chan cliproxyexecutor.StreamChunk)		stream = out		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()				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)				} 				chunks := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), translated, bytes.Clone(payload), &param)				for i := range chunks {					out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunks[i])}				}			}			tail := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), translated, []byte("[DONE]"), &param)			for i := range tail {				out <- cliproxyexecutor.StreamChunk{Payload: []byte(tail[i])}			}			if errScan := scanner.Err(); errScan != nil {				recordAPIResponseError(ctx, e.cfg, errScan)				reporter.publishFailure(ctx)				out <- cliproxyexecutor.StreamChunk{Err: errScan}			} else {				reporter.ensurePublished(ctx)			}		}(httpResp)		return stream, nil	} 	switch {	case lastStatus != 0:		err = statusErr{code: lastStatus, msg: string(lastBody)}	case lastErr != nil:		err = lastErr	default:		err = statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"}	}	return nil, err} // Refresh refreshes the OAuth token 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 is not supported for the antigravity provider.func (e *AntigravityExecutor) CountTokens(context.Context, *cliproxyauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {	return cliproxyexecutor.Response{}, statusErr{code: http.StatusNotImplemented, msg: "count tokens not supported"}} // FetchAntigravityModels retrieves available models using the supplied auth.func FetchAntigravityModels(ctx context.Context, auth *cliproxyauth.Auth, cfg *config.Config) []*registry.ModelInfo {	exec := &AntigravityExecutor{cfg: cfg}	token, updatedAuth, errToken := exec.ensureAccessToken(ctx, auth)	if errToken != nil || token == "" {		return nil	}	if updatedAuth != nil {		auth = updatedAuth	} 	baseURLs := antigravityBaseURLFallbackOrder(auth)	httpClient := newProxyAwareHTTPClient(ctx, cfg, auth, 0) 	for idx, baseURL := range baseURLs {		modelsURL := baseURL + antigravityModelsPath		httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, modelsURL, bytes.NewReader([]byte(`{}`)))		if errReq != nil {			return nil		}		httpReq.Header.Set("Content-Type", "application/json")		httpReq.Header.Set("Authorization", "Bearer "+token)		httpReq.Header.Set("User-Agent", resolveUserAgent(auth))		if host := resolveHost(baseURL); host != "" {			httpReq.Host = host		} 		httpResp, errDo := httpClient.Do(httpReq)		if errDo != nil {			if idx+1 < len(baseURLs) {				log.Debugf("antigravity executor: models request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1])				continue			}			return nil		} 		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 {			if idx+1 < len(baseURLs) {				log.Debugf("antigravity executor: models read error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1])				continue			}			return nil		}		if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices {			if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) {				log.Debugf("antigravity executor: models request rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1])				continue			}			return nil		} 		result := gjson.GetBytes(bodyBytes, "models")		if !result.Exists() {			return nil		} 		now := time.Now().Unix()		models := make([]*registry.ModelInfo, 0, len(result.Map()))		for id := range result.Map() {			id = modelName2Alias(id)			if id != "" {				modelInfo := &registry.ModelInfo{					ID:          id,					Name:        id,					Description: id,					DisplayName: id,					Version:     id,					Object:      "model",					Created:     now,					OwnedBy:     antigravityAuthType,					Type:        antigravityAuthType,				}				// Add Thinking support for thinking models				if strings.HasSuffix(id, "-thinking") || strings.Contains(id, "-thinking-") {					modelInfo.Thinking = &registry.ThinkingSupport{						Min:            1024,						Max:            100000,						ZeroAllowed:    false,						DynamicAllowed: true,					}				}				models = append(models, modelInfo)			}		}		return models	}	return nil} 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	}	updated, errRefresh := e.refreshToken(ctx, 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("User-Agent", defaultAntigravityAgent)	httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded") 	httpClient := newProxyAwareHTTPClient(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 {		return auth, statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)}	} 	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	auth.Metadata["timestamp"] = time.Now().UnixMilli()	auth.Metadata["expired"] = time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339)	auth.Metadata["type"] = antigravityAuthType	return auth, 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"}

Lines 1–500 of 774. Load more