From 02e9e9ac2ed193504d5fcc7466e18546b02f4805 Mon Sep 17 00:00:00 2001 From: coderdailyone Date: Sun, 6 Sep 2026 15:34:12 +0000 Subject: [PATCH 1/3] fix(sandbox): keep a UTF-8 sequence split across PTY yield windows whole Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01DFSMrLZZ3oq1dKmJFAofz6 --- src/agents/sandbox/session/pty_output.py | 33 +++++++++++ tests/sandbox/test_pty_output.py | 71 +++++++++++++++++++++++- tests/sandbox/test_unix_local.py | 32 +++++++++++ 3 files changed, 135 insertions(+), 1 deletion(-) diff --git a/src/agents/sandbox/session/pty_output.py b/src/agents/sandbox/session/pty_output.py index 25cbe774e7..b259762331 100644 --- a/src/agents/sandbox/session/pty_output.py +++ b/src/agents/sandbox/session/pty_output.py @@ -45,6 +45,39 @@ async def collect_pty_output( break output_notify.clear() + if not is_done(): + # A multibyte UTF-8 sequence can straddle two yield windows (the producer wrote + # part of it before the deadline). Hold the incomplete tail back for the next + # collection instead of emitting replacement characters on both sides. + tail_length = incomplete_utf8_tail_length(output) + if tail_length: + async with output_lock: + output_chunks.appendleft(bytes(output[-tail_length:])) + del output[-tail_length:] 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 + + +def incomplete_utf8_tail_length(data: bytes | bytearray) -> int: + """Return how many trailing bytes start a UTF-8 sequence that is not yet complete. + + Only a well-formed prefix counts: a lead byte followed by fewer continuation bytes + than it announces. Invalid bytes are left alone so they decode as replacement + characters immediately rather than being held forever. + """ + limit = min(len(data), 3) + for offset in range(1, limit + 1): + byte = data[-offset] + if byte & 0xC0 == 0x80: + continue # continuation byte; keep looking for the lead byte + if byte & 0xE0 == 0xC0: + expected = 2 + elif byte & 0xF0 == 0xE0: + expected = 3 + elif byte & 0xF8 == 0xF0: + expected = 4 + else: + return 0 # ASCII or an invalid lead byte + return offset if offset < expected else 0 + return 0 diff --git a/tests/sandbox/test_pty_output.py b/tests/sandbox/test_pty_output.py index f15bcf85b4..1b516e035d 100644 --- a/tests/sandbox/test_pty_output.py +++ b/tests/sandbox/test_pty_output.py @@ -5,7 +5,10 @@ import pytest -from agents.sandbox.session.pty_output import collect_pty_output +from agents.sandbox.session.pty_output import ( + collect_pty_output, + incomplete_utf8_tail_length, +) @pytest.mark.asyncio @@ -57,3 +60,69 @@ 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"\xe4\xb8\xad\xe4", 1), + ], +) +def test_incomplete_utf8_tail_length(data: bytes, expected: int) -> None: + assert incomplete_utf8_tail_length(data) == expected diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index fb13be3c51..7bee6f99d1 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -285,6 +285,38 @@ 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 @pytest.mark.requires_native_macos_sandbox async def test_pty_ctrl_c_interrupts_long_running_process(self, tmp_path: Path) -> None: From 7e0211a4c5ba7d8cc8637979d323ebd54f9a1c63 Mon Sep 17 00:00:00 2001 From: coderdailyone Date: Sun, 6 Sep 2026 15:46:36 +0000 Subject: [PATCH 2/3] fix(sandbox): let the UTF-8 decoder define a valid prefix and flush held-back bytes at PTY exit Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01DFSMrLZZ3oq1dKmJFAofz6 --- src/agents/sandbox/sandboxes/docker.py | 7 +++- src/agents/sandbox/sandboxes/unix_local.py | 7 +++- src/agents/sandbox/session/pty_output.py | 42 ++++++++++++---------- tests/sandbox/test_pty_output.py | 32 +++++++++++++++++ tests/sandbox/test_unix_local.py | 27 ++++++++++++++ 5 files changed, 94 insertions(+), 21 deletions(-) diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index 8ca4febe85..f06b5f9c3f 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, drain_pty_output_chunks from ..session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -1267,6 +1267,11 @@ 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 flush whatever is still queued. + deferred = await drain_pty_output_chunks(entry.output_chunks, entry.output_lock) + if deferred: + output += deferred.decode("utf-8", errors="replace").encode("utf-8") async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index ea7f83e9d8..392eccf936 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, drain_pty_output_chunks from ..session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -555,6 +555,11 @@ 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 flush whatever is still queued. + deferred = await drain_pty_output_chunks(entry.output_chunks, entry.output_lock) + if deferred: + output += deferred.decode("utf-8", errors="replace").encode("utf-8") async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) diff --git a/src/agents/sandbox/session/pty_output.py b/src/agents/sandbox/session/pty_output.py index b259762331..2eb663c07d 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 @@ -60,24 +61,27 @@ async def collect_pty_output( def incomplete_utf8_tail_length(data: bytes | bytearray) -> int: - """Return how many trailing bytes start a UTF-8 sequence that is not yet complete. + """Return how many trailing bytes form a valid but not yet complete UTF-8 sequence. - Only a well-formed prefix counts: a lead byte followed by fewer continuation bytes - than it announces. Invalid bytes are left alone so they decode as replacement - characters immediately rather than being held forever. + 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, surrogates, 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. """ - limit = min(len(data), 3) - for offset in range(1, limit + 1): - byte = data[-offset] - if byte & 0xC0 == 0x80: - continue # continuation byte; keep looking for the lead byte - if byte & 0xE0 == 0xC0: - expected = 2 - elif byte & 0xF0 == 0xE0: - expected = 3 - elif byte & 0xF8 == 0xF0: - expected = 4 - else: - return 0 # ASCII or an invalid lead byte - return offset if offset < expected else 0 - return 0 + 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 1b516e035d..48ad6731dd 100644 --- a/tests/sandbox/test_pty_output.py +++ b/tests/sandbox/test_pty_output.py @@ -7,6 +7,7 @@ from agents.sandbox.session.pty_output import ( collect_pty_output, + drain_pty_output_chunks, incomplete_utf8_tail_length, ) @@ -121,8 +122,39 @@ async def test_collect_pty_output_flushes_an_incomplete_sequence_when_done() -> (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 7bee6f99d1..3d3ad980a8 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -317,6 +317,33 @@ async def test_pty_output_keeps_a_utf8_sequence_split_across_yields_whole( 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, + output=b"prefix ", + original_token_count=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 @pytest.mark.requires_native_macos_sandbox async def test_pty_ctrl_c_interrupts_long_running_process(self, tmp_path: Path) -> None: From 3f8fb95483dcdda89e6f39532edb8eed27a1685c Mon Sep 17 00:00:00 2001 From: coderdailyone Date: Mon, 7 Sep 2026 04:29:36 +0000 Subject: [PATCH 3/3] fix(sandbox): apply the PTY token limit once, after draining exit output Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01DFSMrLZZ3oq1dKmJFAofz6 --- src/agents/sandbox/sandboxes/docker.py | 38 +++++++------- src/agents/sandbox/sandboxes/unix_local.py | 38 +++++++------- src/agents/sandbox/session/pty_output.py | 59 +++++++++++++++------- tests/sandbox/test_unix_local.py | 29 ++++++++++- 4 files changed, 106 insertions(+), 58 deletions(-) diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index f06b5f9c3f..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, drain_pty_output_chunks +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) @@ -1268,10 +1268,9 @@ async def _finalize_pty_update( 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 flush whatever is still queued. - deferred = await drain_pty_output_chunks(entry.output_chunks, entry.output_lock) - if deferred: - output += deferred.decode("utf-8", errors="replace").encode("utf-8") + # 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) @@ -1279,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 392eccf936..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, drain_pty_output_chunks +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,18 +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 flush whatever is still queued. - deferred = await drain_pty_output_chunks(entry.output_chunks, entry.output_lock) - if deferred: - output += deferred.decode("utf-8", errors="replace").encode("utf-8") + # 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) @@ -567,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 2eb663c07d..8d7b63f0dc 100644 --- a/src/agents/sandbox/session/pty_output.py +++ b/src/agents/sandbox/session/pty_output.py @@ -19,43 +19,66 @@ 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() - - if not is_done(): - # A multibyte UTF-8 sequence can straddle two yield windows (the producer wrote - # part of it before the deadline). Hold the incomplete tail back for the next - # collection instead of emitting replacement characters on both sides. - tail_length = incomplete_utf8_tail_length(output) - if tail_length: - async with output_lock: + 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:] - text = output.decode("utf-8", errors="replace") + del output[-tail_length:] + return bytes(output) + + +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 @@ -65,9 +88,9 @@ def incomplete_utf8_tail_length(data: bytes | bytearray) -> int: 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, surrogates, 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. + 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: diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index 3d3ad980a8..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 @@ -335,8 +336,8 @@ async def test_finalize_pty_update_flushes_bytes_held_back_when_the_process_exit update = await session._finalize_pty_update( process_id=process_id, entry=entry, - output=b"prefix ", - original_token_count=None, + raw_output=b"prefix ", + max_output_tokens=None, ) assert update.exit_code == 0 @@ -344,6 +345,30 @@ async def test_finalize_pty_update_flushes_bytes_held_back_when_the_process_exit 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: