From 9dc0e6d08b90de6424092b4df38efb5729df453c Mon Sep 17 00:00:00 2001 From: sususu98 Date: Thu, 29 Jan 2026 11:16:00 +0800 Subject: [PATCH] fix(translator): restore usageMetadata in Gemini responses from Antigravity When using Gemini API format with Antigravity backend, the executor renames usageMetadata to cpaUsageMetadata in non-terminal chunks. The Gemini translator was returning this internal field name directly to clients instead of the standard usageMetadata field. Add restoreUsageMetadata() to rename cpaUsageMetadata back to usageMetadata before returning responses to clients. --- .../gemini/antigravity_gemini_response.go | 16 +++- .../antigravity_gemini_response_test.go | 95 +++++++++++++++++++ 2 files changed, 110 insertions(+), 1 deletion(-) create mode 100644 internal/translator/antigravity/gemini/antigravity_gemini_response_test.go diff --git a/internal/translator/antigravity/gemini/antigravity_gemini_response.go b/internal/translator/antigravity/gemini/antigravity_gemini_response.go index 6f9d9791..874dc283 100644 --- a/internal/translator/antigravity/gemini/antigravity_gemini_response.go +++ b/internal/translator/antigravity/gemini/antigravity_gemini_response.go @@ -41,6 +41,7 @@ func ConvertAntigravityResponseToGemini(ctx context.Context, _ string, originalR responseResult := gjson.GetBytes(rawJSON, "response") if responseResult.Exists() { chunk = []byte(responseResult.Raw) + chunk = restoreUsageMetadata(chunk) } } else { chunkTemplate := "[]" @@ -76,7 +77,8 @@ func ConvertAntigravityResponseToGemini(ctx context.Context, _ string, originalR func ConvertAntigravityResponseToGeminiNonStream(_ context.Context, _ string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, _ *any) string { responseResult := gjson.GetBytes(rawJSON, "response") if responseResult.Exists() { - return responseResult.Raw + chunk := restoreUsageMetadata([]byte(responseResult.Raw)) + return string(chunk) } return string(rawJSON) } @@ -84,3 +86,15 @@ func ConvertAntigravityResponseToGeminiNonStream(_ context.Context, _ string, or func GeminiTokenCount(ctx context.Context, count int64) string { return fmt.Sprintf(`{"totalTokens":%d,"promptTokensDetails":[{"modality":"TEXT","tokenCount":%d}]}`, count, count) } + +// restoreUsageMetadata renames cpaUsageMetadata back to usageMetadata. +// The executor renames usageMetadata to cpaUsageMetadata in non-terminal chunks +// to preserve usage data while hiding it from clients that don't expect it. +// When returning standard Gemini API format, we must restore the original name. +func restoreUsageMetadata(chunk []byte) []byte { + if cpaUsage := gjson.GetBytes(chunk, "cpaUsageMetadata"); cpaUsage.Exists() { + chunk, _ = sjson.SetRawBytes(chunk, "usageMetadata", []byte(cpaUsage.Raw)) + chunk, _ = sjson.DeleteBytes(chunk, "cpaUsageMetadata") + } + return chunk +} diff --git a/internal/translator/antigravity/gemini/antigravity_gemini_response_test.go b/internal/translator/antigravity/gemini/antigravity_gemini_response_test.go new file mode 100644 index 00000000..5f96012a --- /dev/null +++ b/internal/translator/antigravity/gemini/antigravity_gemini_response_test.go @@ -0,0 +1,95 @@ +package gemini + +import ( + "context" + "testing" +) + +func TestRestoreUsageMetadata(t *testing.T) { + tests := []struct { + name string + input []byte + expected string + }{ + { + name: "cpaUsageMetadata renamed to usageMetadata", + input: []byte(`{"modelVersion":"gemini-3-pro","cpaUsageMetadata":{"promptTokenCount":100,"candidatesTokenCount":200}}`), + expected: `{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100,"candidatesTokenCount":200}}`, + }, + { + name: "no cpaUsageMetadata unchanged", + input: []byte(`{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100}}`), + expected: `{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100}}`, + }, + { + name: "empty input", + input: []byte(`{}`), + expected: `{}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := restoreUsageMetadata(tt.input) + if string(result) != tt.expected { + t.Errorf("restoreUsageMetadata() = %s, want %s", string(result), tt.expected) + } + }) + } +} + +func TestConvertAntigravityResponseToGeminiNonStream(t *testing.T) { + tests := []struct { + name string + input []byte + expected string + }{ + { + name: "cpaUsageMetadata restored in response", + input: []byte(`{"response":{"modelVersion":"gemini-3-pro","cpaUsageMetadata":{"promptTokenCount":100}}}`), + expected: `{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100}}`, + }, + { + name: "usageMetadata preserved", + input: []byte(`{"response":{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100}}}`), + expected: `{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100}}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := ConvertAntigravityResponseToGeminiNonStream(context.Background(), "", nil, nil, tt.input, nil) + if result != tt.expected { + t.Errorf("ConvertAntigravityResponseToGeminiNonStream() = %s, want %s", result, tt.expected) + } + }) + } +} + +func TestConvertAntigravityResponseToGeminiStream(t *testing.T) { + ctx := context.WithValue(context.Background(), "alt", "") + + tests := []struct { + name string + input []byte + expected string + }{ + { + name: "cpaUsageMetadata restored in streaming response", + input: []byte(`data: {"response":{"modelVersion":"gemini-3-pro","cpaUsageMetadata":{"promptTokenCount":100}}}`), + expected: `{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100}}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + results := ConvertAntigravityResponseToGemini(ctx, "", nil, nil, tt.input, nil) + if len(results) != 1 { + t.Fatalf("expected 1 result, got %d", len(results)) + } + if results[0] != tt.expected { + t.Errorf("ConvertAntigravityResponseToGemini() = %s, want %s", results[0], tt.expected) + } + }) + } +}