feat: track and display LLM token usage across CLI, reports, and extension - #224
Open
arunSunnyKVS wants to merge 1 commit into
Open
feat: track and display LLM token usage across CLI, reports, and extension#224arunSunnyKVS wants to merge 1 commit into
arunSunnyKVS wants to merge 1 commit into
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>
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 issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches📝 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
Comment on lines
+243
to
+245
| const evResult = toEvaluatorResult(evaluatorMeta, attackResults); | ||
| if (evalTracker) evResult.tokenUsage = evalTracker.totals; | ||
| evaluatorResults.push(evResult); |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Preserve per-evaluator usage in partial results.
The normal completion paths add tokenUsage, but stop/error branches push toEvaluatorResult(...) directly. Partial reports therefore retain run totals while omitting the affected evaluator’s usage.
core/src/execute/evaluatorLoop.ts#L243-L245: centralize evaluator-result decoration so the early stop return also receivesevalTracker.totals.core/src/execute/runAllBrowser.ts#L232-L242: apply the same decoration in both attack error/stop branches.
📍 Affects 2 files
core/src/execute/evaluatorLoop.ts#L243-L245(this comment)core/src/execute/runAllBrowser.ts#L232-L242
🤖 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/src/execute/evaluatorLoop.ts` around lines 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.
Comment on lines
+26
to
+38
| /** Record usage from a single LLM call. Safe to call with undefined/partial usage. */ | ||
| record(usage?: { inputTokens?: number; outputTokens?: number; totalTokens?: number }): void { | ||
| if (!usage) return; | ||
| this.input += usage.inputTokens ?? 0; | ||
| this.output += usage.outputTokens ?? 0; | ||
| } | ||
|
|
||
| get totals(): TokenUsage { | ||
| return { | ||
| inputTokens: this.input, | ||
| outputTokens: this.output, | ||
| totalTokens: this.input + this.output, | ||
| }; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Preserve totalTokens end to end. The tracker discards explicit totals, the OpenAI-compatible adapter does not forward total_tokens, and the tests do not cover this contract.
core/src/execute/tokenTracker.ts#L26-L38: accumulate supplied totals and fall back to component sums only when absent.core/src/llm/openaiCompatible.ts#L151-L157: passdata.usage.total_tokensintoTokenTracker.record().core/tests/tokenTracker.test.ts#L10-L22: add total-only and mismatched-total regression cases.
📍 Affects 3 files
core/src/execute/tokenTracker.ts#L26-L38(this comment)core/src/llm/openaiCompatible.ts#L151-L157core/tests/tokenTracker.test.ts#L10-L22
🤖 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/src/execute/tokenTracker.ts` around lines 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.
Comment on lines
+149
to
+152
| if (tokenTracker && result && typeof result === "object" && "usage" in result) { | ||
| tokenTracker.record( | ||
| (result as { usage?: { inputTokens?: number; outputTokens?: number } }).usage | ||
| ); |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Files:"
git ls-files | rg '(^|/)llmRetry\.ts$|openaiCompatible\.ts$|generateNextTurn\.ts$|zod' || true
echo
echo "Relevant snippets:"
for f in core/src/lib/llmRetry.ts core/src/llm/openaiCompatible.ts core/src/generate/generateNextTurn.ts; do
if [ -f "$f" ]; then
echo "===== $f ====="
wc -l "$f"
rg -n "tokenTracker|result\.usage|array.*result|generateText|usage" "$f" -C 4 || true
fi
done
echo
echo "Search tokenTracker implementations/usages:"
rg -n "class .*TokenTracker|record\\(|tokenTracker|Usage|inputTokens|outputTokens" core/src -g '*.ts' || trueRepository: KeyValueSoftwareSystems/agent-opfor
Length of output: 14939
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Dependencies and existing Zod usage:"
if [ -f package.json ]; then cat package.json | sed -n '1,220p'; fi
echo
rg -n "\bzod\b|from \"zod\"|from 'zod'" -g '*.ts' -g '*.tsx' | head -200 || trueRepository: KeyValueSoftwareSystems/agent-opfor
Length of output: 2738
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "tokenTracker implementation:"
cat -n core/src/execute/tokenTracker.ts
echo
echo "openaiCompatible relevant section:"
sed -n '130,170p' core/src/llm/openaiCompatible.ts | cat -n -v
echo
echo "generateNextTurn relevant sections:"
sed -n '130,165p' core/src/generate/generateNextTurn.ts | cat -n -v
sed -n '285,310p' core/src/generate/generateNextTurn.ts | cat -n -v
echo
echo "LLM retry full relevant section:"
sed -n '130,160p' core/src/lib/llmRetry.ts | cat -n -v
echo
echo "Package lock zod versions if present:"
for f in package-lock.json core/package.json core/package-lock.json; do
if [ -f "$f" ]; then
echo "===== $f ====="
rg -n '"zod"|"`@hono/node-server`"|node_modules/zod|zod' "$f" | head -60 || true
fi
done
echo
echo "Behavioral probe: TokenTracker accepts numeric/non-numeric and non-object values with optional inputs:"
node - <<'JS'
const inputs = [
undefined,
null,
123,
"12",
{},
{ inputTokens: "9", outputTokens: "11" },
{ inputTokens: NaN, outputTokens: null },
{ inputTokens: true, outputTokens: "2" },
];
let input = 0, output = 0;
function record(usage) {
input += usage.inputTokens ?? 0;
output += usage.outputTokens ?? 0;
}
for (const usage of inputs) {
try {
record(usage);
console.log(JSON.stringify(usage) + " -> " + input + "/" + output);
} catch (e) {
console.log(JSON.stringify(usage) + " threw " + e.name + ": " + e.message);
}
}
JSRepository: KeyValueSoftwareSystems/agent-opfor
Length of output: 10171
Validate LLM usage before recording. These boundaries pass provider/SDK response objects through casts or unchecked properties into TokenTracker.record(), which can coerce malformed values before persisting token totals. Validate the final usage shape with Zod before recording at llmRetry.ts, openaiCompatible.ts, and both generateNextTurn.ts call sites.
📍 Affects 3 files
core/src/lib/llmRetry.ts#L149-L152(this comment)core/src/llm/openaiCompatible.ts#L149-L157core/src/generate/generateNextTurn.ts#L154-L155core/src/generate/generateNextTurn.ts#L301-L302
🤖 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/src/lib/llmRetry.ts` around lines 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.
Source: Coding guidelines
Comment on lines
+220
to
+224
| ``` | ||
| Results: 5 passed, 2 failed, 0 errors | ||
| Safety score: 71% | ||
| Token usage: 51,323 input / 6,057 output (57,380 total) | ||
| ``` |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Specify the fenced block language.
Add text (or another appropriate language) after the opening fence to satisfy markdownlint MD040.
-```
+```text🧰 Tools
🪛 markdownlint-cli2 (0.23.0)
[warning] 220-220: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🤖 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 `@docs/cli.md` around lines 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.
Source: Linters/SAST tools
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
TokenTrackerclass (core/src/execute/tokenTracker.ts) that meters every LLM call (attacker generation, adaptive follow-ups, judge) and aggregates per-evaluator and per-run totalsrunAll→evaluatorLoop→ attack drivers →withRetry/generateText/chatCompletionJsonContentWhat changed
Core engine
TokenTrackerwith parent-child hierarchy — root tracker per run, child tracker per evaluatorwithRetrywrapper automatically capturesusagefromgenerateText/generateObjectresultsgenerateTextandchatCompletionJsonContentcall sites explicitly record usageTokenUsagetype added toexecute/types.tsandreport/types.tsEvaluatorResult.tokenUsageandUnifiedRunReport.summary.tokenUsagecarry the data throughCLI
Token usage: 51,323 input / 6,057 output (57,380 total)HTML report (core
render.ts)JSON report
summary.tokenUsageand per-evaluatortokenUsagefields for CI consumptionBrowser extension
orchestrator.js: propagatestokenUsagefromrunAllBrowserresult to popuppopup.js/popup.html: Tokens stat on Done screen with input/output breakdown; per-evaluator counts in result rows and downloaded HTML reportpruneRawForHistoryto preservetokenUsage(was being stripped by the allowlist)# Detailsplaceholder → proper styled "5 Detailed Results" section headerDocs
docs/cli.md: new "Token usage tracking" sectiondocs/browser-extension.md: note about token displayAGENTS.md:tokenTracker.tsin key files table + tracking mechanism descriptionTest plan
TokenTrackerunit tests (core/tests/tokenTracker.test.ts) — parent-child aggregation,record(),child(),usagegetteropfor run --config tests/e2e/agents/customer-support/opfor.config.json— verify CLI prints token summary and HTML/JSON reports include token dataMade with Cursor
Summary by CodeRabbit
New Features
Documentation