diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 5cc77e2aff..ea7f83e9d8 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -62,6 +62,7 @@ from ..session.workspace_payloads import coerce_write_payload from ..snapshot import SnapshotBase, SnapshotSpec, resolve_snapshot from ..types import ExecResult, ExposedPortEndpoint, Permissions, User +from ..util.blocking_io import run_blocking_workspace_io from ..util.tar_utils import ( UnsafeTarMemberError, safe_extract_tarfile, @@ -962,7 +963,7 @@ async def rm( try: if normalized.is_dir() and not normalized.is_symlink(): if recursive: - shutil.rmtree(normalized) + await run_blocking_workspace_io(shutil.rmtree, normalized) else: normalized.rmdir() else: @@ -1083,7 +1084,8 @@ async def persist_workspace(self) -> io.IOBase: skip = self._persist_workspace_skip_relpaths() buf = io.BytesIO() - try: + + def _archive_workspace() -> None: with tarfile.open(fileobj=buf, mode="w") as tar: tar.add( root, @@ -1098,6 +1100,9 @@ async def persist_workspace(self) -> io.IOBase: else ti ), ) + + try: + await run_blocking_workspace_io(_archive_workspace) except (tarfile.TarError, OSError) as e: raise WorkspaceArchiveReadError(path=root, cause=e) from e @@ -1106,7 +1111,8 @@ async def persist_workspace(self) -> io.IOBase: async def hydrate_workspace(self, data: io.IOBase) -> None: root = Path(self.state.manifest.root) - try: + + def _extract_workspace() -> None: root.mkdir(parents=True, exist_ok=True) with tarfile.open(fileobj=data, mode="r:*") as tar: safe_extract_tarfile( @@ -1114,6 +1120,9 @@ async def hydrate_workspace(self, data: io.IOBase) -> None: root=root, allow_external_symlink_targets=False, ) + + try: + await run_blocking_workspace_io(_extract_workspace) except UnsafeTarMemberError as e: raise WorkspaceArchiveWriteError( path=root, context={"reason": e.reason, "member": e.member}, cause=e diff --git a/src/agents/sandbox/util/blocking_io.py b/src/agents/sandbox/util/blocking_io.py new file mode 100644 index 0000000000..ff8b58a32e --- /dev/null +++ b/src/agents/sandbox/util/blocking_io.py @@ -0,0 +1,40 @@ +"""Run unbounded blocking workspace I/O off the event loop without abandoning it. + +`asyncio.to_thread()` does not stop its worker when the awaiting task is cancelled, so a +cancelled caller can return while the thread is still writing. Snapshot resume closes the +archive stream and clears the workspace root as soon as its await returns, which would let a +surviving worker extract into a workspace that is being deleted. + +Callers therefore keep waiting for the worker even while cancelled, matching the mutation +semantics the session backends already rely on in `agents.memory.sqlite_session`. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from typing import Any, TypeVar + +_T = TypeVar("_T") + + +async def run_blocking_workspace_io(function: Callable[..., _T], /, *args: Any) -> _T: + """Run `function` in a worker thread and keep ownership until that worker finishes.""" + task = asyncio.ensure_future(asyncio.to_thread(function, *args)) + cancellation: asyncio.CancelledError | None = None + while not task.done(): + try: + await asyncio.wait({task}) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + + try: + result = task.result() + except BaseException: + if cancellation is not None: + raise cancellation from None + raise + if cancellation is not None: + raise cancellation from None + return result diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index c8e9c654a1..fb13be3c51 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -1,7 +1,11 @@ from __future__ import annotations import asyncio +import io import signal +import tarfile +import threading +import time from pathlib import Path from types import SimpleNamespace from typing import cast @@ -464,3 +468,48 @@ async def test_rm_as_user_checks_permissions_then_uses_local_fs( assert session.exec_commands[0][4:6] == ("sh", "-lc") assert session.exec_commands[0][-2:] == (str(target), "0") assert not any(part.startswith("rm ") for part in session.exec_commands[0]) + + +@pytest.mark.asyncio +async def test_hydrate_workspace_cancellation_waits_for_the_extracting_worker( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A cancelled hydrate must not leave a worker writing into the workspace. + + `restore_snapshot_into_workspace_on_resume` closes the archive stream in a `finally` as + soon as its await returns, so if cancellation propagated while the extractor was still + running it would read a closed stream and write into a workspace resume then clears. + """ + workspace = tmp_path / "workspace" + session = _RecordingUnixLocalSession(workspace) + + started = threading.Event() + events: list[str] = [] + + def _slow_extract(tar: object, **kwargs: object) -> None: + _ = tar, kwargs + events.append("extract-start") + started.set() + time.sleep(0.2) + events.append("extract-end") + + monkeypatch.setattr(unix_local_module, "safe_extract_tarfile", _slow_extract) + + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w"): + pass + buf.seek(0) + + task = asyncio.create_task(session.hydrate_workspace(buf)) + while not started.is_set(): + await asyncio.sleep(0.005) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + # The worker finished before the caller observed cancellation, so the archive stream and + # the workspace root are only released once nothing is still writing to them. + assert events == ["extract-start", "extract-end"] + assert not buf.closed