mirror of
https://github.com/usestrix/strix.git
synced 2026-08-24 20:02:39 +02:00
fix(context): page stored output by byte offset so no single line overflows
The line-based pager always kept at least one whole line, so a stored line larger than the page byte ceiling was returned in full and could overflow history (retrieval bypasses the result-bounding wrapper). Page by byte offset instead: every page is bounded by the byte budget even inside one oversized line, a partial UTF-8 char at the window edge is dropped and re-read on the next page, and paging forward reconstructs the output byte-for-byte.
This commit is contained in:
+59
-34
@@ -10,7 +10,6 @@ truncated detail is bounded in history but never lost.
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import itertools
|
|
||||||
import logging
|
import logging
|
||||||
import re
|
import re
|
||||||
import uuid
|
import uuid
|
||||||
@@ -28,10 +27,10 @@ _SPILL_NOTICE = (
|
|||||||
_OUTPUT_ID_RE = re.compile(r"^[0-9a-f]{32}$")
|
_OUTPUT_ID_RE = re.compile(r"^[0-9a-f]{32}$")
|
||||||
_DEFAULT_STORE_DIR = Path.home() / ".strix" / "tool-output"
|
_DEFAULT_STORE_DIR = Path.home() / ".strix" / "tool-output"
|
||||||
|
|
||||||
# Ceilings for a single retrieval page, so a page of very long lines can't
|
# Byte ceiling for a single retrieval page so retrieval itself can never
|
||||||
# itself overflow history. Applied without spilling (no new output_id), or
|
# overflow history — even for one very long line. Retrieval pages by byte
|
||||||
# paging would loop forever.
|
# offset (not by line) precisely so an oversized line is split across pages
|
||||||
_PAGE_MAX_LINES = 2_000
|
# instead of returned whole.
|
||||||
_PAGE_MAX_BYTES = 50 * 1024
|
_PAGE_MAX_BYTES = 50 * 1024
|
||||||
|
|
||||||
# Single-key holder so the configured path can be swapped per scan without a
|
# 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)
|
return _join(head, tail, notice)
|
||||||
|
|
||||||
|
|
||||||
def read_stored_output(output_id: str, *, offset: int = 0, limit: int = 2_000) -> str:
|
def _trim_incomplete_utf8_tail(chunk: bytes) -> bytes:
|
||||||
"""Return a bounded page of a stored output starting at line ``offset``.
|
"""Drop a trailing partial UTF-8 sequence so ``chunk`` decodes cleanly.
|
||||||
|
|
||||||
``output_id`` must be a token previously returned in a truncation notice;
|
A fixed byte window can land in the middle of a multi-byte character; the
|
||||||
it is validated to prevent path traversal. A page is bounded by both a line
|
incomplete tail bytes are dropped here and re-read on the next page (the
|
||||||
count (``limit``, capped at ``_PAGE_MAX_LINES``) and a byte budget
|
caller advances the offset by the *kept* length), so nothing is lost.
|
||||||
(``_PAGE_MAX_BYTES``) so it can't overflow history. The byte budget is
|
"""
|
||||||
honoured by returning *fewer whole lines* — never by dropping content from
|
# A UTF-8 char is 1-4 bytes; scan back over continuation bytes (0b10xxxxxx)
|
||||||
within the page — so paging forward with the printed ``offset`` hint
|
# to the last lead byte, then keep it only if the whole char is present.
|
||||||
reconstructs the full output losslessly.
|
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):
|
if not _OUTPUT_ID_RE.match(output_id):
|
||||||
return f"Invalid output_id: {output_id!r}"
|
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)."
|
return f"No stored output for output_id={output_id!r} (it may have expired)."
|
||||||
|
|
||||||
start = max(0, offset)
|
start = max(0, offset)
|
||||||
max_lines = min(max(1, limit), _PAGE_MAX_LINES)
|
size = path.stat().st_size
|
||||||
window: list[str] = []
|
if start >= size:
|
||||||
budget = 0
|
return ""
|
||||||
has_more = False
|
# Floor at 4 bytes (the max UTF-8 char length) so a page always makes
|
||||||
with path.open(encoding="utf-8") as handle:
|
# progress past a single multi-byte character.
|
||||||
# islice skips to ``start`` lazily instead of materialising the file.
|
budget = min(max(4, limit), _PAGE_MAX_BYTES)
|
||||||
for raw in itertools.islice(handle, start, None):
|
with path.open("rb") as handle:
|
||||||
line = raw.rstrip("\n")
|
handle.seek(start)
|
||||||
size = _byte_len(line) + 1
|
chunk = handle.read(budget)
|
||||||
# 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
|
|
||||||
|
|
||||||
shown = "\n".join(window)
|
has_more = start + len(chunk) < size
|
||||||
if has_more:
|
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 += (
|
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 ...]'
|
f'output_id="{output_id}", offset={next_offset}) to continue ...]'
|
||||||
)
|
)
|
||||||
return shown
|
return shown
|
||||||
|
|||||||
@@ -8,7 +8,7 @@ from strix.tools.output_store import read_stored_output
|
|||||||
|
|
||||||
|
|
||||||
@function_tool(timeout=10)
|
@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.
|
"""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
|
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:
|
Args:
|
||||||
output_id: The id from the truncation notice (a 32-char hex token).
|
output_id: The id from the truncation notice (a 32-char hex token).
|
||||||
offset: Zero-based line number to start reading from.
|
offset: Zero-based byte offset to start reading from.
|
||||||
limit: Maximum number of lines to return. Page forward by increasing
|
limit: Maximum number of bytes to return (capped at 50 KiB). Page
|
||||||
``offset`` using the hint printed at the end of each 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)
|
return read_stored_output(output_id, offset=offset, limit=limit)
|
||||||
|
|||||||
+43
-14
@@ -5,7 +5,6 @@ from __future__ import annotations
|
|||||||
import re
|
import re
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from strix.tools import output_store as _output_store
|
|
||||||
from strix.tools.output_store import (
|
from strix.tools.output_store import (
|
||||||
bound_and_store,
|
bound_and_store,
|
||||||
bound_text,
|
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)
|
output_id = match.group(1)
|
||||||
|
|
||||||
# The full, untruncated output round-trips through the store.
|
# 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()
|
assert full.splitlines() == text.splitlines()
|
||||||
# A buried line elided from the preview is retrievable.
|
# A buried line elided from the preview is retrievable.
|
||||||
assert "secret-line-500" not in bounded
|
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
|
assert output_id is not None
|
||||||
page = read_stored_output(output_id.group(1), offset=0, limit=10)
|
page = read_stored_output(output_id.group(1), offset=0, limit=10)
|
||||||
assert page.startswith("0\n1")
|
assert page.startswith("0\n1")
|
||||||
assert "more lines" in page
|
assert "more;" in page
|
||||||
assert "offset=10" in page
|
assert "offset=10" in page
|
||||||
|
|
||||||
|
|
||||||
def test_read_stored_output_pages_long_lines_losslessly(tmp_path: Path) -> None:
|
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*,
|
# A single line far larger than the page budget must be split across pages
|
||||||
# never by dropping content, so paging forward reconstructs everything and
|
# (never returned whole), and paging forward must reconstruct the output
|
||||||
# never mints a fresh spill id.
|
# byte-for-byte without ever minting a fresh spill id.
|
||||||
configure_output_store(tmp_path)
|
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(
|
output_id = re.search(
|
||||||
r'output_id="([0-9a-f]{32})"',
|
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
|
assert output_id is not None
|
||||||
oid = output_id.group(1)
|
oid = output_id.group(1)
|
||||||
|
|
||||||
collected: list[str] = []
|
collected = ""
|
||||||
offset = 0
|
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)
|
page = read_stored_output(oid, offset=offset, limit=2_000)
|
||||||
body, _sep, hint = page.partition("\n\n[... more lines;")
|
body, _sep, hint = page.partition("\n\n[... more;")
|
||||||
assert len(body.encode("utf-8")) <= _output_store._PAGE_MAX_BYTES
|
# 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])
|
assert re.findall(r'output_id="([0-9a-f]{32})"', body) in ([], [oid])
|
||||||
collected.extend(body.split("\n"))
|
collected += body
|
||||||
if not hint:
|
if not hint:
|
||||||
break
|
break
|
||||||
match = re.search(r"offset=(\d+)", hint)
|
match = re.search(r"offset=(\d+)", hint)
|
||||||
assert match is not None
|
assert match is not None
|
||||||
offset = int(match.group(1))
|
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:
|
def test_read_stored_output_rejects_traversal(tmp_path: Path) -> None:
|
||||||
|
|||||||
Reference in New Issue
Block a user