diff --git a/strix/tools/output_store.py b/strix/tools/output_store.py index bbe814c1..d0a7cf9b 100644 --- a/strix/tools/output_store.py +++ b/strix/tools/output_store.py @@ -179,10 +179,15 @@ def bound_and_store(text: str, *, max_lines: int, max_bytes: int) -> str: def read_stored_output(output_id: str, *, offset: int = 0, limit: int = 2_000) -> str: - """Return up to ``limit`` lines of a stored output starting at ``offset``. + """Return a bounded page of a stored output starting at line ``offset``. ``output_id`` must be a token previously returned in a truncation notice; - it is validated to prevent path traversal. + 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. """ if not _OUTPUT_ID_RE.match(output_id): return f"Invalid output_id: {output_id!r}" @@ -191,18 +196,29 @@ 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) - count = min(max(1, limit), _PAGE_MAX_LINES) + 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: - # Stream to the window instead of materialising the whole file per page. - for _ in itertools.islice(handle, start): - pass - window = [line.rstrip("\n") for line in itertools.islice(handle, count)] - remaining = sum(1 for _ in 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 - # Bound the page's byte size with a plain notice (never a spill id) so a - # page of very long lines stays within history without re-triggering spill. - shown = bound_text("\n".join(window), max_lines=_PAGE_MAX_LINES, max_bytes=_PAGE_MAX_BYTES) - if remaining > 0: - shown += f"\n\n[... {remaining} more lines; call read_tool_output(output_id=" - shown += f'"{output_id}", offset={start + len(window)}) to continue ...]' + shown = "\n".join(window) + if has_more: + next_offset = start + len(window) + shown += ( + "\n\n[... more lines; call read_tool_output(" + f'output_id="{output_id}", offset={next_offset}) to continue ...]' + ) return shown diff --git a/tests/test_output_store.py b/tests/test_output_store.py index 78224e7c..19642d04 100644 --- a/tests/test_output_store.py +++ b/tests/test_output_store.py @@ -5,6 +5,7 @@ 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, @@ -114,23 +115,34 @@ def test_read_stored_output_paginates(tmp_path: Path) -> None: assert "offset=10" in page -def test_read_stored_output_page_is_bounded_without_new_spill(tmp_path: Path) -> None: - # A page of very long lines must be byte-capped in place, never re-spilled - # under a fresh output_id (which would make paging loop forever). +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. configure_output_store(tmp_path) - text = "\n".join("z" * 5_000 for _ in range(50)) + lines = [f"{i}-" + "z" * 5_000 for i in range(50)] output_id = re.search( r'output_id="([0-9a-f]{32})"', - bound_and_store(text, max_lines=4, max_bytes=1_000), + bound_and_store("\n".join(lines), max_lines=4, max_bytes=1_000), ) assert output_id is not None + oid = output_id.group(1) - page = read_stored_output(output_id.group(1), offset=0, limit=2_000) + collected: list[str] = [] + offset = 0 + for _ in range(200): # 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 + assert re.findall(r'output_id="([0-9a-f]{32})"', body) in ([], [oid]) + collected.extend(body.split("\n")) + if not hint: + break + match = re.search(r"offset=(\d+)", hint) + assert match is not None + offset = int(match.group(1)) - assert len(page.encode("utf-8")) <= 60 * 1024 - # No brand-new spill id in the returned page. - ids = re.findall(r'output_id="([0-9a-f]{32})"', page) - assert ids == [output_id.group(1)] or ids == [] + assert collected == lines def test_read_stored_output_rejects_traversal(tmp_path: Path) -> None: