fix(ai): persist chat token usage on early stream close - #430
Merged
Conversation
…m at message_stop The agent loop (stream_generator) stops consuming the provider stream as soon as message_stop arrives, so event_stream_with_context_retry never runs to exhaustion and its post-loop tracker.save() was dead code. Every chat turn's captured usage was discarded: model_usage has no 'chat' purpose rows even though title_generation and compaction records land fine. Move save() into a finally block so it fires on generator close (early break, error, cancellation) as well as normal completion. Tests: - Extend the happy-path chat stream integration test to assert the model_usage row (purpose=chat, correct token counts) is persisted; red-green verified against the unfixed code. - Harden seeded_chat teardown against the fire-and-forget persist task racing the chat delete (FK from model_usage to chats). - test_providers_api (live LLM): assert usage actually flows through streamed message_delta events instead of only reporting it.
prvnsmpth
changed the title
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
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.
tracker.save()into afinallyblock in the chat streaming generator: the agent loop stops consuming the stream atmessage_stop, so the post-loop save never ran and every chat turn's usage was silently dropped frommodel_usage.model_usagerow lands withpurpose=chatand the expected token counts; verified red (unfixed) / green (fixed).message_deltaevents instead of only reporting it, so usage-delivery regressions fail CI.