luispater commited on
Commit
33ad0c6
·
unverified ·
1 Parent(s): fb26717

feat(antigravity): enable token counting via API with resilient routing

Browse files

Introduces the capability to count tokens for Antigravity-backed requests. This implementation leverages the `countTokens` endpoint of the Antigravity API, replacing the prior unsupported stub.

Key aspects of this update include:

- **API Integration**: Direct integration with the Antigravity `countTokens` API, including necessary request payload translation and authentication.
- **Resilient Infrastructure**: A fallback mechanism has been established, allowing the system to attempt connections across multiple Antigravity base URLs to ensure request success even in the event of temporary service interruptions.
- **Model Aliasing**: Added mappings for `gemini-3-flash` and `gemini-3-flash-preview` to ensure compatibility with the latest model variants.
- **Robust Error Handling**: Comprehensive error handling and logging are in place to manage failures during API interactions.

internal/runtime/executor/antigravity_executor.go CHANGED
@@ -32,15 +32,16 @@ import (
32
  const (
33
  antigravityBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com"
34
  // antigravityBaseURLAutopush = "https://autopush-cloudcode-pa.sandbox.googleapis.com"
35
- antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com"
36
- antigravityStreamPath = "/v1internal:streamGenerateContent"
37
- antigravityGeneratePath = "/v1internal:generateContent"
38
- antigravityModelsPath = "/v1internal:fetchAvailableModels"
39
- antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
40
- antigravityClientSecret = "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"
41
- defaultAntigravityAgent = "antigravity/1.11.5 windows/amd64"
42
- antigravityAuthType = "antigravity"
43
- refreshSkew = 3000 * time.Second
 
44
  )
45
 
46
  var randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
@@ -646,9 +647,131 @@ func (e *AntigravityExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Au
646
  return updated, nil
647
  }
648
 
649
- // CountTokens counts tokens for the given request (not supported for Antigravity).
650
- func (e *AntigravityExecutor) CountTokens(context.Context, *cliproxyauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
651
- return cliproxyexecutor.Response{}, statusErr{code: http.StatusNotImplemented, msg: "count tokens not supported"}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
652
  }
653
 
654
  // FetchAntigravityModels retrieves available models using the supplied auth.
@@ -1114,6 +1237,8 @@ func modelName2Alias(modelName string) string {
1114
  return "gemini-3-pro-image-preview"
1115
  case "gemini-3-pro-high":
1116
  return "gemini-3-pro-preview"
 
 
1117
  case "claude-sonnet-4-5":
1118
  return "gemini-claude-sonnet-4-5"
1119
  case "claude-sonnet-4-5-thinking":
@@ -1135,6 +1260,8 @@ func alias2ModelName(modelName string) string {
1135
  return "gemini-3-pro-image"
1136
  case "gemini-3-pro-preview":
1137
  return "gemini-3-pro-high"
 
 
1138
  case "gemini-claude-sonnet-4-5":
1139
  return "claude-sonnet-4-5"
1140
  case "gemini-claude-sonnet-4-5-thinking":
 
32
  const (
33
  antigravityBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com"
34
  // antigravityBaseURLAutopush = "https://autopush-cloudcode-pa.sandbox.googleapis.com"
35
+ antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com"
36
+ antigravityCountTokensPath = "/v1internal:countTokens"
37
+ antigravityStreamPath = "/v1internal:streamGenerateContent"
38
+ antigravityGeneratePath = "/v1internal:generateContent"
39
+ antigravityModelsPath = "/v1internal:fetchAvailableModels"
40
+ antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
41
+ antigravityClientSecret = "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"
42
+ defaultAntigravityAgent = "antigravity/1.11.5 windows/amd64"
43
+ antigravityAuthType = "antigravity"
44
+ refreshSkew = 3000 * time.Second
45
  )
46
 
47
  var randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
 
647
  return updated, nil
648
  }
649
 
650
+ // CountTokens counts tokens for the given request using the Antigravity API.
651
+ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
652
+ token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
653
+ if errToken != nil {
654
+ return cliproxyexecutor.Response{}, errToken
655
+ }
656
+ if updatedAuth != nil {
657
+ auth = updatedAuth
658
+ }
659
+ if strings.TrimSpace(token) == "" {
660
+ return cliproxyexecutor.Response{}, statusErr{code: http.StatusUnauthorized, msg: "missing access token"}
661
+ }
662
+
663
+ from := opts.SourceFormat
664
+ to := sdktranslator.FromString("antigravity")
665
+ respCtx := context.WithValue(ctx, "alt", opts.Alt)
666
+
667
+ baseURLs := antigravityBaseURLFallbackOrder(auth)
668
+ httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
669
+
670
+ var authID, authLabel, authType, authValue string
671
+ if auth != nil {
672
+ authID = auth.ID
673
+ authLabel = auth.Label
674
+ authType, authValue = auth.AccountInfo()
675
+ }
676
+
677
+ var lastStatus int
678
+ var lastBody []byte
679
+ var lastErr error
680
+
681
+ for idx, baseURL := range baseURLs {
682
+ payload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
683
+ payload = applyThinkingMetadataCLI(payload, req.Metadata, req.Model)
684
+ payload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, payload)
685
+ payload = normalizeAntigravityThinking(req.Model, payload)
686
+ payload = deleteJSONField(payload, "project")
687
+ payload = deleteJSONField(payload, "model")
688
+ payload = deleteJSONField(payload, "request.safetySettings")
689
+
690
+ base := strings.TrimSuffix(baseURL, "/")
691
+ if base == "" {
692
+ base = buildBaseURL(auth)
693
+ }
694
+
695
+ var requestURL strings.Builder
696
+ requestURL.WriteString(base)
697
+ requestURL.WriteString(antigravityCountTokensPath)
698
+ if opts.Alt != "" {
699
+ requestURL.WriteString("?$alt=")
700
+ requestURL.WriteString(url.QueryEscape(opts.Alt))
701
+ }
702
+
703
+ httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bytes.NewReader(payload))
704
+ if errReq != nil {
705
+ return cliproxyexecutor.Response{}, errReq
706
+ }
707
+ httpReq.Header.Set("Content-Type", "application/json")
708
+ httpReq.Header.Set("Authorization", "Bearer "+token)
709
+ httpReq.Header.Set("User-Agent", resolveUserAgent(auth))
710
+ httpReq.Header.Set("Accept", "application/json")
711
+ if host := resolveHost(base); host != "" {
712
+ httpReq.Host = host
713
+ }
714
+
715
+ recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
716
+ URL: requestURL.String(),
717
+ Method: http.MethodPost,
718
+ Headers: httpReq.Header.Clone(),
719
+ Body: payload,
720
+ Provider: e.Identifier(),
721
+ AuthID: authID,
722
+ AuthLabel: authLabel,
723
+ AuthType: authType,
724
+ AuthValue: authValue,
725
+ })
726
+
727
+ httpResp, errDo := httpClient.Do(httpReq)
728
+ if errDo != nil {
729
+ recordAPIResponseError(ctx, e.cfg, errDo)
730
+ lastStatus = 0
731
+ lastBody = nil
732
+ lastErr = errDo
733
+ if idx+1 < len(baseURLs) {
734
+ log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1])
735
+ continue
736
+ }
737
+ return cliproxyexecutor.Response{}, errDo
738
+ }
739
+
740
+ recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
741
+ bodyBytes, errRead := io.ReadAll(httpResp.Body)
742
+ if errClose := httpResp.Body.Close(); errClose != nil {
743
+ log.Errorf("antigravity executor: close response body error: %v", errClose)
744
+ }
745
+ if errRead != nil {
746
+ recordAPIResponseError(ctx, e.cfg, errRead)
747
+ return cliproxyexecutor.Response{}, errRead
748
+ }
749
+ appendAPIResponseChunk(ctx, e.cfg, bodyBytes)
750
+
751
+ if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices {
752
+ count := gjson.GetBytes(bodyBytes, "totalTokens").Int()
753
+ translated := sdktranslator.TranslateTokenCount(respCtx, to, from, count, bodyBytes)
754
+ return cliproxyexecutor.Response{Payload: []byte(translated)}, nil
755
+ }
756
+
757
+ lastStatus = httpResp.StatusCode
758
+ lastBody = append([]byte(nil), bodyBytes...)
759
+ lastErr = nil
760
+ if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) {
761
+ log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1])
762
+ continue
763
+ }
764
+ return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)}
765
+ }
766
+
767
+ switch {
768
+ case lastStatus != 0:
769
+ return cliproxyexecutor.Response{}, statusErr{code: lastStatus, msg: string(lastBody)}
770
+ case lastErr != nil:
771
+ return cliproxyexecutor.Response{}, lastErr
772
+ default:
773
+ return cliproxyexecutor.Response{}, statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"}
774
+ }
775
  }
776
 
777
  // FetchAntigravityModels retrieves available models using the supplied auth.
 
1237
  return "gemini-3-pro-image-preview"
1238
  case "gemini-3-pro-high":
1239
  return "gemini-3-pro-preview"
1240
+ case "gemini-3-flash":
1241
+ return "gemini-3-flash-preview"
1242
  case "claude-sonnet-4-5":
1243
  return "gemini-claude-sonnet-4-5"
1244
  case "claude-sonnet-4-5-thinking":
 
1260
  return "gemini-3-pro-image"
1261
  case "gemini-3-pro-preview":
1262
  return "gemini-3-pro-high"
1263
+ case "gemini-3-flash-preview":
1264
+ return "gemini-3-flash"
1265
  case "gemini-claude-sonnet-4-5":
1266
  return "claude-sonnet-4-5"
1267
  case "gemini-claude-sonnet-4-5-thinking":