feat(context): bound per-tool output before it enters agent history

Cap the size of every tool result so a single verbose command (recursive
find, noisy scanner, full page dump) can't pin the conversation near the
model's context window for the rest of a scan.

- New ContextSettings config group with env-tunable caps.
- Default the SDK shell tools' max_output_tokens so exec_command /
  write_stdin truncate head+tail instead of returning unbounded output.
- Bound Strix's own FunctionTool/CustomTool results (line + UTF-8 byte
  head+tail preview with a truncation notice) and cap error strings.
This commit is contained in:
Ahmed Allam 2026-07-25 21:02:08 +00:00 committed by Ahmed Allam
parent d2fbcb726d
commit a70a87f272
6 changed files with 214 additions and 11 deletions

View file

@ -16,6 +16,7 @@ from agents.tool import CustomTool, FunctionTool, Tool
from pydantic import ValidationError
from strix.agents.prompt import render_system_prompt
from strix.config import load_settings
from strix.tools.agents_graph.tools import (
agent_finish,
create_agent,
@ -33,6 +34,7 @@ from strix.tools.notes.tools import (
list_notes,
update_note,
)
from strix.tools.output_store import bound_text
from strix.tools.proxy.tools import (
list_requests,
list_sitemap,
@ -108,8 +110,41 @@ def _extract_custom_input(tool: CustomTool, raw_input: str | dict[str, Any]) ->
return value if isinstance(value, str) else ""
def _tool_output_limits() -> tuple[int, int]:
context = load_settings().context
return context.tool_output_max_lines, context.tool_output_max_bytes
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)
def _format_tool_error(exc: Exception) -> str:
return str(exc) or exc.__class__.__name__
message = str(exc) or exc.__class__.__name__
max_lines, max_bytes = _tool_output_limits()
return bound_text(message, max_lines=max_lines, max_bytes=max_bytes)
def _with_bounded_result(tool: FunctionTool) -> FunctionTool:
"""Cap the size of a tool's result before it enters agent history.
Idempotent: base tools are shared singletons reused across every agent, so
the guard prevents stacking the wrapper on repeated ``build_strix_agent``
calls.
"""
if getattr(tool, "_strix_bounded", False):
return tool
invoke_tool = tool.on_invoke_tool
async def invoke(ctx: Any, raw_input: str) -> Any:
return _bound_result(await invoke_tool(ctx, raw_input))
tool.on_invoke_tool = invoke
tool._strix_bounded = True # type: ignore[attr-defined]
return tool
def _function_tool_with_error_result(tool: FunctionTool) -> FunctionTool:
@ -117,7 +152,7 @@ def _function_tool_with_error_result(tool: FunctionTool) -> FunctionTool:
async def invoke(ctx: Any, raw_input: str) -> Any:
try:
return await invoke_tool(ctx, raw_input)
return _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)
@ -132,7 +167,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 await tool.on_invoke_tool(ctx, custom_input)
return _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)
@ -210,6 +245,15 @@ def _format_validation_error(tool_name: str, exc: ValidationError) -> str:
return f"{tool_name}: invalid arguments — " + "; ".join(parts)
def _apply_shell_output_cap(parsed: dict[str, Any]) -> None:
"""Default the SDK shell tools' own token cap so a single command can't
dump unbounded output into history. The SDK truncates head+tail when the
model omits the field; respect an explicit model-supplied value.
"""
if parsed.get("max_output_tokens") is None:
parsed["max_output_tokens"] = load_settings().context.tool_output_max_tokens
def _wrap_exec_command(tool: FunctionTool) -> FunctionTool:
invoke_tool = tool.on_invoke_tool
@ -218,8 +262,10 @@ def _wrap_exec_command(tool: FunctionTool) -> FunctionTool:
parsed = json.loads(raw_input)
except (json.JSONDecodeError, TypeError):
parsed = None
if isinstance(parsed, dict) and "shell" not in parsed:
parsed["shell"] = "bash"
if isinstance(parsed, dict):
if "shell" not in parsed:
parsed["shell"] = "bash"
_apply_shell_output_cap(parsed)
raw_input = json.dumps(parsed)
try:
return await invoke_tool(ctx, raw_input)
@ -245,8 +291,10 @@ def _wrap_write_stdin(tool: FunctionTool) -> FunctionTool:
parsed = json.loads(raw_input)
except json.JSONDecodeError:
parsed = None
if isinstance(parsed, dict) and isinstance(parsed.get("chars"), str):
parsed["chars"] = _decode_chars_escape(parsed["chars"])
if isinstance(parsed, dict):
if isinstance(parsed.get("chars"), str):
parsed["chars"] = _decode_chars_escape(parsed["chars"])
_apply_shell_output_cap(parsed)
raw_input = json.dumps(parsed)
try:
return await invoke_tool(ctx, raw_input)
@ -447,6 +495,9 @@ def build_strix_agent(
else:
tools = [*_BASE_TOOLS, *agent_tools, agent_finish]
_ensure_unique_tool_names(tools)
tools = [
_with_bounded_result(tool) if isinstance(tool, FunctionTool) else tool for tool in tools
]
logger.info(
"Built %s agent '%s' (skills=%d, tools=%d, scan_mode=%s, whitebox=%s)",

View file

@ -17,6 +17,7 @@ from strix.config.loader import (
persist_current,
)
from strix.config.settings import (
ContextSettings,
DedupeSettings,
IntegrationSettings,
LlmSettings,
@ -27,6 +28,7 @@ from strix.config.settings import (
__all__ = [
"ContextSettings",
"DedupeSettings",
"IntegrationSettings",
"LlmSettings",

View file

@ -55,6 +55,23 @@ class DedupeSettings(BaseSettings):
api_base: str | None = Field(default=None, alias="DEDUPE_LLM_API_BASE")
class ContextSettings(BaseSettings):
"""Context-window management: per-tool-output caps and history compaction."""
model_config = _BASE_CONFIG
auto_compact: bool = Field(default=True, alias="STRIX_CONTEXT_AUTO_COMPACT")
compact_buffer_tokens: int = Field(default=20_000, gt=0, alias="STRIX_CONTEXT_BUFFER_TOKENS")
keep_tokens: int = Field(default=8_000, gt=0, alias="STRIX_CONTEXT_KEEP_TOKENS")
fallback_context_tokens: int = Field(
default=200_000, gt=0, alias="STRIX_CONTEXT_FALLBACK_TOKENS"
)
summary_max_tokens: int = Field(default=4_096, gt=0, alias="STRIX_CONTEXT_SUMMARY_TOKENS")
tool_output_max_tokens: int = Field(default=8_000, gt=0, alias="STRIX_TOOL_OUTPUT_MAX_TOKENS")
tool_output_max_lines: int = Field(default=2_000, gt=0, alias="STRIX_TOOL_OUTPUT_MAX_LINES")
tool_output_max_bytes: int = Field(default=50 * 1024, gt=0, alias="STRIX_TOOL_OUTPUT_MAX_BYTES")
class RuntimeSettings(BaseSettings):
model_config = _BASE_CONFIG
@ -99,6 +116,7 @@ class Settings(BaseSettings):
llm: LlmSettings = Field(default_factory=LlmSettings)
dedupe: DedupeSettings = Field(default_factory=DedupeSettings)
runtime: RuntimeSettings = Field(default_factory=RuntimeSettings)
context: ContextSettings = Field(default_factory=ContextSettings)
telemetry: TelemetrySettings = Field(default_factory=TelemetrySettings)
integrations: IntegrationSettings = Field(default_factory=IntegrationSettings)
viewer: ViewerSettings = Field(default_factory=ViewerSettings)

View file

@ -0,0 +1,72 @@
"""Bound oversized tool results before they enter agent history.
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.
"""
from __future__ import annotations
_TRUNCATION_NOTICE = "[... {lines} lines ({bytes} bytes) truncated ...]"
def _byte_len(text: str) -> int:
return len(text.encode("utf-8"))
def _take_prefix(text: str, max_bytes: int) -> str:
budget = 0
out: list[str] = []
for char in text:
size = len(char.encode("utf-8"))
if budget + size > max_bytes:
break
out.append(char)
budget += size
return "".join(out)
def _take_suffix(text: str, max_bytes: int) -> str:
budget = 0
out: list[str] = []
for char in reversed(text):
size = len(char.encode("utf-8"))
if budget + size > max_bytes:
break
out.append(char)
budget += size
out.reverse()
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.
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.
"""
lines = text.split("\n")
total_bytes = _byte_len(text)
if len(lines) <= max_lines and total_bytes <= max_bytes:
return text
head_lines = max(1, max_lines // 2)
tail_lines = max_lines - head_lines
head = "\n".join(lines[:head_lines])
tail = "\n".join(lines[len(lines) - tail_lines :]) if tail_lines > 0 else ""
# Enforce the byte budget even when the line count alone was fine.
half_bytes = max(1, max_bytes // 2)
if _byte_len(head) > half_bytes:
head = _take_prefix(head, half_bytes)
if tail and _byte_len(tail) > half_bytes:
tail = _take_suffix(tail, half_bytes)
dropped_lines = max(0, len(lines) - head_lines - tail_lines)
dropped_bytes = max(0, total_bytes - _byte_len(head) - _byte_len(tail))
notice = _TRUNCATION_NOTICE.format(lines=dropped_lines, bytes=dropped_bytes)
return f"{head}\n\n{notice}\n\n{tail}" if tail else f"{head}\n\n{notice}"

View file

@ -9,6 +9,7 @@ import pytest
from agents.tool import FunctionTool
from strix.agents import factory
from strix.config import load_settings
def _capturing_exec_tool(captured: dict[str, str]) -> FunctionTool:
@ -32,10 +33,23 @@ async def test_wrap_exec_command_defaults_shell_to_bash() -> None:
result = await wrapped.on_invoke_tool(cast("Any", None), json.dumps({"cmd": "source /tmp/env"}))
assert result == "ok"
assert json.loads(captured["raw_input"]) == {
"cmd": "source /tmp/env",
"shell": "bash",
}
parsed = json.loads(captured["raw_input"])
assert parsed["cmd"] == "source /tmp/env"
assert parsed["shell"] == "bash"
expected_cap = load_settings().context.tool_output_max_tokens
assert parsed["max_output_tokens"] == expected_cap
@pytest.mark.asyncio
async def test_wrap_exec_command_preserves_explicit_output_cap() -> None:
captured: dict[str, str] = {}
wrapped = factory._wrap_exec_command(_capturing_exec_tool(captured))
await wrapped.on_invoke_tool(
cast("Any", None), json.dumps({"cmd": "echo hi", "max_output_tokens": 42})
)
assert json.loads(captured["raw_input"])["max_output_tokens"] == 42
@pytest.mark.asyncio

View file

@ -0,0 +1,46 @@
"""Tests for per-tool-output bounding."""
from __future__ import annotations
from strix.tools.output_store import bound_text
def test_small_output_passes_through_unchanged() -> None:
text = "line 1\nline 2\nline 3"
assert bound_text(text, max_lines=100, max_bytes=10_000) == text
def test_line_limit_keeps_head_and_tail() -> None:
text = "\n".join(str(i) for i in range(1000))
bounded = bound_text(text, max_lines=10, max_bytes=1_000_000)
assert bounded.startswith("0\n1\n2\n3\n4")
assert bounded.rstrip().endswith("999")
assert "truncated" in bounded
# Head + tail only, far fewer than the original 1000 lines.
assert len(bounded.splitlines()) < 30
def test_byte_limit_enforced_on_single_long_line() -> None:
text = "x" * 100_000
bounded = bound_text(text, max_lines=2_000, max_bytes=1_000)
assert "truncated" in bounded
assert len(bounded.encode("utf-8")) < 3_000
def test_multibyte_characters_not_split() -> None:
text = "😀" * 50_000
bounded = bound_text(text, max_lines=2_000, max_bytes=1_000)
# Must remain valid UTF-8 (no broken surrogate halves from a mid-char cut).
assert bounded == bounded.encode("utf-8").decode("utf-8")
assert "truncated" in bounded
def test_notice_reports_dropped_counts() -> None:
text = "\n".join("y" * 10 for _ in range(500))
bounded = bound_text(text, max_lines=10, max_bytes=1_000_000)
assert "lines" in bounded
assert "bytes" in bounded