diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index aee3211fbd..81fe388a63 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -46,7 +46,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_output import collect_pty_output, flush_pty_tail from ....sandbox.session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -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, source_text = 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,8 @@ async def pty_exec_start( entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) async def pty_write_stdin( @@ -890,7 +892,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, source_text = 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 +905,8 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) async def pty_terminate_all(self) -> None: @@ -964,7 +968,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, str]: return await collect_pty_output( output_chunks=entry.output_chunks, output_lock=entry.output_lock, @@ -981,6 +985,8 @@ async def _finalize_pty_update( entry: _BlaxelPtySessionEntry, output: bytes, original_token_count: int | None, + source_text: str = "", + max_output_tokens: int | None = None, ) -> PtyExecUpdate: exit_code = entry.exit_code if entry.done else None live_process_id: int | None = process_id @@ -989,8 +995,23 @@ async def _finalize_pty_update( async with self._pty_lock: removed = self._pty_sessions.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) - if removed is not None: - await self._terminate_pty_entry(removed) + # Draining is destructive and the tail lives on the entry, so the removal has to + # commit first. Cancelled the other way round, the session stays registered with + # its last bytes already gone and a later call cannot get them back. Once it is out + # of the map nothing else can reach it either, pty_terminate_all included, so its + # sockets and sessions have to be closed whatever happens to the drain. + try: + output, original_token_count = await flush_pty_tail( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output=output, + source_text=source_text, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) + finally: + if removed is not None: + await self._terminate_pty_entry(removed) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index bb8d7c37e6..1c9bc23133 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, flush_pty_tail 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, str]: + 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,6 +1050,8 @@ async def _finalize_pty_update( entry: _CloudflarePtyProcessEntry, output: bytes, original_token_count: int | None, + source_text: str = "", + max_output_tokens: int | None = None, ) -> PtyExecUpdate: exit_code = entry.exit_code if entry.output_closed.is_set() else None live_process_id: int | None = process_id @@ -1076,8 +1059,23 @@ async def _finalize_pty_update( async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) - if removed is not None: - await self._terminate_pty_entry(removed) + # Draining is destructive and the tail lives on the entry, so the removal has to + # commit first. Cancelled the other way round, the session stays registered with + # its last bytes already gone and a later call cannot get them back. Once it is out + # of the map nothing else can reach it either, pty_terminate_all included, so its + # sockets and sessions have to be closed whatever happens to the drain. + try: + output, original_token_count = await flush_pty_tail( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output=output, + source_text=source_text, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) + finally: + if removed is not None: + await self._terminate_pty_entry(removed) live_process_id = None return PtyExecUpdate( @@ -1220,7 +1218,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, source_text = 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 +1228,8 @@ async def pty_exec_start( entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) async def pty_write_stdin( @@ -1253,7 +1253,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, source_text = await self._collect_pty_output( entry=entry, yield_time_ms=resolve_pty_write_yield_time_ms( yield_time_ms=yield_time_ms, @@ -1267,6 +1267,8 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) 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..717a14d508 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -45,7 +45,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_output import collect_pty_output, flush_pty_tail from ....sandbox.session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -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, source_text = 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,8 @@ async def _on_data(chunk: bytes | str) -> None: entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) async def _run_pty_waiter(self, entry: _DaytonaPtySessionEntry) -> None: @@ -832,7 +834,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, source_text = 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 +847,8 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) async def _finalize_pty_update( @@ -854,6 +858,8 @@ async def _finalize_pty_update( entry: _DaytonaPtySessionEntry, output: bytes, original_token_count: int | None, + source_text: str = "", + max_output_tokens: int | None = None, ) -> PtyExecUpdate: exit_code = entry.exit_code if entry.done else None live_process_id: int | None = process_id @@ -862,8 +868,23 @@ async def _finalize_pty_update( async with self._pty_lock: removed = self._pty_sessions.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) - if removed is not None: - await self._terminate_pty_entry(removed) + # Draining is destructive and the tail lives on the entry, so the removal has to + # commit first. Cancelled the other way round, the session stays registered with + # its last bytes already gone and a later call cannot get them back. Once it is out + # of the map nothing else can reach it either, pty_terminate_all included, so its + # sockets and sessions have to be closed whatever happens to the drain. + try: + output, original_token_count = await flush_pty_tail( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output=output, + source_text=source_text, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) + finally: + if removed is not None: + await self._terminate_pty_entry(removed) live_process_id = None return PtyExecUpdate( @@ -887,7 +908,7 @@ async def _collect_pty_output( entry: _DaytonaPtySessionEntry, yield_time_ms: int, max_output_tokens: int | None, - ) -> tuple[bytes, int | None]: + ) -> tuple[bytes, int | None, str]: return await collect_pty_output( output_chunks=entry.output_chunks, output_lock=entry.output_lock, diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index 389b665c44..c6f95175e6 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, flush_pty_tail 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 @@ -1050,7 +1050,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, source_text = 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 +1060,8 @@ async def _append_output(payload: bytes | bytearray | str | object) -> None: entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) 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, source_text = 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,8 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) async def pty_terminate_all(self) -> None: @@ -1211,37 +1215,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, str]: + return await collect_pty_output( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output_notify=entry.output_notify, + is_done=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: @@ -1268,6 +1250,8 @@ async def _finalize_pty_update( entry: _E2BPtyProcessEntry, output: bytes, original_token_count: int | None, + source_text: str = "", + max_output_tokens: int | None = None, ) -> PtyExecUpdate: exit_code = self._entry_exit_code(entry) live_process_id: int | None = process_id @@ -1276,8 +1260,23 @@ async def _finalize_pty_update( async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) - if removed is not None: - await self._terminate_pty_entry(removed) + # Draining is destructive and the tail lives on the entry, so the removal has to + # commit first. Cancelled the other way round, the session stays registered with + # its last bytes already gone and a later call cannot get them back. Once it is out + # of the map nothing else can reach it either, pty_terminate_all included, so its + # sockets and sessions have to be closed whatever happens to the drain. + try: + output, original_token_count = await flush_pty_tail( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output=output, + source_text=source_text, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) + finally: + if removed is not None: + await self._terminate_pty_entry(removed) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index 848b8e12a3..9f630c7405 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -64,6 +64,7 @@ _settle_mount_transition, _terminate_ambiguous_mount_session, ) +from ....sandbox.session.pty_output import close_pty_tail, decode_pty_window from ....sandbox.session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -489,6 +490,7 @@ class _ModalPtyProcessEntry: stderr_iter: AsyncIterator[object] | None = None stdout_read_task: asyncio.Task[object] | None = None stderr_read_task: asyncio.Task[object] | None = None + pending_output: bytes = b"" class ModalSandboxSession(BaseSandboxSession): @@ -907,7 +909,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, source_text = 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 +919,8 @@ async def pty_exec_start( entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) async def pty_write_stdin( @@ -940,7 +944,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, source_text = 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 +957,8 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) async def pty_terminate_all(self) -> None: @@ -981,38 +987,55 @@ async def _collect_pty_output( entry: _ModalPtyProcessEntry, yield_time_ms: int, max_output_tokens: int | None, - ) -> tuple[bytes, int | None]: + ) -> tuple[bytes, int | None, str]: 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 + # a character split across two windows starts in the tail the last one held back. the + # field is left alone until the decode below commits, so a cancelled call keeps it + chunks = bytearray(entry.pending_output) - 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 + try: + 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 - if not stdout_chunk and not stderr_chunk: - remaining_s = deadline - time.monotonic() - if remaining_s <= 0: + 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 - await asyncio.sleep(min(_PTY_POLL_INTERVAL_S, remaining_s)) - text = chunks.decode("utf-8", errors="replace") + 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)) + + exited = await self._peek_exit_code(entry.process) is not None + except asyncio.CancelledError: + # A stream item is gone from the stream once it has been read, so anything already + # taken lives only in this buffer. There is no deque to hand it back to, and the + # session outlives a cancelled call, so it goes on the entry for the next window. + # Left behind, the carried lead byte would pair with whatever arrived after the + # continuation this call swallowed. + if chunks: + entry.pending_output = bytes(chunks) + raise + text, entry.pending_output = decode_pty_window(chunks, is_final=exited) truncated_text, original_token_count = truncate_text_by_tokens(text, max_output_tokens) - return truncated_text.encode("utf-8", errors="replace"), original_token_count + return truncated_text.encode("utf-8", errors="replace"), original_token_count, text async def _drain_modal_stream( self, @@ -1110,6 +1133,8 @@ async def _finalize_pty_update( entry: _ModalPtyProcessEntry, output: bytes, original_token_count: int | None, + source_text: str = "", + max_output_tokens: int | None = None, ) -> PtyExecUpdate: exit_code = await self._peek_exit_code(entry.process) live_process_id: int | None = process_id @@ -1117,8 +1142,25 @@ async def _finalize_pty_update( async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) - if removed is not None: - await self._terminate_pty_entry(removed) + + # Reading the tail clears it, so the removal has to commit first. Cancelled the + # other way round, the session stays registered with its last bytes already gone. + # The entry is out of the map now, so nothing else can reach it and + # pty_terminate_all cannot clean it up later. Closing the tail happens to be + # synchronous here, but its cleanup is settled the same way as the others so an + # await added in front of it later cannot start leaking sessions. + try: + output, original_token_count = close_pty_tail( + leftover=entry.pending_output, + output=output, + source_text=source_text, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) + entry.pending_output = b"" + finally: + if removed is not None: + await self._terminate_pty_entry(removed) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index 8ca4febe85..79dff8cf6e 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -55,7 +55,7 @@ from ..session.base_sandbox_session import BaseSandboxSession from ..session.dependencies import Dependencies from ..session.manager import Instrumentation -from ..session.pty_output import collect_pty_output +from ..session.pty_output import collect_pty_output, flush_pty_tail from ..session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -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, source_text = 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,8 @@ async def pty_exec_start( entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) async def pty_write_stdin( @@ -1126,7 +1128,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, source_text = 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 +1141,8 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) async def pty_terminate_all(self) -> None: @@ -1242,7 +1246,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, str]: return await collect_pty_output( output_chunks=entry.output_chunks, output_lock=entry.output_lock, @@ -1259,19 +1263,39 @@ async def _finalize_pty_update( entry: _DockerPtyProcessEntry, output: bytes, original_token_count: int | None, + source_text: str = "", + max_output_tokens: int | None = None, ) -> PtyExecUpdate: if entry.output_closed.is_set() and entry.exit_code is None: await self._refresh_pty_exit_code(entry) - exit_code = entry.exit_code + # _watch_pty_exit can set exit_code before _pump_pty_socket reaches its finally, so + # finalizing on exit_code alone would remove the session while the pump still holds + # output. Collection already waits for output_closed, so this matches it. + exit_code = entry.exit_code if entry.output_closed.is_set() else None live_process_id: int | None = process_id if exit_code is not None: async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) - if removed is not None: - await self._terminate_pty_entry(removed) + # Draining is destructive and the tail lives on the entry, so the removal has to + # commit first. Cancelled the other way round, the session stays registered with + # its last bytes already gone and a later call cannot get them back. Once it is out + # of the map nothing else can reach it either, pty_terminate_all included, so its + # sockets and sessions have to be closed whatever happens to the drain. + try: + output, original_token_count = await flush_pty_tail( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output=output, + source_text=source_text, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) + finally: + if removed is not None: + await self._terminate_pty_entry(removed) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index ea7f83e9d8..34af50d879 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -48,7 +48,7 @@ from ..session.base_sandbox_session import BaseSandboxSession from ..session.dependencies import Dependencies from ..session.manager import Instrumentation -from ..session.pty_output import collect_pty_output +from ..session.pty_output import collect_pty_output, flush_pty_tail from ..session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -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, source_text = 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,8 @@ def _preexec() -> None: entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) async def pty_write_stdin( @@ -442,7 +444,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, source_text = 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 +457,8 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, + max_output_tokens=max_output_tokens, ) async def pty_terminate_all(self) -> None: @@ -533,7 +537,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, str]: return await collect_pty_output( output_chunks=entry.output_chunks, output_lock=entry.output_lock, @@ -550,16 +554,37 @@ async def _finalize_pty_update( entry: _UnixPtyProcessEntry, output: bytes, original_token_count: int | None, + source_text: str = "", + max_output_tokens: int | None = None, ) -> PtyExecUpdate: - exit_code: int | None = entry.process.returncode + # Collection treats the session as finished on output_closed, which is set only after + # the process is reaped and every pump task has drained. Finalizing on returncode + # alone removes the session, and terminating it cancels a pump that still holds the + # rest of a character, so the two have to agree on what finished means. + exit_code: int | None = entry.process.returncode if entry.output_closed.is_set() else None live_process_id: int | None = process_id if exit_code is not None: async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) - if removed is not None: - await self._terminate_pty_entry(removed) + # Draining is destructive and the tail lives on the entry, so the removal has to + # commit first. Cancelled the other way round, the session stays registered with + # its last bytes already gone and a later call cannot get them back. Once it is out + # of the map nothing else can reach it either, pty_terminate_all included, so its + # sockets and sessions have to be closed whatever happens to the drain. + try: + output, original_token_count = await flush_pty_tail( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output=output, + source_text=source_text, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) + finally: + if removed is not None: + await self._terminate_pty_entry(removed) live_process_id = None return PtyExecUpdate( diff --git a/src/agents/sandbox/session/pty_output.py b/src/agents/sandbox/session/pty_output.py index 25cbe774e7..40433762d9 100644 --- a/src/agents/sandbox/session/pty_output.py +++ b/src/agents/sandbox/session/pty_output.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import codecs import time from collections import deque from collections.abc import Callable @@ -8,6 +9,84 @@ from .pty_types import truncate_text_by_tokens +def decode_pty_window(data: bytes | bytearray, *, is_final: bool) -> tuple[str, bytes]: + """Decode one collection window, and return what has to wait for the next one. + + PTY output arrives in repeated windows, so decoding each one with ``errors="replace"`` + destroys any character whose bytes straddle a boundary. The returned bytes are the tail + the caller has to put back in front of the next window. + """ + decoder = codecs.getincrementaldecoder("utf-8")("replace") + text = decoder.decode(data, final=is_final) + pending = bytes(decoder.getstate()[0]) + + # The decoder holds ED A0..BF, which leads a surrogate, even though no third byte can + # complete it. Those 32 prefixes are the only thing it ever buffers that cannot become a + # character, so handing them back would keep the output hidden for as long as the process + # runs. Replace them here instead, the same way the decoder would once it is closed. + if len(pending) == 2 and pending[0] == 0xED and pending[1] >= 0xA0: + text += pending.decode("utf-8", errors="replace") + pending = b"" + + return text, pending + + +def close_pty_tail( + *, + leftover: bytes | bytearray, + output: bytes, + source_text: str, + original_token_count: int | None, + max_output_tokens: int | None, +) -> tuple[bytes, int | None]: + """Fold a tail a collection window left behind into the output of a finished session. + + A window hands an unfinished character back while the stream still looks open, but each + backend decides separately that the process has gone, and it can decide that after the last + collection. Whatever is still waiting then has no later window to complete it, so it is + replaced here rather than leaving with the session. + + ``source_text`` is what the window decoded before it applied ``max_output_tokens``. The tail + is folded into that and truncated once, rather than appended to a rendered result. Truncation + keeps the start and the end of the source, so a tail belongs in the part that is kept: adding + it to the display instead would truncate twice, recount the shortened text rather than the + source, and leave a stale ending that hides whatever the process said last. + """ + if not leftover: + return output, original_token_count + + tail, _ = decode_pty_window(leftover, is_final=True) + if not tail: + return output, original_token_count + + truncated, counted = truncate_text_by_tokens(source_text + tail, max_output_tokens) + return truncated.encode("utf-8", errors="replace"), counted + + +async def flush_pty_tail( + *, + output_chunks: deque[bytes], + output_lock: asyncio.Lock, + output: bytes, + source_text: str, + original_token_count: int | None, + max_output_tokens: int | None, +) -> tuple[bytes, int | None]: + """Drain what a session still holds and close it with :func:`close_pty_tail`.""" + leftover = bytearray() + async with output_lock: + while output_chunks: + leftover.extend(output_chunks.popleft()) + + return close_pty_tail( + leftover=leftover, + output=output, + source_text=source_text, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) + + async def collect_pty_output( *, output_chunks: deque[bytes], @@ -16,35 +95,63 @@ 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.""" +) -> tuple[bytes, int | None, str]: + """Collect and truncate PTY output until the deadline or provider completion. + + Also returns the decoded window before truncation, so that a backend which later finds the + session finished can fold a remaining tail into the real source rather than into the + rendered result. It is internal, no public field carries it. + """ deadline = time.monotonic() + (yield_time_ms / 1000) output = bytearray() - while True: - async with output_lock: - while output_chunks: - output.extend(output_chunks.popleft()) - - if time.monotonic() >= deadline: - break - - if is_done(): + try: + while True: async with output_lock: while output_chunks: output.extend(output_chunks.popleft()) - break - remaining_s = deadline - time.monotonic() - if remaining_s <= 0: - break + if time.monotonic() >= deadline: + break + + if is_done(): + async with output_lock: + while output_chunks: + output.extend(output_chunks.popleft()) + break + + remaining_s = deadline - time.monotonic() + if remaining_s <= 0: + break + + try: + await asyncio.wait_for(output_notify.wait(), timeout=remaining_s) + except asyncio.TimeoutError: + break + output_notify.clear() + except asyncio.CancelledError: + # Everything drained so far lives only in this buffer, and the session outlives a + # cancelled call, so put it back before the cancellation goes on. Otherwise the next + # window reads a continuation whose lead byte went with the abandoned call and reports + # a replacement character for output that did arrive. + if output: + output_chunks.appendleft(bytes(output)) + raise - try: - await asyncio.wait_for(output_notify.wait(), timeout=remaining_s) - except asyncio.TimeoutError: - break - output_notify.clear() + # Output is collected in repeated windows over one persistent deque, so a character + # whose bytes straddle a window boundary used to be replaced twice and lost. An + # incremental decoder keeps that trailing partial sequence instead of replacing it, + # and it is handed back for the next window to finish. Bytes that cannot begin a + # character are not held, they are replaced straight away as before. Completing the + # decoder once the provider is done replaces a tail that no later window will finish. + text, pending = decode_pty_window(output, is_final=is_done()) + if pending: + # Deliberately not under ``output_lock``. Those bytes have already left the deque, so + # this is the only copy, and awaiting the lock is a cancellation point: a caller + # cancelled here would drop the character while the session lives on to read its + # continuation. ``appendleft`` is one synchronous call on a deque, so nothing can + # interleave with the drain loops the lock exists to protect. + output_chunks.appendleft(pending) - 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, text diff --git a/tests/extensions/sandbox/test_blaxel.py b/tests/extensions/sandbox/test_blaxel.py index 3fe1d0d93a..c1ab7d2268 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, "")), ), ): update = await session.pty_write_stdin( @@ -2413,7 +2413,7 @@ 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, _ = await session._collect_pty_output( entry=entry, yield_time_ms=100, max_output_tokens=None ) assert b"final output" in output @@ -2429,7 +2429,7 @@ 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, _ = await session._collect_pty_output( entry=entry, yield_time_ms=1, max_output_tokens=None ) assert output == b"" @@ -2817,7 +2817,7 @@ 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, _, _ = await session._collect_pty_output( entry=entry, yield_time_ms=1, max_output_tokens=None ) assert b"some data" in output @@ -2840,7 +2840,7 @@ 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, _, _ = await session._collect_pty_output( entry=entry, yield_time_ms=5000, max_output_tokens=None ) assert b"chunk1" in output diff --git a/tests/extensions/sandbox/test_cloudflare.py b/tests/extensions/sandbox/test_cloudflare.py index bcd47257e1..a8464dd43e 100644 --- a/tests/extensions/sandbox/test_cloudflare.py +++ b/tests/extensions/sandbox/test_cloudflare.py @@ -1730,6 +1730,33 @@ async def test_cloudflare_pty_exec_start_opens_websocket_and_sends_command() -> assert fake_http.fake_ws.closed is True +@pytest.mark.asyncio +async def test_cloudflare_pty_finalize_flushes_a_partial_character_left_in_the_deque() -> None: + fake_ws = _FakeWebSocket() + sess = _make_session(fake_http=_FakeHttp()) + process_id = await _register_pty_entry(sess, ws=fake_ws, tty=True) + entry = sess._pty_processes[process_id] + + # collect_pty_output handed the first byte of a two byte character back while the stream + # still looked open. the close lands before the finalizer looks, so nothing will drain it + async with entry.output_lock: + entry.output_chunks.append("\u00e9".encode()[:1]) + entry.exit_code = 0 + entry.output_closed.set() + + update = await sess._finalize_pty_update( + process_id=process_id, + entry=entry, + output=b"hi ", + source_text="hi ", + original_token_count=None, + ) + + assert update.process_id is None + assert update.output.decode("utf-8") == "hi \ufffd" + assert not entry.output_chunks + + @pytest.mark.asyncio async def test_cloudflare_pty_write_stdin_sends_input_and_collects_output() -> None: fake_ws = _FakeWebSocket() diff --git a/tests/extensions/sandbox/test_modal.py b/tests/extensions/sandbox/test_modal.py index 44a3fa5c72..13816e75dc 100644 --- a/tests/extensions/sandbox/test_modal.py +++ b/tests/extensions/sandbox/test_modal.py @@ -4444,6 +4444,272 @@ def _exec(self, *command: object, **kwargs: object) -> object: await session.pty_terminate_all() +@pytest.mark.asyncio +async def test_modal_pty_collection_keeps_bytes_it_already_read_when_cancelled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + modal_module, _create_calls, _registry_tags = _load_modal_module(monkeypatch) + + class _OneChunkThenBlocks: + def __init__(self, chunks: list[bytes]) -> None: + self._chunks = chunks + + def __aiter__(self) -> _OneChunkThenBlocks: + return self + + async def __anext__(self) -> bytes: + if self._chunks: + return self._chunks.pop(0) + await asyncio.sleep(3600) + raise AssertionError("unreachable") + + class _FakeProcess: + def __init__(self) -> None: + # the continuation arrives, then the call is cancelled at a later await + self.stdout = _OneChunkThenBlocks(["\u00e9".encode()[1:]]) + self.stderr = _OneChunkThenBlocks([]) + self.poll = _with_aio(lambda: None) + self.terminate = _with_aio(lambda: None) + + class _FakeSandbox: + object_id = "sb-read-cancel" + + 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) + + entry = modal_module._ModalPtyProcessEntry(process=sandbox.process, tty=True) + entry.pending_output = "\u00e9".encode()[:1] + + task = asyncio.create_task( + session._collect_pty_output(entry=entry, yield_time_ms=60_000, max_output_tokens=None) + ) + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + # the stream item is gone, so both halves have to be on the entry or the character is lost + assert entry.pending_output == "\u00e9".encode() + + +@pytest.mark.asyncio +async def test_modal_pty_collection_keeps_its_tail_when_the_call_is_cancelled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + modal_module, _create_calls, _registry_tags = _load_modal_module(monkeypatch) + + class _NeverEndingStream: + def __aiter__(self) -> _NeverEndingStream: + return self + + async def __anext__(self) -> bytes: + await asyncio.sleep(3600) + raise AssertionError("unreachable") + + class _FakeProcess: + def __init__(self) -> None: + self.stdout = _NeverEndingStream() + self.stderr = _NeverEndingStream() + self.poll = _with_aio(lambda: None) + self.terminate = _with_aio(lambda: None) + + class _FakeSandbox: + object_id = "sb-cancel" + + 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) + + entry = modal_module._ModalPtyProcessEntry(process=sandbox.process, tty=True) + held = "\u00e9".encode()[:1] + entry.pending_output = held + + task = asyncio.create_task( + session._collect_pty_output(entry=entry, yield_time_ms=60_000, max_output_tokens=None) + ) + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + # the process is still registered, so the next call has to find the lead byte still there + assert entry.pending_output == held + + +@pytest.mark.asyncio +async def test_modal_pty_finalize_flushes_a_partial_character_when_the_process_has_exited( + monkeypatch: pytest.MonkeyPatch, +) -> None: + modal_module, _create_calls, _registry_tags = _load_modal_module(monkeypatch) + + class _FakeProcess: + def __init__(self) -> None: + self.poll = _with_aio(lambda: 0) + self.terminate = _with_aio(lambda: None) + + class _FakeSandbox: + object_id = "sb-finalize" + + 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) + + # the collector polled while the process was still running, so it held the first byte of a + # two byte character back. by the time the finalizer polls, the process has gone and the + # entry is about to be dropped + entry = modal_module._ModalPtyProcessEntry(process=sandbox.process, tty=True) + entry.pending_output = "\u00e9".encode()[:1] + session._pty_processes[7] = entry + + update = await session._finalize_pty_update( + process_id=7, + entry=entry, + output=b"hi ", + source_text="hi ", + original_token_count=None, + ) + + assert update.exit_code == 0 + assert update.process_id is None + # replaced rather than vanishing with the entry + assert update.output.decode("utf-8") == "hi \ufffd" + assert entry.pending_output == b"" + + +@pytest.mark.asyncio +async def test_modal_pty_output_keeps_a_character_split_across_windows( + monkeypatch: pytest.MonkeyPatch, +) -> None: + modal_module, _create_calls, _registry_tags = _load_modal_module(monkeypatch) + + class _FakeStream: + def __init__(self, chunks: list[bytes]) -> None: + self._chunks = chunks + self._chunk_event = asyncio.Event() + if self._chunks: + self._chunk_event.set() + self.read = _with_aio(self._read) + + def __aiter__(self) -> _FakeStream: + return self + + async def __anext__(self) -> bytes: + while not self._chunks: + self._chunk_event.clear() + await self._chunk_event.wait() + chunk = self._chunks.pop(0) + if not self._chunks: + self._chunk_event.clear() + return chunk + + def append(self, chunk: bytes) -> None: + self._chunks.append(chunk) + self._chunk_event.set() + + def _read(self, size: int | None = None) -> bytes: + if size is None: + raise AssertionError("PTY polling should not call read() with no size") + if self._chunks: + return self._chunks.pop(0) + return b"" + + class _FakeStdin: + def __init__(self, stdout: _FakeStream) -> None: + self.writes: list[bytes] = [] + self._stdout = stdout + self.write = _with_aio(self._write) + self.drain = _with_aio(lambda: None) + + def _write(self, payload: bytes) -> None: + self.writes.append(payload) + # the rest of the character, plus what follows it + self._stdout.append("\u00e9llo".encode()[1:]) + + class _FakeProcess: + def __init__(self) -> None: + # ends mid character, the second byte of e acute only arrives in the next window + self.stdout = _FakeStream([b"h" + "\u00e9".encode()[:1]]) + self.stderr = _FakeStream([]) + self.stdin = _FakeStdin(self.stdout) + self.poll = _with_aio(lambda: None) + self.terminate = _with_aio(lambda: None) + + class _FakeSandbox: + object_id = "sb-split" + + 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.05) + assert started.process_id is not None + + updated = await session.pty_write_stdin( + session_id=started.process_id, + chars="go\n", + yield_time_s=0.05, + ) + + combined = (started.output + updated.output).decode("utf-8") + assert combined == "h\u00e9llo" + assert "\ufffd" not in combined + + await session.pty_terminate_all() + + @pytest.mark.asyncio async def test_modal_pty_start_drains_all_buffered_output_after_exit( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/sandbox/test_pty_output.py b/tests/sandbox/test_pty_output.py index f15bcf85b4..a5e207cf60 100644 --- a/tests/sandbox/test_pty_output.py +++ b/tests/sandbox/test_pty_output.py @@ -1,11 +1,17 @@ from __future__ import annotations import asyncio +import contextlib from collections import deque import pytest -from agents.sandbox.session.pty_output import collect_pty_output +from agents.sandbox.session.pty_output import ( + close_pty_tail, + collect_pty_output, + flush_pty_tail, +) +from agents.sandbox.session.pty_types import truncate_text_by_tokens @pytest.mark.asyncio @@ -24,7 +30,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, _ = await collect_pty_output( output_chunks=output_chunks, output_lock=output_lock, output_notify=output_notify, @@ -46,7 +52,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, _ = await collect_pty_output( output_chunks=output_chunks, output_lock=asyncio.Lock(), output_notify=asyncio.Event(), @@ -57,3 +63,333 @@ def mark_done() -> bool: assert output == b"before done after done" assert original_token_count is None + + +async def _one_window( + chunks: deque[bytes], + lock: asyncio.Lock, + notify: asyncio.Event, + done: dict[str, bool], +) -> bytes: + notify.set() + collected, _, _ = await collect_pty_output( + output_chunks=chunks, + output_lock=lock, + output_notify=notify, + is_done=lambda: done["value"], + yield_time_ms=1, + max_output_tokens=None, + ) + return collected + + +# one, two, three and four byte characters, so every sequence width is split +SPLIT_TEXT = "aé☃\U0001d11eb" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("split", range(1, len(SPLIT_TEXT.encode("utf-8")))) +async def test_collect_pty_output_keeps_a_character_split_across_windows(split: int) -> None: + text = SPLIT_TEXT + raw = text.encode("utf-8") + + chunks: deque[bytes] = deque() + lock = asyncio.Lock() + notify = asyncio.Event() + done = {"value": False} + + chunks.append(raw[:split]) + first = await _one_window(chunks, lock, notify, done) + chunks.append(raw[split:]) + done["value"] = True + second = await _one_window(chunks, lock, notify, done) + + assert (first + second).decode("utf-8") == text + + +@pytest.mark.asyncio +async def test_collect_pty_output_replaces_a_truncated_character_once_done() -> None: + # the stream ends mid character, so there is no later window to complete it + chunks: deque[bytes] = deque([b"hi " + "é".encode()[:1]]) + + collected, _, _ = await collect_pty_output( + output_chunks=chunks, + output_lock=asyncio.Lock(), + output_notify=asyncio.Event(), + is_done=lambda: True, + yield_time_ms=1, + max_output_tokens=None, + ) + + assert collected.decode("utf-8") == "hi �" + assert not chunks + + +@pytest.mark.asyncio +async def test_collect_pty_output_leaves_complete_multibyte_output_alone() -> None: + chunks: deque[bytes] = deque(["héllo".encode()]) + + collected, _, _ = await collect_pty_output( + output_chunks=chunks, + output_lock=asyncio.Lock(), + output_notify=asyncio.Event(), + is_done=lambda: True, + yield_time_ms=1, + max_output_tokens=None, + ) + + assert collected.decode("utf-8") == "héllo" + + +# a lead byte that no character can start with, an overlong form, a value past the +# end of the range, a lone continuation byte, and half of a surrogate pair +@pytest.mark.asyncio +@pytest.mark.parametrize("garbage", [b"\xff", b"\xc0", b"\xc1", b"\xf5", b"\x80", b"\xed\xa0\x80"]) +async def test_collect_pty_output_does_not_hold_back_bytes_that_start_no_character( + garbage: bytes, +) -> None: + # the process is still running, but these bytes will never be completed by a later + # window, so holding them would hide the output until it exits + chunks: deque[bytes] = deque([b"ok " + garbage]) + lock = asyncio.Lock() + notify = asyncio.Event() + done = {"value": False} + + collected = await _one_window(chunks, lock, notify, done) + + assert collected.decode("utf-8").startswith("ok �") + assert not chunks + + +# ED A0..BF leads a surrogate, so no third byte can complete it. the decoder still buffers +# these, and they are the only prefixes it buffers that can never become a character +@pytest.mark.asyncio +@pytest.mark.parametrize("second", [0xA0, 0xAF, 0xBF]) +async def test_collect_pty_output_does_not_hold_back_a_surrogate_lead(second: int) -> None: + chunks: deque[bytes] = deque([b"ok " + bytes([0xED, second])]) + lock = asyncio.Lock() + notify = asyncio.Event() + done = {"value": False} + + collected = await _one_window(chunks, lock, notify, done) + + assert collected.decode("utf-8") == "ok \ufffd\ufffd" + assert not chunks + + +@pytest.mark.asyncio +@pytest.mark.parametrize("second", [0x80, 0x9F]) +async def test_collect_pty_output_still_holds_a_valid_lead_below_the_surrogates( + second: int, +) -> None: + # ED 80..9F is U+D000..U+D7FF, which is a real character, so it must still be waited for + raw = bytes([0xED, second, 0x80]) + chunks: deque[bytes] = deque([raw[:2]]) + lock = asyncio.Lock() + notify = asyncio.Event() + done = {"value": False} + + first = await _one_window(chunks, lock, notify, done) + assert first == b"" + + chunks.append(raw[2:]) + done["value"] = True + second_window = await _one_window(chunks, lock, notify, done) + + assert (first + second_window).decode("utf-8") == raw.decode("utf-8") + + +def test_close_pty_tail_replaces_the_leftover_and_leaves_a_clean_session_alone() -> None: + finished, count = close_pty_tail( + leftover="\u00e9".encode()[:1], + output=b"hi ", + source_text="hi ", + original_token_count=None, + max_output_tokens=None, + ) + assert finished.decode("utf-8") == "hi \ufffd" + + unchanged, same = close_pty_tail( + leftover=b"", + output=b"hi ", + source_text="hi ", + original_token_count=count, + max_output_tokens=None, + ) + assert unchanged == b"hi " + assert same == count + + +def test_close_pty_tail_applies_the_token_cap_to_what_it_adds() -> None: + # the window already truncated to the cap, so the tail cannot be appended past it + capped, _ = close_pty_tail( + leftover="\u00e9".encode()[:1], + output=b"", + source_text="", + original_token_count=None, + max_output_tokens=0, + ) + uncapped, _ = close_pty_tail( + leftover="\u00e9".encode()[:1], + output=b"", + source_text="", + original_token_count=None, + max_output_tokens=None, + ) + + assert uncapped.decode("utf-8") == "\ufffd" + assert len(capped) <= len(uncapped) + + +@pytest.mark.asyncio +async def test_flush_pty_tail_drains_what_the_session_still_holds() -> None: + chunks: deque[bytes] = deque(["\u00e9".encode()[:1]]) + lock = asyncio.Lock() + + flushed, _ = await flush_pty_tail( + output_chunks=chunks, + output_lock=lock, + output=b"hi ", + source_text="hi ", + original_token_count=None, + max_output_tokens=None, + ) + + assert flushed.decode("utf-8") == "hi \ufffd" + assert not chunks + + +@pytest.mark.asyncio +async def test_collect_pty_output_keeps_the_tail_when_a_window_is_cancelled() -> None: + # the lead byte has already left the deque by the time the window decodes, so this is the + # only copy of it. a producer holding the lock must not be able to turn a cancelled call + # into a lost character for the session that carries on + raw = "\u00e9".encode() + chunks: deque[bytes] = deque([raw[:1]]) + lock = asyncio.Lock() + notify = asyncio.Event() + done = {"value": False} + + async def window(yield_time_ms: int) -> bytes: + notify.set() + collected, _, _ = await collect_pty_output( + output_chunks=chunks, + output_lock=lock, + output_notify=notify, + is_done=lambda: done["value"], + yield_time_ms=yield_time_ms, + max_output_tokens=None, + ) + return collected + + task = asyncio.create_task(window(120)) + await asyncio.sleep(0.02) + assert not chunks + + await lock.acquire() + await asyncio.sleep(0.2) + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + lock.release() + + assert list(chunks) == [raw[:1]] + + chunks.append(raw[1:]) + done["value"] = True + assert (await window(60)).decode("utf-8") == "\u00e9" + + +@pytest.mark.asyncio +async def test_collect_pty_output_puts_a_drained_window_back_when_cancelled() -> None: + # the window drains the lead byte a previous one requeued, then the call is cancelled while + # it waits. the session lives on, so those bytes have to go back or its next read reports a + # replacement character for output that did arrive + raw = "\u00e9".encode() + chunks: deque[bytes] = deque([raw[:1]]) + lock = asyncio.Lock() + notify = asyncio.Event() + done = {"value": False} + + async def window(yield_time_ms: int) -> bytes: + notify.set() + collected, _, _ = await collect_pty_output( + output_chunks=chunks, + output_lock=lock, + output_notify=notify, + is_done=lambda: done["value"], + yield_time_ms=yield_time_ms, + max_output_tokens=None, + ) + return collected + + task = asyncio.create_task(window(60_000)) + await asyncio.sleep(0.05) + assert not chunks + + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + assert list(chunks) == [raw[:1]] + + chunks.append(raw[1:]) + done["value"] = True + assert (await window(60)).decode("utf-8") == "\u00e9" + + +def test_close_pty_tail_keeps_the_last_thing_the_process_said() -> None: + # truncation keeps the start and the end, so a tail arriving at the end of a session belongs + # in the part that is kept. folding it into the rendered display instead leaves the old + # ending in place and hides it + source = "A" * 100 + display, count = truncate_text_by_tokens(source, 10) + assert count is not None + + output, recounted = close_pty_tail( + leftover=b"<<>>", + output=display.encode(), + source_text=source, + original_token_count=count, + max_output_tokens=10, + ) + + expected, expected_count = truncate_text_by_tokens(source + "<<>>", 10) + assert output.decode("utf-8") == expected + assert recounted == expected_count + assert "ERROR>>>" in output.decode("utf-8") + assert output != display.encode() + + +def test_close_pty_tail_counts_the_source_and_not_the_shortened_display() -> None: + source = "a" * 100 + display, count = truncate_text_by_tokens(source, 10) + assert count is not None + + _, recounted = close_pty_tail( + leftover="\u00e9".encode()[:1], + output=display.encode(), + source_text=source, + original_token_count=count, + max_output_tokens=10, + ) + + # recounting the display gave fewer tokens here than the window had already measured + assert recounted is not None + assert recounted >= count + + +def test_close_pty_tail_still_folds_the_tail_into_an_untruncated_window() -> None: + display, count = truncate_text_by_tokens("hi ", 10) + assert count is None + + output, recounted = close_pty_tail( + leftover="\u00e9".encode()[:1], + output=display.encode(), + source_text="hi ", + original_token_count=count, + max_output_tokens=10, + ) + + assert output.decode("utf-8") == "hi \ufffd" + assert recounted is None diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index fb13be3c51..1c6674c39c 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import contextlib import io import signal import tarfile @@ -216,6 +217,158 @@ def _unexpected_mkdtemp(*args: object, **kwargs: object) -> str: @pytest.mark.review_optional class TestUnixLocalPty: + @pytest.mark.asyncio + async def test_finalize_still_cleans_up_when_the_drain_is_cancelled( + self, + tmp_path: Path, + ) -> None: + # once the entry is out of the map, pty_terminate_all can no longer reach it, so a + # cancelled drain must not be able to skip its cleanup + session = _RecordingUnixLocalSession(tmp_path) + process = cast( + asyncio.subprocess.Process, + SimpleNamespace(returncode=0, pid=None), + ) + entry = _UnixPtyProcessEntry(process=process, tty=True) + entry.output_closed.set() + entry.output_chunks.append(b"x") + + terminated: list[_UnixPtyProcessEntry] = [] + + async def record_terminate(target: _UnixPtyProcessEntry) -> None: + terminated.append(target) + + session._terminate_pty_entry = record_terminate # type: ignore[method-assign] + + async with session._pty_lock: + session._pty_processes[1] = entry + session._reserved_pty_process_ids.add(1) + + # hold the output lock so the drain blocks after the removal has committed + await entry.output_lock.acquire() + task = asyncio.create_task( + session._finalize_pty_update( + process_id=1, + entry=entry, + output=b"", + original_token_count=None, + source_text="", + ) + ) + await asyncio.sleep(0.05) + assert 1 not in session._pty_processes + + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + entry.output_lock.release() + + assert terminated == [entry] + + @pytest.mark.asyncio + async def test_finalize_does_not_consume_the_tail_before_removal_commits( + self, + tmp_path: Path, + ) -> None: + # the drain empties entry owned state, so if a cancelled finalise can get between it + # and the removal, the session stays registered with its last bytes already gone + session = _RecordingUnixLocalSession(tmp_path) + process = cast( + asyncio.subprocess.Process, + SimpleNamespace(returncode=0, pid=None), + ) + entry = _UnixPtyProcessEntry(process=process, tty=True) + entry.output_closed.set() + raw = "\u00e9".encode() + entry.output_chunks.append(raw[:1]) + + async with session._pty_lock: + session._pty_processes[1] = entry + session._reserved_pty_process_ids.add(1) + + # hold the session map so finalisation blocks on it + await session._pty_lock.acquire() + task = asyncio.create_task( + session._finalize_pty_update( + process_id=1, + entry=entry, + output=b"hi ", + original_token_count=None, + source_text="hi ", + ) + ) + await asyncio.sleep(0.05) + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + session._pty_lock.release() + + # nothing was consumed, so the session can still be finalised properly afterwards + assert list(entry.output_chunks) == [raw[:1]] + assert 1 in session._pty_processes + + @pytest.mark.asyncio + async def test_session_is_not_finalized_while_a_pump_still_holds_output( + self, + tmp_path: Path, + ) -> None: + # the process is reaped before its pump has drained. finalizing on returncode alone + # would drop the session and cancel the pump that still holds the rest of a character + session = _RecordingUnixLocalSession(tmp_path) + process = cast( + asyncio.subprocess.Process, + SimpleNamespace(returncode=0, pid=None), + ) + entry = _UnixPtyProcessEntry(process=process, tty=True) + raw = "\u00e9".encode() + + async with session._pty_lock: + session._pty_processes[1] = entry + session._reserved_pty_process_ids.add(1) + + # a previous window handed the lead byte back, and the continuation is still behind + # the pump, so output_closed is not set yet + entry.output_chunks.append(raw[:1]) + + collected, count, source = await session._collect_pty_output( + entry=entry, yield_time_ms=20, max_output_tokens=None + ) + first = await session._finalize_pty_update( + process_id=1, + entry=entry, + output=collected, + original_token_count=count, + source_text=source, + ) + + # the session has to stay alive, and the lead byte has to stay queued + assert first.process_id == 1 + assert first.exit_code is None + assert first.output == b"" + assert list(entry.output_chunks) == [raw[:1]] + assert 1 in session._pty_processes + + # now the pump delivers the rest and closes + async with entry.output_lock: + entry.output_chunks.append(raw[1:]) + entry.output_notify.set() + entry.output_closed.set() + + collected, count, source = await session._collect_pty_output( + entry=entry, yield_time_ms=20, max_output_tokens=None + ) + final = await session._finalize_pty_update( + process_id=1, + entry=entry, + output=collected, + original_token_count=count, + source_text=source, + ) + + assert final.output.decode("utf-8") == "\u00e9" + assert final.exit_code == 0 + assert final.process_id is None + @pytest.mark.asyncio async def test_tty_fd_close_is_owned_without_blocking_termination( self,