mirror of
https://github.com/usestrix/strix.git
synced 2026-08-25 04:12:37 +02:00
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.
This commit is contained in:
@@ -34,7 +34,7 @@ from strix.tools.notes.tools import (
|
|||||||
list_notes,
|
list_notes,
|
||||||
update_note,
|
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 (
|
from strix.tools.proxy.tools import (
|
||||||
list_requests,
|
list_requests,
|
||||||
list_sitemap,
|
list_sitemap,
|
||||||
@@ -53,6 +53,7 @@ from strix.tools.todo.tools import (
|
|||||||
mark_todo_pending,
|
mark_todo_pending,
|
||||||
update_todo,
|
update_todo,
|
||||||
)
|
)
|
||||||
|
from strix.tools.tool_output import read_tool_output
|
||||||
from strix.tools.web_search.tool import web_search
|
from strix.tools.web_search.tool import web_search
|
||||||
|
|
||||||
|
|
||||||
@@ -114,7 +115,7 @@ def _bound_result(result: Any) -> Any:
|
|||||||
if not isinstance(result, str):
|
if not isinstance(result, str):
|
||||||
return result
|
return result
|
||||||
max_lines, max_bytes = _tool_output_limits()
|
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:
|
def _format_tool_error(exc: Exception) -> str:
|
||||||
@@ -413,6 +414,7 @@ def _finish_tool_use_behavior(
|
|||||||
|
|
||||||
_BASE_TOOLS: tuple[Tool, ...] = (
|
_BASE_TOOLS: tuple[Tool, ...] = (
|
||||||
think,
|
think,
|
||||||
|
read_tool_output,
|
||||||
load_skill,
|
load_skill,
|
||||||
create_todo,
|
create_todo,
|
||||||
list_todos,
|
list_todos,
|
||||||
|
|||||||
@@ -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.core.sessions import open_agent_session
|
||||||
from strix.runtime import session_manager
|
from strix.runtime import session_manager
|
||||||
from strix.telemetry.logging import set_scan_id, setup_scan_logging
|
from strix.telemetry.logging import set_scan_id, setup_scan_logging
|
||||||
|
from strix.tools.output_store import configure_output_store
|
||||||
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
@@ -134,6 +135,7 @@ async def run_strix_scan(
|
|||||||
|
|
||||||
agents_path = state_dir / "agents.json"
|
agents_path = state_dir / "agents.json"
|
||||||
agents_db = state_dir / "agents.db"
|
agents_db = state_dir / "agents.db"
|
||||||
|
configure_output_store(state_dir / "tool-output")
|
||||||
is_resume = agents_path.exists()
|
is_resume = agents_path.exists()
|
||||||
|
|
||||||
logger.info(
|
logger.info(
|
||||||
|
|||||||
+125
-17
@@ -2,15 +2,45 @@
|
|||||||
|
|
||||||
A single verbose tool result (a recursive ``find``, a noisy scanner, a full
|
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
|
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
|
limit for the rest of the scan. Oversized results are spilled to a per-scan
|
||||||
and drops the middle, mirroring how the shell capability truncates its own
|
store on disk; what the agent sees is a head + tail slice plus an id it can
|
||||||
output — the agent still sees the start and end plus how much was removed.
|
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
|
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 ...]"
|
_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:
|
def _byte_len(text: str) -> int:
|
||||||
@@ -42,26 +72,30 @@ def _take_suffix(text: str, max_bytes: int) -> str:
|
|||||||
return "".join(out)
|
return "".join(out)
|
||||||
|
|
||||||
|
|
||||||
def bound_text(text: str, *, max_lines: int, max_bytes: int) -> str:
|
def _head_tail(
|
||||||
"""Return ``text`` unchanged when small, else a head+tail preview.
|
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
|
``max_bytes`` bounds the *entire* joined result, so the notice and its two
|
||||||
byte size). The removed middle is replaced with a notice recording how
|
blank-line separators are reserved out of the byte budget before slicing —
|
||||||
many lines and bytes were dropped so the agent knows output was elided.
|
otherwise ``head + tail`` alone could already consume the whole budget and
|
||||||
``max_bytes`` bounds the *entire* joined result, notice and separators
|
the appended metadata would push the persisted value over ``max_bytes``.
|
||||||
included, and must be large enough to hold the notice itself (guaranteed by
|
``max_bytes`` must be large enough to hold the notice itself (guaranteed by
|
||||||
the ``tool_output_max_bytes`` config floor).
|
the ``tool_output_max_bytes`` config floor).
|
||||||
"""
|
"""
|
||||||
lines = text.split("\n")
|
lines = text.split("\n")
|
||||||
total_bytes = _byte_len(text)
|
total_bytes = _byte_len(text)
|
||||||
if len(lines) <= max_lines and total_bytes <= max_bytes:
|
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
|
# Upper-bound the notice size using the largest possible counts (and a
|
||||||
# head+tail slices can't consume the whole budget and push the persisted
|
# full-length id); the real notice is never longer. ``+ 4`` covers the two
|
||||||
# value over max_bytes. Upper-bound the notice with the largest possible
|
# ``\n\n`` separators added by ``_join``.
|
||||||
# counts; the real notice is never longer. ``+ 4`` covers the separators.
|
notice_overhead = (
|
||||||
notice_overhead = _byte_len(_TRUNCATION_NOTICE.format(lines=len(lines), bytes=total_bytes)) + 4
|
_byte_len(notice_template.format(lines=len(lines), bytes=total_bytes, output_id="0" * 32))
|
||||||
|
+ 4
|
||||||
|
)
|
||||||
byte_budget = max(2, max_bytes - notice_overhead)
|
byte_budget = max(2, max_bytes - notice_overhead)
|
||||||
|
|
||||||
head_lines = max(1, max_lines // 2)
|
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)
|
kept_lines = len(head.split("\n")) + (len(tail.split("\n")) if tail else 0)
|
||||||
dropped_lines = max(0, len(lines) - kept_lines)
|
dropped_lines = max(0, len(lines) - kept_lines)
|
||||||
dropped_bytes = max(0, total_bytes - _byte_len(head) - _byte_len(tail))
|
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}"
|
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
|
||||||
|
|||||||
@@ -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"]
|
||||||
@@ -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)
|
||||||
@@ -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
|
from __future__ import annotations
|
||||||
|
|
||||||
import re
|
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:
|
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)
|
assert dropped == 200 - len(kept)
|
||||||
# The naive middle-only count (max_lines split evenly) would under-report.
|
# The naive middle-only count (max_lines split evenly) would under-report.
|
||||||
assert dropped > 200 - 20
|
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)
|
||||||
|
|||||||
Reference in New Issue
Block a user