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

## Problem

The Context Usage UI was showing inflated `cachedInputTokens` for plan
messages with multi-step tool calls (e.g., ~150k instead of ~50k),
causing the context bar to show >200k tokens which is impossible since
the user can still send messages.

## Root Cause

For multi-step tool calls, `cachedInputTokens` was being summed across
all steps:
- Step 1: reads 50k cached tokens
- Step 2: reads same 50k cached tokens  
- Step 3: reads same 50k cached tokens
- Cumulative: 150k (wrong for context display)

The frontend was falling back to cumulative `usage` when `contextUsage`
was undefined, which had the incorrectly summed values.

## Changes

**Backend (`streamManager.ts`)**:
- Refactor `getStreamMetadata()` to fetch both `totalUsage` (sum of all
steps, for costs) and `contextUsage` (last step only, for context
window) from AI SDK
- Add `contextProviderMetadata` from `streamResult.providerMetadata` for
accurate cache creation token display
- Falls back to tracked `lastStepUsage`/`lastStepProviderMetadata` if
streamResult times out

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

## Result

Context window now correctly shows last step's `inputTokens` (actual
context size) while cost calculation still uses cumulative totals.

## Migration Note

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

---
_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)