diff --git a/strix/tools/output_store.py b/strix/tools/output_store.py index d0a7cf9b..ffa1198f 100644 --- a/strix/tools/output_store.py +++ b/strix/tools/output_store.py @@ -10,7 +10,6 @@ truncated detail is bounded in history but never lost. from __future__ import annotations -import itertools import logging import re import uuid @@ -28,10 +27,10 @@ _SPILL_NOTICE = ( _OUTPUT_ID_RE = re.compile(r"^[0-9a-f]{32}$") _DEFAULT_STORE_DIR = Path.home() / ".strix" / "tool-output" -# Ceilings for a single retrieval page, so a page of very long lines can't -# itself overflow history. Applied without spilling (no new output_id), or -# paging would loop forever. -_PAGE_MAX_LINES = 2_000 +# Byte ceiling for a single retrieval page so retrieval itself can never +# overflow history — even for one very long line. Retrieval pages by byte +# offset (not by line) precisely so an oversized line is split across pages +# instead of returned whole. _PAGE_MAX_BYTES = 50 * 1024 # Single-key holder so the configured path can be swapped per scan without a @@ -178,16 +177,47 @@ def bound_and_store(text: str, *, max_lines: int, max_bytes: int) -> str: return _join(head, tail, notice) -def read_stored_output(output_id: str, *, offset: int = 0, limit: int = 2_000) -> str: - """Return a bounded page of a stored output starting at line ``offset``. +def _trim_incomplete_utf8_tail(chunk: bytes) -> bytes: + """Drop a trailing partial UTF-8 sequence so ``chunk`` decodes cleanly. - ``output_id`` must be a token previously returned in a truncation notice; - it is validated to prevent path traversal. A page is bounded by both a line - count (``limit``, capped at ``_PAGE_MAX_LINES``) and a byte budget - (``_PAGE_MAX_BYTES``) so it can't overflow history. The byte budget is - honoured by returning *fewer whole lines* — never by dropping content from - within the page — so paging forward with the printed ``offset`` hint - reconstructs the full output losslessly. + A fixed byte window can land in the middle of a multi-byte character; the + incomplete tail bytes are dropped here and re-read on the next page (the + caller advances the offset by the *kept* length), so nothing is lost. + """ + # A UTF-8 char is 1-4 bytes; scan back over continuation bytes (0b10xxxxxx) + # to the last lead byte, then keep it only if the whole char is present. + index = len(chunk) - 1 + steps = 0 + while index >= 0 and chunk[index] & 0xC0 == 0x80 and steps < 3: + index -= 1 + steps += 1 + if index < 0: + return chunk + lead = chunk[index] + if lead & 0x80 == 0x00: + expected = 1 + elif lead & 0xE0 == 0xC0: + expected = 2 + elif lead & 0xF0 == 0xE0: + expected = 3 + elif lead & 0xF8 == 0xF0: + expected = 4 + else: + return chunk # invalid lead byte; leave for errors="replace" to handle + if len(chunk) - index < expected: + return chunk[:index] + return chunk + + +def read_stored_output(output_id: str, *, offset: int = 0, limit: int = _PAGE_MAX_BYTES) -> str: + """Return a bounded byte-window of a stored output starting at byte ``offset``. + + ``output_id`` must be a token previously returned in a truncation notice; it + is validated to prevent path traversal. The page is bounded by a UTF-8 byte + budget (``limit``, capped at ``_PAGE_MAX_BYTES``) so it can never overflow + history — even a single very long line is split across pages rather than + returned whole. Paging forward with the printed ``offset`` hint reconstructs + the full output byte-for-byte. """ if not _OUTPUT_ID_RE.match(output_id): return f"Invalid output_id: {output_id!r}" @@ -196,29 +226,24 @@ def read_stored_output(output_id: str, *, offset: int = 0, limit: int = 2_000) - return f"No stored output for output_id={output_id!r} (it may have expired)." start = max(0, offset) - max_lines = min(max(1, limit), _PAGE_MAX_LINES) - window: list[str] = [] - budget = 0 - has_more = False - with path.open(encoding="utf-8") as handle: - # islice skips to ``start`` lazily instead of materialising the file. - for raw in itertools.islice(handle, start, None): - line = raw.rstrip("\n") - size = _byte_len(line) + 1 - # Stop before a line that would breach either budget (always keep at - # least one line). The line we just read is left for the next page, - # so nothing is lost. - if window and (len(window) >= max_lines or budget + size > _PAGE_MAX_BYTES): - has_more = True - break - window.append(line) - budget += size + size = path.stat().st_size + if start >= size: + return "" + # Floor at 4 bytes (the max UTF-8 char length) so a page always makes + # progress past a single multi-byte character. + budget = min(max(4, limit), _PAGE_MAX_BYTES) + with path.open("rb") as handle: + handle.seek(start) + chunk = handle.read(budget) - shown = "\n".join(window) + has_more = start + len(chunk) < size if has_more: - next_offset = start + len(window) + chunk = _trim_incomplete_utf8_tail(chunk) + shown = chunk.decode("utf-8", errors="replace") + if has_more: + next_offset = start + len(chunk) shown += ( - "\n\n[... more lines; call read_tool_output(" + "\n\n[... more; call read_tool_output(" f'output_id="{output_id}", offset={next_offset}) to continue ...]' ) return shown diff --git a/strix/tools/tool_output/tool.py b/strix/tools/tool_output/tool.py index 6feec5ab..2daa704c 100644 --- a/strix/tools/tool_output/tool.py +++ b/strix/tools/tool_output/tool.py @@ -8,7 +8,7 @@ from strix.tools.output_store import read_stored_output @function_tool(timeout=10) -async def read_tool_output(output_id: str, offset: int = 0, limit: int = 2000) -> str: +async def read_tool_output(output_id: str, offset: int = 0, limit: int = 51200) -> str: """Read the full content of an earlier tool result that was truncated. When a tool's output is too large it is trimmed to a head+tail preview in @@ -18,8 +18,9 @@ async def read_tool_output(output_id: str, offset: int = 0, limit: int = 2000) - Args: output_id: The id from the truncation notice (a 32-char hex token). - offset: Zero-based line number to start reading from. - limit: Maximum number of lines to return. Page forward by increasing - ``offset`` using the hint printed at the end of each page. + offset: Zero-based byte offset to start reading from. + limit: Maximum number of bytes to return (capped at 50 KiB). Page + forward by calling again with the ``offset`` printed in the hint at + the end of each page until no hint remains. """ return read_stored_output(output_id, offset=offset, limit=limit) diff --git a/tests/test_output_store.py b/tests/test_output_store.py index 4940dca5..a884ab5a 100644 --- a/tests/test_output_store.py +++ b/tests/test_output_store.py @@ -5,7 +5,6 @@ from __future__ import annotations import re from typing import TYPE_CHECKING -from strix.tools import output_store as _output_store from strix.tools.output_store import ( bound_and_store, bound_text, @@ -105,7 +104,7 @@ def test_bound_and_store_spills_full_output_and_is_retrievable(tmp_path: Path) - output_id = match.group(1) # The full, untruncated output round-trips through the store. - full = read_stored_output(output_id, offset=0, limit=10_000) + full = read_stored_output(output_id, offset=0, limit=1_000_000) assert full.splitlines() == text.splitlines() # A buried line elided from the preview is retrievable. assert "secret-line-500" not in bounded @@ -122,38 +121,68 @@ def test_read_stored_output_paginates(tmp_path: Path) -> None: assert output_id is not None page = read_stored_output(output_id.group(1), offset=0, limit=10) assert page.startswith("0\n1") - assert "more lines" in page + assert "more;" in page assert "offset=10" in page def test_read_stored_output_pages_long_lines_losslessly(tmp_path: Path) -> None: - # A page of very long lines is bounded by returning *fewer whole lines*, - # never by dropping content, so paging forward reconstructs everything and - # never mints a fresh spill id. + # A single line far larger than the page budget must be split across pages + # (never returned whole), and paging forward must reconstruct the output + # byte-for-byte without ever minting a fresh spill id. configure_output_store(tmp_path) - lines = [f"{i}-" + "z" * 5_000 for i in range(50)] + text = "\n".join(f"{i}-" + "z" * 5_000 for i in range(50)) output_id = re.search( r'output_id="([0-9a-f]{32})"', - bound_and_store("\n".join(lines), max_lines=4, max_bytes=1_000), + bound_and_store(text, max_lines=4, max_bytes=1_000), ) assert output_id is not None oid = output_id.group(1) - collected: list[str] = [] + collected = "" offset = 0 - for _ in range(200): # guard against a paging loop + for _ in range(500): # guard against a paging loop page = read_stored_output(oid, offset=offset, limit=2_000) - body, _sep, hint = page.partition("\n\n[... more lines;") - assert len(body.encode("utf-8")) <= _output_store._PAGE_MAX_BYTES + body, _sep, hint = page.partition("\n\n[... more;") + # Every page honours the byte budget, even inside one oversized line. + assert len(body.encode("utf-8")) <= 2_000 assert re.findall(r'output_id="([0-9a-f]{32})"', body) in ([], [oid]) - collected.extend(body.split("\n")) + collected += body if not hint: break match = re.search(r"offset=(\d+)", hint) assert match is not None offset = int(match.group(1)) - assert collected == lines + assert collected == text + + +def test_read_stored_output_pages_multibyte_without_corruption(tmp_path: Path) -> None: + # A byte window can split a 4-byte char; paging must never emit a broken + # (replacement) character and must still reconstruct the text exactly. + configure_output_store(tmp_path) + text = "😀" * 20_000 + output_id = re.search( + r'output_id="([0-9a-f]{32})"', + bound_and_store(text, max_lines=4, max_bytes=1_000), + ) + assert output_id is not None + oid = output_id.group(1) + + collected = "" + offset = 0 + for _ in range(500): # guard against a paging loop + page = read_stored_output(oid, offset=offset, limit=1_002) # not a multiple of 4 + body, _sep, hint = page.partition("\n\n[... more;") + assert "\ufffd" not in body + assert len(body.encode("utf-8")) <= 1_002 + collected += body + if not hint: + break + match = re.search(r"offset=(\d+)", hint) + assert match is not None + offset = int(match.group(1)) + + assert collected == text def test_read_stored_output_rejects_traversal(tmp_path: Path) -> None: