fix(translator): restore usageMetadata in Gemini responses from Antigravity
Unified · Split
| @@ -41,6 +41,7 @@ func ConvertAntigravityResponseToGemini(ctx context.Context, _ string, originalR · 40 unmodified lines | ||
| 41 | 41 | responseResult := gjson.GetBytes(rawJSON, "response") |
| 42 | 42 | if responseResult.Exists() { |
| 43 | 43 | chunk = []byte(responseResult.Raw) |
| 44 | + chunk = restoreUsageMetadata(chunk) | |
| 44 | 45 | } |
| 45 | 46 | } else { |
| 46 | 47 | chunkTemplate := "[]" |
| @@ -76,7 +77,8 @@ func ConvertAntigravityResponseToGemini(ctx context.Context, _ string, originalR · 29 unmodified lines | ||
| 76 | 77 | func ConvertAntigravityResponseToGeminiNonStream(_ context.Context, _ string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, _ *any) string { |
| 77 | 78 | responseResult := gjson.GetBytes(rawJSON, "response") |
| 78 | 79 | if responseResult.Exists() { |
| 79 | − return responseResult.Raw | |
| 80 | + chunk := restoreUsageMetadata([]byte(responseResult.Raw)) | |
| 81 | + return string(chunk) | |
| 80 | 82 | } |
| 81 | 83 | return string(rawJSON) |
| 82 | 84 | } |
| @@ -84,3 +86,15 @@ func ConvertAntigravityResponseToGeminiNonStream(_ context.Context, _ string, or · 1 unmodified line | ||
| 84 | 86 | func GeminiTokenCount(ctx context.Context, count int64) string { |
| 85 | 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. | |
| 94 | +func 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 | +} | |
| @@ -0,0 +1,95 @@ | ||
| 1 | +package gemini | |
| 2 | + | |
| 3 | +import ( | |
| 4 | + "context" | |
| 5 | + "testing" | |
| 6 | +) | |
| 7 | + | |
| 8 | +func 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 | + | |
| 41 | +func 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 | + | |
| 69 | +func 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 | +} | |