Spaces:
Paused
Paused
| // 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/cache" | |
| "github.com/router-for-me/CLIProxyAPI/v6/internal/config" | |
| "github.com/router-for-me/CLIProxyAPI/v6/internal/misc" | |
| "github.com/router-for-me/CLIProxyAPI/v6/internal/registry" | |
| "github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor/helps" | |
| "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking" | |
| antigravityclaude "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/claude" | |
| "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.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 instance | |
| func 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 { | |
| 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 | |
| } | |
| originalPayload := originalPayloadSource | |
| originalPayload, errValidate := validateAntigravityRequestSignatures(from, originalPayload) | |
| if errValidate != nil { | |
| return resp, errValidate | |
| } | |
| req.Payload = originalPayload | |
| token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) | |
| if errToken != nil { | |
| return resp, errToken | |
| } | |
| if updatedAuth != nil { | |
| auth = updatedAuth | |
| } | |
| 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 := helps.PayloadRequestedModel(opts, req.Model) | |
| translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel) | |
| useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg) | |
| 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 { | |
| requestPayload := translated | |
| if useCredits { | |
| if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { | |
| requestPayload = cp | |
| helps.MarkCreditsUsed(ctx) | |
| } | |
| } | |
| httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, false, opts.Alt, baseURL) | |
| if errReq != nil { | |
| err = errReq | |
| return resp, err | |
| } | |
| httpResp, errDo := httpClient.Do(httpReq) | |
| if errDo != nil { | |
| helps.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 | |
| } | |
| helps.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 { | |
| helps.RecordAPIResponseError(ctx, e.cfg, errRead) | |
| err = errRead | |
| return resp, err | |
| } | |
| helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) | |
| if httpResp.StatusCode == http.StatusTooManyRequests { | |
| decision := decideAntigravity429(bodyBytes) | |
| switch decision.kind { | |
| case antigravity429DecisionInstantRetrySameAuth: | |
| if attempt+1 < attempts { | |
| if decision.retryAfter != nil && *decision.retryAfter > 0 { | |
| wait := antigravityInstantRetryDelay(*decision.retryAfter) | |
| log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait) | |
| if errWait := antigravityWait(ctx, wait); errWait != nil { | |
| return resp, errWait | |
| } | |
| } | |
| continue attemptLoop | |
| } | |
| case antigravity429DecisionShortCooldownSwitchAuth: | |
| if decision.retryAfter != nil && *decision.retryAfter > 0 { | |
| markAntigravityShortCooldown(auth, baseModel, time.Now(), *decision.retryAfter) | |
| log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel) | |
| } | |
| case antigravity429DecisionFullQuotaExhausted: | |
| if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { | |
| markAntigravityCreditsPermanentlyDisabled(auth) | |
| } | |
| // No credits logic - just fall through to error return below | |
| } | |
| } | |
| if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { | |
| log.Debugf("antigravity executor: upstream error status: %d, body: %s", httpResp.StatusCode, helps.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 antigravityShouldRetryTransientResourceExhausted429(httpResp.StatusCode, bodyBytes) && attempt+1 < attempts { | |
| delay := antigravityTransient429RetryDelay(attempt) | |
| log.Debugf("antigravity executor: transient 429 resource exhausted 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 | |
| } | |
| 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 | |
| } | |
| } | |
| if antigravityShouldRetrySoftRateLimit(httpResp.StatusCode, bodyBytes) { | |
| if attempt+1 < attempts { | |
| delay := antigravitySoftRateLimitDelay(attempt) | |
| log.Debugf("antigravity executor: soft rate limit 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 | |
| } | |
| } | |
| err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) | |
| return resp, err | |
| } | |
| // Success | |
| if useCredits { | |
| clearAntigravityCreditsFailureState(auth) | |
| } | |
| reporter.Publish(ctx, helps.ParseAntigravityUsage(bodyBytes)) | |
| var param any | |
| converted := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, bodyBytes, ¶m) | |
| resp = cliproxyexecutor.Response{Payload: converted, Headers: httpResp.Header.Clone()} | |
| reporter.EnsurePublished(ctx) | |
| return resp, nil | |
| } | |
| switch { | |
| case lastStatus != 0: | |
| err = newAntigravityStatusErr(lastStatus, lastBody) | |
| 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 | |
| if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown { | |
| 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} | |
| } | |
| 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 | |
| } | |
| originalPayload := originalPayloadSource | |
| originalPayload, errValidate := validateAntigravityRequestSignatures(from, originalPayload) | |
| if errValidate != nil { | |
| return resp, errValidate | |
| } | |
| req.Payload = originalPayload | |
| token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) | |
| if errToken != nil { | |
| return resp, errToken | |
| } | |
| if updatedAuth != nil { | |
| auth = updatedAuth | |
| } | |
| 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 := helps.PayloadRequestedModel(opts, req.Model) | |
| translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel) | |
| useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg) | |
| 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 { | |
| requestPayload := translated | |
| if useCredits { | |
| if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { | |
| requestPayload = cp | |
| helps.MarkCreditsUsed(ctx) | |
| } | |
| } | |
| httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL) | |
| if errReq != nil { | |
| err = errReq | |
| return resp, err | |
| } | |
| httpResp, errDo := httpClient.Do(httpReq) | |
| if errDo != nil { | |
| helps.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 | |
| } | |
| helps.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 { | |
| helps.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 | |
| } | |
| helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) | |
| if httpResp.StatusCode == http.StatusTooManyRequests { | |
| decision := decideAntigravity429(bodyBytes) | |
| switch decision.kind { | |
| case antigravity429DecisionInstantRetrySameAuth: | |
| if attempt+1 < attempts { | |
| if decision.retryAfter != nil && *decision.retryAfter > 0 { | |
| wait := antigravityInstantRetryDelay(*decision.retryAfter) | |
| log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait) | |
| if errWait := antigravityWait(ctx, wait); errWait != nil { | |
| return resp, errWait | |
| } | |
| } | |
| continue attemptLoop | |
| } | |
| case antigravity429DecisionShortCooldownSwitchAuth: | |
| if decision.retryAfter != nil && *decision.retryAfter > 0 { | |
| markAntigravityShortCooldown(auth, baseModel, time.Now(), *decision.retryAfter) | |
| log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel) | |
| } | |
| case antigravity429DecisionFullQuotaExhausted: | |
| if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { | |
| markAntigravityCreditsPermanentlyDisabled(auth) | |
| } | |
| // No credits logic - just fall through to error return below | |
| } | |
| } | |
| 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 antigravityShouldRetryTransientResourceExhausted429(httpResp.StatusCode, bodyBytes) && attempt+1 < attempts { | |
| delay := antigravityTransient429RetryDelay(attempt) | |
| log.Debugf("antigravity executor: transient 429 resource exhausted 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 | |
| } | |
| 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 | |
| } | |
| } | |
| if antigravityShouldRetrySoftRateLimit(httpResp.StatusCode, bodyBytes) { | |
| if attempt+1 < attempts { | |
| delay := antigravitySoftRateLimitDelay(attempt) | |
| log.Debugf("antigravity executor: soft rate limit 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 | |
| } | |
| } | |
| err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) | |
| return resp, err | |
| } | |
| // Stream success | |
| if useCredits { | |
| clearAntigravityCreditsFailureState(auth) | |
| } | |
| 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() | |
| helps.AppendAPIResponseChunk(ctx, e.cfg, line) | |
| // Filter usage metadata for all models | |
| // Only retain usage statistics in the terminal chunk | |
| line = helps.FilterSSEUsageMetadata(line) | |
| payload := helps.JSONPayload(line) | |
| if payload == nil { | |
| continue | |
| } | |
| if detail, ok := helps.ParseAntigravityStreamUsage(payload); ok { | |
| reporter.Publish(ctx, detail) | |
| } | |
| out <- cliproxyexecutor.StreamChunk{Payload: payload} | |
| } | |
| if errScan := scanner.Err(); errScan != nil { | |
| helps.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, helps.ParseAntigravityUsage(resp.Payload)) | |
| var param any | |
| converted := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, resp.Payload, ¶m) | |
| resp = cliproxyexecutor.Response{Payload: converted, Headers: httpResp.Header.Clone()} | |
| reporter.EnsurePublished(ctx) | |
| return resp, nil | |
| } | |
| switch { | |
| case lastStatus != 0: | |
| err = newAntigravityStatusErr(lastStatus, lastBody) | |
| case lastErr != nil: | |
| err = lastErr | |
| default: | |
| err = statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} | |
| } | |
| return resp, err | |
| } | |
| 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 usageMetadataResult := root.Get("usageMetadata"); usageMetadataResult.Exists() { | |
| usageRaw = usageMetadataResult.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 | |
| } | |
| if thought || part.Get("text").Exists() { | |
| kind := "text" | |
| if thought { | |
| kind = "thought" | |
| } | |
| if pendingKind != "" && pendingKind != kind { | |
| flushPending() | |
| } | |
| pendingKind = kind | |
| pendingText.WriteString(text) | |
| if kind == "thought" && sig != "" { | |
| pendingThoughtSig = sig | |
| } | |
| continue | |
| } | |
| flushPending() | |
| parts = append(parts, normalizePart(part)) | |
| } | |
| } | |
| } | |
| flushPending() | |
| if responseTemplate == "" { | |
| responseTemplate = `{"candidates":[{"content":{"role":"model","parts":[]}}]}` | |
| } | |
| partsJSON, _ := json.Marshal(parts) | |
| updatedTemplate, _ := sjson.SetRawBytes([]byte(responseTemplate), "candidates.0.content.parts", partsJSON) | |
| responseTemplate = string(updatedTemplate) | |
| if role != "" { | |
| updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "candidates.0.content.role", role) | |
| responseTemplate = string(updatedTemplate) | |
| } | |
| if finishReason != "" { | |
| updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "candidates.0.finishReason", finishReason) | |
| responseTemplate = string(updatedTemplate) | |
| } | |
| if modelVersion != "" { | |
| updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "modelVersion", modelVersion) | |
| responseTemplate = string(updatedTemplate) | |
| } | |
| if responseID != "" { | |
| updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "responseId", responseID) | |
| responseTemplate = string(updatedTemplate) | |
| } | |
| if usageRaw != "" { | |
| updatedTemplate, _ = sjson.SetRawBytes([]byte(responseTemplate), "usageMetadata", []byte(usageRaw)) | |
| responseTemplate = string(updatedTemplate) | |
| } else if !gjson.Get(responseTemplate, "usageMetadata").Exists() { | |
| updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.promptTokenCount", 0) | |
| responseTemplate = string(updatedTemplate) | |
| updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.candidatesTokenCount", 0) | |
| responseTemplate = string(updatedTemplate) | |
| updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.totalTokenCount", 0) | |
| responseTemplate = string(updatedTemplate) | |
| } | |
| output := `{"response":{},"traceId":""}` | |
| updatedOutput, _ := sjson.SetRawBytes([]byte(output), "response", []byte(responseTemplate)) | |
| output = string(updatedOutput) | |
| if traceID != "" { | |
| updatedOutput, _ = sjson.SetBytes([]byte(output), "traceId", traceID) | |
| output = string(updatedOutput) | |
| } | |
| return []byte(output) | |
| } | |
| // ExecuteStream performs a streaming request to the Antigravity API. | |
| func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) { | |
| if opts.Alt == "responses/compact" { | |
| return nil, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} | |
| } | |
| baseModel := thinking.ParseSuffix(req.Model).ModelName | |
| ctx = context.WithValue(ctx, "alt", "") | |
| if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown { | |
| 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 nil, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d} | |
| } | |
| 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 | |
| } | |
| originalPayload := originalPayloadSource | |
| originalPayload, errValidate := validateAntigravityRequestSignatures(from, originalPayload) | |
| if errValidate != nil { | |
| return nil, errValidate | |
| } | |
| req.Payload = originalPayload | |
| token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) | |
| if errToken != nil { | |
| return nil, errToken | |
| } | |
| if updatedAuth != nil { | |
| auth = updatedAuth | |
| } | |
| 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 nil, err | |
| } | |
| requestedModel := helps.PayloadRequestedModel(opts, req.Model) | |
| translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel) | |
| useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg) | |
| 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 { | |
| requestPayload := translated | |
| if useCredits { | |
| if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { | |
| requestPayload = cp | |
| helps.MarkCreditsUsed(ctx) | |
| } | |
| } | |
| httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL) | |
| if errReq != nil { | |
| err = errReq | |
| return nil, err | |
| } | |
| httpResp, errDo := httpClient.Do(httpReq) | |
| if errDo != nil { | |
| helps.RecordAPIResponseError(ctx, e.cfg, errDo) | |
| if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { | |
| return nil, 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 | |
| } | |
| helps.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 { | |
| helps.RecordAPIResponseError(ctx, e.cfg, errRead) | |
| if errors.Is(errRead, context.Canceled) || errors.Is(errRead, context.DeadlineExceeded) { | |
| err = errRead | |
| return nil, err | |
| } | |
| if errCtx := ctx.Err(); errCtx != nil { | |
| err = errCtx | |
| return nil, 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 nil, err | |
| } | |
| helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) | |
| if httpResp.StatusCode == http.StatusTooManyRequests { | |
| decision := decideAntigravity429(bodyBytes) | |
| switch decision.kind { | |
| case antigravity429DecisionInstantRetrySameAuth: | |
| if attempt+1 < attempts { | |
| if decision.retryAfter != nil && *decision.retryAfter > 0 { | |
| wait := antigravityInstantRetryDelay(*decision.retryAfter) | |
| log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait) | |
| if errWait := antigravityWait(ctx, wait); errWait != nil { | |
| return nil, errWait | |
| } | |
| } | |
| continue attemptLoop | |
| } | |
| case antigravity429DecisionShortCooldownSwitchAuth: | |
| if decision.retryAfter != nil && *decision.retryAfter > 0 { | |
| markAntigravityShortCooldown(auth, baseModel, time.Now(), *decision.retryAfter) | |
| log.Debugf("antigravity executor: short quota cooldown (%s) for model %s recorded", *decision.retryAfter, baseModel) | |
| } | |
| case antigravity429DecisionFullQuotaExhausted: | |
| if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { | |
| markAntigravityCreditsPermanentlyDisabled(auth) | |
| } | |
| // No credits logic - just fall through to error return below | |
| } | |
| } | |
| 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 antigravityShouldRetryTransientResourceExhausted429(httpResp.StatusCode, bodyBytes) && attempt+1 < attempts { | |
| delay := antigravityTransient429RetryDelay(attempt) | |
| log.Debugf("antigravity executor: transient 429 resource exhausted for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) | |
| if errWait := antigravityWait(ctx, delay); errWait != nil { | |
| return nil, errWait | |
| } | |
| continue attemptLoop | |
| } | |
| 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 nil, errWait | |
| } | |
| continue attemptLoop | |
| } | |
| } | |
| if antigravityShouldRetrySoftRateLimit(httpResp.StatusCode, bodyBytes) { | |
| if attempt+1 < attempts { | |
| delay := antigravitySoftRateLimitDelay(attempt) | |
| log.Debugf("antigravity executor: soft rate limit for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) | |
| if errWait := antigravityWait(ctx, delay); errWait != nil { | |
| return nil, errWait | |
| } | |
| continue attemptLoop | |
| } | |
| } | |
| err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) | |
| return nil, err | |
| } | |
| // Stream success | |
| if useCredits { | |
| clearAntigravityCreditsFailureState(auth) | |
| } | |
| 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) | |
| var param any | |
| for scanner.Scan() { | |
| line := scanner.Bytes() | |
| helps.AppendAPIResponseChunk(ctx, e.cfg, line) | |
| // Filter usage metadata for all models | |
| // Only retain usage statistics in the terminal chunk | |
| line = helps.FilterSSEUsageMetadata(line) | |
| payload := helps.JSONPayload(line) | |
| if payload == nil { | |
| continue | |
| } | |
| if detail, ok := helps.ParseAntigravityStreamUsage(payload); ok { | |
| reporter.Publish(ctx, detail) | |
| } | |
| chunks := sdktranslator.TranslateStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, bytes.Clone(payload), ¶m) | |
| for i := range chunks { | |
| out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]} | |
| } | |
| } | |
| tail := sdktranslator.TranslateStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, []byte("[DONE]"), ¶m) | |
| for i := range tail { | |
| out <- cliproxyexecutor.StreamChunk{Payload: tail[i]} | |
| } | |
| if errScan := scanner.Err(); errScan != nil { | |
| helps.RecordAPIResponseError(ctx, e.cfg, errScan) | |
| reporter.PublishFailure(ctx) | |
| out <- cliproxyexecutor.StreamChunk{Err: errScan} | |
| } else { | |
| reporter.EnsurePublished(ctx) | |
| } | |
| }(httpResp) | |
| return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil | |
| } | |
| switch { | |
| case lastStatus != 0: | |
| err = newAntigravityStatusErr(lastStatus, lastBody) | |
| case lastErr != nil: | |
| err = lastErr | |
| default: | |
| err = statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} | |
| } | |
| return nil, err | |
| } | |
| return nil, err | |
| } | |
| // Refresh refreshes the authentication credentials 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 counts tokens for the given request using the Antigravity API. | |
| func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { | |
| baseModel := thinking.ParseSuffix(req.Model).ModelName | |
| from := opts.SourceFormat | |
| to := sdktranslator.FromString("antigravity") | |
| respCtx := context.WithValue(ctx, "alt", opts.Alt) | |
| originalPayloadSource := req.Payload | |
| if len(opts.OriginalRequest) > 0 { | |
| originalPayloadSource = opts.OriginalRequest | |
| } | |
| originalPayloadSource, errValidate := validateAntigravityRequestSignatures(from, originalPayloadSource) | |
| if errValidate != nil { | |
| return cliproxyexecutor.Response{}, errValidate | |
| } | |
| req.Payload = originalPayloadSource | |
| token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) | |
| if errToken != nil { | |
| return cliproxyexecutor.Response{}, errToken | |
| } | |
| if updatedAuth != nil { | |
| auth = updatedAuth | |
| } | |
| if strings.TrimSpace(token) == "" { | |
| return cliproxyexecutor.Response{}, statusErr{code: http.StatusUnauthorized, msg: "missing access token"} | |
| } | |
| // Prepare payload once (doesn't depend on baseURL) | |
| payload := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false) | |
| payload, err := thinking.ApplyThinking(payload, req.Model, from.String(), to.String(), e.Identifier()) | |
| if err != nil { | |
| return cliproxyexecutor.Response{}, err | |
| } | |
| payload = deleteJSONField(payload, "project") | |
| payload = deleteJSONField(payload, "model") | |
| payload = deleteJSONField(payload, "request.safetySettings") | |
| baseURLs := antigravityBaseURLFallbackOrder(auth) | |
| httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) | |
| var authID, authLabel, authType, authValue string | |
| if auth != nil { | |
| authID = auth.ID | |
| authLabel = auth.Label | |
| authType, authValue = auth.AccountInfo() | |
| } | |
| var lastStatus int | |
| var lastBody []byte | |
| var lastErr error | |
| for idx, baseURL := range baseURLs { | |
| base := strings.TrimSuffix(baseURL, "/") | |
| if base == "" { | |
| base = buildBaseURL(auth) | |
| } | |
| var requestURL strings.Builder | |
| requestURL.WriteString(base) | |
| requestURL.WriteString(antigravityCountTokensPath) | |
| if opts.Alt != "" { | |
| requestURL.WriteString("?$alt=") | |
| requestURL.WriteString(url.QueryEscape(opts.Alt)) | |
| } | |
| httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bytes.NewReader(payload)) | |
| if errReq != nil { | |
| return cliproxyexecutor.Response{}, errReq | |
| } | |
| httpReq.Close = true | |
| httpReq.Header.Set("Content-Type", "application/json") | |
| httpReq.Header.Set("Authorization", "Bearer "+token) | |
| httpReq.Header.Set("User-Agent", resolveUserAgent(auth)) | |
| if host := resolveHost(base); host != "" { | |
| httpReq.Host = host | |
| } | |
| var attrs map[string]string | |
| if auth != nil { | |
| attrs = auth.Attributes | |
| } | |
| util.ApplyCustomHeadersFromAttrs(httpReq, attrs) | |
| helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ | |
| URL: requestURL.String(), | |
| Method: http.MethodPost, | |
| Headers: httpReq.Header.Clone(), | |
| Body: payload, | |
| Provider: e.Identifier(), | |
| AuthID: authID, | |
| AuthLabel: authLabel, | |
| AuthType: authType, | |
| AuthValue: authValue, | |
| }) | |
| httpResp, errDo := httpClient.Do(httpReq) | |
| if errDo != nil { | |
| helps.RecordAPIResponseError(ctx, e.cfg, errDo) | |
| if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { | |
| return cliproxyexecutor.Response{}, 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 | |
| } | |
| return cliproxyexecutor.Response{}, errDo | |
| } | |
| helps.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 { | |
| helps.RecordAPIResponseError(ctx, e.cfg, errRead) | |
| return cliproxyexecutor.Response{}, errRead | |
| } | |
| helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) | |
| if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices { | |
| count := gjson.GetBytes(bodyBytes, "totalTokens").Int() | |
| translated := sdktranslator.TranslateTokenCount(respCtx, to, from, count, bodyBytes) | |
| return cliproxyexecutor.Response{Payload: translated, Headers: httpResp.Header.Clone()}, nil | |
| } | |
| 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 | |
| } | |
| } | |
| return cliproxyexecutor.Response{}, sErr | |
| } | |
| 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 | |
| } | |
| } | |
| return cliproxyexecutor.Response{}, sErr | |
| case lastErr != nil: | |
| return cliproxyexecutor.Response{}, lastErr | |
| default: | |
| return cliproxyexecutor.Response{}, statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} | |
| } | |
| } | |
| 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)) { | |
| e.maybeRefreshAntigravityCreditsHint(ctx, auth, accessToken) | |
| return accessToken, nil, nil | |
| } | |
| refreshCtx := context.Background() | |
| if ctx != nil { | |
| if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil { | |
| refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt) | |
| } | |
| } | |
| updated, errRefresh := e.refreshToken(refreshCtx, auth.Clone()) | |
| if errRefresh != nil { | |
| return "", nil, errRefresh | |
| } | |
| return metaStringValue(updated.Metadata, "access_token"), updated, nil | |
| } | |
| func (e *AntigravityExecutor) maybeRefreshAntigravityCreditsHint(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) { | |
| if e == nil || auth == nil || !antigravityCreditsRetryEnabled(e.cfg) { | |
| return | |
| } | |
| if ctx != nil && ctx.Err() != nil { | |
| return | |
| } | |
| authID := strings.TrimSpace(auth.ID) | |
| if authID == "" { | |
| return | |
| } | |
| if hint, ok := cliproxyauth.GetAntigravityCreditsHint(authID); ok && hint.Known { | |
| return | |
| } | |
| if strings.TrimSpace(accessToken) == "" { | |
| accessToken = metaStringValue(auth.Metadata, "access_token") | |
| } | |
| if strings.TrimSpace(accessToken) == "" { | |
| return | |
| } | |
| state := &antigravityCreditsHintRefreshState{} | |
| if existing, loaded := antigravityCreditsHintRefreshByID.LoadOrStore(authID, state); loaded { | |
| if cast, ok := existing.(*antigravityCreditsHintRefreshState); ok && cast != nil { | |
| state = cast | |
| } else { | |
| antigravityCreditsHintRefreshByID.Delete(authID) | |
| antigravityCreditsHintRefreshByID.Store(authID, state) | |
| } | |
| } | |
| now := time.Now() | |
| if !state.mu.TryLock() { | |
| return | |
| } | |
| if !state.lastAttempt.IsZero() && now.Sub(state.lastAttempt) < antigravityCreditsHintRefreshInterval { | |
| state.mu.Unlock() | |
| return | |
| } | |
| state.lastAttempt = now | |
| refreshCtx := context.Background() | |
| if ctx != nil { | |
| if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil { | |
| refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt) | |
| } | |
| } | |
| refreshCtx, cancel := context.WithTimeout(refreshCtx, antigravityCreditsHintRefreshTimeout) | |
| authCopy := auth.Clone() | |
| go func(state *antigravityCreditsHintRefreshState, auth *cliproxyauth.Auth, token string) { | |
| defer cancel() | |
| defer state.mu.Unlock() | |
| e.updateAntigravityCreditsBalance(refreshCtx, auth, token) | |
| }(state, authCopy, accessToken) | |
| } | |
| 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("Content-Type", "application/x-www-form-urlencoded") | |
| // Real Antigravity uses Go's default User-Agent for OAuth token refresh | |
| httpReq.Header.Set("User-Agent", "Go-http-client/2.0") | |
| httpClient := newAntigravityHTTPClient(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 { | |
| 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 | |
| } | |
| } | |
| return auth, sErr | |
| } | |
| 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 | |
| now := time.Now() | |
| auth.Metadata["timestamp"] = now.UnixMilli() | |
| auth.Metadata["expired"] = now.Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339) | |
| auth.Metadata["type"] = antigravityAuthType | |
| if errProject := e.ensureAntigravityProjectID(ctx, auth, tokenResp.AccessToken); errProject != nil { | |
| log.Warnf("antigravity executor: ensure project id failed: %v", errProject) | |
| } | |
| e.updateAntigravityCreditsBalance(ctx, auth, tokenResp.AccessToken) | |
| return auth, nil | |
| } | |
| func (e *AntigravityExecutor) ensureAntigravityProjectID(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) error { | |
| if auth == nil { | |
| return nil | |
| } | |
| if auth.Metadata["project_id"] != nil { | |
| return nil | |
| } | |
| token := strings.TrimSpace(accessToken) | |
| if token == "" { | |
| token = metaStringValue(auth.Metadata, "access_token") | |
| } | |
| if token == "" { | |
| return nil | |
| } | |
| httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) | |
| projectID, errFetch := sdkAuth.FetchAntigravityProjectID(ctx, token, httpClient) | |
| if errFetch != nil { | |
| return errFetch | |
| } | |
| if strings.TrimSpace(projectID) == "" { | |
| return nil | |
| } | |
| if auth.Metadata == nil { | |
| auth.Metadata = make(map[string]any) | |
| } | |
| auth.Metadata["project_id"] = strings.TrimSpace(projectID) | |
| return nil | |
| } | |
| func (e *AntigravityExecutor) updateAntigravityCreditsBalance(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) { | |
| if auth == nil || strings.TrimSpace(auth.ID) == "" { | |
| return | |
| } | |
| token := strings.TrimSpace(accessToken) | |
| if token == "" { | |
| token = metaStringValue(auth.Metadata, "access_token") | |
| } | |
| if token == "" { | |
| return | |
| } | |
| loadReqBody := `{"metadata":{"ideType":"ANTIGRAVITY","platform":"PLATFORM_UNSPECIFIED","pluginType":"GEMINI"}}` | |
| endpointURL := "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist" | |
| httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, strings.NewReader(loadReqBody)) | |
| if errReq != nil { | |
| log.Debugf("antigravity executor: create loadCodeAssist request error: %v", errReq) | |
| return | |
| } | |
| httpReq.Header.Set("Authorization", "Bearer "+token) | |
| httpReq.Header.Set("Content-Type", "application/json") | |
| httpReq.Header.Set("User-Agent", "google-api-nodejs-client/9.15.1") | |
| httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) | |
| httpResp, errDo := httpClient.Do(httpReq) | |
| if errDo != nil { | |
| log.Debugf("antigravity executor: loadCodeAssist request error: %v", errDo) | |
| return | |
| } | |
| defer func() { | |
| if errClose := httpResp.Body.Close(); errClose != nil { | |
| log.Errorf("antigravity executor: close loadCodeAssist response body error: %v", errClose) | |
| } | |
| }() | |
| bodyBytes, errRead := io.ReadAll(httpResp.Body) | |
| if errRead != nil || httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { | |
| log.Debugf("antigravity executor: loadCodeAssist returned status %d, err=%v", httpResp.StatusCode, errRead) | |
| return | |
| } | |
| authID := strings.TrimSpace(auth.ID) | |
| paidTierID := strings.TrimSpace(gjson.GetBytes(bodyBytes, "paidTier.id").String()) | |
| credits := gjson.GetBytes(bodyBytes, "paidTier.availableCredits") | |
| if !credits.IsArray() { | |
| cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{ | |
| Known: true, | |
| Available: false, | |
| PaidTierID: paidTierID, | |
| UpdatedAt: time.Now(), | |
| }) | |
| return | |
| } | |
| for _, credit := range credits.Array() { | |
| if !strings.EqualFold(credit.Get("creditType").String(), "GOOGLE_ONE_AI") { | |
| continue | |
| } | |
| creditAmount, errCA := strconv.ParseFloat(strings.TrimSpace(credit.Get("creditAmount").String()), 64) | |
| if errCA != nil { | |
| continue | |
| } | |
| minAmount, errMA := strconv.ParseFloat(strings.TrimSpace(credit.Get("minimumCreditAmountForUsage").String()), 64) | |
| if errMA != nil { | |
| continue | |
| } | |
| bal := antigravityCreditsBalance{ | |
| CreditAmount: creditAmount, | |
| MinCreditAmount: minAmount, | |
| PaidTierID: paidTierID, | |
| Known: true, | |
| } | |
| antigravityCreditsBalanceByAuth.Store(authID, bal) | |
| cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{ | |
| Known: true, | |
| Available: creditAmount >= minAmount, | |
| CreditAmount: creditAmount, | |
| MinCreditAmount: minAmount, | |
| PaidTierID: paidTierID, | |
| UpdatedAt: time.Now(), | |
| }) | |
| if creditAmount >= minAmount { | |
| clearAntigravityCreditsPermanentlyDisabled(auth) | |
| } | |
| return | |
| } | |
| } | |
| 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"} | |
| } | |
| base := strings.TrimSuffix(baseURL, "/") | |
| if base == "" { | |
| base = buildBaseURL(auth) | |
| } | |
| path := antigravityGeneratePath | |
| if stream { | |
| path = antigravityStreamPath | |
| } | |
| var requestURL strings.Builder | |
| requestURL.WriteString(base) | |
| requestURL.WriteString(path) | |
| if stream { | |
| if alt != "" { | |
| requestURL.WriteString("?$alt=") | |
| requestURL.WriteString(url.QueryEscape(alt)) | |
| } else { | |
| requestURL.WriteString("?alt=sse") | |
| } | |
| } else if alt != "" { | |
| requestURL.WriteString("?$alt=") | |
| requestURL.WriteString(url.QueryEscape(alt)) | |
| } | |
| // Extract project_id from auth metadata if available | |
| projectID := "" | |
| if auth != nil && auth.Metadata != nil { | |
| if pid, ok := auth.Metadata["project_id"].(string); ok { | |
| projectID = strings.TrimSpace(pid) | |
| } | |
| } | |
| payload = geminiToAntigravity(modelName, payload, projectID) | |
| payload, _ = sjson.SetBytes(payload, "model", modelName) | |
| // Cap maxOutputTokens to model's max_completion_tokens from registry | |
| if maxOut := gjson.GetBytes(payload, "request.generationConfig.maxOutputTokens"); maxOut.Exists() && maxOut.Type == gjson.Number { | |
| if modelInfo := registry.LookupModelInfo(modelName, "antigravity"); modelInfo != nil && modelInfo.MaxCompletionTokens > 0 { | |
| if int(maxOut.Int()) > modelInfo.MaxCompletionTokens { | |
| payload, _ = sjson.SetBytes(payload, "request.generationConfig.maxOutputTokens", modelInfo.MaxCompletionTokens) | |
| } | |
| } | |
| } | |
| useAntigravitySchema := strings.Contains(modelName, "claude") || strings.Contains(modelName, "gemini-3-pro") || strings.Contains(modelName, "gemini-3.1-pro") | |
| var ( | |
| bodyReader io.Reader | |
| payloadLog []byte | |
| ) | |
| if antigravityRequestNeedsSchemaSanitization(payload) { | |
| payloadStr := string(payload) | |
| paths := make([]string, 0) | |
| util.Walk(gjson.Parse(payloadStr), "", "parametersJsonSchema", &paths) | |
| for _, p := range paths { | |
| payloadStr, _ = util.RenameKey(payloadStr, p, p[:len(p)-len("parametersJsonSchema")]+"parameters") | |
| } | |
| if useAntigravitySchema { | |
| payloadStr = util.CleanJSONSchemaForAntigravity(payloadStr) | |
| } else { | |
| payloadStr = util.CleanJSONSchemaForGemini(payloadStr) | |
| } | |
| if strings.Contains(modelName, "claude") { | |
| updated, _ := sjson.SetBytes([]byte(payloadStr), "request.toolConfig.functionCallingConfig.mode", "VALIDATED") | |
| payloadStr = string(updated) | |
| } else { | |
| payloadStr, _ = sjson.Delete(payloadStr, "request.generationConfig.maxOutputTokens") | |
| } | |
| bodyReader = strings.NewReader(payloadStr) | |
| if e.cfg != nil && e.cfg.RequestLog { | |
| payloadLog = []byte(payloadStr) | |
| } | |
| } else { | |
| if strings.Contains(modelName, "claude") { | |
| payload, _ = sjson.SetBytes(payload, "request.toolConfig.functionCallingConfig.mode", "VALIDATED") | |
| } else { | |
| payload, _ = sjson.DeleteBytes(payload, "request.generationConfig.maxOutputTokens") | |
| } | |
| bodyReader = bytes.NewReader(payload) | |
| if e.cfg != nil && e.cfg.RequestLog { | |
| payloadLog = append([]byte(nil), payload...) | |
| } | |
| } | |
| // if useAntigravitySchema { | |
| // systemInstructionPartsResult := gjson.Get(payloadStr, "request.systemInstruction.parts") | |
| // payloadStr, _ = sjson.SetBytes([]byte(payloadStr), "request.systemInstruction.role", "user") | |
| // payloadStr, _ = sjson.SetBytes([]byte(payloadStr), "request.systemInstruction.parts.0.text", systemInstruction) | |
| // payloadStr, _ = sjson.SetBytes([]byte(payloadStr), "request.systemInstruction.parts.1.text", fmt.Sprintf("Please ignore following [ignore]%s[/ignore]", systemInstruction)) | |
| // if systemInstructionPartsResult.Exists() && systemInstructionPartsResult.IsArray() { | |
| // for _, partResult := range systemInstructionPartsResult.Array() { | |
| // payloadStr, _ = sjson.SetRawBytes([]byte(payloadStr), "request.systemInstruction.parts.-1", []byte(partResult.Raw)) | |
| // } | |
| // } | |
| // } | |
| httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bodyReader) | |
| if errReq != nil { | |
| return nil, errReq | |
| } | |
| httpReq.Close = true | |
| httpReq.Header.Set("Content-Type", "application/json") | |
| httpReq.Header.Set("Authorization", "Bearer "+token) | |
| httpReq.Header.Set("User-Agent", resolveUserAgent(auth)) | |
| if host := resolveHost(base); host != "" { | |
| httpReq.Host = host | |
| } | |
| var attrs map[string]string | |
| if auth != nil { | |
| attrs = auth.Attributes | |
| } | |
| util.ApplyCustomHeadersFromAttrs(httpReq, attrs) | |
| var authID, authLabel, authType, authValue string | |
| if auth != nil { | |
| authID = auth.ID | |
| authLabel = auth.Label | |
| authType, authValue = auth.AccountInfo() | |
| } | |
| helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ | |
| URL: requestURL.String(), | |
| Method: http.MethodPost, | |
| Headers: httpReq.Header.Clone(), | |
| Body: payloadLog, | |
| Provider: e.Identifier(), | |
| AuthID: authID, | |
| AuthLabel: authLabel, | |
| AuthType: authType, | |
| AuthValue: authValue, | |
| }) | |
| return httpReq, nil | |
| } | |
| func antigravityRequestNeedsSchemaSanitization(payload []byte) bool { | |
| if gjson.GetBytes(payload, "request.tools.0").Exists() { | |
| return true | |
| } | |
| if gjson.GetBytes(payload, "request.generationConfig.responseJsonSchema").Exists() { | |
| return true | |
| } | |
| if gjson.GetBytes(payload, "request.generationConfig.responseSchema").Exists() { | |
| return true | |
| } | |
| return false | |
| } | |
| func tokenExpiry(metadata map[string]any) time.Time { | |
| if metadata == nil { | |
| return time.Time{} | |
| } | |
| if expStr, ok := metadata["expired"].(string); ok { | |
| expStr = strings.TrimSpace(expStr) | |
| if expStr != "" { | |
| if parsed, errParse := time.Parse(time.RFC3339, expStr); errParse == nil { | |
| return parsed | |
| } | |
| } | |
| } | |
| expiresIn, hasExpires := int64Value(metadata["expires_in"]) | |
| tsMs, hasTimestamp := int64Value(metadata["timestamp"]) | |
| if hasExpires && hasTimestamp { | |
| return time.Unix(0, tsMs*int64(time.Millisecond)).Add(time.Duration(expiresIn) * time.Second) | |
| } | |
| return time.Time{} | |
| } | |
| func metaStringValue(metadata map[string]any, key string) string { | |
| if metadata == nil { | |
| return "" | |
| } | |
| if v, ok := metadata[key]; ok { | |
| switch typed := v.(type) { | |
| case string: | |
| return strings.TrimSpace(typed) | |
| case []byte: | |
| return strings.TrimSpace(string(typed)) | |
| } | |
| } | |
| return "" | |
| } | |
| func int64Value(value any) (int64, bool) { | |
| switch typed := value.(type) { | |
| case int: | |
| return int64(typed), true | |
| case int64: | |
| return typed, true | |
| case float64: | |
| return int64(typed), true | |
| case json.Number: | |
| if i, errParse := typed.Int64(); errParse == nil { | |
| return i, true | |
| } | |
| case string: | |
| if strings.TrimSpace(typed) == "" { | |
| return 0, false | |
| } | |
| if i, errParse := strconv.ParseInt(strings.TrimSpace(typed), 10, 64); errParse == nil { | |
| return i, true | |
| } | |
| } | |
| return 0, false | |
| } | |
| func buildBaseURL(auth *cliproxyauth.Auth) string { | |
| if baseURLs := antigravityBaseURLFallbackOrder(auth); len(baseURLs) > 0 { | |
| return baseURLs[0] | |
| } | |
| return antigravityBaseURLDaily | |
| } | |
| func resolveHost(base string) string { | |
| parsed, errParse := url.Parse(base) | |
| if errParse != nil { | |
| return "" | |
| } | |
| if parsed.Host != "" { | |
| return parsed.Host | |
| } | |
| return strings.TrimPrefix(strings.TrimPrefix(base, "https://"), "http://") | |
| } | |
| func resolveUserAgent(auth *cliproxyauth.Auth) string { | |
| if auth != nil { | |
| if auth.Attributes != nil { | |
| if ua := strings.TrimSpace(auth.Attributes["user_agent"]); ua != "" { | |
| return ua | |
| } | |
| } | |
| if auth.Metadata != nil { | |
| if ua, ok := auth.Metadata["user_agent"].(string); ok && strings.TrimSpace(ua) != "" { | |
| return strings.TrimSpace(ua) | |
| } | |
| } | |
| } | |
| return misc.AntigravityUserAgent() | |
| } | |
| func antigravityRetryAttempts(auth *cliproxyauth.Auth, cfg *config.Config) int { | |
| retry := 0 | |
| if cfg != nil { | |
| retry = cfg.RequestRetry | |
| } | |
| if auth != nil { | |
| if override, ok := auth.RequestRetryOverride(); ok { | |
| retry = override | |
| } | |
| } | |
| if retry < 0 { | |
| retry = 0 | |
| } | |
| attempts := retry + 1 | |
| if attempts < 1 { | |
| return 1 | |
| } | |
| return attempts | |
| } | |
| func antigravityShouldRetryNoCapacity(statusCode int, body []byte) bool { | |
| if statusCode != http.StatusServiceUnavailable { | |
| return false | |
| } | |
| if len(body) == 0 { | |
| return false | |
| } | |
| msg := strings.ToLower(string(body)) | |
| return strings.Contains(msg, "no capacity available") | |
| } | |
| func antigravityShouldRetryTransientResourceExhausted429(statusCode int, body []byte) bool { | |
| if statusCode != http.StatusTooManyRequests { | |
| return false | |
| } | |
| if len(body) == 0 { | |
| return false | |
| } | |
| if classifyAntigravity429(body) != antigravity429Unknown { | |
| return false | |
| } | |
| status := strings.TrimSpace(gjson.GetBytes(body, "error.status").String()) | |
| if !strings.EqualFold(status, "RESOURCE_EXHAUSTED") { | |
| return false | |
| } | |
| msg := strings.ToLower(string(body)) | |
| return strings.Contains(msg, "resource has been exhausted") | |
| } | |
| func antigravityShouldRetrySoftRateLimit(statusCode int, body []byte) bool { | |
| if statusCode != http.StatusTooManyRequests { | |
| return false | |
| } | |
| return decideAntigravity429(body).kind == antigravity429DecisionSoftRetry | |
| } | |
| func antigravitySoftRateLimitDelay(attempt int) time.Duration { | |
| if attempt < 0 { | |
| attempt = 0 | |
| } | |
| base := time.Duration(attempt+1) * 500 * time.Millisecond | |
| if base > 3*time.Second { | |
| base = 3 * time.Second | |
| } | |
| return base | |
| } | |
| func antigravityShortCooldownKey(auth *cliproxyauth.Auth, modelName string) string { | |
| if auth == nil { | |
| return "" | |
| } | |
| authID := strings.TrimSpace(auth.ID) | |
| modelName = strings.TrimSpace(modelName) | |
| if authID == "" || modelName == "" { | |
| return "" | |
| } | |
| return authID + "|" + modelName + "|sc" | |
| } | |
| func antigravityIsInShortCooldown(auth *cliproxyauth.Auth, modelName string, now time.Time) (bool, time.Duration) { | |
| key := antigravityShortCooldownKey(auth, modelName) | |
| if key == "" { | |
| return false, 0 | |
| } | |
| value, ok := antigravityShortCooldownByAuth.Load(key) | |
| if !ok { | |
| return false, 0 | |
| } | |
| until, ok := value.(time.Time) | |
| if !ok || until.IsZero() { | |
| antigravityShortCooldownByAuth.Delete(key) | |
| return false, 0 | |
| } | |
| remaining := until.Sub(now) | |
| if remaining <= 0 { | |
| antigravityShortCooldownByAuth.Delete(key) | |
| return false, 0 | |
| } | |
| return true, remaining | |
| } | |
| func markAntigravityShortCooldown(auth *cliproxyauth.Auth, modelName string, now time.Time, duration time.Duration) { | |
| key := antigravityShortCooldownKey(auth, modelName) | |
| if key == "" { | |
| return | |
| } | |
| antigravityShortCooldownByAuth.Store(key, now.Add(duration)) | |
| } | |
| func antigravityNoCapacityRetryDelay(attempt int) time.Duration { | |
| if attempt < 0 { | |
| attempt = 0 | |
| } | |
| delay := time.Duration(attempt+1) * 250 * time.Millisecond | |
| if delay > 2*time.Second { | |
| delay = 2 * time.Second | |
| } | |
| return delay | |
| } | |
| func antigravityTransient429RetryDelay(attempt int) time.Duration { | |
| if attempt < 0 { | |
| attempt = 0 | |
| } | |
| delay := time.Duration(attempt+1) * 100 * time.Millisecond | |
| if delay > 500*time.Millisecond { | |
| delay = 500 * time.Millisecond | |
| } | |
| return delay | |
| } | |
| func antigravityInstantRetryDelay(wait time.Duration) time.Duration { | |
| if wait <= 0 { | |
| return 0 | |
| } | |
| return wait + 800*time.Millisecond | |
| } | |
| func antigravityWait(ctx context.Context, wait time.Duration) error { | |
| if wait <= 0 { | |
| return nil | |
| } | |
| timer := time.NewTimer(wait) | |
| defer timer.Stop() | |
| select { | |
| case <-ctx.Done(): | |
| return ctx.Err() | |
| case <-timer.C: | |
| return nil | |
| } | |
| } | |
| var antigravityBaseURLFallbackOrder = func(auth *cliproxyauth.Auth) []string { | |
| if base := resolveCustomAntigravityBaseURL(auth); base != "" { | |
| return []string{base} | |
| } | |
| return []string{ | |
| antigravityBaseURLProd, | |
| antigravityBaseURLDaily, | |
| antigravitySandboxBaseURLDaily, | |
| } | |
| } | |
| func resolveCustomAntigravityBaseURL(auth *cliproxyauth.Auth) string { | |
| if auth == nil { | |
| return "" | |
| } | |
| if auth.Attributes != nil { | |
| if v := strings.TrimSpace(auth.Attributes["base_url"]); v != "" { | |
| return strings.TrimSuffix(v, "/") | |
| } | |
| } | |
| if auth.Metadata != nil { | |
| if v, ok := auth.Metadata["base_url"].(string); ok { | |
| v = strings.TrimSpace(v) | |
| if v != "" { | |
| return strings.TrimSuffix(v, "/") | |
| } | |
| } | |
| } | |
| return "" | |
| } | |
| func geminiToAntigravity(modelName string, payload []byte, projectID string) []byte { | |
| template := payload | |
| template, _ = sjson.SetBytes(template, "model", modelName) | |
| template, _ = sjson.SetBytes(template, "userAgent", "antigravity") | |
| isImageModel := strings.Contains(modelName, "image") | |
| var reqType string | |
| if isImageModel { | |
| reqType = "image_gen" | |
| } else { | |
| reqType = "agent" | |
| } | |
| template, _ = sjson.SetBytes(template, "requestType", reqType) | |
| // Use real project ID from auth if available, otherwise generate random (legacy fallback) | |
| if projectID != "" { | |
| template, _ = sjson.SetBytes(template, "project", projectID) | |
| } else { | |
| template, _ = sjson.SetBytes(template, "project", generateProjectID()) | |
| } | |
| if isImageModel { | |
| template, _ = sjson.SetBytes(template, "requestId", generateImageGenRequestID()) | |
| } else { | |
| template, _ = sjson.SetBytes(template, "requestId", generateRequestID()) | |
| template, _ = sjson.SetBytes(template, "request.sessionId", generateStableSessionID(payload)) | |
| } | |
| template, _ = sjson.DeleteBytes(template, "request.safetySettings") | |
| if toolConfig := gjson.GetBytes(template, "toolConfig"); toolConfig.Exists() && !gjson.GetBytes(template, "request.toolConfig").Exists() { | |
| template, _ = sjson.SetRawBytes(template, "request.toolConfig", []byte(toolConfig.Raw)) | |
| template, _ = sjson.DeleteBytes(template, "toolConfig") | |
| } | |
| return template | |
| } | |
| func generateRequestID() string { | |
| return "agent-" + uuid.NewString() | |
| } | |
| func generateImageGenRequestID() string { | |
| return fmt.Sprintf("image_gen/%d/%s/12", time.Now().UnixMilli(), uuid.NewString()) | |
| } | |
| func generateSessionID() string { | |
| randSourceMutex.Lock() | |
| n := randSource.Int63n(9_000_000_000_000_000_000) | |
| randSourceMutex.Unlock() | |
| return "-" + strconv.FormatInt(n, 10) | |
| } | |
| func generateStableSessionID(payload []byte) string { | |
| contents := gjson.GetBytes(payload, "request.contents") | |
| if contents.IsArray() { | |
| for _, content := range contents.Array() { | |
| if content.Get("role").String() == "user" { | |
| text := content.Get("parts.0.text").String() | |
| if text != "" { | |
| h := sha256.Sum256([]byte(text)) | |
| n := int64(binary.BigEndian.Uint64(h[:8])) & 0x7FFFFFFFFFFFFFFF | |
| return "-" + strconv.FormatInt(n, 10) | |
| } | |
| } | |
| } | |
| } | |
| return generateSessionID() | |
| } | |
| func generateProjectID() string { | |
| adjectives := []string{"useful", "bright", "swift", "calm", "bold"} | |
| nouns := []string{"fuze", "wave", "spark", "flow", "core"} | |
| randSourceMutex.Lock() | |
| adj := adjectives[randSource.Intn(len(adjectives))] | |
| noun := nouns[randSource.Intn(len(nouns))] | |
| randSourceMutex.Unlock() | |
| randomPart := strings.ToLower(uuid.NewString())[:5] | |
| return adj + "-" + noun + "-" + randomPart | |
| } | |