From ea133bf8d30a0872f416ae64b8fa1e0a2cf8fead Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 02:00:06 +0500 Subject: [PATCH 01/15] keep a utf-8 character that straddles two pty collection windows pty output is collected in repeated windows over one persistent deque, and each window decoded with errors=replace. a character whose bytes land either side of a window boundary was therefore replaced twice, once per half, and the bytes were gone before any later window could put it back together. no error, just wrong text the incomplete tail is now handed back to the deque for the next window. once the provider is done there is no next window so those bytes really are invalid and still get replaced --- src/agents/sandbox/session/pty_output.py | 32 ++++++++++ tests/sandbox/test_pty_output.py | 76 ++++++++++++++++++++++++ 2 files changed, 108 insertions(+) diff --git a/src/agents/sandbox/session/pty_output.py b/src/agents/sandbox/session/pty_output.py index 25cbe774e7..61d6126750 100644 --- a/src/agents/sandbox/session/pty_output.py +++ b/src/agents/sandbox/session/pty_output.py @@ -8,6 +8,26 @@ from .pty_types import truncate_text_by_tokens +def _incomplete_utf8_suffix_len(data: bytes | bytearray) -> int: + """Length of a trailing UTF-8 sequence that is still waiting for its remaining bytes. + + Returns 0 when the buffer does not end mid character. + """ + for back in range(1, min(4, len(data)) + 1): + byte = data[-back] + if byte < 0x80: + return 0 + if byte >= 0xC0: + if byte < 0xE0: + expected = 2 + elif byte < 0xF0: + expected = 3 + else: + expected = 4 + return back if back < expected else 0 + return 0 + + async def collect_pty_output( *, output_chunks: deque[bytes], @@ -45,6 +65,18 @@ async def collect_pty_output( break output_notify.clear() + # PTY output is collected in repeated windows over one persistent deque, so a + # character whose bytes straddle a window boundary would be replaced twice and + # lost. Hand the incomplete tail back for the next window to finish. Once the + # provider is done there is no next window, so the bytes are genuinely invalid. + if not is_done(): + held_back = _incomplete_utf8_suffix_len(output) + if held_back: + tail = bytes(output[-held_back:]) + del output[-held_back:] + async with output_lock: + output_chunks.appendleft(tail) + 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 diff --git a/tests/sandbox/test_pty_output.py b/tests/sandbox/test_pty_output.py index f15bcf85b4..f6f465f704 100644 --- a/tests/sandbox/test_pty_output.py +++ b/tests/sandbox/test_pty_output.py @@ -57,3 +57,79 @@ 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" From f077787878081b0a0ac9aa70525bf46c649e8760 Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 02:07:07 +0500 Subject: [PATCH 02/15] only hold back bytes that can actually start a character my first version read the lead byte and worked out the width by hand, which accepted lead bytes that are not legal. ff c0 c1 and f5 all looked like the start of something longer so they were held back and requeued every window, and the replacement character that used to appear right away stayed hidden until the process exited using an incremental decoder instead. it buffers only a real partial sequence and replaces anything that cannot begin a character straight away, and passing final when the provider is done closes it so a truncated tail is replaced too. it also knows about overlong forms and surrogates which my table did not added the invalid lead bytes as cases. four of the six fail on the old version --- src/agents/sandbox/session/pty_output.py | 45 +++++++----------------- tests/sandbox/test_pty_output.py | 20 +++++++++++ 2 files changed, 33 insertions(+), 32 deletions(-) diff --git a/src/agents/sandbox/session/pty_output.py b/src/agents/sandbox/session/pty_output.py index 61d6126750..02bdade123 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,26 +9,6 @@ from .pty_types import truncate_text_by_tokens -def _incomplete_utf8_suffix_len(data: bytes | bytearray) -> int: - """Length of a trailing UTF-8 sequence that is still waiting for its remaining bytes. - - Returns 0 when the buffer does not end mid character. - """ - for back in range(1, min(4, len(data)) + 1): - byte = data[-back] - if byte < 0x80: - return 0 - if byte >= 0xC0: - if byte < 0xE0: - expected = 2 - elif byte < 0xF0: - expected = 3 - else: - expected = 4 - return back if back < expected else 0 - return 0 - - async def collect_pty_output( *, output_chunks: deque[bytes], @@ -65,18 +46,18 @@ async def collect_pty_output( break output_notify.clear() - # PTY output is collected in repeated windows over one persistent deque, so a - # character whose bytes straddle a window boundary would be replaced twice and - # lost. Hand the incomplete tail back for the next window to finish. Once the - # provider is done there is no next window, so the bytes are genuinely invalid. - if not is_done(): - held_back = _incomplete_utf8_suffix_len(output) - if held_back: - tail = bytes(output[-held_back:]) - del output[-held_back:] - async with output_lock: - output_chunks.appendleft(tail) + # 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. + decoder = codecs.getincrementaldecoder("utf-8")("replace") + text = decoder.decode(output, final=is_done()) + pending = decoder.getstate()[0] + if pending: + async with output_lock: + output_chunks.appendleft(bytes(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 diff --git a/tests/sandbox/test_pty_output.py b/tests/sandbox/test_pty_output.py index f6f465f704..88836451c2 100644 --- a/tests/sandbox/test_pty_output.py +++ b/tests/sandbox/test_pty_output.py @@ -133,3 +133,23 @@ async def test_collect_pty_output_leaves_complete_multibyte_output_alone() -> No ) 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 From 144efe19978399147d691484dbccba0e93a037e3 Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 02:58:09 +0500 Subject: [PATCH 03/15] replace a surrogate lead instead of waiting for a byte that cannot come the decoder buffers ED A0..BF even though that leads a surrogate and no third byte can finish it. held back it would go round the deque again every window and the output would stay hidden until the process exits, which is the same thing the last commit was meant to stop i checked every one and two and three byte prefix the decoder buffers. those 32 are the only ones that can never become a character, everything else it holds is still waiting on a real continuation. so the check is only for them ED 80..9F is U+D000 to U+D7FF and is a real character, so that still waits --- src/agents/sandbox/session/pty_output.py | 13 ++++++-- tests/sandbox/test_pty_output.py | 38 ++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 2 deletions(-) diff --git a/src/agents/sandbox/session/pty_output.py b/src/agents/sandbox/session/pty_output.py index 02bdade123..794473a392 100644 --- a/src/agents/sandbox/session/pty_output.py +++ b/src/agents/sandbox/session/pty_output.py @@ -54,10 +54,19 @@ async def collect_pty_output( # decoder once the provider is done replaces a tail that no later window will finish. decoder = codecs.getincrementaldecoder("utf-8")("replace") text = decoder.decode(output, final=is_done()) - pending = decoder.getstate()[0] + 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"" + if pending: async with output_lock: - output_chunks.appendleft(bytes(pending)) + output_chunks.appendleft(pending) truncated, original_token_count = truncate_text_by_tokens(text, max_output_tokens) return truncated.encode("utf-8", errors="replace"), original_token_count diff --git a/tests/sandbox/test_pty_output.py b/tests/sandbox/test_pty_output.py index 88836451c2..53a7a871bb 100644 --- a/tests/sandbox/test_pty_output.py +++ b/tests/sandbox/test_pty_output.py @@ -153,3 +153,41 @@ async def test_collect_pty_output_does_not_hold_back_bytes_that_start_no_charact 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") From 91f5fb97d111286f7766ad32de151a4e90bad95a Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 03:12:51 +0500 Subject: [PATCH 04/15] fix the same split character on the other three pty backends fixing the shared helper only covered unix_local docker blaxel and daytona. cloudflare e2b and modal each carry their own copy of the collector and each one decoded its window with errors=replace, so the character was still lost there cloudflare and e2b were near copies of the shared helper down to the deque and the lock and the notify event, so they just call it now the way daytona and blaxel already did. that deletes the copies rather than fixing them twice modal reads a stream instead of a deque so it cannot call the helper. pulled the window decode out into decode_pty_window and modal keeps its own tail on the entry, so there is still one place that knows the utf-8 rules added a split character test for modal. it gives h??llo before this --- .../extensions/sandbox/cloudflare/sandbox.py | 37 ++------ src/agents/extensions/sandbox/e2b/sandbox.py | 40 ++------ .../extensions/sandbox/modal/sandbox.py | 9 +- src/agents/sandbox/session/pty_output.py | 35 ++++--- tests/extensions/sandbox/test_modal.py | 93 +++++++++++++++++++ 5 files changed, 141 insertions(+), 73 deletions(-) diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index bb8d7c37e6..f39063142f 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -59,6 +59,7 @@ _settle_mount_transition, with_ephemeral_mounts_removed, ) +from ....sandbox.session.pty_output import collect_pty_output from ....sandbox.session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -67,7 +68,6 @@ clamp_pty_yield_time_ms, process_id_to_prune_from_meta, resolve_pty_write_yield_time_ms, - truncate_text_by_tokens, ) from ....sandbox.session.runtime_helpers import RESOLVE_WORKSPACE_PATH_HELPER, RuntimeHelperScript from ....sandbox.session.sandbox_client import BaseSandboxClient, BaseSandboxClientOptions @@ -1034,33 +1034,14 @@ async def _collect_pty_output( 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 + 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, diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index 389b665c44..7e91619733 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -53,6 +53,7 @@ from ....sandbox.session.base_sandbox_session import BaseSandboxSession from ....sandbox.session.dependencies import Dependencies from ....sandbox.session.manager import Instrumentation +from ....sandbox.session.pty_output import collect_pty_output from ....sandbox.session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -61,7 +62,6 @@ clamp_pty_yield_time_ms, process_id_to_prune_from_meta, resolve_pty_write_yield_time_ms, - truncate_text_by_tokens, ) from ....sandbox.session.runtime_helpers import RESOLVE_WORKSPACE_PATH_HELPER, RuntimeHelperScript from ....sandbox.session.sandbox_client import BaseSandboxClient, BaseSandboxClientOptions @@ -1212,36 +1212,14 @@ async def _collect_pty_output( 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 + 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: diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index 848b8e12a3..246bbbac30 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 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): @@ -983,7 +985,9 @@ async def _collect_pty_output( max_output_tokens: int | None, ) -> tuple[bytes, int | None]: deadline = time.monotonic() + (yield_time_ms / 1000) - chunks = bytearray() + # a character split across two windows starts in the tail the last one held back + chunks = bytearray(entry.pending_output) + entry.pending_output = b"" while True: stdout_chunk = await self._read_modal_stream(entry=entry, stream_name="stdout") @@ -1010,7 +1014,8 @@ async def _collect_pty_output( break await asyncio.sleep(min(_PTY_POLL_INTERVAL_S, remaining_s)) - text = chunks.decode("utf-8", errors="replace") + exited = await self._peek_exit_code(entry.process) is not None + 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 diff --git a/src/agents/sandbox/session/pty_output.py b/src/agents/sandbox/session/pty_output.py index 794473a392..eec2f299f2 100644 --- a/src/agents/sandbox/session/pty_output.py +++ b/src/agents/sandbox/session/pty_output.py @@ -9,6 +9,28 @@ 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 + + async def collect_pty_output( *, output_chunks: deque[bytes], @@ -52,18 +74,7 @@ async def collect_pty_output( # 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. - decoder = codecs.getincrementaldecoder("utf-8")("replace") - text = decoder.decode(output, final=is_done()) - 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"" - + text, pending = decode_pty_window(output, is_final=is_done()) if pending: async with output_lock: output_chunks.appendleft(pending) diff --git a/tests/extensions/sandbox/test_modal.py b/tests/extensions/sandbox/test_modal.py index 44a3fa5c72..f06e99b73f 100644 --- a/tests/extensions/sandbox/test_modal.py +++ b/tests/extensions/sandbox/test_modal.py @@ -4444,6 +4444,99 @@ def _exec(self, *command: object, **kwargs: object) -> object: await session.pty_terminate_all() +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, From fa487b9275e9731c07cf4a00efdab54f6c54d3b8 Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 03:20:38 +0500 Subject: [PATCH 05/15] flush the modal tail when the finalizer is the one that sees the exit the collector polls for the exit code and the finalizer polls again after it. a process that ends between those two polls means the collector kept a partial character back for a next window, and then the finalizer drops the entry and the bytes go with it that is worse than what was there before my change. the old code replaced those bytes in the same window so you saw a question mark. this way the output just loses them quietly so the finalizer closes the tail before it removes the entry. added a test for it, which gives hi instead of hi ? on the last commit also put back the pytest.mark.asyncio i dropped off the split test last time. asyncio_mode is auto so it still ran, but every other test here carries it --- .../extensions/sandbox/modal/sandbox.py | 8 +++ tests/extensions/sandbox/test_modal.py | 53 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index 246bbbac30..9424ae8917 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -1119,6 +1119,14 @@ async def _finalize_pty_update( exit_code = await self._peek_exit_code(entry.process) live_process_id: int | None = process_id if exit_code is not None: + # The collector polled before this and may have seen the process still running, so + # it kept a partial character back for a window that is never going to come. Close + # it here, otherwise those bytes go out with the entry and the output loses them. + if entry.pending_output: + tail, _ = decode_pty_window(entry.pending_output, is_final=True) + entry.pending_output = b"" + output += tail.encode("utf-8", errors="replace") + async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) diff --git a/tests/extensions/sandbox/test_modal.py b/tests/extensions/sandbox/test_modal.py index f06e99b73f..9176d97767 100644 --- a/tests/extensions/sandbox/test_modal.py +++ b/tests/extensions/sandbox/test_modal.py @@ -4444,6 +4444,59 @@ def _exec(self, *command: object, **kwargs: object) -> object: await session.pty_terminate_all() +@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 ", + 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: From e983ef268946ba8c20e27117f1da7f6ab77399b6 Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 03:23:39 +0500 Subject: [PATCH 06/15] flush the cloudflare and e2b tails too i said on the pr that these two did not have the race modal had. that was wrong and i checked it properly after saying it their tail goes back on entry.output_chunks rather than a field, but the shape is the same. collect_pty_output asks is_done one last time, the finalizer asks again, and if it flipped in between the entry is popped with the bytes still sitting in the deque so both finalizers drain what is left and close it before removing the entry. the cloudflare test gives hi instead of hi ? on the last commit --- .../extensions/sandbox/cloudflare/sandbox.py | 13 +++++++++- src/agents/extensions/sandbox/e2b/sandbox.py | 13 +++++++++- tests/extensions/sandbox/test_cloudflare.py | 26 +++++++++++++++++++ 3 files changed, 50 insertions(+), 2 deletions(-) diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index f39063142f..0b3b7164d4 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -59,7 +59,7 @@ _settle_mount_transition, with_ephemeral_mounts_removed, ) -from ....sandbox.session.pty_output import collect_pty_output +from ....sandbox.session.pty_output import collect_pty_output, decode_pty_window from ....sandbox.session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -1054,6 +1054,17 @@ async def _finalize_pty_update( exit_code = entry.exit_code if entry.output_closed.is_set() else None live_process_id: int | None = process_id if entry.output_closed.is_set(): + # collect_pty_output may have handed a partial character back to the deque while + # this looked like it was still running. nothing is going to drain it now, so close + # it here rather than let it leave with the entry + leftover = bytearray() + async with entry.output_lock: + while entry.output_chunks: + leftover.extend(entry.output_chunks.popleft()) + if leftover: + tail, _ = decode_pty_window(leftover, is_final=True) + output += tail.encode("utf-8", errors="replace") + 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/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index 7e91619733..215d073765 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -53,7 +53,7 @@ from ....sandbox.session.base_sandbox_session import BaseSandboxSession from ....sandbox.session.dependencies import Dependencies from ....sandbox.session.manager import Instrumentation -from ....sandbox.session.pty_output import collect_pty_output +from ....sandbox.session.pty_output import collect_pty_output, decode_pty_window from ....sandbox.session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -1251,6 +1251,17 @@ async def _finalize_pty_update( live_process_id: int | None = process_id if exit_code is not None: + # collect_pty_output may have handed a partial character back to the deque while + # this looked like it was still running. nothing is going to drain it now, so close + # it here rather than let it leave with the entry + leftover = bytearray() + async with entry.output_lock: + while entry.output_chunks: + leftover.extend(entry.output_chunks.popleft()) + if leftover: + tail, _ = decode_pty_window(leftover, is_final=True) + output += tail.encode("utf-8", errors="replace") + async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) diff --git a/tests/extensions/sandbox/test_cloudflare.py b/tests/extensions/sandbox/test_cloudflare.py index bcd47257e1..17470d87c3 100644 --- a/tests/extensions/sandbox/test_cloudflare.py +++ b/tests/extensions/sandbox/test_cloudflare.py @@ -1730,6 +1730,32 @@ 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 ", + 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() From 23b2189665d6032255adb3f8f549c3cbbe48927a Mon Sep 17 00:00:00 2001 From: HuzaifaChaudary Date: Sat, 29 Aug 2026 04:29:00 +0500 Subject: [PATCH 07/15] openai --- .../extensions/sandbox/blaxel/sandbox.py | 12 ++++- .../extensions/sandbox/cloudflare/sandbox.py | 22 ++++---- .../extensions/sandbox/daytona/sandbox.py | 12 ++++- src/agents/extensions/sandbox/e2b/sandbox.py | 22 ++++---- .../extensions/sandbox/modal/sandbox.py | 23 +++++---- src/agents/sandbox/sandboxes/docker.py | 12 ++++- src/agents/sandbox/sandboxes/unix_local.py | 12 ++++- src/agents/sandbox/session/pty_output.py | 51 +++++++++++++++++++ 8 files changed, 130 insertions(+), 36 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index aee3211fbd..508c3553eb 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, @@ -866,6 +866,7 @@ async def pty_exec_start( entry=entry, output=output, original_token_count=original_token_count, + max_output_tokens=max_output_tokens, ) async def pty_write_stdin( @@ -903,6 +904,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + max_output_tokens=max_output_tokens, ) async def pty_terminate_all(self) -> None: @@ -981,11 +983,19 @@ async def _finalize_pty_update( entry: _BlaxelPtySessionEntry, output: bytes, original_token_count: int | None, + max_output_tokens: int | None, ) -> PtyExecUpdate: exit_code = entry.exit_code if entry.done else None live_process_id: int | None = process_id if entry.done: + output, original_token_count = await flush_pty_tail( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output=output, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) async with self._pty_lock: removed = self._pty_sessions.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index 0b3b7164d4..fc25bf7fe1 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -59,7 +59,7 @@ _settle_mount_transition, with_ephemeral_mounts_removed, ) -from ....sandbox.session.pty_output import collect_pty_output, decode_pty_window +from ....sandbox.session.pty_output import collect_pty_output, flush_pty_tail from ....sandbox.session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -1050,20 +1050,18 @@ async def _finalize_pty_update( entry: _CloudflarePtyProcessEntry, output: bytes, original_token_count: int | None, + max_output_tokens: int | None, ) -> PtyExecUpdate: exit_code = entry.exit_code if entry.output_closed.is_set() else None live_process_id: int | None = process_id if entry.output_closed.is_set(): - # collect_pty_output may have handed a partial character back to the deque while - # this looked like it was still running. nothing is going to drain it now, so close - # it here rather than let it leave with the entry - leftover = bytearray() - async with entry.output_lock: - while entry.output_chunks: - leftover.extend(entry.output_chunks.popleft()) - if leftover: - tail, _ = decode_pty_window(leftover, is_final=True) - output += tail.encode("utf-8", errors="replace") + output, original_token_count = await flush_pty_tail( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output=output, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) @@ -1222,6 +1220,7 @@ async def pty_exec_start( entry=entry, output=output, original_token_count=original_token_count, + max_output_tokens=max_output_tokens, ) async def pty_write_stdin( @@ -1259,6 +1258,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + 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..aeffea1642 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, @@ -765,6 +765,7 @@ async def _on_data(chunk: bytes | str) -> None: entry=entry, output=output, original_token_count=original_token_count, + max_output_tokens=max_output_tokens, ) async def _run_pty_waiter(self, entry: _DaytonaPtySessionEntry) -> None: @@ -845,6 +846,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + max_output_tokens=max_output_tokens, ) async def _finalize_pty_update( @@ -854,11 +856,19 @@ async def _finalize_pty_update( entry: _DaytonaPtySessionEntry, output: bytes, original_token_count: int | None, + max_output_tokens: int | None, ) -> PtyExecUpdate: exit_code = entry.exit_code if entry.done else None live_process_id: int | None = process_id if entry.done: + output, original_token_count = await flush_pty_tail( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output=output, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) async with self._pty_lock: removed = self._pty_sessions.pop(process_id, None) self._reserved_pty_process_ids.discard(process_id) diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index 215d073765..ee7323720b 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -53,7 +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, decode_pty_window +from ....sandbox.session.pty_output import collect_pty_output, flush_pty_tail from ....sandbox.session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -1060,6 +1060,7 @@ async def _append_output(payload: bytes | bytearray | str | object) -> None: entry=entry, output=output, original_token_count=original_token_count, + max_output_tokens=max_output_tokens, ) async def pty_write_stdin( @@ -1100,6 +1101,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + max_output_tokens=max_output_tokens, ) async def pty_terminate_all(self) -> None: @@ -1246,21 +1248,19 @@ async def _finalize_pty_update( entry: _E2BPtyProcessEntry, output: bytes, original_token_count: int | None, + max_output_tokens: int | None, ) -> PtyExecUpdate: exit_code = self._entry_exit_code(entry) live_process_id: int | None = process_id if exit_code is not None: - # collect_pty_output may have handed a partial character back to the deque while - # this looked like it was still running. nothing is going to drain it now, so close - # it here rather than let it leave with the entry - leftover = bytearray() - async with entry.output_lock: - while entry.output_chunks: - leftover.extend(entry.output_chunks.popleft()) - if leftover: - tail, _ = decode_pty_window(leftover, is_final=True) - output += tail.encode("utf-8", errors="replace") + output, original_token_count = await flush_pty_tail( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output=output, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index 9424ae8917..b7efd2a545 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -64,7 +64,7 @@ _settle_mount_transition, _terminate_ambiguous_mount_session, ) -from ....sandbox.session.pty_output import decode_pty_window +from ....sandbox.session.pty_output import close_pty_tail, decode_pty_window from ....sandbox.session.pty_types import ( PTY_PROCESSES_MAX, PTY_PROCESSES_WARNING, @@ -919,6 +919,7 @@ async def pty_exec_start( entry=entry, output=output, original_token_count=original_token_count, + max_output_tokens=max_output_tokens, ) async def pty_write_stdin( @@ -955,6 +956,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + max_output_tokens=max_output_tokens, ) async def pty_terminate_all(self) -> None: @@ -985,9 +987,9 @@ async def _collect_pty_output( max_output_tokens: int | None, ) -> tuple[bytes, int | None]: deadline = time.monotonic() + (yield_time_ms / 1000) - # a character split across two windows starts in the tail the last one held back + # 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) - entry.pending_output = b"" while True: stdout_chunk = await self._read_modal_stream(entry=entry, stream_name="stdout") @@ -1115,17 +1117,18 @@ async def _finalize_pty_update( entry: _ModalPtyProcessEntry, output: bytes, original_token_count: int | None, + max_output_tokens: int | None, ) -> PtyExecUpdate: exit_code = await self._peek_exit_code(entry.process) live_process_id: int | None = process_id if exit_code is not None: - # The collector polled before this and may have seen the process still running, so - # it kept a partial character back for a window that is never going to come. Close - # it here, otherwise those bytes go out with the entry and the output loses them. - if entry.pending_output: - tail, _ = decode_pty_window(entry.pending_output, is_final=True) - entry.pending_output = b"" - output += tail.encode("utf-8", errors="replace") + output, original_token_count = close_pty_tail( + leftover=entry.pending_output, + output=output, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) + entry.pending_output = b"" async with self._pty_lock: removed = self._pty_processes.pop(process_id, None) diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index 8ca4febe85..48d27705ee 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, @@ -1090,6 +1090,7 @@ async def pty_exec_start( entry=entry, output=output, original_token_count=original_token_count, + max_output_tokens=max_output_tokens, ) async def pty_write_stdin( @@ -1139,6 +1140,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + max_output_tokens=max_output_tokens, ) async def pty_terminate_all(self) -> None: @@ -1259,6 +1261,7 @@ async def _finalize_pty_update( entry: _DockerPtyProcessEntry, output: bytes, original_token_count: int | None, + 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 +1270,13 @@ async def _finalize_pty_update( live_process_id: int | None = process_id if exit_code is not None: + output, original_token_count = await flush_pty_tail( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output=output, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) 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..5b9693e1db 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, @@ -410,6 +410,7 @@ def _preexec() -> None: entry=entry, output=output, original_token_count=original_token_count, + max_output_tokens=max_output_tokens, ) async def pty_write_stdin( @@ -455,6 +456,7 @@ async def pty_write_stdin( entry=entry, output=output, original_token_count=original_token_count, + max_output_tokens=max_output_tokens, ) async def pty_terminate_all(self) -> None: @@ -550,11 +552,19 @@ async def _finalize_pty_update( entry: _UnixPtyProcessEntry, output: bytes, original_token_count: int | None, + 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: + output, original_token_count = await flush_pty_tail( + output_chunks=entry.output_chunks, + output_lock=entry.output_lock, + output=output, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) 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 eec2f299f2..1a0949f12c 100644 --- a/src/agents/sandbox/session/pty_output.py +++ b/src/agents/sandbox/session/pty_output.py @@ -31,6 +31,57 @@ def decode_pty_window(data: bytes | bytearray, *, is_final: bool) -> tuple[str, return text, pending +def close_pty_tail( + *, + leftover: bytes | bytearray, + output: bytes, + 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. + + The result is truncated again because the tail is added after the window already applied + ``max_output_tokens``, and the returned count covers the text including it. + """ + if not leftover: + return output, original_token_count + + tail, _ = decode_pty_window(leftover, is_final=True) + if not tail: + return output, original_token_count + + text = output.decode("utf-8", errors="replace") + tail + truncated, recounted = truncate_text_by_tokens(text, max_output_tokens) + return truncated.encode("utf-8", errors="replace"), recounted + + +async def flush_pty_tail( + *, + output_chunks: deque[bytes], + output_lock: asyncio.Lock, + output: bytes, + 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, + original_token_count=original_token_count, + max_output_tokens=max_output_tokens, + ) + + async def collect_pty_output( *, output_chunks: deque[bytes], From f479af8d0c79cb6baf7137312ea8db227831ced6 Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 04:32:34 +0500 Subject: [PATCH 08/15] one contract for closing a tail instead of a patch per backend this is the complexity reset AGENTS.md asks for. i had written the same flush by hand into three finalizers and was about to write a fourth, and every round of review found the next backend i had missed the root cause of all of it is one thing. holding a partial character back puts it in state that outlives the call, and then every path that ends a session has to know to close it. there are seven of those and they each decide the process is gone with their own predicate. unix_local reads process.returncode while its collector reads output_closed, so they genuinely disagree so there is now close_pty_tail for the rule and flush_pty_tail for draining a deque, and all seven finalizers call one of them. that also fixes the token cap, since the tail used to be appended after the window had already truncated modal no longer clears its tail before the decode commits, so a cancelled call keeps it rather than losing the lead byte for the session that survives --- .../extensions/sandbox/blaxel/sandbox.py | 2 +- .../extensions/sandbox/cloudflare/sandbox.py | 2 +- .../extensions/sandbox/daytona/sandbox.py | 2 +- src/agents/extensions/sandbox/e2b/sandbox.py | 2 +- .../extensions/sandbox/modal/sandbox.py | 2 +- src/agents/sandbox/sandboxes/docker.py | 2 +- src/agents/sandbox/sandboxes/unix_local.py | 2 +- tests/extensions/sandbox/test_modal.py | 57 +++++++++++++++++ tests/sandbox/test_pty_output.py | 61 ++++++++++++++++++- 9 files changed, 124 insertions(+), 8 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index 508c3553eb..71c6604aa7 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -983,7 +983,7 @@ async def _finalize_pty_update( entry: _BlaxelPtySessionEntry, output: bytes, original_token_count: int | None, - max_output_tokens: int | None, + max_output_tokens: int | None = None, ) -> PtyExecUpdate: exit_code = entry.exit_code if entry.done else None live_process_id: int | None = process_id diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index fc25bf7fe1..be42e17120 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -1050,7 +1050,7 @@ async def _finalize_pty_update( entry: _CloudflarePtyProcessEntry, output: bytes, original_token_count: int | None, - max_output_tokens: int | None, + 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 diff --git a/src/agents/extensions/sandbox/daytona/sandbox.py b/src/agents/extensions/sandbox/daytona/sandbox.py index aeffea1642..c56023f348 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -856,7 +856,7 @@ async def _finalize_pty_update( entry: _DaytonaPtySessionEntry, output: bytes, original_token_count: int | None, - max_output_tokens: int | None, + max_output_tokens: int | None = None, ) -> PtyExecUpdate: exit_code = entry.exit_code if entry.done else None live_process_id: int | None = process_id diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index ee7323720b..28e4bafee3 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -1248,7 +1248,7 @@ async def _finalize_pty_update( entry: _E2BPtyProcessEntry, output: bytes, original_token_count: int | None, - max_output_tokens: int | None, + max_output_tokens: int | None = None, ) -> PtyExecUpdate: exit_code = self._entry_exit_code(entry) live_process_id: int | None = process_id diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index b7efd2a545..14662f3e28 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -1117,7 +1117,7 @@ async def _finalize_pty_update( entry: _ModalPtyProcessEntry, output: bytes, original_token_count: int | None, - max_output_tokens: int | None, + max_output_tokens: int | None = None, ) -> PtyExecUpdate: exit_code = await self._peek_exit_code(entry.process) live_process_id: int | None = process_id diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index 48d27705ee..b8e9283f45 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -1261,7 +1261,7 @@ async def _finalize_pty_update( entry: _DockerPtyProcessEntry, output: bytes, original_token_count: int | None, - max_output_tokens: int | None, + 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) diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 5b9693e1db..00783fead3 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -552,7 +552,7 @@ async def _finalize_pty_update( entry: _UnixPtyProcessEntry, output: bytes, original_token_count: int | None, - max_output_tokens: int | None, + max_output_tokens: int | None = None, ) -> PtyExecUpdate: exit_code: int | None = entry.process.returncode live_process_id: int | None = process_id diff --git a/tests/extensions/sandbox/test_modal.py b/tests/extensions/sandbox/test_modal.py index 9176d97767..03c876b0aa 100644 --- a/tests/extensions/sandbox/test_modal.py +++ b/tests/extensions/sandbox/test_modal.py @@ -4444,6 +4444,63 @@ def _exec(self, *command: object, **kwargs: object) -> object: await session.pty_terminate_all() +@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, diff --git a/tests/sandbox/test_pty_output.py b/tests/sandbox/test_pty_output.py index 53a7a871bb..d655b8ab20 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 ( + close_pty_tail, + collect_pty_output, + flush_pty_tail, +) @pytest.mark.asyncio @@ -191,3 +195,58 @@ async def test_collect_pty_output_still_holds_a_valid_lead_below_the_surrogates( 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 ", + 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 ", + 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"", + original_token_count=None, + max_output_tokens=0, + ) + uncapped, _ = close_pty_tail( + leftover="\u00e9".encode()[:1], + output=b"", + 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 ", + original_token_count=None, + max_output_tokens=None, + ) + + assert flushed.decode("utf-8") == "hi \ufffd" + assert not chunks From 5e124bce410cc9ce716d21eaf5e2d95624ccbe7f Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 05:17:49 +0500 Subject: [PATCH 09/15] do not await a lock to put the tail back the bytes have already left the deque when the window decodes, so the local pending is the only copy of them. awaiting output_lock to put it back is a cancellation point, and a caller cancelled while a producer holds the lock loses the character while the session carries on to read its continuation appendleft is one synchronous call and deque is documented safe for it. every holder of that lock is either a single append or a drain loop with no await inside, so there is nothing for this to interleave with. the lock is there for the multi step drain, not for one append doing it without the lock removes the cancellation point rather than recovering from it, so there is no window left to get wrong test follows the reported repro. it drains the lead byte, holds the lock across the deadline, cancels, and checks the byte is still there and the next window still gives e acute --- src/agents/sandbox/session/pty_output.py | 8 +++-- tests/sandbox/test_pty_output.py | 42 ++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/src/agents/sandbox/session/pty_output.py b/src/agents/sandbox/session/pty_output.py index 1a0949f12c..7f5be27818 100644 --- a/src/agents/sandbox/session/pty_output.py +++ b/src/agents/sandbox/session/pty_output.py @@ -127,8 +127,12 @@ async def collect_pty_output( # 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: - async with output_lock: - output_chunks.appendleft(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) truncated, original_token_count = truncate_text_by_tokens(text, max_output_tokens) return truncated.encode("utf-8", errors="replace"), original_token_count diff --git a/tests/sandbox/test_pty_output.py b/tests/sandbox/test_pty_output.py index d655b8ab20..51643e3717 100644 --- a/tests/sandbox/test_pty_output.py +++ b/tests/sandbox/test_pty_output.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import contextlib from collections import deque import pytest @@ -250,3 +251,44 @@ async def test_flush_pty_tail_drains_what_the_session_still_holds() -> 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" From 2070572ee89dc983ebf9de67dfba2bfdfaabefdb Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 05:52:54 +0500 Subject: [PATCH 10/15] put a drained window back on cancellation and stop recounting a capped one two things from review a cancelled collection loses everything it had drained, because that only lives in the local buffer. before this pr that was just the window going with the call that asked for it, but now a lead byte a previous window requeued can be in there, so the session that carries on reads the continuation alone and reports a replacement character for output that did arrive. the loop now puts the buffer back before the cancellation goes on close_pty_tail was folding the tail into output the window had already truncated. that truncates it a second time and recounts the shortened display instead of the source. for 100 bytes at a cap of 10 the window reports 25 tokens and this turned it into 11, with the visible text cut down again as well. a window that already hit the cap is now left alone, since the tail is past the cap like everything else that got dropped and the count already says the output is short both tests fail on 5e124bce --- src/agents/sandbox/session/pty_output.py | 62 +++++++++++++-------- tests/sandbox/test_pty_output.py | 71 ++++++++++++++++++++++++ 2 files changed, 110 insertions(+), 23 deletions(-) diff --git a/src/agents/sandbox/session/pty_output.py b/src/agents/sandbox/session/pty_output.py index 7f5be27818..27b13604f4 100644 --- a/src/agents/sandbox/session/pty_output.py +++ b/src/agents/sandbox/session/pty_output.py @@ -45,8 +45,11 @@ def close_pty_tail( collection. Whatever is still waiting then has no later window to complete it, so it is replaced here rather than leaving with the session. - The result is truncated again because the tail is added after the window already applied - ``max_output_tokens``, and the returned count covers the text including it. + A window that already hit ``max_output_tokens`` is left alone. Its output is at the cap, so + the tail sits past it like the rest of what was dropped, and ``original_token_count`` has + already told the caller the output is short. Folding it in there would truncate the text a + second time and recount the shortened display instead of the source, reporting fewer tokens + than the window measured. """ if not leftover: return output, original_token_count @@ -55,9 +58,13 @@ def close_pty_tail( if not tail: return output, original_token_count + if original_token_count is not None: + return output, original_token_count + + # Nothing was truncated, so this really is the whole output and the count still fits it. text = output.decode("utf-8", errors="replace") + tail - truncated, recounted = truncate_text_by_tokens(text, max_output_tokens) - return truncated.encode("utf-8", errors="replace"), recounted + truncated, counted = truncate_text_by_tokens(text, max_output_tokens) + return truncated.encode("utf-8", errors="replace"), counted async def flush_pty_tail( @@ -95,29 +102,38 @@ async def collect_pty_output( 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 - try: - await asyncio.wait_for(output_notify.wait(), timeout=remaining_s) - except asyncio.TimeoutError: - break - output_notify.clear() + 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 # 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 diff --git a/tests/sandbox/test_pty_output.py b/tests/sandbox/test_pty_output.py index 51643e3717..1a44317bef 100644 --- a/tests/sandbox/test_pty_output.py +++ b/tests/sandbox/test_pty_output.py @@ -11,6 +11,7 @@ collect_pty_output, flush_pty_tail, ) +from agents.sandbox.session.pty_types import truncate_text_by_tokens @pytest.mark.asyncio @@ -292,3 +293,73 @@ async def window(yield_time_ms: int) -> bytes: 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_leaves_a_window_that_already_hit_the_cap_alone() -> None: + # the window truncated, so its output is at the cap and the count describes the source. + # folding a tail in here would truncate a second time and recount the shortened display + display, count = truncate_text_by_tokens("a" * 100, 10) + assert count is not None + + output, kept = close_pty_tail( + leftover="\u00e9".encode()[:1], + output=display.encode(), + original_token_count=count, + max_output_tokens=10, + ) + + assert output == display.encode() + assert kept == 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(), + original_token_count=count, + max_output_tokens=10, + ) + + assert output.decode("utf-8") == "hi \ufffd" + assert recounted is None From ff58b7572c35f6a0fc229109fa0134cf66f5dc0a Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 06:16:18 +0500 Subject: [PATCH 11/15] fold a tail into the source text and agree on when a session is finished two more from review, both correct truncation keeps the start and the end of the output, which i had wrong. i was treating a tail found at session end as past the cap and leaving it out, but it belongs in the end that truncation is meant to keep, so that branch could hide the last thing a process said. the window now carries its decoded text before truncation, and a tail is folded into that and truncated once. for the reported case the display goes from ending AAAAAAAA to ending ERROR>>> and the count from 25 to 30, which is what truncating the real source gives unix_local finalised on process.returncode while collection waited for output_closed, which is only set after the process is reaped and the pumps have drained. so the finaliser could remove the session and cancel a pump that still held the rest of a character. docker had the same split through exit_code. both now use the same drained predicate collection uses, which is the shape cloudflare already had the carried text is internal, no public field holds it --- .../extensions/sandbox/blaxel/sandbox.py | 10 +++- .../extensions/sandbox/cloudflare/sandbox.py | 10 +++- .../extensions/sandbox/daytona/sandbox.py | 10 +++- src/agents/extensions/sandbox/e2b/sandbox.py | 10 +++- .../extensions/sandbox/modal/sandbox.py | 12 ++-- src/agents/sandbox/sandboxes/docker.py | 15 +++-- src/agents/sandbox/sandboxes/unix_local.py | 16 +++-- src/agents/sandbox/session/pty_output.py | 31 +++++----- tests/extensions/sandbox/test_blaxel.py | 10 ++-- tests/extensions/sandbox/test_cloudflare.py | 1 + tests/extensions/sandbox/test_modal.py | 1 + tests/sandbox/test_pty_output.py | 58 ++++++++++++++----- 12 files changed, 127 insertions(+), 57 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index 71c6604aa7..477c7c9fcf 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -856,7 +856,7 @@ async def pty_exec_start( ) yield_time_ms = 10_000 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, 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,7 @@ 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, ) @@ -891,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 == "" @@ -904,6 +905,7 @@ 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, ) @@ -966,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, @@ -983,6 +985,7 @@ 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 @@ -993,6 +996,7 @@ async def _finalize_pty_update( 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, ) diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index be42e17120..7ec40bbb78 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -1033,7 +1033,7 @@ async def _collect_pty_output( entry: _CloudflarePtyProcessEntry, 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, @@ -1050,6 +1050,7 @@ 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 @@ -1059,6 +1060,7 @@ async def _finalize_pty_update( 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, ) @@ -1210,7 +1212,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, @@ -1220,6 +1222,7 @@ 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, ) @@ -1244,7 +1247,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, @@ -1258,6 +1261,7 @@ 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, ) diff --git a/src/agents/extensions/sandbox/daytona/sandbox.py b/src/agents/extensions/sandbox/daytona/sandbox.py index c56023f348..8126bc2f6e 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -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,7 @@ 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, ) @@ -833,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 == "" @@ -846,6 +847,7 @@ 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, ) @@ -856,6 +858,7 @@ 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 @@ -866,6 +869,7 @@ async def _finalize_pty_update( 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, ) @@ -897,7 +901,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 28e4bafee3..3147f3ccb1 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -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,7 @@ 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, ) @@ -1088,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 == "" @@ -1101,6 +1102,7 @@ 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, ) @@ -1213,7 +1215,7 @@ async def _collect_pty_output( entry: _E2BPtyProcessEntry, 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, @@ -1248,6 +1250,7 @@ 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) @@ -1258,6 +1261,7 @@ async def _finalize_pty_update( 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, ) diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index 14662f3e28..a2b95604e5 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -909,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, @@ -919,6 +919,7 @@ 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, ) @@ -943,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 == "" @@ -956,6 +957,7 @@ 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, ) @@ -985,7 +987,7 @@ 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) # 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 @@ -1019,7 +1021,7 @@ async def _collect_pty_output( exited = await self._peek_exit_code(entry.process) is not None 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, @@ -1117,6 +1119,7 @@ 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) @@ -1125,6 +1128,7 @@ async def _finalize_pty_update( 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, ) diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index b8e9283f45..3e1b86e18c 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -1080,7 +1080,7 @@ async def pty_exec_start( ) yield_time_ms = 10_000 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, 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,7 @@ 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, ) @@ -1127,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 == "" @@ -1140,6 +1141,7 @@ 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, ) @@ -1244,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, @@ -1261,12 +1263,16 @@ 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: @@ -1274,6 +1280,7 @@ async def _finalize_pty_update( 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, ) diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 00783fead3..1fde2f5669 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -400,7 +400,7 @@ def _preexec() -> None: ) yield_time_ms = 10_000 if yield_time_s is None else int(yield_time_s * 1000) - output, original_token_count = await self._collect_pty_output( + output, original_token_count, 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,7 @@ def _preexec() -> None: entry=entry, output=output, original_token_count=original_token_count, + source_text=source_text, max_output_tokens=max_output_tokens, ) @@ -443,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 == "" @@ -456,6 +457,7 @@ 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, ) @@ -535,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, @@ -552,9 +554,14 @@ 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: @@ -562,6 +569,7 @@ async def _finalize_pty_update( 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, ) diff --git a/src/agents/sandbox/session/pty_output.py b/src/agents/sandbox/session/pty_output.py index 27b13604f4..40433762d9 100644 --- a/src/agents/sandbox/session/pty_output.py +++ b/src/agents/sandbox/session/pty_output.py @@ -35,6 +35,7 @@ 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]: @@ -45,11 +46,11 @@ def close_pty_tail( collection. Whatever is still waiting then has no later window to complete it, so it is replaced here rather than leaving with the session. - A window that already hit ``max_output_tokens`` is left alone. Its output is at the cap, so - the tail sits past it like the rest of what was dropped, and ``original_token_count`` has - already told the caller the output is short. Folding it in there would truncate the text a - second time and recount the shortened display instead of the source, reporting fewer tokens - than the window measured. + ``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 @@ -58,12 +59,7 @@ def close_pty_tail( if not tail: return output, original_token_count - if original_token_count is not None: - return output, original_token_count - - # Nothing was truncated, so this really is the whole output and the count still fits it. - text = output.decode("utf-8", errors="replace") + tail - truncated, counted = truncate_text_by_tokens(text, max_output_tokens) + truncated, counted = truncate_text_by_tokens(source_text + tail, max_output_tokens) return truncated.encode("utf-8", errors="replace"), counted @@ -72,6 +68,7 @@ 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]: @@ -84,6 +81,7 @@ async def flush_pty_tail( return close_pty_tail( leftover=leftover, output=output, + source_text=source_text, original_token_count=original_token_count, max_output_tokens=max_output_tokens, ) @@ -97,8 +95,13 @@ 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() @@ -151,4 +154,4 @@ async def collect_pty_output( output_chunks.appendleft(pending) 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 17470d87c3..a8464dd43e 100644 --- a/tests/extensions/sandbox/test_cloudflare.py +++ b/tests/extensions/sandbox/test_cloudflare.py @@ -1748,6 +1748,7 @@ async def test_cloudflare_pty_finalize_flushes_a_partial_character_left_in_the_d process_id=process_id, entry=entry, output=b"hi ", + source_text="hi ", original_token_count=None, ) diff --git a/tests/extensions/sandbox/test_modal.py b/tests/extensions/sandbox/test_modal.py index 03c876b0aa..cda182c7cd 100644 --- a/tests/extensions/sandbox/test_modal.py +++ b/tests/extensions/sandbox/test_modal.py @@ -4543,6 +4543,7 @@ def _exec(self, *command: object, **kwargs: object) -> object: process_id=7, entry=entry, output=b"hi ", + source_text="hi ", original_token_count=None, ) diff --git a/tests/sandbox/test_pty_output.py b/tests/sandbox/test_pty_output.py index 1a44317bef..a5e207cf60 100644 --- a/tests/sandbox/test_pty_output.py +++ b/tests/sandbox/test_pty_output.py @@ -30,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, @@ -52,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(), @@ -72,7 +72,7 @@ async def _one_window( done: dict[str, bool], ) -> bytes: notify.set() - collected, _ = await collect_pty_output( + collected, _, _ = await collect_pty_output( output_chunks=chunks, output_lock=lock, output_notify=notify, @@ -112,7 +112,7 @@ async def test_collect_pty_output_replaces_a_truncated_character_once_done() -> # 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( + collected, _, _ = await collect_pty_output( output_chunks=chunks, output_lock=asyncio.Lock(), output_notify=asyncio.Event(), @@ -129,7 +129,7 @@ async def test_collect_pty_output_replaces_a_truncated_character_once_done() -> async def test_collect_pty_output_leaves_complete_multibyte_output_alone() -> None: chunks: deque[bytes] = deque(["héllo".encode()]) - collected, _ = await collect_pty_output( + collected, _, _ = await collect_pty_output( output_chunks=chunks, output_lock=asyncio.Lock(), output_notify=asyncio.Event(), @@ -203,6 +203,7 @@ def test_close_pty_tail_replaces_the_leftover_and_leaves_a_clean_session_alone() finished, count = close_pty_tail( leftover="\u00e9".encode()[:1], output=b"hi ", + source_text="hi ", original_token_count=None, max_output_tokens=None, ) @@ -211,6 +212,7 @@ def test_close_pty_tail_replaces_the_leftover_and_leaves_a_clean_session_alone() unchanged, same = close_pty_tail( leftover=b"", output=b"hi ", + source_text="hi ", original_token_count=count, max_output_tokens=None, ) @@ -223,12 +225,14 @@ def test_close_pty_tail_applies_the_token_cap_to_what_it_adds() -> None: 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, ) @@ -246,6 +250,7 @@ async def test_flush_pty_tail_drains_what_the_session_still_holds() -> None: output_chunks=chunks, output_lock=lock, output=b"hi ", + source_text="hi ", original_token_count=None, max_output_tokens=None, ) @@ -267,7 +272,7 @@ async def test_collect_pty_output_keeps_the_tail_when_a_window_is_cancelled() -> async def window(yield_time_ms: int) -> bytes: notify.set() - collected, _ = await collect_pty_output( + collected, _, _ = await collect_pty_output( output_chunks=chunks, output_lock=lock, output_notify=notify, @@ -308,7 +313,7 @@ async def test_collect_pty_output_puts_a_drained_window_back_when_cancelled() -> async def window(yield_time_ms: int) -> bytes: notify.set() - collected, _ = await collect_pty_output( + collected, _, _ = await collect_pty_output( output_chunks=chunks, output_lock=lock, output_notify=notify, @@ -333,21 +338,45 @@ async def window(yield_time_ms: int) -> bytes: assert (await window(60)).decode("utf-8") == "\u00e9" -def test_close_pty_tail_leaves_a_window_that_already_hit_the_cap_alone() -> None: - # the window truncated, so its output is at the cap and the count describes the source. - # folding a tail in here would truncate a second time and recount the shortened display - display, count = truncate_text_by_tokens("a" * 100, 10) +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, kept = close_pty_tail( + 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, ) - assert output == display.encode() - assert kept == count + # 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: @@ -357,6 +386,7 @@ def test_close_pty_tail_still_folds_the_tail_into_an_untruncated_window() -> Non output, recounted = close_pty_tail( leftover="\u00e9".encode()[:1], output=display.encode(), + source_text="hi ", original_token_count=count, max_output_tokens=10, ) From bf682c80cdd387a6522c6e8ab68bd666b92de13e Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 06:23:42 +0500 Subject: [PATCH 12/15] keep modal bytes that were already read off the stream when cancelled a stream item is gone from the stream once it has been read, so anything the window has taken lives only in the local buffer until the decode commits. any await after that read can be cancelled, the stderr read, the exit poll, the sleep, the drains or the final poll not clearing pending_output only saved the carried lead byte. the continuation this call had already swallowed was still lost, and the next window would then pair that lead byte with whatever arrived after it, which is worse than dropping it so the whole buffer goes on the entry before the cancellation goes on, carried tail and new reads together, and the next window picks up where this one stopped test fails on ff58b757 --- .../extensions/sandbox/modal/sandbox.py | 60 +++++++++++------- tests/extensions/sandbox/test_modal.py | 62 +++++++++++++++++++ 2 files changed, 99 insertions(+), 23 deletions(-) diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index a2b95604e5..2ca5fec1bc 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -993,32 +993,46 @@ async def _collect_pty_output( # field is left alone until the decode below commits, so a cancelled call keeps it chunks = bytearray(entry.pending_output) - while True: - stdout_chunk = await self._read_modal_stream(entry=entry, stream_name="stdout") - stderr_chunk = await self._read_modal_stream(entry=entry, stream_name="stderr") - if stdout_chunk: - chunks.extend(stdout_chunk) - if stderr_chunk: - chunks.extend(stderr_chunk) - - if time.monotonic() >= deadline: - break - - exit_code = await self._peek_exit_code(entry.process) - if exit_code is not None: - stdout_chunks = await self._drain_modal_stream(entry=entry, stream_name="stdout") - stderr_chunks = await self._drain_modal_stream(entry=entry, stream_name="stderr") - chunks.extend(stdout_chunks) - chunks.extend(stderr_chunks) - break + 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)) - exited = await self._peek_exit_code(entry.process) is not None + 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, text diff --git a/tests/extensions/sandbox/test_modal.py b/tests/extensions/sandbox/test_modal.py index cda182c7cd..13816e75dc 100644 --- a/tests/extensions/sandbox/test_modal.py +++ b/tests/extensions/sandbox/test_modal.py @@ -4444,6 +4444,68 @@ 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, From 590a28b782235c04710ce44055af2f20a32db679 Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 07:26:18 +0500 Subject: [PATCH 13/15] pin the unix session lifecycle with a blocked pump the predicate fix had no test, so a later cleanup could put returncode back and nothing would say anything. this holds a process that is already reaped while its pump still has the rest of a character, and checks the first finalise leaves the session alive with the lead byte still queued, then releases the pump and checks the next one gives e acute and closes against the old returncode only predicate it fails the way it should, closing the session early with the replacement character as its output --- tests/sandbox/test_unix_local.py | 62 ++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index fb13be3c51..2b295925ce 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -216,6 +216,68 @@ def _unexpected_mkdtemp(*args: object, **kwargs: object) -> str: @pytest.mark.review_optional class TestUnixLocalPty: + @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, From bdef1730a6cc6f65e85089571e14dc1caa15b387 Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 07:42:18 +0500 Subject: [PATCH 14/15] commit the removal before consuming the tail draining empties entry owned state, and it was happening before the finaliser took the session map lock. cancelled while that lock is contended, the session is still registered but its last bytes are already gone, and a later call has nowhere to get them from so the removal goes first now and the drain follows it. cancelled after that point the session is gone anyway, so the bytes go with a session nobody can read from rather than with one that is still listed. same for modal, where reading the tail clears the field test holds the map lock, cancels a finalise, and checks nothing was consumed and the session is still there to finalise properly. it fails on 590a28b7 --- .../extensions/sandbox/blaxel/sandbox.py | 9 ++-- .../extensions/sandbox/cloudflare/sandbox.py | 10 +++-- .../extensions/sandbox/daytona/sandbox.py | 9 ++-- src/agents/extensions/sandbox/e2b/sandbox.py | 10 +++-- .../extensions/sandbox/modal/sandbox.py | 10 +++-- src/agents/sandbox/sandboxes/docker.py | 9 ++-- src/agents/sandbox/sandboxes/unix_local.py | 9 ++-- tests/sandbox/test_unix_local.py | 43 +++++++++++++++++++ 8 files changed, 85 insertions(+), 24 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index 477c7c9fcf..8d89e64331 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -992,6 +992,12 @@ async def _finalize_pty_update( live_process_id: int | None = process_id if entry.done: + async with self._pty_lock: + removed = self._pty_sessions.pop(process_id, None) + self._reserved_pty_process_ids.discard(process_id) + # 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. output, original_token_count = await flush_pty_tail( output_chunks=entry.output_chunks, output_lock=entry.output_lock, @@ -1000,9 +1006,6 @@ async def _finalize_pty_update( original_token_count=original_token_count, max_output_tokens=max_output_tokens, ) - 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) live_process_id = None diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index 7ec40bbb78..34ecb6a5d7 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -1056,6 +1056,12 @@ async def _finalize_pty_update( exit_code = entry.exit_code if entry.output_closed.is_set() else None live_process_id: int | None = process_id if entry.output_closed.is_set(): + async with self._pty_lock: + removed = self._pty_processes.pop(process_id, None) + self._reserved_pty_process_ids.discard(process_id) + # 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. output, original_token_count = await flush_pty_tail( output_chunks=entry.output_chunks, output_lock=entry.output_lock, @@ -1064,10 +1070,6 @@ async def _finalize_pty_update( original_token_count=original_token_count, max_output_tokens=max_output_tokens, ) - - 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) live_process_id = None diff --git a/src/agents/extensions/sandbox/daytona/sandbox.py b/src/agents/extensions/sandbox/daytona/sandbox.py index 8126bc2f6e..5067bf988a 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -865,6 +865,12 @@ async def _finalize_pty_update( live_process_id: int | None = process_id if entry.done: + async with self._pty_lock: + removed = self._pty_sessions.pop(process_id, None) + self._reserved_pty_process_ids.discard(process_id) + # 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. output, original_token_count = await flush_pty_tail( output_chunks=entry.output_chunks, output_lock=entry.output_lock, @@ -873,9 +879,6 @@ async def _finalize_pty_update( original_token_count=original_token_count, max_output_tokens=max_output_tokens, ) - 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) live_process_id = None diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index 3147f3ccb1..13fb8e4c37 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -1257,6 +1257,12 @@ async def _finalize_pty_update( 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) + # 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. output, original_token_count = await flush_pty_tail( output_chunks=entry.output_chunks, output_lock=entry.output_lock, @@ -1265,10 +1271,6 @@ async def _finalize_pty_update( original_token_count=original_token_count, max_output_tokens=max_output_tokens, ) - - 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) live_process_id = None diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index 2ca5fec1bc..0b391e3f4b 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -1139,6 +1139,12 @@ async def _finalize_pty_update( exit_code = await self._peek_exit_code(entry.process) 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) + + # 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. output, original_token_count = close_pty_tail( leftover=entry.pending_output, output=output, @@ -1147,10 +1153,6 @@ async def _finalize_pty_update( max_output_tokens=max_output_tokens, ) entry.pending_output = b"" - - 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) live_process_id = None diff --git a/src/agents/sandbox/sandboxes/docker.py b/src/agents/sandbox/sandboxes/docker.py index 3e1b86e18c..eb06503575 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -1276,6 +1276,12 @@ async def _finalize_pty_update( 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) + # 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. output, original_token_count = await flush_pty_tail( output_chunks=entry.output_chunks, output_lock=entry.output_lock, @@ -1284,9 +1290,6 @@ async def _finalize_pty_update( original_token_count=original_token_count, max_output_tokens=max_output_tokens, ) - 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) live_process_id = None diff --git a/src/agents/sandbox/sandboxes/unix_local.py b/src/agents/sandbox/sandboxes/unix_local.py index 1fde2f5669..247f19129d 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -565,6 +565,12 @@ async def _finalize_pty_update( 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) + # 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. output, original_token_count = await flush_pty_tail( output_chunks=entry.output_chunks, output_lock=entry.output_lock, @@ -573,9 +579,6 @@ async def _finalize_pty_update( original_token_count=original_token_count, max_output_tokens=max_output_tokens, ) - 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) live_process_id = None diff --git a/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index 2b295925ce..87ddb0e6e5 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,48 @@ def _unexpected_mkdtemp(*args: object, **kwargs: object) -> str: @pytest.mark.review_optional class TestUnixLocalPty: + @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, From 78b5c4d019ad7e1abb2863fb31542d5540790547 Mon Sep 17 00:00:00 2001 From: Huzaifa Iftikhar Date: Sat, 29 Aug 2026 07:58:04 +0500 Subject: [PATCH 15/15] close the entry even when the drain is cancelled putting the removal first left a gap i did not think about. once the entry is out of the map nothing can reach it again, pty_terminate_all included, so a cancellation while the drain waits on the output lock skipped _terminate_pty_entry and the entry leaked. on blaxel that is a websocket and an aiohttp session left open the drain is in a try now with the terminate in the finally, so the cleanup happens on both paths. modal has no await between the pop and the terminate so it could not lose it today, but it is written the same way so an await added in front of it later does not start leaking test holds the output lock, cancels the finalise once the removal has gone through, and checks the entry was still terminated. it fails on bdef1730 --- .../extensions/sandbox/blaxel/sandbox.py | 26 +++++----- .../extensions/sandbox/cloudflare/sandbox.py | 26 +++++----- .../extensions/sandbox/daytona/sandbox.py | 26 +++++----- src/agents/extensions/sandbox/e2b/sandbox.py | 26 +++++----- .../extensions/sandbox/modal/sandbox.py | 26 ++++++---- src/agents/sandbox/sandboxes/docker.py | 26 +++++----- src/agents/sandbox/sandboxes/unix_local.py | 26 +++++----- tests/sandbox/test_unix_local.py | 48 +++++++++++++++++++ 8 files changed, 154 insertions(+), 76 deletions(-) diff --git a/src/agents/extensions/sandbox/blaxel/sandbox.py b/src/agents/extensions/sandbox/blaxel/sandbox.py index 8d89e64331..81fe388a63 100644 --- a/src/agents/extensions/sandbox/blaxel/sandbox.py +++ b/src/agents/extensions/sandbox/blaxel/sandbox.py @@ -997,17 +997,21 @@ async def _finalize_pty_update( self._reserved_pty_process_ids.discard(process_id) # 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. - 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, - ) - if removed is not None: - await self._terminate_pty_entry(removed) + # 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 34ecb6a5d7..1c9bc23133 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -1061,17 +1061,21 @@ async def _finalize_pty_update( self._reserved_pty_process_ids.discard(process_id) # 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. - 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, - ) - if removed is not None: - await self._terminate_pty_entry(removed) + # 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/daytona/sandbox.py b/src/agents/extensions/sandbox/daytona/sandbox.py index 5067bf988a..717a14d508 100644 --- a/src/agents/extensions/sandbox/daytona/sandbox.py +++ b/src/agents/extensions/sandbox/daytona/sandbox.py @@ -870,17 +870,21 @@ async def _finalize_pty_update( self._reserved_pty_process_ids.discard(process_id) # 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. - 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, - ) - if removed is not None: - await self._terminate_pty_entry(removed) + # 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/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index 13fb8e4c37..c6f95175e6 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -1262,17 +1262,21 @@ async def _finalize_pty_update( self._reserved_pty_process_ids.discard(process_id) # 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. - 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, - ) - if removed is not None: - await self._terminate_pty_entry(removed) + # 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 0b391e3f4b..9f630c7405 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -1145,16 +1145,22 @@ async def _finalize_pty_update( # 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. - 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"" - if removed is not None: - await self._terminate_pty_entry(removed) + # 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 eb06503575..79dff8cf6e 100644 --- a/src/agents/sandbox/sandboxes/docker.py +++ b/src/agents/sandbox/sandboxes/docker.py @@ -1281,17 +1281,21 @@ async def _finalize_pty_update( self._reserved_pty_process_ids.discard(process_id) # 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. - 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, - ) - if removed is not None: - await self._terminate_pty_entry(removed) + # 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 247f19129d..34af50d879 100644 --- a/src/agents/sandbox/sandboxes/unix_local.py +++ b/src/agents/sandbox/sandboxes/unix_local.py @@ -570,17 +570,21 @@ async def _finalize_pty_update( self._reserved_pty_process_ids.discard(process_id) # 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. - 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, - ) - if removed is not None: - await self._terminate_pty_entry(removed) + # 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/tests/sandbox/test_unix_local.py b/tests/sandbox/test_unix_local.py index 87ddb0e6e5..1c6674c39c 100644 --- a/tests/sandbox/test_unix_local.py +++ b/tests/sandbox/test_unix_local.py @@ -217,6 +217,54 @@ 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,