fix(sandbox): settle PTY output before cleanup (#4738)

Co-authored-by: Henry Su <henrysu4707@gmail.com>
Co-authored-by: ayaangazali <ayaangazali.work@gmail.com>

Author: 3 people

src/agents/extensions/sandbox/blaxel/sandbox.py

@@ -864,7 +864,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, output_closed = await self._collect_pty_output(
             entry=entry,
             yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms),
             max_output_tokens=max_output_tokens,
@@ -874,6 +874,7 @@ async def pty_exec_start(
             entry=entry,
             output=output,
             original_token_count=original_token_count,
+            output_closed=output_closed,
         )
 
     async def pty_write_stdin(
@@ -898,7 +899,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, output_closed = 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 == ""
@@ -911,6 +912,7 @@ async def pty_write_stdin(
             entry=entry,
             output=output,
             original_token_count=original_token_count,
+            output_closed=output_closed,
         )
 
     async def pty_terminate_all(self) -> None:
@@ -972,7 +974,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, bool]:
         return await collect_pty_output(
             output_chunks=entry.output_chunks,
             output_lock=entry.output_lock,
@@ -989,11 +991,12 @@ async def _finalize_pty_update(
         entry: _BlaxelPtySessionEntry,
         output: bytes,
         original_token_count: int | None,
+        output_closed: bool,
     ) -> PtyExecUpdate:
-        exit_code = entry.exit_code if entry.done else None
+        exit_code = entry.exit_code if output_closed else None
         live_process_id: int | None = process_id
 
-        if entry.done:
+        if output_closed:
             async with self._pty_lock:
                 removed = self._pty_sessions.pop(process_id, None)
                 self._reserved_pty_process_ids.discard(process_id)

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
@@ -1033,34 +1033,15 @@ async def _collect_pty_output(
         entry: _CloudflarePtyProcessEntry,
         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
+    ) -> tuple[bytes, int | None, bool]:
+        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,
@@ -1069,10 +1050,11 @@ async def _finalize_pty_update(
         entry: _CloudflarePtyProcessEntry,
         output: bytes,
         original_token_count: int | None,
+        output_closed: bool,
     ) -> PtyExecUpdate:
-        exit_code = entry.exit_code if entry.output_closed.is_set() else None
+        exit_code = entry.exit_code if output_closed else None
         live_process_id: int | None = process_id
-        if entry.output_closed.is_set():
+        if output_closed:
             async with self._pty_lock:
                 removed = self._pty_processes.pop(process_id, None)
                 self._reserved_pty_process_ids.discard(process_id)
@@ -1220,7 +1202,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, output_closed = await self._collect_pty_output(
             entry=entry,
             yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms),
             max_output_tokens=max_output_tokens,
@@ -1230,6 +1212,7 @@ async def pty_exec_start(
             entry=entry,
             output=output,
             original_token_count=original_token_count,
+            output_closed=output_closed,
         )
 
     async def pty_write_stdin(
@@ -1253,7 +1236,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, output_closed = await self._collect_pty_output(
             entry=entry,
             yield_time_ms=resolve_pty_write_yield_time_ms(
                 yield_time_ms=yield_time_ms,
@@ -1267,6 +1250,7 @@ async def pty_write_stdin(
             entry=entry,
             output=output,
             original_token_count=original_token_count,
+            output_closed=output_closed,
         )
 
     async def pty_terminate_all(self) -> None:

src/agents/extensions/sandbox/daytona/sandbox.py

@@ -387,8 +387,8 @@ class _DaytonaPtySessionEntry:
     output_chunks: deque[bytes] = field(default_factory=deque)
     output_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
     output_notify: asyncio.Event = field(default_factory=asyncio.Event)
+    output_closed: asyncio.Event = field(default_factory=asyncio.Event)
     last_used: float = field(default_factory=time.monotonic)
-    done: bool = False
     exit_code: int | None = None
     worker_task: asyncio.Task[None] | None = None
 
@@ -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, output_closed = 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,
+            output_closed=output_closed,
         )
 
     async def _run_pty_waiter(self, entry: _DaytonaPtySessionEntry) -> None:
@@ -776,7 +777,10 @@ async def _run_pty_waiter(self, entry: _DaytonaPtySessionEntry) -> None:
         except Exception:
             pass
         finally:
-            entry.done = True
+            # AsyncPtyHandle.wait() completes only after its WebSocket reader exits.
+            # That reader awaits every async on_data callback before it can finish,
+            # so this is Daytona's authoritative output-stream close boundary.
+            entry.output_closed.set()
             entry.output_notify.set()
 
     async def _run_session_reader(
@@ -801,11 +805,13 @@ async def _run_session_reader(
                 cmd = await self._sandbox.process.get_session_command(session_id, cmd_id)
                 if cmd.exit_code is not None:
                     entry.exit_code = int(cmd.exit_code)
-                    entry.done = True
             except Exception:
                 pass
-            if not logs_failed:
-                entry.done = True
+            # Once the log callback stream has returned, or has failed after the
+            # provider reports a final exit code, this worker is the only output
+            # producer and no later callback can append bytes.
+            if not logs_failed or entry.exit_code is not None:
+                entry.output_closed.set()
             entry.output_notify.set()
 
     async def pty_write_stdin(
@@ -832,7 +838,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, output_closed = 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 == ""
@@ -845,6 +851,7 @@ async def pty_write_stdin(
             entry=entry,
             output=output,
             original_token_count=original_token_count,
+            output_closed=output_closed,
         )
 
     async def _finalize_pty_update(
@@ -854,11 +861,12 @@ async def _finalize_pty_update(
         entry: _DaytonaPtySessionEntry,
         output: bytes,
         original_token_count: int | None,
+        output_closed: bool,
     ) -> PtyExecUpdate:
-        exit_code = entry.exit_code if entry.done else None
+        exit_code = entry.exit_code if output_closed else None
         live_process_id: int | None = process_id
 
-        if entry.done:
+        if output_closed:
             async with self._pty_lock:
                 removed = self._pty_sessions.pop(process_id, None)
                 self._reserved_pty_process_ids.discard(process_id)
@@ -887,12 +895,12 @@ async def _collect_pty_output(
         entry: _DaytonaPtySessionEntry,
         yield_time_ms: int,
         max_output_tokens: int | None,
-    ) -> tuple[bytes, int | None]:
+    ) -> tuple[bytes, int | None, bool]:
         return await collect_pty_output(
             output_chunks=entry.output_chunks,
             output_lock=entry.output_lock,
             output_notify=entry.output_notify,
-            is_done=lambda: entry.done,
+            is_done=entry.output_closed.is_set,
             yield_time_ms=yield_time_ms,
             max_output_tokens=max_output_tokens,
         )
@@ -901,7 +909,8 @@ def _prune_pty_sessions_if_needed(self) -> _DaytonaPtySessionEntry | None:
         if len(self._pty_sessions) < PTY_PROCESSES_MAX:
             return None
         meta: list[tuple[int, float, bool]] = [
-            (pid, entry.last_used, entry.done) for pid, entry in self._pty_sessions.items()
+            (pid, entry.last_used, entry.output_closed.is_set())
+            for pid, entry in self._pty_sessions.items()
         ]
         pid = process_id_to_prune_from_meta(meta)
         if pid is None: