Fix token usage accounting across tool loops and on errors - #35
Open
symphony-stream wants to merge 1 commit into
Open
Fix token usage accounting across tool loops and on errors#35symphony-stream wants to merge 1 commit into
symphony-stream wants to merge 1 commit into
Conversation
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.
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
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.
Defect
Token usage accounting in
BaseAgentis wrong in two related ways whenever an agent run involves more than a single provider request (a tool loop), and whenever it fails:Undercounting on success.
capture_responserecords 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 ininput_tokens/output_tokens/ costs on the execution row.Total loss on failure. Usage is only captured after
execute_llm_callreturns. Ifaskraises mid-loop (e.g. a tool error, a provider error on round 3, or a timeout),capture_responsenever runs and the execution persists 0 tokens / $0 — even though every completed round was billed by the provider.Reproduction
The success path has the mirror-image problem: with rounds of 1000+2000 input tokens,
context.input_tokensends up2000instead of3000.Fix
One seam, in
BaseAgent:execute_llm_callbaselines the client history size at call start. Conversation replay (themessagesoption) seeds prior turns into the client, so only assistant messages appended after the baseline are this attempt's billable usage.recover_usage_from_historyrunscapture_responseon the last billed response (the same patternTool::Haltrecovery uses) before re-raising — the original error always propagates, and recovery itself is best-effort.capture_responsenow callsaccumulate_attempt_usage, which sums input/output tokens across the attempt's appended assistant responses, stashes the last request's own usage incontext[: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):askraises mid-loop, including model id andllm_usagemetadatacalculate_costsCancelledError) inside the loop#executepathFull suite: 4906 examples, 0 failures (Ruby 3.3.6).
standardrbclean on changed files.