From 26ffd7de911cd78dd4755177fe1d3d4be08e7447 Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Sat, 29 Aug 2026 13:26:55 +0800 Subject: [PATCH 01/12] fix(sandbox): clean up PTY startup cancellation --- .../extensions/sandbox/blaxel/sandbox.py | 9 ++- src/agents/sandbox/sandboxes/unix_local.py | 26 ++++--- src/agents/sandbox/session/pty_types.py | 13 +++- tests/extensions/sandbox/test_blaxel.py | 46 +++++++++++ tests/sandbox/test_unix_local.py | 78 +++++++++++++++++++ 5 files changed, 160 insertions(+), 12 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index aee3211fbd..9cba5c865e 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/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,15 @@ 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: + if not registered: + await _settle_pty_cleanup(self._terminate_pty_entry(entry)) + 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: diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index ea7f83e9d8..af15469a78 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -53,6 +53,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, @@ -354,7 +355,7 @@ def _preexec() -> None: env=env, preexec_fn=_preexec, ) - except Exception: + except BaseException: with suppress(OSError): os.close(primary_fd) with suppress(OSError): @@ -380,15 +381,22 @@ def _preexec() -> None: asyncio.create_task(self._pump_process_stream(entry, process.stderr)), ] - entry.wait_task = asyncio.create_task(self._watch_process_exit(entry)) + registered = False + try: + 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) + 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 = True + except BaseException: + if not registered: + await _settle_pty_cleanup(self._terminate_pty_entry(entry)) + raise if pruned_entry is not None: await self._terminate_pty_entry(pruned_entry) diff --git a/src/agents/sandbox/session/pty_types.py b/src/agents/sandbox/session/pty_types.py index 3f4dab04b0..6e78d2e1ae 100644 --- a/src/agents/sandbox/session/pty_types.py +++ b/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,16 @@ PTY_PROCESS_ID_MAX_EXCLUSIVE = 100_000 +async def _settle_pty_cleanup(cleanup: Awaitable[None]) -> None: + cleanup_task = asyncio.ensure_future(cleanup) + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + continue + await cleanup_task + + @dataclass(frozen=True) class PtyExecUpdate: process_id: int | None diff --git a/tests/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index 3fe1d0d93a..19d5548dec 100644 --- a/tests/extensions/sandbox/test_blaxel.py +++ b/tests/extensions/sandbox/test_blaxel.py @@ -1745,6 +1745,52 @@ 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() + + 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"), [ diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index fb13be3c51..08d6745302 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -216,6 +216,84 @@ def _unexpected_mkdtemp(*args: object, **kwargs: object) -> str: @pytest.mark.review_optional class TestUnixLocalPty: + @pytest.mark.asyncio + async def test_pty_start_cancellation_cleans_up_before_registration( + 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) + process_started = asyncio.Event() + killpg_calls: list[tuple[int, signal.Signals]] = [] + + class _Process: + pid = 1234 + returncode = None + stdout = None + stderr = None + + async def wait(self) -> None: + return None + + async def create_subprocess(*args: object, **kwargs: object) -> _Process: + _ = (args, kwargs) + process_started.set() + return _Process() + + def killpg(pid: int, signum: signal.Signals) -> None: + killpg_calls.append((pid, signum)) + + monkeypatch.setattr(unix_local_module.asyncio, "create_subprocess_exec", create_subprocess) + monkeypatch.setattr(unix_local_module.os, "killpg", killpg) + await session._pty_lock.acquire() + try: + task = asyncio.create_task( + session.pty_exec_start("echo", "hello", shell=False, yield_time_s=0.01) + ) + await process_started.wait() + await asyncio.sleep(0) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert killpg_calls == [(1234, signal.SIGKILL)] + assert session._pty_processes == {} + assert session._reserved_pty_process_ids == set() + finally: + session._pty_lock.release() + + @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, From 5cd36bf6ce071b4c1f2aace46c3f8a0631412102 Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Sat, 29 Aug 2026 13:42:54 +0800 Subject: [PATCH 02/12] fix(sandbox): preserve cancellation during PTY cleanup --- src/agents/sandbox/session/pty_types.py | 16 ++++++-- tests/extensions/sandbox/test_blaxel.py | 49 +++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 3 deletions(-) diff --git a/src/agents/sandbox/session/pty_types.py b/src/agents/sandbox/session/pty_types.py index 6e78d2e1ae..ae940ada09 100644 --- a/src/agents/sandbox/session/pty_types.py +++ b/src/agents/sandbox/session/pty_types.py @@ -21,12 +21,22 @@ async def _settle_pty_cleanup(cleanup: Awaitable[None]) -> None: cleanup_task = asyncio.ensure_future(cleanup) + cancellation: asyncio.CancelledError | None = None while not cleanup_task.done(): try: await asyncio.shield(cleanup_task) - except asyncio.CancelledError: - continue - await cleanup_task + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + + 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) diff --git a/tests/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index 19d5548dec..3d81bbb293 100644 --- a/tests/extensions/sandbox/test_blaxel.py +++ b/tests/extensions/sandbox/test_blaxel.py @@ -1791,6 +1791,55 @@ def ClientSession(self) -> _BlockingSession: 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() + 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"), [ From 2720f1a5a18492c109a8470e6572ac05ac48ac56 Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Sat, 29 Aug 2026 13:43:30 +0800 Subject: [PATCH 03/12] test(sandbox): assert PTY cancellation is preserved --- tests/extensions/sandbox/test_blaxel.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index 3d81bbb293..0cab0f0968 100644 --- a/tests/extensions/sandbox/test_blaxel.py +++ b/tests/extensions/sandbox/test_blaxel.py @@ -1786,6 +1786,8 @@ def ClientSession(self) -> _BlockingSession: with pytest.raises(asyncio.CancelledError): await task + assert task.cancelled() + assert fake_aiohttp.session is not None assert fake_aiohttp.session._closed assert session._pty_sessions == {} From 9a94bd7b23e1b2e1eddcdc1541790781a8fb31a7 Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Sat, 29 Aug 2026 15:12:42 +0800 Subject: [PATCH 04/12] fix(sandbox): preserve cancellation after cleanup failure --- src/agents/sandbox/session/pty_types.py | 21 ++++++----- tests/extensions/sandbox/test_blaxel.py | 50 +++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 10 deletions(-) diff --git a/src/agents/sandbox/session/pty_types.py b/src/agents/sandbox/session/pty_types.py index ae940ada09..c167077020 100644 --- a/src/agents/sandbox/session/pty_types.py +++ b/src/agents/sandbox/session/pty_types.py @@ -21,22 +21,23 @@ async def _settle_pty_cleanup(cleanup: Awaitable[None]) -> None: cleanup_task = asyncio.ensure_future(cleanup) - cancellation: asyncio.CancelledError | None = None - while not cleanup_task.done(): + completion = asyncio.create_task(asyncio.wait((cleanup_task,))) + caller_cancelled = False + while not completion.done(): try: - await asyncio.shield(cleanup_task) - except asyncio.CancelledError as exc: - if cancellation is None: - cancellation = exc + await asyncio.shield(completion) + except asyncio.CancelledError: + caller_cancelled = True + completion.result() try: cleanup_task.result() except BaseException: - if cancellation is not None: - raise cancellation from None + if caller_cancelled: + raise asyncio.CancelledError() from None raise - if cancellation is not None: - raise cancellation from None + if caller_cancelled: + raise asyncio.CancelledError() from None @dataclass(frozen=True) diff --git a/tests/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index 0cab0f0968..eeaf553294 100644 --- a/tests/extensions/sandbox/test_blaxel.py +++ b/tests/extensions/sandbox/test_blaxel.py @@ -1842,6 +1842,56 @@ def ClientSession(self) -> _TimeoutSession: 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"), [ From cbaa66e2de21c7d61110ddeec9a022d6100bf650 Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Sat, 29 Aug 2026 15:49:13 +0800 Subject: [PATCH 05/12] fix(sandbox): preserve PTY cancellation payloads --- .../extensions/sandbox/blaxel/sandbox.py | 7 +++- src/agents/sandbox/sandboxes/unix_local.py | 7 ++++ src/agents/sandbox/session/pty_types.py | 21 ++++++---- tests/extensions/sandbox/test_blaxel.py | 10 +++-- tests/sandbox/test_pty_types.py | 40 +++++++++++++++++++ tests/sandbox/test_unix_local.py | 6 ++- 6 files changed, 75 insertions(+), 16 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index 9cba5c865e..a54bc47905 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -842,9 +842,12 @@ async def pty_exec_start( if not registered: await _settle_pty_cleanup(self._terminate_pty_entry(entry)) raise ExecTimeoutError(command=command, timeout_s=exec_timeout, cause=e) from e - except asyncio.CancelledError: + except asyncio.CancelledError as cancellation: if not registered: - await _settle_pty_cleanup(self._terminate_pty_entry(entry)) + await _settle_pty_cleanup( + self._terminate_pty_entry(entry), + initial_cancellation=cancellation, + ) raise except Exception as e: if not registered: diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index af15469a78..060523f32e 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -393,6 +393,13 @@ def _preexec() -> None: self._pty_processes[process_id] = entry process_count = len(self._pty_processes) registered = True + except asyncio.CancelledError as cancellation: + if not registered: + await _settle_pty_cleanup( + self._terminate_pty_entry(entry), + initial_cancellation=cancellation, + ) + raise except BaseException: if not registered: await _settle_pty_cleanup(self._terminate_pty_entry(entry)) diff --git a/src/agents/sandbox/session/pty_types.py b/src/agents/sandbox/session/pty_types.py index c167077020..65d8163534 100644 --- a/src/agents/sandbox/session/pty_types.py +++ b/src/agents/sandbox/session/pty_types.py @@ -19,25 +19,30 @@ PTY_PROCESS_ID_MAX_EXCLUSIVE = 100_000 -async def _settle_pty_cleanup(cleanup: Awaitable[None]) -> None: +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,))) - caller_cancelled = False + cancellation = initial_cancellation while not completion.done(): try: await asyncio.shield(completion) - except asyncio.CancelledError: - caller_cancelled = True + except asyncio.CancelledError as error: + if cancellation is None: + cancellation = error completion.result() try: cleanup_task.result() except BaseException: - if caller_cancelled: - raise asyncio.CancelledError() from None + if cancellation is not None: + raise cancellation from None raise - if caller_cancelled: - raise asyncio.CancelledError() from None + if cancellation is not None: + raise cancellation from None @dataclass(frozen=True) diff --git a/tests/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index eeaf553294..45fdaac60c 100644 --- a/tests/extensions/sandbox/test_blaxel.py +++ b/tests/extensions/sandbox/test_blaxel.py @@ -1781,11 +1781,12 @@ def ClientSession(self) -> _BlockingSession: 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() + task.cancel("connect-cancel") - with pytest.raises(asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError) as exc_info: await task + assert exc_info.value.args == ("connect-cancel",) assert task.cancelled() assert fake_aiohttp.session is not None @@ -1831,12 +1832,13 @@ def ClientSession(self) -> _TimeoutSession: 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() + task.cancel("cleanup-cancel") allow_cleanup.set() - with pytest.raises(asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError) as exc_info: await task + assert exc_info.value.args == ("cleanup-cancel",) assert fake_aiohttp.session is not None assert fake_aiohttp.session._closed assert session._pty_sessions == {} diff --git a/tests/sandbox/test_pty_types.py b/tests/sandbox/test_pty_types.py index a8c6db2820..85e621c6a3 100644 --- a/tests/sandbox/test_pty_types.py +++ b/tests/sandbox/test_pty_types.py @@ -1,8 +1,13 @@ from __future__ import annotations +import asyncio + +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 +42,38 @@ 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 + + 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",) diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index 08d6745302..872b804bee 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -245,6 +245,7 @@ async def create_subprocess(*args: object, **kwargs: object) -> _Process: def killpg(pid: int, signum: signal.Signals) -> None: killpg_calls.append((pid, signum)) + raise PermissionError("synthetic cleanup failure") monkeypatch.setattr(unix_local_module.asyncio, "create_subprocess_exec", create_subprocess) monkeypatch.setattr(unix_local_module.os, "killpg", killpg) @@ -255,11 +256,12 @@ def killpg(pid: int, signum: signal.Signals) -> None: ) await process_started.wait() await asyncio.sleep(0) - task.cancel() + task.cancel("startup-cancel") - with pytest.raises(asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError) as exc_info: await task + assert exc_info.value.args == ("startup-cancel",) assert killpg_calls == [(1234, signal.SIGKILL)] assert session._pty_processes == {} assert session._reserved_pty_process_ids == set() From 9ef5f9d97158dbc7bc8f7535607ee6eff20125e9 Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Mon, 31 Aug 2026 13:57:32 +0800 Subject: [PATCH 06/12] test(sandbox): account for Python 3.10 cancellation messages --- tests/extensions/sandbox/test_blaxel.py | 7 +++++-- tests/sandbox/test_pty_types.py | 4 +++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index 45fdaac60c..4916a3a970 100644 --- a/tests/extensions/sandbox/test_blaxel.py +++ b/tests/extensions/sandbox/test_blaxel.py @@ -5,6 +5,7 @@ import json import logging import shlex +import sys import tarfile import time import uuid @@ -1786,7 +1787,8 @@ def ClientSession(self) -> _BlockingSession: with pytest.raises(asyncio.CancelledError) as exc_info: await task - assert exc_info.value.args == ("connect-cancel",) + if sys.version_info >= (3, 11): + assert exc_info.value.args == ("connect-cancel",) assert task.cancelled() assert fake_aiohttp.session is not None @@ -1838,7 +1840,8 @@ def ClientSession(self) -> _TimeoutSession: with pytest.raises(asyncio.CancelledError) as exc_info: await task - assert exc_info.value.args == ("cleanup-cancel",) + 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 == {} diff --git a/tests/sandbox/test_pty_types.py b/tests/sandbox/test_pty_types.py index 85e621c6a3..4b09515e65 100644 --- a/tests/sandbox/test_pty_types.py +++ b/tests/sandbox/test_pty_types.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import sys import pytest @@ -62,7 +63,8 @@ async def cleanup() -> None: with pytest.raises(asyncio.CancelledError) as exc_info: await task - assert exc_info.value.args == ("route-A",) + if sys.version_info >= (3, 11): + assert exc_info.value.args == ("route-A",) @pytest.mark.asyncio From 45dc7c35fef6ecefe75909784e2ae416862a5fb7 Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Mon, 31 Aug 2026 14:50:32 +0800 Subject: [PATCH 07/12] fix(sandbox): continue PTY cleanup after kill failure --- src/agents/sandbox/sandboxes/unix_local.py | 2 +- tests/sandbox/test_unix_local.py | 41 +++++++++++++++++++--- 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 060523f32e..befd4e636e 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -605,7 +605,7 @@ async def _terminate_pty_entry(self, entry: _UnixPtyProcessEntry) -> None: entry.primary_fd = None if process.returncode is None and process.pid is not None: - with suppress(ProcessLookupError): + with suppress(OSError): os.killpg(process.pid, signal.SIGKILL) for task in entry.pump_tasks: diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index 872b804bee..d5c568f2ab 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -3,6 +3,7 @@ import asyncio import io import signal +import sys import tarfile import threading import time @@ -227,16 +228,42 @@ async def test_pty_start_cancellation_cleans_up_before_registration( workspace.mkdir() session = _RecordingUnixLocalSession(workspace) process_started = asyncio.Event() + pump_tasks_started = asyncio.Event() + pump_tasks_cancelled = asyncio.Event() + process_wait_started = asyncio.Event() + process_wait_cancelled = asyncio.Event() + pump_task_count = 0 + cancelled_pump_task_count = 0 killpg_calls: list[tuple[int, signal.Signals]] = [] + class _Stream: + async def read(self, size: int) -> bytes: + nonlocal pump_task_count, cancelled_pump_task_count + _ = size + pump_task_count += 1 + if pump_task_count == 2: + pump_tasks_started.set() + try: + return await asyncio.Future[bytes]() + except asyncio.CancelledError: + cancelled_pump_task_count += 1 + if cancelled_pump_task_count == 2: + pump_tasks_cancelled.set() + raise + class _Process: pid = 1234 returncode = None - stdout = None - stderr = None + stdout = _Stream() + stderr = _Stream() async def wait(self) -> None: - return None + process_wait_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + process_wait_cancelled.set() + raise async def create_subprocess(*args: object, **kwargs: object) -> _Process: _ = (args, kwargs) @@ -255,14 +282,18 @@ def killpg(pid: int, signum: signal.Signals) -> None: session.pty_exec_start("echo", "hello", shell=False, yield_time_s=0.01) ) await process_started.wait() - await asyncio.sleep(0) + await pump_tasks_started.wait() + await process_wait_started.wait() task.cancel("startup-cancel") with pytest.raises(asyncio.CancelledError) as exc_info: await task - assert exc_info.value.args == ("startup-cancel",) + if sys.version_info >= (3, 11): + assert exc_info.value.args == ("startup-cancel",) assert killpg_calls == [(1234, signal.SIGKILL)] + assert pump_tasks_cancelled.is_set() + assert process_wait_cancelled.is_set() assert session._pty_processes == {} assert session._reserved_pty_process_ids == set() finally: From 15b40c668fd610946fc066cd6c52985dd388f50c Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Mon, 31 Aug 2026 15:18:02 +0800 Subject: [PATCH 08/12] fix(sandbox): retain failed PTY cleanup ownership --- src/agents/sandbox/sandboxes/unix_local.py | 29 +++++++++++++++++++--- tests/sandbox/test_unix_local.py | 23 +++++++++++++++-- 2 files changed, 47 insertions(+), 5 deletions(-) diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index befd4e636e..06081a71bd 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -163,6 +163,7 @@ class UnixLocalSandboxSession(BaseSandboxSession): _running: bool _pty_lock: asyncio.Lock _pty_processes: dict[int, _UnixPtyProcessEntry] + _unregistered_pty_processes: dict[int, _UnixPtyProcessEntry] _reserved_pty_process_ids: set[int] _fd_close_tasks: set[asyncio.Task[None]] _host_environment_allowlist: frozenset[str] | None @@ -172,6 +173,7 @@ def __init__(self, *, state: UnixLocalSandboxSessionState) -> None: self._running = False self._pty_lock = asyncio.Lock() self._pty_processes = {} + self._unregistered_pty_processes = {} self._reserved_pty_process_ids = set() self._fd_close_tasks = set() self._host_environment_allowlist = None @@ -381,6 +383,7 @@ def _preexec() -> None: asyncio.create_task(self._pump_process_stream(entry, process.stderr)), ] + self._unregistered_pty_processes[id(entry)] = entry registered = False try: entry.wait_task = asyncio.create_task(self._watch_process_exit(entry)) @@ -391,6 +394,7 @@ def _preexec() -> None: self._reserved_pty_process_ids.add(process_id) pruned_entry = self._prune_pty_processes_if_needed() self._pty_processes[process_id] = entry + self._unregistered_pty_processes.pop(id(entry), None) process_count = len(self._pty_processes) registered = True except asyncio.CancelledError as cancellation: @@ -474,11 +478,17 @@ async def pty_write_stdin( async def pty_terminate_all(self) -> None: async with self._pty_lock: - entries = list(self._pty_processes.values()) + entries_by_id = { + id(entry): entry + for entry in ( + *self._pty_processes.values(), + *self._unregistered_pty_processes.values(), + ) + } self._pty_processes.clear() self._reserved_pty_process_ids.clear() - for entry in entries: + for entry in entries_by_id.values(): await self._terminate_pty_entry(entry) async def _resolved_exec_context(self) -> tuple[dict[str, str], str]: @@ -604,9 +614,22 @@ async def _terminate_pty_entry(self, entry: _UnixPtyProcessEntry) -> None: primary_fd = entry.primary_fd entry.primary_fd = None + termination_failed = process.returncode is None and process.pid is None if process.returncode is None and process.pid is not None: - with suppress(OSError): + try: os.killpg(process.pid, signal.SIGKILL) + except OSError: + try: + process.kill() + except ProcessLookupError: + pass + except OSError: + termination_failed = True + + if termination_failed: + self._unregistered_pty_processes[id(entry)] = entry + else: + self._unregistered_pty_processes.pop(id(entry), None) for task in entry.pump_tasks: task.cancel() diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index d5c568f2ab..78745d0383 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -234,6 +234,7 @@ async def test_pty_start_cancellation_cleans_up_before_registration( process_wait_cancelled = asyncio.Event() pump_task_count = 0 cancelled_pump_task_count = 0 + process_kill_calls = 0 killpg_calls: list[tuple[int, signal.Signals]] = [] class _Stream: @@ -265,14 +266,22 @@ async def wait(self) -> None: process_wait_cancelled.set() raise + def kill(self) -> None: + nonlocal process_kill_calls + process_kill_calls += 1 + raise PermissionError("synthetic direct cleanup failure") + + process = _Process() + async def create_subprocess(*args: object, **kwargs: object) -> _Process: _ = (args, kwargs) process_started.set() - return _Process() + return process def killpg(pid: int, signum: signal.Signals) -> None: killpg_calls.append((pid, signum)) - raise PermissionError("synthetic cleanup failure") + if len(killpg_calls) == 1: + raise PermissionError("synthetic cleanup failure") monkeypatch.setattr(unix_local_module.asyncio, "create_subprocess_exec", create_subprocess) monkeypatch.setattr(unix_local_module.os, "killpg", killpg) @@ -292,13 +301,23 @@ def killpg(pid: int, signum: signal.Signals) -> None: if sys.version_info >= (3, 11): assert exc_info.value.args == ("startup-cancel",) assert killpg_calls == [(1234, signal.SIGKILL)] + assert process_kill_calls == 1 assert pump_tasks_cancelled.is_set() assert process_wait_cancelled.is_set() assert session._pty_processes == {} assert session._reserved_pty_process_ids == set() + retained_entries = tuple(session._unregistered_pty_processes.values()) + assert len(retained_entries) == 1 + assert retained_entries[0].process is process finally: session._pty_lock.release() + await session.pty_terminate_all() + + assert killpg_calls == [(1234, signal.SIGKILL), (1234, signal.SIGKILL)] + assert process_kill_calls == 1 + assert session._unregistered_pty_processes == {} + @pytest.mark.asyncio async def test_tty_start_cancellation_closes_open_file_descriptors( self, From 0b86ca2cb36682b3e72d13ffefb2e2efe33f0458 Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Mon, 31 Aug 2026 15:28:01 +0800 Subject: [PATCH 09/12] fix(sandbox): report PTY termination failures --- src/agents/sandbox/sandboxes/unix_local.py | 43 ++++++++++++++-------- tests/sandbox/test_unix_local.py | 21 +++++++++-- 2 files changed, 45 insertions(+), 19 deletions(-) diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 06081a71bd..0510e41884 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -487,9 +487,17 @@ async def pty_terminate_all(self) -> None: } self._pty_processes.clear() self._reserved_pty_process_ids.clear() + self._unregistered_pty_processes.update(entries_by_id) + first_error: Exception | None = None for entry in entries_by_id.values(): - await self._terminate_pty_entry(entry) + try: + await self._terminate_pty_entry(entry) + except Exception as error: + if first_error is None: + first_error = error + if first_error is not None: + raise first_error async def _resolved_exec_context(self) -> tuple[dict[str, str], str]: if self._host_environment_allowlist is None: @@ -613,23 +621,23 @@ async def _terminate_pty_entry(self, entry: _UnixPtyProcessEntry) -> None: process = entry.process primary_fd = entry.primary_fd entry.primary_fd = None + self._unregistered_pty_processes[id(entry)] = entry - termination_failed = process.returncode is None and process.pid is None + termination_error: OSError | None = None + termination_cause: OSError | None = None + if process.returncode is None and process.pid is None: + termination_error = OSError("Cannot terminate PTY process without a PID") if process.returncode is None and process.pid is not None: try: os.killpg(process.pid, signal.SIGKILL) - except OSError: + except OSError as group_error: try: process.kill() except ProcessLookupError: pass - except OSError: - termination_failed = True - - if termination_failed: - self._unregistered_pty_processes[id(entry)] = entry - else: - self._unregistered_pty_processes.pop(id(entry), None) + except OSError as process_error: + termination_error = process_error + termination_cause = group_error for task in entry.pump_tasks: task.cancel() @@ -643,13 +651,16 @@ async def _terminate_pty_entry(self, entry: _UnixPtyProcessEntry) -> None: self._schedule_fd_close(primary_fd) entry.output_closed.set() entry.output_notify.set() - return + else: + if primary_fd is not None: + _close_fd_quietly(primary_fd) + await asyncio.gather(*entry.pump_tasks, return_exceptions=True) + if entry.wait_task is not None: + await asyncio.gather(entry.wait_task, return_exceptions=True) - if primary_fd is not None: - _close_fd_quietly(primary_fd) - await asyncio.gather(*entry.pump_tasks, return_exceptions=True) - if entry.wait_task is not None: - await asyncio.gather(entry.wait_task, return_exceptions=True) + if termination_error is not None: + raise termination_error from termination_cause + self._unregistered_pty_processes.pop(id(entry), None) def _schedule_fd_close(self, fd: int) -> None: task = asyncio.create_task(asyncio.to_thread(_close_fd_quietly, fd)) diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index 78745d0383..f7e66717f9 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -234,6 +234,7 @@ async def test_pty_start_cancellation_cleans_up_before_registration( process_wait_cancelled = asyncio.Event() pump_task_count = 0 cancelled_pump_task_count = 0 + allow_group_kill = False process_kill_calls = 0 killpg_calls: list[tuple[int, signal.Signals]] = [] @@ -280,7 +281,7 @@ async def create_subprocess(*args: object, **kwargs: object) -> _Process: def killpg(pid: int, signum: signal.Signals) -> None: killpg_calls.append((pid, signum)) - if len(killpg_calls) == 1: + if not allow_group_kill: raise PermissionError("synthetic cleanup failure") monkeypatch.setattr(unix_local_module.asyncio, "create_subprocess_exec", create_subprocess) @@ -312,10 +313,24 @@ def killpg(pid: int, signum: signal.Signals) -> None: finally: session._pty_lock.release() - await session.pty_terminate_all() + with pytest.raises(PermissionError, match="synthetic direct cleanup failure"): + await session.pty_terminate_all() assert killpg_calls == [(1234, signal.SIGKILL), (1234, signal.SIGKILL)] - assert process_kill_calls == 1 + assert process_kill_calls == 2 + retained_entries = tuple(session._unregistered_pty_processes.values()) + assert len(retained_entries) == 1 + assert retained_entries[0].process is process + + allow_group_kill = True + await session.pty_terminate_all() + + assert killpg_calls == [ + (1234, signal.SIGKILL), + (1234, signal.SIGKILL), + (1234, signal.SIGKILL), + ] + assert process_kill_calls == 2 assert session._unregistered_pty_processes == {} @pytest.mark.asyncio From 60c2d7e1ac8150536a58dc8db8ab488cc5d50c4f Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Mon, 31 Aug 2026 15:37:06 +0800 Subject: [PATCH 10/12] fix(sandbox): preserve process group cleanup errors --- src/agents/sandbox/sandboxes/unix_local.py | 3 +- tests/sandbox/test_unix_local.py | 34 ++++++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 0510e41884..244cc07afa 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -634,7 +634,8 @@ async def _terminate_pty_entry(self, entry: _UnixPtyProcessEntry) -> None: try: process.kill() except ProcessLookupError: - pass + if not isinstance(group_error, ProcessLookupError): + termination_error = group_error except OSError as process_error: termination_error = process_error termination_cause = group_error diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index f7e66717f9..12dc0213da 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -396,6 +396,40 @@ async def blocked_to_thread(*args: object, **kwargs: object) -> None: assert session._fd_close_tasks == set() + @pytest.mark.asyncio + async def test_pty_termination_preserves_group_error_when_leader_disappears( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + session = _RecordingUnixLocalSession(tmp_path) + allow_group_kill = False + + def killpg(pid: int, signum: signal.Signals) -> None: + _ = (pid, signum) + if not allow_group_kill: + raise PermissionError("synthetic group cleanup failure") + + def kill() -> None: + raise ProcessLookupError("synthetic missing leader") + + monkeypatch.setattr(unix_local_module.os, "killpg", killpg) + process = cast( + asyncio.subprocess.Process, + SimpleNamespace(returncode=None, pid=1234, kill=kill), + ) + entry = _UnixPtyProcessEntry(process=process, tty=False) + + with pytest.raises(PermissionError, match="synthetic group cleanup failure"): + await session._terminate_pty_entry(entry) + + assert tuple(session._unregistered_pty_processes.values()) == (entry,) + + allow_group_kill = True + await session.pty_terminate_all() + + assert session._unregistered_pty_processes == {} + @pytest.mark.asyncio @pytest.mark.requires_native_macos_sandbox async def test_pty_exec_write_poll_and_unknown_session_errors(self, tmp_path: Path) -> None: From f1fd97090b299dc0aadc1d1a0cf9cfbace5f9d1c Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Mon, 31 Aug 2026 15:55:22 +0800 Subject: [PATCH 11/12] fix(sandbox): retry retained PTY process groups --- src/agents/sandbox/sandboxes/unix_local.py | 14 ++++++++++++-- tests/sandbox/test_unix_local.py | 3 ++- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 244cc07afa..a363d11593 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -144,6 +144,7 @@ class _UnixPtyProcessEntry: process: asyncio.subprocess.Process tty: bool primary_fd: int | None = None + group_termination_pending: bool = False last_used: float = field(default_factory=time.monotonic) output_chunks: deque[bytes] = field(default_factory=deque) output_lock: asyncio.Lock = field(default_factory=asyncio.Lock) @@ -625,9 +626,11 @@ async def _terminate_pty_entry(self, entry: _UnixPtyProcessEntry) -> None: termination_error: OSError | None = None termination_cause: OSError | None = None - if process.returncode is None and process.pid is None: + should_terminate_group = process.returncode is None or entry.group_termination_pending + if should_terminate_group and process.pid is None: termination_error = OSError("Cannot terminate PTY process without a PID") - if process.returncode is None and process.pid is not None: + if should_terminate_group and process.pid is not None: + entry.group_termination_pending = False try: os.killpg(process.pid, signal.SIGKILL) except OSError as group_error: @@ -635,10 +638,17 @@ async def _terminate_pty_entry(self, entry: _UnixPtyProcessEntry) -> None: process.kill() except ProcessLookupError: if not isinstance(group_error, ProcessLookupError): + entry.group_termination_pending = True termination_error = group_error except OSError as process_error: + if not isinstance(group_error, ProcessLookupError): + entry.group_termination_pending = True termination_error = process_error termination_cause = group_error + else: + entry.group_termination_pending = False + else: + entry.group_termination_pending = False for task in entry.pump_tasks: task.cancel() diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index 12dc0213da..a5ebe1c357 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -255,7 +255,7 @@ async def read(self, size: int) -> bytes: class _Process: pid = 1234 - returncode = None + returncode: int | None = None stdout = _Stream() stderr = _Stream() @@ -322,6 +322,7 @@ def killpg(pid: int, signum: signal.Signals) -> None: assert len(retained_entries) == 1 assert retained_entries[0].process is process + process.returncode = -signal.SIGKILL allow_group_kill = True await session.pty_terminate_all() From d9e614cf5d6a0cacf425ad9fc33e127bde8531e2 Mon Sep 17 00:00:00 2001 From: HughhhhCoder Date: Mon, 31 Aug 2026 16:08:27 +0800 Subject: [PATCH 12/12] fix(sandbox): remove unreachable PTY cancellation state --- src/agents/sandbox/sandboxes/unix_local.py | 100 +++----------- tests/sandbox/test_unix_local.py | 152 --------------------- 2 files changed, 20 insertions(+), 232 deletions(-) diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index a363d11593..308f6038ed 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -53,7 +53,6 @@ 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, @@ -144,7 +143,6 @@ class _UnixPtyProcessEntry: process: asyncio.subprocess.Process tty: bool primary_fd: int | None = None - group_termination_pending: bool = False last_used: float = field(default_factory=time.monotonic) output_chunks: deque[bytes] = field(default_factory=deque) output_lock: asyncio.Lock = field(default_factory=asyncio.Lock) @@ -164,7 +162,6 @@ class UnixLocalSandboxSession(BaseSandboxSession): _running: bool _pty_lock: asyncio.Lock _pty_processes: dict[int, _UnixPtyProcessEntry] - _unregistered_pty_processes: dict[int, _UnixPtyProcessEntry] _reserved_pty_process_ids: set[int] _fd_close_tasks: set[asyncio.Task[None]] _host_environment_allowlist: frozenset[str] | None @@ -174,7 +171,6 @@ def __init__(self, *, state: UnixLocalSandboxSessionState) -> None: self._running = False self._pty_lock = asyncio.Lock() self._pty_processes = {} - self._unregistered_pty_processes = {} self._reserved_pty_process_ids = set() self._fd_close_tasks = set() self._host_environment_allowlist = None @@ -384,31 +380,15 @@ def _preexec() -> None: asyncio.create_task(self._pump_process_stream(entry, process.stderr)), ] - self._unregistered_pty_processes[id(entry)] = entry - registered = False - try: - entry.wait_task = asyncio.create_task(self._watch_process_exit(entry)) + 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 - self._unregistered_pty_processes.pop(id(entry), None) - process_count = len(self._pty_processes) - registered = True - except asyncio.CancelledError as cancellation: - if not registered: - await _settle_pty_cleanup( - self._terminate_pty_entry(entry), - initial_cancellation=cancellation, - ) - raise - except BaseException: - if not registered: - await _settle_pty_cleanup(self._terminate_pty_entry(entry)) - raise + 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) if pruned_entry is not None: await self._terminate_pty_entry(pruned_entry) @@ -479,26 +459,12 @@ async def pty_write_stdin( async def pty_terminate_all(self) -> None: async with self._pty_lock: - entries_by_id = { - id(entry): entry - for entry in ( - *self._pty_processes.values(), - *self._unregistered_pty_processes.values(), - ) - } + entries = list(self._pty_processes.values()) self._pty_processes.clear() self._reserved_pty_process_ids.clear() - self._unregistered_pty_processes.update(entries_by_id) - first_error: Exception | None = None - for entry in entries_by_id.values(): - try: - await self._terminate_pty_entry(entry) - except Exception as error: - if first_error is None: - first_error = error - if first_error is not None: - raise first_error + for entry in entries: + await self._terminate_pty_entry(entry) async def _resolved_exec_context(self) -> tuple[dict[str, str], str]: if self._host_environment_allowlist is None: @@ -622,33 +588,10 @@ async def _terminate_pty_entry(self, entry: _UnixPtyProcessEntry) -> None: process = entry.process primary_fd = entry.primary_fd entry.primary_fd = None - self._unregistered_pty_processes[id(entry)] = entry - - termination_error: OSError | None = None - termination_cause: OSError | None = None - should_terminate_group = process.returncode is None or entry.group_termination_pending - if should_terminate_group and process.pid is None: - termination_error = OSError("Cannot terminate PTY process without a PID") - if should_terminate_group and process.pid is not None: - entry.group_termination_pending = False - try: + + if process.returncode is None and process.pid is not None: + with suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) - except OSError as group_error: - try: - process.kill() - except ProcessLookupError: - if not isinstance(group_error, ProcessLookupError): - entry.group_termination_pending = True - termination_error = group_error - except OSError as process_error: - if not isinstance(group_error, ProcessLookupError): - entry.group_termination_pending = True - termination_error = process_error - termination_cause = group_error - else: - entry.group_termination_pending = False - else: - entry.group_termination_pending = False for task in entry.pump_tasks: task.cancel() @@ -662,16 +605,13 @@ async def _terminate_pty_entry(self, entry: _UnixPtyProcessEntry) -> None: self._schedule_fd_close(primary_fd) entry.output_closed.set() entry.output_notify.set() - else: - if primary_fd is not None: - _close_fd_quietly(primary_fd) - await asyncio.gather(*entry.pump_tasks, return_exceptions=True) - if entry.wait_task is not None: - await asyncio.gather(entry.wait_task, return_exceptions=True) + return - if termination_error is not None: - raise termination_error from termination_cause - self._unregistered_pty_processes.pop(id(entry), None) + if primary_fd is not None: + _close_fd_quietly(primary_fd) + await asyncio.gather(*entry.pump_tasks, return_exceptions=True) + if entry.wait_task is not None: + await asyncio.gather(entry.wait_task, return_exceptions=True) def _schedule_fd_close(self, fd: int) -> None: task = asyncio.create_task(asyncio.to_thread(_close_fd_quietly, fd)) diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index a5ebe1c357..f2482be90f 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -3,7 +3,6 @@ import asyncio import io import signal -import sys import tarfile import threading import time @@ -217,123 +216,6 @@ def _unexpected_mkdtemp(*args: object, **kwargs: object) -> str: @pytest.mark.review_optional class TestUnixLocalPty: - @pytest.mark.asyncio - async def test_pty_start_cancellation_cleans_up_before_registration( - 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) - process_started = asyncio.Event() - pump_tasks_started = asyncio.Event() - pump_tasks_cancelled = asyncio.Event() - process_wait_started = asyncio.Event() - process_wait_cancelled = asyncio.Event() - pump_task_count = 0 - cancelled_pump_task_count = 0 - allow_group_kill = False - process_kill_calls = 0 - killpg_calls: list[tuple[int, signal.Signals]] = [] - - class _Stream: - async def read(self, size: int) -> bytes: - nonlocal pump_task_count, cancelled_pump_task_count - _ = size - pump_task_count += 1 - if pump_task_count == 2: - pump_tasks_started.set() - try: - return await asyncio.Future[bytes]() - except asyncio.CancelledError: - cancelled_pump_task_count += 1 - if cancelled_pump_task_count == 2: - pump_tasks_cancelled.set() - raise - - class _Process: - pid = 1234 - returncode: int | None = None - stdout = _Stream() - stderr = _Stream() - - async def wait(self) -> None: - process_wait_started.set() - try: - await asyncio.Future() - except asyncio.CancelledError: - process_wait_cancelled.set() - raise - - def kill(self) -> None: - nonlocal process_kill_calls - process_kill_calls += 1 - raise PermissionError("synthetic direct cleanup failure") - - process = _Process() - - async def create_subprocess(*args: object, **kwargs: object) -> _Process: - _ = (args, kwargs) - process_started.set() - return process - - def killpg(pid: int, signum: signal.Signals) -> None: - killpg_calls.append((pid, signum)) - if not allow_group_kill: - raise PermissionError("synthetic cleanup failure") - - monkeypatch.setattr(unix_local_module.asyncio, "create_subprocess_exec", create_subprocess) - monkeypatch.setattr(unix_local_module.os, "killpg", killpg) - await session._pty_lock.acquire() - try: - task = asyncio.create_task( - session.pty_exec_start("echo", "hello", shell=False, yield_time_s=0.01) - ) - await process_started.wait() - await pump_tasks_started.wait() - await process_wait_started.wait() - task.cancel("startup-cancel") - - with pytest.raises(asyncio.CancelledError) as exc_info: - await task - - if sys.version_info >= (3, 11): - assert exc_info.value.args == ("startup-cancel",) - assert killpg_calls == [(1234, signal.SIGKILL)] - assert process_kill_calls == 1 - assert pump_tasks_cancelled.is_set() - assert process_wait_cancelled.is_set() - assert session._pty_processes == {} - assert session._reserved_pty_process_ids == set() - retained_entries = tuple(session._unregistered_pty_processes.values()) - assert len(retained_entries) == 1 - assert retained_entries[0].process is process - finally: - session._pty_lock.release() - - with pytest.raises(PermissionError, match="synthetic direct cleanup failure"): - await session.pty_terminate_all() - - assert killpg_calls == [(1234, signal.SIGKILL), (1234, signal.SIGKILL)] - assert process_kill_calls == 2 - retained_entries = tuple(session._unregistered_pty_processes.values()) - assert len(retained_entries) == 1 - assert retained_entries[0].process is process - - process.returncode = -signal.SIGKILL - allow_group_kill = True - await session.pty_terminate_all() - - assert killpg_calls == [ - (1234, signal.SIGKILL), - (1234, signal.SIGKILL), - (1234, signal.SIGKILL), - ] - assert process_kill_calls == 2 - assert session._unregistered_pty_processes == {} - @pytest.mark.asyncio async def test_tty_start_cancellation_closes_open_file_descriptors( self, @@ -397,40 +279,6 @@ async def blocked_to_thread(*args: object, **kwargs: object) -> None: assert session._fd_close_tasks == set() - @pytest.mark.asyncio - async def test_pty_termination_preserves_group_error_when_leader_disappears( - self, - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - session = _RecordingUnixLocalSession(tmp_path) - allow_group_kill = False - - def killpg(pid: int, signum: signal.Signals) -> None: - _ = (pid, signum) - if not allow_group_kill: - raise PermissionError("synthetic group cleanup failure") - - def kill() -> None: - raise ProcessLookupError("synthetic missing leader") - - monkeypatch.setattr(unix_local_module.os, "killpg", killpg) - process = cast( - asyncio.subprocess.Process, - SimpleNamespace(returncode=None, pid=1234, kill=kill), - ) - entry = _UnixPtyProcessEntry(process=process, tty=False) - - with pytest.raises(PermissionError, match="synthetic group cleanup failure"): - await session._terminate_pty_entry(entry) - - assert tuple(session._unregistered_pty_processes.values()) == (entry,) - - allow_group_kill = True - await session.pty_terminate_all() - - assert session._unregistered_pty_processes == {} - @pytest.mark.asyncio @pytest.mark.requires_native_macos_sandbox async def test_pty_exec_write_poll_and_unknown_session_errors(self, tmp_path: Path) -> None: