From 62c03d70d850fddb47c60e4d034cd84b4c9d4cbc Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Tue, 8 Sep 2026 00:05:41 +0900 Subject: [PATCH] fix(sandbox): preserve PTY output through settlement Co-authored-by: Kazuhiro Sera Co-authored-by: Henry Su Co-authored-by: ayaangazali --- .../extensions/sandbox/blaxel/sandbox.py | 13 +- .../extensions/sandbox/cloudflare/sandbox.py | 50 +- .../extensions/sandbox/daytona/sandbox.py | 33 +- src/agents/extensions/sandbox/e2b/sandbox.py | 62 +-- .../extensions/sandbox/modal/sandbox.py | 149 ++++-- src/agents/sandbox/sandboxes/docker.py | 15 +- src/agents/sandbox/sandboxes/unix_local.py | 13 +- src/agents/sandbox/session/pty_output.py | 154 +++++- tests/extensions/sandbox/test_blaxel.py | 15 +- tests/extensions/sandbox/test_daytona.py | 77 ++- tests/extensions/sandbox/test_e2b.py | 171 ++++++- tests/extensions/sandbox/test_modal.py | 441 ++++++++++++++++++ tests/sandbox/test_pty_output.py | 287 +++++++++++- tests/sandbox/test_unix_local.py | 69 +++ 14 files changed, 1369 insertions(+), 180 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index a54bc47905..0ba8ad0d5b 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -864,7 +864,7 @@ async def pty_exec_start( ) yield_time_ms = 10_000 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms), max_output_tokens=max_output_tokens, @@ -874,6 +874,7 @@ async def pty_exec_start( entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def pty_write_stdin( @@ -898,7 +899,7 @@ async def pty_write_stdin( await asyncio.sleep(0.1) yield_time_ms = 250 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=resolve_pty_write_yield_time_ms( yield_time_ms=yield_time_ms, input_empty=chars == "" @@ -911,6 +912,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def pty_terminate_all(self) -> None: @@ -972,7 +974,7 @@ async def _collect_pty_output( entry: _BlaxelPtySessionEntry, yield_time_ms: int, max_output_tokens: int | None, - ) -> tuple[bytes, int | None]: + ) -> tuple[bytes, int | None, bool]: return await collect_pty_output( output_chunks=entry.output_chunks, output_lock=entry.output_lock, @@ -989,11 +991,12 @@ async def _finalize_pty_update( entry: _BlaxelPtySessionEntry, output: bytes, original_token_count: int | None, + output_closed: bool, ) -> PtyExecUpdate: - exit_code = entry.exit_code if entry.done else None + exit_code = entry.exit_code if output_closed else None live_process_id: int | None = process_id - if entry.done: + if output_closed: async with self._pty_lock: removed = self._pty_sessions.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index bb8d7c37e6..fc2f6c0696 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -59,6 +59,7 @@ _settle_mount_transition, with_ephemeral_mounts_removed, ) +from ....sandbox.session.pty_output import collect_pty_output from ....sandbox.session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -67,7 +68,6 @@ clamp_pty_yield_time_ms, process_id_to_prune_from_meta, resolve_pty_write_yield_time_ms, - truncate_text_by_tokens, ) from ....sandbox.session.runtime_helpers import RESOLVE_WORKSPACE_PATH_HELPER, RuntimeHelperScript from ....sandbox.session.sandbox_client import BaseSandboxClient, BaseSandboxClientOptions @@ -1033,34 +1033,15 @@ async def _collect_pty_output( entry: _CloudflarePtyProcessEntry, yield_time_ms: int, max_output_tokens: int | None, - ) -> tuple[bytes, int | None]: - deadline = time.monotonic() + (yield_time_ms / 1000) - output = bytearray() - - while True: - async with entry.output_lock: - while entry.output_chunks: - output.extend(entry.output_chunks.popleft()) - - if entry.output_closed.is_set(): - async with entry.output_lock: - while entry.output_chunks: - output.extend(entry.output_chunks.popleft()) - break - - remaining_s = deadline - time.monotonic() - if remaining_s <= 0: - break - - try: - await asyncio.wait_for(entry.output_notify.wait(), timeout=remaining_s) - except asyncio.TimeoutError: - break - entry.output_notify.clear() - - text = output.decode("utf-8", errors="replace") - truncated_text, original_token_count = truncate_text_by_tokens(text, max_output_tokens) - return truncated_text.encode("utf-8", errors="replace"), original_token_count + ) -> tuple[bytes, int | None, bool]: + return await collect_pty_output( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output_notify=entry.output_notify, + is_done=entry.output_closed.is_set, + yield_time_ms=yield_time_ms, + max_output_tokens=max_output_tokens, + ) async def _finalize_pty_update( self, @@ -1069,10 +1050,11 @@ async def _finalize_pty_update( entry: _CloudflarePtyProcessEntry, output: bytes, original_token_count: int | None, + output_closed: bool, ) -> PtyExecUpdate: - exit_code = entry.exit_code if entry.output_closed.is_set() else None + exit_code = entry.exit_code if output_closed else None live_process_id: int | None = process_id - if entry.output_closed.is_set(): + if output_closed: async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) @@ -1220,7 +1202,7 @@ async def pty_exec_start( ) yield_time_ms = 10_000 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms), max_output_tokens=max_output_tokens, @@ -1230,6 +1212,7 @@ async def pty_exec_start( entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def pty_write_stdin( @@ -1253,7 +1236,7 @@ async def pty_write_stdin( await asyncio.sleep(0.1) yield_time_ms = 250 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=resolve_pty_write_yield_time_ms( yield_time_ms=yield_time_ms, @@ -1267,6 +1250,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def pty_terminate_all(self) -> None: diff --git a/src/agents/extensions/sandbox/daytona/sandbox.py b/src/agents/extensions/sandbox/daytona/sandbox.py index d62c5021ad..e7c9cb8918 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -387,8 +387,8 @@ class _DaytonaPtySessionEntry: output_chunks: deque[bytes] = field(default_factory=deque) output_lock: asyncio.Lock = field(default_factory=asyncio.Lock) output_notify: asyncio.Event = field(default_factory=asyncio.Event) + output_closed: asyncio.Event = field(default_factory=asyncio.Event) last_used: float = field(default_factory=time.monotonic) - done: bool = False exit_code: int | None = None worker_task: asyncio.Task[None] | None = None @@ -755,7 +755,7 @@ async def _on_data(chunk: bytes | str) -> None: ) yield_time_ms = 10_000 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms), max_output_tokens=max_output_tokens, @@ -765,6 +765,7 @@ async def _on_data(chunk: bytes | str) -> None: entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def _run_pty_waiter(self, entry: _DaytonaPtySessionEntry) -> None: @@ -776,7 +777,10 @@ async def _run_pty_waiter(self, entry: _DaytonaPtySessionEntry) -> None: except Exception: pass finally: - entry.done = True + # AsyncPtyHandle.wait() completes only after its WebSocket reader exits. + # That reader awaits every async on_data callback before it can finish, + # so this is Daytona's authoritative output-stream close boundary. + entry.output_closed.set() entry.output_notify.set() async def _run_session_reader( @@ -801,11 +805,13 @@ async def _run_session_reader( cmd = await self._sandbox.process.get_session_command(session_id, cmd_id) if cmd.exit_code is not None: entry.exit_code = int(cmd.exit_code) - entry.done = True except Exception: pass - if not logs_failed: - entry.done = True + # Once the log callback stream has returned, or has failed after the + # provider reports a final exit code, this worker is the only output + # producer and no later callback can append bytes. + if not logs_failed or entry.exit_code is not None: + entry.output_closed.set() entry.output_notify.set() async def pty_write_stdin( @@ -832,7 +838,7 @@ async def pty_write_stdin( await asyncio.sleep(0.1) yield_time_ms = 250 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=resolve_pty_write_yield_time_ms( yield_time_ms=yield_time_ms, input_empty=chars == "" @@ -845,6 +851,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def _finalize_pty_update( @@ -854,11 +861,12 @@ async def _finalize_pty_update( entry: _DaytonaPtySessionEntry, output: bytes, original_token_count: int | None, + output_closed: bool, ) -> PtyExecUpdate: - exit_code = entry.exit_code if entry.done else None + exit_code = entry.exit_code if output_closed else None live_process_id: int | None = process_id - if entry.done: + if output_closed: async with self._pty_lock: removed = self._pty_sessions.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) @@ -887,12 +895,12 @@ async def _collect_pty_output( entry: _DaytonaPtySessionEntry, yield_time_ms: int, max_output_tokens: int | None, - ) -> tuple[bytes, int | None]: + ) -> tuple[bytes, int | None, bool]: return await collect_pty_output( output_chunks=entry.output_chunks, output_lock=entry.output_lock, output_notify=entry.output_notify, - is_done=lambda: entry.done, + is_done=entry.output_closed.is_set, yield_time_ms=yield_time_ms, max_output_tokens=max_output_tokens, ) @@ -901,7 +909,8 @@ def _prune_pty_sessions_if_needed(self) -> _DaytonaPtySessionEntry | None: if len(self._pty_sessions) < PTY_PROCESSES_MAX: return None meta: list[tuple[int, float, bool]] = [ - (pid, entry.last_used, entry.done) for pid, entry in self._pty_sessions.items() + (pid, entry.last_used, entry.output_closed.is_set()) + for pid, entry in self._pty_sessions.items() ] pid = process_id_to_prune_from_meta(meta) if pid is None: diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index 389b665c44..e54c238801 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -53,6 +53,7 @@ from ....sandbox.session.base_sandbox_session import BaseSandboxSession from ....sandbox.session.dependencies import Dependencies from ....sandbox.session.manager import Instrumentation +from ....sandbox.session.pty_output import collect_pty_output from ....sandbox.session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -61,7 +62,6 @@ clamp_pty_yield_time_ms, process_id_to_prune_from_meta, resolve_pty_write_yield_time_ms, - truncate_text_by_tokens, ) from ....sandbox.session.runtime_helpers import RESOLVE_WORKSPACE_PATH_HELPER, RuntimeHelperScript from ....sandbox.session.sandbox_client import BaseSandboxClient, BaseSandboxClientOptions @@ -687,6 +687,7 @@ class _E2BPtyProcessEntry: output_chunks: deque[bytes] = field(default_factory=deque) output_lock: asyncio.Lock = field(default_factory=asyncio.Lock) output_notify: asyncio.Event = field(default_factory=asyncio.Event) + output_closed: asyncio.Event = field(default_factory=asyncio.Event) last_used: float = field(default_factory=time.monotonic) exit_code: int | None = None wait_task: asyncio.Task[None] | None = None @@ -1050,7 +1051,7 @@ async def _append_output(payload: bytes | bytearray | str | object) -> None: ) yield_time_ms = 10_000 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms), max_output_tokens=max_output_tokens, @@ -1060,6 +1061,7 @@ async def _append_output(payload: bytes | bytearray | str | object) -> None: entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def pty_write_stdin( @@ -1087,7 +1089,7 @@ async def pty_write_stdin( await asyncio.sleep(0.1) yield_time_ms = 250 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=resolve_pty_write_yield_time_ms( yield_time_ms=yield_time_ms, input_empty=chars == "" @@ -1100,6 +1102,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def pty_terminate_all(self) -> None: @@ -1211,37 +1214,15 @@ async def _collect_pty_output( entry: _E2BPtyProcessEntry, yield_time_ms: int, max_output_tokens: int | None, - ) -> tuple[bytes, int | None]: - deadline = time.monotonic() + (yield_time_ms / 1000) - output = bytearray() - - while True: - async with entry.output_lock: - while entry.output_chunks: - output.extend(entry.output_chunks.popleft()) - - if time.monotonic() >= deadline: - break - - if self._entry_exit_code(entry) is not None: - async with entry.output_lock: - while entry.output_chunks: - output.extend(entry.output_chunks.popleft()) - break - - remaining_s = deadline - time.monotonic() - if remaining_s <= 0: - break - - try: - await asyncio.wait_for(entry.output_notify.wait(), timeout=remaining_s) - except asyncio.TimeoutError: - break - entry.output_notify.clear() - - text = output.decode("utf-8", errors="replace") - truncated_text, original_token_count = truncate_text_by_tokens(text, max_output_tokens) - return truncated_text.encode("utf-8", errors="replace"), original_token_count + ) -> tuple[bytes, int | None, bool]: + return await collect_pty_output( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output_notify=entry.output_notify, + is_done=entry.output_closed.is_set, + yield_time_ms=yield_time_ms, + max_output_tokens=max_output_tokens, + ) async def _run_pty_waiter(self, entry: _E2BPtyProcessEntry) -> None: try: @@ -1258,7 +1239,13 @@ async def _run_pty_waiter(self, entry: _E2BPtyProcessEntry) -> None: entry.exit_code = int(value) except (TypeError, ValueError): pass - finally: + if entry.exit_code is not None: + # E2B delivers output through async callbacks that append under this lock. + # Wait behind callbacks already appending terminal bytes before publishing + # the close signal that authorizes collector settlement and PTY removal. + async with entry.output_lock: + pass + entry.output_closed.set() entry.output_notify.set() async def _finalize_pty_update( @@ -1268,8 +1255,9 @@ async def _finalize_pty_update( entry: _E2BPtyProcessEntry, output: bytes, original_token_count: int | None, + output_closed: bool, ) -> PtyExecUpdate: - exit_code = self._entry_exit_code(entry) + exit_code = self._entry_exit_code(entry) if output_closed else None live_process_id: int | None = process_id if exit_code is not None: @@ -1292,7 +1280,7 @@ def _prune_pty_processes_if_needed(self) -> _E2BPtyProcessEntry | None: return None meta: list[tuple[int, float, bool]] = [ - (process_id, entry.last_used, self._entry_exit_code(entry) is not None) + (process_id, entry.last_used, entry.output_closed.is_set()) for process_id, entry in self._pty_processes.items() ] process_id = process_id_to_prune_from_meta(meta) diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index 848b8e12a3..f0fe85e589 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -23,6 +23,7 @@ import shlex import time import uuid +from collections import deque from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import asynccontextmanager from dataclasses import dataclass, field @@ -64,6 +65,7 @@ _settle_mount_transition, _terminate_ambiguous_mount_session, ) +from ....sandbox.session.pty_output import collect_pty_output from ....sandbox.session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -72,7 +74,6 @@ clamp_pty_yield_time_ms, process_id_to_prune_from_meta, resolve_pty_write_yield_time_ms, - truncate_text_by_tokens, ) from ....sandbox.session.runtime_helpers import RESOLVE_WORKSPACE_PATH_HELPER, RuntimeHelperScript from ....sandbox.session.sandbox_client import BaseSandboxClient, BaseSandboxClientOptions @@ -489,6 +490,12 @@ class _ModalPtyProcessEntry: stderr_iter: AsyncIterator[object] | None = None stdout_read_task: asyncio.Task[object] | None = None stderr_read_task: asyncio.Task[object] | None = None + stdout_closed: bool = False + stderr_closed: bool = False + output_chunks: deque[bytes] = field(default_factory=deque) + output_lock: asyncio.Lock = field(default_factory=asyncio.Lock) + output_notify: asyncio.Event = field(default_factory=asyncio.Event) + output_closed: asyncio.Event = field(default_factory=asyncio.Event) class ModalSandboxSession(BaseSandboxSession): @@ -907,7 +914,7 @@ async def pty_exec_start( ) yield_time_ms = 10_000 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms), max_output_tokens=max_output_tokens, @@ -917,6 +924,7 @@ async def pty_exec_start( entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def pty_write_stdin( @@ -940,7 +948,7 @@ async def pty_write_stdin( await asyncio.sleep(0.1) yield_time_ms = 250 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=resolve_pty_write_yield_time_ms( yield_time_ms=yield_time_ms, input_empty=chars == "" @@ -953,6 +961,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def pty_terminate_all(self) -> None: @@ -981,56 +990,80 @@ async def _collect_pty_output( entry: _ModalPtyProcessEntry, yield_time_ms: int, max_output_tokens: int | None, - ) -> tuple[bytes, int | None]: - deadline = time.monotonic() + (yield_time_ms / 1000) - chunks = bytearray() - - while True: - stdout_chunk = await self._read_modal_stream(entry=entry, stream_name="stdout") - stderr_chunk = await self._read_modal_stream(entry=entry, stream_name="stderr") - if stdout_chunk: - chunks.extend(stdout_chunk) - if stderr_chunk: - chunks.extend(stderr_chunk) - - if time.monotonic() >= deadline: - break - - exit_code = await self._peek_exit_code(entry.process) - if exit_code is not None: - stdout_chunks = await self._drain_modal_stream(entry=entry, stream_name="stdout") - stderr_chunks = await self._drain_modal_stream(entry=entry, stream_name="stderr") - chunks.extend(stdout_chunks) - chunks.extend(stderr_chunks) - break - - if not stdout_chunk and not stderr_chunk: - remaining_s = deadline - time.monotonic() - if remaining_s <= 0: - break - await asyncio.sleep(min(_PTY_POLL_INTERVAL_S, remaining_s)) - - text = chunks.decode("utf-8", errors="replace") - truncated_text, original_token_count = truncate_text_by_tokens(text, max_output_tokens) - return truncated_text.encode("utf-8", errors="replace"), original_token_count + ) -> tuple[bytes, int | None, bool]: + async def poll_output( + *, + deadline: float | None, + allow_new_read: bool, + check_exit: bool, + ) -> None: + for stream_name in ("stdout", "stderr"): + chunk = await self._read_modal_stream( + entry=entry, + stream_name=stream_name, + allow_new_read=allow_new_read, + deadline=deadline, + ) + if chunk: + # Commit consumed bytes before another provider call can suspend. + entry.output_chunks.append(chunk) + entry.output_notify.set() + + if ( + check_exit + and deadline is not None + and time.monotonic() < deadline + and await self._peek_exit_code(entry.process) is not None + ): + await self._drain_modal_stream(entry=entry, stream_name="stdout", deadline=deadline) + await self._drain_modal_stream(entry=entry, stream_name="stderr", deadline=deadline) + if entry.stdout_closed and entry.stderr_closed: + entry.output_closed.set() + + async def wait_for_output(remaining_s: float) -> None: + await asyncio.sleep(min(_PTY_POLL_INTERVAL_S, remaining_s)) + + async def settle_output() -> None: + # The deadline path may only collect already-started reads. Checking + # process status here would add a slow provider RPC after the caller's + # requested yield window has already elapsed. + await poll_output(deadline=None, allow_new_read=False, check_exit=False) + + return await collect_pty_output( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output_notify=entry.output_notify, + is_done=entry.output_closed.is_set, + yield_time_ms=yield_time_ms, + max_output_tokens=max_output_tokens, + poll_output=lambda deadline: poll_output( + deadline=deadline, + allow_new_read=True, + check_exit=True, + ), + settle_output=settle_output, + wait_for_output=wait_for_output, + ) async def _drain_modal_stream( self, *, entry: _ModalPtyProcessEntry, stream_name: Literal["stdout", "stderr"], - ) -> bytes: - chunks = bytearray() + deadline: float, + ) -> None: while True: chunk = await self._read_modal_stream( entry=entry, stream_name=stream_name, await_pending=True, + settle_after_exit=True, + deadline=deadline, ) if not chunk: break - chunks.extend(chunk) - return bytes(chunks) + entry.output_chunks.append(chunk) + entry.output_notify.set() async def _read_modal_stream( self, @@ -1038,13 +1071,25 @@ async def _read_modal_stream( entry: _ModalPtyProcessEntry, stream_name: Literal["stdout", "stderr"], await_pending: bool = False, + allow_new_read: bool = True, + settle_after_exit: bool = False, + deadline: float | None = None, ) -> bytes: stream = entry.process.stdout if stream_name == "stdout" else entry.process.stderr + closed_attr = "stdout_closed" if stream_name == "stdout" else "stderr_closed" + if getattr(entry, closed_attr): + return b"" if stream is None: + setattr(entry, closed_attr, True) return b"" iter_attr = "stdout_iter" if stream_name == "stdout" else "stderr_iter" task_attr = "stdout_read_task" if stream_name == "stdout" else "stderr_read_task" + task = getattr(entry, task_attr) + remaining_s = max(0.0, deadline - time.monotonic()) if deadline is not None else 0.0 + if task is None and (not allow_new_read or remaining_s <= 0): + return b"" + stream_iter = getattr(entry, iter_attr) if stream_iter is None: aiter_method = getattr(stream, "__aiter__", None) @@ -1056,13 +1101,12 @@ async def _read_modal_stream( else: setattr(entry, iter_attr, stream_iter) - task = getattr(entry, task_attr) if task is None and stream_iter is not None: task = asyncio.create_task(stream_iter.__anext__()) setattr(entry, task_attr, task) if task is not None: - wait_timeout = 0.2 if await_pending else 0 + wait_timeout = min(0.2, remaining_s) if await_pending else 0 done, _pending = await asyncio.wait({task}, timeout=wait_timeout) if not done: return b"" @@ -1072,9 +1116,12 @@ async def _read_modal_stream( value = task.result() except StopAsyncIteration: setattr(entry, iter_attr, None) + setattr(entry, closed_attr, True) return b"" except Exception: setattr(entry, iter_attr, None) + if settle_after_exit: + setattr(entry, closed_attr, True) return b"" return self._coerce_modal_stream_chunk(value) @@ -1084,13 +1131,23 @@ async def _read_modal_stream( return b"" try: - value = await self._call_modal(read, 16_384, call_timeout=0.2) + value = await self._call_modal(read, 16_384, call_timeout=min(0.2, remaining_s)) + except asyncio.TimeoutError: + # Reaching this window's deadline is not evidence of stream EOF. + return b"" except TypeError: + if settle_after_exit: + setattr(entry, closed_attr, True) return b"" except Exception: + if settle_after_exit: + setattr(entry, closed_attr, True) return b"" - return self._coerce_modal_stream_chunk(value) + chunk = self._coerce_modal_stream_chunk(value) + if not chunk and settle_after_exit: + setattr(entry, closed_attr, True) + return chunk def _coerce_modal_stream_chunk(self, value: object) -> bytes: if value is None: @@ -1110,8 +1167,9 @@ async def _finalize_pty_update( entry: _ModalPtyProcessEntry, output: bytes, original_token_count: int | None, + output_closed: bool, ) -> PtyExecUpdate: - exit_code = await self._peek_exit_code(entry.process) + exit_code = await self._peek_exit_code(entry.process) if output_closed else None live_process_id: int | None = process_id if exit_code is not None: async with self._pty_lock: @@ -1134,8 +1192,7 @@ async def _prune_pty_processes_if_needed(self) -> _ModalPtyProcessEntry | None: meta: list[tuple[int, float, bool]] = [] for process_id, entry in self._pty_processes.items(): - exit_code = await self._peek_exit_code(entry.process) - meta.append((process_id, entry.last_used, exit_code is not None)) + meta.append((process_id, entry.last_used, entry.output_closed.is_set())) process_id_to_prune = process_id_to_prune_from_meta(meta) if process_id_to_prune is None: return None diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index 8ca4febe85..3bc69a043b 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -1080,7 +1080,7 @@ async def pty_exec_start( ) yield_time_ms = 10_000 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms), max_output_tokens=max_output_tokens, @@ -1090,6 +1090,7 @@ async def pty_exec_start( entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def pty_write_stdin( @@ -1126,7 +1127,7 @@ async def pty_write_stdin( await asyncio.sleep(0.1) yield_time_ms = 250 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=resolve_pty_write_yield_time_ms( yield_time_ms=yield_time_ms, input_empty=chars == "" @@ -1139,6 +1140,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def pty_terminate_all(self) -> None: @@ -1242,7 +1244,7 @@ async def _collect_pty_output( entry: _DockerPtyProcessEntry, yield_time_ms: int, max_output_tokens: int | None, - ) -> tuple[bytes, int | None]: + ) -> tuple[bytes, int | None, bool]: return await collect_pty_output( output_chunks=entry.output_chunks, output_lock=entry.output_lock, @@ -1259,11 +1261,12 @@ async def _finalize_pty_update( entry: _DockerPtyProcessEntry, output: bytes, original_token_count: int | None, + output_closed: bool, ) -> PtyExecUpdate: - if entry.output_closed.is_set() and entry.exit_code is None: + if output_closed and entry.exit_code is None: await self._refresh_pty_exit_code(entry) - exit_code = entry.exit_code + exit_code = entry.exit_code if output_closed else None live_process_id: int | None = process_id if exit_code is not None: @@ -1286,7 +1289,7 @@ def _prune_pty_processes_if_needed(self) -> _DockerPtyProcessEntry | None: return None meta = [ - (process_id, entry.last_used, entry.exit_code is not None) + (process_id, entry.last_used, entry.output_closed.is_set()) for process_id, entry in self._pty_processes.items() ] process_id = process_id_to_prune_from_meta(meta) diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 308f6038ed..28eeb265ef 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -400,7 +400,7 @@ def _preexec() -> None: ) yield_time_ms = 10_000 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms), max_output_tokens=max_output_tokens, @@ -410,6 +410,7 @@ def _preexec() -> None: entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def pty_write_stdin( @@ -442,7 +443,7 @@ async def pty_write_stdin( await asyncio.sleep(0.1) yield_time_ms = 250 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, output_closed = await self._collect_pty_output( entry=entry, yield_time_ms=resolve_pty_write_yield_time_ms( yield_time_ms=yield_time_ms, input_empty=chars == "" @@ -455,6 +456,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + output_closed=output_closed, ) async def pty_terminate_all(self) -> None: @@ -533,7 +535,7 @@ async def _collect_pty_output( entry: _UnixPtyProcessEntry, yield_time_ms: int, max_output_tokens: int | None, - ) -> tuple[bytes, int | None]: + ) -> tuple[bytes, int | None, bool]: return await collect_pty_output( output_chunks=entry.output_chunks, output_lock=entry.output_lock, @@ -550,8 +552,9 @@ async def _finalize_pty_update( entry: _UnixPtyProcessEntry, output: bytes, original_token_count: int | None, + output_closed: bool, ) -> PtyExecUpdate: - exit_code: int | None = entry.process.returncode + exit_code: int | None = entry.process.returncode if output_closed else None live_process_id: int | None = process_id if exit_code is not None: @@ -574,7 +577,7 @@ def _prune_pty_processes_if_needed(self) -> _UnixPtyProcessEntry | None: return None meta = [ - (process_id, entry.last_used, entry.process.returncode is not None) + (process_id, entry.last_used, entry.output_closed.is_set()) for process_id, entry in self._pty_processes.items() ] process_id = process_id_to_prune_from_meta(meta) diff --git a/src/agents/sandbox/session/pty_output.py b/src/agents/sandbox/session/pty_output.py index 25cbe774e7..b0f0b0fe3c 100644 --- a/src/agents/sandbox/session/pty_output.py +++ b/src/agents/sandbox/session/pty_output.py @@ -3,11 +3,78 @@ import asyncio import time from collections import deque -from collections.abc import Callable +from collections.abc import Awaitable, Callable from .pty_types import truncate_text_by_tokens +def _incomplete_utf8_suffix_length(data: bytes | bytearray) -> int: + """Return the trailing byte count that can still become one valid UTF-8 scalar.""" + continuation_count = 0 + for byte in reversed(data[-3:]): + if 0x80 <= byte <= 0xBF: + continuation_count += 1 + continue + break + + lead_index = len(data) - continuation_count - 1 + if lead_index < 0: + return 0 + + lead = data[lead_index] + if 0xC2 <= lead <= 0xDF: + expected_length = 2 + elif 0xE0 <= lead <= 0xEF: + expected_length = 3 + elif 0xF0 <= lead <= 0xF4: + expected_length = 4 + else: + return 0 + + suffix_length = continuation_count + 1 + if suffix_length >= expected_length: + return 0 + + if continuation_count: + second = data[lead_index + 1] + if ( + (lead == 0xE0 and second < 0xA0) + or (lead == 0xED and second > 0x9F) + or (lead == 0xF0 and second < 0x90) + or (lead == 0xF4 and second > 0x8F) + ): + return 0 + + return suffix_length + + +async def _drain_output_chunks( + output_chunks: deque[bytes], + output_lock: asyncio.Lock, + output: bytearray, +) -> None: + async with output_lock: + while output_chunks: + output.extend(output_chunks.popleft()) + + +async def _drain_and_carry_incomplete_suffix( + output_chunks: deque[bytes], + output_lock: asyncio.Lock, + output: bytearray, +) -> None: + """Drain and restore a carryable suffix without yielding between ownership changes.""" + async with output_lock: + while output_chunks: + output.extend(output_chunks.popleft()) + + carry = _incomplete_utf8_suffix_length(output) + if carry: + tail = bytes(output[-carry:]) + del output[-carry:] + output_chunks.appendleft(tail) + + async def collect_pty_output( *, output_chunks: deque[bytes], @@ -16,35 +83,76 @@ async def collect_pty_output( is_done: Callable[[], bool], yield_time_ms: int, max_output_tokens: int | None, -) -> tuple[bytes, int | None]: - """Collect and truncate PTY output until the deadline or provider completion.""" + poll_output: Callable[[float], Awaitable[None]] | None = None, + settle_output: Callable[[], Awaitable[None]] | None = None, + wait_for_output: Callable[[float], Awaitable[None]] | None = None, +) -> tuple[bytes, int | None, bool]: + """Collect raw PTY bytes until the deadline or producer completion. + + poll_output adapts pull-based providers into output_chunks. Queue draining, + timeout settlement, UTF-8 carry, and decoding remain shared for every backend. + """ deadline = time.monotonic() + (yield_time_ms / 1000) output = bytearray() + output_closed = False + + try: + while True: + if time.monotonic() >= deadline: + break + + if poll_output is not None: + await poll_output(deadline) + await _drain_output_chunks(output_chunks, output_lock, output) + + if time.monotonic() >= deadline: + break + + if is_done(): + output_closed = True + if settle_output is not None: + await settle_output() + elif poll_output is not None: + await poll_output(deadline) + await _drain_output_chunks(output_chunks, output_lock, output) + break - while True: - async with output_lock: - while output_chunks: - output.extend(output_chunks.popleft()) + remaining_s = deadline - time.monotonic() + if remaining_s <= 0: + break - if time.monotonic() >= deadline: - break + if wait_for_output is not None: + await wait_for_output(remaining_s) + else: + try: + await asyncio.wait_for(output_notify.wait(), timeout=remaining_s) + except asyncio.TimeoutError: + break + output_notify.clear() - if is_done(): - async with output_lock: - while output_chunks: - output.extend(output_chunks.popleft()) - break + # Settle bytes that were queued around the final deadline or completion check. + if settle_output is not None: + await settle_output() + elif poll_output is not None: + await poll_output(deadline) - remaining_s = deadline - time.monotonic() - if remaining_s <= 0: - break + if not output_closed and is_done(): + output_closed = True + if settle_output is not None: + await settle_output() + elif poll_output is not None: + await poll_output(deadline) - try: - await asyncio.wait_for(output_notify.wait(), timeout=remaining_s) - except asyncio.TimeoutError: - break - output_notify.clear() + if output_closed: + await _drain_output_chunks(output_chunks, output_lock, output) + else: + await _drain_and_carry_incomplete_suffix(output_chunks, output_lock, output) + except asyncio.CancelledError: + # Queue operations contain no awaits, so restore ownership synchronously. + if output: + output_chunks.appendleft(bytes(output)) + raise text = output.decode("utf-8", errors="replace") truncated, original_token_count = truncate_text_by_tokens(text, max_output_tokens) - return truncated.encode("utf-8", errors="replace"), original_token_count + return truncated.encode("utf-8", errors="replace"), original_token_count, output_closed diff --git a/tests/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index 4916a3a970..4028cbabd3 100644 --- a/tests/extensions/sandbox/test_blaxel.py +++ b/tests/extensions/sandbox/test_blaxel.py @@ -2092,7 +2092,7 @@ async def test_pty_write_stdin_sends_only_nonempty_input( patch.object( session, "_collect_pty_output", - new=AsyncMock(return_value=(b"", None)), + new=AsyncMock(return_value=(b"", None, False)), ), ): update = await session.pty_write_stdin( @@ -2200,6 +2200,7 @@ async def test_pty_finalize_done_session(self, fake_sandbox: _FakeSandboxInstanc entry=entry, output=b"done output", original_token_count=None, + output_closed=True, ) assert result.process_id is None assert result.exit_code == 0 @@ -2565,10 +2566,11 @@ async def test_collect_output_entry_done_immediately( done=True, ) entry.output_chunks.append(b"final output") - output, token_count = await session._collect_pty_output( + output, token_count, output_closed = await session._collect_pty_output( entry=entry, yield_time_ms=100, max_output_tokens=None ) assert b"final output" in output + assert output_closed is True @pytest.mark.asyncio async def test_collect_output_timeout_path(self, fake_sandbox: _FakeSandboxInstance) -> None: @@ -2581,10 +2583,11 @@ async def test_collect_output_timeout_path(self, fake_sandbox: _FakeSandboxInsta http_session=None, ) # Very short yield time, no output, not done. - output, token_count = await session._collect_pty_output( + output, token_count, output_closed = await session._collect_pty_output( entry=entry, yield_time_ms=1, max_output_tokens=None ) assert output == b"" + assert output_closed is False # --------------------------------------------------------------------------- @@ -2969,10 +2972,11 @@ async def test_collect_output_deadline_break(self, fake_sandbox: _FakeSandboxIns entry.output_chunks.append(b"some data") # yield_time_ms=1 means very short deadline, should hit deadline break. - output, _ = await session._collect_pty_output( + output, _, output_closed = await session._collect_pty_output( entry=entry, yield_time_ms=1, max_output_tokens=None ) assert b"some data" in output + assert output_closed is False @pytest.mark.asyncio async def test_collect_output_done_with_remaining_chunks( @@ -2992,11 +2996,12 @@ async def test_collect_output_done_with_remaining_chunks( entry.output_chunks.append(b"chunk1") entry.output_chunks.append(b"chunk2") - output, _ = await session._collect_pty_output( + output, _, output_closed = await session._collect_pty_output( entry=entry, yield_time_ms=5000, max_output_tokens=None ) assert b"chunk1" in output assert b"chunk2" in output + assert output_closed is True # --------------------------------------------------------------------------- diff --git a/tests/extensions/sandbox/test_daytona.py b/tests/extensions/sandbox/test_daytona.py index 7f2df5bb4f..5fa75dd900 100644 --- a/tests/extensions/sandbox/test_daytona.py +++ b/tests/extensions/sandbox/test_daytona.py @@ -1538,8 +1538,83 @@ async def test_session_reader_keeps_entry_live_when_logs_fail_without_exit_code( lambda _chunk: None, ) - assert entry.done is False assert entry.exit_code is None + assert entry.output_closed.is_set() is False + + @pytest.mark.asyncio + async def test_session_reader_closes_entry_when_logs_fail_after_known_exit( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + daytona_module = _load_daytona_module(monkeypatch) + sandbox = _FakeDaytonaSandbox() + sandbox.process.get_session_command_logs_error = RuntimeError("logs failed") + sandbox.process.session_command_exit_code = 7 + state = daytona_module.DaytonaSandboxSessionState( + manifest=Manifest(root=daytona_module.DEFAULT_DAYTONA_WORKSPACE_ROOT), + snapshot=NoopSnapshot(id="snapshot"), + sandbox_id=sandbox.id, + ) + session = daytona_module.DaytonaSandboxSession.from_state(state, sandbox=sandbox) + entry = daytona_module._DaytonaPtySessionEntry( # noqa: SLF001 + daytona_session_id="session-123", + pty_handle=object(), + tty=False, + cmd_id="cmd-123", + ) + + await session._run_session_reader( # noqa: SLF001 + entry, + "session-123", + "cmd-123", + lambda _chunk: None, + ) + + assert entry.exit_code == 7 + assert entry.output_closed.is_set() is True + + @pytest.mark.asyncio + async def test_tty_waiter_closes_output_after_stream_callback_finishes( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + daytona_module = _load_daytona_module(monkeypatch) + sandbox = _FakeDaytonaSandbox() + state = daytona_module.DaytonaSandboxSessionState( + manifest=Manifest(root=daytona_module.DEFAULT_DAYTONA_WORKSPACE_ROOT), + snapshot=NoopSnapshot(id="snapshot"), + sandbox_id=sandbox.id, + ) + session = daytona_module.DaytonaSandboxSession.from_state(state, sandbox=sandbox) + + callback_started = asyncio.Event() + release_callback = asyncio.Event() + + async def append_terminal_tail() -> None: + callback_started.set() + await release_callback.wait() + entry.output_chunks.append(b"tail") + + class _ExitedPtyHandle: + exit_code = 0 + + async def wait(self) -> None: + await append_terminal_tail() + + entry = daytona_module._DaytonaPtySessionEntry( # noqa: SLF001 + daytona_session_id="session-123", + pty_handle=_ExitedPtyHandle(), + ) + waiter_task = asyncio.create_task(session._run_pty_waiter(entry)) # noqa: SLF001 + await callback_started.wait() + + assert entry.output_closed.is_set() is False + + release_callback.set() + await waiter_task + + assert entry.output_closed.is_set() is True + assert list(entry.output_chunks) == [b"tail"] @pytest.mark.asyncio async def test_terminate_pty_entry_awaits_worker_finalizer( diff --git a/tests/extensions/sandbox/test_e2b.py b/tests/extensions/sandbox/test_e2b.py index 67dc301eef..0149ab3d76 100644 --- a/tests/extensions/sandbox/test_e2b.py +++ b/tests/extensions/sandbox/test_e2b.py @@ -30,6 +30,7 @@ E2BSandboxClientOptions, E2BSandboxSession, E2BSandboxSessionState, + _E2BPtyProcessEntry, ) from agents.sandbox import Manifest from agents.sandbox.entries import ( @@ -2102,7 +2103,7 @@ async def test_e2b_pty_start_non_tty_wakes_on_nonzero_wait_exit() -> None: @pytest.mark.asyncio -async def test_e2b_pty_start_non_tty_exited_command_preserves_waiter() -> None: +async def test_e2b_pty_start_non_tty_keeps_session_until_waiter_closes_output() -> None: sandbox = _FakeE2BSandbox() handle = _FakeE2BAsyncCommandHandle(initial_exit_code=0, wait_until_released=True) sandbox.commands.next_async_command_handle = handle @@ -2116,12 +2117,12 @@ async def test_e2b_pty_start_non_tty_exited_command_preserves_waiter() -> None: session = E2BSandboxSession.from_state(state, sandbox=sandbox) started = await asyncio.wait_for( - session.pty_exec_start("true", shell=False, tty=False, yield_time_s=10), + session.pty_exec_start("true", shell=False, tty=False, yield_time_s=0.25), timeout=1, ) - assert started.process_id is None - assert started.exit_code == 0 + assert started.process_id is not None + assert started.exit_code is None assert started.output == b"" assert handle.kill_calls == 0 @@ -2135,6 +2136,90 @@ async def test_e2b_pty_start_non_tty_exited_command_preserves_waiter() -> None: handle.release_wait() await asyncio.sleep(0) + finished = await session.pty_write_stdin( + session_id=started.process_id, + chars="", + yield_time_s=0, + ) + + assert finished.process_id is None + assert finished.exit_code == 0 + assert finished.output == b"" + + +@pytest.mark.asyncio +async def test_e2b_waiter_closes_output_after_pending_callback_append() -> None: + entry = _E2BPtyProcessEntry(handle=_FakeE2BAsyncCommandHandle(), tty=False) + await entry.output_lock.acquire() + + async def append_terminal_tail() -> None: + async with entry.output_lock: + entry.output_chunks.append(b"tail") + + callback_task = asyncio.create_task(append_terminal_tail()) + await asyncio.sleep(0) + session = E2BSandboxSession.from_state( + E2BSandboxSessionState( + session_id=uuid.uuid4(), + manifest=Manifest(root="/workspace"), + snapshot=NoopSnapshot(id="snapshot"), + sandbox_id="sandbox-id", + workspace_root_ready=True, + ), + sandbox=_FakeE2BSandbox(), + ) + waiter_task = asyncio.create_task(session._run_pty_waiter(entry)) # noqa: SLF001 + await asyncio.sleep(0) + + assert entry.output_closed.is_set() is False + + entry.output_lock.release() + await callback_task + await waiter_task + + assert entry.output_closed.is_set() is True + assert list(entry.output_chunks) == [b"tail"] + + +def test_e2b_prune_prefers_settled_output_over_exit_visible_entry() -> None: + from agents.sandbox.session.pty_types import PTY_PROCESSES_MAX + + sandbox = _FakeE2BSandbox() + session = E2BSandboxSession.from_state( + E2BSandboxSessionState( + session_id=uuid.uuid4(), + manifest=Manifest(root="/workspace"), + snapshot=NoopSnapshot(id="snapshot"), + sandbox_id=sandbox.sandbox_id, + workspace_root_ready=True, + ), + sandbox=sandbox, + ) + exit_visible = _E2BPtyProcessEntry( + handle=_FakeE2BAsyncCommandHandle(initial_exit_code=0), + tty=False, + last_used=0, + ) + settled = _E2BPtyProcessEntry( + handle=_FakeE2BAsyncCommandHandle(initial_exit_code=0), + tty=False, + last_used=1, + ) + settled.output_closed.set() + session._pty_processes = {1: exit_visible, 2: settled} # noqa: SLF001 + for process_id in range(3, PTY_PROCESSES_MAX + 1): + session._pty_processes[process_id] = _E2BPtyProcessEntry( # noqa: SLF001 + handle=_FakeE2BAsyncCommandHandle(), + tty=False, + last_used=float(process_id), + ) + session._reserved_pty_process_ids = set(session._pty_processes) # noqa: SLF001 + + removed = session._prune_pty_processes_if_needed() # noqa: SLF001 + + assert removed is settled + assert 1 in session._pty_processes # noqa: SLF001 + assert 2 not in session._pty_processes # noqa: SLF001 @pytest.mark.asyncio @@ -2668,3 +2753,81 @@ class _FakeNotFound(Exception): assert exc_info.value.context["provider_error"] == "The sandbox was not found: request failed" assert exc_info.value.context["reason"] == "_FakeNotFound" assert exc_info.value.retryable is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "wait_error", [RuntimeError("transport unavailable"), asyncio.CancelledError()] +) +async def test_e2b_indeterminate_wait_does_not_close_surviving_output( + wait_error: BaseException, +) -> None: + handle = _FakeE2BAsyncCommandHandle(wait_error=wait_error) + entry = _E2BPtyProcessEntry(handle=handle, tty=False) + session = E2BSandboxSession.from_state( + E2BSandboxSessionState( + session_id=uuid.uuid4(), + manifest=Manifest(root="/workspace"), + snapshot=NoopSnapshot(id="snapshot"), + sandbox_id="sandbox-id", + workspace_root_ready=True, + ), + sandbox=_FakeE2BSandbox(), + ) + session._pty_processes[1] = entry # noqa: SLF001 + if isinstance(wait_error, asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError): + await session._run_pty_waiter(entry) # noqa: SLF001 + else: + await session._run_pty_waiter(entry) # noqa: SLF001 + assert not entry.output_closed.is_set() + entry.output_chunks.append(b"\xc3") + # Inspect the zero-length collection window without the public five-second clamp. + output, _, closed = await session._collect_pty_output( # noqa: SLF001 + entry=entry, + yield_time_ms=0, + max_output_tokens=None, + ) + assert output == b"" + assert closed is False + assert list(entry.output_chunks) == [b"\xc3"] + entry.output_chunks.append(b"\xa9") + handle.wait_error = None + await session._run_pty_waiter(entry) # noqa: SLF001 + final = await session.pty_write_stdin(session_id=1, chars="", yield_time_s=0) + assert final.output == "é".encode() + assert final.exit_code == 0 + assert final.process_id is None + + +@pytest.mark.asyncio +async def test_e2b_exit_visible_polling_allows_terminal_waiter_to_run() -> None: + sandbox = _FakeE2BSandbox() + handle = _FakeE2BAsyncCommandHandle(initial_exit_code=0) + sandbox.commands.next_async_command_handle = handle + session = E2BSandboxSession.from_state( + E2BSandboxSessionState( + session_id=uuid.uuid4(), + manifest=Manifest(root="/workspace"), + snapshot=NoopSnapshot(id="snapshot"), + sandbox_id=sandbox.sandbox_id, + workspace_root_ready=True, + ), + sandbox=sandbox, + ) + try: + update = await session.pty_exec_start("true", shell=False, tty=False, yield_time_s=0.25) + # A sequential caller must not need to insert sleeps to schedule the waiter. + for _ in range(3): + if update.process_id is None: + break + update = await session.pty_write_stdin( + session_id=update.process_id, + chars="", + yield_time_s=0.25, + ) + assert update.process_id is None + assert update.exit_code == 0 + assert handle.wait_calls == 1 + finally: + await session.pty_terminate_all() diff --git a/tests/extensions/sandbox/test_modal.py b/tests/extensions/sandbox/test_modal.py index 44a3fa5c72..a84a79d75e 100644 --- a/tests/extensions/sandbox/test_modal.py +++ b/tests/extensions/sandbox/test_modal.py @@ -4500,6 +4500,328 @@ def _exec(self, *command: object, **kwargs: object) -> object: assert started.output == b"out-1err-1out-2out-3err-2" +@pytest.mark.asyncio +@pytest.mark.parametrize("fallback_read", [False, True]) +async def test_modal_pty_keeps_session_live_until_delayed_exit_tail_reaches_eof( + monkeypatch: pytest.MonkeyPatch, + fallback_read: bool, +) -> None: + modal_module, _create_calls, _registry_tags = _load_modal_module(monkeypatch) + release_tail = asyncio.Event() + + class _DelayedStream: + def __init__(self, *, tail: bytes | None) -> None: + self._tail = tail + self._returned_tail = False + + def __aiter__(self) -> _DelayedStream: + return self + + async def __anext__(self) -> bytes: + if self._tail is not None and not self._returned_tail: + await release_tail.wait() + self._returned_tail = True + return self._tail + raise StopAsyncIteration + + if fallback_read: + + async def read_aio(stream: _DelayedStream, _size: int) -> bytes: + try: + return await stream.__anext__() + except StopAsyncIteration: + return b"" + + def install_read(stream: _DelayedStream) -> None: + stream.read = _with_aio(lambda _size: b"") + stream.read.aio = lambda size: read_aio(stream, size) + + monkeypatch.delattr(_DelayedStream, "__aiter__") + else: + + def install_read(stream: _DelayedStream) -> None: + pass + + class _FakeProcess: + def __init__(self) -> None: + self.stdout = _DelayedStream(tail=b"tail") + self.stderr = _DelayedStream(tail=None) + install_read(self.stdout) + install_read(self.stderr) + self.poll = _with_aio(lambda: 0) + self.terminate = _with_aio(lambda: None) + + class _FakeSandbox: + object_id = "sb-delayed-tail" + + def __init__(self) -> None: + self.process = _FakeProcess() + self.exec = _with_aio(self._exec) + + def _exec(self, *command: object, **kwargs: object) -> object: + _ = (command, kwargs) + return self.process + + sandbox = _FakeSandbox() + state = modal_module.ModalSandboxSessionState( + manifest=Manifest(root="/workspace"), + snapshot=modal_module.resolve_snapshot(None, "snapshot"), + app_name="sandbox-tests", + sandbox_id=sandbox.object_id, + ) + session = modal_module.ModalSandboxSession.from_state(state, sandbox=sandbox) + + started = await session.pty_exec_start("python3", shell=False, tty=True, yield_time_s=0.01) + + assert started.process_id is not None + assert started.exit_code is None + assert started.output == b"" + + release_tail.set() + finished = await session.pty_write_stdin( + session_id=started.process_id, + chars="", + yield_time_s=0.01, + ) + + assert finished.process_id is None + assert finished.exit_code == 0 + assert finished.output == b"tail" + + +@pytest.mark.asyncio +async def test_modal_pty_closes_failed_stream_after_known_exit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + modal_module, _create_calls, _registry_tags = _load_modal_module(monkeypatch) + + class _FailedStream: + def __aiter__(self) -> _FailedStream: + return self + + async def __anext__(self) -> bytes: + raise RuntimeError("stream failed") + + class _EmptyStream: + def __aiter__(self) -> _EmptyStream: + return self + + async def __anext__(self) -> bytes: + raise StopAsyncIteration + + class _FakeProcess: + def __init__(self) -> None: + self.stdout = _FailedStream() + self.stderr = _EmptyStream() + self.poll = _with_aio(lambda: 0) + self.terminate = _with_aio(lambda: None) + + class _FakeSandbox: + object_id = "sb-failed-stream" + + def __init__(self) -> None: + self.process = _FakeProcess() + self.exec = _with_aio(self._exec) + + def _exec(self, *command: object, **kwargs: object) -> object: + _ = (command, kwargs) + return self.process + + sandbox = _FakeSandbox() + state = modal_module.ModalSandboxSessionState( + manifest=Manifest(root="/workspace"), + snapshot=modal_module.resolve_snapshot(None, "snapshot"), + app_name="sandbox-tests", + sandbox_id=sandbox.object_id, + ) + session = modal_module.ModalSandboxSession.from_state(state, sandbox=sandbox) + + finished = await session.pty_exec_start("python3", shell=False, tty=True, yield_time_s=0.01) + + assert finished.process_id is None + assert finished.exit_code == 0 + assert finished.output == b"" + + +@pytest.mark.asyncio +async def test_modal_pty_does_not_poll_status_after_yield_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + modal_module, _create_calls, _registry_tags = _load_modal_module(monkeypatch) + pending = asyncio.Event() + poll_calls = 0 + + class _PendingStream: + def __aiter__(self) -> _PendingStream: + return self + + async def __anext__(self) -> bytes: + await pending.wait() + raise StopAsyncIteration + + class _FakeProcess: + def __init__(self) -> None: + self.stdout = _PendingStream() + self.stderr = _PendingStream() + self.terminate = _with_aio(lambda: None) + + def poll() -> None: + return None + + async def poll_aio() -> None: + nonlocal poll_calls + poll_calls += 1 + await asyncio.sleep(0.2) + + poll.aio = poll_aio # type: ignore[attr-defined] + self.poll = poll + + class _FakeSandbox: + object_id = "sb-slow-poll" + + def __init__(self) -> None: + self.process = _FakeProcess() + self.exec = _with_aio(self._exec) + + def _exec(self, *command: object, **kwargs: object) -> object: + _ = (command, kwargs) + return self.process + + sandbox = _FakeSandbox() + state = modal_module.ModalSandboxSessionState( + manifest=Manifest(root="/workspace"), + snapshot=modal_module.resolve_snapshot(None, "snapshot"), + app_name="sandbox-tests", + sandbox_id=sandbox.object_id, + ) + session = modal_module.ModalSandboxSession.from_state(state, sandbox=sandbox) + + started = await session.pty_exec_start("python3", shell=False, tty=True, yield_time_s=0.25) + + assert started.process_id is not None + assert started.exit_code is None + assert poll_calls == 1 + + await session.pty_terminate_all() + + +@pytest.mark.asyncio +async def test_modal_pty_skips_status_when_fallback_reads_cross_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + modal_module, _create_calls, _registry_tags = _load_modal_module(monkeypatch) + poll_calls = 0 + + class _SlowReadStream: + def __init__(self) -> None: + def read(_size: int) -> bytes: + return b"" + + async def read_aio(_size: int) -> bytes: + await asyncio.sleep(0.2) + return b"" + + read.aio = read_aio # type: ignore[attr-defined] + self.read = read + + class _FakeProcess: + def __init__(self) -> None: + self.stdout = _SlowReadStream() + self.stderr = _SlowReadStream() + self.terminate = _with_aio(lambda: None) + + def poll() -> None: + return None + + async def poll_aio() -> None: + nonlocal poll_calls + poll_calls += 1 + await asyncio.sleep(0.2) + + poll.aio = poll_aio # type: ignore[attr-defined] + self.poll = poll + + class _FakeSandbox: + object_id = "sb-slow-read" + + def __init__(self) -> None: + self.process = _FakeProcess() + self.exec = _with_aio(self._exec) + + def _exec(self, *command: object, **kwargs: object) -> object: + _ = (command, kwargs) + return self.process + + sandbox = _FakeSandbox() + state = modal_module.ModalSandboxSessionState( + manifest=Manifest(root="/workspace"), + snapshot=modal_module.resolve_snapshot(None, "snapshot"), + app_name="sandbox-tests", + sandbox_id=sandbox.object_id, + ) + session = modal_module.ModalSandboxSession.from_state(state, sandbox=sandbox) + + started = await session.pty_exec_start("python3", shell=False, tty=True, yield_time_s=0.25) + + assert started.process_id is not None + assert started.exit_code is None + assert poll_calls == 0 + + await session.pty_terminate_all() + + +@pytest.mark.asyncio +async def test_modal_pty_keeps_pre_exit_empty_fallback_read_live_for_tail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + modal_module, _create_calls, _registry_tags = _load_modal_module(monkeypatch) + + class _FallbackStream: + def __init__(self, chunks: list[bytes]) -> None: + self._chunks = chunks + self.read = _with_aio(self._read) + + def _read(self, _size: int) -> bytes: + return self._chunks.pop(0) if self._chunks else b"" + + class _FakeProcess: + def __init__(self) -> None: + self.stdout = _FallbackStream([b"", b"tail", b""]) + self.stderr = _FallbackStream([b"", b"", b""]) + self._poll_results: list[int | None] = [None, 0] + self.poll = _with_aio(self._poll) + self.terminate = _with_aio(lambda: None) + + def _poll(self) -> int | None: + return self._poll_results.pop(0) if self._poll_results else 0 + + class _FakeSandbox: + object_id = "sb-fallback-tail" + + def __init__(self) -> None: + self.process = _FakeProcess() + self.exec = _with_aio(self._exec) + + def _exec(self, *command: object, **kwargs: object) -> object: + _ = (command, kwargs) + return self.process + + sandbox = _FakeSandbox() + state = modal_module.ModalSandboxSessionState( + manifest=Manifest(root="/workspace"), + snapshot=modal_module.resolve_snapshot(None, "snapshot"), + app_name="sandbox-tests", + sandbox_id=sandbox.object_id, + ) + session = modal_module.ModalSandboxSession.from_state(state, sandbox=sandbox) + + finished = await session.pty_exec_start("python3", shell=False, tty=True, yield_time_s=0.25) + + assert finished.process_id is None + assert finished.exit_code == 0 + assert finished.output == b"tail" + + @pytest.mark.asyncio async def test_modal_pty_start_wraps_startup_failures( monkeypatch: pytest.MonkeyPatch, @@ -4774,3 +5096,122 @@ async def fail_persist() -> io.IOBase: assert exc_info.value.__context__ is None assert source_error.args == () assert source_error.__traceback__ is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel_during_exit_drain", [False, True]) +async def test_modal_pty_cancellation_preserves_consumed_utf8( + monkeypatch: pytest.MonkeyPatch, cancel_during_exit_drain: bool +) -> None: + modal_module, _, _ = _load_modal_module(monkeypatch) + blocked = asyncio.Event() + release = asyncio.Event() + + class Stream: + def __init__(self, chunks: list[bytes], *, block: bool) -> None: + self.chunks = chunks + + async def read_aio(_size: int) -> bytes: + if self.chunks: + return self.chunks.pop(0) + if block: + blocked.set() + await release.wait() + return b"" + + self.read = _with_aio(lambda _size: b"") + self.read.aio = read_aio + + # Provider doubles control cancellation after consumption, before the next read. + stdout = Stream( + [b"", b"\xa9"] if cancel_during_exit_drain else [b"\xa9"], block=cancel_during_exit_drain + ) + stderr = Stream([b""] if cancel_during_exit_drain else [], block=not cancel_during_exit_drain) + process = types.SimpleNamespace( + stdout=stdout, + stderr=stderr, + poll=_with_aio(lambda: 0 if cancel_during_exit_drain else None), + terminate=_with_aio(lambda: None), + ) + session = modal_module.ModalSandboxSession.from_state( + modal_module.ModalSandboxSessionState( + manifest=Manifest(root="/workspace"), + snapshot=modal_module.resolve_snapshot(None, "snapshot"), + app_name="sandbox-tests", + sandbox_id="sb-cancel-output", + ), + sandbox=types.SimpleNamespace(object_id="sb-cancel-output"), + ) + entry = modal_module._ModalPtyProcessEntry(process=process, tty=False) + entry.output_chunks.append(b"\xc3") + session._pty_processes[1] = entry + session._reserved_pty_process_ids.add(1) + task = asyncio.create_task(session.pty_write_stdin(session_id=1, chars="", yield_time_s=1)) + try: + await asyncio.wait_for(blocked.wait(), timeout=1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert b"".join(entry.output_chunks) == "é".encode() + assert session._pty_processes[1] is entry + release.set() + process.poll = _with_aio(lambda: 0) + update = await session.pty_write_stdin(session_id=1, chars="", yield_time_s=0) + assert update.output == "é".encode() + assert update.process_id is None + assert not entry.output_chunks + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await session.pty_terminate_all() + + +@pytest.mark.asyncio +async def test_modal_pty_bounds_each_fallback_read_by_remaining_window( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from agents.sandbox.session import pty_output + + modal_module, _, _ = _load_modal_module(monkeypatch) + now = 0.0 + timeouts: list[float] = [] + clock = types.SimpleNamespace(monotonic=lambda: now) + monkeypatch.setattr(modal_module, "time", clock) + monkeypatch.setattr(pty_output, "time", clock) + process = types.SimpleNamespace( + stdout=types.SimpleNamespace(read=lambda _size: b""), + stderr=types.SimpleNamespace(read=lambda _size: b""), + poll=lambda: pytest.fail("Status must not be polled after the deadline"), + ) + session = modal_module.ModalSandboxSession.from_state( + modal_module.ModalSandboxSessionState( + manifest=Manifest(root="/workspace"), + snapshot=modal_module.resolve_snapshot(None, "snapshot"), + app_name="sandbox-tests", + sandbox_id="sb-deadline", + ), + sandbox=types.SimpleNamespace(object_id="sb-deadline"), + ) + + async def timed_out_read( + fn: object, *args: object, call_timeout: float, **kwargs: object + ) -> None: + nonlocal now + assert fn in (process.stdout.read, process.stderr.read) + timeouts.append(call_timeout) + now += call_timeout + raise asyncio.TimeoutError + + monkeypatch.setattr(session, "_call_modal", timed_out_read) + entry = modal_module._ModalPtyProcessEntry(process=process, tty=False) + # This boundary owns the deadline; empty stdin polls publicly clamp to five seconds. + output, _, closed = await session._collect_pty_output( + entry=entry, + yield_time_ms=250, + max_output_tokens=None, + ) + assert timeouts == pytest.approx([0.2, 0.05]) + assert now == pytest.approx(0.25) + assert output == b"" + assert closed is False diff --git a/tests/sandbox/test_pty_output.py b/tests/sandbox/test_pty_output.py index f15bcf85b4..22ce8653b1 100644 --- a/tests/sandbox/test_pty_output.py +++ b/tests/sandbox/test_pty_output.py @@ -5,7 +5,11 @@ import pytest -from agents.sandbox.session.pty_output import collect_pty_output +from agents.sandbox.session import pty_output as pty_output_module +from agents.sandbox.session.pty_output import ( + _incomplete_utf8_suffix_length, + collect_pty_output, +) @pytest.mark.asyncio @@ -24,7 +28,7 @@ async def produce_output() -> None: output_notify.set() producer_task = asyncio.create_task(produce_output()) - output, original_token_count = await collect_pty_output( + output, original_token_count, output_closed = await collect_pty_output( output_chunks=output_chunks, output_lock=output_lock, output_notify=output_notify, @@ -36,6 +40,7 @@ async def produce_output() -> None: assert output == b"notified output" assert original_token_count is None + assert output_closed is True @pytest.mark.asyncio @@ -46,7 +51,7 @@ def mark_done() -> bool: output_chunks.append(b" after done") return True - output, original_token_count = await collect_pty_output( + output, original_token_count, output_closed = await collect_pty_output( output_chunks=output_chunks, output_lock=asyncio.Lock(), output_notify=asyncio.Event(), @@ -57,3 +62,279 @@ def mark_done() -> bool: assert output == b"before done after done" assert original_token_count is None + assert output_closed is True + + +@pytest.mark.asyncio +async def test_collect_pty_output_drains_chunks_queued_when_wait_times_out() -> None: + output_chunks: deque[bytes] = deque() + + class TimeoutAfterQueueing: + async def wait(self) -> None: + output_chunks.append(b"queued at timeout") + raise asyncio.TimeoutError + + def clear(self) -> None: + pass + + output, original_token_count, output_closed = await collect_pty_output( + output_chunks=output_chunks, + output_lock=asyncio.Lock(), + output_notify=TimeoutAfterQueueing(), # type: ignore[arg-type] + is_done=lambda: False, + yield_time_ms=500, + max_output_tokens=None, + ) + + assert output == b"queued at timeout" + assert original_token_count is None + assert output_closed is False + assert not output_chunks + + +@pytest.mark.parametrize( + ("character", "split"), + [ + pytest.param("é", 1, id="two-byte-1"), + pytest.param("€", 1, id="three-byte-1"), + pytest.param("€", 2, id="three-byte-2"), + pytest.param("😀", 1, id="four-byte-1"), + pytest.param("😀", 2, id="four-byte-2"), + pytest.param("😀", 3, id="four-byte-3"), + ], +) +@pytest.mark.asyncio +async def test_collect_pty_output_preserves_valid_utf8_at_every_split( + character: str, + split: int, +) -> None: + encoded = character.encode("utf-8") + output_chunks: deque[bytes] = deque([b"a" + encoded[:split]]) + output_lock = asyncio.Lock() + output_notify = asyncio.Event() + done = False + + first, _, first_closed = await collect_pty_output( + output_chunks=output_chunks, + output_lock=output_lock, + output_notify=output_notify, + is_done=lambda: done, + yield_time_ms=0, + max_output_tokens=None, + ) + + output_chunks.append(encoded[split:] + b"b") + done = True + second, _, second_closed = await collect_pty_output( + output_chunks=output_chunks, + output_lock=output_lock, + output_notify=output_notify, + is_done=lambda: done, + yield_time_ms=0, + max_output_tokens=None, + ) + + assert first == b"a" + assert first_closed is False + assert first + second == ("a" + character + "b").encode("utf-8") + assert second_closed is True + assert not output_chunks + + +@pytest.mark.parametrize( + "invalid_prefix", + [ + pytest.param(b"\xe0\x80", id="e0-overlong"), + pytest.param(b"\xed\xa0", id="ed-surrogate"), + pytest.param(b"\xf0\x80", id="f0-overlong"), + pytest.param(b"\xf4\x90", id="f4-out-of-range"), + ], +) +@pytest.mark.asyncio +async def test_collect_pty_output_replaces_restricted_utf8_prefixes_without_carry( + invalid_prefix: bytes, +) -> None: + output_chunks: deque[bytes] = deque([b"prompt" + invalid_prefix]) + + output, _, output_closed = await collect_pty_output( + output_chunks=output_chunks, + output_lock=asyncio.Lock(), + output_notify=asyncio.Event(), + is_done=lambda: False, + yield_time_ms=0, + max_output_tokens=None, + ) + + assert output.decode("utf-8") == "prompt��" + assert output_closed is False + assert not output_chunks + + +@pytest.mark.parametrize( + ("data", "expected"), + [ + pytest.param(b"", 0, id="empty"), + pytest.param(b"abc", 0, id="ascii"), + pytest.param(b"\xc3", 1, id="two-byte-lead"), + pytest.param(b"\xe2\x82", 2, id="three-byte-prefix"), + pytest.param(b"\xf0\x9f\x98", 3, id="four-byte-prefix"), + pytest.param(b"\xe0\x80", 0, id="e0-restricted"), + pytest.param(b"\xe0\xa0", 2, id="e0-valid"), + pytest.param(b"\xed\xa0", 0, id="ed-restricted"), + pytest.param(b"\xed\x9f", 2, id="ed-valid"), + pytest.param(b"\xf0\x80", 0, id="f0-restricted"), + pytest.param(b"\xf0\x90", 2, id="f0-valid"), + pytest.param(b"\xf4\x90", 0, id="f4-restricted"), + pytest.param(b"\xf4\x8f", 2, id="f4-valid"), + pytest.param(b"\x80\x80\x80", 0, id="orphan-continuations"), + ], +) +def test_incomplete_utf8_suffix_length_accepts_only_completable_sequences( + data: bytes, + expected: int, +) -> None: + assert _incomplete_utf8_suffix_length(data) == expected + + +@pytest.mark.asyncio +async def test_collect_pty_output_checks_deadline_before_next_poll( + monkeypatch: pytest.MonkeyPatch, +) -> None: + now = 0.0 + poll_count = 0 + settle_count = 0 + + def monotonic() -> float: + return now + + async def poll_output(_deadline: float) -> None: + nonlocal now, poll_count + poll_count += 1 + now = 0.1 + + async def settle_output() -> None: + nonlocal settle_count + settle_count += 1 + + async def wait_for_output(_remaining_s: float) -> None: + nonlocal now + now = 0.3 + + monkeypatch.setattr(pty_output_module.time, "monotonic", monotonic) + + output, _, output_closed = await collect_pty_output( + output_chunks=deque(), + output_lock=asyncio.Lock(), + output_notify=asyncio.Event(), + is_done=lambda: False, + yield_time_ms=250, + max_output_tokens=None, + poll_output=poll_output, + settle_output=settle_output, + wait_for_output=wait_for_output, + ) + + assert output == b"" + assert output_closed is False + assert poll_count == 1 + assert settle_count == 1 + + +@pytest.mark.asyncio +async def test_collect_pty_output_settles_terminal_carry_once_across_repeated_reads() -> None: + output_chunks: deque[bytes] = deque([b"tail\xe2\x82"]) + output_lock = asyncio.Lock() + output_notify = asyncio.Event() + done = False + + first, _, first_closed = await collect_pty_output( + output_chunks=output_chunks, + output_lock=output_lock, + output_notify=output_notify, + is_done=lambda: done, + yield_time_ms=0, + max_output_tokens=None, + ) + done = True + terminal, _, terminal_closed = await collect_pty_output( + output_chunks=output_chunks, + output_lock=output_lock, + output_notify=output_notify, + is_done=lambda: done, + yield_time_ms=0, + max_output_tokens=None, + ) + repeated, _, repeated_closed = await collect_pty_output( + output_chunks=output_chunks, + output_lock=output_lock, + output_notify=output_notify, + is_done=lambda: done, + yield_time_ms=0, + max_output_tokens=None, + ) + + assert first == b"tail" + assert first_closed is False + assert terminal.decode("utf-8") == "�" + assert terminal_closed is True + assert repeated == b"" + assert repeated_closed is True + assert not output_chunks + + +@pytest.mark.asyncio +async def test_collect_pty_output_restores_carry_when_next_collection_is_cancelled() -> None: + output_chunks: deque[bytes] = deque([b"\xc3"]) + output_lock = asyncio.Lock() + output_notify = asyncio.Event() + done = False + + first, _, first_closed = await collect_pty_output( + output_chunks=output_chunks, + output_lock=output_lock, + output_notify=output_notify, + is_done=lambda: done, + yield_time_ms=0, + max_output_tokens=None, + ) + + wait_started = asyncio.Event() + + async def wait_for_output(_remaining_s: float) -> None: + wait_started.set() + await asyncio.Event().wait() + + cancelled = asyncio.create_task( + collect_pty_output( + output_chunks=output_chunks, + output_lock=output_lock, + output_notify=output_notify, + is_done=lambda: done, + yield_time_ms=1_000, + max_output_tokens=None, + wait_for_output=wait_for_output, + ) + ) + await wait_started.wait() + cancelled.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled + + assert first == b"" + assert first_closed is False + assert list(output_chunks) == [b"\xc3"] + + output_chunks.append(b"\xa9") + done = True + terminal, _, terminal_closed = await collect_pty_output( + output_chunks=output_chunks, + output_lock=output_lock, + output_notify=output_notify, + is_done=lambda: done, + yield_time_ms=0, + max_output_tokens=None, + ) + + assert terminal == "é".encode() + assert terminal_closed is True + assert not output_chunks diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index f2482be90f..9188b1fc33 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -279,6 +279,75 @@ async def blocked_to_thread(*args: object, **kwargs: object) -> None: assert session._fd_close_tasks == set() + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("prefix", "tail", "first_output", "final_output"), + [ + (b"before close", b" terminal", b"before close", b" terminal"), + (b"\xc3", b"\xa9", b"", "é".encode()), + ], + ) + async def test_pty_exit_waits_for_output_close_before_terminal_cleanup( + self, + tmp_path: Path, + prefix: bytes, + tail: bytes, + first_output: bytes, + final_output: bytes, + ) -> None: + session = _RecordingUnixLocalSession(tmp_path) + process = cast( + asyncio.subprocess.Process, + SimpleNamespace(returncode=0, pid=None), + ) + entry = _UnixPtyProcessEntry(process=process, tty=False) + process_id = 1234 + session._pty_processes[process_id] = entry + session._reserved_pty_process_ids.add(process_id) + + entry.output_chunks.append(prefix) + output, token_count, output_closed = await session._collect_pty_output( + entry=entry, + yield_time_ms=0, + max_output_tokens=None, + ) + # The producer can close and queue a terminal tail after collection returns but + # before finalization observes the entry. Removal must follow the collector's + # settled result, not a later read of the mutable close event. + entry.output_chunks.append(tail) + entry.output_closed.set() + still_live = await session._finalize_pty_update( + process_id=process_id, + entry=entry, + output=output, + original_token_count=token_count, + output_closed=output_closed, + ) + + assert still_live.process_id == process_id + assert still_live.exit_code is None + assert still_live.output == first_output + assert process_id in session._pty_processes + + terminal_output, terminal_token_count, terminal_closed = await session._collect_pty_output( + entry=entry, + yield_time_ms=0, + max_output_tokens=None, + ) + terminal = await session._finalize_pty_update( + process_id=process_id, + entry=entry, + output=terminal_output, + original_token_count=terminal_token_count, + output_closed=terminal_closed, + ) + + assert terminal.process_id is None + assert terminal.exit_code == 0 + assert terminal.output == final_output + assert process_id not in session._pty_processes + assert process_id not in session._reserved_pty_process_ids + @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: