Skip to content

fix(agent): update session usage when a request fails, not just on success - #3395

Open
albatrossflyon-coder wants to merge 1 commit into
charmbracelet:mainfrom
albatrossflyon-coder:fix/context-window-usage-on-error
Open

fix(agent): update session usage when a request fails, not just on success#3395
albatrossflyon-coder wants to merge 1 commit into
charmbracelet:mainfrom
albatrossflyon-coder:fix/context-window-usage-on-error

Conversation

@albatrossflyon-coder

Copy link
Copy Markdown

Fixes #3384.

The bug

The sidebar's context-window percentage is driven by session.PromptTokens/session.CompletionTokens, which only get written in OnStepFinish (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 a StepResult, so OnStepFinish never 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 (after currentAssistant != nil, so PrepareStep already ran and stepMessages holds what was about to be sent), record what was actually attempted:

  1. Prefer the provider's own number. charm.land/fantasy's Anthropic/OpenAI/Google providers already parse a context-length-exceeded error body into ProviderError.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.
  2. Fall back to the existing local estimator (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.
  3. Always call 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 real sessionAgent.Run() through a fake fantasy.LanguageModel that 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 current main, green with the fix.

  • go test ./internal/agent/... ./internal/ui/... — 1134 passed (1 unrelated pre-existing failure on main too: TestGlobFilesScopedPrefixMatchesUnscoped, a Windows temp-path casing issue, confirmed unrelated by reproducing it on a clean checkout)
  • go build ./... — clean

…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Crush's context window calculations seem to be wrong

1 participant