From b915d56734e4c6e90c0e634e56b5f74202cd4214 Mon Sep 17 00:00:00 2001 From: Atish Patel Date: Mon, 3 Aug 2026 11:02:22 -0500 Subject: [PATCH] fix: settle usage on the timeout path, not only after DONE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_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 Co-authored-by: Claude Code --- .../container_runtime.py | 40 +++++++++---- .../tests/test_container_runtime.py | 57 +++++++++++++++++++ 2 files changed, 85 insertions(+), 12 deletions(-) diff --git a/benchmarks/harbor-buzz-orchestra/src/harbor_buzz_orchestra/container_runtime.py b/benchmarks/harbor-buzz-orchestra/src/harbor_buzz_orchestra/container_runtime.py index 4d2c19114e..43dd675cf4 100644 --- a/benchmarks/harbor-buzz-orchestra/src/harbor_buzz_orchestra/container_runtime.py +++ b/benchmarks/harbor-buzz-orchestra/src/harbor_buzz_orchestra/container_runtime.py @@ -381,12 +381,26 @@ async def run( ), timeout=manifest.trial_budget.timeout_seconds, ) - # `DONE:` is published before the turn's usage notification is, so - # teardown has to wait or the trial's tokens are lost. Inside the - # try, and only on this path: a timeout has no usage to flush. - await self._settle_usage(environment, agents) await self._verify_m1_output(environment, manifest) finally: + # `DONE:` is published before the turn's usage notification is, so + # teardown has to wait or the trial's tokens are lost. + # + # In the `finally`, so the TIMEOUT path settles too. That path used + # to fall straight through to the kill on the reasoning that "a turn + # that never completed has no usage to flush", which was true while + # buzz-agent reported once per turn and is not any more: it now + # reports after every provider round, so an unfinished turn has + # reported everything but its in-flight request. Skipping the settle + # here is what made `continue_until_timeout` runs uncostable — + # every phase but the last ends on this path, and 97% of one + # measured run's receipt rows came back all zeros. + # + # Cheap in the common case: this returns on the first poll once a + # usage line exists, which after the first round it does. Only a + # phase that never completed a single round can spend the full + # budget, and that phase has nothing to report anyway. + await self._settle_usage(environment, agents) await self._stop_agents(environment, agents + infra) # Logs first, and ahead of anything that touches the network. The # verifier pre-install below used to run first and, when the proxy @@ -1021,10 +1035,10 @@ async def _settle_usage( ) -> None: """Give each agent the moment it needs to report what it spent. - buzz-agent emits its `_goose/unstable/session/update` usage notification - once per turn, immediately *before* returning the `session/prompt` - response (buzz-agent/src/lib.rs:708). A solo agent gets exactly one turn - per trial, so that single notification is the only record of the trial's + buzz-agent emits a `_goose/unstable/session/update` usage notification + after every provider round, and once more immediately *before* returning + the `session/prompt` response. A solo agent gets exactly one turn per + trial, so the final notification is the complete record of the trial's tokens — and it is written after the agent has already published `DONE:` as a tool call. @@ -1039,10 +1053,12 @@ async def _settle_usage( `DONE:` reaches the channel from inside a tool call, so the turn may still owe one model round-trip before it ends and reports, and for a thinking model that is tens of seconds, not milliseconds. - Waiting is pointless on the timeout path — a turn that never completed - has no usage to flush — so callers only invoke this once `DONE:` is seen. - A miss is not fatal; the accounting note already reports an unpriced - trial, and losing the tokens is better than hanging the sweep. + Called on the timeout path too, not just after `DONE:`. An interrupted + turn has still reported every round it finished, and under + `continue_until_timeout` the interrupted turn is the normal case, not + the exception. A miss is not fatal; the accounting note already reports + an unpriced trial, and losing the tokens is better than hanging the + sweep. """ deadline = asyncio.get_running_loop().time() + self.usage_settle_seconds pending = {agent.credential.agent_id: agent for agent in agents} diff --git a/benchmarks/harbor-buzz-orchestra/tests/test_container_runtime.py b/benchmarks/harbor-buzz-orchestra/tests/test_container_runtime.py index b7a42c542b..1f1b0717d3 100644 --- a/benchmarks/harbor-buzz-orchestra/tests/test_container_runtime.py +++ b/benchmarks/harbor-buzz-orchestra/tests/test_container_runtime.py @@ -1,5 +1,6 @@ """The container runtime must launch the production stack, unmodified.""" +import asyncio import hashlib import json import re @@ -783,6 +784,62 @@ async def done(*args, **kwargs): assert "kill" in order +async def test_usage_is_settled_when_the_trial_times_out(tmp_path, monkeypatch): + """The timeout path has to settle too, and it is the path that matters most. + + It used to fall straight through to the kill, on the reasoning that a turn + which never completed had nothing to flush. buzz-agent now reports after + every provider round, so an interrupted turn HAS reported — and under + `continue_until_timeout` every phase but the last ends here. Skipping the + settle on this path is what left 97% of one measured run's receipt rows at + all zeros while the provider billed it in full. + """ + manifest = write_manifest(tmp_path) + trial = trial_handle((credential("orch-1", "orchestrator", "orch-model", "lead"),)) + rt = runtime(tmp_path, poll_seconds=0, usage_settle_seconds=5) + order = [] + + class TimingOutEnvironment(Environment): + async def exec(self, command, env=None, **kwargs): + if "buzz-acp" in command: + return ExecResult(stdout="99\n", stderr="", return_code=0) + if command.startswith("cat "): + order.append("usage") + return ExecResult( + stdout="goose usage update input=1 output=2\n", + stderr="", + return_code=0, + ) + if "/proc/[0-9]*" in command: + order.append("kill") + return ExecResult(stdout="", stderr="", return_code=0) + + async def never_done(*args, **kwargs): + await asyncio.sleep(3600) + + monkeypatch.setattr(rt, "_install_stack", lambda env: _noop()) + monkeypatch.setattr(rt, "_wait_for_agents_ready", lambda *a, **k: _noop()) + monkeypatch.setattr(rt, "_wait_for_done", never_done) + monkeypatch.setattr(rt, "_buzz_json", lambda *a, **k: _value([])) + # TrialBudget is frozen, so swap in a zero-budget copy rather than mutating. + budget = manifest.trial_budget.model_copy(update={"timeout_seconds": 0}) + manifest = manifest.model_copy(update={"trial_budget": budget}) + + with pytest.raises(asyncio.TimeoutError): + await rt.run( + instruction="do the thing", + environment=TimingOutEnvironment(), + manifest=manifest, + trial=trial, + ) + + assert "usage" in order, "the timeout path never settled usage" + assert order.index("usage") < order.index("kill"), ( + "usage was settled after teardown killed the agent, which is the same " + "as not settling at all" + ) + + async def test_wait_for_agents_ready_requires_every_channel_subscription(tmp_path): rt = runtime(tmp_path, poll_seconds=0) logs = {"orch-1": "", "worker-1": ""}