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..b1d7cb2f82 100644 --- a/src/agents/extensions/sandbox/modal/sandbox.py +++ b/src/agents/extensions/sandbox/modal/sandbox.py @@ -1010,6 +1010,24 @@ 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", + 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: + 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 @@ -1038,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: @@ -1045,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) @@ -1056,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/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/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, 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