diff --git a/strix/agents/factory.py b/strix/agents/factory.py index d284dfd3..e6b54323 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 @@ -424,6 +425,76 @@ def _apply_shell_output_cap(parsed: dict[str, Any]) -> None: ) +_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." +) + + +_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. The agent + # asks for a longer yield itself when it expects a command to run longer. + if "yield_time_ms" not in parsed: + parsed["yield_time_ms"] = load_settings().shell_tools.exec_yield_ms + 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 @@ -432,13 +503,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: @@ -448,6 +518,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 @@ -462,9 +535,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 66eaa7b8..60083c1f 100644 --- a/strix/agents/prompts/system_prompt.jinja +++ b/strix/agents/prompts/system_prompt.jinja @@ -196,7 +196,19 @@ 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. +- When you expect a command to take a while (a scan like `nmap`/`nuclei`/`ffuf`, + a build, a long crawl), pass the time you expect to need as `yield_time_ms` + on the first `exec_command` — one call that waits beats a backgrounded + process you then poll for many turns. The shell yields at most 30s per call + (the default), so anything slower than that still backgrounds: harvest it + with one `write_stdin(chars="")` poll per 30s rather than many short ones. - 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..8b8144e3 100644 --- a/strix/config/settings.py +++ b/strix/config/settings.py @@ -4,6 +4,7 @@ from __future__ import annotations from typing import Literal +from agents.sandbox.session.pty_types import PTY_YIELD_TIME_MS_MAX from pydantic import AliasChoices, Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -102,6 +103,40 @@ 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. + + The SDK's PTY layer clamps every yield to ``PTY_YIELD_TIME_MS_MAX``, so both + yields are validated against that ceiling instead of being silently reduced. + """ + + model_config = _BASE_CONFIG + + # Default yield for exec_command when the model omits yield_time_ms. The + # default sits at the ceiling: waiting is cheaper than another poll turn. + exec_yield_ms: int = Field( + default=PTY_YIELD_TIME_MS_MAX, + gt=0, + le=PTY_YIELD_TIME_MS_MAX, + alias="STRIX_SHELL_EXEC_YIELD_MS", + ) + # Default yield for an empty (polling) write_stdin call. The SDK already + # floors an empty poll at 5s; this trades a little latency for far fewer turns. + write_stdin_poll_yield_ms: int = Field( + default=20_000, + gt=0, + le=PTY_YIELD_TIME_MS_MAX, + 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 +186,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..ba913813 100644 --- a/tests/test_agent_factory_shell.py +++ b/tests/test_agent_factory_shell.py @@ -7,10 +7,13 @@ from types import SimpleNamespace from typing import Any, cast import pytest +from agents.sandbox.errors import InvalidManifestPathError +from agents.sandbox.session.pty_types import PTY_YIELD_TIME_MS_MAX from agents.tool import CustomTool, FunctionTool +from pydantic import BaseModel, ValidationError from strix.agents import factory -from strix.config import load_settings +from strix.config import ShellSettings, load_settings def _capturing_exec_tool(captured: dict[str, str]) -> FunctionTool: @@ -26,6 +29,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 +131,245 @@ 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_default_does_not_depend_on_the_binary(cmd: str) -> None: + """The wrapper never guesses a command's runtime: the agent asks for a + longer yield itself when it expects one.""" + 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_yield_ms + assert json.loads(captured["raw_input"])["yield_time_ms"] == expected + + +@pytest.mark.asyncio +async def test_wrap_exec_command_preserves_longer_explicit_yield() -> None: + """A slow command gets the yield the agent asked for, not a guessed one. + + The SDK's PTY layer clamps anything above 30s, so a longer wait than that + cannot be bought with a bigger argument.""" + 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 -p- example.com", "yield_time_ms": 25_000}), + ) + + assert json.loads(captured["raw_input"])["yield_time_ms"] == 25_000 + + +@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_still_gets_the_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" + + +# --- error formatting -------------------------------------------------------- + + +class _ExecArgs(BaseModel): + cmd: str + + +def _raising_tool(name: str, exc: Exception) -> FunctionTool: + async def invoke(_ctx: Any, _raw_input: str) -> str: + raise exc + + return FunctionTool( + name=name, + description="test tool", + params_json_schema={"type": "object", "properties": {}}, + on_invoke_tool=invoke, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("name", "wrap"), + [("exec_command", factory._wrap_exec_command), ("write_stdin", factory._wrap_write_stdin)], +) +async def test_validation_error_is_rendered_as_a_message(name: str, wrap: Any) -> None: + try: + _ExecArgs.model_validate({}) + except ValidationError as exc: + validation_error = exc + + wrapped = wrap(_raising_tool(name, validation_error)) + result = await wrapped.on_invoke_tool(cast("Any", None), json.dumps({"cmd": "echo hi"})) + + assert isinstance(result, str) + assert result.startswith(f"{name}: invalid arguments — ") + assert "cmd" in result + + +@pytest.mark.asyncio +async def test_invalid_workdir_is_rendered_as_a_message() -> None: + exc = InvalidManifestPathError(rel="../etc", reason="escape_root") + wrapped = factory._wrap_exec_command(_raising_tool("exec_command", exc)) + + result = await wrapped.on_invoke_tool(cast("Any", None), json.dumps({"cmd": "ls"})) + + assert isinstance(result, str) + assert "workdir must be a path inside /workspace" in result + assert "'../etc'" in result + + +@pytest.mark.parametrize("field", ["exec_yield_ms", "write_stdin_poll_yield_ms"]) +def test_shell_settings_reject_a_yield_above_the_pty_ceiling(field: str) -> None: + """A yield the PTY layer would clamp is a misconfiguration, not a longer wait.""" + with pytest.raises(ValidationError): + ShellSettings(**{field: PTY_YIELD_TIME_MS_MAX + 1}) + + assert getattr(ShellSettings(**{field: PTY_YIELD_TIME_MS_MAX}), field) == PTY_YIELD_TIME_MS_MAX