From 3225a21cba2fca7ae863e88192242a11cdfbafb3 Mon Sep 17 00:00:00 2001 From: kywch Date: Sun, 30 Aug 2026 21:56:35 -0700 Subject: [PATCH 1/4] fix(acp): preserve timeout usage evidence --- src/benchflow/acp/runtime.py | 91 +++++++++--- src/benchflow/rollout/__init__.py | 3 + tests/test_acp.py | 220 ++++++++++++++++++++++++++++-- tests/test_native_acp_usage.py | 33 +++++ 4 files changed, 315 insertions(+), 32 deletions(-) diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index 770262de1..8626fb973 100644 --- a/src/benchflow/acp/runtime.py +++ b/src/benchflow/acp/runtime.py @@ -66,6 +66,7 @@ _ACP_CONNECT_MAX_RETRIES = 3 _ACP_CONNECT_BASE_DELAY = 2.0 +_PROMPT_TIMEOUT_CLEANUP_TOTAL_SEC = 5.0 _PROMPT_CANCEL_DRAIN_TIMEOUT_SEC = 0.25 _ACP_HANDSHAKE_TIMEOUT_ENV = "BENCHFLOW_ACP_HANDSHAKE_TIMEOUT" _ACP_HANDSHAKE_TIMEOUT_DEFAULT_SEC = 60.0 @@ -745,29 +746,76 @@ async def execute_prompts( return trajectory, len(session.tool_calls) +def _consume_task_result(task: asyncio.Task) -> None: + with contextlib.suppress(BaseException): + task.result() + + +async def _cancel_and_drain_tasks( + tasks: set[asyncio.Task], timeout: float +) -> None: + pending = {task for task in tasks if not task.done()} + for task in pending: + task.cancel() + _done, pending = ( + await asyncio.wait(pending, timeout=timeout) if pending else (set(), set()) + ) + for task in tasks - pending: + _consume_task_result(task) + for task in pending: + task.add_done_callback(_consume_task_result) + + async def _cancel_and_drain_prompt_task(prompt_task: asyncio.Task) -> bool: if prompt_task.done(): + _consume_task_result(prompt_task) return True - prompt_task.cancel() - done, _pending = await asyncio.wait( - {prompt_task}, timeout=_PROMPT_CANCEL_DRAIN_TIMEOUT_SEC - ) - if done: - with contextlib.suppress(BaseException): - prompt_task.result() + await _cancel_and_drain_tasks({prompt_task}, _PROMPT_CANCEL_DRAIN_TIMEOUT_SEC) + if prompt_task.done(): return True logger.warning( "ACP prompt task did not finish within %.2fs after cancellation", _PROMPT_CANCEL_DRAIN_TIMEOUT_SEC, ) + return False - def _consume_prompt_result(task: asyncio.Task) -> None: - with contextlib.suppress(BaseException): - task.result() - prompt_task.add_done_callback(_consume_prompt_result) - return False +async def _cancel_prompt_after_timeout( + acp_client: ACPClient, prompt_task: asyncio.Task +) -> bool: + """Request peer cancellation, then bound cleanup without adding a reader.""" + cancel = getattr(acp_client, "cancel", None) + if not callable(cancel): + return await _cancel_and_drain_prompt_task(prompt_task) + if prompt_task.done(): + _consume_task_result(prompt_task) + return True + + async def _cancel_then_wait_for_prompt(): + await cancel() + return await asyncio.shield(prompt_task) + + loop = asyncio.get_running_loop() + deadline = loop.time() + _PROMPT_TIMEOUT_CLEANUP_TOTAL_SEC + supervisor_task = asyncio.create_task(_cancel_then_wait_for_prompt()) + try: + cooperative_timeout = max( + 0.0, + deadline - loop.time() - _PROMPT_CANCEL_DRAIN_TIMEOUT_SEC, + ) + await asyncio.wait( + {prompt_task, supervisor_task}, + timeout=cooperative_timeout, + return_when=asyncio.FIRST_COMPLETED, + ) + finally: + remaining = max(0.0, deadline - loop.time()) + await _cancel_and_drain_tasks( + {prompt_task, supervisor_task}, + min(_PROMPT_CANCEL_DRAIN_TIMEOUT_SEC, remaining), + ) + return prompt_task.done() def _agent_prompt_timeout_error(session, timeout: int) -> AgentPromptTimeoutError: @@ -801,15 +849,18 @@ async def _prompt_with_wall_clock_budget( ): """Run a prompt until either it finishes or BenchFlow's budget expires.""" prompt_task = asyncio.create_task(acp_client.prompt(prompt)) + cleanup_attempted = False try: done, _pending = await asyncio.wait({prompt_task}, timeout=timeout) if done: return prompt_task.result() - if await _cancel_and_drain_prompt_task(prompt_task): + cleanup_attempted = True + if await _cancel_prompt_after_timeout(acp_client, prompt_task): raise _agent_prompt_timeout_error(session, timeout) raise TimeoutError(f"Agent prompt exceeded wall-clock budget {timeout}s") finally: - if not prompt_task.done(): + if not prompt_task.done() and not cleanup_attempted: + cleanup_attempted = True await _cancel_and_drain_prompt_task(prompt_task) @@ -838,6 +889,7 @@ def _activity_count() -> int: ) prompt_task = asyncio.create_task(acp_client.prompt(prompt)) + cleanup_attempted = False last_progress = asyncio.get_event_loop().time() last_activity_at = datetime.now(UTC) last_count = _activity_count() @@ -887,15 +939,19 @@ def _activity_count() -> int: n_thought_chunks=len(session.thought_chunks), last_activity_at=last_activity_at.isoformat(), ) - raise IdleTimeoutError( + error = IdleTimeoutError( f"Agent idle for {idle_timeout}s with no new tool call, " f"message, or thought " f"(last activity {int(now - last_progress)}s ago, " f"{len(session.tool_calls)} tool calls so far)", diag, ) + cleanup_attempted = True + await _cancel_prompt_after_timeout(acp_client, prompt_task) + raise error if now > deadline: - if await _cancel_and_drain_prompt_task(prompt_task): + cleanup_attempted = True + if await _cancel_prompt_after_timeout(acp_client, prompt_task): raise _agent_prompt_timeout_error(session, timeout) raise TimeoutError( f"Agent prompt exceeded wall-clock budget {timeout}s" @@ -907,5 +963,6 @@ def _activity_count() -> int: # external-cancellation path (CancelledError from sleep). Bound the # drain so a non-cooperative Daytona/ACP read cannot hide the watchdog # timeout forever; cleanup will tear down the live process. - if not prompt_task.done(): + if not prompt_task.done() and not cleanup_attempted: + cleanup_attempted = True await _cancel_and_drain_prompt_task(prompt_task) diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 8e73b0f96..9a2c9c538 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -1371,6 +1371,9 @@ async def disconnect(self) -> None: self._capture_partial_session_factory_trajectory() else: self._capture_partial_acp_trajectory() + collect_native_usage = getattr(self, "_collect_native_acp_usage", None) + if callable(collect_native_usage): + collect_native_usage() if self._acp_client: try: await self._acp_client.close() diff --git a/tests/test_acp.py b/tests/test_acp.py index c20819c61..56f2f769e 100644 --- a/tests/test_acp.py +++ b/tests/test_acp.py @@ -11,13 +11,108 @@ from benchflow.acp.client import ACPClient, ACPError from benchflow.acp.container_transport import ContainerTransport from benchflow.acp.session import ACPSession -from benchflow.acp.transport import StdioTransport +from benchflow.acp.transport import StdioTransport, Transport from benchflow.acp.types import StopReason, ToolCallStatus MOCK_AGENT = str(Path(__file__).parent / "fixtures" / "mock_acp_agent.py") MOCK_AGENT_INTERLEAVED = str( Path(__file__).parent / "fixtures" / "mock_acp_agent_interleaved.py" ) +_TIMEOUT_USAGE = { + "input_tokens": 10, + "output_tokens": 4, + "total_tokens": 14, + "cached_read_tokens": None, + "cached_write_tokens": None, + "thought_tokens": None, +} + + +class _CooperativeCancellationTransport(Transport): + """Minimal ACP peer that returns terminal evidence after session/cancel.""" + + def __init__( + self, + *, + emit_pending_tool: bool, + stall_cancel_after_response: bool = False, + ) -> None: + self._messages: asyncio.Queue[dict] = asyncio.Queue() + self._prompt_request_id: int | None = None + self._emit_pending_tool = emit_pending_tool + self._stall_cancel_after_response = stall_cancel_after_response + self.cancel_calls = 0 + self.cancel_send_finished = asyncio.Event() + self.receive_calls = 0 + + async def start(self) -> None: + pass + + def _queue_update(self, update: dict) -> None: + self._messages.put_nowait( + { + "jsonrpc": "2.0", + "method": "session/update", + "params": {"sessionId": "timeout-session", "update": update}, + } + ) + + def _queue_response(self) -> None: + assert self._prompt_request_id is not None + self._messages.put_nowait( + { + "jsonrpc": "2.0", + "id": self._prompt_request_id, + "result": { + "stopReason": "cancelled", + "usage": { + "inputTokens": 10, + "outputTokens": 4, + "totalTokens": 14, + }, + }, + } + ) + + async def send(self, message: dict) -> None: + if message.get("method") == "session/prompt": + self._prompt_request_id = message["id"] + if self._emit_pending_tool: + self._queue_update( + { + "sessionUpdate": "tool_call", + "toolCallId": "tool-1", + "title": "work before timeout", + "kind": "bash", + } + ) + return + + if message.get("method") == "session/cancel": + self.cancel_calls += 1 + try: + assert message["params"] == {"sessionId": "timeout-session"} + assert self._prompt_request_id is not None + if self._emit_pending_tool: + self._queue_update( + { + "sessionUpdate": "tool_call_update", + "toolCallId": "tool-1", + "status": "cancelled", + } + ) + self._queue_response() + if self._stall_cancel_after_response: + await asyncio.Future() + finally: + self.cancel_send_finished.set() + + async def receive(self) -> dict: + self.receive_calls += 1 + return await self._messages.get() + + async def close(self) -> None: + pass class TestACPClient: @@ -830,25 +925,91 @@ async def test_prompt_with_interleaved_notifications_and_request(self) -> None: await client.close() +class TestACPTimeoutCancellation: + @pytest.mark.asyncio + async def test_wall_timeout_requests_cooperative_acp_cancellation(self) -> None: + """Guards issue #933 against evidence loss demonstrated by PR #1051.""" + from benchflow.acp.runtime import AgentPromptTimeoutError, execute_prompts + + transport = _CooperativeCancellationTransport( + emit_pending_tool=True, + stall_cancel_after_response=True, + ) + client = ACPClient(transport) + session = ACPSession("timeout-session") + client._session = session + + with pytest.raises(AgentPromptTimeoutError) as exc_info: + await asyncio.wait_for( + execute_prompts( + client, + session, + ["solve"], + timeout=0.05, # type: ignore[arg-type] + idle_timeout=None, + ), + timeout=10, + ) + + await asyncio.sleep(0) + assert transport.cancel_calls == 1 + assert transport.cancel_send_finished.is_set() + assert transport.receive_calls == 3 + assert session.latest_usage_totals() == _TIMEOUT_USAGE + assert session.stop_reason == StopReason.CANCELLED + tool_event = next( + event + for event in exc_info.value.trajectory + if event["type"] == "tool_call" + ) + assert tool_event["status"] == ToolCallStatus.CANCELLED.value + assert exc_info.value.terminal_trajectory_complete is True + + @pytest.mark.asyncio + async def test_idle_timeout_requests_cooperative_acp_cancellation(self) -> None: + """Guards issue #933 against idle usage loss demonstrated by PR #1051.""" + from benchflow.acp.runtime import IdleTimeoutError, execute_prompts + + transport = _CooperativeCancellationTransport(emit_pending_tool=False) + client = ACPClient(transport) + session = ACPSession("timeout-session") + client._session = session + + with pytest.raises(IdleTimeoutError): + await asyncio.wait_for( + execute_prompts( + client, + session, + ["solve"], + timeout=30, + idle_timeout=1, + ), + timeout=10, + ) + + await asyncio.sleep(0) + assert transport.cancel_calls == 1 + assert transport.receive_calls == 1 + assert session.latest_usage_totals() == _TIMEOUT_USAGE + assert session.stop_reason == StopReason.CANCELLED + + class TestACPIdleWatchdog: @pytest.mark.asyncio async def test_idle_watchdog_returns_even_when_prompt_cancel_drain_stalls( self, + monkeypatch, ) -> None: - """Guards the 2026-05-22 Daytona/Gemini blocker fix against stuck cancel drain. + """Guards issue #933 against stuck cleanup demonstrated by PR #1051. - When the idle watchdog fires it cancels the prompt task and drains it. - A non-cooperative agent — here one that swallows the cancellation and - blocks on an event that never arrives — must not be able to wedge the - watchdog: it bounds the drain and raises ``IdleTimeoutError`` anyway. + When the idle watchdog fires, both the ACP cancellation send and prompt + can remain non-cooperative. Neither task may wedge the watchdog. Determinism: this asserts the *behaviour* (idle error raised; the stuck prompt task abandoned while still pending) rather than a wall-clock upper bound, so it cannot be squeezed into a spurious failure when the full - suite loads the event loop. The outer ``wait_for`` is only a hang guard; - its timeout is deliberately loose (far above the real ~1.25s runtime), so - load can never trip it, while a regression to an unbounded drain hangs - past it and fails. + suite loads the event loop. The outer ``wait_for`` is only a loose hang + guard; a regression to an unbounded drain hangs past it and fails. """ from benchflow.acp.runtime import IdleTimeoutError, execute_prompts @@ -858,6 +1019,18 @@ class StubbornPromptClient: def __init__(self) -> None: self.release = asyncio.Event() self.task: asyncio.Task | None = None + self.cancel_release = asyncio.Event() + self.cancel_task: asyncio.Task | None = None + self.cancel_finished = asyncio.Event() + self.cancel_calls = 0 + + async def cancel(self) -> None: + self.cancel_calls += 1 + self.cancel_task = asyncio.current_task() + try: + await self.cancel_release.wait() + finally: + self.cancel_finished.set() async def prompt(self, _prompt: str): self.task = asyncio.current_task() @@ -869,9 +1042,15 @@ async def prompt(self, _prompt: str): client = StubbornPromptClient() session = ACPSession("idle-session") + monkeypatch.setattr( + "benchflow.acp.runtime._PROMPT_TIMEOUT_CLEANUP_TOTAL_SEC", + 0.1, + ) + monkeypatch.setattr( + "benchflow.acp.runtime._PROMPT_CANCEL_DRAIN_TIMEOUT_SEC", 0.05 + ) - # Loose hang guard: ~4x the real runtime (1s idle detect + 0.25s bounded - # drain). Not a tight assertion — it only catches a true hang/regression. + # Loose hang guard: far above 1s idle detection plus patched cleanup. hang_guard_sec = 5.0 try: @@ -889,13 +1068,24 @@ async def prompt(self, _prompt: str): # The watchdog returned while the stuck prompt task is still wedged on # release.wait(): proves it bounded the drain instead of waiting for a # cancellation that never completes. Load-independent — no clock math. + assert client.cancel_calls == 1 + assert client.cancel_finished.is_set() + assert client.cancel_task is not None + assert client.cancel_task.cancelled() assert client.task is not None assert not client.task.done() finally: + client.cancel_release.set() client.release.set() - if client.task is not None: - with pytest.raises(asyncio.CancelledError): - await client.task + tasks = [ + task + for task in (client.cancel_task, client.task) + if task is not None + ] + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) class TestIdleTimeoutDiagnostics: diff --git a/tests/test_native_acp_usage.py b/tests/test_native_acp_usage.py index ebfea42f2..7bf0d04e6 100644 --- a/tests/test_native_acp_usage.py +++ b/tests/test_native_acp_usage.py @@ -3,6 +3,7 @@ import json from datetime import datetime from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest @@ -47,6 +48,38 @@ async def fake_send_request(method, params): } +@pytest.mark.asyncio +async def test_disconnect_preserves_native_usage_in_final_metrics(): + """Guards issue #933 against usage loss demonstrated by PR #1051.""" + from benchflow.acp.session import ACPSession + from benchflow.rollout import Rollout + + session = ACPSession("timeout-session") + session.record_prompt_usage( + SimpleNamespace(input_tokens=10, output_tokens=4, total_tokens=14) + ) + client = SimpleNamespace(session=session, close=AsyncMock()) + rollout = Rollout.__new__(Rollout) + rollout._acp_client = client + rollout._session = session + rollout._session_adapter = None + rollout._is_session_factory = False + rollout._trajectory = [] + rollout._session_traj_count = 0 + rollout._native_usage_checkpoint = None + rollout._native_usage_metrics = {"total_tokens": 0} + rollout._usage_metrics = {"usage_source": "unavailable"} + rollout._agent_launch = "" + rollout._env = None + rollout._phase = "connected" + + await rollout.disconnect() + rollout._finalize_usage_metrics() + + assert rollout._usage_metrics["usage_source"] == "agent_native_acp" + assert rollout._usage_metrics["total_tokens"] == 14 + + def test_rollout_native_acp_usage_uses_cumulative_deltas(): """Guards PR #613 follow-up: ACP cumulative usage is not double-counted.""" from benchflow.acp.session import ACPSession From 89f1d43d3b2cae1a946ee4372ed7d491639b24f5 Mon Sep 17 00:00:00 2001 From: kywch Date: Sun, 30 Aug 2026 22:28:00 -0700 Subject: [PATCH 2/4] fix(rollout): reset native usage checkpoint per session --- src/benchflow/acp/runtime.py | 4 +--- src/benchflow/rollout/__init__.py | 1 + tests/test_acp.py | 8 ++------ tests/test_connect_as_env.py | 28 ++++++++++++++++++++++++++++ 4 files changed, 32 insertions(+), 9 deletions(-) diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index 8626fb973..1f0974be1 100644 --- a/src/benchflow/acp/runtime.py +++ b/src/benchflow/acp/runtime.py @@ -751,9 +751,7 @@ def _consume_task_result(task: asyncio.Task) -> None: task.result() -async def _cancel_and_drain_tasks( - tasks: set[asyncio.Task], timeout: float -) -> None: +async def _cancel_and_drain_tasks(tasks: set[asyncio.Task], timeout: float) -> None: pending = {task for task in tasks if not task.done()} for task in pending: task.cancel() diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 9a2c9c538..f7f6695ca 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -2388,6 +2388,7 @@ async def connect_as(self, role: Role) -> None: role.agent, getattr(self, "_task", None), agent_cfg ), ) + self._native_usage_checkpoint = None self._reapply_ask_user_handler() self._attach_trajectory_writer(rollout_dir) self._active_role = role diff --git a/tests/test_acp.py b/tests/test_acp.py index 56f2f769e..fdd80992c 100644 --- a/tests/test_acp.py +++ b/tests/test_acp.py @@ -958,9 +958,7 @@ async def test_wall_timeout_requests_cooperative_acp_cancellation(self) -> None: assert session.latest_usage_totals() == _TIMEOUT_USAGE assert session.stop_reason == StopReason.CANCELLED tool_event = next( - event - for event in exc_info.value.trajectory - if event["type"] == "tool_call" + event for event in exc_info.value.trajectory if event["type"] == "tool_call" ) assert tool_event["status"] == ToolCallStatus.CANCELLED.value assert exc_info.value.terminal_trajectory_complete is True @@ -1078,9 +1076,7 @@ async def prompt(self, _prompt: str): client.cancel_release.set() client.release.set() tasks = [ - task - for task in (client.cancel_task, client.task) - if task is not None + task for task in (client.cancel_task, client.task) if task is not None ] for task in tasks: if not task.done(): diff --git a/tests/test_connect_as_env.py b/tests/test_connect_as_env.py index 129e1d86b..68dbc8f28 100644 --- a/tests/test_connect_as_env.py +++ b/tests/test_connect_as_env.py @@ -140,6 +140,34 @@ def fake_resolve(agent, model, env): assert captured["env"] == {} + @pytest.mark.asyncio + async def test_native_usage_accumulates_across_fresh_sessions(self, _mock_trial): + """Guards PR #1080 / issue #933 across restarted ACP usage counters.""" + from benchflow.acp.session import ACPSession + + first_session = ACPSession("first") + first_session.record_prompt_usage( + {"inputTokens": 10, "outputTokens": 4, "totalTokens": 14} + ) + second_session = ACPSession("second") + second_session.record_prompt_usage( + {"inputTokens": 5, "outputTokens": 2, "totalTokens": 7} + ) + _mock_trial._planes.connect_acp.side_effect = [ + (AsyncMock(), first_session, AsyncMock(), "agent"), + (AsyncMock(), second_session, AsyncMock(), "agent"), + ] + _mock_trial._native_usage_metrics = {"total_tokens": 0} + _mock_trial._native_usage_checkpoint = None + role = _mock_trial._config.scenes[0].roles[0] + + await _mock_trial.connect_as(role) + _mock_trial._collect_native_acp_usage() + await _mock_trial.connect_as(role) + _mock_trial._collect_native_acp_usage() + + assert _mock_trial._native_usage_metrics["total_tokens"] == 21 + @pytest.mark.asyncio async def test_same_agent_different_model_refreshes_credentials(self, _mock_trial): """Guards ENG-91 P0 same-agent role credential refresh regression.""" From ee209502ae74aee7ee84c4836ff8971e320c1e5c Mon Sep 17 00:00:00 2001 From: kywch Date: Tue, 1 Sep 2026 18:37:41 -0700 Subject: [PATCH 3/4] fix(acp): preserve timeout evidence when cancel fails --- src/benchflow/acp/runtime.py | 17 ++++++++++------- tests/test_acp.py | 23 ++++++++++++++++++++--- 2 files changed, 30 insertions(+), 10 deletions(-) diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index 1f0974be1..e8c3d6a0d 100644 --- a/src/benchflow/acp/runtime.py +++ b/src/benchflow/acp/runtime.py @@ -790,27 +790,30 @@ async def _cancel_prompt_after_timeout( _consume_task_result(prompt_task) return True - async def _cancel_then_wait_for_prompt(): - await cancel() - return await asyncio.shield(prompt_task) + async def _request_cancel() -> None: + try: + await cancel() + except asyncio.CancelledError: + raise + except Exception: + logger.warning("Failed to request ACP prompt cancellation", exc_info=True) loop = asyncio.get_running_loop() deadline = loop.time() + _PROMPT_TIMEOUT_CLEANUP_TOTAL_SEC - supervisor_task = asyncio.create_task(_cancel_then_wait_for_prompt()) + cancel_task = asyncio.create_task(_request_cancel()) try: cooperative_timeout = max( 0.0, deadline - loop.time() - _PROMPT_CANCEL_DRAIN_TIMEOUT_SEC, ) await asyncio.wait( - {prompt_task, supervisor_task}, + {prompt_task}, timeout=cooperative_timeout, - return_when=asyncio.FIRST_COMPLETED, ) finally: remaining = max(0.0, deadline - loop.time()) await _cancel_and_drain_tasks( - {prompt_task, supervisor_task}, + {prompt_task, cancel_task}, min(_PROMPT_CANCEL_DRAIN_TIMEOUT_SEC, remaining), ) return prompt_task.done() diff --git a/tests/test_acp.py b/tests/test_acp.py index fdd80992c..a03304a17 100644 --- a/tests/test_acp.py +++ b/tests/test_acp.py @@ -36,11 +36,13 @@ def __init__( *, emit_pending_tool: bool, stall_cancel_after_response: bool = False, + cancel_failure_response_delay: float | None = None, ) -> None: self._messages: asyncio.Queue[dict] = asyncio.Queue() self._prompt_request_id: int | None = None self._emit_pending_tool = emit_pending_tool self._stall_cancel_after_response = stall_cancel_after_response + self._cancel_failure_response_delay = cancel_failure_response_delay self.cancel_calls = 0 self.cancel_send_finished = asyncio.Event() self.receive_calls = 0 @@ -101,6 +103,11 @@ async def send(self, message: dict) -> None: "status": "cancelled", } ) + if self._cancel_failure_response_delay is not None: + asyncio.get_running_loop().call_later( + self._cancel_failure_response_delay, self._queue_response + ) + raise ConnectionError("cancel send failed") self._queue_response() if self._stall_cancel_after_response: await asyncio.Future() @@ -926,14 +933,24 @@ async def test_prompt_with_interleaved_notifications_and_request(self) -> None: class TestACPTimeoutCancellation: + @pytest.mark.parametrize( + "transport_kwargs", + [ + {"stall_cancel_after_response": True}, + {"cancel_failure_response_delay": 0.01}, + ], + ids=["stalled-cancel-send", "failed-cancel-send"], + ) @pytest.mark.asyncio - async def test_wall_timeout_requests_cooperative_acp_cancellation(self) -> None: - """Guards issue #933 against evidence loss demonstrated by PR #1051.""" + async def test_wall_timeout_requests_cooperative_acp_cancellation( + self, transport_kwargs: dict + ) -> None: + """Guards issue #933 and PR #1080's fix for evidence loss shown by PR #1051.""" from benchflow.acp.runtime import AgentPromptTimeoutError, execute_prompts transport = _CooperativeCancellationTransport( emit_pending_tool=True, - stall_cancel_after_response=True, + **transport_kwargs, ) client = ACPClient(transport) session = ACPSession("timeout-session") From 416b3ea22ea5038c74ef5ca4600f809a5b528737 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Wed, 2 Sep 2026 01:33:34 -0700 Subject: [PATCH 4/4] refactor(acp): isolate timeout cleanup lifecycle --- src/benchflow/acp/runtime.py | 89 +-------- src/benchflow/acp/timeout_cleanup.py | 90 +++++++++ tests/test_acp.py | 273 +-------------------------- tests/test_acp_timeout_evidence.py | 259 +++++++++++++++++++++++++ 4 files changed, 359 insertions(+), 352 deletions(-) create mode 100644 src/benchflow/acp/timeout_cleanup.py create mode 100644 tests/test_acp_timeout_evidence.py diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index f29224b46..6fa0c2f15 100644 --- a/src/benchflow/acp/runtime.py +++ b/src/benchflow/acp/runtime.py @@ -26,6 +26,10 @@ from benchflow.acp.client import ACPClient from benchflow.acp.container_transport import ContainerTransport from benchflow.acp.selection import selected_acp_transport +from benchflow.acp.timeout_cleanup import ( + cancel_and_drain_prompt_task, + cancel_prompt_after_timeout, +) from benchflow.acp.types import McpServerSpec from benchflow.acp.watchdog import IdleWatchdog from benchflow.agents.protocol import ACPSessionAdapter @@ -65,8 +69,6 @@ _ACP_CONNECT_MAX_RETRIES = 3 _ACP_CONNECT_BASE_DELAY = 2.0 -_PROMPT_TIMEOUT_CLEANUP_TOTAL_SEC = 5.0 -_PROMPT_CANCEL_DRAIN_TIMEOUT_SEC = 0.25 _ACP_HANDSHAKE_TIMEOUT_ENV = "BENCHFLOW_ACP_HANDSHAKE_TIMEOUT" _ACP_HANDSHAKE_TIMEOUT_DEFAULT_SEC = 60.0 _OPENHANDS_DISABLE_SUBAGENTS_ENV = "BENCHFLOW_OPENHANDS_DISABLE_SUBAGENTS" @@ -745,79 +747,6 @@ async def execute_prompts( return trajectory, len(session.tool_calls) -def _consume_task_result(task: asyncio.Task) -> None: - with contextlib.suppress(BaseException): - task.result() - - -async def _cancel_and_drain_tasks(tasks: set[asyncio.Task], timeout: float) -> None: - pending = {task for task in tasks if not task.done()} - for task in pending: - task.cancel() - _done, pending = ( - await asyncio.wait(pending, timeout=timeout) if pending else (set(), set()) - ) - for task in tasks - pending: - _consume_task_result(task) - for task in pending: - task.add_done_callback(_consume_task_result) - - -async def _cancel_and_drain_prompt_task(prompt_task: asyncio.Task) -> bool: - if prompt_task.done(): - _consume_task_result(prompt_task) - return True - await _cancel_and_drain_tasks({prompt_task}, _PROMPT_CANCEL_DRAIN_TIMEOUT_SEC) - if prompt_task.done(): - return True - - logger.warning( - "ACP prompt task did not finish within %.2fs after cancellation", - _PROMPT_CANCEL_DRAIN_TIMEOUT_SEC, - ) - return False - - -async def _cancel_prompt_after_timeout( - acp_client: ACPClient, prompt_task: asyncio.Task -) -> bool: - """Request peer cancellation, then bound cleanup without adding a reader.""" - cancel = getattr(acp_client, "cancel", None) - if not callable(cancel): - return await _cancel_and_drain_prompt_task(prompt_task) - if prompt_task.done(): - _consume_task_result(prompt_task) - return True - - async def _request_cancel() -> None: - try: - await cancel() - except asyncio.CancelledError: - raise - except Exception: - logger.warning("Failed to request ACP prompt cancellation", exc_info=True) - - loop = asyncio.get_running_loop() - deadline = loop.time() + _PROMPT_TIMEOUT_CLEANUP_TOTAL_SEC - cancel_task = asyncio.create_task(_request_cancel()) - try: - cooperative_timeout = max( - 0.0, - deadline - loop.time() - _PROMPT_CANCEL_DRAIN_TIMEOUT_SEC, - ) - await asyncio.wait( - {prompt_task}, - timeout=cooperative_timeout, - ) - finally: - remaining = max(0.0, deadline - loop.time()) - await _cancel_and_drain_tasks( - {prompt_task, cancel_task}, - min(_PROMPT_CANCEL_DRAIN_TIMEOUT_SEC, remaining), - ) - return prompt_task.done() - - def _agent_prompt_timeout_error(session, timeout: int) -> AgentPromptTimeoutError: session.mark_prompt_end() pending_tool_call_ids = session.pending_tool_call_ids() @@ -855,13 +784,13 @@ async def _prompt_with_wall_clock_budget( if done: return prompt_task.result() cleanup_attempted = True - if await _cancel_prompt_after_timeout(acp_client, prompt_task): + if await cancel_prompt_after_timeout(acp_client, prompt_task): raise _agent_prompt_timeout_error(session, timeout) raise TimeoutError(f"Agent prompt exceeded wall-clock budget {timeout}s") finally: if not prompt_task.done() and not cleanup_attempted: cleanup_attempted = True - await _cancel_and_drain_prompt_task(prompt_task) + await cancel_and_drain_prompt_task(prompt_task) async def _prompt_with_idle_watchdog( @@ -895,11 +824,11 @@ async def _prompt_with_idle_watchdog( if watchdog.idle_expired(now): error = watchdog.timeout_error(session, now=now) cleanup_attempted = True - await _cancel_prompt_after_timeout(acp_client, prompt_task) + await cancel_prompt_after_timeout(acp_client, prompt_task) raise error if watchdog.wall_clock_expired(now): cleanup_attempted = True - if await _cancel_prompt_after_timeout(acp_client, prompt_task): + if await cancel_prompt_after_timeout(acp_client, prompt_task): raise _agent_prompt_timeout_error(session, timeout) raise TimeoutError( f"Agent prompt exceeded wall-clock budget {timeout}s" @@ -913,4 +842,4 @@ async def _prompt_with_idle_watchdog( # timeout forever; cleanup will tear down the live process. if not prompt_task.done() and not cleanup_attempted: cleanup_attempted = True - await _cancel_and_drain_prompt_task(prompt_task) + await cancel_and_drain_prompt_task(prompt_task) diff --git a/src/benchflow/acp/timeout_cleanup.py b/src/benchflow/acp/timeout_cleanup.py new file mode 100644 index 000000000..d5eb2d0a3 --- /dev/null +++ b/src/benchflow/acp/timeout_cleanup.py @@ -0,0 +1,90 @@ +"""Bounded cooperative cleanup for timed-out ACP prompts. + +The original prompt task remains the sole ACP response reader. A best-effort +``session/cancel`` request runs independently while the prompt gets a short +grace period to receive peer-authored terminal usage and tool updates. At the +shared deadline, every surviving task is cancelled and boundedly drained. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +from typing import Any + +logger = logging.getLogger(__name__) + +PROMPT_TIMEOUT_CLEANUP_TOTAL_SEC = 5.0 +PROMPT_CANCEL_DRAIN_TIMEOUT_SEC = 0.25 + + +def _consume_task_result(task: asyncio.Task) -> None: + with contextlib.suppress(BaseException): + task.result() + + +async def _cancel_and_drain_tasks(tasks: set[asyncio.Task], timeout: float) -> None: + pending = {task for task in tasks if not task.done()} + for task in pending: + task.cancel() + _done, pending = ( + await asyncio.wait(pending, timeout=timeout) if pending else (set(), set()) + ) + for task in tasks - pending: + _consume_task_result(task) + for task in pending: + task.add_done_callback(_consume_task_result) + + +async def cancel_and_drain_prompt_task(prompt_task: asyncio.Task) -> bool: + """Hard-cancel a prompt and return whether it completed within the drain.""" + if prompt_task.done(): + _consume_task_result(prompt_task) + return True + await _cancel_and_drain_tasks({prompt_task}, PROMPT_CANCEL_DRAIN_TIMEOUT_SEC) + if prompt_task.done(): + return True + + logger.warning( + "ACP prompt task did not finish within %.2fs after cancellation", + PROMPT_CANCEL_DRAIN_TIMEOUT_SEC, + ) + return False + + +async def cancel_prompt_after_timeout( + acp_client: Any, prompt_task: asyncio.Task +) -> bool: + """Request peer cancellation, then bound cleanup without adding a reader.""" + cancel = getattr(acp_client, "cancel", None) + if not callable(cancel): + return await cancel_and_drain_prompt_task(prompt_task) + if prompt_task.done(): + _consume_task_result(prompt_task) + return True + + async def _request_cancel() -> None: + try: + await cancel() + except asyncio.CancelledError: + raise + except Exception: + logger.warning("Failed to request ACP prompt cancellation", exc_info=True) + + loop = asyncio.get_running_loop() + deadline = loop.time() + PROMPT_TIMEOUT_CLEANUP_TOTAL_SEC + cancel_task = asyncio.create_task(_request_cancel()) + try: + cooperative_timeout = max( + 0.0, + deadline - loop.time() - PROMPT_CANCEL_DRAIN_TIMEOUT_SEC, + ) + await asyncio.wait({prompt_task}, timeout=cooperative_timeout) + finally: + remaining = max(0.0, deadline - loop.time()) + await _cancel_and_drain_tasks( + {prompt_task, cancel_task}, + min(PROMPT_CANCEL_DRAIN_TIMEOUT_SEC, remaining), + ) + return prompt_task.done() diff --git a/tests/test_acp.py b/tests/test_acp.py index 49abb434b..a6a04550e 100644 --- a/tests/test_acp.py +++ b/tests/test_acp.py @@ -11,115 +11,13 @@ from benchflow.acp.client import ACPClient, ACPError from benchflow.acp.container_transport import ContainerTransport from benchflow.acp.session import ACPSession -from benchflow.acp.transport import StdioTransport, Transport +from benchflow.acp.transport import StdioTransport from benchflow.acp.types import StopReason, ToolCallStatus MOCK_AGENT = str(Path(__file__).parent / "fixtures" / "mock_acp_agent.py") MOCK_AGENT_INTERLEAVED = str( Path(__file__).parent / "fixtures" / "mock_acp_agent_interleaved.py" ) -_TIMEOUT_USAGE = { - "input_tokens": 10, - "output_tokens": 4, - "total_tokens": 14, - "cached_read_tokens": None, - "cached_write_tokens": None, - "thought_tokens": None, -} - - -class _CooperativeCancellationTransport(Transport): - """Minimal ACP peer that returns terminal evidence after session/cancel.""" - - def __init__( - self, - *, - emit_pending_tool: bool, - stall_cancel_after_response: bool = False, - cancel_failure_response_delay: float | None = None, - ) -> None: - self._messages: asyncio.Queue[dict] = asyncio.Queue() - self._prompt_request_id: int | None = None - self._emit_pending_tool = emit_pending_tool - self._stall_cancel_after_response = stall_cancel_after_response - self._cancel_failure_response_delay = cancel_failure_response_delay - self.cancel_calls = 0 - self.cancel_send_finished = asyncio.Event() - self.receive_calls = 0 - - async def start(self) -> None: - pass - - def _queue_update(self, update: dict) -> None: - self._messages.put_nowait( - { - "jsonrpc": "2.0", - "method": "session/update", - "params": {"sessionId": "timeout-session", "update": update}, - } - ) - - def _queue_response(self) -> None: - assert self._prompt_request_id is not None - self._messages.put_nowait( - { - "jsonrpc": "2.0", - "id": self._prompt_request_id, - "result": { - "stopReason": "cancelled", - "usage": { - "inputTokens": 10, - "outputTokens": 4, - "totalTokens": 14, - }, - }, - } - ) - - async def send(self, message: dict) -> None: - if message.get("method") == "session/prompt": - self._prompt_request_id = message["id"] - if self._emit_pending_tool: - self._queue_update( - { - "sessionUpdate": "tool_call", - "toolCallId": "tool-1", - "title": "work before timeout", - "kind": "bash", - } - ) - return - - if message.get("method") == "session/cancel": - self.cancel_calls += 1 - try: - assert message["params"] == {"sessionId": "timeout-session"} - assert self._prompt_request_id is not None - if self._emit_pending_tool: - self._queue_update( - { - "sessionUpdate": "tool_call_update", - "toolCallId": "tool-1", - "status": "cancelled", - } - ) - if self._cancel_failure_response_delay is not None: - asyncio.get_running_loop().call_later( - self._cancel_failure_response_delay, self._queue_response - ) - raise ConnectionError("cancel send failed") - self._queue_response() - if self._stall_cancel_after_response: - await asyncio.Future() - finally: - self.cancel_send_finished.set() - - async def receive(self) -> dict: - self.receive_calls += 1 - return await self._messages.get() - - async def close(self) -> None: - pass class TestACPClient: @@ -932,175 +830,6 @@ async def test_prompt_with_interleaved_notifications_and_request(self) -> None: await client.close() -class TestACPTimeoutCancellation: - @pytest.mark.parametrize( - "transport_kwargs", - [ - {"stall_cancel_after_response": True}, - {"cancel_failure_response_delay": 0.01}, - ], - ids=["stalled-cancel-send", "failed-cancel-send"], - ) - @pytest.mark.asyncio - async def test_wall_timeout_requests_cooperative_acp_cancellation( - self, transport_kwargs: dict - ) -> None: - """Guards issue #933 and PR #1080's fix for evidence loss shown by PR #1051.""" - from benchflow.acp.runtime import AgentPromptTimeoutError, execute_prompts - - transport = _CooperativeCancellationTransport( - emit_pending_tool=True, - **transport_kwargs, - ) - client = ACPClient(transport) - session = ACPSession("timeout-session") - client._session = session - - with pytest.raises(AgentPromptTimeoutError) as exc_info: - await asyncio.wait_for( - execute_prompts( - client, - session, - ["solve"], - timeout=0.05, # type: ignore[arg-type] - idle_timeout=None, - ), - timeout=10, - ) - - await asyncio.sleep(0) - assert transport.cancel_calls == 1 - assert transport.cancel_send_finished.is_set() - assert transport.receive_calls == 3 - assert session.latest_usage_totals() == _TIMEOUT_USAGE - assert session.stop_reason == StopReason.CANCELLED - tool_event = next( - event for event in exc_info.value.trajectory if event["type"] == "tool_call" - ) - assert tool_event["status"] == ToolCallStatus.CANCELLED.value - assert exc_info.value.terminal_trajectory_complete is True - - @pytest.mark.asyncio - async def test_idle_timeout_requests_cooperative_acp_cancellation(self) -> None: - """Guards issue #933 against idle usage loss demonstrated by PR #1051.""" - from benchflow.acp.runtime import IdleTimeoutError, execute_prompts - - transport = _CooperativeCancellationTransport(emit_pending_tool=False) - client = ACPClient(transport) - session = ACPSession("timeout-session") - client._session = session - - with pytest.raises(IdleTimeoutError): - await asyncio.wait_for( - execute_prompts( - client, - session, - ["solve"], - timeout=30, - idle_timeout=1, - ), - timeout=10, - ) - - await asyncio.sleep(0) - assert transport.cancel_calls == 1 - assert transport.receive_calls == 1 - assert session.latest_usage_totals() == _TIMEOUT_USAGE - assert session.stop_reason == StopReason.CANCELLED - - -class TestACPIdleWatchdog: - @pytest.mark.asyncio - async def test_idle_watchdog_returns_even_when_prompt_cancel_drain_stalls( - self, - monkeypatch, - ) -> None: - """Guards issue #933 against stuck cleanup demonstrated by PR #1051. - - When the idle watchdog fires, both the ACP cancellation send and prompt - can remain non-cooperative. Neither task may wedge the watchdog. - - Determinism: this asserts the *behaviour* (idle error raised; the stuck - prompt task abandoned while still pending) rather than a wall-clock upper - bound, so it cannot be squeezed into a spurious failure when the full - suite loads the event loop. The outer ``wait_for`` is only a loose hang - guard; a regression to an unbounded drain hangs past it and fails. - """ - from benchflow.acp.runtime import IdleTimeoutError, execute_prompts - - # release is never set while the watchdog runs, so the prompt task stays - # wedged in its cancellation handler — exactly the stuck-drain scenario. - class StubbornPromptClient: - def __init__(self) -> None: - self.release = asyncio.Event() - self.task: asyncio.Task | None = None - self.cancel_release = asyncio.Event() - self.cancel_task: asyncio.Task | None = None - self.cancel_finished = asyncio.Event() - self.cancel_calls = 0 - - async def cancel(self) -> None: - self.cancel_calls += 1 - self.cancel_task = asyncio.current_task() - try: - await self.cancel_release.wait() - finally: - self.cancel_finished.set() - - async def prompt(self, _prompt: str): - self.task = asyncio.current_task() - try: - await asyncio.Future() - except asyncio.CancelledError: - await self.release.wait() - raise - - client = StubbornPromptClient() - session = ACPSession("idle-session") - monkeypatch.setattr( - "benchflow.acp.runtime._PROMPT_TIMEOUT_CLEANUP_TOTAL_SEC", - 0.1, - ) - monkeypatch.setattr( - "benchflow.acp.runtime._PROMPT_CANCEL_DRAIN_TIMEOUT_SEC", 0.05 - ) - - # Loose hang guard: far above 1s idle detection plus patched cleanup. - hang_guard_sec = 5.0 - - try: - with pytest.raises(IdleTimeoutError, match="Agent idle for 1s"): - await asyncio.wait_for( - execute_prompts( - client, # type: ignore[arg-type] - session, - ["solve"], - timeout=30, - idle_timeout=1, - ), - timeout=hang_guard_sec, - ) - # The watchdog returned while the stuck prompt task is still wedged on - # release.wait(): proves it bounded the drain instead of waiting for a - # cancellation that never completes. Load-independent — no clock math. - assert client.cancel_calls == 1 - assert client.cancel_finished.is_set() - assert client.cancel_task is not None - assert client.cancel_task.cancelled() - assert client.task is not None - assert not client.task.done() - finally: - client.cancel_release.set() - client.release.set() - tasks = [ - task for task in (client.cancel_task, client.task) if task is not None - ] - for task in tasks: - if not task.done(): - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - - class TestPendingToolCallGrace: @pytest.mark.asyncio async def test_lost_tool_call_completion_trips_idle_after_grace(self): diff --git a/tests/test_acp_timeout_evidence.py b/tests/test_acp_timeout_evidence.py new file mode 100644 index 000000000..318e3b91b --- /dev/null +++ b/tests/test_acp_timeout_evidence.py @@ -0,0 +1,259 @@ +"""ACP timeout cancellation and terminal-evidence regression coverage.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from benchflow.acp.client import ACPClient +from benchflow.acp.session import ACPSession +from benchflow.acp.transport import Transport +from benchflow.acp.types import StopReason, ToolCallStatus + +_TIMEOUT_USAGE = { + "input_tokens": 10, + "output_tokens": 4, + "total_tokens": 14, + "cached_read_tokens": None, + "cached_write_tokens": None, + "thought_tokens": None, +} + + +class _CooperativeCancellationTransport(Transport): + """Minimal ACP peer that returns terminal evidence after session/cancel.""" + + def __init__( + self, + *, + emit_pending_tool: bool, + stall_cancel_after_response: bool = False, + cancel_failure_response_delay: float | None = None, + ) -> None: + self._messages: asyncio.Queue[dict] = asyncio.Queue() + self._prompt_request_id: int | None = None + self._emit_pending_tool = emit_pending_tool + self._stall_cancel_after_response = stall_cancel_after_response + self._cancel_failure_response_delay = cancel_failure_response_delay + self.cancel_calls = 0 + self.cancel_send_finished = asyncio.Event() + self.receive_calls = 0 + + async def start(self) -> None: + pass + + def _queue_update(self, update: dict) -> None: + self._messages.put_nowait( + { + "jsonrpc": "2.0", + "method": "session/update", + "params": {"sessionId": "timeout-session", "update": update}, + } + ) + + def _queue_response(self) -> None: + assert self._prompt_request_id is not None + self._messages.put_nowait( + { + "jsonrpc": "2.0", + "id": self._prompt_request_id, + "result": { + "stopReason": "cancelled", + "usage": { + "inputTokens": 10, + "outputTokens": 4, + "totalTokens": 14, + }, + }, + } + ) + + async def send(self, message: dict) -> None: + if message.get("method") == "session/prompt": + self._prompt_request_id = message["id"] + if self._emit_pending_tool: + self._queue_update( + { + "sessionUpdate": "tool_call", + "toolCallId": "tool-1", + "title": "work before timeout", + "kind": "bash", + } + ) + return + + if message.get("method") == "session/cancel": + self.cancel_calls += 1 + try: + assert message["params"] == {"sessionId": "timeout-session"} + assert self._prompt_request_id is not None + if self._emit_pending_tool: + self._queue_update( + { + "sessionUpdate": "tool_call_update", + "toolCallId": "tool-1", + "status": "cancelled", + } + ) + if self._cancel_failure_response_delay is not None: + asyncio.get_running_loop().call_later( + self._cancel_failure_response_delay, self._queue_response + ) + raise ConnectionError("cancel send failed") + self._queue_response() + if self._stall_cancel_after_response: + await asyncio.Future() + finally: + self.cancel_send_finished.set() + + async def receive(self) -> dict: + self.receive_calls += 1 + return await self._messages.get() + + async def close(self) -> None: + pass + + +class TestACPTimeoutCancellation: + @pytest.mark.parametrize( + "transport_kwargs", + [ + {"stall_cancel_after_response": True}, + {"cancel_failure_response_delay": 0.01}, + ], + ids=["stalled-cancel-send", "failed-cancel-send"], + ) + @pytest.mark.asyncio + async def test_wall_timeout_requests_cooperative_acp_cancellation( + self, transport_kwargs: dict + ) -> None: + """Guards issue #933 and PR #1080's evidence fix after PR #1051.""" + from benchflow.acp.runtime import AgentPromptTimeoutError, execute_prompts + + transport = _CooperativeCancellationTransport( + emit_pending_tool=True, + **transport_kwargs, + ) + client = ACPClient(transport) + session = ACPSession("timeout-session") + client._session = session + + with pytest.raises(AgentPromptTimeoutError) as exc_info: + await asyncio.wait_for( + execute_prompts( + client, + session, + ["solve"], + timeout=0.05, # type: ignore[arg-type] + idle_timeout=None, + ), + timeout=10, + ) + + await asyncio.sleep(0) + assert transport.cancel_calls == 1 + assert transport.cancel_send_finished.is_set() + assert transport.receive_calls == 3 + assert session.latest_usage_totals() == _TIMEOUT_USAGE + assert session.stop_reason == StopReason.CANCELLED + tool_event = next( + event for event in exc_info.value.trajectory if event["type"] == "tool_call" + ) + assert tool_event["status"] == ToolCallStatus.CANCELLED.value + assert exc_info.value.terminal_trajectory_complete is True + + @pytest.mark.asyncio + async def test_idle_timeout_requests_cooperative_acp_cancellation(self) -> None: + """Guards PR #1080 against idle-timeout usage loss shown by PR #1051.""" + from benchflow.acp.runtime import IdleTimeoutError, execute_prompts + + transport = _CooperativeCancellationTransport(emit_pending_tool=False) + client = ACPClient(transport) + session = ACPSession("timeout-session") + client._session = session + + with pytest.raises(IdleTimeoutError): + await asyncio.wait_for( + execute_prompts( + client, + session, + ["solve"], + timeout=30, + idle_timeout=1, + ), + timeout=10, + ) + + await asyncio.sleep(0) + assert transport.cancel_calls == 1 + assert transport.receive_calls == 1 + assert session.latest_usage_totals() == _TIMEOUT_USAGE + assert session.stop_reason == StopReason.CANCELLED + + @pytest.mark.asyncio + async def test_idle_timeout_bounds_noncooperative_cleanup( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Guards PR #1080: stuck prompt and cancel tasks cannot wedge timeout.""" + from benchflow.acp import timeout_cleanup + from benchflow.acp.runtime import IdleTimeoutError, execute_prompts + + class StubbornPromptClient: + def __init__(self) -> None: + self.release = asyncio.Event() + self.task: asyncio.Task | None = None + self.cancel_release = asyncio.Event() + self.cancel_task: asyncio.Task | None = None + self.cancel_finished = asyncio.Event() + self.cancel_calls = 0 + + async def cancel(self) -> None: + self.cancel_calls += 1 + self.cancel_task = asyncio.current_task() + try: + await self.cancel_release.wait() + finally: + self.cancel_finished.set() + + async def prompt(self, _prompt: str): + self.task = asyncio.current_task() + try: + await asyncio.Future() + except asyncio.CancelledError: + await self.release.wait() + raise + + client = StubbornPromptClient() + session = ACPSession("idle-session") + monkeypatch.setattr(timeout_cleanup, "PROMPT_TIMEOUT_CLEANUP_TOTAL_SEC", 0.1) + monkeypatch.setattr(timeout_cleanup, "PROMPT_CANCEL_DRAIN_TIMEOUT_SEC", 0.05) + + try: + with pytest.raises(IdleTimeoutError, match="Agent idle for 1s"): + await asyncio.wait_for( + execute_prompts( + client, # type: ignore[arg-type] + session, + ["solve"], + timeout=30, + idle_timeout=1, + ), + timeout=5.0, + ) + assert client.cancel_calls == 1 + assert client.cancel_finished.is_set() + assert client.cancel_task is not None + assert client.cancel_task.cancelled() + assert client.task is not None + assert not client.task.done() + finally: + client.cancel_release.set() + client.release.set() + tasks = [ + task for task in (client.cancel_task, client.task) if task is not None + ] + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True)