Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
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
35 changes: 20 additions & 15 deletions src/agents/sandbox/sandboxes/docker.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,11 @@
from ..session.base_sandbox_session import BaseSandboxSession
from ..session.dependencies import Dependencies
from ..session.manager import Instrumentation
from ..session.pty_output import collect_pty_output
from ..session.pty_output import (
collect_pty_output_bytes,
drain_pty_output_chunks,
finish_pty_output,
)
from ..session.pty_types import (
PTY_PROCESSES_MAX,
PTY_PROCESSES_WARNING,
Expand Down Expand Up @@ -1080,16 +1084,15 @@ 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(
raw_output = await self._collect_pty_output(
entry=entry,
yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms),
max_output_tokens=max_output_tokens,
)
return await self._finalize_pty_update(
process_id=process_id,
entry=entry,
output=output,
original_token_count=original_token_count,
raw_output=raw_output,
max_output_tokens=max_output_tokens,
)

async def pty_write_stdin(
Expand Down Expand Up @@ -1126,19 +1129,18 @@ 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(
raw_output = 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 == ""
),
max_output_tokens=max_output_tokens,
)
entry.last_used = time.monotonic()
return await self._finalize_pty_update(
process_id=session_id,
entry=entry,
output=output,
original_token_count=original_token_count,
raw_output=raw_output,
max_output_tokens=max_output_tokens,
)

async def pty_terminate_all(self) -> None:
Expand Down Expand Up @@ -1241,24 +1243,22 @@ async def _collect_pty_output(
*,
entry: _DockerPtyProcessEntry,
yield_time_ms: int,
max_output_tokens: int | None,
) -> tuple[bytes, int | None]:
return await collect_pty_output(
) -> bytes:
return await collect_pty_output_bytes(
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,
*,
process_id: int,
entry: _DockerPtyProcessEntry,
output: bytes,
original_token_count: int | None,
raw_output: bytes,
max_output_tokens: int | None,
) -> PtyExecUpdate:
if entry.output_closed.is_set() and entry.exit_code is None:
await self._refresh_pty_exit_code(entry)
Expand All @@ -1267,13 +1267,18 @@ async def _finalize_pty_update(
live_process_id: int | None = process_id

if exit_code is not None:
# The collector may have held back a partial UTF-8 sequence for a later poll;
# there is none once the entry is removed, so take whatever is still queued
# before the token limit is applied to the whole update.
raw_output += await drain_pty_output_chunks(entry.output_chunks, entry.output_lock)
async with self._pty_lock:
removed = self._pty_processes.pop(process_id, None)
self._reserved_pty_process_ids.discard(process_id)
if removed is not None:
await self._terminate_pty_entry(removed)
live_process_id = None

output, original_token_count = finish_pty_output(raw_output, max_output_tokens)
return PtyExecUpdate(
process_id=live_process_id,
output=output,
Expand Down
35 changes: 20 additions & 15 deletions src/agents/sandbox/sandboxes/unix_local.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,11 @@
from ..session.base_sandbox_session import BaseSandboxSession
from ..session.dependencies import Dependencies
from ..session.manager import Instrumentation
from ..session.pty_output import collect_pty_output
from ..session.pty_output import (
collect_pty_output_bytes,
drain_pty_output_chunks,
finish_pty_output,
)
from ..session.pty_types import (
PTY_PROCESSES_MAX,
PTY_PROCESSES_WARNING,
Expand Down Expand Up @@ -400,16 +404,15 @@ def _preexec() -> 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(
raw_output = await self._collect_pty_output(
entry=entry,
yield_time_ms=clamp_pty_yield_time_ms(yield_time_ms),
max_output_tokens=max_output_tokens,
)
return await self._finalize_pty_update(
process_id=process_id,
entry=entry,
output=output,
original_token_count=original_token_count,
raw_output=raw_output,
max_output_tokens=max_output_tokens,
)

async def pty_write_stdin(
Expand Down Expand Up @@ -442,19 +445,18 @@ 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(
raw_output = 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 == ""
),
max_output_tokens=max_output_tokens,
)
entry.last_used = time.monotonic()
return await self._finalize_pty_update(
process_id=session_id,
entry=entry,
output=output,
original_token_count=original_token_count,
raw_output=raw_output,
max_output_tokens=max_output_tokens,
)

async def pty_terminate_all(self) -> None:
Expand Down Expand Up @@ -532,36 +534,39 @@ async def _collect_pty_output(
*,
entry: _UnixPtyProcessEntry,
yield_time_ms: int,
max_output_tokens: int | None,
) -> tuple[bytes, int | None]:
return await collect_pty_output(
) -> bytes:
return await collect_pty_output_bytes(
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,
*,
process_id: int,
entry: _UnixPtyProcessEntry,
output: bytes,
original_token_count: int | None,
raw_output: bytes,
max_output_tokens: int | None,
) -> PtyExecUpdate:
exit_code: int | None = entry.process.returncode
live_process_id: int | None = process_id

if exit_code is not None:
# The collector may have held back a partial UTF-8 sequence for a later poll;
# there is none once the entry is removed, so take whatever is still queued
# before the token limit is applied to the whole update.
raw_output += await drain_pty_output_chunks(entry.output_chunks, entry.output_lock)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Wait for output pumps before draining exited PTYs

When a quick-exiting process has buffered output but its stdout/stderr pump has not completed its final read, process.returncode can already be non-None; this drains only the bytes currently queued and then _terminate_pty_entry() cancels the pump tasks, so the update returns process_id=None while permanently omitting the unread bytes. Fresh evidence after the prior fix is that the new drain still occurs before pump completion; the Docker path has the same issue when a reader-thread append is scheduled but not queued before the drain. Wait for output_closed or otherwise finish the pumps before the final drain and removal.

AGENTS.md reference: AGENTS.md:L104-L104

Useful? React with 👍 / 👎.

async with self._pty_lock:
removed = self._pty_processes.pop(process_id, None)
self._reserved_pty_process_ids.discard(process_id)
if removed is not None:
await self._terminate_pty_entry(removed)
live_process_id = None

output, original_token_count = finish_pty_output(raw_output, max_output_tokens)
return PtyExecUpdate(
process_id=live_process_id,
output=output,
Expand Down
72 changes: 66 additions & 6 deletions src/agents/sandbox/session/pty_output.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import asyncio
import codecs
import time
from collections import deque
from collections.abc import Callable
Expand All @@ -18,33 +19,92 @@ async def collect_pty_output(
max_output_tokens: int | None,
) -> tuple[bytes, int | None]:
"""Collect and truncate PTY output until the deadline or provider completion."""
raw_output = await collect_pty_output_bytes(
output_chunks=output_chunks,
output_lock=output_lock,
output_notify=output_notify,
is_done=is_done,
yield_time_ms=yield_time_ms,
)
return finish_pty_output(raw_output, max_output_tokens)


async def collect_pty_output_bytes(
*,
output_chunks: deque[bytes],
output_lock: asyncio.Lock,
output_notify: asyncio.Event,
is_done: Callable[[], bool],
yield_time_ms: int,
) -> bytes:
"""Collect raw PTY output until the deadline or provider completion.

A multibyte UTF-8 sequence can straddle two yield windows (the producer wrote part of
it before the deadline). While the provider is still running, the incomplete tail is
held back in `output_chunks` for the next collection instead of being emitted as
replacement characters on both sides.
"""
deadline = time.monotonic() + (yield_time_ms / 1000)
output = bytearray()

while True:
async with output_lock:
while output_chunks:
output.extend(output_chunks.popleft())

if time.monotonic() >= deadline:
break

if is_done():
async with output_lock:
while output_chunks:
output.extend(output_chunks.popleft())
break

remaining_s = deadline - time.monotonic()
if remaining_s <= 0:
break

try:
await asyncio.wait_for(output_notify.wait(), timeout=remaining_s)
except asyncio.TimeoutError:
break
output_notify.clear()
tail_length = incomplete_utf8_tail_length(output)
if tail_length:
async with output_lock:
# Re-check under the lock: once the provider is done nothing else will arrive,
# so the tail must be flushed now rather than parked in the queue.
if not is_done():
output_chunks.appendleft(bytes(output[-tail_length:]))

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Flush deferred bytes before removing an exited PTY

When a command exits after this collector drains a partial character but before its output pump sets output_closed, is_done() remains false and the drained bytes are placed back in the queue. Both UnixLocal and Docker can already observe the exit code in _finalize_pty_update() and immediately remove and terminate that entry, making the deferred bytes unreachable and returning incomplete output. This occurs on the supported quick-exit variant of the split-read scenario, so finalization must wait for output closure or explicitly flush the deferred bytes before removing the entry.

AGENTS.md reference: AGENTS.md:L104-L104

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 7e0211a4: _finalize_pty_update() in both UnixLocal and Docker now drains whatever is still queued (via the new drain_pty_output_chunks()) whenever it observes an exit code, and appends it to the update before the entry is removed, so bytes a collection held back are never stranded. TestUnixLocalPty::test_finalize_pty_update_flushes_bytes_held_back_when_the_process_exited reproduces the window (process exited, output not yet marked closed, the split sequence sitting in the queue) and fails on the previous revision.

del output[-tail_length:]
return bytes(output)


text = output.decode("utf-8", errors="replace")
def finish_pty_output(raw_output: bytes, max_output_tokens: int | None) -> tuple[bytes, int | None]:
"""Decode collected PTY bytes and apply the token limit once."""
text = raw_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


def incomplete_utf8_tail_length(data: bytes | bytearray) -> int:
"""Return how many trailing bytes form a valid but not yet complete UTF-8 sequence.

Python's incremental decoder decides what counts as a valid prefix, so invalid
leaders (``0xC0``, ``0xC1``, ``0xF5`` and up) and ill-formed second bytes (overlong
forms, code points past U+10FFFF) are not held back: they decode to replacement
characters immediately, as before. A pending sequence is at most three bytes long, so
only the tail needs to be inspected.
"""
tail = bytes(data[-3:])
if not tail:
return 0
decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
decoder.decode(tail, final=False)
pending, _ = decoder.getstate()
return len(pending)


async def drain_pty_output_chunks(output_chunks: deque[bytes], output_lock: asyncio.Lock) -> bytes:
"""Take every queued chunk, including bytes a collection held back."""
output = bytearray()
async with output_lock:
while output_chunks:
output.extend(output_chunks.popleft())
return bytes(output)
Loading