fix: include Gemini thoughtsTokenCount in output token billing

Gemini 2.5 Pro/Flash thinking models return thoughtsTokenCount separately
from candidatesTokenCount in usageMetadata, but this field was not parsed
or included in billing calculations, causing thinking tokens to be
unbilled.

- Add ThoughtsTokenCount field to GeminiUsageMetadata struct
- Include thoughtsTokenCount in OutputTokens across all 3 Gemini usage
  parsing paths (non-streaming, streaming, compat layer)
- Add tests covering thinking token scenarios

Closes #554
This commit is contained in:
sususu98
2026-02-11 15:41:54 +08:00
parent ae6fed15cc
commit d21d70a5cf
6 changed files with 144 additions and 4 deletions

View File

@@ -4,6 +4,8 @@ import (
"encoding/json"
"strings"
"testing"
"github.com/stretchr/testify/require"
)
// TestConvertClaudeToolsToGeminiTools_CustomType 测试custom类型工具转换
@@ -203,3 +205,70 @@ func TestEnsureGeminiFunctionCallThoughtSignatures_InsertsWhenMissing(t *testing
t.Fatalf("expected injected thoughtSignature %q, got: %s", geminiDummyThoughtSignature, s)
}
}
func TestExtractGeminiUsage_ThoughtsTokenCount(t *testing.T) {
tests := []struct {
name string
resp map[string]any
wantInput int
wantOutput int
wantCacheRead int
wantNil bool
}{
{
name: "with thoughtsTokenCount",
resp: map[string]any{
"usageMetadata": map[string]any{
"promptTokenCount": float64(100),
"candidatesTokenCount": float64(20),
"thoughtsTokenCount": float64(50),
},
},
wantInput: 100,
wantOutput: 70,
},
{
name: "with thoughtsTokenCount and cache",
resp: map[string]any{
"usageMetadata": map[string]any{
"promptTokenCount": float64(100),
"candidatesTokenCount": float64(20),
"cachedContentTokenCount": float64(30),
"thoughtsTokenCount": float64(50),
},
},
wantInput: 70,
wantOutput: 70,
wantCacheRead: 30,
},
{
name: "without thoughtsTokenCount (old model)",
resp: map[string]any{
"usageMetadata": map[string]any{
"promptTokenCount": float64(100),
"candidatesTokenCount": float64(20),
},
},
wantInput: 100,
wantOutput: 20,
},
{
name: "no usageMetadata",
resp: map[string]any{},
wantNil: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
usage := extractGeminiUsage(tt.resp)
if tt.wantNil {
require.Nil(t, usage)
return
}
require.NotNil(t, usage)
require.Equal(t, tt.wantInput, usage.InputTokens)
require.Equal(t, tt.wantOutput, usage.OutputTokens)
require.Equal(t, tt.wantCacheRead, usage.CacheReadInputTokens)
})
}
}