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:
Ahmed Allam
2026-07-26 00:16:02 +00:00
committed by Devin AI
parent 90d0dfb4f0
commit 08670033e3
2 changed files with 52 additions and 24 deletions
+30 -14
View File
@@ -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
View File
@@ -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: