perf(runtime): chunk captured non-ASCII output on UTF-8 byte boundaries

src/sqlbuild/runtime/output_capture/_helpers/text.py

@@ -4,6 +4,11 @@
 
 import re
 
+from sqlbuild.runtime.output_capture.constants import (
+    UTF8_CONTINUATION_BYTE_MASK,
+    UTF8_CONTINUATION_BYTE_TAG,
+)
+
 _ANSI_ESCAPE: re.Pattern[str] = re.compile(
     r"(?:\x1b\][^\x07\x1b]*(?:\x07|\x1b\\)|\x1b\[[0-?]*[ -/]*[@-~])"
 )
@@ -22,6 +27,48 @@ def chunk_text(*, text: str, max_bytes: int) -> tuple[str, ...]:
         if len(text) <= max_bytes:
             return (text,)
         return tuple(text[index : index + max_bytes] for index in range(0, len(text), max_bytes))
+    try:
+        encoded: bytes = text.encode("utf-8")
+    except UnicodeEncodeError:
+        return _chunk_text_by_character(text=text, max_bytes=max_bytes)
+    if len(encoded) <= max_bytes:
+        return (text,)
+    return _chunk_utf8(encoded=encoded, max_bytes=max_bytes)
+
+
+def _chunk_utf8(*, encoded: bytes, max_bytes: int) -> tuple[str, ...]:
+    """Split valid UTF-8 greedily at the last code point boundary within each limit."""
+
+    chunks: list[str] = []
+    total: int = len(encoded)
+    start: int = 0
+    while start < total:
+        end: int = start + max_bytes
+        if end >= total:
+            end = total
+        else:
+            while (
+                end > start
+                and encoded[end] & UTF8_CONTINUATION_BYTE_MASK == UTF8_CONTINUATION_BYTE_TAG
+            ):
+                end -= 1
+            if end <= start:
+                end = start + 1
+                while (
+                    end < total
+                    and encoded[end] & UTF8_CONTINUATION_BYTE_MASK == UTF8_CONTINUATION_BYTE_TAG
+                ):
+                    end += 1
+        chunks.append(encoded[start:end].decode("utf-8"))
+        start = end
+    if not chunks:
+        chunks.append("")
+    return tuple(chunks)
+
+
+def _chunk_text_by_character(*, text: str, max_bytes: int) -> tuple[str, ...]:
+    """Split text containing lone surrogates, which surrogateescape maps to single bytes."""
+
     if len(text) <= max_bytes and len(text.encode("utf-8", "surrogateescape")) <= max_bytes:
         return (text,)
 

src/sqlbuild/runtime/output_capture/constants.py

@@ -8,6 +8,8 @@
 DEFAULT_OUTPUT_QUEUE_CAPACITY: int = 1024
 DEFAULT_OUTPUT_SHUTDOWN_TIMEOUT_SECONDS: float = 2.0
 MIN_OUTPUT_RECORD_BYTES: int = 4
+UTF8_CONTINUATION_BYTE_MASK: int = 0xC0
+UTF8_CONTINUATION_BYTE_TAG: int = 0x80
 INVOCATION_CONTEXT_ENV: str = "SQLBUILD_INVOCATION_CONTEXT_JSON"
 MAX_INVOCATION_CONTEXT_BYTES: int = 16 * 1024
 MAX_INVOCATION_CONTEXT_DEPTH: int = 8

tests/unit/src/sqlbuild/runtime/output_capture/_helpers/_test_types.py

@@ -0,0 +1,48 @@
+"""Parameterized text chunking test cases."""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+type ChunkTextFunction = Callable[..., tuple[str, ...]]
+
+
+@dataclass(frozen=True)
+class ChunkTextTestCase:
+    """One exact chunk_text expectation."""
+
+    description: str
+    text: str
+    max_bytes: int
+    expected_chunks: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class ChunkTextErrorTestCase:
+    """One input the stored chunk contract rejects."""
+
+    description: str
+    text: str
+    max_bytes: int
+    expected_error: type[Exception]
+
+
+@dataclass(frozen=True)
+class ChunkTextOracleTestCase:
+    """One seeded random comparison against the reference implementation."""
+
+    description: str
+    seed: int
+    max_bytes: int
+    count: int
+    expected_mismatch_count: int
+
+
+@dataclass(frozen=True)
+class ChunkTextPerformanceTestCase:
+    """One large payload chunking time budget."""
+
+    description: str
+    unit: str
+    payload_bytes: int
+    max_bytes: int
+    expected_max_seconds: float

tests/unit/src/sqlbuild/runtime/output_capture/_helpers/conftest.py

@@ -0,0 +1,38 @@
+"""Shared fixtures for output capture text chunking tests."""
+
+import pytest
+
+from tests.unit.src.sqlbuild.runtime.output_capture._helpers._test_types import ChunkTextFunction
+
+
+def _reference_chunk_text(*, text: str, max_bytes: int) -> tuple[str, ...]:
+    """Copy of the original character-at-a-time chunker that defines the stored contract."""
+
+    if text.isascii() and max_bytes > 0:
+        if len(text) <= max_bytes:
+            return (text,)
+        return tuple(text[index : index + max_bytes] for index in range(0, len(text), max_bytes))
+    if len(text) <= max_bytes and len(text.encode("utf-8", "surrogateescape")) <= max_bytes:
+        return (text,)
+
+    chunks: list[str] = []
+    current: list[str] = []
+    current_bytes: int = 0
+    for character in text:
+        character_bytes: int = len(character.encode("utf-8", "surrogateescape"))
+        if current and current_bytes + character_bytes > max_bytes:
+            chunks.append("".join(current))
+            current = []
+            current_bytes = 0
+        current.append(character)
+        current_bytes += character_bytes
+    if current or not chunks:
+        chunks.append("".join(current))
+    return tuple(chunks)
+
+
+@pytest.fixture
+def reference_chunk_text() -> ChunkTextFunction:
+    """Return the reference oracle for chunk boundaries and text."""
+
+    return _reference_chunk_text

tests/unit/src/sqlbuild/runtime/output_capture/_helpers/helpers.py

@@ -0,0 +1,43 @@
+"""Random inputs for output capture text chunking tests."""
+
+import random
+
+_ALPHABET: tuple[str, ...] = (
+    "a",
+    "Z",
+    " ",
+    "\n",
+    "é",
+    "ß",
+    "Ω",
+    "订",
+    "€",
+    "✓",
+    "📦",
+    "🙂",
+    "𝔸",
+    "e\u0301",
+    "\u0308",
+)
+
+
+def random_mixed_width_texts(*, seed: int, max_bytes: int, count: int) -> tuple[str, ...]:
+    """Build texts whose encoded lengths land around multiples of the chunk limit."""
+
+    generator: random.Random = random.Random(seed)
+    return tuple(
+        _random_mixed_width_text(generator=generator, max_bytes=max_bytes) for _ in range(count)
+    )
+
+
+def _random_mixed_width_text(*, generator: random.Random, max_bytes: int) -> str:
+    target_bytes: int = max(
+        0, max_bytes * generator.choice((0, 1, 1, 2, 3)) + generator.randint(-5, 5)
+    )
+    parts: list[str] = []
+    size: int = 0
+    while size < target_bytes:
+        piece: str = generator.choice(_ALPHABET)
+        parts.append(piece)
+        size += len(piece.encode("utf-8"))
+    return "".join(parts)

tests/unit/src/sqlbuild/runtime/output_capture/_helpers/test_text.py

@@ -0,0 +1,128 @@
+"""Tests for deterministic output capture text chunking."""
+
+import time
+
+import pytest
+
+from sqlbuild.runtime.output_capture._helpers.text import chunk_text
+from tests.unit.src.sqlbuild.runtime.output_capture._helpers._test_types import (
+    ChunkTextErrorTestCase,
+    ChunkTextFunction,
+    ChunkTextOracleTestCase,
+    ChunkTextPerformanceTestCase,
+    ChunkTextTestCase,
+)
+from tests.unit.src.sqlbuild.runtime.output_capture._helpers.helpers import (
+    random_mixed_width_texts,
+)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    (
+        ChunkTextTestCase("two_byte_straddles_limit", "aé", 2, ("a", "é")),
+        ChunkTextTestCase("three_byte_straddles_limit", "ab订c", 4, ("ab", "订c")),
+        ChunkTextTestCase("four_byte_straddles_limit", "abc📦", 4, ("abc", "📦")),
+        ChunkTextTestCase("four_byte_exact_limit", "📦📦", 4, ("📦", "📦")),
+        ChunkTextTestCase("combining_mark_split_from_base", "ée\u0301", 4, ("ée", "\u0301")),
+        ChunkTextTestCase("character_wider_than_limit", "a📦b", 2, ("a", "📦", "b")),
+        ChunkTextTestCase("each_character_when_limit_zero", "a订", 0, ("a", "订")),
+        ChunkTextTestCase("empty_text_negative_limit", "", -1, ("",)),
+        ChunkTextTestCase("escaped_surrogate_counts_one_byte", "é\udc80a", 3, ("é\udc80", "a")),
+    ),
+    ids=lambda case: case.description,
+)
+def test_given_multibyte_boundaries_when_chunking_then_code_points_stay_whole(
+    test_case: ChunkTextTestCase,
+    reference_chunk_text: ChunkTextFunction,
+) -> None:
+    result: tuple[str, ...] = chunk_text(text=test_case.text, max_bytes=test_case.max_bytes)
+
+    assert result == test_case.expected_chunks
+    assert result == reference_chunk_text(text=test_case.text, max_bytes=test_case.max_bytes)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    (ChunkTextErrorTestCase("unescapable_surrogate", "é\ud800", 2, UnicodeEncodeError),),
+    ids=lambda case: case.description,
+)
+def test_given_unescapable_surrogate_when_chunking_then_raises_like_reference(
+    test_case: ChunkTextErrorTestCase,
+    reference_chunk_text: ChunkTextFunction,
+) -> None:
+    with pytest.raises(test_case.expected_error):
+        _ = reference_chunk_text(text=test_case.text, max_bytes=test_case.max_bytes)
+    with pytest.raises(test_case.expected_error):
+        _ = chunk_text(text=test_case.text, max_bytes=test_case.max_bytes)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        ChunkTextOracleTestCase(
+            f"limit_{max_bytes}",
+            seed=max_bytes,
+            max_bytes=max_bytes,
+            count=400,
+            expected_mismatch_count=0,
+        )
+        for max_bytes in (1, 2, 3, 4, 5, 6, 7, 8, 13, 64, 257)
+    ],
+    ids=lambda case: case.description,
+)
+def test_given_random_mixed_width_text_when_chunking_then_matches_reference(
+    test_case: ChunkTextOracleTestCase,
+    reference_chunk_text: ChunkTextFunction,
+) -> None:
+    texts: tuple[str, ...] = random_mixed_width_texts(
+        seed=test_case.seed, max_bytes=test_case.max_bytes, count=test_case.count
+    )
+
+    actual: tuple[tuple[str, ...], ...] = tuple(
+        chunk_text(text=text, max_bytes=test_case.max_bytes) for text in texts
+    )
+
+    expected: tuple[tuple[str, ...], ...] = tuple(
+        reference_chunk_text(text=text, max_bytes=test_case.max_bytes) for text in texts
+    )
+    mismatch_count: int = sum(
+        result != reference for result, reference in zip(actual, expected, strict=True)
+    )
+    assert mismatch_count == test_case.expected_mismatch_count
+    assert tuple("".join(chunks) for chunks in actual) == texts
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    (
+        ChunkTextPerformanceTestCase(
+            "non_ascii_ten_megabytes",
+            unit="order ünïcode 订单 📦 e\u0301 ",
+            payload_bytes=10_000_000,
+            max_bytes=64 * 1024,
+            expected_max_seconds=0.1,
+        ),
+    ),
+    ids=lambda case: case.description,
+)
+def test_given_large_non_ascii_payload_when_chunking_then_completes_within_budget(
+    test_case: ChunkTextPerformanceTestCase,
+) -> None:
+    unit_bytes: int = len(test_case.unit.encode("utf-8"))
+    text: str = test_case.unit * (test_case.payload_bytes // unit_bytes + 1)
+    timings: list[float] = []
+    result: tuple[str, ...] = ()
+
+    for _ in range(3):
+        started: float = time.perf_counter()
+        result = chunk_text(text=text, max_bytes=test_case.max_bytes)
+        timings.append(time.perf_counter() - started)
+
+    assert "".join(result) == text
+    assert all(len(chunk.encode("utf-8")) <= test_case.max_bytes for chunk in result)
+    assert min(timings) < test_case.expected_max_seconds
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-vv"])