From 62c4900ce7024e9bf94879c96ed9f53ee4f6763e Mon Sep 17 00:00:00 2001 From: Bingran You Date: Wed, 22 Jul 2026 01:37:07 -0700 Subject: [PATCH 1/2] fix(acp): preserve streamed usage on timeout --- src/benchflow/acp/session.py | 21 +++++-- src/benchflow/rollout/__init__.py | 1 + tests/test_native_acp_usage.py | 91 +++++++++++++++++++++++++++++++ uv.lock | 6 +- 4 files changed, 112 insertions(+), 7 deletions(-) diff --git a/src/benchflow/acp/session.py b/src/benchflow/acp/session.py index 269ef9bff..ffc3b6078 100644 --- a/src/benchflow/acp/session.py +++ b/src/benchflow/acp/session.py @@ -233,11 +233,21 @@ def record_agent_timeout( def record_prompt_usage(self, usage: object | None) -> None: """Record cumulative ACP token usage returned by session/prompt.""" + if self._record_usage_snapshot(usage): + self._notify_change() + + def _record_usage_snapshot(self, usage: object | None) -> bool: snapshot = normalize_acp_usage(usage) if snapshot is None: - return + return False self.usage_snapshots.append(snapshot) - self._notify_change() + return True + + def _record_update_usage(self, update: dict) -> bool: + snapshot_source = update.get("usage") + if snapshot_source is not None and self._record_usage_snapshot(snapshot_source): + return True + return self._record_usage_snapshot(update) def latest_usage_totals(self) -> ACPUsageSnapshot | None: """Return the latest cumulative ACP usage snapshot, if any.""" @@ -274,12 +284,15 @@ def handle_update(self, update: dict) -> None: """Process a session/update notification.""" self._events_active = True update_type = update.get("sessionUpdate") + usage_recorded = self._record_update_usage(update) # Unknown update types (future ACP versions, agent-specific # extensions) mutate no state and must not trigger a no-op # snapshot. Mark events_active so the snapshot path stays on - # the modern branch, but skip _notify_change for unrecognized - # types. + # the modern branch, but only notify for unrecognized types when they + # carried usage telemetry. if update_type not in self._RECOGNIZED_UPDATE_TYPES: + if usage_recorded: + self._notify_change() return if update_type == "tool_call": diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 63909c2dc..b93cd988b 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -1886,6 +1886,7 @@ async def cleanup(self) -> None: finally: self._usage_runtime = None + self._collect_native_acp_usage() self._finalize_usage_metrics() self._enforce_required_usage_tracking() diff --git a/tests/test_native_acp_usage.py b/tests/test_native_acp_usage.py index ebfea42f2..9785b2d2c 100644 --- a/tests/test_native_acp_usage.py +++ b/tests/test_native_acp_usage.py @@ -93,6 +93,97 @@ def test_rollout_native_acp_usage_uses_cumulative_deltas(): } +def test_acp_session_records_usage_from_session_update(): + """Guards #933: ACP timeout accounting can use pre-cancel update usage.""" + from benchflow.acp.session import ACPSession + + session = ACPSession("session-1") + + session.handle_update( + { + "sessionUpdate": "agent_message_chunk", + "content": {"type": "text", "text": "working"}, + "usage": { + "inputTokens": 10, + "outputTokens": 4, + "totalTokens": 16, + "cachedReadTokens": 2, + "cachedWriteTokens": 1, + "thoughtTokens": 1, + }, + } + ) + + assert session.latest_usage_totals() == { + "input_tokens": 10, + "output_tokens": 4, + "total_tokens": 16, + "cached_read_tokens": 2, + "cached_write_tokens": 1, + "thought_tokens": 1, + } + + +@pytest.mark.asyncio +async def test_cleanup_collects_update_usage_after_timed_out_prompt(tmp_path): + """Guards #933: timed-out ACP prompts do not lose streamed token usage.""" + from benchflow.acp.session import ACPSession + from benchflow.rollout import Rollout, RolloutConfig + from benchflow.rollout._usage import _zero_native_acp_usage_metrics + + session = ACPSession("session-1") + session.handle_update( + { + "sessionUpdate": "usage", + "usage": { + "inputTokens": 20, + "outputTokens": 7, + "totalTokens": 30, + "cachedReadTokens": 1, + "cachedWriteTokens": 2, + "thoughtTokens": 3, + }, + } + ) + + rollout = Rollout.__new__(Rollout) + rollout._session = session + rollout._acp_client = SimpleNamespace(session=session) + rollout._trajectory = [] + rollout._session_traj_count = 0 + rollout._session_tool_count = 0 + rollout._terminal_timeout = True + rollout._partial_trajectory = False + rollout._trajectory_source = "none" + rollout._n_tool_calls = 0 + rollout._usage_runtime = None + rollout._usage_metrics = {"usage_source": "unavailable"} + rollout._native_usage_checkpoint = None + rollout._native_usage_metrics = _zero_native_acp_usage_metrics() + rollout._env = None + rollout._environment = None + rollout._config = RolloutConfig(task_path=tmp_path) + + async def fake_disconnect(): + return None + + rollout.disconnect = fake_disconnect + + await rollout.cleanup() + + assert rollout._usage_metrics == { + "n_input_tokens": 20, + "n_output_tokens": 7, + "n_cache_read_tokens": 1, + "n_cache_creation_tokens": 2, + "total_tokens": 30, + "cost_usd": None, + "usage_source": "agent_native_acp", + "price_source": None, + "usage_details": {"thought_tokens": 3}, + } + + def test_rollout_provider_usage_wins_over_native_acp_usage(): """Guards PR #613 follow-up: LiteLLM provider telemetry remains authoritative.""" from benchflow.rollout import Rollout diff --git a/uv.lock b/uv.lock index e55180107..dec20d7f7 100644 --- a/uv.lock +++ b/uv.lock @@ -1982,7 +1982,7 @@ wheels = [ [[package]] name = "mcp" -version = "1.27.2" +version = "1.28.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -2000,9 +2000,9 @@ dependencies = [ { name = "typing-inspection" }, { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/27/3c/347cf965d313f5d41764e7d46bea6ffe7d9ef13b983cc429b0340962a082/mcp-1.27.2.tar.gz", hash = "sha256:8e02db104096d1c25b28e64bde29a5c32b31bc241710213e12fd4d84985bdfef", size = 621116, upload-time = "2026-05-29T17:16:04.039Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6e/77/9450b8f251a13affb6281997d0523c4615f8a8b35d0b21ff30db3a5aac9d/mcp-1.28.1.tar.gz", hash = "sha256:d51e36a5f5644faea4f85ea649bfffa6bc6c26770d42798ad6a3de3d2ba69683", size = 638501, upload-time = "2026-06-26T12:57:29.093Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c9/11/252c6f971dc4f16af1d98a1c469d8ba523aab00d1bb76b4d3bc1ff32eacc/mcp-1.27.2-py3-none-any.whl", hash = "sha256:d6ff5160c6ca65d93013626efb3fc249de683c30b2d8570755ceddd490344de5", size = 220498, upload-time = "2026-05-29T17:16:02.442Z" }, + { url = "https://files.pythonhosted.org/packages/e2/5e/d118fce19f87a2e7d8101c35c8ae0ec289098a4df0ff244cec23e415aca0/mcp-1.28.1-py3-none-any.whl", hash = "sha256:2726bca5e7193f61c5dde8b12500a6de2d9acf6d1a1c0be9e8c2e706437991df", size = 222620, upload-time = "2026-06-26T12:57:27.218Z" }, ] [[package]] From 3bacd9e42cb92de585ffaa35ffa6c9db281a4e7e Mon Sep 17 00:00:00 2001 From: Bingran You Date: Wed, 22 Jul 2026 06:05:24 -0700 Subject: [PATCH 2/2] fix(rollout): collect ACP usage before disconnect --- src/benchflow/rollout/__init__.py | 2 +- tests/test_native_acp_usage.py | 12 +++++------- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index b93cd988b..e1e5b9bdd 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -1828,6 +1828,7 @@ async def soft_verify(self) -> tuple[dict | None, str | None, str | None]: async def cleanup(self) -> None: """Close ACP client and stop the environment.""" self._capture_partial_acp_trajectory() + self._collect_native_acp_usage() await self.disconnect() if self._env and self._config.export_generated_skills_to: @@ -1886,7 +1887,6 @@ async def cleanup(self) -> None: finally: self._usage_runtime = None - self._collect_native_acp_usage() self._finalize_usage_metrics() self._enforce_required_usage_tracking() diff --git a/tests/test_native_acp_usage.py b/tests/test_native_acp_usage.py index 9785b2d2c..49105cb31 100644 --- a/tests/test_native_acp_usage.py +++ b/tests/test_native_acp_usage.py @@ -94,7 +94,7 @@ def test_rollout_native_acp_usage_uses_cumulative_deltas(): def test_acp_session_records_usage_from_session_update(): - """Guards #933: ACP timeout accounting can use pre-cancel update usage.""" + """Guards PR #934 / issue #933: timeout accounting uses update usage.""" from benchflow.acp.session import ACPSession session = ACPSession("session-1") @@ -126,7 +126,7 @@ def test_acp_session_records_usage_from_session_update(): @pytest.mark.asyncio async def test_cleanup_collects_update_usage_after_timed_out_prompt(tmp_path): - """Guards #933: timed-out ACP prompts do not lose streamed token usage.""" + """Guards PR #934 / issue #933: cleanup preserves streamed usage.""" from benchflow.acp.session import ACPSession from benchflow.rollout import Rollout, RolloutConfig from benchflow.rollout._usage import _zero_native_acp_usage_metrics @@ -160,17 +160,15 @@ async def test_cleanup_collects_update_usage_after_timed_out_prompt(tmp_path): rollout._usage_metrics = {"usage_source": "unavailable"} rollout._native_usage_checkpoint = None rollout._native_usage_metrics = _zero_native_acp_usage_metrics() + rollout._agent_launch = "" rollout._env = None rollout._environment = None + rollout._error = None rollout._config = RolloutConfig(task_path=tmp_path) - async def fake_disconnect(): - return None - - rollout.disconnect = fake_disconnect - await rollout.cleanup() + assert rollout._session is None assert rollout._usage_metrics == { "n_input_tokens": 20, "n_output_tokens": 7,