diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index 8ca4febe85..db1286d560 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -55,7 +55,11 @@ 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_bytes, + drain_pty_output_chunks, + finish_pty_output, +) from ..session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -1080,16 +1084,15 @@ 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( + raw_output = await self._collect_pty_output( entry=entry, yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms), - max_output_tokens=max_output_tokens, ) return await self._finalize_pty_update( process_id=process_id, entry=entry, - output=output, - original_token_count=original_token_count, + raw_output=raw_output, + max_output_tokens=max_output_tokens, ) async def pty_write_stdin( @@ -1126,19 +1129,18 @@ 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( + raw_output = 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 == "" ), - max_output_tokens=max_output_tokens, ) entry.last_used = time.monotonic() return await self._finalize_pty_update( process_id=session_id, entry=entry, - output=output, - original_token_count=original_token_count, + raw_output=raw_output, + max_output_tokens=max_output_tokens, ) async def pty_terminate_all(self) -> None: @@ -1241,15 +1243,13 @@ async def _collect_pty_output( *, entry: _DockerPtyProcessEntry, yield_time_ms: int, - max_output_tokens: int | None, - ) -> tuple[bytes, int | None]: - return await collect_pty_output( + ) -> bytes: + return await collect_pty_output_bytes( 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( @@ -1257,8 +1257,8 @@ async def _finalize_pty_update( *, process_id: int, entry: _DockerPtyProcessEntry, - output: bytes, - original_token_count: int | None, + raw_output: bytes, + max_output_tokens: int | None, ) -> PtyExecUpdate: if entry.output_closed.is_set() and entry.exit_code is None: await self._refresh_pty_exit_code(entry) @@ -1267,6 +1267,10 @@ async def _finalize_pty_update( live_process_id: int | None = process_id if exit_code is not None: + # The collector may have held back a partial UTF-8 sequence for a later poll; + # there is none once the entry is removed, so take whatever is still queued + # before the token limit is applied to the whole update. + raw_output += await drain_pty_output_chunks(entry.output_chunks, entry.output_lock) async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) @@ -1274,6 +1278,7 @@ async def _finalize_pty_update( await self._terminate_pty_entry(removed) live_process_id = None + output, original_token_count = finish_pty_output(raw_output, max_output_tokens) return PtyExecUpdate( process_id=live_process_id, output=output, diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index ea7f83e9d8..ddc5f1d7b4 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -48,7 +48,11 @@ 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_bytes, + drain_pty_output_chunks, + finish_pty_output, +) from ..session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -400,16 +404,15 @@ 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( + raw_output = await self._collect_pty_output( entry=entry, yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms), - max_output_tokens=max_output_tokens, ) return await self._finalize_pty_update( process_id=process_id, entry=entry, - output=output, - original_token_count=original_token_count, + raw_output=raw_output, + max_output_tokens=max_output_tokens, ) async def pty_write_stdin( @@ -442,19 +445,18 @@ 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( + raw_output = 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 == "" ), - max_output_tokens=max_output_tokens, ) entry.last_used = time.monotonic() return await self._finalize_pty_update( process_id=session_id, entry=entry, - output=output, - original_token_count=original_token_count, + raw_output=raw_output, + max_output_tokens=max_output_tokens, ) async def pty_terminate_all(self) -> None: @@ -532,15 +534,13 @@ async def _collect_pty_output( *, entry: _UnixPtyProcessEntry, yield_time_ms: int, - max_output_tokens: int | None, - ) -> tuple[bytes, int | None]: - return await collect_pty_output( + ) -> bytes: + return await collect_pty_output_bytes( 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( @@ -548,13 +548,17 @@ async def _finalize_pty_update( *, process_id: int, entry: _UnixPtyProcessEntry, - output: bytes, - original_token_count: int | None, + raw_output: bytes, + max_output_tokens: int | None, ) -> PtyExecUpdate: exit_code: int | None = entry.process.returncode live_process_id: int | None = process_id if exit_code is not None: + # The collector may have held back a partial UTF-8 sequence for a later poll; + # there is none once the entry is removed, so take whatever is still queued + # before the token limit is applied to the whole update. + raw_output += await drain_pty_output_chunks(entry.output_chunks, entry.output_lock) async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) @@ -562,6 +566,7 @@ async def _finalize_pty_update( await self._terminate_pty_entry(removed) live_process_id = None + output, original_token_count = finish_pty_output(raw_output, max_output_tokens) return PtyExecUpdate( process_id=live_process_id, output=output, diff --git a/src/agents/sandbox/session/pty_output.py b/src/agents/sandbox/session/pty_output.py index 25cbe774e7..8d7b63f0dc 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 @@ -18,33 +19,92 @@ async def collect_pty_output( max_output_tokens: int | None, ) -> tuple[bytes, int | None]: """Collect and truncate PTY output until the deadline or provider completion.""" + raw_output = await collect_pty_output_bytes( + output_chunks=output_chunks, + output_lock=output_lock, + output_notify=output_notify, + is_done=is_done, + yield_time_ms=yield_time_ms, + ) + return finish_pty_output(raw_output, max_output_tokens) + + +async def collect_pty_output_bytes( + *, + output_chunks: deque[bytes], + output_lock: asyncio.Lock, + output_notify: asyncio.Event, + is_done: Callable[[], bool], + yield_time_ms: int, +) -> bytes: + """Collect raw PTY output until the deadline or provider completion. + + A multibyte UTF-8 sequence can straddle two yield windows (the producer wrote part of + it before the deadline). While the provider is still running, the incomplete tail is + held back in `output_chunks` for the next collection instead of being emitted as + replacement characters on both sides. + """ 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(): 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() + tail_length = incomplete_utf8_tail_length(output) + if tail_length: + async with output_lock: + # Re-check under the lock: once the provider is done nothing else will arrive, + # so the tail must be flushed now rather than parked in the queue. + if not is_done(): + output_chunks.appendleft(bytes(output[-tail_length:])) + del output[-tail_length:] + return bytes(output) + - text = output.decode("utf-8", errors="replace") +def finish_pty_output(raw_output: bytes, max_output_tokens: int | None) -> tuple[bytes, int | None]: + """Decode collected PTY bytes and apply the token limit once.""" + text = raw_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 + + +def incomplete_utf8_tail_length(data: bytes | bytearray) -> int: + """Return how many trailing bytes form a valid but not yet complete UTF-8 sequence. + + Python's incremental decoder decides what counts as a valid prefix, so invalid + leaders (``0xC0``, ``0xC1``, ``0xF5`` and up) and ill-formed second bytes (overlong + forms, code points past U+10FFFF) are not held back: they decode to replacement + characters immediately, as before. A pending sequence is at most three bytes long, so + only the tail needs to be inspected. + """ + tail = bytes(data[-3:]) + if not tail: + return 0 + decoder = codecs.getincrementaldecoder("utf-8")(errors="replace") + decoder.decode(tail, final=False) + pending, _ = decoder.getstate() + return len(pending) + + +async def drain_pty_output_chunks(output_chunks: deque[bytes], output_lock: asyncio.Lock) -> bytes: + """Take every queued chunk, including bytes a collection held back.""" + output = bytearray() + async with output_lock: + while output_chunks: + output.extend(output_chunks.popleft()) + return bytes(output) diff --git a/tests/sandbox/test_pty_output.py b/tests/sandbox/test_pty_output.py index f15bcf85b4..48ad6731dd 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.pty_output import ( + collect_pty_output, + drain_pty_output_chunks, + incomplete_utf8_tail_length, +) @pytest.mark.asyncio @@ -57,3 +61,100 @@ def mark_done() -> bool: assert output == b"before done after done" assert original_token_count is None + + +@pytest.mark.asyncio +async def test_collect_pty_output_holds_back_a_split_utf8_sequence_until_it_completes() -> None: + output_chunks: deque[bytes] = deque([b"prefix \xe4\xb8"]) + output_lock = asyncio.Lock() + output_notify = asyncio.Event() + + first_output, _ = await collect_pty_output( + output_chunks=output_chunks, + output_lock=output_lock, + output_notify=output_notify, + is_done=lambda: False, + yield_time_ms=10, + max_output_tokens=None, + ) + assert first_output == b"prefix " + assert list(output_chunks) == [b"\xe4\xb8"] + + output_chunks.append(b"\xad\n") + second_output, _ = await collect_pty_output( + output_chunks=output_chunks, + output_lock=output_lock, + output_notify=output_notify, + is_done=lambda: False, + yield_time_ms=10, + max_output_tokens=None, + ) + assert second_output == "中\n".encode() + assert not output_chunks + + +@pytest.mark.asyncio +async def test_collect_pty_output_flushes_an_incomplete_sequence_when_done() -> None: + output_chunks: deque[bytes] = deque([b"tail \xe4\xb8"]) + + output, _ = await collect_pty_output( + output_chunks=output_chunks, + output_lock=asyncio.Lock(), + output_notify=asyncio.Event(), + is_done=lambda: True, + yield_time_ms=10, + max_output_tokens=None, + ) + assert output == "tail \ufffd".encode() + assert not output_chunks + + +@pytest.mark.parametrize( + ("data", "expected"), + [ + (b"", 0), + (b"ascii", 0), + ("中".encode(), 0), + (b"\xe4", 1), + (b"\xe4\xb8", 2), + (b"\xf0\x9f\x98", 3), + (b"\xc3", 1), + (b"ok \xf0\x9f", 2), + (b"\xb8\xad", 0), # stray continuation bytes are invalid, not incomplete + (b"\xff", 0), # invalid lead byte + (b"\xc0", 0), # overlong two-byte leaders are never valid + (b"\xc1", 0), + (b"\xf5", 0), # beyond U+10FFFF + (b"\xe0\x80", 0), # overlong three-byte form + (b"\xf4\x90", 0), # past U+10FFFF + (b"\xf0\x80", 0), # overlong four-byte form + (b"\xe4\xb8\xad\xe4", 1), + ], +) +def test_incomplete_utf8_tail_length(data: bytes, expected: int) -> None: + assert incomplete_utf8_tail_length(data) == expected + + +@pytest.mark.asyncio +async def test_collect_pty_output_returns_an_invalid_leader_immediately() -> None: + output_chunks: deque[bytes] = deque([b"\xc0"]) + + output, _ = await collect_pty_output( + output_chunks=output_chunks, + output_lock=asyncio.Lock(), + output_notify=asyncio.Event(), + is_done=lambda: False, + yield_time_ms=10, + max_output_tokens=None, + ) + assert output == "\ufffd".encode() + assert not output_chunks + + +@pytest.mark.asyncio +async def test_drain_pty_output_chunks_takes_everything_queued() -> None: + output_chunks: deque[bytes] = deque([b"\xe4\xb8", b"\xad\n"]) + + assert await drain_pty_output_chunks(output_chunks, asyncio.Lock()) == "中\n".encode() + assert not output_chunks + assert await drain_pty_output_chunks(output_chunks, asyncio.Lock()) == b"" diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index fb13be3c51..05ecce209d 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -22,6 +22,7 @@ UnixLocalSandboxSessionState, _UnixPtyProcessEntry, ) +from agents.sandbox.session.pty_types import truncate_text_by_tokens from agents.sandbox.snapshot import NoopSnapshot from agents.sandbox.types import ExecResult, User @@ -285,6 +286,89 @@ async def test_pty_exec_write_poll_and_unknown_session_errors(self, tmp_path: Pa with pytest.raises(PtySessionNotFoundError): await session.pty_write_stdin(session_id=999_999, chars="") + @pytest.mark.asyncio + @pytest.mark.requires_native_macos_sandbox + async def test_pty_output_keeps_a_utf8_sequence_split_across_yields_whole( + self, tmp_path: Path + ) -> None: + client = UnixLocalSandboxClient() + manifest = Manifest(root=str(tmp_path / "workspace")) + + async with await client.create(manifest=manifest, snapshot=None, options=None) as session: + # Emit the three bytes of U+4E2D across two writes with a pause in between so the + # first yield window closes in the middle of the sequence. + started = await session.pty_exec_start( + "sh", + "-c", + "printf '\\344\\270'; sleep 0.4; printf '\\255\\n'", + shell=False, + tty=False, + yield_time_s=0.1, + ) + assert started.process_id is not None + chunks = [started.output] + exit_code = started.exit_code + while exit_code is None: + update = await session.pty_write_stdin( + session_id=started.process_id, chars="", yield_time_s=0.1 + ) + chunks.append(update.output) + exit_code = update.exit_code + + assert exit_code == 0 + assert b"".join(chunks).decode("utf-8") == "中\n" + + @pytest.mark.asyncio + async def test_finalize_pty_update_flushes_bytes_held_back_when_the_process_exited( + self, tmp_path: Path + ) -> None: + session = _RecordingUnixLocalSession(tmp_path / "workspace") + process = await asyncio.create_subprocess_exec("true") + await process.wait() + # The process has exited but the output pump has not marked output closed yet, + # so the previous collection held the partial sequence back in the queue. + entry = unix_local_module._UnixPtyProcessEntry(process=process, tty=False) + entry.output_chunks.extend([b"\xe4\xb8", b"\xad\n"]) + process_id = 4242 + async with session._pty_lock: + session._pty_processes[process_id] = entry + + update = await session._finalize_pty_update( + process_id=process_id, + entry=entry, + raw_output=b"prefix ", + max_output_tokens=None, + ) + + assert update.exit_code == 0 + assert update.process_id is None + assert update.output == "prefix 中\n".encode() + assert not entry.output_chunks + + @pytest.mark.asyncio + async def test_finalize_pty_update_applies_the_token_limit_after_flushing( + self, tmp_path: Path + ) -> None: + session = _RecordingUnixLocalSession(tmp_path / "workspace") + process = await asyncio.create_subprocess_exec("true") + await process.wait() + entry = unix_local_module._UnixPtyProcessEntry(process=process, tty=False) + entry.output_chunks.extend([b"\xe4\xb8", b"\xad" * 1 + b"\n" + b"x" * 64]) + async with session._pty_lock: + session._pty_processes[4243] = entry + + update = await session._finalize_pty_update( + process_id=4243, + entry=entry, + raw_output=b"abcd", + max_output_tokens=1, + ) + + expected_text, expected_count = truncate_text_by_tokens("abcd中\n" + "x" * 64, 1) + assert update.output == expected_text.encode() + assert update.original_token_count == expected_count + assert expected_count is not None + @pytest.mark.asyncio @pytest.mark.requires_native_macos_sandbox async def test_pty_ctrl_c_interrupts_long_running_process(self, tmp_path: Path) -> None: