fix(sandbox): retain PTY output queued at timeout

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

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

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,13 +1056,18 @@ 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:
             return b""
 
         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)

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

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,

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