🤖 feat: mid-stream token usage updates (#757)

Emit `usage-delta` events on AI SDK `finish-step`, allowing the UI to
display live token counts as the AI generates multi-step responses.



https://github.com/user-attachments/assets/80cd3a98-e7fa-415f-9ac6-7b7a241a7a5a

## Changes

**Backend:**
- Add `UsageDeltaEvent` type emitted on each `finish-step`
- `StreamManager` emits usage after each tool call step completes
- Forward events through `AIService` → `AgentSession` → IPC

**Frontend:**
- `StreamingMessageAggregator` stores active stream usage per messageId
- `WorkspaceStore` exposes `liveUsage` in `WorkspaceUsageState`
- `CostsTab`, `RightSidebar`, `ChatMetaSidebar` consume `liveUsage` for
real-time context window display

**Tests:**
- Unit tests for usage-delta handling in `StreamingMessageAggregator`
- Integration test for usage-delta events during multi-step streams

## Testing

This feature has been tested with OpenAI, Anthropic & Google.

---
_Generated with `mux`_

Author: ethanndickson

src/browser/components/ChatMetaSidebar.tsx

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

src/browser/components/RightSidebar.tsx

@@ -134,7 +134,7 @@ const RightSidebarComponent: React.FC<RightSidebarProps> = ({
   const costsPanelId = `${baseId}-panel-costs`;
   const reviewPanelId = `${baseId}-panel-review`;
 
-  const lastUsage = usage?.usageHistory[usage.usageHistory.length - 1];
+  const lastUsage = usage?.liveUsage ?? usage?.usageHistory[usage.usageHistory.length - 1];
 
   // Memoize vertical meter data calculation to prevent unnecessary re-renders
   const verticalMeterData = React.useMemo(() => {

src/browser/components/RightSidebar/CostsTab.tsx

@@ -63,8 +63,15 @@ const CostsTabComponent: React.FC<CostsTabProps> = ({ workspaceId }) => {
   const { options } = useProviderOptions();
   const use1M = options.anthropic?.use1MContext ?? false;
 
-  // Check if we have any data to display
-  const hasUsageData = usage && usage.usageHistory.length > 0;
+  // Session usage for cost
+  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);
   const hasConsumerData = consumers && (consumers.totalTokens > 0 || consumers.isCalculating);
   const hasAnyData = hasUsageData || hasConsumerData;
 
@@ -80,28 +87,22 @@ const CostsTabComponent: React.FC<CostsTabProps> = ({ workspaceId }) => {
     );
   }
 
-  // Context Usage always shows Last Request data
-  const lastRequestUsage = hasUsageData
-    ? usage.usageHistory[usage.usageHistory.length - 1]
-    : undefined;
+  // Last Request (for Cost section): always the last completed request
+  const lastRequestUsage = usage.usageHistory[usage.usageHistory.length - 1];
 
   // Cost and Details table use viewMode
-  const displayUsage =
-    viewMode === "last-request"
-      ? usage.usageHistory[usage.usageHistory.length - 1]
-      : sumUsageHistory(usage.usageHistory);
+  const displayUsage = viewMode === "last-request" ? lastRequestUsage : sessionUsage;
 
   return (
     <div className="text-light font-primary text-[13px] leading-relaxed">
       {hasUsageData && (
         <div data-testid="context-usage-section" className="mt-2 mb-5">
           <div data-testid="context-usage-list" className="flex flex-col gap-3">
             {(() => {
-              // Context Usage always uses last request
-              const contextUsage = lastRequestUsage;
-
-              // Get model from last request (for context window display)
-              const model = lastRequestUsage?.model ?? "unknown";
+              // Context usage: live when streaming, else last historical
+              const contextUsage =
+                usage.liveUsage ?? usage.usageHistory[usage.usageHistory.length - 1];
+              const model = contextUsage?.model ?? "unknown";
 
               // Get max tokens for the model from the model stats database
               const modelStats = getModelStats(model);

src/browser/stores/WorkspaceStore.ts

@@ -17,7 +17,7 @@ import {
   isRestoreToInput,
 } from "@/common/types/ipc";
 import { MapStore } from "./MapStore";
-import { collectUsageHistory } from "@/common/utils/tokens/displayUsage";
+import { collectUsageHistory, createDisplayUsage } from "@/common/utils/tokens/displayUsage";
 import { WorkspaceConsumerManager } from "./WorkspaceConsumerManager";
 import type { ChatUsageDisplay } from "@/common/utils/tokens/usageAggregator";
 import type { TokenConsumer } from "@/common/types/chatStats";
@@ -63,6 +63,8 @@ type DerivedState = Record<string, number>;
 export interface WorkspaceUsageState {
   usageHistory: ChatUsageDisplay[];
   totalTokens: number;
+  /** Live usage during streaming (inputTokens = current context window) */
+  liveUsage?: ChatUsageDisplay;
 }
 
 /**
@@ -178,6 +180,10 @@ export class WorkspaceStore {
       aggregator.handleReasoningEnd(data as never);
       this.states.bump(workspaceId);
     },
+    "usage-delta": (workspaceId, aggregator, data) => {
+      aggregator.handleUsageDelta(data as never);
+      this.usageStore.bump(workspaceId);
+    },
     "init-start": (workspaceId, aggregator, data) => {
       aggregator.handleMessage(data);
       this.states.bump(workspaceId);
@@ -449,7 +455,12 @@ export class WorkspaceStore {
         0
       );
 
-      return { usageHistory, totalTokens };
+      // Include active stream usage if currently streaming (already converted)
+      const activeStreamId = aggregator.getActiveStreamMessageId();
+      const rawUsage = activeStreamId ? aggregator.getActiveStreamUsage(activeStreamId) : undefined;
+      const liveUsage = rawUsage && model ? createDisplayUsage(rawUsage, model) : undefined;
+
+      return { usageHistory, totalTokens, liveUsage };
     });
   }
 

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

@@ -379,4 +379,96 @@ describe("StreamingMessageAggregator", () => {
       expect(aggregator.getCurrentTodos()).toHaveLength(0);
     });
   });
+
+  describe("usage-delta handling", () => {
+    test("handleUsageDelta stores usage by messageId", () => {
+      const aggregator = new StreamingMessageAggregator(TEST_CREATED_AT);
+
+      aggregator.handleUsageDelta({
+        type: "usage-delta",
+        workspaceId: "ws-1",
+        messageId: "msg-1",
+        usage: { inputTokens: 1000, outputTokens: 50, totalTokens: 1050 },
+      });
+
+      expect(aggregator.getActiveStreamUsage("msg-1")).toEqual({
+        inputTokens: 1000,
+        outputTokens: 50,
+        totalTokens: 1050,
+      });
+    });
+
+    test("clearTokenState removes usage", () => {
+      const aggregator = new StreamingMessageAggregator(TEST_CREATED_AT);
+
+      aggregator.handleUsageDelta({
+        type: "usage-delta",
+        workspaceId: "ws-1",
+        messageId: "msg-1",
+        usage: { inputTokens: 1000, outputTokens: 50, totalTokens: 1050 },
+      });
+
+      expect(aggregator.getActiveStreamUsage("msg-1")).toBeDefined();
+
+      aggregator.clearTokenState("msg-1");
+
+      expect(aggregator.getActiveStreamUsage("msg-1")).toBeUndefined();
+    });
+
+    test("latest usage-delta replaces previous for same messageId", () => {
+      const aggregator = new StreamingMessageAggregator(TEST_CREATED_AT);
+
+      // First step usage
+      aggregator.handleUsageDelta({
+        type: "usage-delta",
+        workspaceId: "ws-1",
+        messageId: "msg-1",
+        usage: { inputTokens: 1000, outputTokens: 50, totalTokens: 1050 },
+      });
+
+      // Second step usage (larger context after tool result added)
+      aggregator.handleUsageDelta({
+        type: "usage-delta",
+        workspaceId: "ws-1",
+        messageId: "msg-1",
+        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,
+      });
+    });
+
+    test("tracks usage independently per messageId", () => {
+      const aggregator = new StreamingMessageAggregator(TEST_CREATED_AT);
+
+      aggregator.handleUsageDelta({
+        type: "usage-delta",
+        workspaceId: "ws-1",
+        messageId: "msg-1",
+        usage: { inputTokens: 1000, outputTokens: 50, totalTokens: 1050 },
+      });
+
+      aggregator.handleUsageDelta({
+        type: "usage-delta",
+        workspaceId: "ws-1",
+        messageId: "msg-2",
+        usage: { inputTokens: 2000, outputTokens: 100, totalTokens: 2100 },
+      });
+
+      expect(aggregator.getActiveStreamUsage("msg-1")).toEqual({
+        inputTokens: 1000,
+        outputTokens: 50,
+        totalTokens: 1050,
+      });
+      expect(aggregator.getActiveStreamUsage("msg-2")).toEqual({
+        inputTokens: 2000,
+        outputTokens: 100,
+        totalTokens: 2100,
+      });
+    });
+  });
 });

src/browser/utils/messages/StreamingMessageAggregator.ts

@@ -9,6 +9,7 @@ import { createMuxMessage } from "@/common/types/message";
 import type {
   StreamStartEvent,
   StreamDeltaEvent,
+  UsageDeltaEvent,
   StreamEndEvent,
   StreamAbortEvent,
   ToolCallStartEvent,
@@ -17,6 +18,7 @@ import type {
   ReasoningDeltaEvent,
   ReasoningEndEvent,
 } from "@/common/types/stream";
+import type { LanguageModelV2Usage } from "@ai-sdk/provider";
 import type { TodoItem, StatusSetToolResult } from "@/common/types/tools";
 
 import type { WorkspaceChatMessage, StreamErrorMessage, DeleteMessage } from "@/common/types/ipc";
@@ -72,6 +74,10 @@ export class StreamingMessageAggregator {
   // Delta history for token counting and TPS calculation
   private deltaHistory = new Map<string, DeltaRecordStorage>();
 
+  // Active stream step usage (updated on each stream-step event)
+  // Tracks cumulative usage for the currently streaming message
+  private activeStreamStepUsage = new Map<string, LanguageModelV2Usage>();
+
   // Current TODO list (updated when todo_write succeeds, cleared on stream end)
   // Stream-scoped: automatically reset when stream completes
   // On reload: only reconstructed if reconnecting to active stream
@@ -992,5 +998,20 @@ export class StreamingMessageAggregator {
    */
   clearTokenState(messageId: string): void {
     this.deltaHistory.delete(messageId);
+    this.activeStreamStepUsage.delete(messageId);
+  }
+
+  /**
+   * Handle usage-delta event: update cumulative usage for active stream
+   */
+  handleUsageDelta(data: UsageDeltaEvent): void {
+    this.activeStreamStepUsage.set(data.messageId, data.usage);
+  }
+
+  /**
+   * Get active stream usage for a message (if streaming)
+   */
+  getActiveStreamUsage(messageId: string): LanguageModelV2Usage | undefined {
+    return this.activeStreamStepUsage.get(messageId);
   }
 }

src/common/types/ipc.ts

@@ -21,6 +21,7 @@ import type {
   StreamDeltaEvent,
   StreamEndEvent,
   StreamAbortEvent,
+  UsageDeltaEvent,
   ToolCallStartEvent,
   ToolCallDeltaEvent,
   ToolCallEndEvent,
@@ -103,6 +104,7 @@ export type WorkspaceChatMessage =
   | DeleteMessage
   | StreamStartEvent
   | StreamDeltaEvent
+  | UsageDeltaEvent
   | StreamEndEvent
   | StreamAbortEvent
   | ToolCallStartEvent
@@ -149,6 +151,11 @@ export function isStreamAbort(msg: WorkspaceChatMessage): msg is StreamAbortEven
   return "type" in msg && msg.type === "stream-abort";
 }
 
+// Type guard for usage delta events
+export function isUsageDelta(msg: WorkspaceChatMessage): msg is UsageDeltaEvent {
+  return "type" in msg && msg.type === "usage-delta";
+}
+
 // Type guard for tool call start events
 export function isToolCallStart(msg: WorkspaceChatMessage): msg is ToolCallStartEvent {
   return "type" in msg && msg.type === "tool-call-start";

src/common/types/stream.ts

@@ -121,6 +121,17 @@ export interface ReasoningEndEvent {
   messageId: string;
 }
 
+/**
+ * Emitted on each AI SDK finish-step event, providing incremental usage updates.
+ * Allows UI to update token display as steps complete (after each tool call or at stream end).
+ */
+export interface UsageDeltaEvent {
+  type: "usage-delta";
+  workspaceId: string;
+  messageId: string;
+  usage: LanguageModelV2Usage; // This step's usage (inputTokens = full context)
+}
+
 export type AIServiceEvent =
   | StreamStartEvent
   | StreamDeltaEvent
@@ -132,4 +143,5 @@ export type AIServiceEvent =
   | ToolCallEndEvent
   | ReasoningStartEvent
   | ReasoningDeltaEvent
-  | ReasoningEndEvent;
+  | ReasoningEndEvent
+  | UsageDeltaEvent;

src/node/services/agentSession.ts

@@ -464,6 +464,7 @@ export class AgentSession {
     });
     forward("reasoning-delta", (payload) => this.emitChatEvent(payload));
     forward("reasoning-end", (payload) => this.emitChatEvent(payload));
+    forward("usage-delta", (payload) => this.emitChatEvent(payload));
 
     forward("stream-end", async (payload) => {
       const handled = await this.compactionHandler.handleCompletion(payload as StreamEndEvent);

src/node/services/aiService.ts

@@ -203,6 +203,7 @@ export class AIService extends EventEmitter {
     // Forward reasoning events
     this.streamManager.on("reasoning-delta", (data) => this.emit("reasoning-delta", data));
     this.streamManager.on("reasoning-end", (data) => this.emit("reasoning-end", data));
+    this.streamManager.on("usage-delta", (data) => this.emit("usage-delta", data));
   }
 
   private async ensureSessionsDir(): Promise<void> {