-
Notifications
You must be signed in to change notification settings - Fork 4.7k
fix(sandbox): keep a UTF-8 sequence split across PTY yield windows whole #4892
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.
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
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| 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 | ||
|
|
@@ -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:])) | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When a command exits after this collector drains a partial character but before its output pump sets AGENTS.md reference: AGENTS.md:L104-L104 Useful? React with 👍 / 👎.
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Fixed in |
||
| 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) | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When a quick-exiting process has buffered output but its stdout/stderr pump has not completed its final read,
process.returncodecan already be non-None; this drains only the bytes currently queued and then_terminate_pty_entry()cancels the pump tasks, so the update returnsprocess_id=Nonewhile 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 foroutput_closedor otherwise finish the pumps before the final drain and removal.AGENTS.md reference: AGENTS.md:L104-L104
Useful? React with 👍 / 👎.