Skip to content

Fix token usage accounting across tool loops and on errors - #35

Open
symphony-stream wants to merge 1 commit into
adham90:mainfrom
symphony-stream:fix/usage-accounting-tool-loops
Open

Fix token usage accounting across tool loops and on errors#35
symphony-stream wants to merge 1 commit into
adham90:mainfrom
symphony-stream:fix/usage-accounting-tool-loops

Conversation

@symphony-stream

Copy link
Copy Markdown
Contributor

Defect

Token usage accounting in BaseAgent is wrong in two related ways whenever an agent run involves more than a single provider request (a tool loop), and whenever it fails:

  1. Undercounting on success. capture_response records usage from the final response only. A tool loop bills one request per round, so a run with N rounds records roughly 1/N of the real spend in input_tokens / output_tokens / costs on the execution row.

  2. Total loss on failure. Usage is only captured after execute_llm_call returns. If ask raises mid-loop (e.g. a tool error, a provider error on round 3, or a timeout), capture_response never runs and the execution persists 0 tokens / $0 — even though every completed round was billed by the provider.

Reproduction

class LoopAgent < RubyLLM::Agents::BaseAgent
  model "gpt-4o"
  def user_prompt = "do the thing"
end

agent = LoopAgent.new
context = RubyLLM::Agents::Pipeline::Context.new(input: {}, agent_class: LoopAgent, agent_instance: agent)

# A chat client whose ask appends two billable assistant responses
# (round 1 + round 2 of a tool loop) and then raises from the tool.
client = <chat double; messages grows by 2 assistant messages with usage, ask raises>
agent.instance_variable_set(:@client, client)

agent.send(:execute_llm_call, client, context) # raises
context.input_tokens  # => 0 (expected: sum of both rounds)
context.total_cost    # => 0.0

The success path has the mirror-image problem: with rounds of 1000+2000 input tokens, context.input_tokens ends up 2000 instead of 3000.

Fix

One seam, in BaseAgent:

  • execute_llm_call baselines the client history size at call start. Conversation replay (the messages option) seeds prior turns into the client, so only assistant messages appended after the baseline are this attempt's billable usage.
  • On exception, recover_usage_from_history runs capture_response on the last billed response (the same pattern Tool::Halt recovery uses) before re-raising — the original error always propagates, and recovery itself is best-effort.
  • capture_response now calls accumulate_attempt_usage, which sums input/output tokens across the attempt's appended assistant responses, stashes the last request's own usage in context[:llm_usage] metadata (requests, last_input_tokens, last_output_tokens — the prompt size of the final request is the current context fill), and recalculates costs from the sums.

Retries/fallbacks re-baseline per attempt (a fresh client per execute), so persisted numbers always describe the last attempt — the billed one — never a cross-attempt double count. Single-request runs are unchanged: the sum of one appended assistant message equals the response's own usage, and when the history is unreadable the pre-existing single-response numbers stand.

Tests

New spec/lib/base_agent_usage_spec.rb (10 examples):

  • sums usage across the tool loop when ask raises mid-loop, including model id and llm_usage metadata
  • prices the summed usage through calculate_costs
  • recovers usage on cooperative cancellation (CancelledError) inside the loop
  • ignores assistant history seeded before the call (transcript replay)
  • leaves context untouched when the failure precedes any completed response
  • never masks the original error when recovery itself fails
  • sums usage on the successful path
  • keeps single-response numbers when the history is unavailable
  • re-baselines per attempt (retry overwrites instead of double counting)
  • sums through the full #execute path

Full suite: 4906 examples, 0 failures (Ruby 3.3.6). standardrb clean on changed files.

capture_response records usage from the final response only, and only
when execute_llm_call succeeds. A multi-round tool loop undercounts the
real spend (one request is billed per round), and an exception mid-loop
discards the usage entirely — the execution persists 0 tokens / $0.

Baseline the client history when the call starts, then sum usage across
the assistant responses appended during the attempt — after a successful
capture_response, and from the rescue when the call raises. Retries and
fallbacks re-baseline per attempt, so persisted numbers describe the
last (billed) attempt and never double count. The last request's own
usage is stashed in context[:llm_usage] metadata.
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.

1 participant