diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index aee3211fbd..43154f3d14 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -856,7 +856,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, @@ -866,6 +866,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( @@ -890,7 +891,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 == "" @@ -903,6 +904,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: @@ -964,7 +966,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, @@ -981,11 +983,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..c89cfce5bc 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -387,6 +387,7 @@ 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 @@ -755,7 +756,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 +766,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: @@ -777,6 +779,10 @@ async def _run_pty_waiter(self, entry: _DaytonaPtySessionEntry) -> None: 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( @@ -804,8 +810,12 @@ async def _run_session_reader( entry.done = True except Exception: pass - if not logs_failed: + # 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.done = True + entry.output_closed.set() entry.output_notify.set() async def pty_write_stdin( @@ -832,7 +842,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 +855,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 +865,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 +899,13 @@ 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, + should_return=lambda: entry.done, yield_time_ms=yield_time_ms, max_output_tokens=max_output_tokens, ) @@ -901,7 +914,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..882ccf2ca3 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,16 @@ 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, + should_return=lambda: self._entry_exit_code(entry) is not None, + yield_time_ms=yield_time_ms, + max_output_tokens=max_output_tokens, + ) async def _run_pty_waiter(self, entry: _E2BPtyProcessEntry) -> None: try: @@ -1259,6 +1241,12 @@ async def _run_pty_waiter(self, entry: _E2BPtyProcessEntry) -> None: except (TypeError, ValueError): pass finally: + # 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 +1256,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 +1281,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..223dcdd67a 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,38 +990,69 @@ 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") + ) -> tuple[bytes, int | None, bool]: + async def poll_output( + *, + deadline: float | None, + allow_new_read: bool, + check_exit: bool, + ) -> None: + chunks: list[bytes] = [] + stdout_chunk = await self._read_modal_stream( + entry=entry, + stream_name="stdout", + allow_new_read=allow_new_read, + ) + stderr_chunk = await self._read_modal_stream( + entry=entry, + stream_name="stderr", + allow_new_read=allow_new_read, + ) if stdout_chunk: - chunks.extend(stdout_chunk) + chunks.append(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 + chunks.append(stderr_chunk) + + if ( + check_exit + and deadline is not None + and time.monotonic() < deadline + and await self._peek_exit_code(entry.process) is not None + ): + chunks.append(await self._drain_modal_stream(entry=entry, stream_name="stdout")) + chunks.append(await self._drain_modal_stream(entry=entry, stream_name="stderr")) + if entry.stdout_closed and entry.stderr_closed: + entry.output_closed.set() + + if chunks: + async with entry.output_lock: + entry.output_chunks.extend(chunk for chunk in chunks if chunk) + entry.output_notify.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, @@ -1026,6 +1066,7 @@ async def _drain_modal_stream( entry=entry, stream_name=stream_name, await_pending=True, + settle_after_exit=True, ) if not chunk: break @@ -1038,13 +1079,23 @@ 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, ) -> 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) + if task is None and not allow_new_read: + return b"" + stream_iter = getattr(entry, iter_attr) if stream_iter is None: aiter_method = getattr(stream, "__aiter__", None) @@ -1056,7 +1107,6 @@ 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) @@ -1072,9 +1122,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) @@ -1086,11 +1139,18 @@ async def _read_modal_stream( try: value = await self._call_modal(read, 16_384, call_timeout=0.2) 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 +1170,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 +1195,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 ea7f83e9d8..aeae5ebc39 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..2f5ba46c36 100644 --- a/src/agents/sandbox/session/pty_output.py +++ b/src/agents/sandbox/session/pty_output.py @@ -3,48 +3,179 @@ 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 _restore_unreturned_output( + output_chunks: deque[bytes], + output_lock: asyncio.Lock, + output: bytearray, +) -> None: + """Restore bytes owned by a cancelled collection ahead of later queued bytes.""" + if not output: + return + + async with output_lock: + output_chunks.appendleft(bytes(output)) + + async def collect_pty_output( *, output_chunks: deque[bytes], output_lock: asyncio.Lock, output_notify: asyncio.Event, is_done: Callable[[], bool], + should_return: Callable[[], bool] | None = None, 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 + + if should_return is not None and should_return(): + 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: + restore_task = asyncio.create_task( + _restore_unreturned_output(output_chunks, output_lock, output) + ) + while not restore_task.done(): + try: + await asyncio.shield(restore_task) + except asyncio.CancelledError: + continue + restore_task.result() + 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 3fe1d0d93a..d2e87737e8 100644 --- a/tests/extensions/sandbox/test_blaxel.py +++ b/tests/extensions/sandbox/test_blaxel.py @@ -1940,7 +1940,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( @@ -2048,6 +2048,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 @@ -2413,10 +2414,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: @@ -2429,10 +2431,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 # --------------------------------------------------------------------------- @@ -2817,10 +2820,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( @@ -2840,11 +2844,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..2bb60fda02 100644 --- a/tests/extensions/sandbox/test_daytona.py +++ b/tests/extensions/sandbox/test_daytona.py @@ -1540,6 +1540,85 @@ async def test_session_reader_keeps_entry_live_when_logs_fail_without_exit_code( 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.done is True + 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.done is False + assert entry.output_closed.is_set() is False + + release_callback.set() + await waiter_task + + assert entry.done is True + 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..342a390055 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 @@ -2120,8 +2121,8 @@ async def test_e2b_pty_start_non_tty_exited_command_preserves_waiter() -> None: 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 diff --git a/tests/extensions/sandbox/test_modal.py b/tests/extensions/sandbox/test_modal.py index 44a3fa5c72..cc7952289b 100644 --- a/tests/extensions/sandbox/test_modal.py +++ b/tests/extensions/sandbox/test_modal.py @@ -4500,6 +4500,306 @@ def _exec(self, *command: object, **kwargs: object) -> object: assert started.output == b"out-1err-1out-2out-3err-2" +@pytest.mark.asyncio +async def test_modal_pty_keeps_session_live_until_delayed_exit_tail_reaches_eof( + monkeypatch: pytest.MonkeyPatch, +) -> 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 + + class _FakeProcess: + def __init__(self) -> None: + self.stdout = _DelayedStream(tail=b"tail") + self.stderr = _DelayedStream(tail=None) + 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, diff --git a/tests/sandbox/test_pty_output.py b/tests/sandbox/test_pty_output.py index f15bcf85b4..4498f191e2 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,290 @@ 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() + await output_lock.acquire() + cancelled.cancel() + await asyncio.sleep(0) + cancelled.cancel() + await asyncio.sleep(0) + cancelled.cancel() + await asyncio.sleep(0) + + assert cancelled.done() is False + assert not output_chunks + + output_lock.release() + 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 fb13be3c51..4491c3522f 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -251,6 +251,64 @@ async def blocked_to_thread(*args: object, **kwargs: object) -> None: assert session._fd_close_tasks == set() + @pytest.mark.asyncio + async def test_pty_exit_waits_for_output_close_before_terminal_cleanup( + self, + tmp_path: Path, + ) -> 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(b"before close") + 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(b" terminal") + 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 == b"before close" + 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 == b" terminal" + 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: