feat: track and display LLM token usage across CLI, reports, and extension - #224
Merged
Conversation
…nsion Add a TokenTracker class that meters every LLM call (attacker generation, adaptive follow-ups, judge) and aggregates per-evaluator and per-run totals. - Core: thread TokenTracker through runAll → evaluatorLoop → attack drivers; withRetry, generateText, and chatCompletionJsonContent auto-record usage - CLI: print token summary line after run results - Reports: add Token Usage stat card to HTML executive summary, per-evaluator counts in detail headers, and tokenUsage fields in JSON output - Extension: propagate tokenUsage from orchestrator to popup; show Tokens stat on Done screen with in/out breakdown; per-evaluator counts in result rows and downloaded HTML report; fix pruneRawForHistory to preserve tokenUsage; fix "# Details" → proper section 5 header in HTML template - Docs: update cli.md, browser-extension.md, AGENTS.md Co-authored-by: Cursor <cursoragent@cursor.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
WalkthroughThe PR adds nested token tracking across LLM generation and judging, attaches usage to evaluator and run results, and displays token counts in CLI, JSON, HTML, and browser-extension outputs. ChangesToken usage tracking
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
arunSunnyKVS
marked this pull request as ready for review
July 28, 2026 11:54
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
core/tests/tokenTracker.test.ts (1)
10-22: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover the
totalTokensbranch.These tests only verify totals derived from input/output, so they would pass while explicit provider totals are silently ignored. Add a total-only case and a case where
totalTokensdiffers from the component sum.🤖 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 `@core/tests/tokenTracker.test.ts` around lines 10 - 22, Add tests in the TokenTracker record suite for an input containing only explicit totalTokens, and for input where totalTokens differs from inputTokens plus outputTokens; assert that totals preserve the provider-supplied total in both cases while retaining the expected component values.
🤖 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.
Inline comments:
In `@core/src/execute/evaluatorLoop.ts`:
- Around line 243-245: Centralize evaluator-result decoration so partial results
preserve per-evaluator token usage: in core/src/execute/evaluatorLoop.ts lines
243-245, ensure the early-stop return applies evalTracker.totals to
toEvaluatorResult output; in core/src/execute/runAllBrowser.ts lines 232-242,
apply the same decoration in both attack error/stop branches. Reuse the existing
evaluator result and evalTracker symbols without changing normal completion
behavior.
In `@core/src/execute/tokenTracker.ts`:
- Around line 26-38: Preserve explicit total token counts end to end: update
TokenTracker.record in core/src/execute/tokenTracker.ts to accumulate supplied
totalTokens and use input/output sums only when totalTokens is absent; update
the OpenAI-compatible adapter in core/src/llm/openaiCompatible.ts to pass
data.usage.total_tokens; extend core/tests/tokenTracker.test.ts with total-only
and mismatched-total regression cases.
In `@core/src/lib/llmRetry.ts`:
- Around line 149-152: Validate the final LLM usage object with the existing Zod
validation approach before calling TokenTracker.record(), replacing unchecked
casts/property access. Apply this at core/src/lib/llmRetry.ts:149-152,
core/src/llm/openaiCompatible.ts:149-157, and both
core/src/generate/generateNextTurn.ts:154-155 and :301-302 call sites; only
record usage when it matches the expected inputTokens/outputTokens shape.
In `@docs/cli.md`:
- Around line 220-224: Update the fenced code block in the CLI results example
by specifying an appropriate language, such as text, immediately after its
opening fence to satisfy markdownlint MD040; leave the example content
unchanged.
---
Nitpick comments:
In `@core/tests/tokenTracker.test.ts`:
- Around line 10-22: Add tests in the TokenTracker record suite for an input
containing only explicit totalTokens, and for input where totalTokens differs
from inputTokens plus outputTokens; assert that totals preserve the
provider-supplied total in both cases while retaining the expected component
values.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 9a0f354c-09be-4e06-b271-c5d6a75e74b6
⛔ Files ignored due to path filters (1)
package-lock.jsonis excluded by!**/package-lock.json
📒 Files selected for processing (24)
AGENTS.mdcore/src/evaluators/judge.tscore/src/execute/agentAttackDriver.tscore/src/execute/evaluatorLoop.tscore/src/execute/mcpAttackDriver.tscore/src/execute/runAll.tscore/src/execute/runAllBrowser.tscore/src/execute/tokenTracker.tscore/src/execute/types.tscore/src/generate/generateAttacks.tscore/src/generate/generateNextTurn.tscore/src/lib/llmRetry.tscore/src/llm/openaiCompatible.tscore/src/report/buildReport.tscore/src/report/render.tscore/src/report/types.tscore/src/run/judge.tscore/tests/tokenTracker.test.tsdocs/browser-extension.mddocs/cli.mdrunners/cli/src/commands/run.tsrunners/extension/orchestrator.jsrunners/extension/popup.htmlrunners/extension/popup.js
…cstrings Centralize evaluator-result decoration so stop/error branches also receive evalTracker.totals, preventing partial reports from losing the affected evaluator's usage. Add docstrings across all modified files to meet the 80% coverage threshold required by CodeRabbit. Co-authored-by: Cursor <cursoragent@cursor.com>
TokenTracker now accumulates the provider-supplied totalTokens (which may include reasoning tokens) instead of recomputing it from inputTokens + outputTokens. All recording boundaries (withRetry, chatCompletionJsonContent, generateNextTurn) now validate the usage shape with a Zod schema before persisting, rejecting negative, non-integer, or unexpected fields. The OpenAI-compatible adapter now forwards total_tokens into TokenTracker.record(). New test cases cover total-only usage, mismatched totals, and parseUsage rejection of malformed input. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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.
Inline comments:
In `@core/src/execute/tokenTracker.ts`:
- Around line 30-41: Update LlmUsageSchema in core/src/execute/tokenTracker.ts
to strip unknown provider metadata while validating and normalizing inputTokens,
outputTokens, and totalTokens; remove the strict rejection behavior. Update the
regression test in core/tests/tokenTracker.test.ts at lines 119-123 to expect
the normalized tracked usage when extra metadata fields are present.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: bba93145-28f3-4675-b71c-c749198786e6
📒 Files selected for processing (5)
core/src/execute/tokenTracker.tscore/src/generate/generateNextTurn.tscore/src/lib/llmRetry.tscore/src/llm/openaiCompatible.tscore/tests/tokenTracker.test.ts
🚧 Files skipped from review as they are similar to previous changes (3)
- core/src/lib/llmRetry.ts
- core/src/llm/openaiCompatible.ts
- core/src/generate/generateNextTurn.ts
Replace .strict() with .passthrough() in LlmUsageSchema so provider-specific fields (cachedTokens, reasoningTokens, etc.) are silently stripped instead of failing the entire parse. Co-authored-by: Cursor <cursoragent@cursor.com>
jithin23-kv
approved these changes
Jul 29, 2026
This was referenced Jul 30, 2026
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
2 participants
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.
Problem
opfor runmakes multiple LLM calls per run (attacker generation, adaptive follow-ups, judge) via the Vercel AI SDK, which returns token usage on everygenerateText()result. Currently all usage data is discarded — users have no visibility into how many tokens a run consumed, making cost management and budget alerting impossible.Solution
Introduce a
TokenTrackerclass with a parent-child hierarchy: one root tracker per run, one child per evaluator. Every LLM call site (withRetry,generateText,generateObject,chatCompletionJsonContent) now records its usage. Aggregated totals are surfaced in the CLI summary, HTML/JSON reports, and browser extension popup.Changes
core/execute/tokenTracker.ts— NewTokenTracker/ChildTrackerclasses andTokenUsageinterfaceexecute/types.ts— Added optionaltokenUsagefield toEvaluatorResultandUnifiedRunReport.summaryexecute/runAll.ts— ThreadTokenTrackerthroughRunAllOptions, attach run-level totals to report summaryexecute/evaluatorLoop.ts— Create per-evaluator child trackers, attachevalTracker.totalsto every evaluator result (including partial/error/stop branches)execute/runAllBrowser.ts— Same per-evaluator tracking for browser path, centralized partial-result decorationexecute/agentAttackDriver.ts— PasstokenTrackerto judge callsexecute/mcpAttackDriver.ts— PasstokenTrackerto MCP judge callslib/llmRetry.ts—withRetryauto-recordsusagefrom successful results whentokenTrackeris providedllm/openaiCompatible.ts—chatCompletionJsonContentparses and recordsusagefrom raw fetch responsesgenerate/generateAttacks.ts— ThreadtokenTrackerthrough attack generationgenerate/generateNextTurn.ts— ThreadtokenTrackerthrough adaptive follow-up generationevaluators/judge.ts— ThreadtokenTrackerthrough judge callsrun/judge.ts— ThreadtokenTrackerthrough MCP judge helpersreport/buildReport.ts— MaptokenUsageintoEvaluatorViewModelfor HTML renderingreport/render.ts— Display per-evaluator token counts in HTML report detail accordionsreport/types.ts— AddedtokenUsagetoEvaluatorViewModeland report summary typerunners/cli/commands/run.ts— Display token usage summary line after resultsrunners/extension/orchestrator.js— PasstokenUsagefromrunAllBrowserreport to popup result objects (success + error paths)popup.js— Aggregate and display token usage on Done screen (total + input/output breakdown), per-evaluator tokens in result rows, per-evaluator tokens in downloadable HTML reportpopup.html— Token stat card in Done screen, CSS adjustments for 2-column stat gridTests
core/tests/tokenTracker.test.ts— Unit tests forTokenTrackerparent-child hierarchyDocs
docs/cli.md— Token usage tracking sectiondocs/browser-extension.md— Token usage noteAGENTS.md—tokenTracker.tsin key files table, token tracking mechanism docsIssue
Closes #223
How to test
CLI
npm run buildopfor run --config tests/e2e/agents/customer-support/opfor.config.jsonsummary.tokenUsageand per-evaluatortokenUsagefieldsBrowser extension
npm run build(rebuildscore.bundle.js)chrome://extensions→ Load unpacked →runners/extension/)Unit tests
npm test -- --run core/tests/tokenTracker.test.tsScreenshots
N/A
Summary by CodeRabbit