mirror of
https://github.com/usestrix/strix.git
synced 2026-08-16 09:26:39 +02:00
fix(context): clamp shell output cap and count byte-trimmed dropped lines
Treat tool_output_max_tokens as a ceiling so an explicit model-supplied cap can't exceed it, and derive the truncation notice's dropped-line count from the lines actually kept after the byte-trim pass. Also cast the pygments fallback lexer so it satisfies the resolve_lexer return type under the pre-commit mypy hook.
This commit is contained in:
+11
-5
@@ -246,12 +246,18 @@ def _format_validation_error(tool_name: str, exc: ValidationError) -> str:
|
|||||||
|
|
||||||
|
|
||||||
def _apply_shell_output_cap(parsed: dict[str, Any]) -> None:
|
def _apply_shell_output_cap(parsed: dict[str, Any]) -> None:
|
||||||
"""Default the SDK shell tools' own token cap so a single command can't
|
"""Bound the SDK shell tools' own token cap so a single command can't dump
|
||||||
dump unbounded output into history. The SDK truncates head+tail when the
|
unbounded output into history. The SDK truncates head+tail from this value.
|
||||||
model omits the field; respect an explicit model-supplied 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:
|
ceiling = load_settings().context.tool_output_max_tokens
|
||||||
parsed["max_output_tokens"] = 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:
|
def _wrap_exec_command(tool: FunctionTool) -> FunctionTool:
|
||||||
|
|||||||
@@ -10,7 +10,7 @@ import re
|
|||||||
import tempfile
|
import tempfile
|
||||||
from datetime import UTC, datetime
|
from datetime import UTC, datetime
|
||||||
from pathlib import Path
|
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 import PythonLexer, get_lexer_by_name, guess_lexer
|
||||||
from pygments.lexers.special import TextLexer
|
from pygments.lexers.special import TextLexer
|
||||||
@@ -74,10 +74,10 @@ def resolve_lexer(language: str | None, code: str) -> Lexer:
|
|||||||
try:
|
try:
|
||||||
lexer = guess_lexer(code)
|
lexer = guess_lexer(code)
|
||||||
except ClassNotFound:
|
except ClassNotFound:
|
||||||
return PythonLexer()
|
return cast("Lexer", PythonLexer())
|
||||||
# ``guess_lexer`` returns the plain-text lexer when it can't detect anything.
|
# ``guess_lexer`` returns the plain-text lexer when it can't detect anything.
|
||||||
if isinstance(lexer, TextLexer):
|
if isinstance(lexer, TextLexer):
|
||||||
return PythonLexer()
|
return cast("Lexer", PythonLexer())
|
||||||
return lexer
|
return lexer
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -66,7 +66,11 @@ def bound_text(text: str, *, max_lines: int, max_bytes: int) -> str:
|
|||||||
if tail and _byte_len(tail) > half_bytes:
|
if tail and _byte_len(tail) > half_bytes:
|
||||||
tail = _take_suffix(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))
|
dropped_bytes = max(0, total_bytes - _byte_len(head) - _byte_len(tail))
|
||||||
notice = _TRUNCATION_NOTICE.format(lines=dropped_lines, bytes=dropped_bytes)
|
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}"
|
return f"{head}\n\n{notice}\n\n{tail}" if tail else f"{head}\n\n{notice}"
|
||||||
|
|||||||
@@ -41,7 +41,7 @@ async def test_wrap_exec_command_defaults_shell_to_bash() -> None:
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@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] = {}
|
captured: dict[str, str] = {}
|
||||||
wrapped = factory._wrap_exec_command(_capturing_exec_tool(captured))
|
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
|
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.asyncio
|
||||||
@pytest.mark.parametrize("shell", ["/bin/zsh", ""])
|
@pytest.mark.parametrize("shell", ["/bin/zsh", ""])
|
||||||
async def test_wrap_exec_command_preserves_explicit_shell(shell: str) -> None:
|
async def test_wrap_exec_command_preserves_explicit_shell(shell: str) -> None:
|
||||||
|
|||||||
@@ -2,6 +2,8 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
|
|
||||||
from strix.tools.output_store import bound_text
|
from strix.tools.output_store import bound_text
|
||||||
|
|
||||||
|
|
||||||
@@ -44,3 +46,18 @@ def test_notice_reports_dropped_counts() -> None:
|
|||||||
|
|
||||||
assert "lines" in bounded
|
assert "lines" in bounded
|
||||||
assert "bytes" 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
|
||||||
|
|||||||
Reference in New Issue
Block a user