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-26 14:38:12 -07:00
committed by Ahmed Allam
parent d2fbcb726d
commit a70a87f272
6 changed files with 214 additions and 11 deletions
+58 -7
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)",
+2
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",
+18
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)
+72
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}"
+18 -4
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
+46
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