diff --git a/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response.go b/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response.go index 60fe4895..54458e6b 100644 --- a/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response.go +++ b/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response.go @@ -95,9 +95,7 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq // Extract and set usage metadata (token counts). if usageResult := gjson.GetBytes(rawJSON, "response.usageMetadata"); usageResult.Exists() { cachedTokenCount := usageResult.Get("cachedContentTokenCount").Int() - if candidatesTokenCountResult := usageResult.Get("candidatesTokenCount"); candidatesTokenCountResult.Exists() { - template, _ = sjson.SetBytes(template, "usage.completion_tokens", candidatesTokenCountResult.Int()) - } + template, _ = sjson.SetBytes(template, "usage.completion_tokens", usageResult.Get("candidatesTokenCount").Int()) if totalTokenCountResult := usageResult.Get("totalTokenCount"); totalTokenCountResult.Exists() { template, _ = sjson.SetBytes(template, "usage.total_tokens", totalTokenCountResult.Int()) } diff --git a/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response_test.go b/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response_test.go index 18f0f58a..39429e06 100644 --- a/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response_test.go +++ b/internal/translator/antigravity/openai/chat-completions/antigravity_openai_response_test.go @@ -128,6 +128,20 @@ func TestNoFinishReasonOnIntermediateChunks(t *testing.T) { } } +func TestConvertAntigravityResponseToOpenAIIncludesZeroCompletionTokensWhenMissing(t *testing.T) { + var param any + chunk := []byte(`{"response":{"usageMetadata":{"promptTokenCount":16,"thoughtsTokenCount":42,"totalTokenCount":58}}}`) + + result := ConvertAntigravityResponseToOpenAI(context.Background(), "model", nil, nil, chunk, ¶m) + if len(result) != 1 { + t.Fatalf("expected 1 result, got %d", len(result)) + } + completionTokens := gjson.GetBytes(result[0], "usage.completion_tokens") + if !completionTokens.Exists() || completionTokens.Int() != 0 { + t.Fatalf("completion_tokens = %s, want present with value 0. Output: %s", completionTokens.Raw, result[0]) + } +} + func TestConvertAntigravityResponseToOpenAINonStreamRestoresDisambiguatedName(t *testing.T) { first := "mcp__plugin_cloudflare_cloudflare-builds__workers_builds_get_build" second := "mcp__plugin_cloudflare_cloudflare-builds__workers_builds_get_build_logs" diff --git a/internal/translator/gemini/openai/chat-completions/gemini_openai_response.go b/internal/translator/gemini/openai/chat-completions/gemini_openai_response.go index cb44c258..b9c68aa0 100644 --- a/internal/translator/gemini/openai/chat-completions/gemini_openai_response.go +++ b/internal/translator/gemini/openai/chat-completions/gemini_openai_response.go @@ -110,9 +110,7 @@ func ConvertGeminiResponseToOpenAI(_ context.Context, _ string, originalRequestR // Usage is applied to the base template so it appears in the chunks. if usageResult := gjson.GetBytes(rawJSON, "usageMetadata"); usageResult.Exists() { cachedTokenCount := usageResult.Get("cachedContentTokenCount").Int() - if candidatesTokenCountResult := usageResult.Get("candidatesTokenCount"); candidatesTokenCountResult.Exists() { - baseTemplate, _ = sjson.SetBytes(baseTemplate, "usage.completion_tokens", candidatesTokenCountResult.Int()) - } + baseTemplate, _ = sjson.SetBytes(baseTemplate, "usage.completion_tokens", usageResult.Get("candidatesTokenCount").Int()) if totalTokenCountResult := usageResult.Get("totalTokenCount"); totalTokenCountResult.Exists() { baseTemplate, _ = sjson.SetBytes(baseTemplate, "usage.total_tokens", totalTokenCountResult.Int()) } @@ -312,9 +310,7 @@ func ConvertGeminiResponseToOpenAINonStream(_ context.Context, _ string, origina } if usageResult := gjson.GetBytes(rawJSON, "usageMetadata"); usageResult.Exists() { - if candidatesTokenCountResult := usageResult.Get("candidatesTokenCount"); candidatesTokenCountResult.Exists() { - template, _ = sjson.SetBytes(template, "usage.completion_tokens", candidatesTokenCountResult.Int()) - } + template, _ = sjson.SetBytes(template, "usage.completion_tokens", usageResult.Get("candidatesTokenCount").Int()) if totalTokenCountResult := usageResult.Get("totalTokenCount"); totalTokenCountResult.Exists() { template, _ = sjson.SetBytes(template, "usage.total_tokens", totalTokenCountResult.Int()) } diff --git a/internal/translator/gemini/openai/chat-completions/gemini_openai_response_test.go b/internal/translator/gemini/openai/chat-completions/gemini_openai_response_test.go index 177f4082..c2714c49 100644 --- a/internal/translator/gemini/openai/chat-completions/gemini_openai_response_test.go +++ b/internal/translator/gemini/openai/chat-completions/gemini_openai_response_test.go @@ -7,6 +7,30 @@ import ( "github.com/tidwall/gjson" ) +func TestConvertGeminiResponseToOpenAIIncludesZeroCompletionTokensWhenMissing(t *testing.T) { + var param any + chunk := []byte(`{"usageMetadata":{"promptTokenCount":16,"thoughtsTokenCount":42,"totalTokenCount":58}}`) + + result := ConvertGeminiResponseToOpenAI(context.Background(), "model", nil, nil, chunk, ¶m) + if len(result) != 1 { + t.Fatalf("expected 1 result, got %d", len(result)) + } + completionTokens := gjson.GetBytes(result[0], "usage.completion_tokens") + if !completionTokens.Exists() || completionTokens.Int() != 0 { + t.Fatalf("completion_tokens = %s, want present with value 0. Output: %s", completionTokens.Raw, result[0]) + } +} + +func TestConvertGeminiResponseToOpenAINonStreamIncludesZeroCompletionTokensWhenMissing(t *testing.T) { + response := []byte(`{"usageMetadata":{"promptTokenCount":16,"thoughtsTokenCount":42,"totalTokenCount":58}}`) + + result := ConvertGeminiResponseToOpenAINonStream(context.Background(), "model", nil, nil, response, nil) + completionTokens := gjson.GetBytes(result, "usage.completion_tokens") + if !completionTokens.Exists() || completionTokens.Int() != 0 { + t.Fatalf("completion_tokens = %s, want present with value 0. Output: %s", completionTokens.Raw, result) + } +} + func TestGeminiFinishReasonOnlyOnFinalChunk(t *testing.T) { ctx := context.Background() var param any