From d0fb44ca95e8e21ba01df0132a01fe4e00468661 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Wed, 2 Sep 2026 21:08:21 +0800 Subject: [PATCH] fix(antigravity): strip tool config, labels, and session id in token counting - Strip `request.toolConfig`, `request.labels`, and `request.sessionId` from the payload before dispatching count tokens requests. Closes: #4749 --- .../antigravity_executor_signature_test.go | 42 +++++++++++++++++++ .../executor/antigravity_executor_tokens.go | 3 ++ 2 files changed, 45 insertions(+) diff --git a/internal/runtime/executor/antigravity_executor_signature_test.go b/internal/runtime/executor/antigravity_executor_signature_test.go index 3d6597f11..5ab8e1bc8 100644 --- a/internal/runtime/executor/antigravity_executor_signature_test.go +++ b/internal/runtime/executor/antigravity_executor_signature_test.go @@ -1108,3 +1108,45 @@ func TestAntigravityExecutor_CacheModeSkipsPrecheck(t *testing.T) { t.Fatalf("cache mode should skip precheck, got: %v", err) } } + +func TestAntigravityExecutorCountTokensStripsStatefulFields(t *testing.T) { + var upstreamBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != antigravityCountTokensPath { + t.Fatalf("path = %q, want %q", r.URL.Path, antigravityCountTokensPath) + } + body, errRead := io.ReadAll(r.Body) + if errRead != nil { + t.Fatalf("read countTokens body: %v", errRead) + } + upstreamBody = append([]byte(nil), body...) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"totalTokens":42}`)) + })) + defer server.Close() + + payload := []byte(`{"contents":[{"role":"user","parts":[{"text":"hello"}]}],"toolConfig":{"functionCallingConfig":{"mode":"AUTO"}},"labels":{"source":"ide"},"sessionId":"session-123"}`) + exec := NewAntigravityExecutor(&config.Config{RequestRetry: 1}) + _, errCount := exec.CountTokens(context.Background(), testAntigravityAuth(server.URL), cliproxyexecutor.Request{ + Model: "gemini-3.6-flash-high", + Payload: payload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatGemini, + ResponseFormat: sdktranslator.FormatGemini, + }) + if errCount != nil { + t.Fatalf("CountTokens() error = %v", errCount) + } + if len(upstreamBody) == 0 { + t.Fatal("countTokens upstream body was not captured") + } + if gjson.GetBytes(upstreamBody, "request.toolConfig").Exists() { + t.Fatalf("upstream countTokens body retained request.toolConfig: %s", upstreamBody) + } + if gjson.GetBytes(upstreamBody, "request.labels").Exists() { + t.Fatalf("upstream countTokens body retained request.labels: %s", upstreamBody) + } + if gjson.GetBytes(upstreamBody, "request.sessionId").Exists() { + t.Fatalf("upstream countTokens body retained request.sessionId: %s", upstreamBody) + } +} diff --git a/internal/runtime/executor/antigravity_executor_tokens.go b/internal/runtime/executor/antigravity_executor_tokens.go index 78e859c06..80da62cc6 100644 --- a/internal/runtime/executor/antigravity_executor_tokens.go +++ b/internal/runtime/executor/antigravity_executor_tokens.go @@ -65,6 +65,9 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut payload = helps.DeleteJSONField(payload, "project") payload = helps.DeleteJSONField(payload, "model") payload = helps.DeleteJSONField(payload, "request.safetySettings") + payload = helps.DeleteJSONField(payload, "request.toolConfig") + payload = helps.DeleteJSONField(payload, "request.labels") + payload = helps.DeleteJSONField(payload, "request.sessionId") base := resolveAntigravityRequestBaseURL(auth) httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)