Merge pull request #3117 from sususu98/fix/restore-antigravity-ua-handling

fix: restore Antigravity user agent handling

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) == "" {
@@ -118,6 +123,7 @@ func (o *AntigravityAuth) FetchUserInfo(ctx context.Context, accessToken string)
 		return "", fmt.Errorf("antigravity userinfo: create request: %w", err)
 	}
 	req.Header.Set("Authorization", "Bearer "+accessToken)
+	req.Header.Set("User-Agent", o.loadCodeAssistUserAgent())
 
 	resp, errDo := o.httpClient.Do(req)
 	if errDo != nil {
@@ -153,11 +159,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",
+			"ide_type":    "ANTIGRAVITY",
-			"platform":   "PLATFORM_UNSPECIFIED",
+			"ide_version": misc.AntigravityVersionFromUserAgent(userAgent),
-			"pluginType": "GEMINI",
+			"ide_name":    "antigravity",
 		},
 	}
 
@@ -173,9 +180,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("User-Agent", userAgent)
-	req.Header.Set("X-Goog-Api-Client", APIClient)
+	req.Header.Set("X-Goog-Api-Client", misc.AntigravityGoogAPIClientUA)
-	req.Header.Set("Client-Metadata", ClientMetadata)
 
 	resp, errDo := o.httpClient.Do(req)
 	if errDo != nil {
@@ -244,12 +250,13 @@ func (o *AntigravityAuth) FetchProjectID(ctx context.Context, accessToken string
 // OnboardUser attempts to fetch the project ID via onboardUser by polling for completion
 func (o *AntigravityAuth) OnboardUser(ctx context.Context, accessToken, tierID string) (string, error) {
 	log.Infof("Antigravity: onboarding user with tier: %s", tierID)
+	userAgent := o.loadCodeAssistUserAgent()
 	requestBody := map[string]any{
 		"tierId": tierID,
 		"metadata": map[string]string{
-			"ideType":    "ANTIGRAVITY",
+			"ide_type":    "ANTIGRAVITY",
-			"platform":   "PLATFORM_UNSPECIFIED",
+			"ide_version": misc.AntigravityVersionFromUserAgent(userAgent),
-			"pluginType": "GEMINI",
+			"ide_name":    "antigravity",
 		},
 	}
 
@@ -277,9 +284,8 @@ 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", userAgent)
-		req.Header.Set("X-Goog-Api-Client", APIClient)
+		req.Header.Set("X-Goog-Api-Client", misc.AntigravityGoogAPIClientUA)
-		req.Header.Set("Client-Metadata", ClientMetadata)
 
 		resp, errDo := o.httpClient.Do(req)
 		if errDo != nil {

internal/auth/antigravity/constants.go

@@ -21,14 +21,11 @@ var Scopes = []string{
 const (
 	TokenEndpoint    = "https://oauth2.googleapis.com/token"
 	AuthEndpoint     = "https://accounts.google.com/o/oauth2/v2/auth"
-	UserInfoEndpoint = "https://www.googleapis.com/oauth2/v1/userinfo?alt=json"
+	UserInfoEndpoint = "https://www.googleapis.com/oauth2/v2/userinfo?alt=json"
 )
 
 // Antigravity API configuration
 const (
-	APIEndpoint    = "https://cloudcode-pa.googleapis.com"
+	APIEndpoint = "https://cloudcode-pa.googleapis.com"
-	APIVersion     = "v1internal"
+	APIVersion  = "v1internal"
-	APIUserAgent   = "google-api-nodejs-client/9.15.1"
-	APIClient      = "google-cloud-sdk vscode_cloudshelleditor/0.1"
-	ClientMetadata = `{"ideType":"IDE_UNSPECIFIED","platform":"PLATFORM_UNSPECIFIED","pluginType":"GEMINI"}`
 )

internal/misc/antigravity_version.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"strings"
 	"sync"
 	"time"
 
@@ -18,6 +19,8 @@ const (
 	antigravityFallbackVersion = "1.21.9"
 	antigravityVersionCacheTTL = 6 * time.Hour
 	antigravityFetchTimeout    = 10 * time.Second
+	AntigravityNodeAPIClientUA = "google-api-nodejs-client/10.3.0"
+	AntigravityGoogAPIClientUA = "gl-node/22.21.1"
 )
 
 type antigravityRelease struct {
@@ -107,6 +110,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()

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"}}`
+	userAgent := resolveLoadCodeAssistUserAgent(auth)
-	endpointURL := "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"
+	loadReqBody, errMarshal := json.Marshal(map[string]any{
-	httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, strings.NewReader(loadReqBody))
+		"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", misc.AntigravityGoogAPIClientUA)
 
 	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

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