Skip to content
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
114 changes: 101 additions & 13 deletions strix/agents/factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
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,
Expand All @@ -33,6 +34,7 @@
list_notes,
update_note,
)
from strix.tools.output_store import bound_text
from strix.tools.proxy.tools import (
list_requests,
list_sitemap,
Expand Down Expand Up @@ -103,16 +105,49 @@ 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:
invoke_tool = tool.on_invoke_tool

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)
Expand All @@ -127,7 +162,7 @@ async def invoke(ctx: Any, raw_input: str) -> Any:
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)
Expand Down Expand Up @@ -159,12 +194,43 @@ async def approve(ctx: Any, args: dict[str, Any], call_id: str) -> bool:
)


def _configure_chat_completions_filesystem_tools(toolset: Any) -> None:
def _bound_custom_tool(tool: CustomTool) -> CustomTool:
"""Bound a native ``CustomTool`` result in place.

Chat-completions mode converts filesystem ``CustomTool``s to ``FunctionTool``s
(which bounds the result), but the Responses path keeps them native, so a
large ``read_file``/directory listing would otherwise append unbounded text
to history. Wrap ``on_invoke_tool`` so the same head+tail bound applies.
"""
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
return tool


def _configure_filesystem_tools(toolset: Any, *, chat_completions: bool) -> None:
for name, tool in vars(toolset).items():
if isinstance(tool, CustomTool):
setattr(toolset, name, _custom_tool_as_function_tool(tool))
if chat_completions:
if isinstance(tool, CustomTool):
setattr(toolset, name, _custom_tool_as_function_tool(tool))
elif isinstance(tool, FunctionTool):
setattr(toolset, name, _function_tool_with_error_result(tool))
# Responses-API path: keep tools native but still bound their output so
# filesystem reads can't exhaust the context window on later turns.
elif isinstance(tool, CustomTool):
setattr(toolset, name, _bound_custom_tool(tool))
elif isinstance(tool, FunctionTool):
setattr(toolset, name, _function_tool_with_error_result(tool))
setattr(toolset, name, _with_bounded_result(tool))


def _make_filesystem_configurator(*, chat_completions: bool) -> Any:
def configure(toolset: Any) -> None:
_configure_filesystem_tools(toolset, chat_completions=chat_completions)

return configure


_CHARS_ESCAPE_RE = re.compile(r"\\(?:u[0-9a-fA-F]{4}|x[0-9a-fA-F]{2}|[0abtnvfr\\])")
Expand Down Expand Up @@ -205,6 +271,21 @@ 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:
"""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.
"""
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:
invoke_tool = tool.on_invoke_tool

Expand All @@ -213,8 +294,10 @@ async def invoke(ctx: Any, raw_input: str) -> Any:
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)
Expand All @@ -240,8 +323,10 @@ async def invoke(ctx: Any, raw_input: str) -> Any:
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)
Expand Down Expand Up @@ -440,6 +525,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
]
Comment thread
greptile-apps[bot] marked this conversation as resolved.

logger.info(
"Built %s agent '%s' (skills=%d, tools=%d, scan_mode=%s, whitebox=%s)",
Expand All @@ -459,8 +547,8 @@ def build_strix_agent(
model=None,
capabilities=[
Filesystem(
configure_tools=(
_configure_chat_completions_filesystem_tools if chat_completions_tools else None
configure_tools=_make_filesystem_configurator(
chat_completions=chat_completions_tools,
),
),
Shell(
Expand Down
2 changes: 2 additions & 0 deletions strix/config/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
persist_current,
)
from strix.config.settings import (
ContextSettings,
DedupeSettings,
IntegrationSettings,
LlmSettings,
Expand All @@ -27,6 +28,7 @@


__all__ = [
"ContextSettings",
"DedupeSettings",
"IntegrationSettings",
"LlmSettings",
Expand Down
18 changes: 18 additions & 0 deletions strix/config/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
6 changes: 3 additions & 3 deletions strix/report/writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
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
Expand Down Expand Up @@ -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


Expand Down
85 changes: 85 additions & 0 deletions strix/tools/output_store.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
"""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.
``max_bytes`` bounds the *entire* joined result, notice and separators
included.
"""
lines = text.split("\n")
total_bytes = _byte_len(text)
if len(lines) <= max_lines and total_bytes <= max_bytes:
return text

# Reserve room for the notice and its two blank-line separators so the
# head+tail slices can't consume the whole budget and push the persisted
# value over max_bytes. Upper-bound the notice with the largest possible
# counts; the real notice is never longer. ``+ 4`` covers the separators.
notice_overhead = _byte_len(_TRUNCATION_NOTICE.format(lines=len(lines), bytes=total_bytes)) + 4
byte_budget = max(2, max_bytes - notice_overhead)
Comment thread
greptile-apps[bot] marked this conversation as resolved.

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, byte_budget // 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)

# 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}"
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Loading