From ffa221428953d621cd0e436331b0989f79700d1b Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Sat, 29 Aug 2026 14:00:08 +0800 Subject: [PATCH 01/16] fix(sandbox): settle PTY cleanup before cancellation --- .../extensions/sandbox/blaxel/sandbox.py | 6 +- .../extensions/sandbox/cloudflare/sandbox.py | 6 +- .../extensions/sandbox/daytona/sandbox.py | 6 +- src/agents/extensions/sandbox/e2b/sandbox.py | 6 +- .../extensions/sandbox/modal/sandbox.py | 6 +- src/agents/sandbox/sandboxes/docker.py | 6 +- src/agents/sandbox/sandboxes/unix_local.py | 6 +- .../sandbox/session/base_sandbox_session.py | 17 +++++ tests/sandbox/test_base_sandbox_session.py | 73 +++++++++++++++++++ 9 files changed, 111 insertions(+), 21 deletions(-) create mode 100644 tests/sandbox/test_base_sandbox_session.py diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index aee3211fbd..c17c7ddf79 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -847,7 +847,7 @@ async def pty_exec_start( raise _blaxel_exec_transport_error(command=command, cause=e) from e if pruned is not None: - await self._terminate_pty_entry(pruned) + await self._settle_pty_cleanup(self._terminate_pty_entry(pruned)) if process_count >= PTY_PROCESSES_WARNING: logger.warning( @@ -911,7 +911,7 @@ async def pty_terminate_all(self) -> None: self._pty_sessions.clear() self._reserved_pty_process_ids.clear() for entry in entries: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) # -- PTY internals ------------------------------------------------------- @@ -990,7 +990,7 @@ async def _finalize_pty_update( removed = self._pty_sessions.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._terminate_pty_entry(removed) + await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index bb8d7c37e6..4dec9d0b3b 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -1077,7 +1077,7 @@ async def _finalize_pty_update( removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._terminate_pty_entry(removed) + await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) live_process_id = None return PtyExecUpdate( @@ -1211,7 +1211,7 @@ async def pty_exec_start( raise ExecTransportError(command=tuple(str(part) for part in command), cause=e) from e if pruned_entry is not None: - await self._terminate_pty_entry(pruned_entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(pruned_entry)) if process_count >= PTY_PROCESSES_WARNING: logger.warning( @@ -1276,7 +1276,7 @@ async def pty_terminate_all(self) -> None: self._reserved_pty_process_ids.clear() for entry in entries: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) async def read(self, path: Path | str, *, user: str | User | None = None) -> io.IOBase: if user is not None: diff --git a/src/agents/extensions/sandbox/daytona/sandbox.py b/src/agents/extensions/sandbox/daytona/sandbox.py index d62c5021ad..34390f284a 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -746,7 +746,7 @@ async def _on_data(chunk: bytes | str) -> None: raise if pruned is not None: - await self._terminate_pty_entry(pruned) + await self._settle_pty_cleanup(self._terminate_pty_entry(pruned)) if process_count >= PTY_PROCESSES_WARNING: logger.warning( @@ -863,7 +863,7 @@ async def _finalize_pty_update( removed = self._pty_sessions.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._terminate_pty_entry(removed) + await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) live_process_id = None return PtyExecUpdate( @@ -879,7 +879,7 @@ async def pty_terminate_all(self) -> None: self._pty_sessions.clear() self._reserved_pty_process_ids.clear() for entry in entries: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) async def _collect_pty_output( self, diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index 389b665c44..bc8d9c3307 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -1041,7 +1041,7 @@ async def _append_output(payload: bytes | bytearray | str | object) -> None: ) if pruned_entry is not None: - await self._terminate_pty_entry(pruned_entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(pruned_entry)) if process_count >= PTY_PROCESSES_WARNING: logger.warning( @@ -1109,7 +1109,7 @@ async def pty_terminate_all(self) -> None: self._reserved_pty_process_ids.clear() for entry in entries: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase: if user is not None: @@ -1277,7 +1277,7 @@ async def _finalize_pty_update( removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._terminate_pty_entry(removed) + await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index 848b8e12a3..929d9e1bf7 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -898,7 +898,7 @@ async def pty_exec_start( raise _modal_exec_transport_error(command=command, cause=e) from e if pruned_entry is not None: - await self._terminate_pty_entry(pruned_entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(pruned_entry)) if process_count >= PTY_PROCESSES_WARNING: logger.warning( @@ -962,7 +962,7 @@ async def pty_terminate_all(self) -> None: self._reserved_pty_process_ids.clear() for entry in entries: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) async def _write_pty_stdin(self, process: ContainerProcess[bytes], payload: bytes) -> None: stdin = process.stdin @@ -1118,7 +1118,7 @@ async def _finalize_pty_update( removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._terminate_pty_entry(removed) + await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index 8ca4febe85..f3dd1123cf 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -1071,7 +1071,7 @@ async def pty_exec_start( raise if pruned_entry is not None: - await self._terminate_pty_entry(pruned_entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(pruned_entry)) if process_count >= PTY_PROCESSES_WARNING: logger.warning( @@ -1148,7 +1148,7 @@ async def pty_terminate_all(self) -> None: self._reserved_pty_process_ids.clear() for entry in entries: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) def _pump_pty_socket( self, entry: _DockerPtyProcessEntry, loop: asyncio.AbstractEventLoop @@ -1271,7 +1271,7 @@ async def _finalize_pty_update( removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._terminate_pty_entry(removed) + await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index ea7f83e9d8..a4ec7e5e08 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -391,7 +391,7 @@ def _preexec() -> None: process_count = len(self._pty_processes) if pruned_entry is not None: - await self._terminate_pty_entry(pruned_entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(pruned_entry)) if process_count >= PTY_PROCESSES_WARNING: logger.warning( @@ -464,7 +464,7 @@ async def pty_terminate_all(self) -> None: self._reserved_pty_process_ids.clear() for entry in entries: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) async def _resolved_exec_context(self) -> tuple[dict[str, str], str]: if self._host_environment_allowlist is None: @@ -559,7 +559,7 @@ async def _finalize_pty_update( removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._terminate_pty_entry(removed) + await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/sandbox/session/base_sandbox_session.py b/src/agents/sandbox/session/base_sandbox_session.py index d377bea9ef..cb79346cb3 100644 --- a/src/agents/sandbox/session/base_sandbox_session.py +++ b/src/agents/sandbox/session/base_sandbox_session.py @@ -711,6 +711,23 @@ def _resolve_pty_session_entry( raise PtySessionNotFoundError(session_id=session_id) return entry + async def _settle_pty_cleanup(self, operation: Awaitable[None]) -> None: + """Complete cleanup after PTY ownership leaves the session registry.""" + + task = asyncio.create_task(operation, name="agents.pty_cleanup") + completion = asyncio.create_task(asyncio.wait((task,))) + caller_cancelled = False + while not completion.done(): + try: + await asyncio.shield(completion) + except asyncio.CancelledError: + caller_cancelled = True + + completion.result() + task.result() + if caller_cancelled: + raise asyncio.CancelledError() + async def pty_exec_start( self, *command: str | Path, diff --git a/tests/sandbox/test_base_sandbox_session.py b/tests/sandbox/test_base_sandbox_session.py new file mode 100644 index 0000000000..0b5276efdb --- /dev/null +++ b/tests/sandbox/test_base_sandbox_session.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from agents.sandbox.session.base_sandbox_session import BaseSandboxSession + + +class _Session(BaseSandboxSession): + async def _exec_internal(self, *command: str, timeout: float | None = None): + raise NotImplementedError + + async def hydrate_workspace(self, *args, **kwargs): + raise NotImplementedError + + async def persist_workspace(self, *args, **kwargs): + raise NotImplementedError + + async def read(self, *args, **kwargs): + raise NotImplementedError + + async def running(self): + raise NotImplementedError + + async def write(self, *args, **kwargs): + raise NotImplementedError + + +def _session() -> BaseSandboxSession: + return _Session() + + +@pytest.mark.asyncio +async def test_pty_cleanup_completes_before_propagating_cancellation() -> None: + started = asyncio.Event() + release = asyncio.Event() + completed = False + + async def cleanup() -> None: + nonlocal completed + started.set() + await release.wait() + completed = True + + task = asyncio.create_task(_session()._settle_pty_cleanup(cleanup())) + await started.wait() + task.cancel() + task.cancel() + release.set() + + with pytest.raises(asyncio.CancelledError): + await task + assert completed + + +@pytest.mark.asyncio +async def test_pty_cleanup_preserves_cleanup_exception() -> None: + started = asyncio.Event() + release = asyncio.Event() + + async def cleanup() -> None: + started.set() + await release.wait() + raise RuntimeError("cleanup failed") + + task = asyncio.create_task(_session()._settle_pty_cleanup(cleanup())) + await started.wait() + task.cancel() + release.set() + + with pytest.raises(RuntimeError, match="cleanup failed"): + await task From 8fc7aab2103ea0cb9e28a45009a9e85709b7bb5c Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Sat, 29 Aug 2026 14:03:12 +0800 Subject: [PATCH 02/16] fix(sandbox): settle PTY cleanup batches --- .../extensions/sandbox/blaxel/sandbox.py | 8 ++++-- .../extensions/sandbox/cloudflare/sandbox.py | 7 +++-- .../extensions/sandbox/daytona/sandbox.py | 8 ++++-- src/agents/extensions/sandbox/e2b/sandbox.py | 7 +++-- .../extensions/sandbox/modal/sandbox.py | 7 +++-- src/agents/sandbox/sandboxes/docker.py | 7 +++-- src/agents/sandbox/sandboxes/unix_local.py | 7 +++-- tests/sandbox/test_base_sandbox_session.py | 27 +++++++++++++++++++ 8 files changed, 64 insertions(+), 14 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index c17c7ddf79..e7f52b89b9 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -910,8 +910,12 @@ async def pty_terminate_all(self) -> None: entries = list(self._pty_sessions.values()) self._pty_sessions.clear() self._reserved_pty_process_ids.clear() - for entry in entries: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + + async def cleanup_all() -> None: + for entry in entries: + await self._terminate_pty_entry(entry) + + await self._settle_pty_cleanup(cleanup_all()) # -- PTY internals ------------------------------------------------------- diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index 4dec9d0b3b..f56c93ca6c 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -1275,8 +1275,11 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - for entry in entries: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + async def cleanup_all() -> None: + for entry in entries: + await self._terminate_pty_entry(entry) + + await self._settle_pty_cleanup(cleanup_all()) async def read(self, path: Path | str, *, user: str | User | None = None) -> io.IOBase: if user is not None: diff --git a/src/agents/extensions/sandbox/daytona/sandbox.py b/src/agents/extensions/sandbox/daytona/sandbox.py index 34390f284a..5e8ec299f6 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -878,8 +878,12 @@ async def pty_terminate_all(self) -> None: entries = list(self._pty_sessions.values()) self._pty_sessions.clear() self._reserved_pty_process_ids.clear() - for entry in entries: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + + async def cleanup_all() -> None: + for entry in entries: + await self._terminate_pty_entry(entry) + + await self._settle_pty_cleanup(cleanup_all()) async def _collect_pty_output( self, diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index bc8d9c3307..d3874d82ed 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -1108,8 +1108,11 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - for entry in entries: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + async def cleanup_all() -> None: + for entry in entries: + await self._terminate_pty_entry(entry) + + await self._settle_pty_cleanup(cleanup_all()) async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase: if user is not None: diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index 929d9e1bf7..4bebe39d70 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -961,8 +961,11 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - for entry in entries: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + async def cleanup_all() -> None: + for entry in entries: + await self._terminate_pty_entry(entry) + + await self._settle_pty_cleanup(cleanup_all()) async def _write_pty_stdin(self, process: ContainerProcess[bytes], payload: bytes) -> None: stdin = process.stdin diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index f3dd1123cf..5a22b77355 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -1147,8 +1147,11 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - for entry in entries: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + async def cleanup_all() -> None: + for entry in entries: + await self._terminate_pty_entry(entry) + + await self._settle_pty_cleanup(cleanup_all()) def _pump_pty_socket( self, entry: _DockerPtyProcessEntry, loop: asyncio.AbstractEventLoop diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index a4ec7e5e08..26a003229e 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -463,8 +463,11 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - for entry in entries: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + async def cleanup_all() -> None: + for entry in entries: + await self._terminate_pty_entry(entry) + + await self._settle_pty_cleanup(cleanup_all()) async def _resolved_exec_context(self) -> tuple[dict[str, str], str]: if self._host_environment_allowlist is None: diff --git a/tests/sandbox/test_base_sandbox_session.py b/tests/sandbox/test_base_sandbox_session.py index 0b5276efdb..b4d8e34a0b 100644 --- a/tests/sandbox/test_base_sandbox_session.py +++ b/tests/sandbox/test_base_sandbox_session.py @@ -71,3 +71,30 @@ async def cleanup() -> None: with pytest.raises(RuntimeError, match="cleanup failed"): await task + + +@pytest.mark.asyncio +async def test_pty_cleanup_settles_a_sequential_batch() -> None: + started: list[int] = [] + release = asyncio.Event() + completed: list[int] = [] + + async def cleanup_all() -> None: + for entry in (1, 2): + started.append(entry) + await release.wait() + completed.append(entry) + release.clear() + if entry == 1: + release.set() + + task = asyncio.create_task(_session()._settle_pty_cleanup(cleanup_all())) + while started != [1]: + await asyncio.sleep(0) + task.cancel() + release.set() + + with pytest.raises(asyncio.CancelledError): + await task + assert started == [1, 2] + assert completed == [1, 2] From 3e9f727088726da5240f7067a34f0ad017ae939c Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Sat, 29 Aug 2026 15:10:16 +0800 Subject: [PATCH 03/16] test(sandbox): cover Blaxel PTY batch cleanup cancellation --- tests/extensions/sandbox/test_blaxel.py | 47 +++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index 3fe1d0d93a..e6e87ef69f 100644 --- a/tests/extensions/sandbox/test_blaxel.py +++ b/tests/extensions/sandbox/test_blaxel.py @@ -1975,6 +1975,53 @@ async def test_pty_terminate_all(self, fake_sandbox: _FakeSandboxInstance) -> No assert len(session._reserved_pty_process_ids) == 0 assert ws._closed + @pytest.mark.asyncio + async def test_pty_terminate_all_settles_after_registry_clear( + self, fake_sandbox: _FakeSandboxInstance + ) -> None: + from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry + + session = _make_session(fake_sandbox) + close_started = asyncio.Event() + release_close = asyncio.Event() + + class _BlockingCloseWS(_FakeWS): + async def close(self) -> None: + close_started.set() + await release_close.wait() + self._closed = True + + first_ws = _BlockingCloseWS() + second_ws = _FakeWS() + first_http = _FakeHTTPSession(first_ws) + second_http = _FakeHTTPSession(second_ws) + first = _BlaxelPtySessionEntry("first", first_ws, first_http) + second = _BlaxelPtySessionEntry("second", second_ws, second_http) + session._pty_sessions.update({1: first, 2: second}) + session._reserved_pty_process_ids.update({1, 2}) + + task = asyncio.create_task(session.pty_terminate_all()) + await close_started.wait() + + # Ownership has already left the registry before cleanup finishes. + assert session._pty_sessions == {} + assert session._reserved_pty_process_ids == set() + + task.cancel() + await asyncio.sleep(0) + task.cancel() # Exercise repeated caller cancellation. + await asyncio.sleep(0) + assert not first_ws._closed + assert not second_ws._closed + + release_close.set() + with pytest.raises(asyncio.CancelledError): + await task + + assert first_ws._closed and first_http._closed + assert second_ws._closed and second_http._closed + assert session._pty_sessions == {} + @pytest.mark.asyncio async def test_pty_ws_reader_error_message(self, fake_sandbox: _FakeSandboxInstance) -> None: from agents.extensions.sandbox.blaxel import sandbox as mod From 9ce15f766d710695b1e93b48a7e255b3551662ca Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Sat, 29 Aug 2026 15:15:59 +0800 Subject: [PATCH 04/16] fix(sandbox): preserve PTY cleanup cancellation reasons Signed-off-by: mikemikimike <13286568797@163.com> --- .../sandbox/session/base_sandbox_session.py | 10 +++++----- tests/sandbox/test_base_sandbox_session.py | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/src/agents/sandbox/session/base_sandbox_session.py b/src/agents/sandbox/session/base_sandbox_session.py index cb79346cb3..8309e35027 100644 --- a/src/agents/sandbox/session/base_sandbox_session.py +++ b/src/agents/sandbox/session/base_sandbox_session.py @@ -716,17 +716,17 @@ async def _settle_pty_cleanup(self, operation: Awaitable[None]) -> None: task = asyncio.create_task(operation, name="agents.pty_cleanup") completion = asyncio.create_task(asyncio.wait((task,))) - caller_cancelled = False + caller_cancellation: asyncio.CancelledError | None = None while not completion.done(): try: await asyncio.shield(completion) - except asyncio.CancelledError: - caller_cancelled = True + except asyncio.CancelledError as error: + caller_cancellation = caller_cancellation or error completion.result() task.result() - if caller_cancelled: - raise asyncio.CancelledError() + if caller_cancellation is not None: + raise caller_cancellation async def pty_exec_start( self, diff --git a/tests/sandbox/test_base_sandbox_session.py b/tests/sandbox/test_base_sandbox_session.py index b4d8e34a0b..233296fca0 100644 --- a/tests/sandbox/test_base_sandbox_session.py +++ b/tests/sandbox/test_base_sandbox_session.py @@ -54,6 +54,25 @@ async def cleanup() -> None: assert completed +@pytest.mark.asyncio +async def test_pty_cleanup_preserves_cancellation_reason() -> None: + started = asyncio.Event() + release = asyncio.Event() + + async def cleanup() -> None: + started.set() + await release.wait() + + task = asyncio.create_task(_session()._settle_pty_cleanup(cleanup())) + await started.wait() + task.cancel("caller stopped cleanup") + release.set() + + with pytest.raises(asyncio.CancelledError) as exc_info: + await task + assert exc_info.value.args == ("caller stopped cleanup",) + + @pytest.mark.asyncio async def test_pty_cleanup_preserves_cleanup_exception() -> None: started = asyncio.Event() From 68036ee68de81f748a222e3fc64c4a4ee88259ab Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Sat, 29 Aug 2026 15:25:58 +0800 Subject: [PATCH 05/16] fix(sandbox): preserve PTY cleanup cancellation reason --- tests/sandbox/test_base_sandbox_session.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/sandbox/test_base_sandbox_session.py b/tests/sandbox/test_base_sandbox_session.py index 233296fca0..75563461ca 100644 --- a/tests/sandbox/test_base_sandbox_session.py +++ b/tests/sandbox/test_base_sandbox_session.py @@ -45,13 +45,14 @@ async def cleanup() -> None: task = asyncio.create_task(_session()._settle_pty_cleanup(cleanup())) await started.wait() - task.cancel() - task.cancel() + task.cancel("cleanup requested") + task.cancel("cleanup requested again") release.set() - with pytest.raises(asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError) as exc_info: await task assert completed + assert exc_info.value.args == ("cleanup requested",) @pytest.mark.asyncio From 0708d65538ab16eb64981791772821c27d447d94 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Sat, 29 Aug 2026 16:52:53 +0800 Subject: [PATCH 06/16] test(sandbox): address PTY cleanup review feedback --- tests/extensions/sandbox/test_blaxel.py | 33 +++++++++++----------- tests/sandbox/test_base_sandbox_session.py | 7 +++-- 2 files changed, 22 insertions(+), 18 deletions(-) diff --git a/tests/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index e6e87ef69f..335fafa2ce 100644 --- a/tests/extensions/sandbox/test_blaxel.py +++ b/tests/extensions/sandbox/test_blaxel.py @@ -2001,22 +2001,23 @@ async def close(self) -> None: session._reserved_pty_process_ids.update({1, 2}) task = asyncio.create_task(session.pty_terminate_all()) - await close_started.wait() - - # Ownership has already left the registry before cleanup finishes. - assert session._pty_sessions == {} - assert session._reserved_pty_process_ids == set() - - task.cancel() - await asyncio.sleep(0) - task.cancel() # Exercise repeated caller cancellation. - await asyncio.sleep(0) - assert not first_ws._closed - assert not second_ws._closed - - release_close.set() - with pytest.raises(asyncio.CancelledError): - await task + try: + await asyncio.wait_for(close_started.wait(), timeout=5) + + # Ownership has already left the registry before cleanup finishes. + assert session._pty_sessions == {} + assert session._reserved_pty_process_ids == set() + + task.cancel() + await asyncio.sleep(0) + task.cancel() # Exercise repeated caller cancellation. + await asyncio.sleep(0) + assert not first_ws._closed + assert not second_ws._closed + finally: + release_close.set() + with pytest.raises(asyncio.CancelledError): + await task assert first_ws._closed and first_http._closed assert second_ws._closed and second_http._closed diff --git a/tests/sandbox/test_base_sandbox_session.py b/tests/sandbox/test_base_sandbox_session.py index 75563461ca..062fbb4bd5 100644 --- a/tests/sandbox/test_base_sandbox_session.py +++ b/tests/sandbox/test_base_sandbox_session.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import sys import pytest @@ -52,7 +53,8 @@ async def cleanup() -> None: with pytest.raises(asyncio.CancelledError) as exc_info: await task assert completed - assert exc_info.value.args == ("cleanup requested",) + if sys.version_info >= (3, 11): + assert exc_info.value.args == ("cleanup requested",) @pytest.mark.asyncio @@ -71,7 +73,8 @@ async def cleanup() -> None: with pytest.raises(asyncio.CancelledError) as exc_info: await task - assert exc_info.value.args == ("caller stopped cleanup",) + if sys.version_info >= (3, 11): + assert exc_info.value.args == ("caller stopped cleanup",) @pytest.mark.asyncio From 16960bd289c028cd74f88803ff481f70d77543fa Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Sat, 29 Aug 2026 17:04:08 +0800 Subject: [PATCH 07/16] test(sandbox): bound cleanup readiness waits --- tests/sandbox/test_base_sandbox_session.py | 96 ++++++++++++++-------- 1 file changed, 64 insertions(+), 32 deletions(-) diff --git a/tests/sandbox/test_base_sandbox_session.py b/tests/sandbox/test_base_sandbox_session.py index 062fbb4bd5..aaafd2244a 100644 --- a/tests/sandbox/test_base_sandbox_session.py +++ b/tests/sandbox/test_base_sandbox_session.py @@ -2,6 +2,7 @@ import asyncio import sys +from contextlib import suppress import pytest @@ -45,16 +46,23 @@ async def cleanup() -> None: completed = True task = asyncio.create_task(_session()._settle_pty_cleanup(cleanup())) - await started.wait() - task.cancel("cleanup requested") - task.cancel("cleanup requested again") - release.set() - - with pytest.raises(asyncio.CancelledError) as exc_info: - await task - assert completed - if sys.version_info >= (3, 11): - assert exc_info.value.args == ("cleanup requested",) + try: + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel("cleanup requested") + task.cancel("cleanup requested again") + release.set() + + with pytest.raises(asyncio.CancelledError) as exc_info: + await task + assert completed + if sys.version_info >= (3, 11): + assert exc_info.value.args == ("cleanup requested",) + finally: + release.set() + if not task.done(): + task.cancel() + with suppress(BaseException): + await task @pytest.mark.asyncio @@ -67,14 +75,21 @@ async def cleanup() -> None: await release.wait() task = asyncio.create_task(_session()._settle_pty_cleanup(cleanup())) - await started.wait() - task.cancel("caller stopped cleanup") - release.set() - - with pytest.raises(asyncio.CancelledError) as exc_info: - await task - if sys.version_info >= (3, 11): - assert exc_info.value.args == ("caller stopped cleanup",) + try: + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel("caller stopped cleanup") + release.set() + + with pytest.raises(asyncio.CancelledError) as exc_info: + await task + if sys.version_info >= (3, 11): + assert exc_info.value.args == ("caller stopped cleanup",) + finally: + release.set() + if not task.done(): + task.cancel() + with suppress(BaseException): + await task @pytest.mark.asyncio @@ -88,12 +103,19 @@ async def cleanup() -> None: raise RuntimeError("cleanup failed") task = asyncio.create_task(_session()._settle_pty_cleanup(cleanup())) - await started.wait() - task.cancel() - release.set() + try: + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel() + release.set() - with pytest.raises(RuntimeError, match="cleanup failed"): - await task + with pytest.raises(RuntimeError, match="cleanup failed"): + await task + finally: + release.set() + if not task.done(): + task.cancel() + with suppress(BaseException): + await task @pytest.mark.asyncio @@ -112,12 +134,22 @@ async def cleanup_all() -> None: release.set() task = asyncio.create_task(_session()._settle_pty_cleanup(cleanup_all())) - while started != [1]: - await asyncio.sleep(0) - task.cancel() - release.set() - - with pytest.raises(asyncio.CancelledError): - await task - assert started == [1, 2] - assert completed == [1, 2] + try: + async def wait_for_first_entry() -> None: + while started != [1]: + await asyncio.sleep(0) + + await asyncio.wait_for(wait_for_first_entry(), timeout=5) + task.cancel() + release.set() + + with pytest.raises(asyncio.CancelledError): + await task + assert started == [1, 2] + assert completed == [1, 2] + finally: + release.set() + if not task.done(): + task.cancel() + with suppress(BaseException): + await task From 3fb9e81f737e097993ef36ab316459c3751abc14 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Thu, 3 Sep 2026 22:44:17 +0800 Subject: [PATCH 08/16] fix(sandbox): settle every PTY cleanup after failure --- .../extensions/sandbox/blaxel/sandbox.py | 16 ++-- .../extensions/sandbox/cloudflare/sandbox.py | 18 ++-- .../extensions/sandbox/daytona/sandbox.py | 26 ++---- src/agents/extensions/sandbox/e2b/sandbox.py | 12 ++- .../extensions/sandbox/modal/sandbox.py | 14 ++-- src/agents/sandbox/sandboxes/docker.py | 20 +++-- src/agents/sandbox/sandboxes/unix_local.py | 27 +++--- .../sandbox/session/base_sandbox_session.py | 22 ++++- tests/extensions/sandbox/test_blaxel.py | 84 +++++++++++++++++++ tests/sandbox/test_base_sandbox_session.py | 16 ++++ 10 files changed, 184 insertions(+), 71 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index e7f52b89b9..c7326fdff7 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -839,11 +839,15 @@ async def pty_exec_start( registered = True except asyncio.TimeoutError as e: if not registered: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) raise ExecTimeoutError(command=command, timeout_s=exec_timeout, cause=e) from e + except asyncio.CancelledError: + if not registered: + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + raise except Exception as e: if not registered: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) raise _blaxel_exec_transport_error(command=command, cause=e) from e if pruned is not None: @@ -911,11 +915,9 @@ async def pty_terminate_all(self) -> None: self._pty_sessions.clear() self._reserved_pty_process_ids.clear() - async def cleanup_all() -> None: - for entry in entries: - await self._terminate_pty_entry(entry) - - await self._settle_pty_cleanup(cleanup_all()) + await self._settle_pty_cleanup( + self._cleanup_pty_entries(entries, self._terminate_pty_entry) + ) # -- PTY internals ------------------------------------------------------- diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index f56c93ca6c..a75447ab68 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -1119,10 +1119,14 @@ async def _cleanup_unregistered_pty( ) -> None: """Best-effort cleanup of a PTY WebSocket or entry that was never registered.""" if entry is not None and not registered: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) elif ws is not None and not registered: - with suppress(Exception): - await ws.close() + + async def close_ws() -> None: + with suppress(Exception): + await ws.close() + + await self._settle_pty_cleanup(close_ws()) async def pty_exec_start( self, @@ -1275,11 +1279,9 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - async def cleanup_all() -> None: - for entry in entries: - await self._terminate_pty_entry(entry) - - await self._settle_pty_cleanup(cleanup_all()) + await self._settle_pty_cleanup( + self._cleanup_pty_entries(entries, self._terminate_pty_entry) + ) async def read(self, path: Path | str, *, user: str | User | None = None) -> io.IOBase: if user is not None: diff --git a/src/agents/extensions/sandbox/daytona/sandbox.py b/src/agents/extensions/sandbox/daytona/sandbox.py index 5e8ec299f6..1b2ee6dc31 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -720,29 +720,17 @@ async def _on_data(chunk: bytes | str) -> None: registered = True except asyncio.TimeoutError as e: if not registered: - cleanup_task = asyncio.ensure_future(self._terminate_pty_entry(entry)) - try: - await asyncio.shield(cleanup_task) - except BaseException: - await asyncio.shield(cleanup_task) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) raise ExecTimeoutError(command=command, timeout_s=timeout, cause=e) from e except Exception as e: if not registered: - cleanup_task = asyncio.ensure_future(self._terminate_pty_entry(entry)) - try: - await asyncio.shield(cleanup_task) - except BaseException: - await asyncio.shield(cleanup_task) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) if timeout_error_types and isinstance(e, timeout_error_types): raise ExecTimeoutError(command=command, timeout_s=timeout, cause=e) from e raise _daytona_exec_transport_error(command=command, cause=e) from e except BaseException: if not registered: - cleanup_task = asyncio.ensure_future(self._terminate_pty_entry(entry)) - try: - await asyncio.shield(cleanup_task) - except BaseException: - await asyncio.shield(cleanup_task) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) raise if pruned is not None: @@ -879,11 +867,9 @@ async def pty_terminate_all(self) -> None: self._pty_sessions.clear() self._reserved_pty_process_ids.clear() - async def cleanup_all() -> None: - for entry in entries: - await self._terminate_pty_entry(entry) - - await self._settle_pty_cleanup(cleanup_all()) + await self._settle_pty_cleanup( + self._cleanup_pty_entries(entries, self._terminate_pty_entry) + ) async def _collect_pty_output( self, diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index d3874d82ed..e2bc4e3c8d 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -1026,11 +1026,11 @@ async def _append_output(payload: bytes | bytearray | str | object) -> None: registered = True except asyncio.CancelledError: if not registered and entry.handle is not None: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) raise except Exception as e: if not registered and entry.handle is not None: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) if isinstance(e, ExecTransportError): raise _raise_e2b_exec_error( @@ -1108,11 +1108,9 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - async def cleanup_all() -> None: - for entry in entries: - await self._terminate_pty_entry(entry) - - await self._settle_pty_cleanup(cleanup_all()) + await self._settle_pty_cleanup( + self._cleanup_pty_entries(entries, self._terminate_pty_entry) + ) async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase: if user is not None: diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index 4bebe39d70..aff3c9060a 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -884,15 +884,15 @@ async def pty_exec_start( process_count = len(self._pty_processes) except asyncio.TimeoutError as e: if entry is not None and not registered: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) raise ExecTimeoutError(command=command, timeout_s=timeout, cause=e) from e except asyncio.CancelledError: if entry is not None and not registered: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) raise except Exception as e: if entry is not None and not registered: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) if exception_chain_contains_type(e, _modal_exec_timeout_error_types()): raise ExecTimeoutError(command=command, timeout_s=timeout, cause=e) from e raise _modal_exec_transport_error(command=command, cause=e) from e @@ -961,11 +961,9 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - async def cleanup_all() -> None: - for entry in entries: - await self._terminate_pty_entry(entry) - - await self._settle_pty_cleanup(cleanup_all()) + await self._settle_pty_cleanup( + self._cleanup_pty_entries(entries, self._terminate_pty_entry) + ) async def _write_pty_stdin(self, process: ContainerProcess[bytes], payload: bytes) -> None: stdin = process.stdin diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index 5a22b77355..93a49a558c 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -1053,13 +1053,15 @@ async def pty_exec_start( registered = True except asyncio.TimeoutError as e: if entry is not None and not registered: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) elif pty_pid_path is not None: - await self._kill_pty_pid_path(pty_pid_path) + await self._settle_pty_cleanup(self._kill_pty_pid_path(pty_pid_path)) raise ExecTimeoutError(command=command, timeout_s=timeout, cause=e) from e except Exception as e: if entry is not None and not registered: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + elif pty_pid_path is not None: + await self._settle_pty_cleanup(self._kill_pty_pid_path(pty_pid_path)) raise ExecTransportError( command=command, context={"retry_safe": True}, @@ -1067,7 +1069,9 @@ async def pty_exec_start( ) from e except BaseException: if entry is not None and not registered: - await self._terminate_pty_entry(entry) + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + elif pty_pid_path is not None: + await self._settle_pty_cleanup(self._kill_pty_pid_path(pty_pid_path)) raise if pruned_entry is not None: @@ -1147,11 +1151,9 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - async def cleanup_all() -> None: - for entry in entries: - await self._terminate_pty_entry(entry) - - await self._settle_pty_cleanup(cleanup_all()) + await self._settle_pty_cleanup( + self._cleanup_pty_entries(entries, self._terminate_pty_entry) + ) def _pump_pty_socket( self, entry: _DockerPtyProcessEntry, loop: asyncio.AbstractEventLoop diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 26a003229e..058ea6c86a 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -383,12 +383,19 @@ def _preexec() -> None: entry.wait_task = asyncio.create_task(self._watch_process_exit(entry)) pruned_entry: _UnixPtyProcessEntry | None = None - async with self._pty_lock: - process_id = allocate_pty_process_id(self._reserved_pty_process_ids) - self._reserved_pty_process_ids.add(process_id) - pruned_entry = self._prune_pty_processes_if_needed() - self._pty_processes[process_id] = entry - process_count = len(self._pty_processes) + registered = False + try: + async with self._pty_lock: + process_id = allocate_pty_process_id(self._reserved_pty_process_ids) + self._reserved_pty_process_ids.add(process_id) + pruned_entry = self._prune_pty_processes_if_needed() + self._pty_processes[process_id] = entry + process_count = len(self._pty_processes) + registered = True + except BaseException: + if not registered: + await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + raise if pruned_entry is not None: await self._settle_pty_cleanup(self._terminate_pty_entry(pruned_entry)) @@ -463,11 +470,9 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - async def cleanup_all() -> None: - for entry in entries: - await self._terminate_pty_entry(entry) - - await self._settle_pty_cleanup(cleanup_all()) + await self._settle_pty_cleanup( + self._cleanup_pty_entries(entries, self._terminate_pty_entry) + ) async def _resolved_exec_context(self) -> tuple[dict[str, str], str]: if self._host_environment_allowlist is None: diff --git a/src/agents/sandbox/session/base_sandbox_session.py b/src/agents/sandbox/session/base_sandbox_session.py index 8309e35027..570cef4a54 100644 --- a/src/agents/sandbox/session/base_sandbox_session.py +++ b/src/agents/sandbox/session/base_sandbox_session.py @@ -714,7 +714,10 @@ def _resolve_pty_session_entry( async def _settle_pty_cleanup(self, operation: Awaitable[None]) -> None: """Complete cleanup after PTY ownership leaves the session registry.""" - task = asyncio.create_task(operation, name="agents.pty_cleanup") + async def run_operation() -> None: + await operation + + task = asyncio.create_task(run_operation(), name="agents.pty_cleanup") completion = asyncio.create_task(asyncio.wait((task,))) caller_cancellation: asyncio.CancelledError | None = None while not completion.done(): @@ -728,6 +731,23 @@ async def _settle_pty_cleanup(self, operation: Awaitable[None]) -> None: if caller_cancellation is not None: raise caller_cancellation + async def _cleanup_pty_entries( + self, + entries: Sequence[_PtyEntryT], + cleanup_entry: Callable[[_PtyEntryT], Awaitable[None]], + ) -> None: + """Attempt every PTY cleanup and re-raise the first failure.""" + + first_error: BaseException | None = None + for entry in entries: + try: + await cleanup_entry(entry) + except BaseException as error: + if first_error is None: + first_error = error + if first_error is not None: + raise first_error + async def pty_exec_start( self, *command: str | Path, diff --git a/tests/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index 335fafa2ce..056f529d83 100644 --- a/tests/extensions/sandbox/test_blaxel.py +++ b/tests/extensions/sandbox/test_blaxel.py @@ -8,6 +8,7 @@ import tarfile import time import uuid +from contextlib import suppress from dataclasses import FrozenInstanceError from pathlib import Path from typing import Any @@ -1906,6 +1907,54 @@ async def close(self) -> None: with pytest.raises(ExecTransportError): await session.pty_exec_start("echo", "hello") + @pytest.mark.asyncio + async def test_pty_exec_start_cancellation_cleans_unregistered_entry( + self, fake_sandbox: _FakeSandboxInstance + ) -> None: + from agents.extensions.sandbox.blaxel import sandbox as mod + + send_started = asyncio.Event() + release_send = asyncio.Event() + + class _BlockingSendWS(_FakeWS): + async def send_str(self, data: str) -> None: + self._sent.append(data) + send_started.set() + await release_send.wait() + + class _TrackingAiohttp(_FakeAiohttp): + def __init__(self, ws: _FakeWS) -> None: + super().__init__(ws=ws) + self.session: _FakeHTTPSession | None = None + + def ClientSession(self) -> _FakeHTTPSession: + self.session = _FakeHTTPSession(self._ws) + return self.session + + ws = _BlockingSendWS() + fake_aiohttp = _TrackingAiohttp(ws) + session = _make_session(fake_sandbox) + task = asyncio.create_task(session.pty_exec_start("echo", "hello")) + try: + with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp): + await asyncio.wait_for(send_started.wait(), timeout=5) + task.cancel("cancel setup") + release_send.set() + + with pytest.raises(asyncio.CancelledError): + await task + + assert ws._closed + assert fake_aiohttp.session is not None + assert fake_aiohttp.session._closed + assert session._pty_sessions == {} + finally: + release_send.set() + if not task.done(): + task.cancel() + with suppress(BaseException): + await task + @pytest.mark.parametrize( ("chars", "expected_send_count"), [ @@ -2023,6 +2072,41 @@ async def close(self) -> None: assert second_ws._closed and second_http._closed assert session._pty_sessions == {} + @pytest.mark.asyncio + async def test_pty_terminate_all_continues_after_cleanup_failure( + self, fake_sandbox: _FakeSandboxInstance + ) -> None: + from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry + + session = _make_session(fake_sandbox) + first_ws = _FakeWS() + second_ws = _FakeWS() + first_http = _FakeHTTPSession(first_ws) + second_http = _FakeHTTPSession(second_ws) + first = _BlaxelPtySessionEntry("first", first_ws, first_http) + second = _BlaxelPtySessionEntry("second", second_ws, second_http) + session._pty_sessions.update({1: first, 2: second}) + session._reserved_pty_process_ids.update({1, 2}) + cleanup_calls: list[str] = [] + + original_terminate = session._terminate_pty_entry + + async def terminate(entry: _BlaxelPtySessionEntry) -> None: + cleanup_calls.append(entry.ws_session_id) + if entry is first: + raise RuntimeError("first cleanup failed") + await original_terminate(entry) + + with patch.object(session, "_terminate_pty_entry", side_effect=terminate): + with pytest.raises(RuntimeError, match="first cleanup failed"): + await session.pty_terminate_all() + + assert cleanup_calls == ["first", "second"] + assert second_ws._closed + assert second_http._closed + assert session._pty_sessions == {} + assert session._reserved_pty_process_ids == set() + @pytest.mark.asyncio async def test_pty_ws_reader_error_message(self, fake_sandbox: _FakeSandboxInstance) -> None: from agents.extensions.sandbox.blaxel import sandbox as mod diff --git a/tests/sandbox/test_base_sandbox_session.py b/tests/sandbox/test_base_sandbox_session.py index aaafd2244a..d13ac8f7d1 100644 --- a/tests/sandbox/test_base_sandbox_session.py +++ b/tests/sandbox/test_base_sandbox_session.py @@ -135,6 +135,7 @@ async def cleanup_all() -> None: task = asyncio.create_task(_session()._settle_pty_cleanup(cleanup_all())) try: + async def wait_for_first_entry() -> None: while started != [1]: await asyncio.sleep(0) @@ -153,3 +154,18 @@ async def wait_for_first_entry() -> None: task.cancel() with suppress(BaseException): await task + + +@pytest.mark.asyncio +async def test_pty_cleanup_attempts_remaining_entries_after_failure() -> None: + attempted: list[int] = [] + + async def cleanup(entry: int) -> None: + attempted.append(entry) + if entry == 1: + raise RuntimeError("first cleanup failed") + + with pytest.raises(RuntimeError, match="first cleanup failed"): + await _session()._cleanup_pty_entries((1, 2), cleanup) + + assert attempted == [1, 2] From 10d6f2dc627a8959a602356bd67d98cf680e1083 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Thu, 3 Sep 2026 23:10:04 +0800 Subject: [PATCH 09/16] test(sandbox): keep PTY cleanup regression isolated --- .../extensions/sandbox/blaxel/sandbox.py | 8 +-- .../extensions/sandbox/cloudflare/sandbox.py | 10 +--- .../extensions/sandbox/daytona/sandbox.py | 18 +++++- src/agents/extensions/sandbox/e2b/sandbox.py | 4 +- .../extensions/sandbox/modal/sandbox.py | 6 +- src/agents/sandbox/sandboxes/docker.py | 12 ++-- src/agents/sandbox/sandboxes/unix_local.py | 19 ++----- tests/extensions/sandbox/test_blaxel.py | 56 +++---------------- 8 files changed, 42 insertions(+), 91 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index c7326fdff7..48b7933463 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -839,15 +839,11 @@ async def pty_exec_start( registered = True except asyncio.TimeoutError as e: if not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + await self._terminate_pty_entry(entry) raise ExecTimeoutError(command=command, timeout_s=exec_timeout, cause=e) from e - except asyncio.CancelledError: - if not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) - raise except Exception as e: if not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + await self._terminate_pty_entry(entry) raise _blaxel_exec_transport_error(command=command, cause=e) from e if pruned is not None: diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index a75447ab68..3772887227 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -1119,14 +1119,10 @@ async def _cleanup_unregistered_pty( ) -> None: """Best-effort cleanup of a PTY WebSocket or entry that was never registered.""" if entry is not None and not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + await self._terminate_pty_entry(entry) elif ws is not None and not registered: - - async def close_ws() -> None: - with suppress(Exception): - await ws.close() - - await self._settle_pty_cleanup(close_ws()) + with suppress(Exception): + await ws.close() async def pty_exec_start( self, diff --git a/src/agents/extensions/sandbox/daytona/sandbox.py b/src/agents/extensions/sandbox/daytona/sandbox.py index 1b2ee6dc31..613e4c2a48 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -720,17 +720,29 @@ async def _on_data(chunk: bytes | str) -> None: registered = True except asyncio.TimeoutError as e: if not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + cleanup_task = asyncio.ensure_future(self._terminate_pty_entry(entry)) + try: + await asyncio.shield(cleanup_task) + except BaseException: + await asyncio.shield(cleanup_task) raise ExecTimeoutError(command=command, timeout_s=timeout, cause=e) from e except Exception as e: if not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + cleanup_task = asyncio.ensure_future(self._terminate_pty_entry(entry)) + try: + await asyncio.shield(cleanup_task) + except BaseException: + await asyncio.shield(cleanup_task) if timeout_error_types and isinstance(e, timeout_error_types): raise ExecTimeoutError(command=command, timeout_s=timeout, cause=e) from e raise _daytona_exec_transport_error(command=command, cause=e) from e except BaseException: if not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + cleanup_task = asyncio.ensure_future(self._terminate_pty_entry(entry)) + try: + await asyncio.shield(cleanup_task) + except BaseException: + await asyncio.shield(cleanup_task) raise if pruned is not None: diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index e2bc4e3c8d..bc62619f8d 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -1026,11 +1026,11 @@ async def _append_output(payload: bytes | bytearray | str | object) -> None: registered = True except asyncio.CancelledError: if not registered and entry.handle is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + await self._terminate_pty_entry(entry) raise except Exception as e: if not registered and entry.handle is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + await self._terminate_pty_entry(entry) if isinstance(e, ExecTransportError): raise _raise_e2b_exec_error( diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index aff3c9060a..cf00fba437 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -884,15 +884,15 @@ async def pty_exec_start( process_count = len(self._pty_processes) except asyncio.TimeoutError as e: if entry is not None and not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + await self._terminate_pty_entry(entry) raise ExecTimeoutError(command=command, timeout_s=timeout, cause=e) from e except asyncio.CancelledError: if entry is not None and not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + await self._terminate_pty_entry(entry) raise except Exception as e: if entry is not None and not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + await self._terminate_pty_entry(entry) if exception_chain_contains_type(e, _modal_exec_timeout_error_types()): raise ExecTimeoutError(command=command, timeout_s=timeout, cause=e) from e raise _modal_exec_transport_error(command=command, cause=e) from e diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index 93a49a558c..c392f3cf55 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -1053,15 +1053,13 @@ async def pty_exec_start( registered = True except asyncio.TimeoutError as e: if entry is not None and not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) + await self._terminate_pty_entry(entry) elif pty_pid_path is not None: - await self._settle_pty_cleanup(self._kill_pty_pid_path(pty_pid_path)) + await self._kill_pty_pid_path(pty_pid_path) raise ExecTimeoutError(command=command, timeout_s=timeout, cause=e) from e except Exception as e: if entry is not None and not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) - elif pty_pid_path is not None: - await self._settle_pty_cleanup(self._kill_pty_pid_path(pty_pid_path)) + await self._terminate_pty_entry(entry) raise ExecTransportError( command=command, context={"retry_safe": True}, @@ -1069,9 +1067,7 @@ async def pty_exec_start( ) from e except BaseException: if entry is not None and not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) - elif pty_pid_path is not None: - await self._settle_pty_cleanup(self._kill_pty_pid_path(pty_pid_path)) + await self._terminate_pty_entry(entry) raise if pruned_entry is not None: diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 058ea6c86a..6b18f7afe6 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -383,19 +383,12 @@ def _preexec() -> None: entry.wait_task = asyncio.create_task(self._watch_process_exit(entry)) pruned_entry: _UnixPtyProcessEntry | None = None - registered = False - try: - async with self._pty_lock: - process_id = allocate_pty_process_id(self._reserved_pty_process_ids) - self._reserved_pty_process_ids.add(process_id) - pruned_entry = self._prune_pty_processes_if_needed() - self._pty_processes[process_id] = entry - process_count = len(self._pty_processes) - registered = True - except BaseException: - if not registered: - await self._settle_pty_cleanup(self._terminate_pty_entry(entry)) - raise + async with self._pty_lock: + process_id = allocate_pty_process_id(self._reserved_pty_process_ids) + self._reserved_pty_process_ids.add(process_id) + pruned_entry = self._prune_pty_processes_if_needed() + self._pty_processes[process_id] = entry + process_count = len(self._pty_processes) if pruned_entry is not None: await self._settle_pty_cleanup(self._terminate_pty_entry(pruned_entry)) diff --git a/tests/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index 056f529d83..7d943d91d9 100644 --- a/tests/extensions/sandbox/test_blaxel.py +++ b/tests/extensions/sandbox/test_blaxel.py @@ -1907,54 +1907,6 @@ async def close(self) -> None: with pytest.raises(ExecTransportError): await session.pty_exec_start("echo", "hello") - @pytest.mark.asyncio - async def test_pty_exec_start_cancellation_cleans_unregistered_entry( - self, fake_sandbox: _FakeSandboxInstance - ) -> None: - from agents.extensions.sandbox.blaxel import sandbox as mod - - send_started = asyncio.Event() - release_send = asyncio.Event() - - class _BlockingSendWS(_FakeWS): - async def send_str(self, data: str) -> None: - self._sent.append(data) - send_started.set() - await release_send.wait() - - class _TrackingAiohttp(_FakeAiohttp): - def __init__(self, ws: _FakeWS) -> None: - super().__init__(ws=ws) - self.session: _FakeHTTPSession | None = None - - def ClientSession(self) -> _FakeHTTPSession: - self.session = _FakeHTTPSession(self._ws) - return self.session - - ws = _BlockingSendWS() - fake_aiohttp = _TrackingAiohttp(ws) - session = _make_session(fake_sandbox) - task = asyncio.create_task(session.pty_exec_start("echo", "hello")) - try: - with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp): - await asyncio.wait_for(send_started.wait(), timeout=5) - task.cancel("cancel setup") - release_send.set() - - with pytest.raises(asyncio.CancelledError): - await task - - assert ws._closed - assert fake_aiohttp.session is not None - assert fake_aiohttp.session._closed - assert session._pty_sessions == {} - finally: - release_send.set() - if not task.done(): - task.cancel() - with suppress(BaseException): - await task - @pytest.mark.parametrize( ("chars", "expected_send_count"), [ @@ -2063,10 +2015,16 @@ async def close(self) -> None: await asyncio.sleep(0) assert not first_ws._closed assert not second_ws._closed - finally: + release_close.set() with pytest.raises(asyncio.CancelledError): await task + finally: + release_close.set() + if not task.done(): + task.cancel() + with suppress(BaseException): + await task assert first_ws._closed and first_http._closed assert second_ws._closed and second_http._closed From 54a942db5f8580417687df140247655a0868c009 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Tue, 8 Sep 2026 06:23:42 +0800 Subject: [PATCH 10/16] fix(sandbox): bound Docker PTY cleanup waits --- src/agents/sandbox/sandboxes/docker.py | 70 ++++++++++++++---------- tests/sandbox/test_docker.py | 74 ++++++++++++++++++++++++++ 2 files changed, 117 insertions(+), 27 deletions(-) diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index c392f3cf55..3445cac42f 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -97,6 +97,7 @@ # RAM and spill larger ones to a temp file so a big upload can't OOM the process. _STREAM_SPOOL_MAX_SIZE = 16 * 1024 * 1024 _DEFERRED_CLEANUP_TIMEOUT_S = 30.0 +_PTY_CLEANUP_TIMEOUT_S = 5.0 def _measure_stream(stream: io.IOBase) -> tuple[int, io.IOBase, io.IOBase | None]: @@ -500,9 +501,20 @@ async def _stage_workspace_copy( ) return staging_parent, staging_workspace - async def _rm_best_effort(self, path: Path) -> None: + async def _rm_best_effort(self, path: Path, *, timeout: float | None = None) -> None: try: - await self.exec("rm", "-rf", "--", sandbox_path_str(path), shell=False) + if timeout is None: + await self.exec("rm", "-rf", "--", sandbox_path_str(path), shell=False) + else: + command = ["rm", "-rf", "--", sandbox_path_str(path)] + await self._exec_run( + cmd=command, + workdir=self.state.manifest.root if self._workspace_root_ready else None, + user=None, + timeout=timeout, + command_for_errors=tuple(command), + kill_on_timeout=False, + ) except Exception: pass @@ -1223,9 +1235,12 @@ async def _refresh_pty_exit_code(self, entry: _DockerPtyProcessEntry) -> None: api = container_client.api try: - inspect_result = await loop.run_in_executor( - _DOCKER_EXECUTOR, - lambda: api.exec_inspect(entry.exec_id), + inspect_result = await asyncio.wait_for( + loop.run_in_executor( + _DOCKER_EXECUTOR, + lambda: api.exec_inspect(entry.exec_id), + ), + timeout=_PTY_CLEANUP_TIMEOUT_S, ) except Exception: return @@ -1306,7 +1321,7 @@ async def _terminate_pty_entry(self, entry: _DockerPtyProcessEntry) -> None: if entry.exit_code is None: await self._kill_pty_pid_path(entry.pid_path) else: - await self._rm_best_effort(entry.pid_path) + await self._rm_best_effort(entry.pid_path, timeout=_PTY_CLEANUP_TIMEOUT_S) try: cast(Any, entry.sock).close() @@ -1322,32 +1337,33 @@ async def _terminate_pty_entry(self, entry: _DockerPtyProcessEntry) -> None: ) async def _kill_pty_pid_path(self, pid_path: Path) -> None: - loop = asyncio.get_running_loop() + command = [ + "sh", + "-lc", + ( + 'if [ -f "$1" ]; then ' + 'pid="$(cat "$1" 2>/dev/null || true)"; ' + 'if [ -n "$pid" ]; then ' + 'kill -KILL "$pid" >/dev/null 2>&1 || true; ' + "fi; " + "fi" + ), + "sh", + sandbox_path_str(pid_path), + ] try: - await loop.run_in_executor( - _DOCKER_EXECUTOR, - lambda: self._container.exec_run( - cmd=[ - "sh", - "-lc", - ( - 'if [ -f "$1" ]; then ' - 'pid="$(cat "$1" 2>/dev/null || true)"; ' - 'if [ -n "$pid" ]; then ' - 'kill -KILL "$pid" >/dev/null 2>&1 || true; ' - "fi; " - "fi" - ), - "sh", - sandbox_path_str(pid_path), - ], - demux=True, - ), + await self._exec_run( + cmd=command, + workdir=None, + user=None, + timeout=_PTY_CLEANUP_TIMEOUT_S, + command_for_errors=("kill", sandbox_path_str(pid_path)), + kill_on_timeout=False, ) except Exception: pass - await self._rm_best_effort(pid_path) + await self._rm_best_effort(pid_path, timeout=_PTY_CLEANUP_TIMEOUT_S) async def exists(self) -> bool: try: diff --git a/tests/sandbox/test_docker.py b/tests/sandbox/test_docker.py index e4c7cc812f..91aa2b8e32 100644 --- a/tests/sandbox/test_docker.py +++ b/tests/sandbox/test_docker.py @@ -4702,6 +4702,80 @@ async def test_docker_pty_non_tty_rejects_stdin_and_stop_cleans_up( await session.pty_write_stdin(session_id=started.process_id, chars="") +@pytest.mark.asyncio +async def test_docker_pty_cleanup_bounds_stalled_backend_and_continues_batch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + api = _FakePtyApi() + container = _FakePtyContainer(api) + session = DockerSandboxSession( + docker_client=object(), + container=container, + state=DockerSandboxSessionState( + manifest=Manifest(root="/workspace"), + snapshot=NoopSnapshot(id="snapshot"), + image=DEFAULT_PYTHON_SANDBOX_IMAGE, + container_id="container", + workspace_root_ready=True, + ), + ) + first_socket = _FakePtySocket(api) + second_socket = _FakePtySocket(api) + first_entry = docker_sandbox._DockerPtyProcessEntry( + exec_id="exec-first", + sock=first_socket, + raw_sock=first_socket, + pid_path=Path("/tmp/first.pid"), + tty=False, + ) + second_entry = docker_sandbox._DockerPtyProcessEntry( + exec_id="exec-second", + sock=second_socket, + raw_sock=second_socket, + pid_path=Path("/tmp/second.pid"), + tty=False, + ) + session._pty_processes = {1: first_entry, 2: second_entry} + session._reserved_pty_process_ids = {1, 2} + + first_kill_started = threading.Event() + first_kill_finished = threading.Event() + release_first_kill = threading.Event() + original_exec_run = container.exec_run + + def stalled_first_kill( + cmd: list[str], + demux: bool = True, + workdir: str | None = None, + user: str = "", + ) -> object: + if cmd[:2] == ["sh", "-lc"] and cmd[-1] == "/tmp/first.pid": + first_kill_started.set() + try: + release_first_kill.wait() + finally: + first_kill_finished.set() + return original_exec_run(cmd, demux=demux, workdir=workdir, user=user) + + monkeypatch.setattr(docker_sandbox, "_PTY_CLEANUP_TIMEOUT_S", 0.01) + monkeypatch.setattr(container, "exec_run", stalled_first_kill) + + cleanup_task = asyncio.create_task(session.pty_terminate_all()) + await asyncio.wait_for(asyncio.to_thread(first_kill_started.wait), timeout=0.5) + cleanup_task.cancel() + + try: + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.shield(cleanup_task), timeout=0.5) + + assert first_socket.closed is True + assert second_socket.closed is True + assert session._pty_processes == {} + finally: + release_first_kill.set() + await asyncio.wait_for(asyncio.to_thread(first_kill_finished.wait), timeout=0.5) + + @pytest.mark.asyncio @pytest.mark.parametrize("operation", ["exec_create", "exec_start"]) async def test_docker_pty_exec_start_times_out_blocking_docker_startup( From 3678c6364af456f279a750691e7f84a0fde3df9b Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Tue, 8 Sep 2026 20:47:37 +0800 Subject: [PATCH 11/16] fix(sandbox): bound PTY cleanup settlement --- .../extensions/sandbox/blaxel/sandbox.py | 4 +- .../extensions/sandbox/cloudflare/sandbox.py | 4 +- .../extensions/sandbox/daytona/sandbox.py | 4 +- src/agents/extensions/sandbox/e2b/sandbox.py | 4 +- .../extensions/sandbox/modal/sandbox.py | 4 +- src/agents/sandbox/sandboxes/docker.py | 9 +- src/agents/sandbox/sandboxes/unix_local.py | 4 +- .../sandbox/session/base_sandbox_session.py | 83 ++++++++++++++++--- tests/sandbox/test_base_sandbox_session.py | 36 ++++++++ tests/sandbox/test_docker.py | 38 ++++----- 10 files changed, 135 insertions(+), 55 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index 48b7933463..01a13b5c13 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -911,9 +911,7 @@ async def pty_terminate_all(self) -> None: self._pty_sessions.clear() self._reserved_pty_process_ids.clear() - await self._settle_pty_cleanup( - self._cleanup_pty_entries(entries, self._terminate_pty_entry) - ) + await self._cleanup_pty_entries(entries, self._terminate_pty_entry) # -- PTY internals ------------------------------------------------------- diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index 3772887227..ccb007d5d2 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -1275,9 +1275,7 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - await self._settle_pty_cleanup( - self._cleanup_pty_entries(entries, self._terminate_pty_entry) - ) + await self._cleanup_pty_entries(entries, self._terminate_pty_entry) async def read(self, path: Path | str, *, user: str | User | None = None) -> io.IOBase: if user is not None: diff --git a/src/agents/extensions/sandbox/daytona/sandbox.py b/src/agents/extensions/sandbox/daytona/sandbox.py index 613e4c2a48..96d6473f89 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -879,9 +879,7 @@ async def pty_terminate_all(self) -> None: self._pty_sessions.clear() self._reserved_pty_process_ids.clear() - await self._settle_pty_cleanup( - self._cleanup_pty_entries(entries, self._terminate_pty_entry) - ) + await self._cleanup_pty_entries(entries, self._terminate_pty_entry) async def _collect_pty_output( self, diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index bc62619f8d..d75f86171c 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -1108,9 +1108,7 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - await self._settle_pty_cleanup( - self._cleanup_pty_entries(entries, self._terminate_pty_entry) - ) + await self._cleanup_pty_entries(entries, self._terminate_pty_entry) async def read(self, path: Path, *, user: str | User | None = None) -> io.IOBase: if user is not None: diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index cf00fba437..704c479b5a 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -961,9 +961,7 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - await self._settle_pty_cleanup( - self._cleanup_pty_entries(entries, self._terminate_pty_entry) - ) + await self._cleanup_pty_entries(entries, self._terminate_pty_entry) async def _write_pty_stdin(self, process: ContainerProcess[bytes], payload: bytes) -> None: stdin = process.stdin diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index 3445cac42f..acd661ee12 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -1159,8 +1159,10 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - await self._settle_pty_cleanup( - self._cleanup_pty_entries(entries, self._terminate_pty_entry) + await self._cleanup_pty_entries( + entries, + self._terminate_pty_entry, + timeout=_PTY_CLEANUP_TIMEOUT_S, ) def _pump_pty_socket( @@ -1346,6 +1348,7 @@ async def _kill_pty_pid_path(self, pid_path: Path) -> None: 'if [ -n "$pid" ]; then ' 'kill -KILL "$pid" >/dev/null 2>&1 || true; ' "fi; " + 'rm -f -- "$1" >/dev/null 2>&1 || true; ' "fi" ), "sh", @@ -1363,8 +1366,6 @@ async def _kill_pty_pid_path(self, pid_path: Path) -> None: except Exception: pass - await self._rm_best_effort(pid_path, timeout=_PTY_CLEANUP_TIMEOUT_S) - async def exists(self) -> bool: try: self._docker_client.containers.get(self.state.container_id) diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 6b18f7afe6..909c3189e0 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -463,9 +463,7 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - await self._settle_pty_cleanup( - self._cleanup_pty_entries(entries, self._terminate_pty_entry) - ) + await self._cleanup_pty_entries(entries, self._terminate_pty_entry) async def _resolved_exec_context(self) -> tuple[dict[str, str], str]: if self._host_environment_allowlist is None: diff --git a/src/agents/sandbox/session/base_sandbox_session.py b/src/agents/sandbox/session/base_sandbox_session.py index 570cef4a54..3c7fd4ffdb 100644 --- a/src/agents/sandbox/session/base_sandbox_session.py +++ b/src/agents/sandbox/session/base_sandbox_session.py @@ -53,6 +53,7 @@ _PtyEntryT = TypeVar("_PtyEntryT") _RUNTIME_HELPER_CACHE_KEY_UNSET = object() +_DEFAULT_PTY_CLEANUP_TIMEOUT_S = 5.0 _WORKSPACE_ROOT_PROBE_TIMEOUT_S = 10.0 _READ_PATH_PROBE_TIMEOUT_S = 10.0 _READ_PATH_PROBE_SCRIPT = """ @@ -223,6 +224,7 @@ class BaseSandboxSession(abc.ABC): _max_manifest_entry_concurrency: int | None = DEFAULT_MAX_MANIFEST_ENTRY_CONCURRENCY _max_local_dir_file_concurrency: int | None = DEFAULT_MAX_LOCAL_DIR_FILE_CONCURRENCY _archive_limits: SandboxArchiveLimits | None = None + _pty_cleanup_tasks: set[asyncio.Task[None]] | None = None def _runtime_has_protected_mount_authority(self) -> bool: """Return whether SDK-owned runtime state contains live mount authority.""" @@ -711,37 +713,92 @@ def _resolve_pty_session_entry( raise PtySessionNotFoundError(session_id=session_id) return entry - async def _settle_pty_cleanup(self, operation: Awaitable[None]) -> None: - """Complete cleanup after PTY ownership leaves the session registry.""" + def _pty_cleanup_timeout_s(self) -> float: + timeouts = getattr(getattr(self, "state", None), "timeouts", None) + timeout = getattr(timeouts, "cleanup_s", None) + if timeout is not None: + return float(timeout) + return _DEFAULT_PTY_CLEANUP_TIMEOUT_S + + def _track_pty_cleanup_task(self, task: asyncio.Task[None]) -> None: + tasks = self._pty_cleanup_tasks + if tasks is None: + tasks = set() + self._pty_cleanup_tasks = tasks + tasks.add(task) + + def forget_task(done: asyncio.Task[None]) -> None: + tasks.discard(done) + if not done.cancelled(): + done.exception() + + task.add_done_callback(forget_task) + + async def _settle_pty_cleanup( + self, + operation: Awaitable[None], + *, + timeout: float | None = None, + ) -> None: + """Settle cleanup after PTY ownership leaves the session registry. + + The cleanup task is independently owned so caller cancellation cannot + abandon it. A timeout bounds how long the caller waits while leaving + the provider operation running to finish its ordered cleanup. + """ async def run_operation() -> None: await operation task = asyncio.create_task(run_operation(), name="agents.pty_cleanup") + self._track_pty_cleanup_task(task) completion = asyncio.create_task(asyncio.wait((task,))) caller_cancellation: asyncio.CancelledError | None = None - while not completion.done(): - try: - await asyncio.shield(completion) - except asyncio.CancelledError as error: - caller_cancellation = caller_cancellation or error - - completion.result() - task.result() - if caller_cancellation is not None: - raise caller_cancellation + deadline = asyncio.get_running_loop().time() + ( + self._pty_cleanup_timeout_s() if timeout is None else timeout + ) + timed_out = False + try: + while not completion.done(): + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + timed_out = True + break + try: + await asyncio.wait_for(asyncio.shield(completion), timeout=remaining) + except asyncio.CancelledError as error: + caller_cancellation = caller_cancellation or error + except asyncio.TimeoutError: + timed_out = True + break + + if completion.done(): + completion.result() + task.result() + if caller_cancellation is not None: + raise caller_cancellation + elif caller_cancellation is not None: + raise caller_cancellation + elif timed_out: + raise asyncio.TimeoutError() + finally: + if not completion.done(): + completion.cancel() + await asyncio.gather(completion, return_exceptions=True) async def _cleanup_pty_entries( self, entries: Sequence[_PtyEntryT], cleanup_entry: Callable[[_PtyEntryT], Awaitable[None]], + *, + timeout: float | None = None, ) -> None: """Attempt every PTY cleanup and re-raise the first failure.""" first_error: BaseException | None = None for entry in entries: try: - await cleanup_entry(entry) + await self._settle_pty_cleanup(cleanup_entry(entry), timeout=timeout) except BaseException as error: if first_error is None: first_error = error diff --git a/tests/sandbox/test_base_sandbox_session.py b/tests/sandbox/test_base_sandbox_session.py index d13ac8f7d1..c04e3ea0d3 100644 --- a/tests/sandbox/test_base_sandbox_session.py +++ b/tests/sandbox/test_base_sandbox_session.py @@ -92,6 +92,42 @@ async def cleanup() -> None: await task +@pytest.mark.asyncio +async def test_pty_cleanup_timeout_preserves_cancellation_and_owned_task() -> None: + session = _session() + started = asyncio.Event() + release = asyncio.Event() + completed = asyncio.Event() + + async def cleanup() -> None: + started.set() + await release.wait() + completed.set() + + task = asyncio.create_task(session._settle_pty_cleanup(cleanup(), timeout=0.01)) + try: + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel("caller stopped cleanup") + + with pytest.raises(asyncio.CancelledError) as exc_info: + await task + if sys.version_info >= (3, 11): + assert exc_info.value.args == ("caller stopped cleanup",) + assert not completed.is_set() + assert session._pty_cleanup_tasks + + release.set() + await asyncio.wait_for(completed.wait(), timeout=0.5) + await asyncio.sleep(0) + assert session._pty_cleanup_tasks == set() + finally: + release.set() + if not task.done(): + task.cancel() + with suppress(BaseException): + await task + + @pytest.mark.asyncio async def test_pty_cleanup_preserves_cleanup_exception() -> None: started = asyncio.Event() diff --git a/tests/sandbox/test_docker.py b/tests/sandbox/test_docker.py index 91aa2b8e32..b71c9f8f11 100644 --- a/tests/sandbox/test_docker.py +++ b/tests/sandbox/test_docker.py @@ -12,6 +12,7 @@ import time import uuid from collections.abc import Callable, Iterator +from contextlib import suppress from pathlib import Path from typing import cast @@ -3633,6 +3634,7 @@ def _assert_pty_kill_call(call: dict[str, object]) -> None: 'if [ -f "$1" ]; then ' 'pid="$(cat "$1" 2>/dev/null || true)"; ' 'if [ -n "$pid" ]; then kill -KILL "$pid" >/dev/null 2>&1 || true; fi; ' + 'rm -f -- "$1" >/dev/null 2>&1 || true; ' "fi" ), ] @@ -4689,14 +4691,8 @@ async def test_docker_pty_non_tty_rejects_stdin_and_stop_cleans_up( await session.stop() assert api.socket.closed is True - assert len(container.exec_calls) == 2 + assert len(container.exec_calls) == 1 _assert_pty_kill_call(container.exec_calls[0]) - assert container.exec_calls[1]["cmd"] == [ - "rm", - "-rf", - "--", - cast(list[str], api.exec_create_calls[0]["cmd"])[5], - ] with pytest.raises(PtySessionNotFoundError): await session.pty_write_stdin(session_id=started.process_id, chars="") @@ -4760,20 +4756,28 @@ def stalled_first_kill( monkeypatch.setattr(docker_sandbox, "_PTY_CLEANUP_TIMEOUT_S", 0.01) monkeypatch.setattr(container, "exec_run", stalled_first_kill) - cleanup_task = asyncio.create_task(session.pty_terminate_all()) - await asyncio.wait_for(asyncio.to_thread(first_kill_started.wait), timeout=0.5) - cleanup_task.cancel() - + cleanup_task: asyncio.Task[None] | None = None try: + cleanup_task = asyncio.create_task(session.pty_terminate_all()) + await asyncio.wait_for(asyncio.to_thread(first_kill_started.wait), timeout=0.5) + cleanup_task.cancel() + with pytest.raises(asyncio.CancelledError): await asyncio.wait_for(asyncio.shield(cleanup_task), timeout=0.5) - assert first_socket.closed is True assert second_socket.closed is True assert session._pty_processes == {} finally: release_first_kill.set() - await asyncio.wait_for(asyncio.to_thread(first_kill_finished.wait), timeout=0.5) + if first_kill_started.is_set(): + await asyncio.wait_for(asyncio.to_thread(first_kill_finished.wait), timeout=0.5) + if cleanup_task is not None: + if not cleanup_task.done(): + cleanup_task.cancel() + with suppress(BaseException): + await cleanup_task + + assert first_socket.closed is True @pytest.mark.asyncio @@ -4813,14 +4817,8 @@ def _delayed_operation(*args: object, **kwargs: object) -> object: yield_time_s=0.01, ) - assert len(container.exec_calls) == 2 + assert len(container.exec_calls) == 1 _assert_pty_kill_call(container.exec_calls[0]) - assert container.exec_calls[1]["cmd"] == [ - "rm", - "-rf", - "--", - cast(list[str], container.exec_calls[0]["cmd"])[4], - ] @pytest.mark.asyncio From f5653ef8f6b23f5f5d4cedd0f09bf782db636f0c Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Tue, 8 Sep 2026 23:58:04 +0800 Subject: [PATCH 12/16] fix(sandbox): preserve PTY cleanup ownership --- .../extensions/sandbox/blaxel/sandbox.py | 13 +- .../extensions/sandbox/cloudflare/sandbox.py | 13 +- .../extensions/sandbox/daytona/sandbox.py | 13 +- src/agents/extensions/sandbox/e2b/sandbox.py | 13 +- .../extensions/sandbox/modal/sandbox.py | 13 +- src/agents/sandbox/sandboxes/docker.py | 25 ++- src/agents/sandbox/sandboxes/unix_local.py | 13 +- .../sandbox/session/base_sandbox_session.py | 107 ++++++++++-- tests/extensions/sandbox/test_blaxel.py | 2 +- tests/sandbox/test_base_sandbox_session.py | 154 ++++++++++++++++++ tests/sandbox/test_docker.py | 64 +++++++- 11 files changed, 406 insertions(+), 24 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index 01a13b5c13..f97c7df652 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -847,7 +847,18 @@ async def pty_exec_start( raise _blaxel_exec_transport_error(command=command, cause=e) from e if pruned is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(pruned)) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned), propagate_timeout=False + ) + except BaseException: + await self._rollback_pty_start( + process_id, + entry, + self._pty_sessions, + lambda: self._terminate_pty_entry(entry), + ) + raise if process_count >= PTY_PROCESSES_WARNING: logger.warning( diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index ccb007d5d2..f3dfe90882 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -1211,7 +1211,18 @@ async def pty_exec_start( raise ExecTransportError(command=tuple(str(part) for part in command), cause=e) from e if pruned_entry is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(pruned_entry)) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned_entry), propagate_timeout=False + ) + except BaseException: + await self._rollback_pty_start( + process_id, + entry, + self._pty_processes, + lambda: self._terminate_pty_entry(entry), + ) + raise if process_count >= PTY_PROCESSES_WARNING: logger.warning( diff --git a/src/agents/extensions/sandbox/daytona/sandbox.py b/src/agents/extensions/sandbox/daytona/sandbox.py index 96d6473f89..f659e75553 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -746,7 +746,18 @@ async def _on_data(chunk: bytes | str) -> None: raise if pruned is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(pruned)) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned), propagate_timeout=False + ) + except BaseException: + await self._rollback_pty_start( + process_id, + entry, + self._pty_sessions, + lambda: self._terminate_pty_entry(entry), + ) + raise if process_count >= PTY_PROCESSES_WARNING: logger.warning( diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index d75f86171c..05c2cb1933 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -1041,7 +1041,18 @@ async def _append_output(payload: bytes | bytearray | str | object) -> None: ) if pruned_entry is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(pruned_entry)) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned_entry), propagate_timeout=False + ) + except BaseException: + await self._rollback_pty_start( + process_id, + entry, + self._pty_processes, + lambda: self._terminate_pty_entry(entry), + ) + raise if process_count >= PTY_PROCESSES_WARNING: logger.warning( diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index 704c479b5a..7194ff69d5 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -898,7 +898,18 @@ async def pty_exec_start( raise _modal_exec_transport_error(command=command, cause=e) from e if pruned_entry is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(pruned_entry)) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned_entry), propagate_timeout=False + ) + except BaseException: + await self._rollback_pty_start( + process_id, + entry, + self._pty_processes, + lambda: self._terminate_pty_entry(entry), + ) + raise if process_count >= PTY_PROCESSES_WARNING: logger.warning( diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index acd661ee12..91c1818aef 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -100,6 +100,11 @@ _PTY_CLEANUP_TIMEOUT_S = 5.0 +def _consume_future_exception(future: asyncio.Future[Any]) -> None: + if not future.cancelled(): + future.exception() + + def _measure_stream(stream: io.IOBase) -> tuple[int, io.IOBase, io.IOBase | None]: """Return ``(length, readable_stream, spool_to_close)`` for a length-framed write. @@ -565,6 +570,7 @@ async def _exec_run( timeout: float | None, command_for_errors: tuple[str | Path, ...], kill_on_timeout: bool, + keep_running_on_timeout: bool = False, ) -> ExecResult: loop = asyncio.get_running_loop() future = loop.run_in_executor( @@ -576,9 +582,12 @@ async def _exec_run( user=user or "", ), ) + wait_target = asyncio.shield(future) if keep_running_on_timeout else future try: - exec_result = await asyncio.wait_for(future, timeout=timeout) + exec_result = await asyncio.wait_for(wait_target, timeout=timeout) except asyncio.TimeoutError as e: + if keep_running_on_timeout and not future.done(): + future.add_done_callback(_consume_future_exception) if kill_on_timeout: # Best-effort: kill processes matching the command line. # If this fails, the caller still gets a timeout error. @@ -1083,7 +1092,18 @@ async def pty_exec_start( raise if pruned_entry is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(pruned_entry)) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned_entry), propagate_timeout=False + ) + except BaseException: + await self._rollback_pty_start( + process_id, + entry, + self._pty_processes, + lambda: self._terminate_pty_entry(entry), + ) + raise if process_count >= PTY_PROCESSES_WARNING: logger.warning( @@ -1362,6 +1382,7 @@ async def _kill_pty_pid_path(self, pid_path: Path) -> None: timeout=_PTY_CLEANUP_TIMEOUT_S, command_for_errors=("kill", sandbox_path_str(pid_path)), kill_on_timeout=False, + keep_running_on_timeout=True, ) except Exception: pass diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 909c3189e0..ec7efaa11a 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -391,7 +391,18 @@ def _preexec() -> None: process_count = len(self._pty_processes) if pruned_entry is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(pruned_entry)) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned_entry), propagate_timeout=False + ) + except BaseException: + await self._rollback_pty_start( + process_id, + entry, + self._pty_processes, + lambda: self._terminate_pty_entry(entry), + ) + raise if process_count >= PTY_PROCESSES_WARNING: logger.warning( diff --git a/src/agents/sandbox/session/base_sandbox_session.py b/src/agents/sandbox/session/base_sandbox_session.py index 3c7fd4ffdb..4093ffdb50 100644 --- a/src/agents/sandbox/session/base_sandbox_session.py +++ b/src/agents/sandbox/session/base_sandbox_session.py @@ -3,8 +3,9 @@ import io import shlex from collections.abc import Awaitable, Callable, Mapping, Sequence +from contextlib import suppress from pathlib import Path, PurePath -from typing import Literal, NoReturn, TypeVar +from typing import Any, Literal, NoReturn, TypeVar from typing_extensions import Self @@ -225,6 +226,9 @@ class BaseSandboxSession(abc.ABC): _max_local_dir_file_concurrency: int | None = DEFAULT_MAX_LOCAL_DIR_FILE_CONCURRENCY _archive_limits: SandboxArchiveLimits | None = None _pty_cleanup_tasks: set[asyncio.Task[None]] | None = None + _pty_lock: asyncio.Lock + _pty_processes: dict[int, Any] + _reserved_pty_process_ids: set[int] def _runtime_has_protected_mount_authority(self) -> bool: """Return whether SDK-owned runtime state contains live mount authority.""" @@ -400,15 +404,38 @@ async def stop(self) -> None: try: try: await self._before_stop() + except BaseException as before_stop_error: + # Persist before re-raising cancellation or a cleanup deadline/error so the + # backend cannot be deleted with workspace state that exists only remotely. + await self._persist_snapshot_before_stop_error() + if isinstance(before_stop_error, Exception): + wrapped = self._wrap_stop_error(before_stop_error) + if wrapped is not before_stop_error: + raise wrapped from before_stop_error + raise + try: await self._persist_snapshot() - except Exception as e: - wrapped = self._wrap_stop_error(e) - if wrapped is e: + except Exception as error: + wrapped = self._wrap_stop_error(error) + if wrapped is error: raise - raise wrapped from e + raise wrapped from error finally: await self._after_stop() + async def _persist_snapshot_before_stop_error(self) -> None: + """Persist a snapshot even if stop is cancelled again while doing so.""" + + snapshot_task = asyncio.create_task( + self._persist_snapshot(), name="agents.persist_snapshot_after_stop_error" + ) + while not snapshot_task.done(): + try: + await asyncio.shield(snapshot_task) + except asyncio.CancelledError: + continue + snapshot_task.result() + async def _before_stop(self) -> None: """Run transient process cleanup before snapshot persistence.""" @@ -739,6 +766,7 @@ async def _settle_pty_cleanup( operation: Awaitable[None], *, timeout: float | None = None, + propagate_timeout: bool = True, ) -> None: """Settle cleanup after PTY ownership leaves the session registry. @@ -779,13 +807,34 @@ async def run_operation() -> None: raise caller_cancellation elif caller_cancellation is not None: raise caller_cancellation - elif timed_out: + elif timed_out and propagate_timeout: raise asyncio.TimeoutError() finally: if not completion.done(): completion.cancel() await asyncio.gather(completion, return_exceptions=True) + async def _rollback_pty_start( + self, + process_id: int, + entry: Any, + pty_registry: dict[int, Any], + terminate_entry: Callable[[], Awaitable[None]], + ) -> None: + """Remove and terminate a PTY whose start failed after registration.""" + + async with self._pty_lock: + if pty_registry.get(process_id) is not entry: + return + pty_registry.pop(process_id) + self._reserved_pty_process_ids.discard(process_id) + + with suppress(BaseException): + await self._settle_pty_cleanup( + terminate_entry(), + propagate_timeout=False, + ) + async def _cleanup_pty_entries( self, entries: Sequence[_PtyEntryT], @@ -795,15 +844,53 @@ async def _cleanup_pty_entries( ) -> None: """Attempt every PTY cleanup and re-raise the first failure.""" + loop = asyncio.get_running_loop() + batch_timeout = self._pty_cleanup_timeout_s() if timeout is None else timeout + deadline = loop.time() + batch_timeout + cleanup_tasks = [ + asyncio.create_task( + self._settle_pty_cleanup( + cleanup_entry(entry), + timeout=batch_timeout, + propagate_timeout=False, + ), + name="agents.pty_cleanup_batch", + ) + for entry in entries + ] + + def consume_cleanup_task_exception(task: asyncio.Task[None]) -> None: + if not task.cancelled(): + task.exception() + + for task in cleanup_tasks: + task.add_done_callback(consume_cleanup_task_exception) + + pending = set(cleanup_tasks) + caller_cancellation: asyncio.CancelledError | None = None first_error: BaseException | None = None - for entry in entries: + while pending: + remaining = deadline - loop.time() + if remaining <= 0: + break try: - await self._settle_pty_cleanup(cleanup_entry(entry), timeout=timeout) + done, pending = await asyncio.wait(pending, timeout=remaining) except BaseException as error: - if first_error is None: - first_error = error + if isinstance(error, asyncio.CancelledError): + caller_cancellation = caller_cancellation or error + continue + raise + + for task in done: + try: + task.result() + except BaseException as error: + if first_error is None: + first_error = error if first_error is not None: raise first_error + if caller_cancellation is not None: + raise caller_cancellation async def pty_exec_start( self, diff --git a/tests/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index 7d943d91d9..77c457995d 100644 --- a/tests/extensions/sandbox/test_blaxel.py +++ b/tests/extensions/sandbox/test_blaxel.py @@ -2014,7 +2014,7 @@ async def close(self) -> None: task.cancel() # Exercise repeated caller cancellation. await asyncio.sleep(0) assert not first_ws._closed - assert not second_ws._closed + assert second_ws._closed release_close.set() with pytest.raises(asyncio.CancelledError): diff --git a/tests/sandbox/test_base_sandbox_session.py b/tests/sandbox/test_base_sandbox_session.py index c04e3ea0d3..693754273e 100644 --- a/tests/sandbox/test_base_sandbox_session.py +++ b/tests/sandbox/test_base_sandbox_session.py @@ -3,9 +3,11 @@ import asyncio import sys from contextlib import suppress +from types import SimpleNamespace import pytest +from agents.sandbox.session import base_sandbox_session from agents.sandbox.session.base_sandbox_session import BaseSandboxSession @@ -128,6 +130,39 @@ async def cleanup() -> None: await task +@pytest.mark.asyncio +async def test_pty_cleanup_can_detach_after_timeout() -> None: + session = _session() + started = asyncio.Event() + release = asyncio.Event() + completed = asyncio.Event() + + async def cleanup() -> None: + started.set() + await release.wait() + completed.set() + + task = asyncio.create_task( + session._settle_pty_cleanup(cleanup(), timeout=0.01, propagate_timeout=False) + ) + try: + await asyncio.wait_for(started.wait(), timeout=5) + await asyncio.wait_for(task, timeout=0.5) + assert not completed.is_set() + assert session._pty_cleanup_tasks + + release.set() + await asyncio.wait_for(completed.wait(), timeout=0.5) + await asyncio.sleep(0) + assert session._pty_cleanup_tasks == set() + finally: + release.set() + if not task.done(): + task.cancel() + with suppress(BaseException): + await task + + @pytest.mark.asyncio async def test_pty_cleanup_preserves_cleanup_exception() -> None: started = asyncio.Event() @@ -205,3 +240,122 @@ async def cleanup(entry: int) -> None: await _session()._cleanup_pty_entries((1, 2), cleanup) assert attempted == [1, 2] + + +@pytest.mark.asyncio +async def test_pty_cleanup_batch_uses_one_deadline_and_starts_every_entry() -> None: + session = _session() + started: list[int] = [] + release = asyncio.Event() + completed: list[int] = [] + + async def cleanup(entry: int) -> None: + started.append(entry) + await release.wait() + completed.append(entry) + + batch = asyncio.create_task(session._cleanup_pty_entries((1, 2), cleanup, timeout=0.01)) + try: + + async def wait_for_all_started() -> None: + while started != [1, 2]: + await asyncio.sleep(0) + + await asyncio.wait_for(wait_for_all_started(), timeout=0.5) + await asyncio.wait_for(batch, timeout=0.5) + assert completed == [] + + release.set() + cleanup_tasks = tuple(session._pty_cleanup_tasks or ()) + await asyncio.wait_for(asyncio.gather(*cleanup_tasks), timeout=0.5) + finally: + release.set() + if not batch.done(): + batch.cancel() + with suppress(BaseException): + await batch + cleanup_tasks = tuple(session._pty_cleanup_tasks or ()) + if cleanup_tasks: + await asyncio.gather(*cleanup_tasks, return_exceptions=True) + + assert completed == [1, 2] + + +@pytest.mark.asyncio +async def test_pty_start_rollback_removes_and_terminates_exact_entry() -> None: + session = _session() + entry = object() + session._pty_lock = asyncio.Lock() + session._pty_processes = {7: entry} + session._reserved_pty_process_ids = {7} + terminated = False + + async def terminate() -> None: + nonlocal terminated + terminated = True + + await session._rollback_pty_start(7, entry, session._pty_processes, terminate) + + assert session._pty_processes == {} + assert session._reserved_pty_process_ids == set() + assert terminated + + +@pytest.mark.asyncio +async def test_pty_start_rollback_accepts_provider_session_registry() -> None: + session = _session() + entry = object() + registry = {7: entry} + session._pty_lock = asyncio.Lock() + session._reserved_pty_process_ids = {7} + terminated = False + + async def terminate() -> None: + nonlocal terminated + terminated = True + + await session._rollback_pty_start(7, entry, registry, terminate) + + assert registry == {} + assert session._reserved_pty_process_ids == set() + assert terminated + + +@pytest.mark.asyncio +async def test_stop_persists_snapshot_after_cleanup_cancellation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session = _session() + session.state = SimpleNamespace(manifest=object(), type="test") + monkeypatch.setattr( + base_sandbox_session, + "validate_manifest_mount_credential_boundaries", + lambda *args, **kwargs: None, + ) + persisted = False + snapshot_started = asyncio.Event() + release_snapshot = asyncio.Event() + + async def before_stop() -> None: + raise asyncio.CancelledError("cleanup cancelled") + + async def persist_snapshot() -> None: + nonlocal persisted + snapshot_started.set() + await release_snapshot.wait() + persisted = True + + session._before_stop = before_stop + session._persist_snapshot = persist_snapshot + + stop_task = asyncio.create_task(session.stop()) + await asyncio.wait_for(snapshot_started.wait(), timeout=0.5) + stop_task.cancel("second cleanup cancellation") + await asyncio.sleep(0) + assert not stop_task.done() + release_snapshot.set() + + with pytest.raises(asyncio.CancelledError): + await stop_task + + assert persisted diff --git a/tests/sandbox/test_docker.py b/tests/sandbox/test_docker.py index b71c9f8f11..925924c8ea 100644 --- a/tests/sandbox/test_docker.py +++ b/tests/sandbox/test_docker.py @@ -12,6 +12,7 @@ import time import uuid from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor from contextlib import suppress from pathlib import Path from typing import cast @@ -4771,15 +4772,68 @@ def stalled_first_kill( release_first_kill.set() if first_kill_started.is_set(): await asyncio.wait_for(asyncio.to_thread(first_kill_finished.wait), timeout=0.5) - if cleanup_task is not None: - if not cleanup_task.done(): - cleanup_task.cancel() - with suppress(BaseException): - await cleanup_task + if cleanup_task is not None: + if not cleanup_task.done(): + cleanup_task.cancel() + with suppress(BaseException): + await cleanup_task + async def wait_for_first_socket_close() -> None: + while not first_socket.closed: + await asyncio.sleep(0) + + await asyncio.wait_for(wait_for_first_socket_close(), timeout=0.5) assert first_socket.closed is True +@pytest.mark.asyncio +async def test_docker_pty_kill_remains_queued_after_cleanup_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + api = _FakePtyApi() + container = _FakePtyContainer(api) + session = DockerSandboxSession( + docker_client=object(), + container=container, + state=DockerSandboxSessionState( + manifest=Manifest(root="/workspace"), + snapshot=NoopSnapshot(id="snapshot"), + image=DEFAULT_PYTHON_SANDBOX_IMAGE, + container_id="container", + workspace_root_ready=True, + ), + ) + + executor = ThreadPoolExecutor(max_workers=1) + blocker_started = threading.Event() + release_blocker = threading.Event() + + def block_executor() -> None: + blocker_started.set() + release_blocker.wait() + + executor.submit(block_executor) + monkeypatch.setattr(docker_sandbox, "_DOCKER_EXECUTOR", executor) + monkeypatch.setattr(docker_sandbox, "_PTY_CLEANUP_TIMEOUT_S", 0.01) + + try: + await asyncio.wait_for(asyncio.to_thread(blocker_started.wait), timeout=0.5) + await session._kill_pty_pid_path(Path("/tmp/queued.pid")) + assert container.exec_calls == [] + + release_blocker.set() + + async def wait_for_kill() -> None: + while not container.exec_calls: + await asyncio.sleep(0) + + await asyncio.wait_for(wait_for_kill(), timeout=0.5) + _assert_pty_kill_call(container.exec_calls[0]) + finally: + release_blocker.set() + await asyncio.to_thread(executor.shutdown, True) + + @pytest.mark.asyncio @pytest.mark.parametrize("operation", ["exec_create", "exec_start"]) async def test_docker_pty_exec_start_times_out_blocking_docker_startup( From b684ea7a848f7d8610dad0a565ccac52b517edb9 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Wed, 9 Sep 2026 16:50:11 +0800 Subject: [PATCH 13/16] :bug: make PTY cleanup and snapshot fallback bounded --- .../extensions/sandbox/blaxel/sandbox.py | 6 +- .../extensions/sandbox/cloudflare/sandbox.py | 6 +- .../extensions/sandbox/daytona/sandbox.py | 6 +- src/agents/extensions/sandbox/e2b/sandbox.py | 6 +- .../extensions/sandbox/modal/sandbox.py | 6 +- src/agents/sandbox/runtime_session_manager.py | 6 +- src/agents/sandbox/sandboxes/docker.py | 25 ++++-- src/agents/sandbox/sandboxes/unix_local.py | 6 +- .../sandbox/session/base_sandbox_session.py | 79 ++++++++++++++++--- src/agents/sandbox/session/sandbox_session.py | 3 + tests/sandbox/test_base_sandbox_session.py | 36 ++++++++- tests/sandbox/test_docker.py | 10 +-- 12 files changed, 157 insertions(+), 38 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index f97c7df652..1ff6675218 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -849,7 +849,7 @@ async def pty_exec_start( if pruned is not None: try: await self._settle_pty_cleanup( - self._terminate_pty_entry(pruned), propagate_timeout=False + self._terminate_pty_entry(pruned), propagate_timeout=True ) except BaseException: await self._rollback_pty_start( @@ -1001,7 +1001,9 @@ async def _finalize_pty_update( removed = self._pty_sessions.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) + await self._settle_pty_cleanup( + self._terminate_pty_entry(removed), propagate_timeout=False + ) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index f3dfe90882..130ab38ff6 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -1077,7 +1077,9 @@ async def _finalize_pty_update( removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) + await self._settle_pty_cleanup( + self._terminate_pty_entry(removed), propagate_timeout=False + ) live_process_id = None return PtyExecUpdate( @@ -1213,7 +1215,7 @@ async def pty_exec_start( if pruned_entry is not None: try: await self._settle_pty_cleanup( - self._terminate_pty_entry(pruned_entry), propagate_timeout=False + self._terminate_pty_entry(pruned_entry), propagate_timeout=True ) except BaseException: await self._rollback_pty_start( diff --git a/src/agents/extensions/sandbox/daytona/sandbox.py b/src/agents/extensions/sandbox/daytona/sandbox.py index f659e75553..80765204ca 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -748,7 +748,7 @@ async def _on_data(chunk: bytes | str) -> None: if pruned is not None: try: await self._settle_pty_cleanup( - self._terminate_pty_entry(pruned), propagate_timeout=False + self._terminate_pty_entry(pruned), propagate_timeout=True ) except BaseException: await self._rollback_pty_start( @@ -874,7 +874,9 @@ async def _finalize_pty_update( removed = self._pty_sessions.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) + await self._settle_pty_cleanup( + self._terminate_pty_entry(removed), propagate_timeout=False + ) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index 05c2cb1933..616dcdd140 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -1043,7 +1043,7 @@ async def _append_output(payload: bytes | bytearray | str | object) -> None: if pruned_entry is not None: try: await self._settle_pty_cleanup( - self._terminate_pty_entry(pruned_entry), propagate_timeout=False + self._terminate_pty_entry(pruned_entry), propagate_timeout=True ) except BaseException: await self._rollback_pty_start( @@ -1287,7 +1287,9 @@ async def _finalize_pty_update( removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) + await self._settle_pty_cleanup( + self._terminate_pty_entry(removed), propagate_timeout=False + ) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index 7194ff69d5..85f3270fa8 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -900,7 +900,7 @@ async def pty_exec_start( if pruned_entry is not None: try: await self._settle_pty_cleanup( - self._terminate_pty_entry(pruned_entry), propagate_timeout=False + self._terminate_pty_entry(pruned_entry), propagate_timeout=True ) except BaseException: await self._rollback_pty_start( @@ -1128,7 +1128,9 @@ async def _finalize_pty_update( removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) + await self._settle_pty_cleanup( + self._terminate_pty_entry(removed), propagate_timeout=False + ) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/sandbox/runtime_session_manager.py b/src/agents/sandbox/runtime_session_manager.py index bc1a5379e9..565f3ef635 100644 --- a/src/agents/sandbox/runtime_session_manager.py +++ b/src/agents/sandbox/runtime_session_manager.py @@ -102,7 +102,11 @@ async def cleanup(self) -> None: cleanup_error = exc finally: try: - if self._client is not None and isinstance(self._session, SandboxSession): + if ( + self._client is not None + and isinstance(self._session, SandboxSession) + and not self._session._should_preserve_backend_on_cleanup() + ): await self._client.delete(self._session) except BaseException as exc: # pragma: no cover if cleanup_error is None: diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index 91c1818aef..fa35d92236 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -1094,7 +1094,7 @@ async def pty_exec_start( if pruned_entry is not None: try: await self._settle_pty_cleanup( - self._terminate_pty_entry(pruned_entry), propagate_timeout=False + self._terminate_pty_entry(pruned_entry), propagate_timeout=True ) except BaseException: await self._rollback_pty_start( @@ -1309,7 +1309,9 @@ async def _finalize_pty_update( removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) + await self._settle_pty_cleanup( + self._terminate_pty_entry(removed), propagate_timeout=False + ) live_process_id = None return PtyExecUpdate( @@ -1374,8 +1376,11 @@ async def _kill_pty_pid_path(self, pid_path: Path) -> None: "sh", sandbox_path_str(pid_path), ] - try: - await self._exec_run( + # Keep the whole executor operation independently owned. In particular, a kill queued + # behind all Docker workers must still start after this caller's deadline expires; a + # cancelled queued future would otherwise leave both the process and PID file orphaned. + kill_task = asyncio.create_task( + self._exec_run( cmd=command, workdir=None, user=None, @@ -1383,7 +1388,17 @@ async def _kill_pty_pid_path(self, pid_path: Path) -> None: command_for_errors=("kill", sandbox_path_str(pid_path)), kill_on_timeout=False, keep_running_on_timeout=True, - ) + ), + name="agents.docker_pty_kill", + ) + try: + await asyncio.wait_for(asyncio.shield(kill_task), timeout=_PTY_CLEANUP_TIMEOUT_S) + except asyncio.TimeoutError: + self._track_pty_cleanup_task(kill_task) + except asyncio.CancelledError: + if not kill_task.done(): + self._track_pty_cleanup_task(kill_task) + raise except Exception: pass diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index ec7efaa11a..e2924cf36b 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -393,7 +393,7 @@ def _preexec() -> None: if pruned_entry is not None: try: await self._settle_pty_cleanup( - self._terminate_pty_entry(pruned_entry), propagate_timeout=False + self._terminate_pty_entry(pruned_entry), propagate_timeout=True ) except BaseException: await self._rollback_pty_start( @@ -569,7 +569,9 @@ async def _finalize_pty_update( removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) if removed is not None: - await self._settle_pty_cleanup(self._terminate_pty_entry(removed)) + await self._settle_pty_cleanup( + self._terminate_pty_entry(removed), propagate_timeout=False + ) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/sandbox/session/base_sandbox_session.py b/src/agents/sandbox/session/base_sandbox_session.py index 4093ffdb50..2a404c7705 100644 --- a/src/agents/sandbox/session/base_sandbox_session.py +++ b/src/agents/sandbox/session/base_sandbox_session.py @@ -225,7 +225,11 @@ class BaseSandboxSession(abc.ABC): _max_manifest_entry_concurrency: int | None = DEFAULT_MAX_MANIFEST_ENTRY_CONCURRENCY _max_local_dir_file_concurrency: int | None = DEFAULT_MAX_LOCAL_DIR_FILE_CONCURRENCY _archive_limits: SandboxArchiveLimits | None = None - _pty_cleanup_tasks: set[asyncio.Task[None]] | None = None + _pty_cleanup_tasks: set[asyncio.Task[Any]] | None = None + # Set when a failed stop could not prove that the current workspace was persisted. Runner-owned + # cleanup must retain the backend in that case so it can be resumed instead of deleting the + # only remaining copy of the workspace. + _backend_preservation_required: bool = False _pty_lock: asyncio.Lock _pty_processes: dict[int, Any] _reserved_pty_process_ids: set[int] @@ -407,7 +411,14 @@ async def stop(self) -> None: except BaseException as before_stop_error: # Persist before re-raising cancellation or a cleanup deadline/error so the # backend cannot be deleted with workspace state that exists only remotely. - await self._persist_snapshot_before_stop_error() + self._backend_preservation_required = True + snapshot_error = await self._persist_snapshot_before_stop_error() + if snapshot_error is None: + self._backend_preservation_required = False + else: + # Keep the cleanup failure that caused stop() to fail as the primary error, + # while retaining the snapshot failure as diagnostic context. + raise before_stop_error from snapshot_error if isinstance(before_stop_error, Exception): wrapped = self._wrap_stop_error(before_stop_error) if wrapped is not before_stop_error: @@ -415,6 +426,7 @@ async def stop(self) -> None: raise try: await self._persist_snapshot() + self._backend_preservation_required = False except Exception as error: wrapped = self._wrap_stop_error(error) if wrapped is error: @@ -423,18 +435,52 @@ async def stop(self) -> None: finally: await self._after_stop() - async def _persist_snapshot_before_stop_error(self) -> None: - """Persist a snapshot even if stop is cancelled again while doing so.""" + async def _persist_snapshot_before_stop_error(self) -> BaseException | None: + """Persist a snapshot with a deadline without replacing the original stop failure.""" snapshot_task = asyncio.create_task( self._persist_snapshot(), name="agents.persist_snapshot_after_stop_error" ) - while not snapshot_task.done(): - try: - await asyncio.shield(snapshot_task) - except asyncio.CancelledError: - continue - snapshot_task.result() + completion = asyncio.create_task(asyncio.wait((snapshot_task,))) + caller_cancellation: asyncio.CancelledError | None = None + timed_out = False + deadline = asyncio.get_running_loop().time() + self._pty_cleanup_timeout_s() + try: + while not completion.done(): + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + timed_out = True + break + try: + await asyncio.wait_for(asyncio.shield(completion), timeout=remaining) + except asyncio.CancelledError as error: + caller_cancellation = caller_cancellation or error + except asyncio.TimeoutError: + timed_out = True + break + + if completion.done(): + completion.result() + try: + snapshot_task.result() + except BaseException as error: + return error + return None + + if timed_out: + # Keep the operation owned after the caller gives up waiting. The backend is + # retained because a snapshot that is still running cannot be treated as durable. + self._track_pty_cleanup_task(snapshot_task) + return asyncio.TimeoutError() + if caller_cancellation is not None: + # This branch is only reachable if the completion task became done between the + # loop condition and the cancellation; keep the original stop failure primary. + return caller_cancellation + return asyncio.TimeoutError() + finally: + if not completion.done(): + completion.cancel() + await asyncio.gather(completion, return_exceptions=True) async def _before_stop(self) -> None: """Run transient process cleanup before snapshot persistence.""" @@ -464,6 +510,11 @@ def supports_docker_volume_mounts(self) -> bool: def supports_pty(self) -> bool: return False + def _should_preserve_backend_on_cleanup(self) -> bool: + """Return whether cleanup must retain the provider backend for a later resume.""" + + return self._backend_preservation_required + @redact_mount_error_data async def shutdown(self) -> None: """ @@ -747,14 +798,14 @@ def _pty_cleanup_timeout_s(self) -> float: return float(timeout) return _DEFAULT_PTY_CLEANUP_TIMEOUT_S - def _track_pty_cleanup_task(self, task: asyncio.Task[None]) -> None: + def _track_pty_cleanup_task(self, task: asyncio.Task[Any]) -> None: tasks = self._pty_cleanup_tasks if tasks is None: tasks = set() self._pty_cleanup_tasks = tasks tasks.add(task) - def forget_task(done: asyncio.Task[None]) -> None: + def forget_task(done: asyncio.Task[Any]) -> None: tasks.discard(done) if not done.cancelled(): done.exception() @@ -859,7 +910,7 @@ async def _cleanup_pty_entries( for entry in entries ] - def consume_cleanup_task_exception(task: asyncio.Task[None]) -> None: + def consume_cleanup_task_exception(task: asyncio.Task[Any]) -> None: if not task.cancelled(): task.exception() @@ -891,6 +942,8 @@ def consume_cleanup_task_exception(task: asyncio.Task[None]) -> None: raise first_error if caller_cancellation is not None: raise caller_cancellation + if pending: + raise asyncio.TimeoutError() async def pty_exec_start( self, diff --git a/src/agents/sandbox/session/sandbox_session.py b/src/agents/sandbox/session/sandbox_session.py index 923f025857..1640d96783 100644 --- a/src/agents/sandbox/session/sandbox_session.py +++ b/src/agents/sandbox/session/sandbox_session.py @@ -264,6 +264,9 @@ def state(self, value: SandboxSessionState) -> None: # pragma: no cover def _runtime_has_protected_mount_authority(self) -> bool: return self._inner._runtime_has_protected_mount_authority() + def _should_preserve_backend_on_cleanup(self) -> bool: + return self._inner._should_preserve_backend_on_cleanup() + @property def dependencies(self) -> Dependencies: return self._inner.dependencies diff --git a/tests/sandbox/test_base_sandbox_session.py b/tests/sandbox/test_base_sandbox_session.py index 693754273e..0815c23787 100644 --- a/tests/sandbox/test_base_sandbox_session.py +++ b/tests/sandbox/test_base_sandbox_session.py @@ -1,12 +1,14 @@ from __future__ import annotations import asyncio +import inspect import sys from contextlib import suppress from types import SimpleNamespace import pytest +from agents.sandbox.manifest import Manifest from agents.sandbox.session import base_sandbox_session from agents.sandbox.session.base_sandbox_session import BaseSandboxSession @@ -262,7 +264,8 @@ async def wait_for_all_started() -> None: await asyncio.sleep(0) await asyncio.wait_for(wait_for_all_started(), timeout=0.5) - await asyncio.wait_for(batch, timeout=0.5) + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(batch, timeout=0.5) assert completed == [] release.set() @@ -326,7 +329,7 @@ async def test_stop_persists_snapshot_after_cleanup_cancellation( monkeypatch: pytest.MonkeyPatch, ) -> None: session = _session() - session.state = SimpleNamespace(manifest=object(), type="test") + session.state = SimpleNamespace(manifest=Manifest(), type="test") monkeypatch.setattr( base_sandbox_session, "validate_manifest_mount_credential_boundaries", @@ -359,3 +362,32 @@ async def persist_snapshot() -> None: await stop_task assert persisted + + +@pytest.mark.asyncio +async def test_stop_preserves_original_cleanup_failure_when_snapshot_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session = _session() + session.state = SimpleNamespace(manifest=Manifest(), type="test") + monkeypatch.setattr( + base_sandbox_session, + "validate_manifest_mount_credential_boundaries", + lambda *args, **kwargs: None, + ) + + async def before_stop() -> None: + raise RuntimeError("pty cleanup failed") + + async def persist_snapshot() -> None: + raise ValueError("snapshot failed") + + session._before_stop = before_stop + session._persist_snapshot = persist_snapshot + + with pytest.raises(RuntimeError) as exc_info: + await inspect.unwrap(BaseSandboxSession.stop)(session) + + assert str(exc_info.value) == "pty cleanup failed" + assert isinstance(exc_info.value.__cause__, ValueError) + assert session._should_preserve_backend_on_cleanup() diff --git a/tests/sandbox/test_docker.py b/tests/sandbox/test_docker.py index 925924c8ea..7d2bb45b1e 100644 --- a/tests/sandbox/test_docker.py +++ b/tests/sandbox/test_docker.py @@ -4772,11 +4772,11 @@ def stalled_first_kill( release_first_kill.set() if first_kill_started.is_set(): await asyncio.wait_for(asyncio.to_thread(first_kill_finished.wait), timeout=0.5) - if cleanup_task is not None: - if not cleanup_task.done(): - cleanup_task.cancel() - with suppress(BaseException): - await cleanup_task + if cleanup_task is not None: + if not cleanup_task.done(): + cleanup_task.cancel() + with suppress(BaseException): + await cleanup_task async def wait_for_first_socket_close() -> None: while not first_socket.closed: From 9d74e54299f549f079e10d68c104790efa30d547 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Wed, 9 Sep 2026 17:24:10 +0800 Subject: [PATCH 14/16] fix sandbox cleanup rollback and preservation --- src/agents/sandbox/runtime_session_manager.py | 30 ++++---- .../sandbox/session/base_sandbox_session.py | 32 +++++--- tests/sandbox/test_base_sandbox_session.py | 77 +++++++++++++++++++ tests/sandbox/test_runtime.py | 19 +++++ 4 files changed, 133 insertions(+), 25 deletions(-) diff --git a/src/agents/sandbox/runtime_session_manager.py b/src/agents/sandbox/runtime_session_manager.py index 565f3ef635..9e19cc84f0 100644 --- a/src/agents/sandbox/runtime_session_manager.py +++ b/src/agents/sandbox/runtime_session_manager.py @@ -95,28 +95,32 @@ async def cleanup(self) -> None: except BaseException as exc: # pragma: no cover if cleanup_error is None: cleanup_error = exc + preserve_backend = ( + isinstance(self._session, SandboxSession) + and self._session._should_preserve_backend_on_cleanup() + ) + if not preserve_backend: + try: + await self._session.shutdown() + except BaseException as exc: # pragma: no cover + if cleanup_error is None: + cleanup_error = exc try: - await self._session.shutdown() + if ( + self._client is not None + and isinstance(self._session, SandboxSession) + and not self._session._should_preserve_backend_on_cleanup() + ): + await self._client.delete(self._session) except BaseException as exc: # pragma: no cover if cleanup_error is None: cleanup_error = exc finally: try: - if ( - self._client is not None - and isinstance(self._session, SandboxSession) - and not self._session._should_preserve_backend_on_cleanup() - ): - await self._client.delete(self._session) + await self._session._aclose_dependencies() except BaseException as exc: # pragma: no cover if cleanup_error is None: cleanup_error = exc - finally: - try: - await self._session._aclose_dependencies() - except BaseException as exc: # pragma: no cover - if cleanup_error is None: - cleanup_error = exc if cleanup_error is not None: raise cleanup_error diff --git a/src/agents/sandbox/session/base_sandbox_session.py b/src/agents/sandbox/session/base_sandbox_session.py index 2a404c7705..f87bd6a55c 100644 --- a/src/agents/sandbox/session/base_sandbox_session.py +++ b/src/agents/sandbox/session/base_sandbox_session.py @@ -418,6 +418,10 @@ async def stop(self) -> None: else: # Keep the cleanup failure that caused stop() to fail as the primary error, # while retaining the snapshot failure as diagnostic context. + if isinstance(before_stop_error, Exception): + wrapped = self._wrap_stop_error(before_stop_error) + if wrapped is not before_stop_error: + raise wrapped from snapshot_error raise before_stop_error from snapshot_error if isinstance(before_stop_error, Exception): wrapped = self._wrap_stop_error(before_stop_error) @@ -874,15 +878,18 @@ async def _rollback_pty_start( ) -> None: """Remove and terminate a PTY whose start failed after registration.""" - async with self._pty_lock: - if pty_registry.get(process_id) is not entry: - return - pty_registry.pop(process_id) - self._reserved_pty_process_ids.discard(process_id) + async def rollback() -> None: + async with self._pty_lock: + if pty_registry.get(process_id) is not entry: + return + pty_registry.pop(process_id) + self._reserved_pty_process_ids.discard(process_id) + + await terminate_entry() with suppress(BaseException): await self._settle_pty_cleanup( - terminate_entry(), + rollback(), propagate_timeout=False, ) @@ -919,7 +926,7 @@ def consume_cleanup_task_exception(task: asyncio.Task[Any]) -> None: pending = set(cleanup_tasks) caller_cancellation: asyncio.CancelledError | None = None - first_error: BaseException | None = None + cleanup_errors: dict[int, BaseException] = {} while pending: remaining = deadline - loop.time() if remaining <= 0: @@ -932,14 +939,15 @@ def consume_cleanup_task_exception(task: asyncio.Task[Any]) -> None: continue raise - for task in done: + for index, task in enumerate(cleanup_tasks): + if task not in done: + continue try: task.result() except BaseException as error: - if first_error is None: - first_error = error - if first_error is not None: - raise first_error + cleanup_errors.setdefault(index, error) + if cleanup_errors: + raise cleanup_errors[min(cleanup_errors)] if caller_cancellation is not None: raise caller_cancellation if pending: diff --git a/tests/sandbox/test_base_sandbox_session.py b/tests/sandbox/test_base_sandbox_session.py index 0815c23787..5e43a68eab 100644 --- a/tests/sandbox/test_base_sandbox_session.py +++ b/tests/sandbox/test_base_sandbox_session.py @@ -244,6 +244,27 @@ async def cleanup(entry: int) -> None: assert attempted == [1, 2] +@pytest.mark.asyncio +async def test_pty_cleanup_raises_first_entry_error_deterministically() -> None: + second_failed = asyncio.Event() + release_first = asyncio.Event() + + async def cleanup(entry: int) -> None: + if entry == 1: + await release_first.wait() + else: + second_failed.set() + raise RuntimeError(f"cleanup {entry} failed") + + task = asyncio.create_task(_session()._cleanup_pty_entries((1, 2), cleanup)) + await second_failed.wait() + assert not task.done() + release_first.set() + + with pytest.raises(RuntimeError, match="cleanup 1 failed"): + await task + + @pytest.mark.asyncio async def test_pty_cleanup_batch_uses_one_deadline_and_starts_every_entry() -> None: session = _session() @@ -324,6 +345,31 @@ async def terminate() -> None: assert terminated +@pytest.mark.asyncio +async def test_pty_start_rollback_settles_registry_removal_before_cancellation() -> None: + session = _session() + entry = object() + registry = {7: entry} + session._pty_lock = asyncio.Lock() + session._reserved_pty_process_ids = {7} + await session._pty_lock.acquire() + terminated = False + + async def terminate() -> None: + nonlocal terminated + terminated = True + + task = asyncio.create_task(session._rollback_pty_start(7, entry, registry, terminate)) + await asyncio.sleep(0) + task.cancel("rollback cancelled") + session._pty_lock.release() + + await task + assert registry == {} + assert session._reserved_pty_process_ids == set() + assert terminated + + @pytest.mark.asyncio async def test_stop_persists_snapshot_after_cleanup_cancellation( monkeypatch: pytest.MonkeyPatch, @@ -391,3 +437,34 @@ async def persist_snapshot() -> None: assert str(exc_info.value) == "pty cleanup failed" assert isinstance(exc_info.value.__cause__, ValueError) assert session._should_preserve_backend_on_cleanup() + + +@pytest.mark.asyncio +async def test_stop_wraps_cleanup_failure_when_snapshot_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session = _session() + session.state = SimpleNamespace(manifest=Manifest(), type="test") + monkeypatch.setattr( + base_sandbox_session, + "validate_manifest_mount_credential_boundaries", + lambda *args, **kwargs: None, + ) + source_error = RuntimeError("cleanup failed") + snapshot_error = ValueError("snapshot failed") + + async def before_stop() -> None: + raise source_error + + async def persist_snapshot() -> None: + raise snapshot_error + + session._before_stop = before_stop + session._persist_snapshot = persist_snapshot + session._wrap_stop_error = lambda error: RuntimeError("wrapped cleanup failed") + + with pytest.raises(RuntimeError, match="wrapped cleanup failed") as exc_info: + await inspect.unwrap(BaseSandboxSession.stop)(session) + + assert exc_info.value.__cause__ is snapshot_error + assert session._should_preserve_backend_on_cleanup() diff --git a/tests/sandbox/test_runtime.py b/tests/sandbox/test_runtime.py index 7b4d7f540b..3e05c11326 100644 --- a/tests/sandbox/test_runtime.py +++ b/tests/sandbox/test_runtime.py @@ -608,6 +608,25 @@ async def failing_hook() -> None: assert session.close_dependency_calls == 1 +@pytest.mark.asyncio +async def test_runner_owned_cleanup_preserves_backend_before_shutdown() -> None: + inner = _FakeSession(Manifest()) + inner._backend_preservation_required = True + client = _FakeClient(inner) + resources = _SandboxSessionResources( + session=client.session, + client=client, + owns_session=True, + ) + + await resources.cleanup() + + assert inner.stop_calls == 1 + assert inner.shutdown_calls == 0 + assert client.delete_calls == 0 + assert inner.close_dependency_calls == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("runner_owned", [False, True]) async def test_pre_stop_cancellation_skips_persistence_and_completes_cleanup( From 475d3514edff868ed56489975b49ed80ce21f8ac Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Thu, 10 Sep 2026 00:38:30 +0800 Subject: [PATCH 15/16] fix(sandbox): drain detached cleanup before closing dependencies --- .../sandbox/session/base_sandbox_session.py | 26 ++++++- tests/sandbox/test_base_sandbox_session.py | 17 +++-- tests/sandbox/test_runtime.py | 75 ++++++++++++++++++- 3 files changed, 108 insertions(+), 10 deletions(-) diff --git a/src/agents/sandbox/session/base_sandbox_session.py b/src/agents/sandbox/session/base_sandbox_session.py index f87bd6a55c..430a51e861 100644 --- a/src/agents/sandbox/session/base_sandbox_session.py +++ b/src/agents/sandbox/session/base_sandbox_session.py @@ -655,11 +655,29 @@ async def _run_pre_stop_hooks(self) -> None: await self.run_pre_stop_hooks() async def _aclose_dependencies(self) -> None: + caller_cancellation = await self._wait_for_tracked_cleanup_tasks() dependencies = self._dependencies - if dependencies is None or self._dependencies_closed: - return - self._dependencies_closed = True - await dependencies.aclose() + if dependencies is not None and not self._dependencies_closed: + self._dependencies_closed = True + await dependencies.aclose() + if caller_cancellation is not None: + raise caller_cancellation + + async def _wait_for_tracked_cleanup_tasks(self) -> asyncio.CancelledError | None: + """Wait for detached cleanup before closing dependencies it may still use.""" + + caller_cancellation: asyncio.CancelledError | None = None + while True: + tasks = tuple(task for task in (self._pty_cleanup_tasks or ()) if not task.done()) + if not tasks: + break + completion = asyncio.gather(*tasks, return_exceptions=True) + try: + await asyncio.shield(completion) + except asyncio.CancelledError as error: + caller_cancellation = caller_cancellation or error + + return caller_cancellation @staticmethod def _workspace_relpaths_overlap(lhs: Path, rhs: Path) -> bool: diff --git a/tests/sandbox/test_base_sandbox_session.py b/tests/sandbox/test_base_sandbox_session.py index 5e43a68eab..619a269777 100644 --- a/tests/sandbox/test_base_sandbox_session.py +++ b/tests/sandbox/test_base_sandbox_session.py @@ -257,12 +257,19 @@ async def cleanup(entry: int) -> None: raise RuntimeError(f"cleanup {entry} failed") task = asyncio.create_task(_session()._cleanup_pty_entries((1, 2), cleanup)) - await second_failed.wait() - assert not task.done() - release_first.set() + try: + await asyncio.wait_for(second_failed.wait(), timeout=0.5) + assert not task.done() + release_first.set() - with pytest.raises(RuntimeError, match="cleanup 1 failed"): - await task + with pytest.raises(RuntimeError, match="cleanup 1 failed"): + await task + finally: + release_first.set() + if not task.done(): + task.cancel() + with suppress(BaseException): + await task @pytest.mark.asyncio diff --git a/tests/sandbox/test_runtime.py b/tests/sandbox/test_runtime.py index 3e05c11326..84becfc9a7 100644 --- a/tests/sandbox/test_runtime.py +++ b/tests/sandbox/test_runtime.py @@ -99,7 +99,7 @@ from agents.sandbox.session.sandbox_client import BaseSandboxClient from agents.sandbox.session.sandbox_session import SandboxSession from agents.sandbox.session.sandbox_session_state import SandboxSessionState -from agents.sandbox.snapshot import LocalSnapshotSpec, NoopSnapshot, SnapshotBase +from agents.sandbox.snapshot import LocalSnapshotSpec, NoopSnapshot, RemoteSnapshot, SnapshotBase from agents.sandbox.types import ExecResult from agents.stream_events import RunItemStreamEvent from agents.testing import ScriptedModel, scripted_sandbox_session @@ -627,6 +627,79 @@ async def test_runner_owned_cleanup_preserves_backend_before_shutdown() -> None: assert inner.close_dependency_calls == 1 +@pytest.mark.asyncio +async def test_runner_owned_cleanup_waits_for_detached_snapshot_before_closing_dependencies() -> ( + None +): + upload_started = asyncio.Event() + release_upload = asyncio.Event() + upload_finished = asyncio.Event() + + class _RemoteSnapshotClient: + closed = False + closed_before_upload = False + + async def upload(self, snapshot_id: str, data: io.IOBase) -> None: + _ = (snapshot_id, data) + upload_started.set() + await release_upload.wait() + upload_finished.set() + + async def aclose(self) -> None: + self.closed = True + self.closed_before_upload = not upload_finished.is_set() + + class _DetachedSnapshotSession(_FakeSession): + def __init__(self, manifest: Manifest) -> None: + super().__init__(manifest) + self.state.snapshot = RemoteSnapshot( + id="detached", + client_dependency_key="remote_snapshot_client", + ) + + def _pty_cleanup_timeout_s(self) -> float: + return 0.01 + + async def _before_stop(self) -> None: + raise asyncio.CancelledError("pty cleanup failed") + + async def stop(self) -> None: + self.stop_calls += 1 + self._running = False + await BaseSandboxSession.stop(self) + + snapshot_client = _RemoteSnapshotClient() + inner = _DetachedSnapshotSession(Manifest()) + inner.set_dependencies( + Dependencies().bind_factory( + "remote_snapshot_client", + lambda _dependencies: snapshot_client, + owns_result=True, + ) + ) + client = _FakeClient(inner) + resources = _SandboxSessionResources( + session=client.session, + client=client, + owns_session=True, + ) + + cleanup = asyncio.create_task(resources.cleanup()) + await asyncio.wait_for(upload_started.wait(), timeout=0.5) + await asyncio.sleep(0.05) + assert not cleanup.done() + assert not snapshot_client.closed_before_upload + + release_upload.set() + with pytest.raises(asyncio.CancelledError): + await cleanup + + assert upload_finished.is_set() + assert snapshot_client.closed + assert not snapshot_client.closed_before_upload + assert inner.close_dependency_calls == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("runner_owned", [False, True]) async def test_pre_stop_cancellation_skips_persistence_and_completes_cleanup( From 2e9e1e48a42f6ea664b2325ea6933095c94260fb Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Thu, 10 Sep 2026 02:16:48 +0800 Subject: [PATCH 16/16] fix(sandbox): bound deferred cleanup finalization --- src/agents/sandbox/runtime_session_manager.py | 64 ++++++++++++++++++- src/agents/sandbox/sandboxes/docker.py | 33 +++++----- .../sandbox/session/base_sandbox_session.py | 62 ++++++++++++++++-- src/agents/sandbox/session/sandbox_session.py | 9 +++ tests/sandbox/test_runtime.py | 31 ++++++--- 5 files changed, 168 insertions(+), 31 deletions(-) diff --git a/src/agents/sandbox/runtime_session_manager.py b/src/agents/sandbox/runtime_session_manager.py index 9e19cc84f0..b51a18369a 100644 --- a/src/agents/sandbox/runtime_session_manager.py +++ b/src/agents/sandbox/runtime_session_manager.py @@ -57,6 +57,7 @@ def __init__( self._cleanup_lock = asyncio.Lock() self._cleaned = False self._started = False + self._deferred_cleanup_task: asyncio.Task[Any] | None = None @property def session(self) -> BaseSandboxSession: @@ -75,6 +76,50 @@ async def ensure_started(self) -> None: await self._session.start() self._started = True + def _schedule_deferred_cleanup(self) -> None: + task = self._deferred_cleanup_task + if task is not None and not task.done(): + return + + task = asyncio.create_task( + self._finish_deferred_cleanup(), + name="agents.deferred_session_cleanup", + ) + self._deferred_cleanup_task = task + + def consume_task_exception(done: asyncio.Task[Any]) -> None: + if not done.cancelled(): + done.exception() + + task.add_done_callback(consume_task_exception) + + async def _finish_deferred_cleanup(self) -> None: + try: + while True: + await self._session._wait_for_tracked_cleanup_tasks() + if self._session._should_preserve_backend_on_cleanup(): + return + + try: + await self._session.shutdown() + except BaseException: + if self._session._has_pending_pty_cleanup_tasks(): + continue + + if self._session._has_pending_pty_cleanup_tasks(): + continue + if self._session._should_preserve_backend_on_cleanup(): + return + if self._client is not None and isinstance(self._session, SandboxSession): + await self._client.delete(self._session) + return + finally: + if not self._session._has_pending_pty_cleanup_tasks(): + try: + await self._session._aclose_dependencies() + except BaseException: + pass + @redact_mount_error_data async def cleanup(self) -> None: if not self._owns_session: @@ -105,22 +150,39 @@ async def cleanup(self) -> None: except BaseException as exc: # pragma: no cover if cleanup_error is None: cleanup_error = exc + pending_cleanup_after_shutdown = self._session._has_pending_pty_cleanup_tasks() + preserve_backend = ( + preserve_backend + or pending_cleanup_after_shutdown + or ( + isinstance(self._session, SandboxSession) + and self._session._should_preserve_backend_on_cleanup() + ) + ) try: if ( self._client is not None and isinstance(self._session, SandboxSession) - and not self._session._should_preserve_backend_on_cleanup() + and not preserve_backend ): await self._client.delete(self._session) except BaseException as exc: # pragma: no cover if cleanup_error is None: cleanup_error = exc finally: + pending_cleanup_before_dependencies = ( + pending_cleanup_after_shutdown or self._session._has_pending_pty_cleanup_tasks() + ) try: await self._session._aclose_dependencies() except BaseException as exc: # pragma: no cover if cleanup_error is None: cleanup_error = exc + if ( + pending_cleanup_before_dependencies + or self._session._has_pending_pty_cleanup_tasks() + ): + self._schedule_deferred_cleanup() if cleanup_error is not None: raise cleanup_error diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index fa35d92236..445ab8e82a 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -1340,25 +1340,26 @@ async def _terminate_pty_entry(self, entry: _DockerPtyProcessEntry) -> None: if entry.wait_task is not None: entry.wait_task.cancel() - await self._refresh_pty_exit_code(entry) - - if entry.exit_code is None: - await self._kill_pty_pid_path(entry.pid_path) - else: - await self._rm_best_effort(entry.pid_path, timeout=_PTY_CLEANUP_TIMEOUT_S) - try: - cast(Any, entry.sock).close() - except Exception: - pass + await self._refresh_pty_exit_code(entry) + + if entry.exit_code is None: + await self._kill_pty_pid_path(entry.pid_path) + else: + await self._rm_best_effort(entry.pid_path, timeout=_PTY_CLEANUP_TIMEOUT_S) + finally: + try: + cast(Any, entry.sock).close() + except Exception: + pass - if entry.reader_thread is not None: - await asyncio.to_thread(entry.reader_thread.join, 1.0) + if entry.reader_thread is not None: + await asyncio.to_thread(entry.reader_thread.join, 1.0) - await asyncio.gather( - *(task for task in (entry.wait_task,) if task is not None), - return_exceptions=True, - ) + await asyncio.gather( + *(task for task in (entry.wait_task,) if task is not None), + return_exceptions=True, + ) async def _kill_pty_pid_path(self, pid_path: Path) -> None: command = [ diff --git a/src/agents/sandbox/session/base_sandbox_session.py b/src/agents/sandbox/session/base_sandbox_session.py index 430a51e861..a5ceaf1bd7 100644 --- a/src/agents/sandbox/session/base_sandbox_session.py +++ b/src/agents/sandbox/session/base_sandbox_session.py @@ -226,6 +226,7 @@ class BaseSandboxSession(abc.ABC): _max_local_dir_file_concurrency: int | None = DEFAULT_MAX_LOCAL_DIR_FILE_CONCURRENCY _archive_limits: SandboxArchiveLimits | None = None _pty_cleanup_tasks: set[asyncio.Task[Any]] | None = None + _deferred_dependency_close_task: asyncio.Task[Any] | None = None # Set when a failed stop could not prove that the current workspace was persisted. Runner-owned # cleanup must retain the backend in that case so it can be resumed instead of deleting the # only remaining copy of the workspace. @@ -445,6 +446,15 @@ async def _persist_snapshot_before_stop_error(self) -> BaseException | None: snapshot_task = asyncio.create_task( self._persist_snapshot(), name="agents.persist_snapshot_after_stop_error" ) + + def mark_snapshot_durable(task: asyncio.Task[Any]) -> None: + try: + task.result() + except BaseException: + return + self._backend_preservation_required = False + + snapshot_task.add_done_callback(mark_snapshot_durable) completion = asyncio.create_task(asyncio.wait((snapshot_task,))) caller_cancellation: asyncio.CancelledError | None = None timed_out = False @@ -586,9 +596,12 @@ async def _aclose_impl(self) -> None: if cleanup_error is None: cleanup_error = exc finally: + pending_cleanup_before_dependencies = self._has_pending_pty_cleanup_tasks() try: await self._aclose_dependencies() except BaseException as exc: + if pending_cleanup_before_dependencies or self._has_pending_pty_cleanup_tasks(): + self._schedule_deferred_dependency_close() if cleanup_error is None: cleanup_error = exc if cleanup_error is not None: @@ -655,7 +668,13 @@ async def _run_pre_stop_hooks(self) -> None: await self.run_pre_stop_hooks() async def _aclose_dependencies(self) -> None: - caller_cancellation = await self._wait_for_tracked_cleanup_tasks() + caller_cancellation, timed_out = await self._wait_for_tracked_cleanup_tasks( + timeout=self._pty_cleanup_timeout_s() + ) + if timed_out: + if caller_cancellation is not None: + raise caller_cancellation + raise asyncio.TimeoutError() dependencies = self._dependencies if dependencies is not None and not self._dependencies_closed: self._dependencies_closed = True @@ -663,21 +682,52 @@ async def _aclose_dependencies(self) -> None: if caller_cancellation is not None: raise caller_cancellation - async def _wait_for_tracked_cleanup_tasks(self) -> asyncio.CancelledError | None: - """Wait for detached cleanup before closing dependencies it may still use.""" + def _has_pending_pty_cleanup_tasks(self) -> bool: + return any(task for task in (self._pty_cleanup_tasks or ()) if not task.done()) + + async def _wait_for_tracked_cleanup_tasks( + self, *, timeout: float | None = None + ) -> tuple[asyncio.CancelledError | None, bool]: + """Wait for detached cleanup without cancelling tasks that still own provider resources.""" caller_cancellation: asyncio.CancelledError | None = None + deadline = ( + None if timeout is None else asyncio.get_running_loop().time() + max(timeout, 0.0) + ) while True: tasks = tuple(task for task in (self._pty_cleanup_tasks or ()) if not task.done()) if not tasks: break - completion = asyncio.gather(*tasks, return_exceptions=True) + remaining = None if deadline is None else deadline - asyncio.get_running_loop().time() + if remaining is not None and remaining <= 0: + return caller_cancellation, True try: - await asyncio.shield(completion) + await asyncio.wait(tasks, timeout=remaining) except asyncio.CancelledError as error: caller_cancellation = caller_cancellation or error - return caller_cancellation + return caller_cancellation, False + + def _schedule_deferred_dependency_close(self) -> None: + task = self._deferred_dependency_close_task + if task is not None and not task.done(): + return + + task = asyncio.create_task( + self._finish_deferred_dependency_close(), + name="agents.deferred_dependency_close", + ) + self._deferred_dependency_close_task = task + + def consume_task_exception(done: asyncio.Task[Any]) -> None: + if not done.cancelled(): + done.exception() + + task.add_done_callback(consume_task_exception) + + async def _finish_deferred_dependency_close(self) -> None: + await self._wait_for_tracked_cleanup_tasks() + await self._aclose_dependencies() @staticmethod def _workspace_relpaths_overlap(lhs: Path, rhs: Path) -> bool: diff --git a/src/agents/sandbox/session/sandbox_session.py b/src/agents/sandbox/session/sandbox_session.py index 1640d96783..069fa63c85 100644 --- a/src/agents/sandbox/session/sandbox_session.py +++ b/src/agents/sandbox/session/sandbox_session.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import io import ipaddress import time @@ -267,6 +268,14 @@ def _runtime_has_protected_mount_authority(self) -> bool: def _should_preserve_backend_on_cleanup(self) -> bool: return self._inner._should_preserve_backend_on_cleanup() + def _has_pending_pty_cleanup_tasks(self) -> bool: + return self._inner._has_pending_pty_cleanup_tasks() + + async def _wait_for_tracked_cleanup_tasks( + self, *, timeout: float | None = None + ) -> tuple[asyncio.CancelledError | None, bool]: + return await self._inner._wait_for_tracked_cleanup_tasks(timeout=timeout) + @property def dependencies(self) -> Dependencies: return self._inner.dependencies diff --git a/tests/sandbox/test_runtime.py b/tests/sandbox/test_runtime.py index 84becfc9a7..c60c441c75 100644 --- a/tests/sandbox/test_runtime.py +++ b/tests/sandbox/test_runtime.py @@ -12,6 +12,7 @@ import tempfile import uuid from collections.abc import Sequence +from contextlib import suppress from pathlib import Path, PurePosixPath, PureWindowsPath from typing import Any, ClassVar, Literal, TypedDict, cast @@ -685,19 +686,33 @@ async def stop(self) -> None: ) cleanup = asyncio.create_task(resources.cleanup()) - await asyncio.wait_for(upload_started.wait(), timeout=0.5) - await asyncio.sleep(0.05) - assert not cleanup.done() - assert not snapshot_client.closed_before_upload + try: + await asyncio.wait_for(upload_started.wait(), timeout=0.5) + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(cleanup, timeout=0.5) + + assert not upload_finished.is_set() + assert not snapshot_client.closed + assert inner.shutdown_calls == 0 + assert client.delete_calls == 0 + finally: + release_upload.set() + if not cleanup.done(): + with suppress(BaseException): + await asyncio.wait_for(cleanup, timeout=0.5) - release_upload.set() - with pytest.raises(asyncio.CancelledError): - await cleanup + async def wait_for_deferred_cleanup() -> None: + while not snapshot_client.closed: + await asyncio.sleep(0) + + await asyncio.wait_for(wait_for_deferred_cleanup(), timeout=0.5) assert upload_finished.is_set() assert snapshot_client.closed assert not snapshot_client.closed_before_upload - assert inner.close_dependency_calls == 1 + assert inner.shutdown_calls == 1 + assert client.delete_calls == 1 + assert inner.close_dependency_calls == 2 @pytest.mark.asyncio