🤖 fix: accumulate output tokens across steps in cost tracking

In multi-step mode (tool use), the AI SDK emits per-step usage where:
- inputTokens: Current context window (grows each step)
- outputTokens: Tokens generated in that step only

The aggregator was replacing usage on each step, losing earlier output
tokens. Now we accumulate output/reasoning tokens while keeping the
latest input tokens (which correctly represent the growing context).

This fixes the issue where costs appeared to increment by only ~2¢
at a time instead of reflecting cumulative output costs.

_Generated with mux_

Author: ethanndickson

src/browser/utils/messages/StreamingMessageAggregator.test.ts

@@ -415,10 +415,10 @@ describe("StreamingMessageAggregator", () => {
       expect(aggregator.getActiveStreamUsage("msg-1")).toBeUndefined();
     });
 
-    test("latest usage-delta replaces previous for same messageId", () => {
+    test("accumulates output tokens across steps, keeps latest input", () => {
       const aggregator = new StreamingMessageAggregator(TEST_CREATED_AT);
 
-      // First step usage
+      // First step usage (AI calls a tool)
       aggregator.handleUsageDelta({
         type: "usage-delta",
         workspaceId: "ws-1",
@@ -434,12 +434,12 @@ describe("StreamingMessageAggregator", () => {
         usage: { inputTokens: 1500, outputTokens: 100, totalTokens: 1600 },
       });
 
-      // Should have latest values, not summed
-      expect(aggregator.getActiveStreamUsage("msg-1")).toEqual({
-        inputTokens: 1500,
-        outputTokens: 100,
-        totalTokens: 1600,
-      });
+      // Input should be latest (context window grows), output should be summed
+      const usage = aggregator.getActiveStreamUsage("msg-1");
+      expect(usage?.inputTokens).toBe(1500); // Latest context window
+      expect(usage?.outputTokens).toBe(150); // Accumulated: 50 + 100
+      expect(usage?.totalTokens).toBe(1650); // Latest input + accumulated output
+      expect(usage?.reasoningTokens).toBe(0); // Accumulated (none in this test)
     });
 
     test("tracks usage independently per messageId", () => {

src/browser/utils/messages/StreamingMessageAggregator.ts

@@ -1002,10 +1002,31 @@ export class StreamingMessageAggregator {
   }
 
   /**
-   * Handle usage-delta event: update cumulative usage for active stream
+   * Handle usage-delta event: accumulate usage for active stream across steps
+   *
+   * In multi-step mode (tool use), each step emits per-step usage:
+   * - inputTokens: Current context window (grows each step) - use latest
+   * - cachedInputTokens: Current cache state - use latest
+   * - outputTokens: Tokens generated this step only - accumulate
+   * - reasoningTokens: Reasoning tokens this step only - accumulate
    */
   handleUsageDelta(data: UsageDeltaEvent): void {
-    this.activeStreamStepUsage.set(data.messageId, data.usage);
+    const existing = this.activeStreamStepUsage.get(data.messageId);
+    if (existing) {
+      // Accumulate output/reasoning tokens, keep latest input (context grows each step)
+      const outputTokens = (existing.outputTokens ?? 0) + (data.usage.outputTokens ?? 0);
+      const reasoningTokens = (existing.reasoningTokens ?? 0) + (data.usage.reasoningTokens ?? 0);
+      this.activeStreamStepUsage.set(data.messageId, {
+        inputTokens: data.usage.inputTokens, // Latest context window
+        cachedInputTokens: data.usage.cachedInputTokens, // Latest cache state
+        outputTokens,
+        reasoningTokens,
+        totalTokens:
+          (data.usage.inputTokens ?? 0) + (data.usage.cachedInputTokens ?? 0) + outputTokens,
+      });
+    } else {
+      this.activeStreamStepUsage.set(data.messageId, data.usage);
+    }
   }
 
   /**