feat: display token usage after LLM responses and in spinner#157
Merged
Conversation
Show token usage numbers (input/output/cached tokens) in the output: - Display accumulated session usage as a dimmed line after the final LLM response (when no more tool calls are pending) - Show session token totals in the spinner while waiting for LLM - Use emoji format (📥 📤 💾) for compact display - Only show non-zero values; omit caching info if unused - Extract cached token counts from providers (OpenAI, Claude, Gemini/VertexAI, Bedrock) - For OpenAI streaming, only capture usage from the final usage-only chunk to avoid incorrect intermediate values Closes #152
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds token-usage collection and display: new CompletionTokenUsage type, propagation of input/output/cached token counts from providers through streaming and non-streaming paths, session accumulation, spinner label integration, and printing of usage after final LLM responses. Changes
Sequence DiagramsequenceDiagram
participant Provider as LLM Provider
participant Handler as SseHandler
participant Client as call_chat_completions*
participant Main as Main/REPL
participant Config as Config
participant Session as Session
participant Render as Render Pipeline
Provider->>Handler: stream chunks (including usage metadata)
Handler->>Handler: set_usage(input, output, cached)
Handler->>Client: take() -> (text, tool_calls, CompletionTokenUsage)
Client->>Main: return (text, tool_calls, usage)
Main->>Config: after_chat_completion(..., usage)
Config->>Session: add_completion_usage(usage)
Session->>Session: accumulate completion_usage
Main->>Render: render_stream(..., spinner_message_with_usage)
Render->>Render: display spinner including token counts
Main->>Main: print final usage line after response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/client/common.rs (1)
583-596:⚠️ Potential issue | 🟠 MajorDon’t lose captured usage on local post-response failures.
Once
CompletionTokenUsage::new(...)has run, the provider has already returned billable counts. If a later local step fails here—especiallyeval_tool_calls()—start_directive_inner()falls into the zero-usage error path insrc/main.rsLines 484-486 and 507-509, so session totals silently undercount real completions. Keepusageseparate from these fallible local steps or surface it alongside the error.Also applies to: 637-654
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/client/common.rs` around lines 583 - 596, Compute and preserve the billable token usage immediately after calling CompletionTokenUsage::new(...) and do not let subsequent fallible local steps (like eval_tool_calls(...)) discard it; modify the code paths in start_directive_inner and the blocks around CompletionTokenUsage::new (including the similar block at 637-654) so that on any local error you either return a result type that carries both the error and the Usage or wrap/attach the Usage to the propagated error, e.g., call CompletionTokenUsage::new before any fallible work, keep the `usage` variable in scope, and change the error propagation from `Err(e)` to an error variant or context that includes `usage` when eval_tool_calls(...) (or print/markdown steps) fail so session totals always reflect the returned usage.
🧹 Nitpick comments (1)
src/client/common.rs (1)
553-565: Extract the spinner-label formatting into one helper.These blocks are identical. Centralizing the session-usage formatting keeps the streaming and non-streaming paths from drifting the next time the display rules change.
♻️ Suggested extraction
+fn completion_spinner_message(config: &GlobalConfig) -> String { + let config = config.read(); + if let Some(session) = &config.session { + let usage = session.completion_usage(); + if !usage.is_empty() { + return format!("Generating [{}]", usage); + } + } + "Generating".to_string() +} + pub async fn call_chat_completions( input: &Input, print: bool, extract_code: bool, client: &dyn Client, abort_signal: AbortSignal, ) -> Result<(String, Vec<ToolResult>, CompletionTokenUsage)> { - let spinner_message = { - let config = client.global_config().read(); - if let Some(session) = &config.session { - let su = session.completion_usage(); - if !su.is_empty() { - format!("Generating [{}]", su) - } else { - "Generating".to_string() - } - } else { - "Generating".to_string() - } - }; + let spinner_message = completion_spinner_message(client.global_config());Also applies to: 607-619
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/client/common.rs` around lines 553 - 565, The spinner label construction duplicated in the block that builds spinner_message and the identical block later should be extracted into a single helper function (e.g., format_spinner_label or completion_usage_label) that takes a &Client or &GlobalConfig/Option<Session> and returns the computed String using session.completion_usage(); replace both occurrences with calls to that helper so streaming and non‑streaming paths share the same formatting logic and avoid drift; update callers that currently call client.global_config().read() -> session -> completion_usage() to use the new helper.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/client/bedrock.rs`:
- Around line 286-291: The streaming metadata handler and the non-streaming
output path both ignore Bedrock's cacheReadInputTokens by passing None; update
the calls to read data["usage"]["cacheReadInputTokens"] and pass that value
through. Specifically, in the streaming branch where handler.set_usage(...) is
called, replace the third argument None with
data["usage"]["cacheReadInputTokens"].as_u64() (or an Option wrapper if the
field can be missing), and in the non-streaming path set
ChatCompletionsOutput.cached_tokens from data["usage"]["cacheReadInputTokens"]
(using as_u64()/Option) instead of leaving it None so cached token counts flow
through consistently with other providers.
In `@src/client/openai.rs`:
- Around line 196-209: The stream handler is capturing usage via
handler.set_usage(...) but the outgoing streamed chat completion request never
enables OpenAI's usage events; add stream_options: { "include_usage": true } to
the request body used for streamed chat completions so the final SSE chunk
contains usage data. Locate the code that builds/sends the streamed chat
completion request (the same code path that pairs with handler.set_usage) and
inject stream_options.include_usage = true into the JSON request payload for
streaming requests so usage information is emitted and parsed.
In `@src/config/session.rs`:
- Around line 69-70: The completion_usage field is being skipped by serde and
therefore resets on load and is not cleared by clear_messages() or
empty_session(); remove the #[serde(skip)] attribute from the
Session::completion_usage field so it's persisted with the session, and
additionally update clear_messages() and empty_session() methods to explicitly
reset self.completion_usage (e.g., to CompletionTokenUsage::default()) when
clearing to avoid carrying stale totals when clearing in-memory sessions.
In `@src/repl/mod.rs`:
- Around line 1158-1167: When config.read().session is None the code currently
falls back to the last-round `usage`, which misses earlier tool-call rounds;
thread a per-request `CompletionTokenUsage` accumulator through `ask_inner`
(pass it into and update it inside `ask_inner` where tool calls occur) and,
where `display_usage` is computed in the `tool_results.is_empty()` branch, use
the per-request accumulator (if present/initialized) as the fallback instead of
the stale `usage` variable; update references to `CompletionTokenUsage`,
`ask_inner`, `tool_results`, `display_usage`, and `usage` so the accumulator is
created at request start, passed into `ask_inner`, incremented on each
tool-driven round, and used when `config.session` is None.
---
Outside diff comments:
In `@src/client/common.rs`:
- Around line 583-596: Compute and preserve the billable token usage immediately
after calling CompletionTokenUsage::new(...) and do not let subsequent fallible
local steps (like eval_tool_calls(...)) discard it; modify the code paths in
start_directive_inner and the blocks around CompletionTokenUsage::new (including
the similar block at 637-654) so that on any local error you either return a
result type that carries both the error and the Usage or wrap/attach the Usage
to the propagated error, e.g., call CompletionTokenUsage::new before any
fallible work, keep the `usage` variable in scope, and change the error
propagation from `Err(e)` to an error variant or context that includes `usage`
when eval_tool_calls(...) (or print/markdown steps) fail so session totals
always reflect the returned usage.
---
Nitpick comments:
In `@src/client/common.rs`:
- Around line 553-565: The spinner label construction duplicated in the block
that builds spinner_message and the identical block later should be extracted
into a single helper function (e.g., format_spinner_label or
completion_usage_label) that takes a &Client or &GlobalConfig/Option<Session>
and returns the computed String using session.completion_usage(); replace both
occurrences with calls to that helper so streaming and non‑streaming paths share
the same formatting logic and avoid drift; update callers that currently call
client.global_config().read() -> session -> completion_usage() to use the new
helper.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 1b8011f4-ffa3-4c36-af72-5c5b9354abd3
📒 Files selected for processing (15)
.changesets/display-token-usage.mdsrc/acp/server.rssrc/client/bedrock.rssrc/client/claude.rssrc/client/cohere.rssrc/client/common.rssrc/client/openai.rssrc/client/stream.rssrc/client/vertexai.rssrc/config/mod.rssrc/config/session.rssrc/main.rssrc/render/mod.rssrc/render/stream.rssrc/repl/mod.rs
Comment on lines
+69
to
+70
| #[serde(skip)] | ||
| completion_usage: CompletionTokenUsage, |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf "\ncompletion_usage field + methods:\n"
sed -n '60,170p' src/config/session.rs
printf "\nclear_messages implementation:\n"
sed -n '532,540p' src/config/session.rs
printf "\nempty_session call sites / clear paths:\n"
rg -n -C3 --type rust '\bempty_session\s*\(|\bclear_messages\s*\(' src
printf "\nsave/load references:\n"
rg -n -C3 --type rust '\b(save|load)\s*\(' src/config/session.rs src/config/mod.rs
printf "\ncompletion_usage references:\n"
rg -n -C2 --type rust 'completion_usage|add_completion_usage' srcRepository: dobesv/harnx
Length of output: 12324
completion_usage is lost on reload and not cleared on clear_messages().
The field uses #[serde(skip)], so sessions loaded from disk reset completion_usage to its default (empty) state. Meanwhile, clear_messages() and empty_session() do not reset it, leaving stale totals after clearing the session in place. Either persist completion_usage as normal session state (remove #[serde(skip)]) or ensure it resets alongside message history.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/config/session.rs` around lines 69 - 70, The completion_usage field is
being skipped by serde and therefore resets on load and is not cleared by
clear_messages() or empty_session(); remove the #[serde(skip)] attribute from
the Session::completion_usage field so it's persisted with the session, and
additionally update clear_messages() and empty_session() methods to explicitly
reset self.completion_usage (e.g., to CompletionTokenUsage::default()) when
clearing to avoid carrying stale totals when clearing in-memory sessions.
Comment on lines
+1158
to
+1167
| if tool_results.is_empty() { | ||
| let session_usage = config | ||
| .read() | ||
| .session | ||
| .as_ref() | ||
| .map(|s| s.completion_usage().clone()); | ||
| let display_usage = session_usage.as_ref().unwrap_or(&usage); | ||
| if !display_usage.is_empty() { | ||
| eprintln!("{}", dimmed_text(&format!("[{}]", display_usage))); | ||
| } |
Contributor
There was a problem hiding this comment.
Accumulate tool-call rounds when no session is active.
When config.session is None, display_usage falls back to the last round’s usage. A tool-driven turn outside an explicit session will therefore print incomplete totals on the final line, because earlier rounds were never accumulated anywhere. Please thread a per-request CompletionTokenUsage accumulator through ask_inner and use that as the no-session fallback instead.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/repl/mod.rs` around lines 1158 - 1167, When config.read().session is None
the code currently falls back to the last-round `usage`, which misses earlier
tool-call rounds; thread a per-request `CompletionTokenUsage` accumulator
through `ask_inner` (pass it into and update it inside `ask_inner` where tool
calls occur) and, where `display_usage` is computed in the
`tool_results.is_empty()` branch, use the per-request accumulator (if
present/initialized) as the fallback instead of the stale `usage` variable;
update references to `CompletionTokenUsage`, `ask_inner`, `tool_results`,
`display_usage`, and `usage` so the accumulator is created at request start,
passed into `ask_inner`, incremented on each tool-driven round, and used when
`config.session` is None.
- Extract Bedrock cacheReadInputTokens in both streaming and non-streaming paths - Add stream_options.include_usage to OpenAI streaming requests so usage data is actually sent - Reset completion_usage in clear_messages() to avoid stale totals - Extract duplicated spinner label construction into spinner_label() helper
This was referenced Apr 5, 2026
Merged
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.
Summary
Implements #152 — show token usage numbers (input/output/cached tokens) in the CLI output.
What it does
[📥 12345 📤 678 💾 9000]⠋ Generating [📥 5000 📤 200]Emoji format
Uses compact emoji labels instead of verbose text:
Provider support
Extracts cached token counts from providers that support it:
usage.prompt_tokens_details.cached_tokens(streaming: final usage-only chunk)usage.cache_read_input_tokens(streaming: cumulativemessage_delta)usageMetadata.cachedContentTokenCountFiles changed (15)
src/client/common.rs—CompletionTokenUsagestruct, propagate usage from both streaming/non-streaming paths, spinner message with session totalssrc/client/stream.rs— Usage tracking inSseHandler, updatedtake()return typesrc/client/openai.rs— Cached tokens extraction, streaming usage from final chunk onlysrc/client/claude.rs— Cached tokens extraction, streaming usage frommessage_start/message_delta(cumulative)src/client/vertexai.rs— Cached tokens extraction (non-stream + stream)src/client/bedrock.rs— Streaming metadata usage extractionsrc/client/cohere.rs—cached_tokens: Nonesrc/config/session.rs— Session-levelCompletionTokenUsageaccumulatorsrc/config/mod.rs—after_chat_completionaccumulates usage in sessionsrc/render/mod.rs+src/render/stream.rs— Spinner message passthroughsrc/main.rs+src/repl/mod.rs— Display usage after final response, pass usage toafter_chat_completionsrc/acp/server.rs— Handle updatedtake()3-tuple.changesets/display-token-usage.md— ChangesetVerification
cargo fmt --all --check✅cargo clippy --all --all-targets -- -D warnings✅cargo test --all✅ (162 passed; 1 pre-existing flaky test)Closes #152
Summary by CodeRabbit