|
|
|
|
|
|
|
|
|
|
|
package executor |
|
|
|
|
|
import ( |
|
|
"bufio" |
|
|
"bytes" |
|
|
"context" |
|
|
"crypto/sha256" |
|
|
"encoding/binary" |
|
|
"encoding/json" |
|
|
"fmt" |
|
|
"io" |
|
|
"math/rand" |
|
|
"net/http" |
|
|
"net/url" |
|
|
"strconv" |
|
|
"strings" |
|
|
"sync" |
|
|
"time" |
|
|
|
|
|
"github.com/google/uuid" |
|
|
"github.com/router-for-me/CLIProxyAPI/v6/internal/config" |
|
|
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry" |
|
|
"github.com/router-for-me/CLIProxyAPI/v6/internal/util" |
|
|
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" |
|
|
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" |
|
|
sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator" |
|
|
log "github.com/sirupsen/logrus" |
|
|
"github.com/tidwall/gjson" |
|
|
"github.com/tidwall/sjson" |
|
|
) |
|
|
|
|
|
const ( |
|
|
antigravityBaseURLDaily = "https://daily-cloudcode-pa.googleapis.com" |
|
|
antigravitySandboxBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com" |
|
|
antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com" |
|
|
antigravityCountTokensPath = "/v1internal:countTokens" |
|
|
antigravityStreamPath = "/v1internal:streamGenerateContent" |
|
|
antigravityGeneratePath = "/v1internal:generateContent" |
|
|
antigravityModelsPath = "/v1internal:fetchAvailableModels" |
|
|
antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" |
|
|
antigravityClientSecret = "YOUR_ANTIGRAVITY_CLIENT_SECRET" |
|
|
defaultAntigravityAgent = "antigravity/1.104.0 darwin/arm64" |
|
|
antigravityAuthType = "antigravity" |
|
|
refreshSkew = 3000 * time.Second |
|
|
) |
|
|
|
|
|
var ( |
|
|
randSource = rand.New(rand.NewSource(time.Now().UnixNano())) |
|
|
randSourceMutex sync.Mutex |
|
|
) |
|
|
|
|
|
|
|
|
type AntigravityExecutor struct { |
|
|
cfg *config.Config |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
func NewAntigravityExecutor(cfg *config.Config) *AntigravityExecutor { |
|
|
return &AntigravityExecutor{cfg: cfg} |
|
|
} |
|
|
|
|
|
|
|
|
func (e *AntigravityExecutor) Identifier() string { return antigravityAuthType } |
|
|
|
|
|
|
|
|
func (e *AntigravityExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil } |
|
|
|
|
|
|
|
|
func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { |
|
|
isClaude := strings.Contains(strings.ToLower(req.Model), "claude") |
|
|
if isClaude { |
|
|
return e.executeClaudeNonStream(ctx, auth, req, opts) |
|
|
} |
|
|
|
|
|
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) |
|
|
if errToken != nil { |
|
|
return resp, errToken |
|
|
} |
|
|
if updatedAuth != nil { |
|
|
auth = updatedAuth |
|
|
} |
|
|
|
|
|
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth) |
|
|
defer reporter.trackFailure(ctx, &err) |
|
|
|
|
|
from := opts.SourceFormat |
|
|
to := sdktranslator.FromString("antigravity") |
|
|
originalPayload := bytes.Clone(req.Payload) |
|
|
if len(opts.OriginalRequest) > 0 { |
|
|
originalPayload = bytes.Clone(opts.OriginalRequest) |
|
|
} |
|
|
originalTranslated := sdktranslator.TranslateRequest(from, to, req.Model, originalPayload, false) |
|
|
translated := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false) |
|
|
|
|
|
translated = ApplyThinkingMetadataCLI(translated, req.Metadata, req.Model) |
|
|
translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated) |
|
|
translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, translated) |
|
|
translated = normalizeAntigravityThinking(req.Model, translated, isClaude) |
|
|
translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated, originalTranslated) |
|
|
|
|
|
baseURLs := antigravityBaseURLFallbackOrder(auth) |
|
|
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0) |
|
|
|
|
|
var lastStatus int |
|
|
var lastBody []byte |
|
|
var lastErr error |
|
|
|
|
|
for idx, baseURL := range baseURLs { |
|
|
httpReq, errReq := e.buildRequest(ctx, auth, token, req.Model, translated, false, opts.Alt, baseURL) |
|
|
if errReq != nil { |
|
|
err = errReq |
|
|
return resp, err |
|
|
} |
|
|
|
|
|
httpResp, errDo := httpClient.Do(httpReq) |
|
|
if errDo != nil { |
|
|
recordAPIResponseError(ctx, e.cfg, errDo) |
|
|
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, ¶m) |
|
|
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 |
|
|
} |
|
|
|
|
|
|
|
|
func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { |
|
|
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) |
|
|
if errToken != nil { |
|
|
return resp, errToken |
|
|
} |
|
|
if updatedAuth != nil { |
|
|
auth = updatedAuth |
|
|
} |
|
|
|
|
|
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth) |
|
|
defer reporter.trackFailure(ctx, &err) |
|
|
|
|
|
from := opts.SourceFormat |
|
|
to := sdktranslator.FromString("antigravity") |
|
|
originalPayload := bytes.Clone(req.Payload) |
|
|
if len(opts.OriginalRequest) > 0 { |
|
|
originalPayload = bytes.Clone(opts.OriginalRequest) |
|
|
} |
|
|
originalTranslated := sdktranslator.TranslateRequest(from, to, req.Model, originalPayload, true) |
|
|
translated := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true) |
|
|
|
|
|
translated = ApplyThinkingMetadataCLI(translated, req.Metadata, req.Model) |
|
|
translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated) |
|
|
translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, translated) |
|
|
translated = normalizeAntigravityThinking(req.Model, translated, true) |
|
|
translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated, originalTranslated) |
|
|
|
|
|
baseURLs := antigravityBaseURLFallbackOrder(auth) |
|
|
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0) |
|
|
|
|
|
var lastStatus int |
|
|
var lastBody []byte |
|
|
var lastErr error |
|
|
|
|
|
for idx, baseURL := range baseURLs { |
|
|
httpReq, errReq := e.buildRequest(ctx, auth, token, req.Model, translated, true, opts.Alt, baseURL) |
|
|
if errReq != nil { |
|
|
err = errReq |
|
|
return resp, err |
|
|
} |
|
|
|
|
|
httpResp, errDo := httpClient.Do(httpReq) |
|
|
if errDo != nil { |
|
|
recordAPIResponseError(ctx, e.cfg, errDo) |
|
|
lastStatus = 0 |
|
|
lastBody = nil |
|
|
lastErr = errDo |
|
|
if idx+1 < len(baseURLs) { |
|
|
log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) |
|
|
continue |
|
|
} |
|
|
err = errDo |
|
|
return resp, err |
|
|
} |
|
|
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) |
|
|
if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { |
|
|
bodyBytes, errRead := io.ReadAll(httpResp.Body) |
|
|
if errClose := httpResp.Body.Close(); errClose != nil { |
|
|
log.Errorf("antigravity executor: close response body error: %v", errClose) |
|
|
} |
|
|
if errRead != nil { |
|
|
recordAPIResponseError(ctx, e.cfg, errRead) |
|
|
lastStatus = 0 |
|
|
lastBody = nil |
|
|
lastErr = errRead |
|
|
if idx+1 < len(baseURLs) { |
|
|
log.Debugf("antigravity executor: read error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) |
|
|
continue |
|
|
} |
|
|
err = errRead |
|
|
return resp, err |
|
|
} |
|
|
appendAPIResponseChunk(ctx, e.cfg, bodyBytes) |
|
|
lastStatus = httpResp.StatusCode |
|
|
lastBody = append([]byte(nil), bodyBytes...) |
|
|
lastErr = nil |
|
|
if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) { |
|
|
log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) |
|
|
continue |
|
|
} |
|
|
err = statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)} |
|
|
return resp, err |
|
|
} |
|
|
|
|
|
out := make(chan cliproxyexecutor.StreamChunk) |
|
|
go func(resp *http.Response) { |
|
|
defer close(out) |
|
|
defer func() { |
|
|
if errClose := resp.Body.Close(); errClose != nil { |
|
|
log.Errorf("antigravity executor: close response body error: %v", errClose) |
|
|
} |
|
|
}() |
|
|
scanner := bufio.NewScanner(resp.Body) |
|
|
scanner.Buffer(nil, streamScannerBuffer) |
|
|
for scanner.Scan() { |
|
|
line := scanner.Bytes() |
|
|
appendAPIResponseChunk(ctx, e.cfg, line) |
|
|
|
|
|
|
|
|
|
|
|
line = FilterSSEUsageMetadata(line) |
|
|
|
|
|
payload := jsonPayload(line) |
|
|
if payload == nil { |
|
|
continue |
|
|
} |
|
|
|
|
|
if detail, ok := parseAntigravityStreamUsage(payload); ok { |
|
|
reporter.publish(ctx, detail) |
|
|
} |
|
|
|
|
|
out <- cliproxyexecutor.StreamChunk{Payload: payload} |
|
|
} |
|
|
if errScan := scanner.Err(); errScan != nil { |
|
|
recordAPIResponseError(ctx, e.cfg, errScan) |
|
|
reporter.publishFailure(ctx) |
|
|
out <- cliproxyexecutor.StreamChunk{Err: errScan} |
|
|
} else { |
|
|
reporter.ensurePublished(ctx) |
|
|
} |
|
|
}(httpResp) |
|
|
|
|
|
var buffer bytes.Buffer |
|
|
for chunk := range out { |
|
|
if chunk.Err != nil { |
|
|
return resp, chunk.Err |
|
|
} |
|
|
if len(chunk.Payload) > 0 { |
|
|
_, _ = buffer.Write(chunk.Payload) |
|
|
_, _ = buffer.Write([]byte("\n")) |
|
|
} |
|
|
} |
|
|
resp = cliproxyexecutor.Response{Payload: e.convertStreamToNonStream(buffer.Bytes())} |
|
|
|
|
|
reporter.publish(ctx, parseAntigravityUsage(resp.Payload)) |
|
|
var param any |
|
|
converted := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), translated, resp.Payload, ¶m) |
|
|
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 |
|
|
} |
|
|
|
|
|
func (e *AntigravityExecutor) convertStreamToNonStream(stream []byte) []byte { |
|
|
responseTemplate := "" |
|
|
var traceID string |
|
|
var finishReason string |
|
|
var modelVersion string |
|
|
var responseID string |
|
|
var role string |
|
|
var usageRaw string |
|
|
parts := make([]map[string]interface{}, 0) |
|
|
var pendingKind string |
|
|
var pendingText strings.Builder |
|
|
var pendingThoughtSig string |
|
|
|
|
|
flushPending := func() { |
|
|
if pendingKind == "" { |
|
|
return |
|
|
} |
|
|
text := pendingText.String() |
|
|
switch pendingKind { |
|
|
case "text": |
|
|
if strings.TrimSpace(text) == "" { |
|
|
pendingKind = "" |
|
|
pendingText.Reset() |
|
|
pendingThoughtSig = "" |
|
|
return |
|
|
} |
|
|
parts = append(parts, map[string]interface{}{"text": text}) |
|
|
case "thought": |
|
|
if strings.TrimSpace(text) == "" && pendingThoughtSig == "" { |
|
|
pendingKind = "" |
|
|
pendingText.Reset() |
|
|
pendingThoughtSig = "" |
|
|
return |
|
|
} |
|
|
part := map[string]interface{}{"thought": true} |
|
|
part["text"] = text |
|
|
if pendingThoughtSig != "" { |
|
|
part["thoughtSignature"] = pendingThoughtSig |
|
|
} |
|
|
parts = append(parts, part) |
|
|
} |
|
|
pendingKind = "" |
|
|
pendingText.Reset() |
|
|
pendingThoughtSig = "" |
|
|
} |
|
|
|
|
|
normalizePart := func(partResult gjson.Result) map[string]interface{} { |
|
|
var m map[string]interface{} |
|
|
_ = json.Unmarshal([]byte(partResult.Raw), &m) |
|
|
if m == nil { |
|
|
m = map[string]interface{}{} |
|
|
} |
|
|
sig := partResult.Get("thoughtSignature").String() |
|
|
if sig == "" { |
|
|
sig = partResult.Get("thought_signature").String() |
|
|
} |
|
|
if sig != "" { |
|
|
m["thoughtSignature"] = sig |
|
|
delete(m, "thought_signature") |
|
|
} |
|
|
if inlineData, ok := m["inline_data"]; ok { |
|
|
m["inlineData"] = inlineData |
|
|
delete(m, "inline_data") |
|
|
} |
|
|
return m |
|
|
} |
|
|
|
|
|
for _, line := range bytes.Split(stream, []byte("\n")) { |
|
|
trimmed := bytes.TrimSpace(line) |
|
|
if len(trimmed) == 0 || !gjson.ValidBytes(trimmed) { |
|
|
continue |
|
|
} |
|
|
|
|
|
root := gjson.ParseBytes(trimmed) |
|
|
responseNode := root.Get("response") |
|
|
if !responseNode.Exists() { |
|
|
if root.Get("candidates").Exists() { |
|
|
responseNode = root |
|
|
} else { |
|
|
continue |
|
|
} |
|
|
} |
|
|
responseTemplate = responseNode.Raw |
|
|
|
|
|
if traceResult := root.Get("traceId"); traceResult.Exists() && traceResult.String() != "" { |
|
|
traceID = traceResult.String() |
|
|
} |
|
|
|
|
|
if roleResult := responseNode.Get("candidates.0.content.role"); roleResult.Exists() { |
|
|
role = roleResult.String() |
|
|
} |
|
|
|
|
|
if finishResult := responseNode.Get("candidates.0.finishReason"); finishResult.Exists() && finishResult.String() != "" { |
|
|
finishReason = finishResult.String() |
|
|
} |
|
|
|
|
|
if modelResult := responseNode.Get("modelVersion"); modelResult.Exists() && modelResult.String() != "" { |
|
|
modelVersion = modelResult.String() |
|
|
} |
|
|
if responseIDResult := responseNode.Get("responseId"); responseIDResult.Exists() && responseIDResult.String() != "" { |
|
|
responseID = responseIDResult.String() |
|
|
} |
|
|
if usageResult := responseNode.Get("usageMetadata"); usageResult.Exists() { |
|
|
usageRaw = usageResult.Raw |
|
|
} else if usageResult := root.Get("usageMetadata"); usageResult.Exists() { |
|
|
usageRaw = usageResult.Raw |
|
|
} |
|
|
|
|
|
if partsResult := responseNode.Get("candidates.0.content.parts"); partsResult.IsArray() { |
|
|
for _, part := range partsResult.Array() { |
|
|
hasFunctionCall := part.Get("functionCall").Exists() |
|
|
hasInlineData := part.Get("inlineData").Exists() || part.Get("inline_data").Exists() |
|
|
sig := part.Get("thoughtSignature").String() |
|
|
if sig == "" { |
|
|
sig = part.Get("thought_signature").String() |
|
|
} |
|
|
text := part.Get("text").String() |
|
|
thought := part.Get("thought").Bool() |
|
|
|
|
|
if hasFunctionCall || hasInlineData { |
|
|
flushPending() |
|
|
parts = append(parts, normalizePart(part)) |
|
|
continue |
|
|
} |
|
|
|
|
|
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) |
|
|
responseTemplate, _ = sjson.SetRaw(responseTemplate, "candidates.0.content.parts", string(partsJSON)) |
|
|
if role != "" { |
|
|
responseTemplate, _ = sjson.Set(responseTemplate, "candidates.0.content.role", role) |
|
|
} |
|
|
if finishReason != "" { |
|
|
responseTemplate, _ = sjson.Set(responseTemplate, "candidates.0.finishReason", finishReason) |
|
|
} |
|
|
if modelVersion != "" { |
|
|
responseTemplate, _ = sjson.Set(responseTemplate, "modelVersion", modelVersion) |
|
|
} |
|
|
if responseID != "" { |
|
|
responseTemplate, _ = sjson.Set(responseTemplate, "responseId", responseID) |
|
|
} |
|
|
if usageRaw != "" { |
|
|
responseTemplate, _ = sjson.SetRaw(responseTemplate, "usageMetadata", usageRaw) |
|
|
} else if !gjson.Get(responseTemplate, "usageMetadata").Exists() { |
|
|
responseTemplate, _ = sjson.Set(responseTemplate, "usageMetadata.promptTokenCount", 0) |
|
|
responseTemplate, _ = sjson.Set(responseTemplate, "usageMetadata.candidatesTokenCount", 0) |
|
|
responseTemplate, _ = sjson.Set(responseTemplate, "usageMetadata.totalTokenCount", 0) |
|
|
} |
|
|
|
|
|
output := `{"response":{},"traceId":""}` |
|
|
output, _ = sjson.SetRaw(output, "response", responseTemplate) |
|
|
if traceID != "" { |
|
|
output, _ = sjson.Set(output, "traceId", traceID) |
|
|
} |
|
|
return []byte(output) |
|
|
} |
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
isClaude := strings.Contains(strings.ToLower(req.Model), "claude") |
|
|
|
|
|
from := opts.SourceFormat |
|
|
to := sdktranslator.FromString("antigravity") |
|
|
originalPayload := bytes.Clone(req.Payload) |
|
|
if len(opts.OriginalRequest) > 0 { |
|
|
originalPayload = bytes.Clone(opts.OriginalRequest) |
|
|
} |
|
|
originalTranslated := sdktranslator.TranslateRequest(from, to, req.Model, originalPayload, true) |
|
|
translated := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true) |
|
|
|
|
|
translated = ApplyThinkingMetadataCLI(translated, req.Metadata, req.Model) |
|
|
translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated) |
|
|
translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, translated) |
|
|
translated = normalizeAntigravityThinking(req.Model, translated, isClaude) |
|
|
translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated, originalTranslated) |
|
|
|
|
|
baseURLs := antigravityBaseURLFallbackOrder(auth) |
|
|
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0) |
|
|
|
|
|
var lastStatus int |
|
|
var lastBody []byte |
|
|
var lastErr error |
|
|
|
|
|
for idx, baseURL := range baseURLs { |
|
|
httpReq, errReq := e.buildRequest(ctx, auth, token, req.Model, translated, 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) |
|
|
|
|
|
|
|
|
|
|
|
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), ¶m) |
|
|
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]"), ¶m) |
|
|
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 |
|
|
} |
|
|
|
|
|
|
|
|
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 |
|
|
} |
|
|
|
|
|
|
|
|
func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { |
|
|
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"} |
|
|
} |
|
|
|
|
|
from := opts.SourceFormat |
|
|
to := sdktranslator.FromString("antigravity") |
|
|
respCtx := context.WithValue(ctx, "alt", opts.Alt) |
|
|
|
|
|
isClaude := strings.Contains(strings.ToLower(req.Model), "claude") |
|
|
|
|
|
baseURLs := antigravityBaseURLFallbackOrder(auth) |
|
|
httpClient := newProxyAwareHTTPClient(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 { |
|
|
payload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false) |
|
|
payload = ApplyThinkingMetadataCLI(payload, req.Metadata, req.Model) |
|
|
payload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, payload) |
|
|
payload = normalizeAntigravityThinking(req.Model, payload, isClaude) |
|
|
payload = deleteJSONField(payload, "project") |
|
|
payload = deleteJSONField(payload, "model") |
|
|
payload = deleteJSONField(payload, "request.safetySettings") |
|
|
|
|
|
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.Header.Set("Content-Type", "application/json") |
|
|
httpReq.Header.Set("Authorization", "Bearer "+token) |
|
|
httpReq.Header.Set("User-Agent", resolveUserAgent(auth)) |
|
|
httpReq.Header.Set("Accept", "application/json") |
|
|
if host := resolveHost(base); host != "" { |
|
|
httpReq.Host = host |
|
|
} |
|
|
|
|
|
recordAPIRequest(ctx, e.cfg, 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 { |
|
|
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 |
|
|
} |
|
|
return cliproxyexecutor.Response{}, errDo |
|
|
} |
|
|
|
|
|
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) |
|
|
return cliproxyexecutor.Response{}, errRead |
|
|
} |
|
|
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: []byte(translated)}, 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 |
|
|
} |
|
|
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)} |
|
|
} |
|
|
|
|
|
switch { |
|
|
case lastStatus != 0: |
|
|
return cliproxyexecutor.Response{}, statusErr{code: lastStatus, msg: string(lastBody)} |
|
|
case lastErr != nil: |
|
|
return cliproxyexecutor.Response{}, lastErr |
|
|
default: |
|
|
return cliproxyexecutor.Response{}, statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
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() |
|
|
modelConfig := registry.GetAntigravityModelConfig() |
|
|
models := make([]*registry.ModelInfo, 0, len(result.Map())) |
|
|
for originalName := range result.Map() { |
|
|
aliasName := modelName2Alias(originalName) |
|
|
if aliasName != "" { |
|
|
cfg := modelConfig[aliasName] |
|
|
modelName := aliasName |
|
|
if cfg != nil && cfg.Name != "" { |
|
|
modelName = cfg.Name |
|
|
} |
|
|
modelInfo := ®istry.ModelInfo{ |
|
|
ID: aliasName, |
|
|
Name: modelName, |
|
|
Description: aliasName, |
|
|
DisplayName: aliasName, |
|
|
Version: aliasName, |
|
|
Object: "model", |
|
|
Created: now, |
|
|
OwnedBy: antigravityAuthType, |
|
|
Type: antigravityAuthType, |
|
|
} |
|
|
|
|
|
if cfg != nil { |
|
|
if cfg.Thinking != nil { |
|
|
modelInfo.Thinking = cfg.Thinking |
|
|
} |
|
|
if cfg.MaxCompletionTokens > 0 { |
|
|
modelInfo.MaxCompletionTokens = cfg.MaxCompletionTokens |
|
|
} |
|
|
} |
|
|
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"} |
|
|
} |
|
|
|
|
|
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)) |
|
|
} |
|
|
|
|
|
|
|
|
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", alias2ModelName(modelName)) |
|
|
|
|
|
if strings.Contains(modelName, "claude") { |
|
|
strJSON := string(payload) |
|
|
paths := make([]string, 0) |
|
|
util.Walk(gjson.ParseBytes(payload), "", "parametersJsonSchema", &paths) |
|
|
for _, p := range paths { |
|
|
strJSON, _ = util.RenameKey(strJSON, p, p[:len(p)-len("parametersJsonSchema")]+"parameters") |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
strJSON = util.CleanJSONSchemaForAntigravity(strJSON) |
|
|
|
|
|
payload = []byte(strJSON) |
|
|
} |
|
|
|
|
|
httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bytes.NewReader(payload)) |
|
|
if errReq != nil { |
|
|
return nil, errReq |
|
|
} |
|
|
httpReq.Header.Set("Content-Type", "application/json") |
|
|
httpReq.Header.Set("Authorization", "Bearer "+token) |
|
|
httpReq.Header.Set("User-Agent", resolveUserAgent(auth)) |
|
|
if stream { |
|
|
httpReq.Header.Set("Accept", "text/event-stream") |
|
|
} else { |
|
|
httpReq.Header.Set("Accept", "application/json") |
|
|
} |
|
|
if host := resolveHost(base); host != "" { |
|
|
httpReq.Host = host |
|
|
} |
|
|
|
|
|
var authID, authLabel, authType, authValue string |
|
|
if auth != nil { |
|
|
authID = auth.ID |
|
|
authLabel = auth.Label |
|
|
authType, authValue = auth.AccountInfo() |
|
|
} |
|
|
recordAPIRequest(ctx, e.cfg, upstreamRequestLog{ |
|
|
URL: requestURL.String(), |
|
|
Method: http.MethodPost, |
|
|
Headers: httpReq.Header.Clone(), |
|
|
Body: payload, |
|
|
Provider: e.Identifier(), |
|
|
AuthID: authID, |
|
|
AuthLabel: authLabel, |
|
|
AuthType: authType, |
|
|
AuthValue: authValue, |
|
|
}) |
|
|
|
|
|
return httpReq, nil |
|
|
} |
|
|
|
|
|
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 defaultAntigravityAgent |
|
|
} |
|
|
|
|
|
func antigravityBaseURLFallbackOrder(auth *cliproxyauth.Auth) []string { |
|
|
if base := resolveCustomAntigravityBaseURL(auth); base != "" { |
|
|
return []string{base} |
|
|
} |
|
|
return []string{ |
|
|
antigravityBaseURLDaily, |
|
|
antigravitySandboxBaseURLDaily, |
|
|
antigravityBaseURLProd, |
|
|
} |
|
|
} |
|
|
|
|
|
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, _ := sjson.Set(string(payload), "model", modelName) |
|
|
template, _ = sjson.Set(template, "userAgent", "antigravity") |
|
|
|
|
|
|
|
|
if projectID != "" { |
|
|
template, _ = sjson.Set(template, "project", projectID) |
|
|
} else { |
|
|
template, _ = sjson.Set(template, "project", generateProjectID()) |
|
|
} |
|
|
template, _ = sjson.Set(template, "requestId", generateRequestID()) |
|
|
template, _ = sjson.Set(template, "request.sessionId", generateStableSessionID(payload)) |
|
|
|
|
|
template, _ = sjson.Delete(template, "request.safetySettings") |
|
|
template, _ = sjson.Set(template, "request.toolConfig.functionCallingConfig.mode", "VALIDATED") |
|
|
|
|
|
if !strings.HasPrefix(modelName, "gemini-3-") { |
|
|
if thinkingLevel := gjson.Get(template, "request.generationConfig.thinkingConfig.thinkingLevel"); thinkingLevel.Exists() { |
|
|
template, _ = sjson.Delete(template, "request.generationConfig.thinkingConfig.thinkingLevel") |
|
|
template, _ = sjson.Set(template, "request.generationConfig.thinkingConfig.thinkingBudget", -1) |
|
|
} |
|
|
} |
|
|
|
|
|
if strings.Contains(modelName, "claude") { |
|
|
gjson.Get(template, "request.tools").ForEach(func(key, tool gjson.Result) bool { |
|
|
tool.Get("functionDeclarations").ForEach(func(funKey, funcDecl gjson.Result) bool { |
|
|
if funcDecl.Get("parametersJsonSchema").Exists() { |
|
|
template, _ = sjson.SetRaw(template, fmt.Sprintf("request.tools.%d.functionDeclarations.%d.parameters", key.Int(), funKey.Int()), funcDecl.Get("parametersJsonSchema").Raw) |
|
|
template, _ = sjson.Delete(template, fmt.Sprintf("request.tools.%d.functionDeclarations.%d.parameters.$schema", key.Int(), funKey.Int())) |
|
|
template, _ = sjson.Delete(template, fmt.Sprintf("request.tools.%d.functionDeclarations.%d.parametersJsonSchema", key.Int(), funKey.Int())) |
|
|
} |
|
|
return true |
|
|
}) |
|
|
return true |
|
|
}) |
|
|
} else { |
|
|
template, _ = sjson.Delete(template, "request.generationConfig.maxOutputTokens") |
|
|
} |
|
|
|
|
|
return []byte(template) |
|
|
} |
|
|
|
|
|
func generateRequestID() string { |
|
|
return "agent-" + 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 |
|
|
} |
|
|
|
|
|
func modelName2Alias(modelName string) string { |
|
|
switch modelName { |
|
|
case "rev19-uic3-1p": |
|
|
return "gemini-2.5-computer-use-preview-10-2025" |
|
|
case "gemini-3-pro-image": |
|
|
return "gemini-3-pro-image-preview" |
|
|
case "gemini-3-pro-high": |
|
|
return "gemini-3-pro-preview" |
|
|
case "gemini-3-flash": |
|
|
return "gemini-3-flash-preview" |
|
|
case "claude-sonnet-4-5": |
|
|
return "gemini-claude-sonnet-4-5" |
|
|
case "claude-sonnet-4-5-thinking": |
|
|
return "gemini-claude-sonnet-4-5-thinking" |
|
|
case "claude-opus-4-5-thinking": |
|
|
return "gemini-claude-opus-4-5-thinking" |
|
|
case "chat_20706", "chat_23310", "gemini-2.5-flash-thinking", "gemini-3-pro-low", "gemini-2.5-pro": |
|
|
return "" |
|
|
default: |
|
|
return modelName |
|
|
} |
|
|
} |
|
|
|
|
|
func alias2ModelName(modelName string) string { |
|
|
switch modelName { |
|
|
case "gemini-2.5-computer-use-preview-10-2025": |
|
|
return "rev19-uic3-1p" |
|
|
case "gemini-3-pro-image-preview": |
|
|
return "gemini-3-pro-image" |
|
|
case "gemini-3-pro-preview": |
|
|
return "gemini-3-pro-high" |
|
|
case "gemini-3-flash-preview": |
|
|
return "gemini-3-flash" |
|
|
case "gemini-claude-sonnet-4-5": |
|
|
return "claude-sonnet-4-5" |
|
|
case "gemini-claude-sonnet-4-5-thinking": |
|
|
return "claude-sonnet-4-5-thinking" |
|
|
case "gemini-claude-opus-4-5-thinking": |
|
|
return "claude-opus-4-5-thinking" |
|
|
default: |
|
|
return modelName |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
func normalizeAntigravityThinking(model string, payload []byte, isClaude bool) []byte { |
|
|
payload = util.StripThinkingConfigIfUnsupported(model, payload) |
|
|
if !util.ModelSupportsThinking(model) { |
|
|
return payload |
|
|
} |
|
|
budget := gjson.GetBytes(payload, "request.generationConfig.thinkingConfig.thinkingBudget") |
|
|
if !budget.Exists() { |
|
|
return payload |
|
|
} |
|
|
raw := int(budget.Int()) |
|
|
normalized := util.NormalizeThinkingBudget(model, raw) |
|
|
|
|
|
if isClaude { |
|
|
effectiveMax, setDefaultMax := antigravityEffectiveMaxTokens(model, payload) |
|
|
if effectiveMax > 0 && normalized >= effectiveMax { |
|
|
normalized = effectiveMax - 1 |
|
|
} |
|
|
minBudget := antigravityMinThinkingBudget(model) |
|
|
if minBudget > 0 && normalized >= 0 && normalized < minBudget { |
|
|
|
|
|
payload, _ = sjson.DeleteBytes(payload, "request.generationConfig.thinkingConfig") |
|
|
return payload |
|
|
} |
|
|
if setDefaultMax { |
|
|
if res, errSet := sjson.SetBytes(payload, "request.generationConfig.maxOutputTokens", effectiveMax); errSet == nil { |
|
|
payload = res |
|
|
} |
|
|
} |
|
|
} |
|
|
|
|
|
updated, err := sjson.SetBytes(payload, "request.generationConfig.thinkingConfig.thinkingBudget", normalized) |
|
|
if err != nil { |
|
|
return payload |
|
|
} |
|
|
return updated |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
func antigravityEffectiveMaxTokens(model string, payload []byte) (max int, fromModel bool) { |
|
|
if maxTok := gjson.GetBytes(payload, "request.generationConfig.maxOutputTokens"); maxTok.Exists() && maxTok.Int() > 0 { |
|
|
return int(maxTok.Int()), false |
|
|
} |
|
|
if modelInfo := registry.GetGlobalRegistry().GetModelInfo(model); modelInfo != nil && modelInfo.MaxCompletionTokens > 0 { |
|
|
return modelInfo.MaxCompletionTokens, true |
|
|
} |
|
|
return 0, false |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
func antigravityMinThinkingBudget(model string) int { |
|
|
if modelInfo := registry.GetGlobalRegistry().GetModelInfo(model); modelInfo != nil && modelInfo.Thinking != nil { |
|
|
return modelInfo.Thinking.Min |
|
|
} |
|
|
return -1 |
|
|
} |
|
|
|