mirror of
https://github.com/usestrix/strix.git
synced 2026-08-21 10:48:59 +02:00
fix(context): page retrieval by byte budget without dropping content
Bounding a retrieval page with a destructive head+tail preview and then advancing the offset past those lines could permanently lose the elided output the store exists to preserve. Instead honour the page byte budget by returning fewer whole lines and advancing the offset only by lines actually returned, so paging forward reconstructs the full output.
This commit is contained in:
+30
-14
@@ -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
|
||||
|
||||
+22
-10
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user