🤖 fix: accurate cost estimation for multi-step tool usage (#831)

## Problem

The application was severely underestimating costs for conversations
involving tool calls.

### Root Cause

The Vercel AI SDK provides two usage metrics:
- `streamResult.usage` — Token usage from **the last step only**
- `streamResult.totalUsage` — Sum of token usage across **all steps**

The application was using `usage` instead of `totalUsage`. For a
conversation with 10 tool calls, only ~1/10th of actual consumption was
reported. A $5 conversation would display as $0.50.

### The Complicating Factor

Two UI elements use token data with different semantic requirements:

| Display | Needs | Why |
|---------|-------|-----|
| **Cost** | Sum of all steps | If model read context 10 times, you paid
for 10 reads |
| **Context window** | Last step only | Shows "how full is the
conversation now" for the next request |

Simply switching to `totalUsage` would fix costs but break context
display (showing 500% utilization after many tool calls).

### Cache Creation Tokens

Anthropic's cache creation tokens (`cacheCreationInputTokens`) are:
- Only in provider-specific metadata, not normalized usage
- Need to be summed across all steps
- Not automatically aggregated by the AI SDK

Even with `totalUsage`, cache creation costs were lost unless manually
aggregated from each step's provider metadata.

## Solution

Track both values with different semantic purposes:

**For cost calculation:**
- `usage` / `cumulativeUsage` — total across all steps
- `providerMetadata` / `cumulativeProviderMetadata` — aggregated cache
creation tokens

**For context window display:**
- `contextUsage` / `lastContextUsage` — last step only
- `contextProviderMetadata` — last step only

### Key Changes

1. **Backend** (`streamManager.ts`): Use `totalUsage` for cost, track
`lastStepUsage` for context, aggregate provider metadata across steps
2. **Types**: Extended `StreamEndEvent`, `MuxMetadata`,
`UsageDeltaEvent` with dual fields
3. **Frontend**: `StreamingMessageAggregator` tracks both cumulative and
per-step usage
4. **Store**: `WorkspaceUsageState` provides `usageHistory` (cost) and
`lastContextUsage` (context window)
5. **UI**: Components use appropriate field for their purpose

### Also Fixed

- **OpenAI cached token double-counting**: Gateway models
(`mux-gateway:openai/gpt-5.1`) weren't recognized as OpenAI, causing
cached tokens to be counted in both "Cache Read" and "Input". Now
normalizes gateway model strings before provider detection.

- **Google/Gemini cached token double-counting**: Google, like OpenAI,
reports `inputTokens` inclusive of `cachedInputTokens`. Extended the
subtraction logic to handle Google models.

---
_Generated with `mux`_

Author: ethanndickson

src/browser/components/ChatMetaSidebar.tsx

@@ -19,7 +19,8 @@ const ChatMetaSidebarComponent: React.FC<ChatMetaSidebarProps> = ({ workspaceId,
   const use1M = options.anthropic?.use1MContext ?? false;
   const chatAreaSize = useResizeObserver(chatAreaRef);
 
-  const lastUsage = usage?.liveUsage ?? usage?.usageHistory[usage.usageHistory.length - 1];
+  // Use lastContextUsage for context window display (last step = actual context size)
+  const lastUsage = usage?.liveUsage ?? usage?.lastContextUsage;
 
   // Memoize vertical meter data calculation to prevent unnecessary re-renders
   const verticalMeterData = React.useMemo(() => {

src/browser/components/RightSidebar.tsx

@@ -135,7 +135,8 @@ const RightSidebarComponent: React.FC<RightSidebarProps> = ({
   const costsPanelId = `${baseId}-panel-costs`;
   const reviewPanelId = `${baseId}-panel-review`;
 
-  const lastUsage = usage?.liveUsage ?? usage?.usageHistory[usage.usageHistory.length - 1];
+  // Use lastContextUsage for context window display (last step = actual context size)
+  const lastUsage = usage?.liveUsage ?? usage?.lastContextUsage;
   const model = lastUsage?.model ?? null;
 
   // Auto-compaction settings: threshold per-model

src/browser/components/RightSidebar/CostsTab.tsx

@@ -65,22 +65,28 @@ const CostsTabComponent: React.FC<CostsTabProps> = ({ workspaceId }) => {
   const use1M = options.anthropic?.use1MContext ?? false;
 
   // Get model from context usage for per-model threshold storage
-  const contextUsage = usage.liveUsage ?? usage.usageHistory[usage.usageHistory.length - 1];
-  const currentModel = contextUsage?.model ?? null;
+  // Use lastContextUsage for context window display (last step's usage)
+  const contextUsageForModel = usage.liveUsage ?? usage.lastContextUsage;
+  const currentModel = contextUsageForModel?.model ?? null;
 
   // Auto-compaction settings: threshold per-model (100 = disabled)
   const { threshold: autoCompactThreshold, setThreshold: setAutoCompactThreshold } =
     useAutoCompactionSettings(workspaceId, currentModel);
 
-  // Session usage for cost
+  // Session usage for cost calculation
+  // Uses usageHistory (total across all steps) + liveCostUsage (cumulative during streaming)
   const sessionUsage = React.useMemo(() => {
     const historicalSum = sumUsageHistory(usage.usageHistory);
-    if (!usage.liveUsage) return historicalSum;
-    if (!historicalSum) return usage.liveUsage;
-    return sumUsageHistory([historicalSum, usage.liveUsage]);
-  }, [usage.usageHistory, usage.liveUsage]);
-
-  const hasUsageData = usage && (usage.usageHistory.length > 0 || usage.liveUsage !== undefined);
+    if (!usage.liveCostUsage) return historicalSum;
+    if (!historicalSum) return usage.liveCostUsage;
+    return sumUsageHistory([historicalSum, usage.liveCostUsage]);
+  }, [usage.usageHistory, usage.liveCostUsage]);
+
+  const hasUsageData =
+    usage &&
+    (usage.usageHistory.length > 0 ||
+      usage.lastContextUsage !== undefined ||
+      usage.liveUsage !== undefined);
   const hasConsumerData = consumers && (consumers.totalTokens > 0 || consumers.isCalculating);
   const hasAnyData = hasUsageData || hasConsumerData;
 
@@ -109,8 +115,8 @@ const CostsTabComponent: React.FC<CostsTabProps> = ({ workspaceId }) => {
           <div data-testid="context-usage-list" className="flex flex-col gap-3">
             {(() => {
               // Context usage: live when streaming, else last historical
-              const contextUsage =
-                usage.liveUsage ?? usage.usageHistory[usage.usageHistory.length - 1];
+              // Uses lastContextUsage (last step) for accurate context window size
+              const contextUsage = usage.liveUsage ?? usage.lastContextUsage;
               const model = contextUsage?.model ?? "unknown";
 
               // Get max tokens for the model from the model stats database

src/browser/stores/WorkspaceStore.ts

@@ -59,12 +59,21 @@ type DerivedState = Record<string, number>;
 /**
  * Usage metadata extracted from API responses (no tokenization).
  * Updates instantly when usage metadata arrives.
+ *
+ * For multi-step tool calls, cost and context usage differ:
+ * - usageHistory: Total usage per message (sum of all steps) for cost calculation
+ * - lastContextUsage: Last step's usage for context window display (inputTokens = actual context size)
  */
 export interface WorkspaceUsageState {
+  /** Usage history for cost calculation (total across all steps per message) */
   usageHistory: ChatUsageDisplay[];
+  /** Last message's context usage (last step only, for context window display) */
+  lastContextUsage?: ChatUsageDisplay;
   totalTokens: number;
-  /** Live usage during streaming (inputTokens = current context window) */
+  /** Live context usage during streaming (last step's inputTokens = current context window) */
   liveUsage?: ChatUsageDisplay;
+  /** Live cost usage during streaming (cumulative across all steps) */
+  liveCostUsage?: ChatUsageDisplay;
 }
 
 /**
@@ -441,6 +450,8 @@ export class WorkspaceStore {
 
       const messages = aggregator.getAllMessages();
       const model = aggregator.getCurrentModel();
+
+      // Collect usage history for cost calculation (total across all steps per message)
       const usageHistory = collectUsageHistory(messages, model);
 
       // Calculate total from usage history (now includes historical)
@@ -455,12 +466,47 @@ export class WorkspaceStore {
         0
       );
 
-      // Include active stream usage if currently streaming (already converted)
+      // Get last message's context usage for context window display
+      // Uses contextUsage (last step) if available, falls back to usage for old messages
+      const lastContextUsage = (() => {
+        for (let i = messages.length - 1; i >= 0; i--) {
+          const msg = messages[i];
+          if (msg.role === "assistant") {
+            const rawUsage = msg.metadata?.contextUsage ?? msg.metadata?.usage;
+            const providerMeta =
+              msg.metadata?.contextProviderMetadata ?? msg.metadata?.providerMetadata;
+            if (rawUsage) {
+              const msgModel = msg.metadata?.model ?? model ?? "unknown";
+              return createDisplayUsage(rawUsage, msgModel, providerMeta);
+            }
+          }
+        }
+        return undefined;
+      })();
+
+      // Include active stream usage if currently streaming
       const activeStreamId = aggregator.getActiveStreamMessageId();
-      const rawUsage = activeStreamId ? aggregator.getActiveStreamUsage(activeStreamId) : undefined;
-      const liveUsage = rawUsage && model ? createDisplayUsage(rawUsage, model) : undefined;
 
-      return { usageHistory, totalTokens, liveUsage };
+      // Live context usage (last step's inputTokens = current context window)
+      const rawContextUsage = activeStreamId
+        ? aggregator.getActiveStreamUsage(activeStreamId)
+        : undefined;
+      const liveUsage =
+        rawContextUsage && model ? createDisplayUsage(rawContextUsage, model) : undefined;
+
+      // Live cost usage (cumulative across all steps, with accumulated cache creation tokens)
+      const rawCumulativeUsage = activeStreamId
+        ? aggregator.getActiveStreamCumulativeUsage(activeStreamId)
+        : undefined;
+      const rawCumulativeProviderMetadata = activeStreamId
+        ? aggregator.getActiveStreamCumulativeProviderMetadata(activeStreamId)
+        : undefined;
+      const liveCostUsage =
+        rawCumulativeUsage && model
+          ? createDisplayUsage(rawCumulativeUsage, model, rawCumulativeProviderMetadata)
+          : undefined;
+
+      return { usageHistory, lastContextUsage, totalTokens, liveUsage, liveCostUsage };
     });
   }
 

src/browser/utils/compaction/autoCompactionCheck.test.ts

@@ -40,9 +40,11 @@ const createMockUsage = (
   }
 
   // Add recent usage
-  usageHistory.push(createUsageEntry(lastEntryTokens, model));
+  const recentUsage = createUsageEntry(lastEntryTokens, model);
+  usageHistory.push(recentUsage);
 
-  return { usageHistory, totalTokens: 0, liveUsage };
+  // lastContextUsage is the most recent context window state
+  return { usageHistory, lastContextUsage: recentUsage, totalTokens: 0, liveUsage };
 };
 
 describe("checkAutoCompaction", () => {
@@ -136,17 +138,17 @@ describe("checkAutoCompaction", () => {
 
     test("includes all token types in calculation", () => {
       // Create usage with all token types specified
+      const usageEntry = {
+        input: { tokens: 10_000 },
+        cached: { tokens: 5_000 },
+        cacheCreate: { tokens: 2_000 },
+        output: { tokens: 3_000 },
+        reasoning: { tokens: 1_000 },
+        model: KNOWN_MODELS.SONNET.id,
+      };
       const usage: WorkspaceUsageState = {
-        usageHistory: [
-          {
-            input: { tokens: 10_000 },
-            cached: { tokens: 5_000 },
-            cacheCreate: { tokens: 2_000 },
-            output: { tokens: 3_000 },
-            reasoning: { tokens: 1_000 },
-            model: KNOWN_MODELS.SONNET.id,
-          },
-        ],
+        usageHistory: [usageEntry],
+        lastContextUsage: usageEntry,
         totalTokens: 0,
       };
 
@@ -232,17 +234,17 @@ describe("checkAutoCompaction", () => {
     });
 
     test("handles zero tokens gracefully", () => {
+      const zeroEntry = {
+        input: { tokens: 0 },
+        cached: { tokens: 0 },
+        cacheCreate: { tokens: 0 },
+        output: { tokens: 0 },
+        reasoning: { tokens: 0 },
+        model: KNOWN_MODELS.SONNET.id,
+      };
       const usage: WorkspaceUsageState = {
-        usageHistory: [
-          {
-            input: { tokens: 0 },
-            cached: { tokens: 0 },
-            cacheCreate: { tokens: 0 },
-            output: { tokens: 0 },
-            reasoning: { tokens: 0 },
-            model: KNOWN_MODELS.SONNET.id,
-          },
-        ],
+        usageHistory: [zeroEntry],
+        lastContextUsage: zeroEntry,
         totalTokens: 0,
       };
 
@@ -357,7 +359,11 @@ describe("checkAutoCompaction", () => {
     test("shouldForceCompact triggers with empty history but liveUsage near limit", () => {
       // Bug fix: empty history but liveUsage should still trigger
       const liveUsage = createUsageEntry(SONNET_MAX_TOKENS - BUFFER);
-      const usage: WorkspaceUsageState = { usageHistory: [], totalTokens: 0, liveUsage };
+      const usage: WorkspaceUsageState = {
+        usageHistory: [],
+        totalTokens: 0,
+        liveUsage,
+      };
       const result = checkAutoCompaction(usage, KNOWN_MODELS.SONNET.id, false);
 
       expect(result.shouldForceCompact).toBe(true);

src/browser/utils/compaction/autoCompactionCheck.ts

@@ -94,8 +94,9 @@ export function checkAutoCompaction(
     };
   }
 
-  // Current usage: live when streaming, else last historical (pattern from CostsTab)
-  const lastUsage = usage.usageHistory[usage.usageHistory.length - 1];
+  // Current usage: live when streaming, else last historical
+  // Use lastContextUsage (last step) for accurate context window size
+  const lastUsage = usage.lastContextUsage;
   const currentUsage = usage.liveUsage ?? lastUsage;
 
   // Force-compact when approaching context limit (can trigger even with empty history if streaming)