Skip to content

feat: per-run token and cost accounting (#40)#77

Merged
igmarin merged 3 commits into
mainfrom
feat/token-cost-accounting-40
Jun 30, 2026
Merged

feat: per-run token and cost accounting (#40)#77
igmarin merged 3 commits into
mainfrom
feat/token-cost-accounting-40

Conversation

@igmarin

@igmarin igmarin commented Jun 30, 2026

Copy link
Copy Markdown
Owner

Per-run token & cost accounting (Closes #40)

Problem

The client layer already returns token usage — BaseClient#handle_response includes usage: extract_usage(parsed) on both the success and inline-error paths, and Providers::Anthropic#extract_usage normalizes to { prompt_tokens:, completion_tokens:, total_tokens: } — but that data was captured and then thrown away:

  • agent/react_agent/step.rb read only :success/:error/:message and dropped client_result[:usage].
  • agent/react_agent/loop_runner.rb never accumulated usage.
  • services/agent_spawner_service.rb hardcoded usage: {} on both its success and rescue paths.
  • There was no cost model at all (a grep for pricing|cost|price found nothing).

What changed — usage threading (recovered)

The usage hash that handle_response already returns is now threaded all the way up:

  1. Step.call captures client_result[:usage] (defaulting to {}) and includes it as :usage on every returned step result — finish, continue-with-tools, client failure, and empty-response paths.
  2. LoopRunner accumulates each step's usage into a running total and injects the summed { prompt_tokens:, completion_tokens:, total_tokens: } into the final response[:usage]. Missing/empty usage is treated as zero.
  3. AgentSpawnerService now reads the real aggregate from agent_result.dig(:response, :usage) and surfaces it (success path); the mock and rescue paths return an explicit zeroed EMPTY_USAGE.
  4. RunnerService sums the baseline + context agent usage and exposes tokens and cost at the top level of the result (additive — no existing keys renamed).

CostCalculator (new service)

services/cost_calculator.rb — a .call service object that computes USD cost from a usage hash + model name:

  • Small, clearly-approximate per-model price table (USD per 1K input / output tokens), sourced from public OpenAI/Anthropic pricing: 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.
  • Matches by exact name first, then by longest name prefix, so dated variants (e.g. claude-3-5-sonnet-20241022, gpt-4o-mini-2024-07-18) resolve correctly.
  • Unknown model → nil (no crash). Tolerates nil/empty usage and string usage keys.
  • Easy to extend — just add a row to PRICES.

Where tokens/cost surface

  • Human output (OutputFormatter): a Tokens: N | Est. Cost: $X.XXXX line in the delta-report header ( when cost is unknown).
  • JSON output (services/json_formatter.rb): tokens and cost are guaranteed present at the top level (additive defaults — zero usage / null cost when absent) so consumers see a stable shape.

Disjoint from #41

Per the hard constraint, no file under lib/skill_bench/clients/ was modified (including base_client.rb). This PR only consumes the usage hash that handle_response already returns. Changes are confined to step.rb, loop_runner.rb, agent_spawner_service.rb, runner_service.rb, output_formatter.rb, services/json_formatter.rb, and the new services/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 mock provider returns no real usage, so mock runs sum to zero tokens and an unknown (nil) cost — rendered as 0 / . 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.
  • Usage threading: StepLoopRunner (summed across iterations, zero default) → AgentSpawnerService (real aggregate, zero for mock/rescue).
  • Formatters render tokens/cost (human + JSON), including the / zero mock case.
  • RunnerService surfaces tokens/cost (mock → zeros / nil).

Gates (all green from repo root)

  • bundle exec rubocop — no offenses
  • bundle exec reek — clean
  • bundle exec rake test — 782 runs, 0 failures, 0 errors
  • bundle exec rake yard:coverage — 0 undocumented public objects

Closes #40

Summary by CodeRabbit

  • New Features

    • Added token-usage tracking across agent runs and included usage totals in results.
    • Added estimated cost reporting based on token usage and model selection.
    • JSON and human-readable outputs now show token totals and cost information.
  • Bug Fixes

    • Improved handling of missing or differently formatted token data.
    • Defaulted usage values to zero when token data is unavailable.
    • Ensured error and fallback paths still return consistent usage details.

@coderabbitai

coderabbitai Bot commented Jun 30, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Warning

Review limit reached

@igmarin, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 13 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

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 configuration

Configuration used: Repository UI (base), Organization UI (inherited)

Review profile: CHILL

Plan: Pro

Run ID: f56b5b5b-6f47-4d24-99c2-e51217050ff2

📥 Commits

Reviewing files that changed from the base of the PR and between 79b1001 and d03859b.

📒 Files selected for processing (2)
  • lib/skill_bench/services/cost_calculator.rb
  • test/services/cost_calculator_test.rb
📝 Walkthrough

Walkthrough

Token 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.

Changes

Token usage and cost pipeline

Layer / File(s) Summary
Step captures usage from client result
lib/skill_bench/agent/react_agent/step.rb, test/evaluator/react_agent/step_test.rb
Step.call extracts usage from the client result and adds it to every return path (failure, empty response, no tool calls, continue), with doc and tests updated.
LoopRunner accumulates usage across iterations
lib/skill_bench/agent/react_agent/loop_runner.rb, test/evaluator/react_agent/loop_runner_test.rb
LoopRunner.call initializes and updates a running total_usage, replaces merge_iterations with finalize (injecting iterations and usage), and adds empty_usage/add_usage/token_count helpers.
AgentSpawnerService surfaces aggregated usage
lib/skill_bench/services/agent_spawner_service.rb, test/agent_eval/services/agent_spawner_service_test.rb
Adds EMPTY_USAGE constant and returns real usage data for mock, success, and error paths instead of hardcoded empty hashes.
CostCalculator service for model pricing
lib/skill_bench/services/cost_calculator.rb, test/services/cost_calculator_test.rb
New CostCalculator class with PRICES/TOKENS_PER_UNIT constants computes estimated cost from usage and model via longest-prefix matching, returning nil for unknown models.
RunnerService aggregates usage and computes cost
lib/skill_bench/services/runner_service.rb, test/agent_eval/services/runner_service_test.rb
Requires CostCalculator, aggregates baseline/context usage, resolves pricing model, and includes tokens/cost in the success response.
JSON and human formatters render tokens and cost
lib/skill_bench/services/json_formatter.rb, lib/skill_bench/output_formatter.rb, test/services/json_formatter_test.rb, test/agent_eval/output_formatter_test.rb
JsonFormatter ensures tokens/cost fields always exist; OutputFormatter adds a "Tokens / Est. Cost" line to human delta reports.

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)
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

  • igmarin/ruby-skill-bench#13: Both PRs modify LoopRunner.call/Step.call return payloads in the ReAct agent loop, with this PR extending that flow with usage accumulation.
  • igmarin/ruby-skill-bench#9: Both PRs modify runner_service.rb to thread usage/result metadata through the agent-eval pipeline.
  • igmarin/ruby-skill-bench#8: Both PRs modify output_formatter.rb's delta-report human rendering logic.

Poem

Tokens hopped in, one by one,
Summed up neatly till the run was done. 🐇
A calculator priced each byte,
Dollars and cents now shown just right.
No more usage lost in the burrow's floor—
Just clean little numbers, forevermore. ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly summarizes the main change: per-run token and cost accounting.
Linked Issues check ✅ Passed The PR threads usage through step, loop, and spawner, adds CostCalculator, and surfaces tokens/cost in output and JSON as requested.
Out of Scope Changes check ✅ Passed All modified files are directly related to token/cost accounting and the associated formatter and test updates.

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (2)
lib/skill_bench/agent/react_agent/loop_runner.rb (1)

72-100: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Duplicate usage-aggregation logic across services.

empty_usage, add_usage, and token_count here mirror equivalent helpers in RunnerService (per cross-file graph evidence showing add_usage/token_count calls at runner_service.rb lines ~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 with empty_usage/add_usage/token_count that both LoopRunner and RunnerService call.

🤖 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 win

Duplicate usage-normalization helpers across files.

empty_usage, add_usage, and token_count here reimplement the exact same logic already defined in LoopRunner (per the cross-file graph evidence showing identical method bodies). JsonFormatter::EMPTY_USAGE independently re-declares the same zeroed-usage shape, and CostCalculator#token_count has 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::Usage module with .empty, .add, .token_count) used by LoopRunner, RunnerService, JsonFormatter, and CostCalculator.

🤖 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

📥 Commits

Reviewing files that changed from the base of the PR and between 8011f03 and 79b1001.

📒 Files selected for processing (14)
  • lib/skill_bench/agent/react_agent/loop_runner.rb
  • lib/skill_bench/agent/react_agent/step.rb
  • lib/skill_bench/output_formatter.rb
  • lib/skill_bench/services/agent_spawner_service.rb
  • lib/skill_bench/services/cost_calculator.rb
  • lib/skill_bench/services/json_formatter.rb
  • lib/skill_bench/services/runner_service.rb
  • test/agent_eval/output_formatter_test.rb
  • test/agent_eval/services/agent_spawner_service_test.rb
  • test/agent_eval/services/runner_service_test.rb
  • test/evaluator/react_agent/loop_runner_test.rb
  • test/evaluator/react_agent/step_test.rb
  • test/services/cost_calculator_test.rb
  • test/services/json_formatter_test.rb

Comment thread lib/skill_bench/services/cost_calculator.rb
igmarin added 2 commits June 30, 2026 14:53
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.
@igmarin
igmarin merged commit f388709 into main Jun 30, 2026
5 checks passed
@igmarin
igmarin deleted the feat/token-cost-accounting-40 branch June 30, 2026 20:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Token & cost accounting end-to-end (usage captured at client layer then dropped)

1 participant