fix: settle usage on the timeout path, not only after DONE - #4553
Closed
atishpatel wants to merge 1 commit into
Closed
fix: settle usage on the timeout path, not only after DONE#4553atishpatel wants to merge 1 commit into
atishpatel wants to merge 1 commit into
Conversation
`_settle_usage` waited for the agent to report its tokens before teardown killed it, but only on the success path. The timeout path fell straight through to the kill, on the stated reasoning that "a turn that never completed has no usage to flush". That was true while buzz-agent reported once per turn. It is not true now that it reports after every provider round, and it was always the wrong path to skip: under `continue_until_timeout` every phase but the last ends here, so the skip applied to essentially the whole run. It is the harness half of the 97%-zero-rows / 5.9x-undercount measured on the OpenRouter LHTB-46 cell; the agent half is the per-round reporting. Moving the call into the `finally` covers both paths and keeps the ordering that matters — settle, then kill. It stays cheap: it returns on the first poll once a usage line exists, which after the first round it does, so only a phase that never completed a single round can spend the full budget, and that phase has nothing to report anyway. Depends on the per-round reporting in buzz-agent to be worth much on its own. Signed-off-by: Atish Patel <atish@squareup.com> Co-authored-by: Claude Code <noreply@anthropic.com>
atishpatel
added a commit
that referenced
this pull request
Aug 3, 2026
## The bug buzz-agent emitted its `usage_update` notification in exactly one place: after `ctx.run()` returned. Until that moment a turn's token counters lived only in the prompt task's stack frame. **A turn killed mid-flight reported nothing at all** — the provider had already billed every round it completed, and no consumer ever saw any of it. That is not a corner case for anything that ends a turn on a clock. It is the normal case for a long-horizon benchmark run that relaunches its agent between phases. ## How big Measured against a provider's own billing ledger over one run's window: | | provider ledger | what we recorded | |---|---|---| | the relaunched lead seat | $485 / 348M tok | $98.99 / 90.3M tok | | the two seats that were not relaunched | $29.90 / 856M | $25.81 / 765M — reconciles | 97% of that run's usage rows came back all zeros, against 1–4% for comparable runs that never relaunch. In one 450-phase trial exactly 7 phases recorded any usage — and each of those carries 177k–437k input tokens, a whole session's worth landing in the one phase that happened to end gracefully. Worth being precise about what was *not* wrong, since both were plausible and both were checked: - **Not pricing.** The rates were verified against the provider's endpoints API and match what we charge. - **Not a truncation bug.** The usage files were intact and internally consistent. The tokens were never captured in the first place. ## The fix The run loop now emits a session-cumulative `usage_update` after every usage-bearing provider response, so an interrupted turn has reported everything but its single in-flight request. - **Emitting more than once per turn is already part of the contract.** buzz-acp's `UsageTracker` advances its committed baseline only at publish time, and goose behaves the same way — which is why the tracker was written to tolerate it. - **The turn-start session baseline is snapshotted into `RunCtx`** so the mid-turn figure stays *session*-cumulative. A turn-local number would be discarded by a high-water-mark consumer and lose the turn entirely; there is a test for exactly that. - **Snapshot by value, not a session handle.** The loop reports once per round, and taking the sessions lock on each would serialise concurrent sessions behind one another's provider round-trips. Nothing else advances those counters while the turn holds `busy`, so it cannot go stale. - **One shared `wire::usage_update_payload`** for both call sites, so the mid-turn and end-of-turn shapes cannot drift. A drift there would present as tokens silently vanishing, which is the failure this reporting exists to prevent. ## Why not a SIGTERM handler That was the obvious shape and it does not work. At signal time the counters are not sitting anywhere a handler could reach — they are in the turn's stack frame, and the value the handler would need has not been folded into the session yet. Making usage durable *during* the turn is what actually fixes it; once it is, a handler adds nothing beyond the in-flight request, whose cost is unknown until its response lands. ## Tests - `usage_is_reported_after_each_round_not_only_at_turn_end` — two rounds; asserts the **first** notification carries round 1's counts alone, proving it went out before round 2 returned. - `mid_turn_usage_includes_earlier_turns` — a mid-turn report must be session-cumulative, not turn-local. buzz-agent 18/18 on the `fake_llm` suite, 382 unit. `cargo fmt` / `clippy` / `cargo check --workspace --all-targets` clean. ## Scope Agent-side only, against `main`. The matching harness change — settling usage on the timeout path, which was skipped on the reasoning that an incomplete turn has nothing to flush — is **#4553**, against the benchmark branch, since that harness does not exist on `main`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: Atish Patel <atish@squareup.com> Co-authored-by: Claude Code <noreply@anthropic.com>
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.
The harness half of the usage-accounting fix. The agent half is #4545, which targets
main.What was wrong
_settle_usagewaits for the agent to report its tokens before teardown kills it — but it was called only on the success path, afterDONE:. The timeout path fell straight through to the kill, on the stated reasoning that "a turn that never completed has no usage to flush".That reasoning held while buzz-agent reported once per turn. #4545 changes that: it now reports after every provider round, so an interrupted turn has reported. And it was the wrong path to skip regardless — under
continue_until_timeoutevery phase but the last ends on the timeout path, so the skip covered essentially the entire run.Why it matters
Measured on the OpenRouter LHTB-46 cell (Kimi K3 lead + 2x DeepSeek V4 Flash):
In one 450-phase trial exactly 7 phases recorded any usage at all — and each of those carries 177k–437k input tokens, i.e. a whole session's worth landing in the one phase that happened to end gracefully.
The change
Move the call into the
finallyso both paths settle, keeping the ordering that matters: settle, then kill.It stays cheap.
_settle_usagereturns on the first poll once a usage line exists, which after the first round it does, so only a phase that never completed a single round can spend the full 60s budget — and that phase has nothing to report anyway.Also corrects the
_settle_usagedocstring, which described the old once-per-turn emission and cited a line number that has moved.Test
test_usage_is_settled_when_the_trial_times_out— drivesrun()with a_wait_for_donethat never returns and a zero trial budget, then asserts usage was settled and that it happened before teardown killed the agent. Settling after the kill is the same as not settling, so the ordering is the assertion, not the call count.272 pass.
Ordering
Worth little without #4545 — on its own it settles a turn that still only reports at the end. Land or merge the agent side alongside it.
This does not recover the already-collected runs. Every LHTB cost recorded so far stays a floor.
🤖 Generated with Claude Code