diff --git a/src/bub/builtin/shell_manager.py b/src/bub/builtin/shell_manager.py index 9cebde9d..42a4ecfc 100644 --- a/src/bub/builtin/shell_manager.py +++ b/src/bub/builtin/shell_manager.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import codecs import contextlib import uuid from collections.abc import AsyncIterator @@ -228,8 +229,12 @@ async def _drain_stream( shell: ManagedShell, stream: asyncio.StreamReader, ) -> None: + decoder = codecs.getincrementaldecoder("utf-8")(errors="replace") while chunk := await stream.read(4096): - shell.output_chunks.append(chunk.decode("utf-8", errors="replace")) + if text := decoder.decode(chunk): + shell.output_chunks.append(text) + if text := decoder.decode(b"", final=True): + shell.output_chunks.append(text) shell_manager = ShellManager() diff --git a/tests/test_shell_manager.py b/tests/test_shell_manager.py index 234ce8a2..a35b1067 100644 --- a/tests/test_shell_manager.py +++ b/tests/test_shell_manager.py @@ -319,3 +319,29 @@ async def broken_reader() -> None: await asyncio.gather(reader, return_exceptions=True) assert all(shell.returncode is not None for shell in shells) assert manager._shells == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", ["stdout", "stderr"]) +@pytest.mark.parametrize("character", ["é", "中", "😀"]) +async def test_shell_preserves_utf8_across_read_boundaries(stream: str, character: str) -> None: + payload = b"a" * 4095 + character.encode("utf-8") + manager = ShellManager() + async with manager.lifespan(): + shell = await manager.start(cmd=_python(f"import sys; sys.{stream}.buffer.write({payload!r})"), cwd=None) + async with asyncio.timeout(5): + await manager.wait_closed(shell.shell_id) + assert shell.returncode == 0 + assert shell.output == payload.decode("utf-8") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("payload", [b"", b"valid\xfftext", b"incomplete\xf0\x9f"]) +async def test_shell_replaces_invalid_or_incomplete_utf8(payload: bytes) -> None: + manager = ShellManager() + async with manager.lifespan(): + shell = await manager.start(cmd=_python(f"import sys; sys.stdout.buffer.write({payload!r})"), cwd=None) + async with asyncio.timeout(5): + await manager.wait_closed(shell.shell_id) + assert shell.returncode == 0 + assert shell.output == payload.decode("utf-8", errors="replace")