diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 732cdf8f..ea80fb61 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -246,12 +246,18 @@ def _format_validation_error(tool_name: str, exc: ValidationError) -> str: 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. + """Bound the SDK shell tools' own token cap so a single command can't dump + unbounded output into history. The SDK truncates head+tail from this value. + + The configured cap is a ceiling: a missing value defaults to it, and a + larger model-supplied value is clamped down to it. A smaller explicit value + is respected, so the model can still ask for less. """ - if parsed.get("max_output_tokens") is None: - parsed["max_output_tokens"] = load_settings().context.tool_output_max_tokens + ceiling = load_settings().context.tool_output_max_tokens + requested = parsed.get("max_output_tokens") + parsed["max_output_tokens"] = ( + ceiling if not isinstance(requested, int) or requested > ceiling else requested + ) def _wrap_exec_command(tool: FunctionTool) -> FunctionTool: diff --git a/strix/report/writer.py b/strix/report/writer.py index 32fc961f..1b0a8a1d 100644 --- a/strix/report/writer.py +++ b/strix/report/writer.py @@ -10,7 +10,7 @@ import re import tempfile from datetime import UTC, datetime from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from pygments.lexers import PythonLexer, get_lexer_by_name, guess_lexer from pygments.lexers.special import TextLexer @@ -74,10 +74,10 @@ def resolve_lexer(language: str | None, code: str) -> Lexer: try: lexer = guess_lexer(code) except ClassNotFound: - return PythonLexer() + return cast("Lexer", PythonLexer()) # ``guess_lexer`` returns the plain-text lexer when it can't detect anything. if isinstance(lexer, TextLexer): - return PythonLexer() + return cast("Lexer", PythonLexer()) return lexer diff --git a/strix/tools/output_store.py b/strix/tools/output_store.py index 44a9fa9c..92390ce9 100644 --- a/strix/tools/output_store.py +++ b/strix/tools/output_store.py @@ -66,7 +66,11 @@ def bound_text(text: str, *, max_lines: int, max_bytes: int) -> str: if tail and _byte_len(tail) > half_bytes: tail = _take_suffix(tail, half_bytes) - dropped_lines = max(0, len(lines) - head_lines - tail_lines) + # Count kept lines from the final slices: the byte pass above may have + # dropped whole lines from head/tail, so deriving this from the original + # head_lines/tail_lines would undercount what was actually removed. + 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 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 30276562..7165378c 100644 --- a/tests/test_agent_factory_shell.py +++ b/tests/test_agent_factory_shell.py @@ -41,7 +41,7 @@ async def test_wrap_exec_command_defaults_shell_to_bash() -> None: @pytest.mark.asyncio -async def test_wrap_exec_command_preserves_explicit_output_cap() -> None: +async def test_wrap_exec_command_preserves_smaller_explicit_output_cap() -> None: captured: dict[str, str] = {} wrapped = factory._wrap_exec_command(_capturing_exec_tool(captured)) @@ -52,6 +52,20 @@ async def test_wrap_exec_command_preserves_explicit_output_cap() -> None: assert json.loads(captured["raw_input"])["max_output_tokens"] == 42 +@pytest.mark.asyncio +async def test_wrap_exec_command_clamps_oversized_explicit_output_cap() -> None: + captured: dict[str, str] = {} + wrapped = factory._wrap_exec_command(_capturing_exec_tool(captured)) + ceiling = load_settings().context.tool_output_max_tokens + + await wrapped.on_invoke_tool( + cast("Any", None), + json.dumps({"cmd": "echo hi", "max_output_tokens": ceiling * 100}), + ) + + assert json.loads(captured["raw_input"])["max_output_tokens"] == ceiling + + @pytest.mark.asyncio @pytest.mark.parametrize("shell", ["/bin/zsh", ""]) async def test_wrap_exec_command_preserves_explicit_shell(shell: str) -> None: diff --git a/tests/test_output_store.py b/tests/test_output_store.py index 23bcdf30..2dc1906f 100644 --- a/tests/test_output_store.py +++ b/tests/test_output_store.py @@ -2,6 +2,8 @@ from __future__ import annotations +import re + from strix.tools.output_store import bound_text @@ -44,3 +46,18 @@ def test_notice_reports_dropped_counts() -> None: assert "lines" in bounded assert "bytes" in bounded + + +def test_dropped_line_count_accounts_for_byte_trimming() -> None: + # A tight byte budget forces the byte pass to drop whole lines from the + # head/tail slices; the notice must count those, not just the middle. + text = "\n".join(f"line-{i}" for i in range(200)) + bounded = bound_text(text, max_lines=20, max_bytes=40) + + match = re.search(r"\[\.\.\. (\d+) lines", bounded) + assert match is not None, bounded + dropped = int(match.group(1)) + kept = [ln for ln in bounded.splitlines() if ln and "truncated" not in ln] + assert dropped == 200 - len(kept) + # The naive middle-only count (max_lines split evenly) would under-report. + assert dropped > 200 - 20