History

Merge pull request #1298 from sususu98/fix/restore-usageMetadata-in-gemini-translator

fix(translator): restore usageMetadata in Gemini responses from Antigravity

Luis Pater committed 9 months ago

Commit da0bbf22 parents 4eb1e60 9dc0e6d2 files · +110 −1Browse files

Unified · Split

2 files
  1. internal/translator/antigravity/gemini/antigravity_gemini_response.go +15 −1
  2. internal/translator/antigravity/gemini/antigravity_gemini_response_test.go +95 −0
internal/translator/antigravity/gemini/antigravity_gemini_response.go+15−1
@@ -41,6 +41,7 @@ func ConvertAntigravityResponseToGemini(ctx context.Context, _ string, originalR · 40 unmodified lines
41 responseResult := gjson.GetBytes(rawJSON, "response")41 responseResult := gjson.GetBytes(rawJSON, "response")
42 if responseResult.Exists() {42 if responseResult.Exists() {
43 chunk = []byte(responseResult.Raw)43 chunk = []byte(responseResult.Raw)
44 chunk = restoreUsageMetadata(chunk)
44 }45 }
45 } else {46 } else {
46 chunkTemplate := "[]"47 chunkTemplate := "[]"
@@ -76,7 +77,8 @@ func ConvertAntigravityResponseToGemini(ctx context.Context, _ string, originalR · 29 unmodified lines
76func ConvertAntigravityResponseToGeminiNonStream(_ context.Context, _ string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, _ *any) string {77func ConvertAntigravityResponseToGeminiNonStream(_ context.Context, _ string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, _ *any) string {
77 responseResult := gjson.GetBytes(rawJSON, "response")78 responseResult := gjson.GetBytes(rawJSON, "response")
78 if responseResult.Exists() {79 if responseResult.Exists() {
79 return responseResult.Raw80 chunk := restoreUsageMetadata([]byte(responseResult.Raw))
81 return string(chunk)
80 }82 }
81 return string(rawJSON)83 return string(rawJSON)
82}84}
@@ -84,3 +86,15 @@ func ConvertAntigravityResponseToGeminiNonStream(_ context.Context, _ string, or · 1 unmodified line
84func GeminiTokenCount(ctx context.Context, count int64) string {86func GeminiTokenCount(ctx context.Context, count int64) string {
85 return fmt.Sprintf(`{"totalTokens":%d,"promptTokensDetails":[{"modality":"TEXT","tokenCount":%d}]}`, count, count)87 return fmt.Sprintf(`{"totalTokens":%d,"promptTokensDetails":[{"modality":"TEXT","tokenCount":%d}]}`, count, count)
86}88}
89
90// restoreUsageMetadata renames cpaUsageMetadata back to usageMetadata.
91// The executor renames usageMetadata to cpaUsageMetadata in non-terminal chunks
92// to preserve usage data while hiding it from clients that don't expect it.
93// When returning standard Gemini API format, we must restore the original name.
94func restoreUsageMetadata(chunk []byte) []byte {
95 if cpaUsage := gjson.GetBytes(chunk, "cpaUsageMetadata"); cpaUsage.Exists() {
96 chunk, _ = sjson.SetRawBytes(chunk, "usageMetadata", []byte(cpaUsage.Raw))
97 chunk, _ = sjson.DeleteBytes(chunk, "cpaUsageMetadata")
98 }
99 return chunk
100}
internal/translator/antigravity/gemini/antigravity_gemini_response_test.go+95−0
@@ -0,0 +1,95 @@
1package gemini
2
3import (
4 "context"
5 "testing"
6)
7
8func TestRestoreUsageMetadata(t *testing.T) {
9 tests := []struct {
10 name string
11 input []byte
12 expected string
13 }{
14 {
15 name: "cpaUsageMetadata renamed to usageMetadata",
16 input: []byte(`{"modelVersion":"gemini-3-pro","cpaUsageMetadata":{"promptTokenCount":100,"candidatesTokenCount":200}}`),
17 expected: `{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100,"candidatesTokenCount":200}}`,
18 },
19 {
20 name: "no cpaUsageMetadata unchanged",
21 input: []byte(`{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100}}`),
22 expected: `{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100}}`,
23 },
24 {
25 name: "empty input",
26 input: []byte(`{}`),
27 expected: `{}`,
28 },
29 }
30
31 for _, tt := range tests {
32 t.Run(tt.name, func(t *testing.T) {
33 result := restoreUsageMetadata(tt.input)
34 if string(result) != tt.expected {
35 t.Errorf("restoreUsageMetadata() = %s, want %s", string(result), tt.expected)
36 }
37 })
38 }
39}
40
41func TestConvertAntigravityResponseToGeminiNonStream(t *testing.T) {
42 tests := []struct {
43 name string
44 input []byte
45 expected string
46 }{
47 {
48 name: "cpaUsageMetadata restored in response",
49 input: []byte(`{"response":{"modelVersion":"gemini-3-pro","cpaUsageMetadata":{"promptTokenCount":100}}}`),
50 expected: `{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100}}`,
51 },
52 {
53 name: "usageMetadata preserved",
54 input: []byte(`{"response":{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100}}}`),
55 expected: `{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100}}`,
56 },
57 }
58
59 for _, tt := range tests {
60 t.Run(tt.name, func(t *testing.T) {
61 result := ConvertAntigravityResponseToGeminiNonStream(context.Background(), "", nil, nil, tt.input, nil)
62 if result != tt.expected {
63 t.Errorf("ConvertAntigravityResponseToGeminiNonStream() = %s, want %s", result, tt.expected)
64 }
65 })
66 }
67}
68
69func TestConvertAntigravityResponseToGeminiStream(t *testing.T) {
70 ctx := context.WithValue(context.Background(), "alt", "")
71
72 tests := []struct {
73 name string
74 input []byte
75 expected string
76 }{
77 {
78 name: "cpaUsageMetadata restored in streaming response",
79 input: []byte(`data: {"response":{"modelVersion":"gemini-3-pro","cpaUsageMetadata":{"promptTokenCount":100}}}`),
80 expected: `{"modelVersion":"gemini-3-pro","usageMetadata":{"promptTokenCount":100}}`,
81 },
82 }
83
84 for _, tt := range tests {
85 t.Run(tt.name, func(t *testing.T) {
86 results := ConvertAntigravityResponseToGemini(ctx, "", nil, nil, tt.input, nil)
87 if len(results) != 1 {
88 t.Fatalf("expected 1 result, got %d", len(results))
89 }
90 if results[0] != tt.expected {
91 t.Errorf("ConvertAntigravityResponseToGemini() = %s, want %s", results[0], tt.expected)
92 }
93 })
94 }
95}