From 6b9318ae07aff2348c0bc48b1d73b4a7c47d44de Mon Sep 17 00:00:00 2001 From: Alex Schapiro Date: Tue, 25 Aug 2026 15:30:20 +0000 Subject: [PATCH] perf(shell): raise default yield times so agents stop polling backgrounded shells The SDK yields after only 250ms on a `write_stdin` poll and 10s on `exec_command`, so agents burn many turns re-polling a backgrounded process for almost no output. Each poll costs a full LLM turn, which makes even trivial commands take minutes of wall time. Raise the defaults in the existing `exec_command` / `write_stdin` wrappers: - an empty-`chars` `write_stdin` (a poll, not input) yields 20s instead of 250ms, so one poll returns a meaningful result - `exec_command` yields 30s by default, and 120s for known long-running security binaries matched on the leading binary of the command - a bare `sleep N` hand-wait is clamped to 60s and annotated with a hint pointing at `write_stdin(chars="")`, which returns as soon as there is output or the process exits Every override is skipped when the model passes `yield_time_ms` explicitly, and the new values are configurable through `STRIX_SHELL_*` env vars. Command parsing fails open: an unparsable command just gets the plain default, and a `sleep` inside a compound command is never rewritten. --- strix/agents/factory.py | 144 +++++++++++++++++- strix/agents/prompts/system_prompt.jinja | 8 +- strix/config/__init__.py | 2 + strix/config/settings.py | 24 +++ tests/test_agent_factory_shell.py | 177 +++++++++++++++++++++++ 5 files changed, 347 insertions(+), 8 deletions(-) diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 3b56a55f..1efac5dd 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -7,6 +7,7 @@ import inspect import json import logging import re +import shlex from typing import TYPE_CHECKING, Any from agents.agent import ToolsToFinalOutputResult @@ -387,6 +388,135 @@ def _apply_shell_output_cap(parsed: dict[str, Any]) -> None: ) +# Binaries that routinely run for minutes; they get a bigger exec yield so the +# agent gets a result in one call instead of polling a backgrounded process. +_LONG_RUNNING_BINARIES: frozenset[str] = frozenset( + { + "amass", + "dirsearch", + "feroxbuster", + "ffuf", + "gobuster", + "httpx", + "katana", + "masscan", + "nikto", + "nmap", + "nuclei", + "sqlmap", + "subfinder", + "wpscan", + } +) + +# Wrapper commands that precede the real binary; skip them when parsing. +_COMMAND_PREFIXES: frozenset[str] = frozenset( + {"command", "doas", "env", "nice", "nohup", "stdbuf", "sudo", "time"} +) + +_ENV_ASSIGN_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=") + +_SLEEP_HINT = ( + "\n\n[strix] To wait on a background job, prefer " + 'write_stdin(session_id=..., chars="", yield_time_ms=...), which returns as ' + "soon as there is new output or the process exits — better than a blind sleep." +) + + +def _leading_binary(cmd: str) -> str | None: + """Best-effort name of the first real binary in ``cmd``. + + Skips env-var assignments and wrapper prefixes (``sudo``/``env``/…) and + stops at the first token of a pipeline. Returns ``None`` when the command + cannot be parsed so callers can fall back to the plain default.""" + try: + tokens = shlex.split(cmd) + except ValueError: + return None + idx = 0 + while idx < len(tokens): + token = tokens[idx] + if _ENV_ASSIGN_RE.match(token): + idx += 1 + continue + if token in _COMMAND_PREFIXES: + idx += 1 + while idx < len(tokens) and tokens[idx].startswith("-"): + idx += 1 + continue + return token.rsplit("/", 1)[-1] or None + return None + + +def _default_exec_yield_ms(cmd: Any) -> int: + settings = load_settings().shell_tools + if isinstance(cmd, str) and _leading_binary(cmd) in _LONG_RUNNING_BINARIES: + return settings.exec_long_yield_ms + return settings.exec_yield_ms + + +_SLEEP_DURATION_RE = re.compile(r"^(\d+(?:\.\d+)?)([smhd]?)$") +_SLEEP_UNIT_SECONDS = {"": 1, "s": 1, "m": 60, "h": 3600, "d": 86400} + + +def _pure_sleep_seconds(cmd: Any) -> float | None: + """Total seconds a command sleeps, but only when it is *nothing but* a + ``sleep`` (``sleep 30``, ``sleep 1m 30s``). Returns ``None`` for anything + compound so an embedded sleep is never rewritten.""" + if not isinstance(cmd, str): + return None + try: + tokens = shlex.split(cmd) + except ValueError: + return None + if len(tokens) < 2 or tokens[0] != "sleep": + return None + total = 0.0 + for token in tokens[1:]: + match = _SLEEP_DURATION_RE.match(token) + if match is None: + return None + total += float(match.group(1)) * _SLEEP_UNIT_SECONDS[match.group(2)] + return total + + +def _apply_sleep_guard(parsed: dict[str, Any]) -> bool: + """Clamp an absurd bare ``sleep`` to the configured cap. Returns ``True`` + when the command is a pure sleep so the caller can append a hint.""" + seconds = _pure_sleep_seconds(parsed.get("cmd")) + if seconds is None: + return False + cap = load_settings().shell_tools.max_sleep_seconds + if seconds > cap: + parsed["cmd"] = f"sleep {cap}" + return True + + +def _normalize_exec_args(parsed: dict[str, Any]) -> bool: + """Apply Strix's ``exec_command`` defaults. Returns ``True`` when the + command is a bare sleep so the caller can append a hint.""" + if "shell" not in parsed: + parsed["shell"] = "bash" + # Raise the yield above the SDK's 10s so a command returns in one call + # instead of getting backgrounded and then polled turn after turn. + if "yield_time_ms" not in parsed: + parsed["yield_time_ms"] = _default_exec_yield_ms(parsed.get("cmd")) + is_sleep = _apply_sleep_guard(parsed) + _apply_shell_output_cap(parsed) + return is_sleep + + +def _normalize_write_stdin_args(parsed: dict[str, Any]) -> None: + """Apply Strix's ``write_stdin`` defaults.""" + if isinstance(parsed.get("chars"), str): + parsed["chars"] = _decode_chars_escape(parsed["chars"]) + # An empty ``chars`` is a poll, not input: yield long enough for a + # meaningful result unless the model asked for a specific wait. + if not parsed.get("chars") and "yield_time_ms" not in parsed: + parsed["yield_time_ms"] = load_settings().shell_tools.write_stdin_poll_yield_ms + _apply_shell_output_cap(parsed) + + def _wrap_exec_command(tool: FunctionTool) -> FunctionTool: invoke_tool = tool.on_invoke_tool @@ -395,13 +525,12 @@ def _wrap_exec_command(tool: FunctionTool) -> FunctionTool: parsed = json.loads(raw_input) except (json.JSONDecodeError, TypeError): parsed = None + is_sleep = False if isinstance(parsed, dict): - if "shell" not in parsed: - parsed["shell"] = "bash" - _apply_shell_output_cap(parsed) + is_sleep = _normalize_exec_args(parsed) raw_input = json.dumps(parsed) try: - return await invoke_tool(ctx, raw_input) + result = await invoke_tool(ctx, raw_input) except ValidationError as exc: return _format_validation_error(tool.name, exc) except InvalidManifestPathError as exc: @@ -411,6 +540,9 @@ def _wrap_exec_command(tool: FunctionTool) -> FunctionTool: "(or omitted to use the turn's cwd). " f"Got: {rel!r}." ) + if is_sleep and isinstance(result, str): + return result + _SLEEP_HINT + return result tool.on_invoke_tool = invoke return tool @@ -425,9 +557,7 @@ def _wrap_write_stdin(tool: FunctionTool) -> FunctionTool: except json.JSONDecodeError: parsed = None if isinstance(parsed, dict): - if isinstance(parsed.get("chars"), str): - parsed["chars"] = _decode_chars_escape(parsed["chars"]) - _apply_shell_output_cap(parsed) + _normalize_write_stdin_args(parsed) raw_input = json.dumps(parsed) try: return await invoke_tool(ctx, raw_input) diff --git a/strix/agents/prompts/system_prompt.jinja b/strix/agents/prompts/system_prompt.jinja index 95590394..1b0ffe03 100644 --- a/strix/agents/prompts/system_prompt.jinja +++ b/strix/agents/prompts/system_prompt.jinja @@ -180,7 +180,13 @@ EFFICIENCY TACTICS: run them with `python3 script.py`. For one-off snippets, `python3 -c` or a here-document is acceptable, but avoid deeply nested quotes/parentheses — if a snippet needs complex quoting or is more than a few lines, write it to a - file first to prevent syntax errors. + file first to prevent syntax errors. Write scripts with `apply_patch`, not an + interactive blocking heredoc (`cat > exploit.py <, chars="", yield_time_ms=)`, which returns + the instant there is new output or the process exits, or give the original + `exec_command` a bigger `yield_time_ms` up front so it doesn't background. - Before importing a third-party Python library, make sure it is installed. The sandbox's `python3` runs inside a preconfigured virtualenv that ships `requests`, `httpx`, `beautifulsoup4` (bs4), `lxml`, `pyjwt`, and diff --git a/strix/config/__init__.py b/strix/config/__init__.py index f21fdab6..6680f49e 100644 --- a/strix/config/__init__.py +++ b/strix/config/__init__.py @@ -23,6 +23,7 @@ from strix.config.settings import ( LlmSettings, RuntimeSettings, Settings, + ShellSettings, TelemetrySettings, ) @@ -34,6 +35,7 @@ __all__ = [ "LlmSettings", "RuntimeSettings", "Settings", + "ShellSettings", "TelemetrySettings", "apply_config_override", "load_settings", diff --git a/strix/config/settings.py b/strix/config/settings.py index 42a2c97e..645212f0 100644 --- a/strix/config/settings.py +++ b/strix/config/settings.py @@ -102,6 +102,29 @@ class ContextSettings(BaseSettings): ) +class ShellSettings(BaseSettings): + """Yield-time defaults for the SDK shell tools. + + Agents spend many turns polling backgrounded shells because the SDK yields + after only 250ms on a ``write_stdin`` poll and 10s on ``exec_command``. + Raising these defaults lets one call return a meaningful result instead of a + no-op round-trip. An explicit ``yield_time_ms`` from the model always wins. + """ + + model_config = _BASE_CONFIG + + # Default yield for exec_command when the model omits yield_time_ms. + exec_yield_ms: int = Field(default=30_000, gt=0, alias="STRIX_SHELL_EXEC_YIELD_MS") + # Larger default for known long-running security binaries. + exec_long_yield_ms: int = Field(default=120_000, gt=0, alias="STRIX_SHELL_EXEC_LONG_YIELD_MS") + # Default yield for an empty (polling) write_stdin call. + write_stdin_poll_yield_ms: int = Field( + default=20_000, gt=0, alias="STRIX_SHELL_WRITE_STDIN_POLL_YIELD_MS" + ) + # Cap on a bare `sleep N` hand-wait (seconds); larger sleeps are clamped. + max_sleep_seconds: int = Field(default=60, gt=0, alias="STRIX_SHELL_MAX_SLEEP_SECONDS") + + class RuntimeSettings(BaseSettings): model_config = _BASE_CONFIG @@ -151,6 +174,7 @@ class Settings(BaseSettings): dedupe: DedupeSettings = Field(default_factory=DedupeSettings) runtime: RuntimeSettings = Field(default_factory=RuntimeSettings) context: ContextSettings = Field(default_factory=ContextSettings) + shell_tools: ShellSettings = Field(default_factory=ShellSettings) telemetry: TelemetrySettings = Field(default_factory=TelemetrySettings) integrations: IntegrationSettings = Field(default_factory=IntegrationSettings) viewer: ViewerSettings = Field(default_factory=ViewerSettings) diff --git a/tests/test_agent_factory_shell.py b/tests/test_agent_factory_shell.py index 6bef1211..63023769 100644 --- a/tests/test_agent_factory_shell.py +++ b/tests/test_agent_factory_shell.py @@ -26,6 +26,19 @@ def _capturing_exec_tool(captured: dict[str, str]) -> FunctionTool: ) +def _capturing_write_stdin_tool(captured: dict[str, str]) -> FunctionTool: + async def invoke(_ctx: Any, raw_input: str) -> str: + captured["raw_input"] = raw_input + return "ok" + + return FunctionTool( + name="write_stdin", + description="test tool", + params_json_schema={"type": "object", "properties": {}}, + on_invoke_tool=invoke, + ) + + @pytest.mark.asyncio async def test_wrap_exec_command_defaults_shell_to_bash() -> None: captured: dict[str, str] = {} @@ -115,3 +128,167 @@ def test_function_tools_are_result_bounded() -> None: by_name = {t.name: t for t in agent.tools} assert getattr(by_name["think"], "_strix_bounded", False) is True + + +# --- yield-time defaults: exec_command -------------------------------------- + + +@pytest.mark.asyncio +async def test_wrap_exec_command_raises_default_yield_when_omitted() -> 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 ok"})) + + expected = load_settings().shell_tools.exec_yield_ms + assert json.loads(captured["raw_input"])["yield_time_ms"] == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "cmd", + ["nmap -sV example.com", "sudo nmap -sV example.com", "PROXY=1 ffuf -u http://x"], +) +async def test_wrap_exec_command_uses_long_yield_for_known_binaries(cmd: str) -> 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": cmd})) + + expected = load_settings().shell_tools.exec_long_yield_ms + assert json.loads(captured["raw_input"])["yield_time_ms"] == expected + + +@pytest.mark.asyncio +async def test_wrap_exec_command_preserves_explicit_yield() -> 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": "nmap example.com", "yield_time_ms": 500}) + ) + + assert json.loads(captured["raw_input"])["yield_time_ms"] == 500 + + +@pytest.mark.asyncio +async def test_wrap_exec_command_unparsable_command_uses_plain_default() -> 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": 'nmap "unterminated'})) + + expected = load_settings().shell_tools.exec_yield_ms + assert json.loads(captured["raw_input"])["yield_time_ms"] == expected + + +# --- sleep guard ------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_wrap_exec_command_caps_absurd_sleep_and_hints() -> None: + captured: dict[str, str] = {} + wrapped = factory._wrap_exec_command(_capturing_exec_tool(captured)) + cap = load_settings().shell_tools.max_sleep_seconds + + result = await wrapped.on_invoke_tool(cast("Any", None), json.dumps({"cmd": "sleep 3600"})) + + assert json.loads(captured["raw_input"])["cmd"] == f"sleep {cap}" + assert isinstance(result, str) + assert "write_stdin" in result + + +@pytest.mark.asyncio +async def test_wrap_exec_command_short_sleep_kept_but_hinted() -> None: + captured: dict[str, str] = {} + wrapped = factory._wrap_exec_command(_capturing_exec_tool(captured)) + + result = await wrapped.on_invoke_tool(cast("Any", None), json.dumps({"cmd": "sleep 5"})) + + assert json.loads(captured["raw_input"])["cmd"] == "sleep 5" + assert isinstance(result, str) + assert "write_stdin" in result + + +@pytest.mark.asyncio +async def test_wrap_exec_command_leaves_compound_sleep_untouched() -> None: + captured: dict[str, str] = {} + wrapped = factory._wrap_exec_command(_capturing_exec_tool(captured)) + + result = await wrapped.on_invoke_tool( + cast("Any", None), json.dumps({"cmd": "sleep 3600 && curl http://x"}) + ) + + assert json.loads(captured["raw_input"])["cmd"] == "sleep 3600 && curl http://x" + assert result == "ok" + + +@pytest.mark.asyncio +async def test_wrap_exec_command_malformed_input_passes_through() -> None: + captured: dict[str, str] = {} + wrapped = factory._wrap_exec_command(_capturing_exec_tool(captured)) + + assert await wrapped.on_invoke_tool(cast("Any", None), "not json") == "ok" + assert captured["raw_input"] == "not json" + + +# --- yield-time defaults: write_stdin --------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("payload", [{"session_id": 1}, {"session_id": 1, "chars": ""}]) +async def test_wrap_write_stdin_empty_poll_gets_raised_default(payload: dict[str, Any]) -> None: + captured: dict[str, str] = {} + wrapped = factory._wrap_write_stdin(_capturing_write_stdin_tool(captured)) + + await wrapped.on_invoke_tool(cast("Any", None), json.dumps(payload)) + + expected = load_settings().shell_tools.write_stdin_poll_yield_ms + assert json.loads(captured["raw_input"])["yield_time_ms"] == expected + + +@pytest.mark.asyncio +async def test_wrap_write_stdin_empty_poll_preserves_explicit_yield() -> None: + captured: dict[str, str] = {} + wrapped = factory._wrap_write_stdin(_capturing_write_stdin_tool(captured)) + + await wrapped.on_invoke_tool( + cast("Any", None), json.dumps({"session_id": 1, "chars": "", "yield_time_ms": 250}) + ) + + assert json.loads(captured["raw_input"])["yield_time_ms"] == 250 + + +@pytest.mark.asyncio +async def test_wrap_write_stdin_non_empty_chars_keeps_snappy() -> None: + captured: dict[str, str] = {} + wrapped = factory._wrap_write_stdin(_capturing_write_stdin_tool(captured)) + + await wrapped.on_invoke_tool( + cast("Any", None), json.dumps({"session_id": 1, "chars": "print(1)\\n"}) + ) + + parsed = json.loads(captured["raw_input"]) + assert "yield_time_ms" not in parsed + assert parsed["chars"] == "print(1)\n" + + +@pytest.mark.asyncio +async def test_wrap_write_stdin_non_empty_chars_preserves_explicit_yield() -> None: + captured: dict[str, str] = {} + wrapped = factory._wrap_write_stdin(_capturing_write_stdin_tool(captured)) + + await wrapped.on_invoke_tool( + cast("Any", None), json.dumps({"session_id": 1, "chars": "y\\n", "yield_time_ms": 100}) + ) + + assert json.loads(captured["raw_input"])["yield_time_ms"] == 100 + + +@pytest.mark.asyncio +async def test_wrap_write_stdin_malformed_input_passes_through() -> None: + captured: dict[str, str] = {} + wrapped = factory._wrap_write_stdin(_capturing_write_stdin_tool(captured)) + + assert await wrapped.on_invoke_tool(cast("Any", None), "not json") == "ok" + assert captured["raw_input"] == "not json"