Repository navigation
[usage] Surface provider-reported token usage instead of only estimating it - #175
Merged
Merged
Conversation
The status bar's `tokens:` counter was an estimate: bytes/4 over the streamed
output text. The request already asked providers for the real figure
(`stream_options: {include_usage: true}`), but every answer was discarded — the
SSE loop skipped the choices-less usage chunk, and it broke on
`finish_reason`, so the chunk carrying usage was never even read off the
socket. Ollama's `prompt_eval_count`/`eval_count` were dropped the same way.
- `types.Usage` and `ChatResponse.Usage` — a pointer, because "the endpoint
reported no usage" must stay distinguishable from "it reported zero" — plus
`StreamCallbacks.OnUsage`.
- OpenAI-compatible provider: parse usage from the streaming tail, which now
runs to `[DONE]` instead of stopping at `finish_reason`, and from the batch
body; derive `total_tokens` when an endpoint omits it.
- Ollama: read `prompt_eval_count`/`eval_count` from the stream's done line
before the loop breaks, and from the batch body's top level.
- Agent: `UsageTotal` accumulates every reported call and `OnUsage` fires once
per reported call, including the max-steps summary call, which builds its own
callbacks.
- TUI: a reported usage replaces the estimate charged to that call
(`callTokens`); an endpoint that reports nothing sends no event, so its
estimate stays in the counter. One counter, not two.
Reading past `finish_reason` is required for the usage chunk. The cost is that a
server which sends `finish_reason` and then holds the connection open without
`[DONE]` is now read until the provider's own 120 s timeout instead of
returning at `finish_reason`; the SSE contract ends the stream with `[DONE]`.
Roadmap entry A2 (`docs/roadmap.md`), picked up as issue #174.
Tests: `go test ./... -count=1` and `-race` green (13 packages); staticcheck
v0.8.1 and gofmt clean. The new fixtures place the usage chunk after the
`finish_reason` chunk, so a parser that stops early fails them.
The `tokens:` counter changed meaning: it is the provider level now, not len(text)/4, and it includes the prompt. A user reading the status bar should know why the number jumped, and the changelog is where this repository records that. Roadmap entry A2, issue #174.
This was referenced Oct 10, 2026
[agent] Bound the SSE read after finish_reason so a stream without [DONE] cannot stall the call
#176
Closed
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.
Refs #174, roadmap entry A2 (tracker: #168).
What changed
types.Usage(PromptTokens,CompletionTokens,TotalTokens,Add) andChatResponse.Usage *Usage— a pointer, so "the endpoint reported nothing" stays distinguishablefrom "it reported zero" — plus
StreamCallbacks.OnUsage.agent/provider_openai.go:usageis parsed from the SSE tail, which now reads todata: [DONE]instead of breaking at
finish_reason— that break is what made the choices-less usage chunkunreachable — and from the batch body. A
total_tokensthe endpoint omits is derived from the twohalves.
agent/provider_ollama.go:prompt_eval_count/eval_countfrom the stream's done line (readbefore the loop breaks) and from the batch body's top level. Counters that are absent stay
nil;counters that are present and zero are a reported zero.
agent/agent.go:UsageTotalaccumulates every reported call;recordUsagefiresOnUsageonceper reported call, including the max-steps summary call, which builds its own callbacks.
tui/:UsageMsgcarries a report to the UI, and a report replaces the estimate charged to thatcall (
callTokens) with the provider's number.sessionTokensstays the single displayedcounter.
agent/usage_test.go(13 functions) andtui/usage_test.go(5).docs/roadmap.mdcarries its issue link, and the document now states thepickup convention;
CODEBASE.mdupdated for the new fields, events and provider behaviour.Why
The status bar's
tokens:counter waslen(text)/4over streamed output, while the request alreadyasked every OpenAI-compatible endpoint for real usage (
stream_options: {include_usage: true}).Two separate spots threw that answer away: the choices check skipped the usage-only chunk, and the
loop broke at
finish_reasonbefore the chunk was ever read off the socket. Reading pastfinish_reasonis the only way to reach it.Rejected alternative: a second counter for real usage. The roadmap asks to extend the existing one,
and two numbers on screen would disagree. The estimate is therefore kept as the fallback for an
endpoint that reports nothing, and a report replaces exactly the share of the counter that call
estimated.
Cost and pricing are deliberately not here: no price table exists in the tree, and the roadmap keeps
that with A3's cumulative-budget work.
Verification
go test ./... -count=1→ 13 packages ok, 0 failed.go test ./... -count=1 -race→ same.make staticcheck(v0.8.1) → exit 0.gofmt -l .→ no output.The new gate was mutation-checked six ways; every mutation failed on an assertion, not a build
error:
usage chunk after finish_reason was droppedbreakat the finish chunkusage chunk after finish_reason was droppedtokens: 922, want the reported total 912callTokensreset at the step boundarycallTokens = 10, want 0 at the step boundaryOnUsage fired 1 times, want 2 (step + summary)done-line counters were droppedEvery new fixture places the usage chunk after the
finish_reasonchunk, which is the order thebug hid behind.
Honest scope
documented chunk order (OpenAI's
include_usagetail, Ollama's done line).finish_reasonand then holds the connection openwithout
[DONE]is now read until the provider's own 120 s timeout, where the old code returned atfinish_reason. The SSE contract ends a stream with[DONE], and a stall beforefinish_reasonalready blocked for the same bound.
tokens:now includes prompt tokens, so the number jumps well above the old estimate on the firstreported call. That is the point of the change rather than a regression, but it is visible.