From 4da0366eed481a4e81c3d6eb6ee5aec061e77c8a Mon Sep 17 00:00:00 2001 From: Yongrui Lin Date: Tue, 10 Mar 2026 10:39:04 -0700 Subject: [PATCH] feat(core): add per-model token usage to stream-json output (#21839) --- .../nonInteractiveCli.test.ts.snap | 6 +-- packages/cli/src/utils/errors.test.ts | 1 + .../src/output/stream-json-formatter.test.ts | 39 +++++++++++++++ .../core/src/output/stream-json-formatter.ts | 50 +++++++++++++------ packages/core/src/output/types.ts | 9 ++++ 5 files changed, 87 insertions(+), 18 deletions(-) diff --git a/packages/cli/src/__snapshots__/nonInteractiveCli.test.ts.snap b/packages/cli/src/__snapshots__/nonInteractiveCli.test.ts.snap index 8c1a85cdd7..92f396a59c 100644 --- a/packages/cli/src/__snapshots__/nonInteractiveCli.test.ts.snap +++ b/packages/cli/src/__snapshots__/nonInteractiveCli.test.ts.snap @@ -4,7 +4,7 @@ exports[`runNonInteractive > should emit appropriate error event in streaming JS "{"type":"init","timestamp":"","session_id":"test-session-id","model":"test-model"} {"type":"message","timestamp":"","role":"user","content":"Loop test"} {"type":"error","timestamp":"","severity":"warning","message":"Loop detected, stopping execution"} -{"type":"result","timestamp":"","status":"success","stats":{"total_tokens":0,"input_tokens":0,"output_tokens":0,"cached":0,"input":0,"duration_ms":,"tool_calls":0}} +{"type":"result","timestamp":"","status":"success","stats":{"total_tokens":0,"input_tokens":0,"output_tokens":0,"cached":0,"input":0,"duration_ms":,"tool_calls":0,"models":{}}} " `; @@ -12,7 +12,7 @@ exports[`runNonInteractive > should emit appropriate error event in streaming JS "{"type":"init","timestamp":"","session_id":"test-session-id","model":"test-model"} {"type":"message","timestamp":"","role":"user","content":"Max turns test"} {"type":"error","timestamp":"","severity":"error","message":"Maximum session turns exceeded"} -{"type":"result","timestamp":"","status":"success","stats":{"total_tokens":0,"input_tokens":0,"output_tokens":0,"cached":0,"input":0,"duration_ms":,"tool_calls":0}} +{"type":"result","timestamp":"","status":"success","stats":{"total_tokens":0,"input_tokens":0,"output_tokens":0,"cached":0,"input":0,"duration_ms":,"tool_calls":0,"models":{}}} " `; @@ -23,7 +23,7 @@ exports[`runNonInteractive > should emit appropriate events for streaming JSON o {"type":"tool_use","timestamp":"","tool_name":"testTool","tool_id":"tool-1","parameters":{"arg1":"value1"}} {"type":"tool_result","timestamp":"","tool_id":"tool-1","status":"success","output":"Tool executed successfully"} {"type":"message","timestamp":"","role":"assistant","content":"Final answer","delta":true} -{"type":"result","timestamp":"","status":"success","stats":{"total_tokens":0,"input_tokens":0,"output_tokens":0,"cached":0,"input":0,"duration_ms":,"tool_calls":0}} +{"type":"result","timestamp":"","status":"success","stats":{"total_tokens":0,"input_tokens":0,"output_tokens":0,"cached":0,"input":0,"duration_ms":,"tool_calls":0,"models":{}}} " `; diff --git a/packages/cli/src/utils/errors.test.ts b/packages/cli/src/utils/errors.test.ts index c5b7a7e7fe..38ee059bbe 100644 --- a/packages/cli/src/utils/errors.test.ts +++ b/packages/cli/src/utils/errors.test.ts @@ -74,6 +74,7 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { input: 0, duration_ms: 0, tool_calls: 0, + models: {}, }), })), uiTelemetryService: { diff --git a/packages/core/src/output/stream-json-formatter.test.ts b/packages/core/src/output/stream-json-formatter.test.ts index c911a9dbc2..f4f3ae07a0 100644 --- a/packages/core/src/output/stream-json-formatter.test.ts +++ b/packages/core/src/output/stream-json-formatter.test.ts @@ -154,6 +154,7 @@ describe('StreamJsonFormatter', () => { input: 50, duration_ms: 1200, tool_calls: 2, + models: {}, }, }; @@ -180,6 +181,7 @@ describe('StreamJsonFormatter', () => { input: 50, duration_ms: 1200, tool_calls: 0, + models: {}, }, }; @@ -304,6 +306,15 @@ describe('StreamJsonFormatter', () => { input: 50, duration_ms: 1200, tool_calls: 2, + models: { + 'gemini-2.0-flash': { + total_tokens: 80, + input_tokens: 50, + output_tokens: 30, + cached: 0, + input: 50, + }, + }, }); }); @@ -347,6 +358,22 @@ describe('StreamJsonFormatter', () => { input: 150, duration_ms: 3000, tool_calls: 5, + models: { + 'gemini-pro': { + total_tokens: 80, + input_tokens: 50, + output_tokens: 30, + cached: 0, + input: 50, + }, + 'gemini-ultra': { + total_tokens: 170, + input_tokens: 100, + output_tokens: 70, + cached: 0, + input: 100, + }, + }, }); }); @@ -376,6 +403,15 @@ describe('StreamJsonFormatter', () => { input: 20, duration_ms: 1200, tool_calls: 0, + models: { + 'gemini-pro': { + total_tokens: 80, + input_tokens: 50, + output_tokens: 30, + cached: 30, + input: 20, + }, + }, }); }); @@ -392,6 +428,7 @@ describe('StreamJsonFormatter', () => { input: 0, duration_ms: 100, tool_calls: 0, + models: {}, }); }); @@ -521,6 +558,7 @@ describe('StreamJsonFormatter', () => { input: 0, duration_ms: 0, tool_calls: 0, + models: {}, }, } as ResultEvent, ]; @@ -544,6 +582,7 @@ describe('StreamJsonFormatter', () => { input: 50, duration_ms: 1200, tool_calls: 2, + models: {}, }, }; diff --git a/packages/core/src/output/stream-json-formatter.ts b/packages/core/src/output/stream-json-formatter.ts index 585dbb0789..6475e6d482 100644 --- a/packages/core/src/output/stream-json-formatter.ts +++ b/packages/core/src/output/stream-json-formatter.ts @@ -4,7 +4,11 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { JsonStreamEvent, StreamStats } from './types.js'; +import type { + JsonStreamEvent, + ModelStreamStats, + StreamStats, +} from './types.js'; import type { SessionMetrics } from '../telemetry/uiTelemetry.js'; /** @@ -31,7 +35,7 @@ export class StreamJsonFormatter { /** * Converts SessionMetrics to simplified StreamStats format. - * Aggregates token counts across all models. + * Includes per-model token breakdowns and aggregated totals. * @param metrics - The session metrics from telemetry * @param durationMs - The session duration in milliseconds * @returns Simplified stats for streaming output @@ -40,20 +44,35 @@ export class StreamJsonFormatter { metrics: SessionMetrics, durationMs: number, ): StreamStats { - let totalTokens = 0; - let inputTokens = 0; - let outputTokens = 0; - let cached = 0; - let input = 0; + const { totalTokens, inputTokens, outputTokens, cached, input, models } = + Object.entries(metrics.models).reduce( + (acc, [modelName, modelMetrics]) => { + const modelStats: ModelStreamStats = { + total_tokens: modelMetrics.tokens.total, + input_tokens: modelMetrics.tokens.prompt, + output_tokens: modelMetrics.tokens.candidates, + cached: modelMetrics.tokens.cached, + input: modelMetrics.tokens.input, + }; - // Aggregate token counts across all models - for (const modelMetrics of Object.values(metrics.models)) { - totalTokens += modelMetrics.tokens.total; - inputTokens += modelMetrics.tokens.prompt; - outputTokens += modelMetrics.tokens.candidates; - cached += modelMetrics.tokens.cached; - input += modelMetrics.tokens.input; - } + acc.models[modelName] = modelStats; + acc.totalTokens += modelStats.total_tokens; + acc.inputTokens += modelStats.input_tokens; + acc.outputTokens += modelStats.output_tokens; + acc.cached += modelStats.cached; + acc.input += modelStats.input; + + return acc; + }, + { + totalTokens: 0, + inputTokens: 0, + outputTokens: 0, + cached: 0, + input: 0, + models: {} as Record, + }, + ); return { total_tokens: totalTokens, @@ -63,6 +82,7 @@ export class StreamJsonFormatter { input, duration_ms: durationMs, tool_calls: metrics.tools.totalCalls, + models, }; } } diff --git a/packages/core/src/output/types.ts b/packages/core/src/output/types.ts index 0c129eac93..c67c8afe99 100644 --- a/packages/core/src/output/types.ts +++ b/packages/core/src/output/types.ts @@ -77,6 +77,14 @@ export interface ErrorEvent extends BaseJsonStreamEvent { message: string; } +export interface ModelStreamStats { + total_tokens: number; + input_tokens: number; + output_tokens: number; + cached: number; + input: number; +} + export interface StreamStats { total_tokens: number; input_tokens: number; @@ -86,6 +94,7 @@ export interface StreamStats { input: number; duration_ms: number; tool_calls: number; + models: Record; } export interface ResultEvent extends BaseJsonStreamEvent {