diff --git a/internal/runtime/executor/gemini_cli_executor.go b/internal/runtime/executor/gemini_cli_executor.go index a18f824a6..375989839 100644 --- a/internal/runtime/executor/gemini_cli_executor.go +++ b/internal/runtime/executor/gemini_cli_executor.go @@ -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 } diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index d3093de18..c5e258c86 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -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") diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index 3708b7317..c77335fd6 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -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",