From 1da3140d8409e5651c5446c65aa3eccd2e30e4fc Mon Sep 17 00:00:00 2001 From: Ahmed Allam Date: Sat, 25 Jul 2026 21:38:43 +0000 Subject: [PATCH] feat(context): spill oversized tool output to disk with a retrieval tool Truncating a large tool result to a head+tail preview loses the middle, which may hold the one line that matters (a buried match, a stack frame, a credential). Instead of dropping it, persist the full output and let the agent page back to it on demand. - output_store.py: bound_and_store() writes the complete output to a per-scan store and embeds an output_id in the truncation notice; read_stored_output() serves it back in validated, paginated chunks (output_id is a 32-char hex token, guarding against path traversal). - read_tool_output tool: lets the agent retrieve any elided output by id, paging via offset/limit. - factory.py: bounded tool results now spill via bound_and_store; the new tool is registered for every scan agent. - runner.py: point the store at the run's .state/tool-output directory so spilled output lives beside the rest of the scan state. Falls back to a plain head+tail preview if the spill write fails. --- strix/agents/factory.py | 6 +- strix/core/runner.py | 2 + strix/tools/output_store.py | 142 ++++++++++++++++++++++++---- strix/tools/tool_output/__init__.py | 6 ++ strix/tools/tool_output/tool.py | 25 +++++ tests/test_output_store.py | 63 +++++++++++- 6 files changed, 223 insertions(+), 21 deletions(-) create mode 100644 strix/tools/tool_output/__init__.py create mode 100644 strix/tools/tool_output/tool.py diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 4b880de4..c411c2ea 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -34,7 +34,7 @@ from strix.tools.notes.tools import ( list_notes, update_note, ) -from strix.tools.output_store import bound_text +from strix.tools.output_store import bound_and_store, bound_text from strix.tools.proxy.tools import ( list_requests, list_sitemap, @@ -53,6 +53,7 @@ from strix.tools.todo.tools import ( mark_todo_pending, update_todo, ) +from strix.tools.tool_output import read_tool_output from strix.tools.web_search.tool import web_search @@ -114,7 +115,7 @@ def _bound_result(result: Any) -> Any: if not isinstance(result, str): return result max_lines, max_bytes = _tool_output_limits() - return bound_text(result, max_lines=max_lines, max_bytes=max_bytes) + return bound_and_store(result, max_lines=max_lines, max_bytes=max_bytes) def _format_tool_error(exc: Exception) -> str: @@ -413,6 +414,7 @@ def _finish_tool_use_behavior( _BASE_TOOLS: tuple[Tool, ...] = ( think, + read_tool_output, load_skill, create_todo, list_todos, diff --git a/strix/core/runner.py b/strix/core/runner.py index 23b58709..f9473e2b 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -40,6 +40,7 @@ from strix.core.paths import run_dir_for, runtime_state_dir from strix.core.sessions import open_agent_session from strix.runtime import session_manager from strix.telemetry.logging import set_scan_id, setup_scan_logging +from strix.tools.output_store import configure_output_store if TYPE_CHECKING: @@ -134,6 +135,7 @@ async def run_strix_scan( agents_path = state_dir / "agents.json" agents_db = state_dir / "agents.db" + configure_output_store(state_dir / "tool-output") is_resume = agents_path.exists() logger.info( diff --git a/strix/tools/output_store.py b/strix/tools/output_store.py index 6e6e6202..51fffa9e 100644 --- a/strix/tools/output_store.py +++ b/strix/tools/output_store.py @@ -2,15 +2,45 @@ A single verbose tool result (a recursive ``find``, a noisy scanner, a full page dump) can otherwise pin the whole conversation near the model's context -limit for the rest of the scan. This keeps a head + tail slice of the output -and drops the middle, mirroring how the shell capability truncates its own -output — the agent still sees the start and end plus how much was removed. +limit for the rest of the scan. Oversized results are spilled to a per-scan +store on disk; what the agent sees is a head + tail slice plus an id it can +pass to ``read_tool_output`` to page through the full content on demand — so +truncated detail is bounded in history but never lost. """ from __future__ import annotations +import logging +import re +import uuid +from pathlib import Path + + +logger = logging.getLogger(__name__) _TRUNCATION_NOTICE = "[... {lines} lines ({bytes} bytes) truncated ...]" +_SPILL_NOTICE = ( + "[... {lines} lines ({bytes} bytes) truncated — full output saved as " + 'output_id="{output_id}"; read it with read_tool_output(output_id="{output_id}") ...]' +) + +_OUTPUT_ID_RE = re.compile(r"^[0-9a-f]{32}$") +_DEFAULT_STORE_DIR = Path.home() / ".strix" / "tool-output" + +# Single-key holder so the configured path can be swapped per scan without a +# module-level ``global`` rebind. +_config: dict[str, Path] = {} + + +def configure_output_store(directory: Path) -> None: + """Point the tool-output store at ``directory`` (created on demand).""" + _config["dir"] = directory + + +def _active_store_dir() -> Path: + directory = _config.get("dir", _DEFAULT_STORE_DIR) + directory.mkdir(parents=True, exist_ok=True) + return directory def _byte_len(text: str) -> int: @@ -42,26 +72,30 @@ def _take_suffix(text: str, max_bytes: int) -> str: return "".join(out) -def bound_text(text: str, *, max_lines: int, max_bytes: int) -> str: - """Return ``text`` unchanged when small, else a head+tail preview. +def _head_tail( + text: str, max_lines: int, max_bytes: int, *, notice_template: str = _TRUNCATION_NOTICE +) -> tuple[str, str, int, int] | None: + """Head/tail slices plus dropped line/byte counts, or ``None`` if small. - Truncation happens on whichever limit is hit first (line count or UTF-8 - byte size). The removed middle is replaced with a notice recording how - many lines and bytes were dropped so the agent knows output was elided. - ``max_bytes`` bounds the *entire* joined result, notice and separators - included, and must be large enough to hold the notice itself (guaranteed by + ``max_bytes`` bounds the *entire* joined result, so the notice and its two + blank-line separators are reserved out of the byte budget before slicing — + otherwise ``head + tail`` alone could already consume the whole budget and + the appended metadata would push the persisted value over ``max_bytes``. + ``max_bytes`` must be large enough to hold the notice itself (guaranteed by the ``tool_output_max_bytes`` config floor). """ lines = text.split("\n") total_bytes = _byte_len(text) if len(lines) <= max_lines and total_bytes <= max_bytes: - return text + return None - # Reserve room for the notice and its two blank-line separators so the - # head+tail slices can't consume the whole budget and push the persisted - # value over max_bytes. Upper-bound the notice with the largest possible - # counts; the real notice is never longer. ``+ 4`` covers the separators. - notice_overhead = _byte_len(_TRUNCATION_NOTICE.format(lines=len(lines), bytes=total_bytes)) + 4 + # Upper-bound the notice size using the largest possible counts (and a + # full-length id); the real notice is never longer. ``+ 4`` covers the two + # ``\n\n`` separators added by ``_join``. + notice_overhead = ( + _byte_len(notice_template.format(lines=len(lines), bytes=total_bytes, output_id="0" * 32)) + + 4 + ) byte_budget = max(2, max_bytes - notice_overhead) head_lines = max(1, max_lines // 2) @@ -82,5 +116,79 @@ def bound_text(text: str, *, max_lines: int, max_bytes: int) -> str: kept_lines = len(head.split("\n")) + (len(tail.split("\n")) if tail else 0) dropped_lines = max(0, len(lines) - kept_lines) dropped_bytes = max(0, total_bytes - _byte_len(head) - _byte_len(tail)) - notice = _TRUNCATION_NOTICE.format(lines=dropped_lines, bytes=dropped_bytes) + return head, tail, dropped_lines, dropped_bytes + + +def _join(head: str, tail: str, notice: str) -> str: return f"{head}\n\n{notice}\n\n{tail}" if tail else f"{head}\n\n{notice}" + + +def bound_text(text: str, *, max_lines: int, max_bytes: int) -> str: + """Return ``text`` unchanged when small, else a head+tail preview. + + Truncation happens on whichever limit is hit first (line count or UTF-8 + byte size). The removed middle is replaced with a notice recording how + many lines and bytes were dropped so the agent knows output was elided. + Nothing is persisted; use :func:`bound_and_store` to keep the full output. + """ + parts = _head_tail(text, max_lines, max_bytes) + if parts is None: + return text + head, tail, dropped_lines, dropped_bytes = parts + return _join(head, tail, _TRUNCATION_NOTICE.format(lines=dropped_lines, bytes=dropped_bytes)) + + +def store_full_output(text: str) -> str | None: + """Persist ``text`` and return its output id, or ``None`` if writing fails.""" + output_id = uuid.uuid4().hex + try: + (_active_store_dir() / f"{output_id}.txt").write_text(text, encoding="utf-8") + except OSError: + logger.exception("failed to persist oversized tool output") + return None + return output_id + + +def bound_and_store(text: str, *, max_lines: int, max_bytes: int) -> str: + """Like :func:`bound_text`, but spill the full output and reference its id. + + When truncation happens the complete output is written to the store and the + preview's notice tells the agent the ``output_id`` to pass to + ``read_tool_output``. Falls back to a plain preview if the spill fails. + """ + # Reserve for the (longer) spill notice so the preview honours max_bytes + # whether or not the spill succeeds. + parts = _head_tail(text, max_lines, max_bytes, notice_template=_SPILL_NOTICE) + if parts is None: + return text + head, tail, dropped_lines, dropped_bytes = parts + output_id = store_full_output(text) + notice = ( + _SPILL_NOTICE.format(lines=dropped_lines, bytes=dropped_bytes, output_id=output_id) + if output_id is not None + else _TRUNCATION_NOTICE.format(lines=dropped_lines, bytes=dropped_bytes) + ) + return _join(head, tail, notice) + + +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``. + + ``output_id`` must be a token previously returned in a truncation notice; + it is validated to prevent path traversal. + """ + if not _OUTPUT_ID_RE.match(output_id): + return f"Invalid output_id: {output_id!r}" + path = _active_store_dir() / f"{output_id}.txt" + if not path.is_file(): + return f"No stored output for output_id={output_id!r} (it may have expired)." + + lines = path.read_text(encoding="utf-8").split("\n") + start = max(0, offset) + window = lines[start : start + max(1, limit)] + shown = "\n".join(window) + remaining = len(lines) - (start + len(window)) + 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 ...]' + return shown diff --git a/strix/tools/tool_output/__init__.py b/strix/tools/tool_output/__init__.py new file mode 100644 index 00000000..db4e8933 --- /dev/null +++ b/strix/tools/tool_output/__init__.py @@ -0,0 +1,6 @@ +"""Retrieval of oversized tool outputs spilled to the tool-output store.""" + +from strix.tools.tool_output.tool import read_tool_output + + +__all__ = ["read_tool_output"] diff --git a/strix/tools/tool_output/tool.py b/strix/tools/tool_output/tool.py new file mode 100644 index 00000000..6feec5ab --- /dev/null +++ b/strix/tools/tool_output/tool.py @@ -0,0 +1,25 @@ +"""``read_tool_output`` — page through a previously truncated tool result.""" + +from __future__ import annotations + +from agents import function_tool + +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: + """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 + the conversation and the complete text is saved with an ``output_id`` shown + in the truncation notice. Use this to retrieve the parts that were elided + (e.g. a specific match buried in the middle of a long scan or file dump). + + 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. + """ + 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 14bff19c..818572a3 100644 --- a/tests/test_output_store.py +++ b/tests/test_output_store.py @@ -1,10 +1,20 @@ -"""Tests for per-tool-output bounding.""" +"""Tests for per-tool-output bounding and the durable spill store.""" from __future__ import annotations import re +from typing import TYPE_CHECKING -from strix.tools.output_store import bound_text +from strix.tools.output_store import ( + bound_and_store, + bound_text, + configure_output_store, + read_stored_output, +) + + +if TYPE_CHECKING: + from pathlib import Path def test_small_output_passes_through_unchanged() -> None: @@ -63,3 +73,52 @@ def test_dropped_line_count_accounts_for_byte_trimming() -> None: assert dropped == 200 - len(kept) # The naive middle-only count (max_lines split evenly) would under-report. assert dropped > 200 - 20 + + +def test_bound_and_store_small_output_not_spilled(tmp_path: Path) -> None: + configure_output_store(tmp_path) + text = "just a few lines\nsecond line" + assert bound_and_store(text, max_lines=100, max_bytes=10_000) == text + assert list(tmp_path.iterdir()) == [] + + +def test_bound_and_store_spills_full_output_and_is_retrievable(tmp_path: Path) -> None: + configure_output_store(tmp_path) + text = "\n".join(f"secret-line-{i}" for i in range(1000)) + + bounded = bound_and_store(text, max_lines=10, max_bytes=1_000_000) + + match = re.search(r'output_id="([0-9a-f]{32})"', bounded) + assert match is not None, bounded + 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) + assert full.splitlines() == text.splitlines() + # A buried line elided from the preview is retrievable. + assert "secret-line-500" not in bounded + assert "secret-line-500" in full + + +def test_read_stored_output_paginates(tmp_path: Path) -> None: + configure_output_store(tmp_path) + text = "\n".join(str(i) for i in range(100)) + output_id = re.search( + r'output_id="([0-9a-f]{32})"', + bound_and_store(text, max_lines=4, max_bytes=1_000_000), + ) + 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 "offset=10" in page + + +def test_read_stored_output_rejects_traversal(tmp_path: Path) -> None: + configure_output_store(tmp_path) + assert "Invalid output_id" in read_stored_output("../../etc/passwd") + + +def test_read_stored_output_missing_id(tmp_path: Path) -> None: + configure_output_store(tmp_path) + assert "No stored output" in read_stored_output("0" * 32)