diff --git a/bun.lock b/bun.lock index 899b1121ae..d75d88c206 100644 --- a/bun.lock +++ b/bun.lock @@ -1,5 +1,6 @@ { "lockfileVersion": 1, + "configVersion": 0, "workspaces": { "": { "name": "@coder/cmux", diff --git a/src/browser/stores/WorkspaceStore.ts b/src/browser/stores/WorkspaceStore.ts index 295d0acd49..c7b2c86c8b 100644 --- a/src/browser/stores/WorkspaceStore.ts +++ b/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); diff --git a/src/common/types/stream.ts b/src/common/types/stream.ts index 306abfc7f8..6407ced589 100644 --- a/src/common/types/stream.ts +++ b/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; + // Last step's provider metadata (for context window cache display) + contextProviderMetadata?: Record; duration?: number; }; abandonPartial?: boolean; diff --git a/src/node/services/streamManager.ts b/src/node/services/streamManager.ts index c6912c1c51..cae73395d0 100644 --- a/src/node/services/streamManager.ts +++ b/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, });