diff --git a/src/art/tau_bench/client.py b/src/art/tau_bench/client.py index 4beb08351..55b8b5dfd 100644 --- a/src/art/tau_bench/client.py +++ b/src/art/tau_bench/client.py @@ -386,9 +386,21 @@ async def _delete_after_error( self, env_id: str, primary_error: BaseException ) -> None: cleanup = asyncio.create_task(self.delete_environment(env_id)) + deadline = asyncio.get_running_loop().time() + DEFAULT_CLEANUP_TIMEOUT + interrupted: asyncio.CancelledError | None = None try: - async with asyncio.timeout(DEFAULT_CLEANUP_TIMEOUT): - await asyncio.shield(cleanup) + while not cleanup.done(): + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + raise TimeoutError + try: + done, _ = await asyncio.wait({cleanup}, timeout=remaining) + except asyncio.CancelledError as error: + interrupted = error + continue + if not done: + raise TimeoutError + cleanup.result() except BaseException as cleanup_error: if not cleanup.done(): cleanup.cancel() @@ -401,6 +413,13 @@ async def _delete_after_error( env_id, exc_info=True, ) + if interrupted is not None: + if isinstance(primary_error, asyncio.CancelledError): + raise primary_error + raise BaseExceptionGroup( + "Tau environment failure during cancelled cleanup", + [primary_error, interrupted], + ) from None def _auth_headers(self) -> dict[str, str]: if self.api_key is None: diff --git a/tests/unit/test_tau_bench_client.py b/tests/unit/test_tau_bench_client.py index 4cac9a85d..b8e61c230 100644 --- a/tests/unit/test_tau_bench_client.py +++ b/tests/unit/test_tau_bench_client.py @@ -1079,3 +1079,117 @@ async def test_rollout_stops_before_next_turn_exceeds_context( assert trajectory.metrics["num_turns"] == 1 assert client.steps == 1 assert client.deleted == ["env-1"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel_primary", [False, True]) +@pytest.mark.parametrize("cleanup_fails", [False, True]) +async def test_environment_joins_cooperative_delete_despite_repeated_parent_cancel( + cancel_primary: bool, + cleanup_fails: bool, +) -> None: + before = asyncio.all_tasks() + entered, release, closed = (asyncio.Event() for _ in range(3)) + primary = ( + asyncio.CancelledError("original") if cancel_primary else ValueError("original") + ) + + class Client(FakeTauBenchClient): + async def delete_environment(self, env_id: str) -> DeleteEnvironmentResponse: + entered.set() + try: + await release.wait() + if cleanup_fails: + raise RuntimeError("public deletion failure") + return await super().delete_environment(env_id) + finally: + closed.set() + + client = Client() + + async def owner() -> None: + async with client.environment(domain="telecom", task_id="task_001"): + raise primary + + task = asyncio.create_task(owner()) + try: + await asyncio.wait_for(entered.wait(), 1) + task.cancel() + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + assert not task.done() and not closed.is_set() + release.set() + with pytest.raises(BaseException) as caught: + await task + if cancel_primary: + assert caught.value is primary + else: + assert isinstance(caught.value, BaseExceptionGroup) + assert primary in caught.value.exceptions + assert any( + isinstance(e, asyncio.CancelledError) for e in caught.value.exceptions + ) + assert closed.is_set() + assert bool(getattr(primary, "__notes__", [])) == cleanup_fails + assert client.deleted == ([] if cleanup_fails else ["env-1"]) + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await asyncio.sleep(0) + assert asyncio.all_tasks() == before + + +@pytest.mark.asyncio +async def test_environment_parent_cancel_does_not_renew_delete_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + before = asyncio.all_tasks() + monkeypatch.setattr(client_module, "DEFAULT_CLEANUP_TIMEOUT", 0.03) + entered, closed = asyncio.Event(), asyncio.Event() + timeouts: list[float] = [] + real_wait = asyncio.wait + + async def observe_wait(*args: Any, **kwargs: Any) -> Any: + timeouts.append(kwargs["timeout"]) + return await real_wait(*args, **kwargs) + + monkeypatch.setattr(client_module.asyncio, "wait", observe_wait) + + class Client(FakeTauBenchClient): + async def delete_environment(self, env_id: str) -> DeleteEnvironmentResponse: + entered.set() + try: + await asyncio.Event().wait() + finally: + closed.set() + raise AssertionError("unreachable") + + primary = ValueError("original") + + async def owner() -> None: + async with Client().environment(domain="telecom", task_id="task_001"): + raise primary + + task = asyncio.create_task(owner()) + try: + await asyncio.wait_for(entered.wait(), 1) + for _ in range(2): + task.cancel() + await asyncio.sleep(0) + with pytest.raises(BaseExceptionGroup) as caught: + await asyncio.wait_for(task, 1) + assert primary in caught.value.exceptions + assert any( + isinstance(e, asyncio.CancelledError) for e in caught.value.exceptions + ) + assert any("TimeoutError" in note for note in primary.__notes__) + assert len(timeouts) >= 3 and 0 < timeouts[-1] < timeouts[0] <= 0.03 + assert timeouts == sorted(timeouts, reverse=True) + await asyncio.wait_for(closed.wait(), 1) + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await asyncio.sleep(0) + assert asyncio.all_tasks() == before