🤖 fix: preserve usage data when stream is interrupted #837
+25
−3
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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.
github-merge-queue bot
pushed a commit
that referenced
this pull request
Dec 2, 2025
## Problem When a message is queued during streaming and then sent at `tool-call-end` (which interrupts the current stream), the usage data (breakdown by type and context window) resets to 0. ### Root Causes **1. AI SDK's `totalUsage` returns zeros on abort** In `cancelStreamSafely()`, we called `getStreamMetadata()` which awaits the AI SDK's `totalUsage` promise. When a stream is aborted mid-execution, this promise resolves with **zeros** (not `undefined`), so our fallback logic (`totalUsage ?? cumulativeUsage`) never triggered. **2. `usageStore` cache not invalidated on `stream-start`** The `MapStore` caches computed usage state. When a new stream starts after an abort: - `stream-abort` bumps `usageStore` ✓ - `stream-start` only bumped `states`, NOT `usageStore` ✗ - The stale cached value showed `liveUsage` as undefined ## Solution 1. **Use tracked `cumulativeUsage` directly** instead of AI SDK's unreliable `totalUsage` on abort. This is updated on each `finish-step` event (before tool execution), so it has accurate data even when interrupted mid-tool-call. 2. **Bump `usageStore` on `stream-start`** to invalidate the cache when a new stream begins. ### Trade-off Usage from the **interrupted step** is abandoned. The AI SDK's `finish-step` event fires *before* tool execution begins, so if we abort during tool execution, that step's usage was already recorded. However, if we abort during the model's response generation (before `finish-step`), that partial step's usage is lost. This is acceptable since: - We can't get reliable data from the SDK for interrupted steps - Cumulative usage from all *completed* steps is preserved - The alternative (zeros everywhere) is worse --- _Generated with `mux`_
ethanndickson
changed the title
5 tasks
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Problem
When a stream is interrupted (e.g., by a queued message, user cancellation, or starting a new message), the usage data (breakdown by type and context window) resets to 0.
Root Causes
1. AI SDK's
totalUsagereturns zeros on abortIn
cancelStreamSafely(), we calledgetStreamMetadata()which awaits the AI SDK'stotalUsagepromise. When a stream is aborted mid-execution, this promise resolves with zeros (notundefined), so our fallback logic (totalUsage ?? cumulativeUsage) never triggered.2.
usageStorecache not invalidated onstream-startThe
MapStorecaches computed usage state. When a new stream starts after an abort:stream-abortbumpsusageStore✓stream-startonly bumpedstates, NOTusageStore✗liveUsageas undefinedSolution
Use tracked
cumulativeUsagedirectly instead of AI SDK's unreliabletotalUsageon abort. This is updated on eachfinish-stepevent (before tool execution), so it has accurate data even when interrupted mid-tool-call.Bump
usageStoreonstream-startto invalidate the cache when a new stream begins.Trade-off
Usage from the interrupted step is abandoned. The AI SDK's
finish-stepevent fires before tool execution begins, so if we abort during tool execution, that step's usage was already recorded. However, if we abort during the model's response generation (beforefinish-step), that partial step's usage is lost. This is acceptable since:Generated with
mux