github-mirrors/CLIProxyAPI
1
0
Fork 0
mirror of https://mirror.skon.top/github.com/router-for-me/CLIProxyAPI synced 2026-04-30 16:20:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: parse gemini cli usage metadata variants

Browse Source
This commit is contained in:
sususu98
2026-04-29 12:46:53 +08:00
parent 2ea8f77efb
commit 4982512da2
3 changed files with 80 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
internal/runtime/executor/gemini_cli_executor.go
View File

@@ -422,7 +422,9 @@ func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut
helps.RecordAPIResponseError(ctx, e.cfg, errScan)
reporter.PublishFailure(ctx)
out <- cliproxyexecutor.StreamChunk{Err: errScan}
return
}
reporter.EnsurePublished(ctx)
return
}

42
internal/runtime/executor/helps/usage_helpers.go
View File

@@ -370,12 +370,22 @@ func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail {
return detail
}
func hasGeminiFamilyUsageTokenFields(node gjson.Result) bool {
return node.Get("promptTokenCount").Exists() ||
node.Get("candidatesTokenCount").Exists() ||
node.Get("thoughtsTokenCount").Exists() ||
node.Get("totalTokenCount").Exists() ||
node.Get("cachedContentTokenCount").Exists()
}
func ParseGeminiCLIUsage(data []byte) usage.Detail {
usageNode := gjson.ParseBytes(data)
node := usageNode.Get("response.usageMetadata")
if !node.Exists() {
node = usageNode.Get("response.usage_metadata")
}
node := firstExistingUsageNode(usageNode,
"response.usageMetadata",
"response.usage_metadata",
"usageMetadata",
"usage_metadata",
)
if !node.Exists() {
return usage.Detail{}
}
@@ -414,16 +424,32 @@ func ParseGeminiCLIStreamUsage(line []byte) (usage.Detail, bool) {
if len(payload) == 0 || !gjson.ValidBytes(payload) {
return usage.Detail{}, false
}
node := gjson.GetBytes(payload, "response.usageMetadata")
if !node.Exists() {
node = gjson.GetBytes(payload, "usage_metadata")
}
root := gjson.ParseBytes(payload)
node := firstExistingUsageNode(root,
"response.usageMetadata",
"response.usage_metadata",
"usageMetadata",
"usage_metadata",
)
if !node.Exists() {
return usage.Detail{}, false
}
if !hasGeminiFamilyUsageTokenFields(node) {
return usage.Detail{}, false
}
return parseGeminiFamilyUsageDetail(node), true
}
func firstExistingUsageNode(root gjson.Result, paths ...string) gjson.Result {
for _, path := range paths {
node := root.Get(path)
if node.Exists() {
return node
}
}
return gjson.Result{}
}
func ParseAntigravityUsage(data []byte) usage.Detail {
usageNode := gjson.ParseBytes(data)
node := usageNode.Get("response.usageMetadata")

44
internal/runtime/executor/helps/usage_helpers_test.go
View File

@@ -47,6 +47,50 @@ func TestParseOpenAIUsageResponses(t *testing.T) {
}
}
func TestParseGeminiCLIUsage_TopLevelUsageMetadata(t *testing.T) {
data := []byte(`{"usageMetadata":{"promptTokenCount":11,"candidatesTokenCount":7,"thoughtsTokenCount":3,"totalTokenCount":21,"cachedContentTokenCount":5}}`)
detail := ParseGeminiCLIUsage(data)
if detail.InputTokens != 11 {
t.Fatalf("input tokens = %d, want %d", detail.InputTokens, 11)
}
if detail.OutputTokens != 7 {
t.Fatalf("output tokens = %d, want %d", detail.OutputTokens, 7)
}
if detail.ReasoningTokens != 3 {
t.Fatalf("reasoning tokens = %d, want %d", detail.ReasoningTokens, 3)
}
if detail.TotalTokens != 21 {
t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 21)
}
if detail.CachedTokens != 5 {
t.Fatalf("cached tokens = %d, want %d", detail.CachedTokens, 5)
}
}
func TestParseGeminiCLIStreamUsage_ResponseSnakeCaseUsageMetadata(t *testing.T) {
line := []byte(`data: {"response":{"usage_metadata":{"promptTokenCount":13,"candidatesTokenCount":2,"totalTokenCount":15}}}`)
detail, ok := ParseGeminiCLIStreamUsage(line)
if !ok {
t.Fatal("ParseGeminiCLIStreamUsage() ok = false, want true")
}
if detail.InputTokens != 13 {
t.Fatalf("input tokens = %d, want %d", detail.InputTokens, 13)
}
if detail.OutputTokens != 2 {
t.Fatalf("output tokens = %d, want %d", detail.OutputTokens, 2)
}
if detail.TotalTokens != 15 {
t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 15)
}
}
func TestParseGeminiCLIStreamUsage_IgnoresTrafficTypeOnlyUsageMetadata(t *testing.T) {
line := []byte(`data: {"response":{"usageMetadata":{"trafficType":"ON_DEMAND"}}}`)
if detail, ok := ParseGeminiCLIStreamUsage(line); ok {
t.Fatalf("ParseGeminiCLIStreamUsage() = (%+v, true), want false for traffic-only usage metadata", detail)
}
}
func TestUsageReporterBuildRecordIncludesLatency(t *testing.T) {
reporter := &UsageReporter{
provider: "openai",