Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
seratch wants to merge 1 commit into
base: main
Choose a base branch
from
+1,069 −157
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 8 additions & 5 deletions src/agents/extensions/sandbox/blaxel/sandbox.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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, 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,
Expand All @@ -866,6 +866,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(
Expand All @@ -890,7 +891,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 == ""
Expand All @@ -903,6 +904,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:
Expand Down Expand Up @@ -964,7 +966,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,
Expand All @@ -981,11 +983,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)
Expand Down
50 changes: 17 additions & 33 deletions src/agents/extensions/sandbox/cloudflare/sandbox.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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)
Expand Down Expand Up @@ -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,
Expand All @@ -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(
Expand All @@ -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,
Expand All @@ -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:
Expand Down
30 changes: 22 additions & 8 deletions src/agents/extensions/sandbox/daytona/sandbox.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -387,6 +387,7 @@ 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
Expand Down Expand Up @@ -755,7 +756,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,
Expand All @@ -765,6 +766,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:
Expand All @@ -777,6 +779,10 @@ async def _run_pty_waiter(self, entry: _DaytonaPtySessionEntry) -> None:
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(
Expand Down Expand Up @@ -804,8 +810,12 @@ async def _run_session_reader(
entry.done = True
except Exception:
pass
if not logs_failed:
# 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.done = True
entry.output_closed.set()
entry.output_notify.set()

async def pty_write_stdin(
Expand All @@ -832,7 +842,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 == ""
Expand All @@ -845,6 +855,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(
Expand All @@ -854,11 +865,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)
Expand Down Expand Up @@ -887,12 +899,13 @@ 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,
should_return=lambda: entry.done,
yield_time_ms=yield_time_ms,
max_output_tokens=max_output_tokens,
)
Expand All @@ -901,7 +914,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:
Expand Down
Loading
Loading