From e78d45acc90bf623ad1bd8f0bc8cabcc7929322f Mon Sep 17 00:00:00 2001 From: sususu98 Date: Fri, 24 Apr 2026 19:46:52 +0800 Subject: [PATCH] fix antigravity user agent handling --- internal/auth/antigravity/auth.go | 25 +++++--- internal/misc/antigravity_version.go | 61 +++++++++++++++++++ .../runtime/executor/antigravity_executor.go | 48 +++++++++++---- .../antigravity_executor_credits_test.go | 45 ++++++++++++++ 4 files changed, 158 insertions(+), 21 deletions(-) diff --git a/internal/auth/antigravity/auth.go b/internal/auth/antigravity/auth.go index 449f413fc..12d112c4e 100644 --- a/internal/auth/antigravity/auth.go +++ b/internal/auth/antigravity/auth.go @@ -12,6 +12,7 @@ import ( "time" "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/util" log "github.com/sirupsen/logrus" ) @@ -36,17 +37,21 @@ type AntigravityAuth struct { // NewAntigravityAuth creates a new Antigravity auth service. func NewAntigravityAuth(cfg *config.Config, httpClient *http.Client) *AntigravityAuth { - if httpClient != nil { - return &AntigravityAuth{httpClient: httpClient} - } if cfg == nil { cfg = &config.Config{} } + if httpClient != nil { + return &AntigravityAuth{httpClient: httpClient} + } return &AntigravityAuth{ httpClient: util.SetProxy(&cfg.SDKConfig, &http.Client{}), } } +func (o *AntigravityAuth) loadCodeAssistUserAgent() string { + return misc.AntigravityLoadCodeAssistUserAgent("") +} + // BuildAuthURL generates the OAuth authorization URL. func (o *AntigravityAuth) BuildAuthURL(state, redirectURI string) string { if strings.TrimSpace(redirectURI) == "" { @@ -153,11 +158,12 @@ func (o *AntigravityAuth) FetchUserInfo(ctx context.Context, accessToken string) // FetchProjectID retrieves the project ID for the authenticated user via loadCodeAssist func (o *AntigravityAuth) FetchProjectID(ctx context.Context, accessToken string) (string, error) { + userAgent := o.loadCodeAssistUserAgent() loadReqBody := map[string]any{ "metadata": map[string]string{ - "ideType": "ANTIGRAVITY", - "platform": "PLATFORM_UNSPECIFIED", - "pluginType": "GEMINI", + "ide_type": "ANTIGRAVITY", + "ide_version": misc.AntigravityVersionFromUserAgent(userAgent), + "ide_name": "antigravity", }, } @@ -173,9 +179,8 @@ func (o *AntigravityAuth) FetchProjectID(ctx context.Context, accessToken string } req.Header.Set("Authorization", "Bearer "+accessToken) req.Header.Set("Content-Type", "application/json") - req.Header.Set("User-Agent", APIUserAgent) - req.Header.Set("X-Goog-Api-Client", APIClient) - req.Header.Set("Client-Metadata", ClientMetadata) + req.Header.Set("User-Agent", userAgent) + req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1") resp, errDo := o.httpClient.Do(req) if errDo != nil { @@ -277,7 +282,7 @@ func (o *AntigravityAuth) OnboardUser(ctx context.Context, accessToken, tierID s } req.Header.Set("Authorization", "Bearer "+accessToken) req.Header.Set("Content-Type", "application/json") - req.Header.Set("User-Agent", APIUserAgent) + req.Header.Set("User-Agent", o.loadCodeAssistUserAgent()) req.Header.Set("X-Goog-Api-Client", APIClient) req.Header.Set("Client-Metadata", ClientMetadata) diff --git a/internal/misc/antigravity_version.go b/internal/misc/antigravity_version.go index 595cfefd9..1f05073ee 100644 --- a/internal/misc/antigravity_version.go +++ b/internal/misc/antigravity_version.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "net/http" + "strings" "sync" "time" @@ -18,6 +19,7 @@ const ( antigravityFallbackVersion = "1.21.9" antigravityVersionCacheTTL = 6 * time.Hour antigravityFetchTimeout = 10 * time.Second + AntigravityNodeAPIClientUA = "google-api-nodejs-client/10.3.0" ) type antigravityRelease struct { @@ -107,6 +109,65 @@ func AntigravityUserAgent() string { return fmt.Sprintf("antigravity/%s darwin/arm64", AntigravityLatestVersion()) } +func antigravityBaseUserAgent(userAgent string) string { + userAgent = strings.TrimSpace(userAgent) + if userAgent == "" { + return AntigravityUserAgent() + } + lower := strings.ToLower(userAgent) + if strings.HasPrefix(lower, "antigravity/") { + if idx := strings.Index(lower, " google-api-nodejs-client/"); idx >= 0 { + trimmed := strings.TrimSpace(userAgent[:idx]) + if trimmed != "" { + return trimmed + } + } + } + return userAgent +} + +// AntigravityRequestUserAgent returns the short Antigravity runtime UA used by +// generate/stream/model-list requests. +func AntigravityRequestUserAgent(userAgent string) string { + return antigravityBaseUserAgent(userAgent) +} + +// AntigravityLoadCodeAssistUserAgent returns the long Antigravity control-plane +// UA used by loadCodeAssist requests. +func AntigravityLoadCodeAssistUserAgent(userAgent string) string { + userAgent = strings.TrimSpace(userAgent) + if userAgent == "" { + return AntigravityUserAgent() + " " + AntigravityNodeAPIClientUA + } + lower := strings.ToLower(userAgent) + if !strings.HasPrefix(lower, "antigravity/") { + return userAgent + } + if strings.Contains(lower, "google-api-nodejs-client/") { + return userAgent + } + return antigravityBaseUserAgent(userAgent) + " " + AntigravityNodeAPIClientUA +} + +// AntigravityVersionFromUserAgent extracts the Antigravity version prefix from +// either the short or long Antigravity UA forms. +func AntigravityVersionFromUserAgent(userAgent string) string { + base := antigravityBaseUserAgent(userAgent) + lower := strings.ToLower(base) + if !strings.HasPrefix(lower, "antigravity/") { + return AntigravityLatestVersion() + } + rest := base[len("antigravity/"):] + if idx := strings.IndexAny(rest, " \t"); idx >= 0 { + rest = rest[:idx] + } + rest = strings.TrimSpace(rest) + if rest == "" { + return AntigravityLatestVersion() + } + return rest +} + func fetchAntigravityLatestVersion(ctx context.Context) (string, error) { if ctx == nil { ctx = context.Background() diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index 5cc93448d..3b3943b8a 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -478,7 +478,7 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au 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 { + if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { 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} @@ -680,7 +680,7 @@ attemptLoop: // 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 { + if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { 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} @@ -1139,7 +1139,7 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya baseModel := thinking.ParseSuffix(req.Model).ModelName ctx = context.WithValue(ctx, "alt", "") - if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown { + if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { 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} @@ -1763,16 +1763,29 @@ func (e *AntigravityExecutor) updateAntigravityCreditsBalance(ctx context.Contex 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)) + userAgent := resolveLoadCodeAssistUserAgent(auth) + loadReqBody, errMarshal := json.Marshal(map[string]any{ + "metadata": map[string]string{ + "ide_type": "ANTIGRAVITY", + "ide_version": misc.AntigravityVersionFromUserAgent(userAgent), + "ide_name": "antigravity", + }, + }) + if errMarshal != nil { + log.Debugf("antigravity executor: marshal loadCodeAssist request error: %v", errMarshal) + return + } + baseURL := buildBaseURL(auth) + endpointURL := strings.TrimSuffix(baseURL, "/") + "/v1internal:loadCodeAssist" + httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, bytes.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", resolveUserAgent(auth)) + httpReq.Header.Set("User-Agent", userAgent) + httpReq.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1") httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) httpResp, errDo := httpClient.Do(httpReq) @@ -2070,19 +2083,28 @@ func resolveHost(base string) string { } func resolveUserAgent(auth *cliproxyauth.Auth) string { + return misc.AntigravityRequestUserAgent(antigravityConfiguredUserAgent(auth)) +} + +func resolveLoadCodeAssistUserAgent(auth *cliproxyauth.Auth) string { + return misc.AntigravityLoadCodeAssistUserAgent(antigravityConfiguredUserAgent(auth)) +} + +func antigravityConfiguredUserAgent(auth *cliproxyauth.Auth) string { + raw := "" if auth != nil { if auth.Attributes != nil { if ua := strings.TrimSpace(auth.Attributes["user_agent"]); ua != "" { - return ua + raw = ua } } - if auth.Metadata != nil { + if raw == "" && auth.Metadata != nil { if ua, ok := auth.Metadata["user_agent"].(string); ok && strings.TrimSpace(ua) != "" { - return strings.TrimSpace(ua) + raw = strings.TrimSpace(ua) } } } - return misc.AntigravityUserAgent() + return raw } func antigravityRetryAttempts(auth *cliproxyauth.Auth, cfg *config.Config) int { @@ -2141,6 +2163,10 @@ func antigravityShouldRetrySoftRateLimit(statusCode int, body []byte) bool { return decideAntigravity429(body).kind == antigravity429DecisionSoftRetry } +func antigravityShouldBypassShortCooldown(ctx context.Context, cfg *config.Config) bool { + return cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(cfg) +} + func antigravitySoftRateLimitDelay(attempt int) time.Duration { if attempt < 0 { attempt = 0 diff --git a/internal/runtime/executor/antigravity_executor_credits_test.go b/internal/runtime/executor/antigravity_executor_credits_test.go index 6e38223e5..4569f5dfd 100644 --- a/internal/runtime/executor/antigravity_executor_credits_test.go +++ b/internal/runtime/executor/antigravity_executor_credits_test.go @@ -216,6 +216,11 @@ func TestAntigravityExecute_CreditsInjectedWhenConductorRequests(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, _ := io.ReadAll(r.Body) _ = r.Body.Close() + if r.URL.Path == "/v1internal:loadCodeAssist" { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"paidTier":{"id":"tier-1","availableCredits":[{"creditType":"GOOGLE_ONE_AI","creditAmount":"25000","minimumCreditAmountForUsage":"50"}]}}`)) + return + } requestBodies = append(requestBodies, string(body)) if !strings.Contains(string(body), `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) { @@ -269,6 +274,11 @@ func TestAntigravityExecute_NoCreditsWithoutConductorFlag(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, _ := io.ReadAll(r.Body) _ = r.Body.Close() + if r.URL.Path == "/v1internal:loadCodeAssist" { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"paidTier":{"id":"tier-1","availableCredits":[{"creditType":"GOOGLE_ONE_AI","creditAmount":"25000","minimumCreditAmountForUsage":"50"}]}}`)) + return + } requestBodies = append(requestBodies, string(body)) w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(`{"error":{"status":"RESOURCE_EXHAUSTED","message":"QUOTA_EXHAUSTED"}}`)) @@ -429,6 +439,41 @@ func TestEnsureAccessToken_WarmTokenLoadsCreditsHint(t *testing.T) { } } +func TestUpdateAntigravityCreditsBalance_LoadCodeAssistUserAgent(t *testing.T) { + resetAntigravityCreditsRetryState() + t.Cleanup(resetAntigravityCreditsRetryState) + + exec := NewAntigravityExecutor(&config.Config{}) + const userAgent = "antigravity/1.23.2 windows/amd64 google-api-nodejs-client/10.3.0" + auth := &cliproxyauth.Auth{ + ID: "auth-load-code-assist-ua", + Attributes: map[string]string{"user_agent": userAgent}, + } + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", roundTripperFunc(func(req *http.Request) (*http.Response, error) { + if req.URL.String() != "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist" { + t.Fatalf("unexpected request url %s", req.URL.String()) + } + if got := req.Header.Get("User-Agent"); got != userAgent { + t.Fatalf("User-Agent = %q, want %q", got, userAgent) + } + if got := req.Header.Get("X-Goog-Api-Client"); got != "gl-node/22.21.1" { + t.Fatalf("X-Goog-Api-Client = %q, want %q", got, "gl-node/22.21.1") + } + body, _ := io.ReadAll(req.Body) + _ = req.Body.Close() + if string(body) != `{"metadata":{"ide_name":"antigravity","ide_type":"ANTIGRAVITY","ide_version":"1.23.2"}}` { + t.Fatalf("loadCodeAssist body = %s", string(body)) + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(`{"paidTier":{"id":"tier-1","availableCredits":[{"creditType":"GOOGLE_ONE_AI","creditAmount":"25000","minimumCreditAmountForUsage":"50"}]}}`)), + }, nil + })) + + exec.updateAntigravityCreditsBalance(ctx, auth, "token") +} + func TestParseMetaFloat(t *testing.T) { tests := []struct { name string