fix(agent): update session usage when a request fails, not just on success - #3395
Open
albatrossflyon-coder wants to merge 1 commit into
Open
Conversation
…ccess Fixes charmbracelet#3384. session.PromptTokens/CompletionTokens (the numbers behind the sidebar's context-window percentage) were only ever updated in OnStepFinish, which fires on a successful step. A request rejected outright for exceeding the model's context window never reaches OnStepFinish, so the UI kept showing the last successful (smaller) usage instead of the size that was actually attempted and rejected. On such a failure, record what was actually sent: prefer the provider's own reported token count (charm.land/fantasy's Anthropic/OpenAI/Google providers already parse ContextUsedTokens/ContextMaxTokens out of the context-length error body — crush just never read them) and fall back to the existing local message-token estimator otherwise. Always pass estimated=true regardless of source, since a rejected request was never billed and must not add phantom cost or fire a usage-telemetry event. Added TestRun_ContextTooLargeErrorUpdatesSessionUsage, which drives a real sessionAgent.Run() through a fake model that succeeds once (small usage) then fails with a context-too-large ProviderError, and asserts the session reflects the larger, provider-reported figure afterward. Confirmed red against current main, green with the fix.
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
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.
Fixes #3384.
The bug
The sidebar's context-window percentage is driven by
session.PromptTokens/session.CompletionTokens, which only get written inOnStepFinish(internal/agent/agent.go) — i.e. only on a successful step. A request that OpenRouter (or any provider) rejects outright for exceeding the model's context window never produces aStepResult, soOnStepFinishnever fires, and the session keeps whatever usage the last successful step reported. That's exactly the report: crush showed 62% (165k/262k) — stale, from before the conversation grew — while OpenRouter's own error said the actual request exceeded 262k.The fix
In the existing error-handling branch of
Run(aftercurrentAssistant != nil, soPrepareStepalready ran andstepMessagesholds what was about to be sent), record what was actually attempted:charm.land/fantasy's Anthropic/OpenAI/Google providers already parse a context-length-exceeded error body intoProviderError.ContextUsedTokens/ContextMaxTokens(IsContextTooLarge()) — crush just never read them. This is the exact figure the provider rejected the request for, more precise than any local estimate.estimateMessageTokens, already used elsewhere in this file for the "successful step, zero usage" case) when the error isn't a parsed context-too-large error.updateSessionUsage(..., estimated=true)regardless of which source won — the request was never billed by the provider, so this must never add phantom cost or fire the usage-telemetry event that a real successful step does.Testing
Added
TestRun_ContextTooLargeErrorUpdatesSessionUsage: drives a realsessionAgent.Run()through a fakefantasy.LanguageModelthat succeeds once (small usage) and then fails with a*fantasy.ProviderError{ContextTooLargeErr: true, ContextUsedTokens: 270000, ...}, and asserts the session reflects the larger, provider-reported figure afterward (not the stale successful-step number). Confirmed red against currentmain, green with the fix.go test ./internal/agent/... ./internal/ui/...— 1134 passed (1 unrelated pre-existing failure onmaintoo:TestGlobFilesScopedPrefixMatchesUnscoped, a Windows temp-path casing issue, confirmed unrelated by reproducing it on a clean checkout)go build ./...— clean