feat: token usage tracking - #10
Conversation
Captures input/output/total token counts from Vertex AI responses: - Provider.Review() reads UsageMetadata from genai response - Reviewer aggregates across chunks - MultiProvider mergeResults aggregates across models - Terminal summary box shows: 'Tokens: 1500 in / 200 out / 1700 total' - JSON output includes 'usage' field - slog.Info logs usage for CI job visibility Tests: - TestColorTerminalOutput_TokenUsage: verifies counts in output - TestColorTerminalOutput_NoTokenUsage: hidden when zero - 147 tests total
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThis PR adds token usage tracking from provider responses through result aggregation and terminal rendering. It also adds tests for merged usage counts and output formatting with and without token data. ChangesToken Usage Tracking
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Provider
participant Reviewer
participant mergeResults
participant Terminal
Provider->>Provider: read UsageMetadata
Provider->>Reviewer: ReviewResult with Usage
Reviewer->>Reviewer: accumulate totalUsage per chunk
Reviewer->>mergeResults: per-model ReviewResults
mergeResults->>mergeResults: sum InputTokens, OutputTokens, TotalTokens
mergeResults-->>Reviewer: merged ReviewResult with Usage
Reviewer->>Terminal: ReviewResult with aggregated Usage
Terminal->>Terminal: render token usage row if TotalTokens > 0
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the Comment |
![coderabbitai[bot]](https://avatars.github.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
internal/model/provider.go (1)
21-26: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win
omitemptyon struct fieldUsagehas no effect.Go's
encoding/jsonnever treats a non-pointer struct as "empty" foromitemptypurposes — onlyfalse,0, nil pointers/interfaces, and empty arrays/slices/maps/strings qualify.TokenUsageis a plain struct, soUsagewill always be serialized in JSON output, even when all counts are zero, despite theomitemptytag suggesting otherwise.If the intent is to hide the usage block when zero (matching the terminal's conditional behavior), use a pointer (
*TokenUsage) and only set it whenTotalTokens > 0, or use Go 1.24'somitzerotag if the module's Go version supports it.Proposed fix using a pointer field
type ReviewResult struct { Summary string `json:"summary"` Findings []Finding `json:"findings"` - Usage TokenUsage `json:"usage,omitempty"` + Usage *TokenUsage `json:"usage,omitempty"` }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/model/provider.go` around lines 21 - 26, The `ReviewResult` struct’s `Usage` field uses `omitempty` on a non-pointer `TokenUsage`, so it will still always appear in JSON. Update `ReviewResult` in `provider.go` so `Usage` can actually be omitted when empty, preferably by changing it to a pointer `*TokenUsage` and only populating it when `TotalTokens > 0` in the code that builds `ReviewResult`; alternatively, if the project targets Go 1.24+, use `omitzero` instead of `omitempty`.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@internal/model/provider.go`:
- Around line 21-26: The `ReviewResult` struct’s `Usage` field uses `omitempty`
on a non-pointer `TokenUsage`, so it will still always appear in JSON. Update
`ReviewResult` in `provider.go` so `Usage` can actually be omitted when empty,
preferably by changing it to a pointer `*TokenUsage` and only populating it when
`TotalTokens > 0` in the code that builds `ReviewResult`; alternatively, if the
project targets Go 1.24+, use `omitzero` instead of `omitempty`.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 58d073c0-c5b6-4673-8482-761120c4dcce
📒 Files selected for processing (5)
internal/model/multi.gointernal/model/provider.gointernal/reviewer/output_test.gointernal/reviewer/reviewer.gointernal/reviewer/terminal.go
Fixes from deep review:
- TokenUsage fields: int32 → int64 (defensive for multi-chunk × multi-model)
- Usage field: TokenUsage → *TokenUsage (omitempty now properly suppresses
zero-value JSON output instead of emitting {0,0,0})
- All nil-pointer checks added (terminal, reviewer, multi aggregation)
- Plain-text TerminalOutput() now shows token usage (was missing)
Tests added:
- TestMergeResults_UsageAggregation: verifies summing across models
- TestMergeResults_UsageNilWhenNoUsage: nil when no results have usage
- TestTerminalOutput_PlainText_TokenUsage: plain-text path coverage
- TestTerminalOutput_PlainText_NoTokenUsage: hidden when nil
- 151 tests total
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
internal/reviewer/output.go (1)
16-54: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winToken usage line unreachable when there are no findings.
The early return at line 21-24 (
len(result.Findings) == 0) exits before the token usage block at lines 50-54 is reached, so plain-textTerminalOutputnever prints tokens when there are no findings. This contradictsColorTerminalOutput, which places its equivalent usage block before its own no-findings return, so it does show tokens even with zero findings. This breaks parity between the two renderers and doesn't match the PR's goal of surfacing token usage in terminal output regardless of finding count.🐛 Proposed fix
if len(result.Findings) == 0 { sb.WriteString("✅ No issues found. Code looks clean and ready to merge.\n") + if result.Usage != nil && result.Usage.TotalTokens > 0 { + fmt.Fprintf(&sb, "---\nTokens: %d in / %d out / %d total\n", + result.Usage.InputTokens, result.Usage.OutputTokens, result.Usage.TotalTokens) + } return sb.String() }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/reviewer/output.go` around lines 16 - 54, Move the token usage rendering in TerminalOutput so it runs before the len(result.Findings) == 0 early return, matching the behavior of ColorTerminalOutput. Keep the no-findings message, but ensure the usage block in TerminalOutput is emitted whenever result.Usage.TotalTokens is set, even when there are no findings. Use TerminalOutput and the existing result.Usage block as the main touchpoints.internal/reviewer/output_test.go (1)
235-270: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winMissing coverage for "usage present, no findings" case.
None of the added tests cover
TerminalOutput/ColorTerminalOutputwith a non-nilUsageand zero findings, which is exactly the scenario where plain-text output currently drops the token line (see comment onoutput.go). Adding this case would have caught the early-return bug.Want me to add a test like
TestTerminalOutput_PlainText_TokenUsage_NoFindingsasserting the tokens line still appears whenFindingsis empty butUsageis set?🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/reviewer/output_test.go` around lines 235 - 270, Add coverage for the case where ReviewResult.Usage is set but Findings is empty, since TerminalOutput and ColorTerminalOutput currently skip the token line in that path. Update output_test.go by adding a focused test around TerminalOutput and/or ColorTerminalOutput that uses non-nil Usage with zero findings and asserts the token counts still appear. Keep the test aligned with the existing TerminalOutput_PlainText_TokenUsage and TestColorTerminalOutput_NoTokenUsage patterns so it exercises the early-return behavior in the output helpers.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@internal/reviewer/output_test.go`:
- Around line 235-270: Add coverage for the case where ReviewResult.Usage is set
but Findings is empty, since TerminalOutput and ColorTerminalOutput currently
skip the token line in that path. Update output_test.go by adding a focused test
around TerminalOutput and/or ColorTerminalOutput that uses non-nil Usage with
zero findings and asserts the token counts still appear. Keep the test aligned
with the existing TerminalOutput_PlainText_TokenUsage and
TestColorTerminalOutput_NoTokenUsage patterns so it exercises the early-return
behavior in the output helpers.
In `@internal/reviewer/output.go`:
- Around line 16-54: Move the token usage rendering in TerminalOutput so it runs
before the len(result.Findings) == 0 early return, matching the behavior of
ColorTerminalOutput. Keep the no-findings message, but ensure the usage block in
TerminalOutput is emitted whenever result.Usage.TotalTokens is set, even when
there are no findings. Use TerminalOutput and the existing result.Usage block as
the main touchpoints.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: a117cd97-c879-4d33-8d80-4fb24a60060d
📒 Files selected for processing (7)
internal/model/multi.gointernal/model/multi_test.gointernal/model/provider.gointernal/reviewer/output.gointernal/reviewer/output_test.gointernal/reviewer/reviewer.gointernal/reviewer/terminal.go
🚧 Files skipped from review as they are similar to previous changes (4)
- internal/reviewer/terminal.go
- internal/model/provider.go
- internal/reviewer/reviewer.go
- internal/model/multi.go
…odeRabbit) TerminalOutput() early-returned before reaching the usage block when Findings was empty. Added usage rendering before the early return, matching ColorTerminalOutput's behavior. Tests: - TestTerminalOutput_PlainText_UsageWithNoFindings - TestColorTerminalOutput_UsageWithNoFindings - 153 tests total
There was a problem hiding this comment.
🧹 Nitpick comments (1)
internal/reviewer/output_test.go (1)
272-300: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffConsider table-driven tests for these usage-rendering cases.
These two tests (and the similarly-shaped ones at lines 216-270) could be consolidated into a single table-driven test over
{usage *model.TokenUsage, findings []model.Finding, wantContains []string}cases, covering plain-text/color and with/without findings. Given the existing file already uses this standalone-function style throughout, this is optional and can be deferred.As per path instructions, "
**/*_test.go: Verify table-driven tests, proper cleanup, and race condition safety."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/reviewer/output_test.go` around lines 272 - 300, The usage-rendering tests in TerminalOutput and ColorTerminalOutput are duplicated across plain-text and colored cases, so consolidate them into one table-driven test that iterates over combinations of ReviewResult.Usage and Findings and checks the expected substrings for both output functions. Keep the existing assertions for token totals and “no findings” behavior, but drive them from a shared cases table keyed by TerminalOutput, ColorTerminalOutput, and model.TokenUsage to reduce repetition and make future cases easier to add.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@internal/reviewer/output_test.go`:
- Around line 272-300: The usage-rendering tests in TerminalOutput and
ColorTerminalOutput are duplicated across plain-text and colored cases, so
consolidate them into one table-driven test that iterates over combinations of
ReviewResult.Usage and Findings and checks the expected substrings for both
output functions. Keep the existing assertions for token totals and “no
findings” behavior, but drive them from a shared cases table keyed by
TerminalOutput, ColorTerminalOutput, and model.TokenUsage to reduce repetition
and make future cases easier to add.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 695a3259-8c8d-4115-8e04-1d86a943335e
📒 Files selected for processing (2)
internal/reviewer/output.gointernal/reviewer/output_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
- internal/reviewer/output.go
Replaces 6 standalone token usage tests with 1 table-driven test covering 3 cases × 2 renderers = 6 subtests. Easier to extend.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
internal/reviewer/output_test.go (1)
226-245: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMissing test case: findings present + nil usage.
TerminalOutput(output.go) has two separateresult.Usage != nilchecks guarding the "Tokens:" line — one in the no-findings early-return branch, one after the by-file loop for the has-findings branch. The current table only exercises the has-findings branch with usage present (case 1) and the no-findings branch with usage absent (case 2). There's no case verifying that the has-findings branch correctly omits "Tokens:" whenUsageis nil, so a regression dropping that nil-check in the second location would go undetected.♻️ Suggested additional test case
{ name: "no findings but usage shown", findings: nil, usage: &model.TokenUsage{InputTokens: 3000, OutputTokens: 100, TotalTokens: 3100}, wantContains: []string{"3000", "3100"}, }, + { + name: "findings present but nil usage hidden", + findings: []model.Finding{finding}, + usage: nil, + wantAbsent: []string{"Tokens:"}, + }, }As per path instructions for
**/*_test.go: "Verify table-driven tests, proper cleanup, and race condition safety."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/reviewer/output_test.go` around lines 226 - 245, Add a table-driven case for TerminalOutput that covers findings present with nil Usage so the has-findings path is exercised without printing "Tokens:". Update the existing output tests in output_test.go alongside the current TerminalOutput scenarios, and assert that the output contains the findings but omits the token line when result.Usage is nil. This will verify the second result.Usage != nil guard in TerminalOutput and prevent regressions in the by-file branch.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@internal/reviewer/output_test.go`:
- Around line 226-245: Add a table-driven case for TerminalOutput that covers
findings present with nil Usage so the has-findings path is exercised without
printing "Tokens:". Update the existing output tests in output_test.go alongside
the current TerminalOutput scenarios, and assert that the output contains the
findings but omits the token line when result.Usage is nil. This will verify the
second result.Usage != nil guard in TerminalOutput and prevent regressions in
the by-file branch.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 8f4e92e7-ebac-4247-bac8-e680f4063dcd
📒 Files selected for processing (1)
internal/reviewer/output_test.go
1 participant
Summary
Adds input/output/total token tracking from Vertex AI responses. Zero external deps.
Where it shows up
Terminal (in the summary box):
JSON (
--json):{"usage": {"input_tokens": 12847, "output_tokens": 1203, "total_tokens": 14050}}CI logs (via slog):
Changes
internal/model/provider.goTokenUsagestruct, capture fromUsageMetadatainternal/model/multi.gomergeResultsinternal/reviewer/reviewer.gointernal/reviewer/terminal.gointernal/reviewer/output_test.go147 tests, all passing.
Summary by CodeRabbit