🤖 fix: preserve usage data when stream is interrupted

bun.lock

@@ -1,5 +1,6 @@
 {
   "lockfileVersion": 1,
+  "configVersion": 0,
   "workspaces": {
     "": {
       "name": "@coder/cmux",

src/browser/stores/WorkspaceStore.ts

@@ -144,6 +144,8 @@ export class WorkspaceStore {
       // Don't reset retry state here - stream might still fail after starting
       // Retry state will be reset on stream-end (successful completion)
       this.states.bump(workspaceId);
+      // Bump usage store so liveUsage is recomputed with new activeStreamId
+      this.usageStore.bump(workspaceId);
     },
     "stream-delta": (workspaceId, aggregator, data) => {
       aggregator.handleStreamDelta(data as never);

src/common/types/stream.ts

@@ -59,7 +59,14 @@ export interface StreamAbortEvent {
   messageId: string;
   // Metadata may contain usage if abort occurred after stream completed processing
   metadata?: {
+    // Total usage across all steps (for cost calculation)
     usage?: LanguageModelV2Usage;
+    // Last step's usage (for context window display - inputTokens = current context size)
+    contextUsage?: LanguageModelV2Usage;
+    // Provider metadata for cost calculation (cache tokens, etc.)
+    providerMetadata?: Record<string, unknown>;
+    // Last step's provider metadata (for context window cache display)
+    contextProviderMetadata?: Record<string, unknown>;
     duration?: number;
   };
   abandonPartial?: boolean;

src/node/services/streamManager.ts

@@ -500,15 +500,27 @@ export class StreamManager extends EventEmitter {
       // while a new stream starts (e.g., old stream writing to partial.json)
       await streamInfo.processingPromise;
 
-      // Get usage and duration metadata (usage may be undefined if aborted early)
-      const { usage, duration } = await this.getStreamMetadata(streamInfo);
+      // For aborts, use our tracked cumulativeUsage directly instead of AI SDK's totalUsage.
+      // cumulativeUsage is updated on each finish-step event (before tool execution),
+      // so it has accurate data even when the stream is interrupted mid-tool-call.
+      // AI SDK's totalUsage may return zeros or stale data when aborted.
+      const duration = Date.now() - streamInfo.startTime;
+      const hasCumulativeUsage = (streamInfo.cumulativeUsage.totalTokens ?? 0) > 0;
+      const usage = hasCumulativeUsage ? streamInfo.cumulativeUsage : undefined;
+
+      // For context window display, use last step's usage (inputTokens = current context size)
+      const contextUsage = streamInfo.lastStepUsage;
+      const contextProviderMetadata = streamInfo.lastStepProviderMetadata;
+
+      // Include provider metadata for accurate cost calculation
+      const providerMetadata = streamInfo.cumulativeProviderMetadata;
 
       // Emit abort event with usage if available
       this.emit("stream-abort", {
         type: "stream-abort",
         workspaceId: workspaceId as string,
         messageId: streamInfo.messageId,
-        metadata: { usage, duration },
+        metadata: { usage, contextUsage, duration, providerMetadata, contextProviderMetadata },
         abandonPartial,
       });