fix(sandbox): clean up PTY startup cancellation (#4750)

Author: Hughhhhcoder

src/agents/extensions/sandbox/blaxel/sandbox.py

@@ -51,6 +51,7 @@
     PTY_PROCESSES_MAX,
     PTY_PROCESSES_WARNING,
     PtyExecUpdate,
+    _settle_pty_cleanup,
     allocate_pty_process_id,
     clamp_pty_yield_time_ms,
     process_id_to_prune_from_meta,
@@ -839,11 +840,18 @@ async def pty_exec_start(
                 registered = True
         except asyncio.TimeoutError as e:
             if not registered:
-                await self._terminate_pty_entry(entry)
+                await _settle_pty_cleanup(self._terminate_pty_entry(entry))
             raise ExecTimeoutError(command=command, timeout_s=exec_timeout, cause=e) from e
+        except asyncio.CancelledError as cancellation:
+            if not registered:
+                await _settle_pty_cleanup(
+                    self._terminate_pty_entry(entry),
+                    initial_cancellation=cancellation,
+                )
+            raise
         except Exception as e:
             if not registered:
-                await self._terminate_pty_entry(entry)
+                await _settle_pty_cleanup(self._terminate_pty_entry(entry))
             raise _blaxel_exec_transport_error(command=command, cause=e) from e
 
         if pruned is not None:

src/agents/sandbox/sandboxes/unix_local.py

@@ -354,7 +354,7 @@ def _preexec() -> None:
                     env=env,
                     preexec_fn=_preexec,
                 )
-            except Exception:
+            except BaseException:
                 with suppress(OSError):
                     os.close(primary_fd)
                 with suppress(OSError):

src/agents/sandbox/session/pty_types.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+import asyncio
 import random
-from collections.abc import Sequence
+from collections.abc import Awaitable, Sequence
 from dataclasses import dataclass
 
 from ..util.token_truncation import formatted_truncate_text_with_token_count
@@ -18,6 +19,32 @@
 PTY_PROCESS_ID_MAX_EXCLUSIVE = 100_000
 
 
+async def _settle_pty_cleanup(
+    cleanup: Awaitable[None],
+    *,
+    initial_cancellation: asyncio.CancelledError | None = None,
+) -> None:
+    cleanup_task = asyncio.ensure_future(cleanup)
+    completion = asyncio.create_task(asyncio.wait((cleanup_task,)))
+    cancellation = initial_cancellation
+    while not completion.done():
+        try:
+            await asyncio.shield(completion)
+        except asyncio.CancelledError as error:
+            if cancellation is None:
+                cancellation = error
+
+    completion.result()
+    try:
+        cleanup_task.result()
+    except BaseException:
+        if cancellation is not None:
+            raise cancellation from None
+        raise
+    if cancellation is not None:
+        raise cancellation from None
+
+
 @dataclass(frozen=True)
 class PtyExecUpdate:
     process_id: int | None

tests/extensions/sandbox/test_blaxel.py

@@ -5,6 +5,7 @@
 import json
 import logging
 import shlex
+import sys
 import tarfile
 import time
 import uuid
@@ -1745,6 +1746,157 @@ def ClientSession(self) -> _FakeHTTPSession:
 
 
 class TestPtyExec:
+    @pytest.mark.asyncio
+    async def test_pty_exec_start_cancellation_closes_unregistered_http_session(
+        self, fake_sandbox: _FakeSandboxInstance
+    ) -> None:
+        from agents.extensions.sandbox.blaxel import sandbox as mod
+
+        connect_started = asyncio.Event()
+
+        class _BlockingSession:
+            def __init__(self) -> None:
+                self._closed = False
+
+            async def ws_connect(self, url: str) -> None:
+                _ = url
+                connect_started.set()
+                await asyncio.Event().wait()
+
+            async def close(self) -> None:
+                self._closed = True
+
+        class _BlockingAiohttp:
+            WSMsgType = _FakeAiohttp.WSMsgType
+
+            def __init__(self) -> None:
+                self.session: _BlockingSession | None = None
+
+            def ClientSession(self) -> _BlockingSession:
+                self.session = _BlockingSession()
+                return self.session
+
+        fake_aiohttp = _BlockingAiohttp()
+        session = _make_session(fake_sandbox)
+
+        with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
+            task = asyncio.create_task(session.pty_exec_start("echo", "hello"))
+            await connect_started.wait()
+            task.cancel("connect-cancel")
+
+            with pytest.raises(asyncio.CancelledError) as exc_info:
+                await task
+
+            if sys.version_info >= (3, 11):
+                assert exc_info.value.args == ("connect-cancel",)
+            assert task.cancelled()
+
+        assert fake_aiohttp.session is not None
+        assert fake_aiohttp.session._closed
+        assert session._pty_sessions == {}
+        assert session._reserved_pty_process_ids == set()
+
+    @pytest.mark.asyncio
+    async def test_pty_exec_start_preserves_cancellation_during_cleanup(
+        self, fake_sandbox: _FakeSandboxInstance
+    ) -> None:
+        from agents.extensions.sandbox.blaxel import sandbox as mod
+
+        cleanup_started = asyncio.Event()
+        allow_cleanup = asyncio.Event()
+
+        class _TimeoutSession:
+            def __init__(self) -> None:
+                self._closed = False
+
+            async def ws_connect(self, url: str) -> None:
+                _ = url
+                raise asyncio.TimeoutError()
+
+            async def close(self) -> None:
+                cleanup_started.set()
+                await allow_cleanup.wait()
+                self._closed = True
+
+        class _TimeoutAiohttp:
+            WSMsgType = _FakeAiohttp.WSMsgType
+
+            def __init__(self) -> None:
+                self.session: _TimeoutSession | None = None
+
+            def ClientSession(self) -> _TimeoutSession:
+                self.session = _TimeoutSession()
+                return self.session
+
+        fake_aiohttp = _TimeoutAiohttp()
+        session = _make_session(fake_sandbox)
+
+        with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
+            task = asyncio.create_task(session.pty_exec_start("echo", "hello"))
+            await cleanup_started.wait()
+            task.cancel("cleanup-cancel")
+            allow_cleanup.set()
+
+            with pytest.raises(asyncio.CancelledError) as exc_info:
+                await task
+
+            if sys.version_info >= (3, 11):
+                assert exc_info.value.args == ("cleanup-cancel",)
+        assert fake_aiohttp.session is not None
+        assert fake_aiohttp.session._closed
+        assert session._pty_sessions == {}
+        assert session._reserved_pty_process_ids == set()
+
+    @pytest.mark.asyncio
+    async def test_pty_exec_start_preserves_cancellation_when_cleanup_fails(
+        self, fake_sandbox: _FakeSandboxInstance
+    ) -> None:
+        from agents.extensions.sandbox.blaxel import sandbox as mod
+
+        cleanup_started = asyncio.Event()
+        allow_cleanup = asyncio.Event()
+
+        class _FailingCleanupSession:
+            def __init__(self) -> None:
+                self._closed = False
+
+            async def ws_connect(self, url: str) -> None:
+                _ = url
+                raise asyncio.TimeoutError()
+
+            async def close(self) -> None:
+                cleanup_started.set()
+                await allow_cleanup.wait()
+                self._closed = True
+                raise RuntimeError("synthetic cleanup failure")
+
+        class _FailingCleanupAiohttp:
+            WSMsgType = _FakeAiohttp.WSMsgType
+
+            def __init__(self) -> None:
+                self.session: _FailingCleanupSession | None = None
+
+            def ClientSession(self) -> _FailingCleanupSession:
+                self.session = _FailingCleanupSession()
+                return self.session
+
+        fake_aiohttp = _FailingCleanupAiohttp()
+        session = _make_session(fake_sandbox)
+
+        with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
+            task = asyncio.create_task(session.pty_exec_start("echo", "hello"))
+            await cleanup_started.wait()
+            task.cancel()
+            allow_cleanup.set()
+
+            with pytest.raises(asyncio.CancelledError):
+                await task
+
+        assert fake_aiohttp.session is not None
+        assert fake_aiohttp.session._closed
+        assert session._pty_sessions == {}
+        assert session._reserved_pty_process_ids == set()
+
     @pytest.mark.parametrize(
         ("messages", "expected_output"),
         [

tests/sandbox/test_pty_types.py

@@ -1,8 +1,14 @@
 from __future__ import annotations
 
+import asyncio
+import sys
+
+import pytest
+
 from agents.sandbox.session.pty_types import (
     PTY_EMPTY_YIELD_TIME_MS_MIN,
     PTY_YIELD_TIME_MS_MIN,
+    _settle_pty_cleanup,
     allocate_pty_process_id,
     clamp_pty_yield_time_ms,
     process_id_to_prune_from_meta,
@@ -37,3 +43,39 @@ def test_process_id_to_prune_from_meta_prefers_exited_unprotected_sessions() ->
     meta.append((2002, 2.0, False))
 
     assert process_id_to_prune_from_meta(meta) == 2001
+
+
+@pytest.mark.asyncio
+async def test_settle_pty_cleanup_preserves_cancel_reason_when_cleanup_fails() -> None:
+    cleanup_started = asyncio.Event()
+    cleanup_release = asyncio.Event()
+
+    async def cleanup() -> None:
+        cleanup_started.set()
+        await cleanup_release.wait()
+        raise RuntimeError("synthetic cleanup failure")
+
+    task = asyncio.create_task(_settle_pty_cleanup(cleanup()))
+    await cleanup_started.wait()
+    task.cancel("route-A")
+    cleanup_release.set()
+
+    with pytest.raises(asyncio.CancelledError) as exc_info:
+        await task
+
+    if sys.version_info >= (3, 11):
+        assert exc_info.value.args == ("route-A",)
+
+
+@pytest.mark.asyncio
+async def test_settle_pty_cleanup_preserves_initial_cancel_reason_when_cleanup_fails() -> None:
+    async def cleanup() -> None:
+        raise RuntimeError("synthetic cleanup failure")
+
+    with pytest.raises(asyncio.CancelledError) as exc_info:
+        await _settle_pty_cleanup(
+            cleanup(),
+            initial_cancellation=asyncio.CancelledError("startup-cancel"),
+        )
+
+    assert exc_info.value.args == ("startup-cancel",)

tests/sandbox/test_unix_local.py

@@ -216,6 +216,34 @@ def _unexpected_mkdtemp(*args: object, **kwargs: object) -> str:
 
 @pytest.mark.review_optional
 class TestUnixLocalPty:
+    @pytest.mark.asyncio
+    async def test_tty_start_cancellation_closes_open_file_descriptors(
+        self,
+        tmp_path: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        monkeypatch.setattr(unix_local_module.sys, "platform", "linux")
+        workspace = tmp_path / "workspace"
+        workspace.mkdir()
+        session = _RecordingUnixLocalSession(workspace)
+        close_calls: list[int] = []
+
+        def openpty() -> tuple[int, int]:
+            return 101, 102
+
+        async def create_subprocess(*args: object, **kwargs: object) -> None:
+            _ = (args, kwargs)
+            raise asyncio.CancelledError()
+
+        monkeypatch.setattr(unix_local_module.os, "openpty", openpty)
+        monkeypatch.setattr(unix_local_module.os, "close", close_calls.append)
+        monkeypatch.setattr(unix_local_module.asyncio, "create_subprocess_exec", create_subprocess)
+
+        with pytest.raises(asyncio.CancelledError):
+            await session.pty_exec_start("echo", "hello", shell=False, tty=True)
+
+        assert close_calls == [101, 102]
+
     @pytest.mark.asyncio
     async def test_tty_fd_close_is_owned_without_blocking_termination(
         self,