From 3b79e97f000aa65461e61839c647a70f7e754554 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 26 Jul 2026 14:42:22 -0700 Subject: [PATCH] feat(context): spill oversized tool output into the sandbox workspace (#882) Co-authored-by: Ahmed Allam --- strix/agents/factory.py | 14 ++-- strix/core/runner.py | 21 ++++++ strix/tools/output_store.py | 114 +++++++++++++++++++++++++++--- tests/test_agent_factory_shell.py | 7 ++ tests/test_output_store.py | 100 +++++++++++++++++++++++++- 5 files changed, 236 insertions(+), 20 deletions(-) diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 872e3564..5d922b19 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, @@ -115,11 +115,11 @@ def _tool_output_limits() -> tuple[int, int]: return context.tool_output_max_lines, context.tool_output_max_bytes -def _bound_result(result: Any) -> Any: +async 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 await bound_and_store(result, max_lines=max_lines, max_bytes=max_bytes) def _format_tool_error(exc: Exception) -> str: @@ -135,7 +135,7 @@ def _with_bounded_result(tool: FunctionTool) -> FunctionTool: invoke_tool = tool.on_invoke_tool async def invoke(ctx: Any, raw_input: str) -> Any: - return _bound_result(await invoke_tool(ctx, raw_input)) + return await _bound_result(await invoke_tool(ctx, raw_input)) tool.on_invoke_tool = invoke tool._strix_bounded = True # type: ignore[attr-defined] @@ -147,7 +147,7 @@ def _function_tool_with_error_result(tool: FunctionTool) -> FunctionTool: async def invoke(ctx: Any, raw_input: str) -> Any: try: - return _bound_result(await invoke_tool(ctx, raw_input)) + return await _bound_result(await invoke_tool(ctx, raw_input)) except Exception as exc: # noqa: BLE001 - tool errors should be model-visible results. logger.debug("Tool %s failed; returning error as result", tool.name, exc_info=True) return _format_tool_error(exc) @@ -162,7 +162,7 @@ def _custom_tool_as_function_tool(tool: CustomTool) -> FunctionTool: if not custom_input: return f"`{_custom_tool_input_field(tool)}` must be a non-empty string." try: - return _bound_result(await tool.on_invoke_tool(ctx, custom_input)) + return await _bound_result(await tool.on_invoke_tool(ctx, custom_input)) except Exception as exc: # noqa: BLE001 - matches SDK CustomTool error-as-result behavior. logger.debug("Tool %s failed; returning error as result", tool.name, exc_info=True) return _format_tool_error(exc) @@ -199,7 +199,7 @@ def _bound_custom_tool(tool: CustomTool) -> CustomTool: invoke_tool = tool.on_invoke_tool async def invoke(ctx: Any, raw_input: str) -> Any: - return _bound_result(await invoke_tool(ctx, raw_input)) + return await _bound_result(await invoke_tool(ctx, raw_input)) tool.on_invoke_tool = invoke return tool diff --git a/strix/core/runner.py b/strix/core/runner.py index 23b58709..5776ce42 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -3,10 +3,12 @@ from __future__ import annotations import contextlib +import io import json import logging import uuid from collections.abc import Callable +from pathlib import Path from typing import TYPE_CHECKING, Any from agents import RunConfig @@ -40,6 +42,10 @@ 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 ( + WORKSPACE_SPILL_DIR, + configure_spill_writer, +) if TYPE_CHECKING: @@ -203,6 +209,20 @@ async def run_strix_scan( ) logger.info("Sandbox ready for scan %s", scan_id) + sandbox_session = bundle["session"] + + async def _spill_to_workspace(output_id: str, text: str) -> str | None: + """Write an oversized tool result into the sandbox; return its path or None.""" + path = f"{WORKSPACE_SPILL_DIR}/{output_id}.txt" + try: + await sandbox_session.write(Path(path), io.BytesIO(text.encode("utf-8"))) + except Exception: + logger.exception("failed to spill tool output to sandbox workspace") + return None + return path + + configure_spill_writer(_spill_to_workspace) + sessions_to_close: list[SQLiteSession] = [] try: @@ -399,6 +419,7 @@ async def run_strix_scan( await coordinator.set_status(root_id, "failed") raise finally: + configure_spill_writer(None) for s in sessions_to_close: with contextlib.suppress(Exception): s.close() diff --git a/strix/tools/output_store.py b/strix/tools/output_store.py index ed65dfa8..a525563c 100644 --- a/strix/tools/output_store.py +++ b/strix/tools/output_store.py @@ -1,13 +1,47 @@ """Bound oversized tool results before they enter agent history. -Keeps a head + tail slice and drops the middle, replacing it with a notice of -how much was removed. +Oversized results are spilled into the sandbox at +``/workspace/.strix/tool-output/.txt``; the agent sees a head + tail slice +plus the path and reads the rest back with its own file tools. The spill writer +is injected by the runner via :func:`configure_spill_writer`. """ from __future__ import annotations +import logging +import uuid +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + +logger = logging.getLogger(__name__) _TRUNCATION_NOTICE = "[... {lines} lines ({bytes} bytes) truncated ...]" +_WORKSPACE_SPILL_NOTICE = ( + "[... {lines} lines ({bytes} bytes) truncated — full output saved to {path} " + "in the sandbox; read it with exec_command (e.g. `sed -n`, `grep`, `cat`) ...]" +) + +WORKSPACE_SPILL_DIR = "/workspace/.strix/tool-output" + +# Longest possible workspace path, used only to reserve notice bytes. +_SAMPLE_WORKSPACE_PATH = f"{WORKSPACE_SPILL_DIR}/{'0' * 32}.txt" + +if TYPE_CHECKING: + SpillWriter = Callable[[str, str], Awaitable[str | None]] + +_spill: dict[str, SpillWriter] = {} + + +def configure_spill_writer(writer: SpillWriter | None) -> None: + """Install (or clear) the sandbox-workspace spill writer.""" + if writer is None: + _spill.pop("writer", None) + else: + _spill["writer"] = writer def _byte_len(text: str) -> int: @@ -39,20 +73,37 @@ 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_templates: tuple[str, ...] = (_TRUNCATION_NOTICE,), +) -> tuple[str, str, int, int] | None: + """Head/tail slices plus dropped line/byte counts, or ``None`` if small. - Truncates on whichever limit is hit first (line count or UTF-8 byte size). - ``max_bytes`` bounds the entire joined result, notice and separators - included. + ``max_bytes`` bounds the entire joined result; the largest of + ``notice_templates`` (plus separators) is reserved before slicing. """ lines = text.split("\n") total_bytes = _byte_len(text) if len(lines) <= max_lines and total_bytes <= max_bytes: - return text + return None - # Reserve notice + separator bytes up front; ``+ 4`` covers the two "\n\n". - notice_overhead = _byte_len(_TRUNCATION_NOTICE.format(lines=len(lines), bytes=total_bytes)) + 4 + # Reserve using the largest counts/path; ``+ 4`` covers the two "\n\n". + notice_overhead = ( + max( + _byte_len( + template.format( + lines=len(lines), + bytes=total_bytes, + path=_SAMPLE_WORKSPACE_PATH, + ) + ) + for template in notice_templates + ) + + 4 + ) byte_budget = max(2, max_bytes - notice_overhead) head_lines = max(1, max_lines // 2) @@ -70,5 +121,46 @@ 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. + + 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)) + + +async def bound_and_store(text: str, *, max_lines: int, max_bytes: int) -> str: + """Like :func:`bound_text`, but spill the full output into the sandbox and + point the agent at its path. Degrades to a plain preview if the spill fails. + """ + parts = _head_tail( + text, + max_lines, + max_bytes, + notice_templates=(_WORKSPACE_SPILL_NOTICE, _TRUNCATION_NOTICE), + ) + if parts is None: + return text + head, tail, dropped_lines, dropped_bytes = parts + + writer = _spill.get("writer") + if writer is not None: + path = await writer(uuid.uuid4().hex, text) + if path is not None: + notice = _WORKSPACE_SPILL_NOTICE.format( + lines=dropped_lines, bytes=dropped_bytes, path=path + ) + return _join(head, tail, notice) + + return _join(head, tail, _TRUNCATION_NOTICE.format(lines=dropped_lines, bytes=dropped_bytes)) diff --git a/tests/test_agent_factory_shell.py b/tests/test_agent_factory_shell.py index 814da60c..6bef1211 100644 --- a/tests/test_agent_factory_shell.py +++ b/tests/test_agent_factory_shell.py @@ -108,3 +108,10 @@ async def test_chat_completions_filesystem_custom_tool_becomes_function_tool() - factory._configure_filesystem_tools(toolset, chat_completions=True) assert isinstance(toolset.read_file, FunctionTool) + + +def test_function_tools_are_result_bounded() -> None: + agent = factory.build_strix_agent(is_root=True) + by_name = {t.name: t for t in agent.tools} + + assert getattr(by_name["think"], "_strix_bounded", False) is True diff --git a/tests/test_output_store.py b/tests/test_output_store.py index 4db8ccf7..a2570918 100644 --- a/tests/test_output_store.py +++ b/tests/test_output_store.py @@ -1,10 +1,22 @@ -"""Tests for per-tool-output bounding.""" +"""Tests for per-tool-output bounding and sandbox-workspace spill.""" from __future__ import annotations import re -from strix.tools.output_store import bound_text +import pytest + +from strix.tools.output_store import ( + WORKSPACE_SPILL_DIR, + bound_and_store, + bound_text, + configure_spill_writer, +) + + +@pytest.fixture(autouse=True) +def _clear_spill_writer() -> None: + configure_spill_writer(None) def test_small_output_passes_through_unchanged() -> None: @@ -58,3 +70,87 @@ def test_dropped_line_count_accounts_for_byte_trimming() -> None: kept = [ln for ln in bounded.splitlines() if ln and "truncated" not in ln] assert dropped == 200 - len(kept) assert dropped > 200 - 20 + + +async def test_bound_and_store_small_output_not_spilled() -> None: + written: dict[str, str] = {} + + async def writer(output_id: str, text: str) -> str | None: + written[output_id] = text + return f"{WORKSPACE_SPILL_DIR}/{output_id}.txt" + + configure_spill_writer(writer) + text = "just a few lines\nsecond line" + assert await bound_and_store(text, max_lines=100, max_bytes=10_000) == text + assert written == {} + + +async def test_bound_and_store_spills_full_output_to_workspace() -> None: + written: dict[str, str] = {} + + async def writer(output_id: str, text: str) -> str | None: + written[output_id] = text + return f"{WORKSPACE_SPILL_DIR}/{output_id}.txt" + + configure_spill_writer(writer) + text = "\n".join(f"secret-line-{i}" for i in range(1000)) + bounded = await bound_and_store(text, max_lines=10, max_bytes=1_000_000) + + assert WORKSPACE_SPILL_DIR in bounded + assert "exec_command" in bounded + assert "read_tool_output" not in bounded + assert len(bounded.encode("utf-8")) <= 1_000_000 + assert list(written.values()) == [text] + stored = next(iter(written.values())) + assert stored.splitlines() == text.splitlines() + # A buried line elided from the preview is still present in the spilled file. + assert "secret-line-500" not in bounded + assert "secret-line-500" in stored + + +async def test_workspace_notice_carries_the_returned_path() -> None: + async def writer(output_id: str, _text: str) -> str | None: + return f"{WORKSPACE_SPILL_DIR}/{output_id}.txt" + + configure_spill_writer(writer) + text = "\n".join(f"line-{i}" for i in range(1000)) + bounded = await bound_and_store(text, max_lines=10, max_bytes=1_000_000) + + match = re.search(rf"{re.escape(WORKSPACE_SPILL_DIR)}/([0-9a-f]{{32}})\.txt", bounded) + assert match is not None, bounded + + +async def test_no_writer_degrades_to_plain_preview() -> None: + text = "\n".join(f"line-{i}" for i in range(1000)) + bounded = await bound_and_store(text, max_lines=10, max_bytes=1_000_000) + + assert "truncated" in bounded + assert WORKSPACE_SPILL_DIR not in bounded + assert "read_tool_output" not in bounded + assert len(bounded.encode("utf-8")) <= 1_000_000 + + +async def test_writer_failure_degrades_to_plain_preview() -> None: + async def failing_writer(_output_id: str, _text: str) -> str | None: + return None + + configure_spill_writer(failing_writer) + text = "\n".join(f"line-{i}" for i in range(1000)) + bounded = await bound_and_store(text, max_lines=10, max_bytes=1_000_000) + + assert "truncated" in bounded + assert WORKSPACE_SPILL_DIR not in bounded + assert len(bounded.encode("utf-8")) <= 1_000_000 + + +async def test_workspace_preview_honours_byte_budget() -> None: + # The workspace notice is longer than a plain notice; the preview reserves for it. + async def writer(output_id: str, _text: str) -> str | None: + return f"{WORKSPACE_SPILL_DIR}/{output_id}.txt" + + configure_spill_writer(writer) + text = "\n".join("x" * 500 for _ in range(200)) + bounded = await bound_and_store(text, max_lines=2_000, max_bytes=2_000) + + assert WORKSPACE_SPILL_DIR in bounded + assert len(bounded.encode("utf-8")) <= 2_000