feat: per-run token and cost accounting (#40)#77
Conversation
|
Warning Review limit reached
Next review available in: 13 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Repository UI (base), Organization UI (inherited) Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughToken usage is captured at the LLM client/step layer and threaded through ReactAgent's LoopRunner, AgentSpawnerService, and RunnerService. A new CostCalculator service prices usage per model. JsonFormatter and OutputFormatter render token counts and estimated cost. Tests cover all new paths. ChangesToken usage and cost pipeline
Sequence Diagram(s)sequenceDiagram
participant LoopRunner
participant Step
participant AgentSpawnerService
participant RunnerService
participant CostCalculator
participant OutputFormatter
LoopRunner->>Step: call(iteration)
Step-->>LoopRunner: step_result with usage
LoopRunner->>LoopRunner: accumulate total_usage
LoopRunner-->>AgentSpawnerService: response with usage
AgentSpawnerService-->>RunnerService: result with usage
RunnerService->>RunnerService: aggregate_usage(baseline, context)
RunnerService->>CostCalculator: call(usage:, model:)
CostCalculator-->>RunnerService: cost
RunnerService-->>OutputFormatter: result with tokens, cost
OutputFormatter->>OutputFormatter: build_usage_line(result)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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 |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
lib/skill_bench/agent/react_agent/loop_runner.rb (1)
72-100: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate usage-aggregation logic across services.
empty_usage,add_usage, andtoken_counthere mirror equivalent helpers inRunnerService(per cross-file graph evidence showingadd_usage/token_countcalls atrunner_service.rblines ~204-218). Having the same token-summing logic implemented twice risks the two diverging (e.g. a future fix to tolerant key-reading lands in one file but not the other).Consider extracting a shared
Services::UsageAggregator(or similar) module withempty_usage/add_usage/token_countthat bothLoopRunnerandRunnerServicecall.🤖 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 `@lib/skill_bench/agent/react_agent/loop_runner.rb` around lines 72 - 100, The token-usage helpers in LoopRunner duplicate the same aggregation logic already used in RunnerService, so move empty_usage, add_usage, and token_count into a shared usage-aggregation module or service and have both LoopRunner and RunnerService call it. Keep the existing behavior for nil and string-keyed usage hashes, and update the references in LoopRunner to use the shared helper methods so the summing logic stays consistent in one place.lib/skill_bench/services/runner_service.rb (1)
177-220: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate usage-normalization helpers across files.
empty_usage,add_usage, andtoken_counthere reimplement the exact same logic already defined inLoopRunner(per the cross-file graph evidence showing identical method bodies).JsonFormatter::EMPTY_USAGEindependently re-declares the same zeroed-usage shape, andCostCalculator#token_counthas yet another (single-arg) variant of the same "tolerate string keys" logic. Four near-identical implementations of the same small concept is a DRY smell that will diverge over time (e.g., if a key name or fallback rule changes, it now needs updating in 4 places).Consider extracting a small shared utility (e.g.
SkillBench::Usagemodule with.empty,.add,.token_count) used byLoopRunner,RunnerService,JsonFormatter, andCostCalculator.🤖 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 `@lib/skill_bench/services/runner_service.rb` around lines 177 - 220, The usage-normalization logic is duplicated across RunnerService and other classes, so centralize it in one shared helper instead of keeping separate copies of empty_usage, add_usage, and token_count. Extract a common usage utility/module (for example, a SkillBench::Usage helper) and update RunnerService’s aggregate_usage path, along with LoopRunner, JsonFormatter::EMPTY_USAGE, and CostCalculator#token_count, to use the shared implementation so the token-count rules stay consistent in one place.
🤖 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 `@lib/skill_bench/services/cost_calculator.rb`:
- Around line 6-26: The CostCalculator::PRICES table is missing the newer model
families used by the project, so lookups in CostCalculator#price_for fall
through to nil for models like claude-sonnet-4-20250514 and claude-opus-4-7.
Update PRICES to include the current GPT and Claude families/aliases the
codebase actually passes in, preserving the longest-prefix matching behavior so
dated model variants still resolve correctly.
---
Nitpick comments:
In `@lib/skill_bench/agent/react_agent/loop_runner.rb`:
- Around line 72-100: The token-usage helpers in LoopRunner duplicate the same
aggregation logic already used in RunnerService, so move empty_usage, add_usage,
and token_count into a shared usage-aggregation module or service and have both
LoopRunner and RunnerService call it. Keep the existing behavior for nil and
string-keyed usage hashes, and update the references in LoopRunner to use the
shared helper methods so the summing logic stays consistent in one place.
In `@lib/skill_bench/services/runner_service.rb`:
- Around line 177-220: The usage-normalization logic is duplicated across
RunnerService and other classes, so centralize it in one shared helper instead
of keeping separate copies of empty_usage, add_usage, and token_count. Extract a
common usage utility/module (for example, a SkillBench::Usage helper) and update
RunnerService’s aggregate_usage path, along with LoopRunner,
JsonFormatter::EMPTY_USAGE, and CostCalculator#token_count, to use the shared
implementation so the token-count rules stay consistent in one place.
🪄 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: Repository UI (base), Organization UI (inherited)
Review profile: CHILL
Plan: Pro
Run ID: 5346215a-99d9-4bff-b479-d04b268a81f7
📒 Files selected for processing (14)
lib/skill_bench/agent/react_agent/loop_runner.rblib/skill_bench/agent/react_agent/step.rblib/skill_bench/output_formatter.rblib/skill_bench/services/agent_spawner_service.rblib/skill_bench/services/cost_calculator.rblib/skill_bench/services/json_formatter.rblib/skill_bench/services/runner_service.rbtest/agent_eval/output_formatter_test.rbtest/agent_eval/services/agent_spawner_service_test.rbtest/agent_eval/services/runner_service_test.rbtest/evaluator/react_agent/loop_runner_test.rbtest/evaluator/react_agent/step_test.rbtest/services/cost_calculator_test.rbtest/services/json_formatter_test.rb
PRICES lacked the claude-sonnet-4 / claude-opus-4 families that the project actually passes in (config/defaults.rb -> claude-sonnet-4-20250514, clients/provider_schemas.rb -> claude-opus-4-7), so price_for fell through to nil and cost was reported unknown for the default Anthropic model. Add both families, preserving longest-prefix matching so dated variants still resolve. Covered by a new test asserting the real default strings.
Per-run token & cost accounting (Closes #40)
Problem
The client layer already returns token usage —
BaseClient#handle_responseincludesusage: extract_usage(parsed)on both the success and inline-error paths, andProviders::Anthropic#extract_usagenormalizes to{ prompt_tokens:, completion_tokens:, total_tokens: }— but that data was captured and then thrown away:agent/react_agent/step.rbread only:success/:error/:messageand droppedclient_result[:usage].agent/react_agent/loop_runner.rbnever accumulated usage.services/agent_spawner_service.rbhardcodedusage: {}on both its success and rescue paths.pricing|cost|pricefound nothing).What changed — usage threading (recovered)
The
usagehash thathandle_responsealready returns is now threaded all the way up:Step.callcapturesclient_result[:usage](defaulting to{}) and includes it as:usageon every returned step result — finish, continue-with-tools, client failure, and empty-response paths.LoopRunneraccumulates each step's usage into a running total and injects the summed{ prompt_tokens:, completion_tokens:, total_tokens: }into the finalresponse[:usage]. Missing/empty usage is treated as zero.AgentSpawnerServicenow reads the real aggregate fromagent_result.dig(:response, :usage)and surfaces it (success path); the mock and rescue paths return an explicit zeroedEMPTY_USAGE.RunnerServicesums the baseline + context agent usage and exposestokensandcostat the top level of the result (additive — no existing keys renamed).CostCalculator (new service)
services/cost_calculator.rb— a.callservice object that computes USD cost from a usage hash + model name:gpt-4o,gpt-4o-mini,gpt-4-turbo,gpt-4,gpt-3.5-turbo,claude-3-5-sonnet,claude-3-5-haiku,claude-3-opus,claude-3-sonnet,claude-3-haiku.claude-3-5-sonnet-20241022,gpt-4o-mini-2024-07-18) resolve correctly.nil(no crash). Toleratesnil/empty usage and string usage keys.PRICES.Where tokens/cost surface
OutputFormatter): aTokens: N | Est. Cost: $X.XXXXline in the delta-report header (—when cost is unknown).services/json_formatter.rb):tokensandcostare guaranteed present at the top level (additive defaults — zero usage /nullcost when absent) so consumers see a stable shape.Disjoint from #41
Per the hard constraint, no file under
lib/skill_bench/clients/was modified (includingbase_client.rb). This PR only consumes theusagehash thathandle_responsealready returns. Changes are confined tostep.rb,loop_runner.rb,agent_spawner_service.rb,runner_service.rb,output_formatter.rb,services/json_formatter.rb, and the newservices/cost_calculator.rb(+ tests).Judge usage — deferred
Judge-side usage is not included: reaching it would require touching the judge/
clients/boundary that PR #41 owns, so it stays out to keep the PRs disjoint. Aggregation is scoped to agent usage (baseline + context); judge-cost is a follow-up.Mock handling
The
mockprovider returns no real usage, so mock runs sum to zero tokens and an unknown (nil) cost — rendered as0/—. Verified by tests.Tests
cost_calculator_test.rb— known model, Anthropic model, dated-variant prefix match, longest-prefix preference, unknown/nil model → nil, empty/nil usage, string keys.Step→LoopRunner(summed across iterations, zero default) →AgentSpawnerService(real aggregate, zero for mock/rescue).tokens/cost(human + JSON), including the—/ zero mock case.RunnerServicesurfacestokens/cost(mock → zeros / nil).Gates (all green from repo root)
bundle exec rubocop— no offensesbundle exec reek— cleanbundle exec rake test— 782 runs, 0 failures, 0 errorsbundle exec rake yard:coverage— 0 undocumented public objectsCloses #40
Summary by CodeRabbit
New Features
Bug Fixes