From e3359bb71c56ca41533a8f4837cff949b930c6e2 Mon Sep 17 00:00:00 2001 From: Henry Su Date: Fri, 21 Aug 2026 15:25:30 -0500 Subject: [PATCH 1/2] fix(sandbox): retain PTY output queued at timeout --- .../extensions/sandbox/cloudflare/sandbox.py | 5 ++++ src/agents/extensions/sandbox/e2b/sandbox.py | 5 ++++ .../extensions/sandbox/modal/sandbox.py | 10 ++++++++ src/agents/sandbox/session/pty_output.py | 6 +++++ tests/sandbox/test_pty_output.py | 25 +++++++++++++++++++ 5 files changed, 51 insertions(+) diff --git a/src/agents/extensions/sandbox/cloudflare/sandbox.py b/src/agents/extensions/sandbox/cloudflare/sandbox.py index 8efad06a2e..150da45a36 100644 --- a/src/agents/extensions/sandbox/cloudflare/sandbox.py +++ b/src/agents/extensions/sandbox/cloudflare/sandbox.py @@ -1059,6 +1059,11 @@ async def _collect_pty_output( break entry.output_notify.clear() + # Capture output queued just before the deadline notification timed out. + async with entry.output_lock: + while entry.output_chunks: + output.extend(entry.output_chunks.popleft()) + 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 diff --git a/src/agents/extensions/sandbox/e2b/sandbox.py b/src/agents/extensions/sandbox/e2b/sandbox.py index 389b665c44..4fb7259c21 100644 --- a/src/agents/extensions/sandbox/e2b/sandbox.py +++ b/src/agents/extensions/sandbox/e2b/sandbox.py @@ -1239,6 +1239,11 @@ async def _collect_pty_output( break entry.output_notify.clear() + # Capture output queued just before the deadline notification timed out. + async with entry.output_lock: + while entry.output_chunks: + output.extend(entry.output_chunks.popleft()) + 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 diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index 848b8e12a3..b04ef200c1 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -1010,6 +1010,16 @@ async def _collect_pty_output( break await asyncio.sleep(min(_PTY_POLL_INTERVAL_S, remaining_s)) + # A stream task can complete between the final poll and the deadline check. + # Perform one last non-blocking read so that completed output is not deferred + # to a later pty_write_stdin call. + 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) + text = chunks.decode("utf-8", errors="replace") truncated_text, original_token_count = truncate_text_by_tokens(text, max_output_tokens) return truncated_text.encode("utf-8", errors="replace"), original_token_count diff --git a/src/agents/sandbox/session/pty_output.py b/src/agents/sandbox/session/pty_output.py index 25cbe774e7..77f7fdb860 100644 --- a/src/agents/sandbox/session/pty_output.py +++ b/src/agents/sandbox/session/pty_output.py @@ -45,6 +45,12 @@ async def collect_pty_output( break output_notify.clear() + # A producer can enqueue output after the last top-of-loop drain but before the + # notification wait times out. Capture that tail before returning the snapshot. + async with output_lock: + while output_chunks: + output.extend(output_chunks.popleft()) + 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..9637c1b8c8 100644 --- a/tests/sandbox/test_pty_output.py +++ b/tests/sandbox/test_pty_output.py @@ -57,3 +57,28 @@ def mark_done() -> bool: assert output == b"before done after done" assert original_token_count is None + + +@pytest.mark.asyncio +async def test_collect_pty_output_drains_chunks_queued_when_wait_times_out() -> None: + output_chunks: deque[bytes] = deque() + + class TimeoutAfterQueueing: + async def wait(self) -> None: + output_chunks.append(b"queued at timeout") + raise asyncio.TimeoutError + + def clear(self) -> None: + pass + + output, original_token_count = await collect_pty_output( + output_chunks=output_chunks, + output_lock=asyncio.Lock(), + output_notify=TimeoutAfterQueueing(), # type: ignore[arg-type] + is_done=lambda: False, + yield_time_ms=500, + max_output_tokens=None, + ) + + assert output == b"queued at timeout" + assert original_token_count is None From 08dd6d052c2dcff4d5a72e0a1c3dff435df63d04 Mon Sep 17 00:00:00 2001 From: Henry Su Date: Fri, 21 Aug 2026 23:49:50 -0500 Subject: [PATCH 2/2] fix(sandbox): keep Modal final probe non-blocking --- .../extensions/sandbox/modal/sandbox.py | 18 ++- tests/extensions/sandbox/test_modal.py | 111 ++++++++++++++++++ 2 files changed, 126 insertions(+), 3 deletions(-) diff --git a/src/agents/extensions/sandbox/modal/sandbox.py b/src/agents/extensions/sandbox/modal/sandbox.py index b04ef200c1..b1d7cb2f82 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -1013,8 +1013,16 @@ async def _collect_pty_output( # A stream task can complete between the final poll and the deadline check. # Perform one last non-blocking read so that completed output is not deferred # to a later pty_write_stdin call. - stdout_chunk = await self._read_modal_stream(entry=entry, stream_name="stdout") - stderr_chunk = await self._read_modal_stream(entry=entry, stream_name="stderr") + stdout_chunk = await self._read_modal_stream( + entry=entry, + stream_name="stdout", + allow_new_read=False, + ) + stderr_chunk = await self._read_modal_stream( + entry=entry, + stream_name="stderr", + allow_new_read=False, + ) if stdout_chunk: chunks.extend(stdout_chunk) if stderr_chunk: @@ -1048,6 +1056,7 @@ async def _read_modal_stream( entry: _ModalPtyProcessEntry, stream_name: Literal["stdout", "stderr"], await_pending: bool = False, + allow_new_read: bool = True, ) -> bytes: stream = entry.process.stdout if stream_name == "stdout" else entry.process.stderr if stream is None: @@ -1055,6 +1064,10 @@ async def _read_modal_stream( iter_attr = "stdout_iter" if stream_name == "stdout" else "stderr_iter" task_attr = "stdout_read_task" if stream_name == "stdout" else "stderr_read_task" + task = getattr(entry, task_attr) + if task is None and not allow_new_read: + return b"" + stream_iter = getattr(entry, iter_attr) if stream_iter is None: aiter_method = getattr(stream, "__aiter__", None) @@ -1066,7 +1079,6 @@ async def _read_modal_stream( else: setattr(entry, iter_attr, stream_iter) - task = getattr(entry, task_attr) if task is None and stream_iter is not None: task = asyncio.create_task(stream_iter.__anext__()) setattr(entry, task_attr, task) diff --git a/tests/extensions/sandbox/test_modal.py b/tests/extensions/sandbox/test_modal.py index 44a3fa5c72..2410c0b8d9 100644 --- a/tests/extensions/sandbox/test_modal.py +++ b/tests/extensions/sandbox/test_modal.py @@ -4500,6 +4500,117 @@ def _exec(self, *command: object, **kwargs: object) -> object: assert started.output == b"out-1err-1out-2out-3err-2" +@pytest.mark.asyncio +async def test_modal_pty_final_probe_skips_blocking_read_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + modal_module, _create_calls, _registry_tags = _load_modal_module(monkeypatch) + + class _FakeStream: + def __init__(self) -> None: + self.read_calls = 0 + self.read = _with_aio(self._read) + + def _read(self, _size: int | None = None) -> bytes: + self.read_calls += 1 + return b"" + + class _FakeProcess: + def __init__(self) -> None: + self.stdout = _FakeStream() + self.stderr = _FakeStream() + self.poll = _with_aio(lambda: None) + + process = _FakeProcess() + state = modal_module.ModalSandboxSessionState( + manifest=Manifest(root="/workspace"), + snapshot=modal_module.resolve_snapshot(None, "snapshot"), + app_name="sandbox-tests", + sandbox_id="sb-final-probe", + ) + session = modal_module.ModalSandboxSession.from_state(state) + entry = modal_module._ModalPtyProcessEntry( # noqa: SLF001 + process=process, + tty=True, + last_used=0.0, + ) + + output, original_token_count = await session._collect_pty_output( # noqa: SLF001 + entry=entry, + yield_time_ms=0, + max_output_tokens=None, + ) + + assert output == b"" + assert original_token_count is None + assert process.stdout.read_calls == 1 + assert process.stderr.read_calls == 1 + + +@pytest.mark.asyncio +async def test_modal_pty_final_probe_only_polls_existing_iterator_task( + monkeypatch: pytest.MonkeyPatch, +) -> None: + modal_module, _create_calls, _registry_tags = _load_modal_module(monkeypatch) + + class _FakeStream: + def __init__(self) -> None: + self.anext_calls = 0 + + def __aiter__(self) -> _FakeStream: + return self + + async def __anext__(self) -> bytes: + self.anext_calls += 1 + return b"unexpected" + + class _FakeProcess: + def __init__(self) -> None: + self.stdout = _FakeStream() + self.stderr = _FakeStream() + + process = _FakeProcess() + entry = modal_module._ModalPtyProcessEntry( # noqa: SLF001 + process=process, + tty=True, + last_used=0.0, + ) + session = modal_module.ModalSandboxSession.from_state( + modal_module.ModalSandboxSessionState( + manifest=Manifest(root="/workspace"), + snapshot=modal_module.resolve_snapshot(None, "snapshot"), + app_name="sandbox-tests", + sandbox_id="sb-final-iterator-probe", + ) + ) + + assert ( + await session._read_modal_stream( # noqa: SLF001 + entry=entry, + stream_name="stdout", + allow_new_read=False, + ) + == b"" + ) + assert process.stdout.anext_calls == 0 + + async def _completed_chunk() -> bytes: + return b"ready" + + entry.stdout_read_task = asyncio.create_task(_completed_chunk()) + await entry.stdout_read_task + + assert ( + await session._read_modal_stream( # noqa: SLF001 + entry=entry, + stream_name="stdout", + allow_new_read=False, + ) + == b"ready" + ) + assert process.stdout.anext_calls == 0 + + @pytest.mark.asyncio async def test_modal_pty_start_wraps_startup_failures( monkeypatch: pytest.MonkeyPatch,