lywsvip/CLIProxyAPI
1
0
Fork 0
mirror of https://github.com/router-for-me/CLIProxyAPI.git synced 2026-09-03 06:35:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(antigravity): strip tool config, labels, and session id in token counting

Browse Source 
- Strip `request.toolConfig`, `request.labels`, and `request.sessionId` from the payload before dispatching count tokens requests.

Closes: #4749
This commit is contained in:
Luis Pater
2026-09-02 21:08:21 +08:00
parent 02c02cda50
commit d0fb44ca95
2 changed files with 45 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
internal/runtime/executor/antigravity_executor_signature_test.go
View File

@@ -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)
}
}

3
internal/runtime/executor/antigravity_executor_tokens.go
View File

@@ -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)