🤖 fix: correct context usage display for multi-step tool calls (#893)

## Problem

The Context Usage UI was sometimes showing inflated token counts for
multi-step tool calls (e.g., plan messages showing ~150k
`cachedInputTokens` instead of ~50k).

## Root Cause

Two issues combined:

1. **Backend**: On stream-end, we only used `lastStepUsage` tracked from
`finish-step` events. If no `finish-step` event was received (edge
cases), `contextUsage` would be undefined. We weren't fetching
`streamResult.usage` (which contains last step's usage) as the primary
source.

2. **Frontend**: When `contextUsage` was undefined, it fell back to
cumulative `usage`. For multi-step requests, cumulative usage sums
`cachedInputTokens` across all steps (each step reads from cache),
inflating the context window display.

<img width="297" height="72" alt="image"
src="https://github.com/user-attachments/assets/c8a58e76-df24-4003-9d3e-369417173cd3"
/>



## Fix

**Backend (`streamManager.ts`)**:
- Fetch `contextUsage` from `streamResult.usage` (last step) on
stream-end as the primary source
- Fall back to tracked `lastStepUsage` only if `streamResult.usage`
times out
- Each metadata fetch has independent timeout/error handling so one
failure doesn't mask others

**Frontend (`WorkspaceStore.ts`)**:
- Remove fallback from `contextUsage` to `usage` - only use
`contextUsage` for context window display

**Schema (`message.ts`)**:
- Add `contextUsage` and `contextProviderMetadata` to Zod schema to
match TypeScript types

## Migration Note

⚠️ **Old workspaces** created before this fix may not have
`contextUsage` stored in their message history. These workspaces will
show no context usage until the next message is sent.

<img width="295" height="62" alt="image"
src="https://github.com/user-attachments/assets/07c7a4c0-c2ea-4f0c-a0b1-9a45fbd789a9"
/>

---
_Generated with `mux`_

Author: ethanndickson

Makefile

@@ -224,7 +224,7 @@ lint: node_modules/.installed ## Run ESLint (typecheck runs in separate target)
 lint-fix: node_modules/.installed ## Run linter with --fix
 	@./scripts/lint.sh --fix
 
-ifeq ($(OS),Windows_NT) 
+ifeq ($(OS),Windows_NT)
 typecheck: node_modules/.installed src/version.ts ## Run TypeScript type checking (uses tsgo for 10x speedup)
 	@# On Windows, use npm run because bun x doesn't correctly pass arguments
 	@npmx concurrently -g \

bun.lock

@@ -1,5 +1,6 @@
 {
   "lockfileVersion": 1,
+  "configVersion": 0,
   "workspaces": {
     "": {
       "name": "mux",

src/browser/stores/WorkspaceStore.ts

@@ -486,7 +486,7 @@ export class WorkspaceStore {
             if (msg.metadata?.compacted) {
               continue;
             }
-            const rawUsage = msg.metadata?.contextUsage ?? msg.metadata?.usage;
+            const rawUsage = msg.metadata?.contextUsage;
             const providerMeta =
               msg.metadata?.contextProviderMetadata ?? msg.metadata?.providerMetadata;
             if (rawUsage) {

src/browser/stories/mockFactory.ts

@@ -199,6 +199,7 @@ export function createAssistantMessage(
       timestamp: opts.timestamp ?? STABLE_TIMESTAMP,
       model: opts.model ?? DEFAULT_MODEL,
       usage: { inputTokens: 100, outputTokens: 50, totalTokens: 150 },
+      contextUsage: { inputTokens: 100, outputTokens: 50, totalTokens: 150 },
       duration: 1000,
     },
   };

src/common/orpc/schemas/message.ts

@@ -76,7 +76,9 @@ export const MuxMessageSchema = z.object({
       timestamp: z.number().optional(),
       model: z.string().optional(),
       usage: z.any().optional(),
+      contextUsage: z.any().optional(),
       providerMetadata: z.record(z.string(), z.unknown()).optional(),
+      contextProviderMetadata: z.record(z.string(), z.unknown()).optional(),
       duration: z.number().optional(),
       systemMessageTokens: z.number().optional(),
       muxMetadata: z.any().optional(),

src/node/services/streamManager.ts

@@ -334,22 +334,34 @@ export class StreamManager extends EventEmitter {
   private async getStreamMetadata(
     streamInfo: WorkspaceStreamInfo,
     timeoutMs = 1000
-  ): Promise<{ usage?: LanguageModelV2Usage; duration: number }> {
-    let usage = undefined;
-    try {
-      // Race usage retrieval against timeout to prevent hanging on abort
-      // CRITICAL: Use totalUsage (sum of all steps) not usage (last step only)
-      // For multi-step tool calls, usage would severely undercount actual token consumption
-      usage = await Promise.race([
-        streamInfo.streamResult.totalUsage,
+  ): Promise<{
+    totalUsage?: LanguageModelV2Usage;
+    contextUsage?: LanguageModelV2Usage;
+    contextProviderMetadata?: Record<string, unknown>;
+    duration: number;
+  }> {
+    // Helper: wrap promise with independent timeout + error handling
+    // Each promise resolves independently - one failure doesn't mask others
+    const withTimeout = <T>(promise: Promise<T>): Promise<T | undefined> =>
+      Promise.race([
+        promise,
         new Promise<undefined>((resolve) => setTimeout(() => resolve(undefined), timeoutMs)),
-      ]);
-    } catch (error) {
-      log.debug("Could not retrieve usage:", error);
-    }
+      ]).catch(() => undefined);
+
+    // Fetch all metadata in parallel with independent timeouts
+    // - totalUsage: sum of all steps (for cost calculation)
+    // - contextUsage: last step only (for context window display)
+    // - contextProviderMetadata: last step (for context window cache display)
+    const [totalUsage, contextUsage, contextProviderMetadata] = await Promise.all([
+      withTimeout(streamInfo.streamResult.totalUsage),
+      withTimeout(streamInfo.streamResult.usage),
+      withTimeout(streamInfo.streamResult.providerMetadata),
+    ]);
 
     return {
-      usage,
+      totalUsage,
+      contextUsage,
+      contextProviderMetadata,
       duration: Date.now() - streamInfo.startTime,
     };
   }
@@ -1045,17 +1057,20 @@ export class StreamManager extends EventEmitter {
 
       // Check if stream completed successfully
       if (!streamInfo.abortController.signal.aborted) {
-        // Get usage, duration, and provider metadata from stream result
-        // CRITICAL: Use totalUsage (via getStreamMetadata) and aggregated providerMetadata
-        // to correctly account for all steps in multi-tool-call conversations
-        const { usage, duration } = await this.getStreamMetadata(streamInfo);
+        // Get all metadata from stream result in one call
+        // - totalUsage: sum of all steps (for cost calculation)
+        // - contextUsage: last step only (for context window display)
+        // - contextProviderMetadata: last step (for context window cache tokens)
+        // Falls back to tracked values from finish-step if streamResult fails/times out
+        const streamMeta = await this.getStreamMetadata(streamInfo);
+        const totalUsage = streamMeta.totalUsage;
+        const contextUsage = streamMeta.contextUsage ?? streamInfo.lastStepUsage;
+        const contextProviderMetadata =
+          streamMeta.contextProviderMetadata ?? streamInfo.lastStepProviderMetadata;
+        const duration = streamMeta.duration;
+        // Aggregated provider metadata across all steps (for cost calculation with cache tokens)
         const providerMetadata = await this.getAggregatedProviderMetadata(streamInfo);
 
-        // For context window display, use last step's usage (inputTokens = current context size)
-        // This is stored in streamInfo during finish-step handling
-        const contextUsage = streamInfo.lastStepUsage;
-        const contextProviderMetadata = streamInfo.lastStepProviderMetadata;
-
         // Emit stream end event with parts preserved in temporal order
         const streamEndEvent: StreamEndEvent = {
           type: "stream-end",
@@ -1064,7 +1079,7 @@ export class StreamManager extends EventEmitter {
           metadata: {
             ...streamInfo.initialMetadata, // AIService-provided metadata (systemMessageTokens, etc)
             model: streamInfo.model,
-            usage, // Total across all steps (for cost calculation)
+            usage: totalUsage, // Total across all steps (for cost calculation)
             contextUsage, // Last step only (for context window display)
             providerMetadata, // Aggregated (for cost calculation)
             contextProviderMetadata, // Last step (for context window display)