Skip to content

Commit 6f6331f

Browse files
conversunconversun
committed
Merge upstream/main: antigravity UA fixes, OpenAI image handler tests
2 parents a2b97ff + c4965be commit 6f6331f

7 files changed

Lines changed: 311 additions & 44 deletions

File tree

‎internal/auth/antigravity/auth.go‎

Lines changed: 21 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"time"
1313

1414
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
15+
"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
1516
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
1617
log "github.com/sirupsen/logrus"
1718
)
@@ -36,17 +37,21 @@ type AntigravityAuth struct {
3637

3738
// NewAntigravityAuth creates a new Antigravity auth service.
3839
func NewAntigravityAuth(cfg *config.Config, httpClient *http.Client) *AntigravityAuth {
39-
if httpClient != nil {
40-
return &AntigravityAuth{httpClient: httpClient}
41-
}
4240
if cfg == nil {
4341
cfg = &config.Config{}
4442
}
43+
if httpClient != nil {
44+
return &AntigravityAuth{httpClient: httpClient}
45+
}
4546
return &AntigravityAuth{
4647
httpClient: util.SetProxy(&cfg.SDKConfig, &http.Client{}),
4748
}
4849
}
4950

51+
func (o *AntigravityAuth) loadCodeAssistUserAgent() string {
52+
return misc.AntigravityLoadCodeAssistUserAgent("")
53+
}
54+
5055
// BuildAuthURL generates the OAuth authorization URL.
5156
func (o *AntigravityAuth) BuildAuthURL(state, redirectURI string) string {
5257
if strings.TrimSpace(redirectURI) == "" {
@@ -118,6 +123,7 @@ func (o *AntigravityAuth) FetchUserInfo(ctx context.Context, accessToken string)
118123
return "", fmt.Errorf("antigravity userinfo: create request: %w", err)
119124
}
120125
req.Header.Set("Authorization", "Bearer "+accessToken)
126+
req.Header.Set("User-Agent", o.loadCodeAssistUserAgent())
121127

122128
resp, errDo := o.httpClient.Do(req)
123129
if errDo != nil {
@@ -153,11 +159,12 @@ func (o *AntigravityAuth) FetchUserInfo(ctx context.Context, accessToken string)
153159

154160
// FetchProjectID retrieves the project ID for the authenticated user via loadCodeAssist
155161
func (o *AntigravityAuth) FetchProjectID(ctx context.Context, accessToken string) (string, error) {
162+
userAgent := o.loadCodeAssistUserAgent()
156163
loadReqBody := map[string]any{
157164
"metadata": map[string]string{
158-
"ideType": "ANTIGRAVITY",
159-
"platform": "PLATFORM_UNSPECIFIED",
160-
"pluginType": "GEMINI",
165+
"ide_type": "ANTIGRAVITY",
166+
"ide_version": misc.AntigravityVersionFromUserAgent(userAgent),
167+
"ide_name": "antigravity",
161168
},
162169
}
163170

@@ -173,9 +180,8 @@ func (o *AntigravityAuth) FetchProjectID(ctx context.Context, accessToken string
173180
}
174181
req.Header.Set("Authorization", "Bearer "+accessToken)
175182
req.Header.Set("Content-Type", "application/json")
176-
req.Header.Set("User-Agent", APIUserAgent)
177-
req.Header.Set("X-Goog-Api-Client", APIClient)
178-
req.Header.Set("Client-Metadata", ClientMetadata)
183+
req.Header.Set("User-Agent", userAgent)
184+
req.Header.Set("X-Goog-Api-Client", misc.AntigravityGoogAPIClientUA)
179185

180186
resp, errDo := o.httpClient.Do(req)
181187
if errDo != nil {
@@ -244,12 +250,13 @@ func (o *AntigravityAuth) FetchProjectID(ctx context.Context, accessToken string
244250
// OnboardUser attempts to fetch the project ID via onboardUser by polling for completion
245251
func (o *AntigravityAuth) OnboardUser(ctx context.Context, accessToken, tierID string) (string, error) {
246252
log.Infof("Antigravity: onboarding user with tier: %s", tierID)
253+
userAgent := o.loadCodeAssistUserAgent()
247254
requestBody := map[string]any{
248255
"tierId": tierID,
249256
"metadata": map[string]string{
250-
"ideType": "ANTIGRAVITY",
251-
"platform": "PLATFORM_UNSPECIFIED",
252-
"pluginType": "GEMINI",
257+
"ide_type": "ANTIGRAVITY",
258+
"ide_version": misc.AntigravityVersionFromUserAgent(userAgent),
259+
"ide_name": "antigravity",
253260
},
254261
}
255262

@@ -277,9 +284,8 @@ func (o *AntigravityAuth) OnboardUser(ctx context.Context, accessToken, tierID s
277284
}
278285
req.Header.Set("Authorization", "Bearer "+accessToken)
279286
req.Header.Set("Content-Type", "application/json")
280-
req.Header.Set("User-Agent", APIUserAgent)
281-
req.Header.Set("X-Goog-Api-Client", APIClient)
282-
req.Header.Set("Client-Metadata", ClientMetadata)
287+
req.Header.Set("User-Agent", userAgent)
288+
req.Header.Set("X-Goog-Api-Client", misc.AntigravityGoogAPIClientUA)
283289

284290
resp, errDo := o.httpClient.Do(req)
285291
if errDo != nil {

‎internal/auth/antigravity/constants.go‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,11 @@ var Scopes = []string{
2121
const (
2222
TokenEndpoint = "https://oauth2.googleapis.com/token"
2323
AuthEndpoint = "https://accounts.google.com/o/oauth2/v2/auth"
24-
UserInfoEndpoint = "https://www.googleapis.com/oauth2/v1/userinfo?alt=json"
24+
UserInfoEndpoint = "https://www.googleapis.com/oauth2/v2/userinfo?alt=json"
2525
)
2626

2727
// Antigravity API configuration
2828
const (
29-
APIEndpoint = "https://cloudcode-pa.googleapis.com"
30-
APIVersion = "v1internal"
31-
APIUserAgent = "google-api-nodejs-client/9.15.1"
32-
APIClient = "google-cloud-sdk vscode_cloudshelleditor/0.1"
33-
ClientMetadata = `{"ideType":"IDE_UNSPECIFIED","platform":"PLATFORM_UNSPECIFIED","pluginType":"GEMINI"}`
29+
APIEndpoint = "https://cloudcode-pa.googleapis.com"
30+
APIVersion = "v1internal"
3431
)

‎internal/misc/antigravity_version.go‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"errors"
88
"fmt"
99
"net/http"
10+
"strings"
1011
"sync"
1112
"time"
1213

@@ -18,6 +19,8 @@ const (
1819
antigravityFallbackVersion = "1.21.9"
1920
antigravityVersionCacheTTL = 6 * time.Hour
2021
antigravityFetchTimeout = 10 * time.Second
22+
AntigravityNodeAPIClientUA = "google-api-nodejs-client/10.3.0"
23+
AntigravityGoogAPIClientUA = "gl-node/22.21.1"
2124
)
2225

2326
type antigravityRelease struct {
@@ -107,6 +110,65 @@ func AntigravityUserAgent() string {
107110
return fmt.Sprintf("antigravity/%s darwin/arm64", AntigravityLatestVersion())
108111
}
109112

113+
func antigravityBaseUserAgent(userAgent string) string {
114+
userAgent = strings.TrimSpace(userAgent)
115+
if userAgent == "" {
116+
return AntigravityUserAgent()
117+
}
118+
lower := strings.ToLower(userAgent)
119+
if strings.HasPrefix(lower, "antigravity/") {
120+
if idx := strings.Index(lower, " google-api-nodejs-client/"); idx >= 0 {
121+
trimmed := strings.TrimSpace(userAgent[:idx])
122+
if trimmed != "" {
123+
return trimmed
124+
}
125+
}
126+
}
127+
return userAgent
128+
}
129+
130+
// AntigravityRequestUserAgent returns the short Antigravity runtime UA used by
131+
// generate/stream/model-list requests.
132+
func AntigravityRequestUserAgent(userAgent string) string {
133+
return antigravityBaseUserAgent(userAgent)
134+
}
135+
136+
// AntigravityLoadCodeAssistUserAgent returns the long Antigravity control-plane
137+
// UA used by loadCodeAssist requests.
138+
func AntigravityLoadCodeAssistUserAgent(userAgent string) string {
139+
userAgent = strings.TrimSpace(userAgent)
140+
if userAgent == "" {
141+
return AntigravityUserAgent() + " " + AntigravityNodeAPIClientUA
142+
}
143+
lower := strings.ToLower(userAgent)
144+
if !strings.HasPrefix(lower, "antigravity/") {
145+
return userAgent
146+
}
147+
if strings.Contains(lower, "google-api-nodejs-client/") {
148+
return userAgent
149+
}
150+
return antigravityBaseUserAgent(userAgent) + " " + AntigravityNodeAPIClientUA
151+
}
152+
153+
// AntigravityVersionFromUserAgent extracts the Antigravity version prefix from
154+
// either the short or long Antigravity UA forms.
155+
func AntigravityVersionFromUserAgent(userAgent string) string {
156+
base := antigravityBaseUserAgent(userAgent)
157+
lower := strings.ToLower(base)
158+
if !strings.HasPrefix(lower, "antigravity/") {
159+
return AntigravityLatestVersion()
160+
}
161+
rest := base[len("antigravity/"):]
162+
if idx := strings.IndexAny(rest, " \t"); idx >= 0 {
163+
rest = rest[:idx]
164+
}
165+
rest = strings.TrimSpace(rest)
166+
if rest == "" {
167+
return AntigravityLatestVersion()
168+
}
169+
return rest
170+
}
171+
110172
func fetchAntigravityLatestVersion(ctx context.Context) (string, error) {
111173
if ctx == nil {
112174
ctx = context.Background()

‎internal/runtime/executor/antigravity_executor.go‎

Lines changed: 37 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -478,7 +478,7 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
478478
return resp, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"}
479479
}
480480
baseModel := thinking.ParseSuffix(req.Model).ModelName
481-
if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown {
481+
if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) {
482482
log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining)
483483
d := remaining
484484
return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
@@ -680,7 +680,7 @@ attemptLoop:
680680
// executeClaudeNonStream performs a claude non-streaming request to the Antigravity API.
681681
func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {
682682
baseModel := thinking.ParseSuffix(req.Model).ModelName
683-
if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown {
683+
if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) {
684684
log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining)
685685
d := remaining
686686
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
11391139
baseModel := thinking.ParseSuffix(req.Model).ModelName
11401140

11411141
ctx = context.WithValue(ctx, "alt", "")
1142-
if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown {
1142+
if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) {
11431143
log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining)
11441144
d := remaining
11451145
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
17631763
return
17641764
}
17651765

1766-
loadReqBody := `{"metadata":{"ideType":"ANTIGRAVITY","platform":"PLATFORM_UNSPECIFIED","pluginType":"GEMINI"}}`
1767-
endpointURL := "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"
1768-
httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, strings.NewReader(loadReqBody))
1766+
userAgent := resolveLoadCodeAssistUserAgent(auth)
1767+
loadReqBody, errMarshal := json.Marshal(map[string]any{
1768+
"metadata": map[string]string{
1769+
"ide_type": "ANTIGRAVITY",
1770+
"ide_version": misc.AntigravityVersionFromUserAgent(userAgent),
1771+
"ide_name": "antigravity",
1772+
},
1773+
})
1774+
if errMarshal != nil {
1775+
log.Debugf("antigravity executor: marshal loadCodeAssist request error: %v", errMarshal)
1776+
return
1777+
}
1778+
baseURL := buildBaseURL(auth)
1779+
endpointURL := strings.TrimSuffix(baseURL, "/") + "/v1internal:loadCodeAssist"
1780+
httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, bytes.NewReader(loadReqBody))
17691781
if errReq != nil {
17701782
log.Debugf("antigravity executor: create loadCodeAssist request error: %v", errReq)
17711783
return
17721784
}
17731785
httpReq.Header.Set("Authorization", "Bearer "+token)
17741786
httpReq.Header.Set("Content-Type", "application/json")
1775-
httpReq.Header.Set("User-Agent", "google-api-nodejs-client/9.15.1")
1787+
httpReq.Header.Set("User-Agent", userAgent)
1788+
httpReq.Header.Set("X-Goog-Api-Client", misc.AntigravityGoogAPIClientUA)
17761789

17771790
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
17781791
httpResp, errDo := httpClient.Do(httpReq)
@@ -2070,19 +2083,28 @@ func resolveHost(base string) string {
20702083
}
20712084

20722085
func resolveUserAgent(auth *cliproxyauth.Auth) string {
2086+
return misc.AntigravityRequestUserAgent(antigravityConfiguredUserAgent(auth))
2087+
}
2088+
2089+
func resolveLoadCodeAssistUserAgent(auth *cliproxyauth.Auth) string {
2090+
return misc.AntigravityLoadCodeAssistUserAgent(antigravityConfiguredUserAgent(auth))
2091+
}
2092+
2093+
func antigravityConfiguredUserAgent(auth *cliproxyauth.Auth) string {
2094+
raw := ""
20732095
if auth != nil {
20742096
if auth.Attributes != nil {
20752097
if ua := strings.TrimSpace(auth.Attributes["user_agent"]); ua != "" {
2076-
return ua
2098+
raw = ua
20772099
}
20782100
}
2079-
if auth.Metadata != nil {
2101+
if raw == "" && auth.Metadata != nil {
20802102
if ua, ok := auth.Metadata["user_agent"].(string); ok && strings.TrimSpace(ua) != "" {
2081-
return strings.TrimSpace(ua)
2103+
raw = strings.TrimSpace(ua)
20822104
}
20832105
}
20842106
}
2085-
return misc.AntigravityUserAgent()
2107+
return raw
20862108
}
20872109

20882110
func antigravityRetryAttempts(auth *cliproxyauth.Auth, cfg *config.Config) int {
@@ -2141,6 +2163,10 @@ func antigravityShouldRetrySoftRateLimit(statusCode int, body []byte) bool {
21412163
return decideAntigravity429(body).kind == antigravity429DecisionSoftRetry
21422164
}
21432165

2166+
func antigravityShouldBypassShortCooldown(ctx context.Context, cfg *config.Config) bool {
2167+
return cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(cfg)
2168+
}
2169+
21442170
func antigravitySoftRateLimitDelay(attempt int) time.Duration {
21452171
if attempt < 0 {
21462172
attempt = 0

‎internal/runtime/executor/antigravity_executor_credits_test.go‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,11 @@ func TestAntigravityExecute_CreditsInjectedWhenConductorRequests(t *testing.T) {
216216
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
217217
body, _ := io.ReadAll(r.Body)
218218
_ = r.Body.Close()
219+
if r.URL.Path == "/v1internal:loadCodeAssist" {
220+
w.Header().Set("Content-Type", "application/json")
221+
_, _ = w.Write([]byte(`{"paidTier":{"id":"tier-1","availableCredits":[{"creditType":"GOOGLE_ONE_AI","creditAmount":"25000","minimumCreditAmountForUsage":"50"}]}}`))
222+
return
223+
}
219224
requestBodies = append(requestBodies, string(body))
220225

221226
if !strings.Contains(string(body), `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) {
@@ -269,6 +274,11 @@ func TestAntigravityExecute_NoCreditsWithoutConductorFlag(t *testing.T) {
269274
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
270275
body, _ := io.ReadAll(r.Body)
271276
_ = r.Body.Close()
277+
if r.URL.Path == "/v1internal:loadCodeAssist" {
278+
w.Header().Set("Content-Type", "application/json")
279+
_, _ = w.Write([]byte(`{"paidTier":{"id":"tier-1","availableCredits":[{"creditType":"GOOGLE_ONE_AI","creditAmount":"25000","minimumCreditAmountForUsage":"50"}]}}`))
280+
return
281+
}
272282
requestBodies = append(requestBodies, string(body))
273283
w.WriteHeader(http.StatusTooManyRequests)
274284
_, _ = w.Write([]byte(`{"error":{"status":"RESOURCE_EXHAUSTED","message":"QUOTA_EXHAUSTED"}}`))
@@ -429,6 +439,41 @@ func TestEnsureAccessToken_WarmTokenLoadsCreditsHint(t *testing.T) {
429439
}
430440
}
431441

442+
func TestUpdateAntigravityCreditsBalance_LoadCodeAssistUserAgent(t *testing.T) {
443+
resetAntigravityCreditsRetryState()
444+
t.Cleanup(resetAntigravityCreditsRetryState)
445+
446+
exec := NewAntigravityExecutor(&config.Config{})
447+
const userAgent = "antigravity/1.23.2 windows/amd64 google-api-nodejs-client/10.3.0"
448+
auth := &cliproxyauth.Auth{
449+
ID: "auth-load-code-assist-ua",
450+
Attributes: map[string]string{"user_agent": userAgent},
451+
}
452+
ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", roundTripperFunc(func(req *http.Request) (*http.Response, error) {
453+
if req.URL.String() != "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist" {
454+
t.Fatalf("unexpected request url %s", req.URL.String())
455+
}
456+
if got := req.Header.Get("User-Agent"); got != userAgent {
457+
t.Fatalf("User-Agent = %q, want %q", got, userAgent)
458+
}
459+
if got := req.Header.Get("X-Goog-Api-Client"); got != "gl-node/22.21.1" {
460+
t.Fatalf("X-Goog-Api-Client = %q, want %q", got, "gl-node/22.21.1")
461+
}
462+
body, _ := io.ReadAll(req.Body)
463+
_ = req.Body.Close()
464+
if string(body) != `{"metadata":{"ide_name":"antigravity","ide_type":"ANTIGRAVITY","ide_version":"1.23.2"}}` {
465+
t.Fatalf("loadCodeAssist body = %s", string(body))
466+
}
467+
return &http.Response{
468+
StatusCode: http.StatusOK,
469+
Header: make(http.Header),
470+
Body: io.NopCloser(strings.NewReader(`{"paidTier":{"id":"tier-1","availableCredits":[{"creditType":"GOOGLE_ONE_AI","creditAmount":"25000","minimumCreditAmountForUsage":"50"}]}}`)),
471+
}, nil
472+
}))
473+
474+
exec.updateAntigravityCreditsBalance(ctx, auth, "token")
475+
}
476+
432477
func TestParseMetaFloat(t *testing.T) {
433478
tests := []struct {
434479
name string

0 commit comments

Comments
 (0)