Skip to content

[usage] Surface provider-reported token usage instead of only estimating it - #175

Merged
yusiwen merged 2 commits into
masterfrom
feat/usage-reported-tokens
Oct 10, 2026
Merged

yusiwen merged 2 commits into
masterfrom
feat/usage-reported-tokens

Conversation

@yusiwen

@yusiwen yusiwen commented Oct 10, 2026

Copy link
Copy Markdown
Owner

Refs #174, roadmap entry A2 (tracker: #168).

What changed

  • types.Usage (PromptTokens, CompletionTokens, TotalTokens, Add) and
    ChatResponse.Usage *Usage — a pointer, so "the endpoint reported nothing" stays distinguishable
    from "it reported zero" — plus StreamCallbacks.OnUsage.
  • agent/provider_openai.go: usage is parsed from the SSE tail, which now reads to data: [DONE]
    instead of breaking at finish_reason — that break is what made the choices-less usage chunk
    unreachable — and from the batch body. A total_tokens the endpoint omits is derived from the two
    halves.
  • agent/provider_ollama.go: prompt_eval_count/eval_count from the stream's done line (read
    before 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: UsageTotal accumulates every reported call; recordUsage fires OnUsage once
    per reported call, including the max-steps summary call, which builds its own callbacks.
  • tui/: UsageMsg carries a report to the UI, and a report replaces the estimate charged to that
    call (callTokens) with the provider's number. sessionTokens stays the single displayed
    counter.
  • Tests: agent/usage_test.go (13 functions) and tui/usage_test.go (5).
  • Docs: the A2 entry in docs/roadmap.md carries its issue link, and the document now states the
    pickup convention; CODEBASE.md updated for the new fields, events and provider behaviour.

Why

The status bar's tokens: counter was len(text)/4 over streamed output, while the request already
asked 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_reason before the chunk was ever read off the socket. Reading past
finish_reason is 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:

    Mutation What the test said
    usage read after the choices check usage chunk after finish_reason was dropped
    break at the finish chunk usage chunk after finish_reason was dropped
    TUI adds the report without taking back the estimate tokens: 922, want the reported total 912
    no callTokens reset at the step boundary callTokens = 10, want 0 at the step boundary
    summary call's usage not recorded OnUsage fired 1 times, want 2 (step + summary)
    Ollama done-line counters dropped done-line counters were dropped
  • Every new fixture places the usage chunk after the finish_reason chunk, which is the order the
    bug hid behind.

Honest scope

  • No live provider call was made from this machine; parsing is pinned by fixtures that mirror the
    documented chunk order (OpenAI's include_usage tail, Ollama's done line).
  • Behaviour change: an SSE stream that sends finish_reason and then holds the connection open
    without [DONE] is now read until the provider's own 120 s timeout, where the old code returned at
    finish_reason. The SSE contract ends a stream with [DONE], and a stall before finish_reason
    already blocked for the same bound.
  • tokens: now includes prompt tokens, so the number jumps well above the old estimate on the first
    reported call. That is the point of the change rather than a regression, but it is visible.
  • Cost/pricing accounting and any cumulative budget or circuit breaker remain open — roadmap A3.
  • Usage is not persisted into the session file.

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.
@yusiwen yusiwen added the enhancement New feature or request label Oct 10, 2026
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant