diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 128c6d6e..732cdf8f 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -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)", diff --git a/strix/config/__init__.py b/strix/config/__init__.py index 6e9dded1..f21fdab6 100644 --- a/strix/config/__init__.py +++ b/strix/config/__init__.py @@ -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", diff --git a/strix/config/settings.py b/strix/config/settings.py index 98d68934..edb0ab09 100644 --- a/strix/config/settings.py +++ b/strix/config/settings.py @@ -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) diff --git a/strix/tools/output_store.py b/strix/tools/output_store.py new file mode 100644 index 00000000..44a9fa9c --- /dev/null +++ b/strix/tools/output_store.py @@ -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}" diff --git a/tests/test_agent_factory_shell.py b/tests/test_agent_factory_shell.py index 6de93219..30276562 100644 --- a/tests/test_agent_factory_shell.py +++ b/tests/test_agent_factory_shell.py @@ -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 diff --git a/tests/test_output_store.py b/tests/test_output_store.py new file mode 100644 index 00000000..23bcdf30 --- /dev/null +++ b/tests/test_output_store.py @@ -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