Skip to content

Commit 7afe084

Browse files
committed
🤖 fix: preserve usage data when queued message interrupts stream
Two bugs caused usage/cost data to reset when a queued message was sent: 1. AI SDK's totalUsage returns zeros on abort - now use tracked cumulativeUsage 2. usageStore cache not invalidated on stream-start - now bump on new stream Usage from the interrupted step is abandoned since finish-step fires before tool execution, but cumulative usage from completed steps is preserved.
1 parent 372f0d9 commit 7afe084

File tree

4 files changed

+25
-3
lines changed

4 files changed

+25
-3
lines changed

bun.lock

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
{
22
"lockfileVersion": 1,
3+
"configVersion": 0,
34
"workspaces": {
45
"": {
56
"name": "@coder/cmux",

src/browser/stores/WorkspaceStore.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,8 @@ export class WorkspaceStore {
144144
// Don't reset retry state here - stream might still fail after starting
145145
// Retry state will be reset on stream-end (successful completion)
146146
this.states.bump(workspaceId);
147+
// Bump usage store so liveUsage is recomputed with new activeStreamId
148+
this.usageStore.bump(workspaceId);
147149
},
148150
"stream-delta": (workspaceId, aggregator, data) => {
149151
aggregator.handleStreamDelta(data as never);

src/common/types/stream.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,14 @@ export interface StreamAbortEvent {
5959
messageId: string;
6060
// Metadata may contain usage if abort occurred after stream completed processing
6161
metadata?: {
62+
// Total usage across all steps (for cost calculation)
6263
usage?: LanguageModelV2Usage;
64+
// Last step's usage (for context window display - inputTokens = current context size)
65+
contextUsage?: LanguageModelV2Usage;
66+
// Provider metadata for cost calculation (cache tokens, etc.)
67+
providerMetadata?: Record<string, unknown>;
68+
// Last step's provider metadata (for context window cache display)
69+
contextProviderMetadata?: Record<string, unknown>;
6370
duration?: number;
6471
};
6572
abandonPartial?: boolean;

src/node/services/streamManager.ts

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -500,15 +500,27 @@ export class StreamManager extends EventEmitter {
500500
// while a new stream starts (e.g., old stream writing to partial.json)
501501
await streamInfo.processingPromise;
502502

503-
// Get usage and duration metadata (usage may be undefined if aborted early)
504-
const { usage, duration } = await this.getStreamMetadata(streamInfo);
503+
// For aborts, use our tracked cumulativeUsage directly instead of AI SDK's totalUsage.
504+
// cumulativeUsage is updated on each finish-step event (before tool execution),
505+
// so it has accurate data even when the stream is interrupted mid-tool-call.
506+
// AI SDK's totalUsage may return zeros or stale data when aborted.
507+
const duration = Date.now() - streamInfo.startTime;
508+
const hasCumulativeUsage = (streamInfo.cumulativeUsage.totalTokens ?? 0) > 0;
509+
const usage = hasCumulativeUsage ? streamInfo.cumulativeUsage : undefined;
510+
511+
// For context window display, use last step's usage (inputTokens = current context size)
512+
const contextUsage = streamInfo.lastStepUsage;
513+
const contextProviderMetadata = streamInfo.lastStepProviderMetadata;
514+
515+
// Include provider metadata for accurate cost calculation
516+
const providerMetadata = streamInfo.cumulativeProviderMetadata;
505517

506518
// Emit abort event with usage if available
507519
this.emit("stream-abort", {
508520
type: "stream-abort",
509521
workspaceId: workspaceId as string,
510522
messageId: streamInfo.messageId,
511-
metadata: { usage, duration },
523+
metadata: { usage, contextUsage, duration, providerMetadata, contextProviderMetadata },
512524
abandonPartial,
513525
});
514526

0 commit comments

Comments
 (0)