feat(core): add per-model token usage to stream-json output (#21839)

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: {},
         },
       };
 

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<string, ModelStreamStats>,
+        },
+      );
 
     return {
       total_tokens: totalTokens,
@@ -63,6 +82,7 @@ export class StreamJsonFormatter {
       input,
       duration_ms: durationMs,
       tool_calls: metrics.tools.totalCalls,
+      models,
     };
   }
 }

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<string, ModelStreamStats>;
 }
 
 export interface ResultEvent extends BaseJsonStreamEvent {