diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index aee3211fbd..1ff6675218 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._terminate_pty_entry(pruned) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned), propagate_timeout=True + ) + 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( @@ -910,8 +921,8 @@ 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._terminate_pty_entry(entry) + + await self._cleanup_pty_entries(entries, self._terminate_pty_entry) # -- PTY internals ------------------------------------------------------- @@ -990,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._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 bb8d7c37e6..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._terminate_pty_entry(removed) + await self._settle_pty_cleanup( + self._terminate_pty_entry(removed), propagate_timeout=False + ) live_process_id = None return PtyExecUpdate( @@ -1211,7 +1213,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._terminate_pty_entry(pruned_entry) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned_entry), propagate_timeout=True + ) + 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( @@ -1275,8 +1288,7 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - for entry in entries: - await self._terminate_pty_entry(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 d62c5021ad..80765204ca 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._terminate_pty_entry(pruned) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned), propagate_timeout=True + ) + 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( @@ -863,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._terminate_pty_entry(removed) + await self._settle_pty_cleanup( + self._terminate_pty_entry(removed), propagate_timeout=False + ) live_process_id = None return PtyExecUpdate( @@ -878,8 +891,8 @@ 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._terminate_pty_entry(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 389b665c44..616dcdd140 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._terminate_pty_entry(pruned_entry) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned_entry), propagate_timeout=True + ) + 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( @@ -1108,8 +1119,7 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - for entry in entries: - await self._terminate_pty_entry(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: @@ -1277,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._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 848b8e12a3..85f3270fa8 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._terminate_pty_entry(pruned_entry) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned_entry), propagate_timeout=True + ) + 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( @@ -961,8 +972,7 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - for entry in entries: - await self._terminate_pty_entry(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 @@ -1118,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._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..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: @@ -95,24 +140,49 @@ 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 + 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: - await self._session.shutdown() + if ( + self._client is not None + and isinstance(self._session, SandboxSession) + 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: - if self._client is not None and isinstance(self._session, SandboxSession): - 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 ( + 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 8ca4febe85..445ab8e82a 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -97,6 +97,12 @@ # 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 _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]: @@ -500,9 +506,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 @@ -553,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( @@ -564,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. @@ -1071,7 +1092,18 @@ async def pty_exec_start( raise if pruned_entry is not None: - await self._terminate_pty_entry(pruned_entry) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned_entry), propagate_timeout=True + ) + 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( @@ -1147,8 +1179,11 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - for entry in entries: - await self._terminate_pty_entry(entry) + await self._cleanup_pty_entries( + entries, + self._terminate_pty_entry, + timeout=_PTY_CLEANUP_TIMEOUT_S, + ) def _pump_pty_socket( self, entry: _DockerPtyProcessEntry, loop: asyncio.AbstractEventLoop @@ -1222,9 +1257,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 @@ -1271,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._terminate_pty_entry(removed) + await self._settle_pty_cleanup( + self._terminate_pty_entry(removed), propagate_timeout=False + ) live_process_id = None return PtyExecUpdate( @@ -1300,54 +1340,69 @@ 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) - 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: - 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; " + 'rm -f -- "$1" >/dev/null 2>&1 || true; ' + "fi" + ), + "sh", + sandbox_path_str(pid_path), + ] + # 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, + timeout=_PTY_CLEANUP_TIMEOUT_S, + 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 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 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 - await self._rm_best_effort(pid_path) - 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 ea7f83e9d8..e2924cf36b 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._terminate_pty_entry(pruned_entry) + try: + await self._settle_pty_cleanup( + self._terminate_pty_entry(pruned_entry), propagate_timeout=True + ) + 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( @@ -463,8 +474,7 @@ async def pty_terminate_all(self) -> None: self._pty_processes.clear() self._reserved_pty_process_ids.clear() - for entry in entries: - await self._terminate_pty_entry(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: @@ -559,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._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 d377bea9ef..a5ceaf1bd7 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 @@ -53,6 +54,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 +225,15 @@ 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[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. + _backend_preservation_required: bool = False + _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.""" @@ -398,15 +409,93 @@ 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. + 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. + 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) + 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: + self._backend_preservation_required = False + 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) -> 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" + ) + + 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 + 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.""" @@ -435,6 +524,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: """ @@ -502,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: @@ -571,11 +668,66 @@ async def _run_pre_stop_hooks(self) -> None: await self.run_pre_stop_hooks() async def _aclose_dependencies(self) -> None: + 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 None or self._dependencies_closed: + 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 + + 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 + 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.wait(tasks, timeout=remaining) + except asyncio.CancelledError as error: + caller_cancellation = caller_cancellation or error + + 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 - self._dependencies_closed = True - await dependencies.aclose() + + 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: @@ -711,6 +863,164 @@ def _resolve_pty_session_entry( raise PtySessionNotFoundError(session_id=session_id) return entry + 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[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[Any]) -> 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, + propagate_timeout: bool = True, + ) -> 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 + 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 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 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( + rollback(), + propagate_timeout=False, + ) + + 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.""" + + 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[Any]) -> 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 + cleanup_errors: dict[int, BaseException] = {} + while pending: + remaining = deadline - loop.time() + if remaining <= 0: + break + try: + done, pending = await asyncio.wait(pending, timeout=remaining) + except BaseException as error: + if isinstance(error, asyncio.CancelledError): + caller_cancellation = caller_cancellation or error + continue + raise + + for index, task in enumerate(cleanup_tasks): + if task not in done: + continue + try: + task.result() + except BaseException as 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: + raise asyncio.TimeoutError() + async def pty_exec_start( self, *command: str | Path, diff --git a/src/agents/sandbox/session/sandbox_session.py b/src/agents/sandbox/session/sandbox_session.py index 923f025857..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 @@ -264,6 +265,17 @@ 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() + + 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/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index 3fe1d0d93a..77c457995d 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 @@ -1975,6 +1976,95 @@ 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()) + 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 second_ws._closed + + 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 + 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 new file mode 100644 index 0000000000..619a269777 --- /dev/null +++ b/tests/sandbox/test_base_sandbox_session.py @@ -0,0 +1,477 @@ +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 + + +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())) + 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 +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())) + 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 +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_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() + 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())) + try: + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel() + release.set() + + 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 +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())) + 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 + + +@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] + + +@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)) + 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 + finally: + release_first.set() + if not task.done(): + task.cancel() + with suppress(BaseException): + await task + + +@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) + with pytest.raises(asyncio.TimeoutError): + 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_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, +) -> None: + session = _session() + session.state = SimpleNamespace(manifest=Manifest(), 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 + + +@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() + + +@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_docker.py b/tests/sandbox/test_docker.py index e4c7cc812f..7d2bb45b1e 100644 --- a/tests/sandbox/test_docker.py +++ b/tests/sandbox/test_docker.py @@ -12,6 +12,8 @@ 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 @@ -3633,6 +3635,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,19 +4692,148 @@ 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="") +@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.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 second_socket.closed is True + assert session._pty_processes == {} + finally: + 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 + + 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( @@ -4739,14 +4871,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 diff --git a/tests/sandbox/test_runtime.py b/tests/sandbox/test_runtime.py index 7b4d7f540b..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 @@ -99,7 +100,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 @@ -608,6 +609,112 @@ 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 +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()) + 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) + + 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.shutdown_calls == 1 + assert client.delete_calls == 1 + assert inner.close_dependency_calls == 2 + + @pytest.mark.asyncio @pytest.mark.parametrize("runner_owned", [False, True]) async def test_pre_stop_cancellation_skips_persistence_and_completes_cleanup(