fix(sandbox): keep unbounded UnixLocal workspace I/O off the event loop

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,14 +1111,18 @@ 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(
                     tar,
                     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

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

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