diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 1efac5dd..6a196294 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -388,34 +388,6 @@ 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 ' @@ -423,38 +395,6 @@ _SLEEP_HINT = ( ) -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} @@ -498,9 +438,10 @@ def _normalize_exec_args(parsed: dict[str, Any]) -> bool: 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. + # 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"] = _default_exec_yield_ms(parsed.get("cmd")) + 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 diff --git a/strix/agents/prompts/system_prompt.jinja b/strix/agents/prompts/system_prompt.jinja index 1b0ffe03..a7d6f8c2 100644 --- a/strix/agents/prompts/system_prompt.jinja +++ b/strix/agents/prompts/system_prompt.jinja @@ -187,6 +187,10 @@ EFFICIENCY TACTICS: `write_stdin(session_id=, 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), estimate its runtime and pass that as + `yield_time_ms` on the first `exec_command` — one call that waits beats a + backgrounded process you then poll for many turns. Default is 30s. - 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/settings.py b/strix/config/settings.py index 645212f0..43a319f5 100644 --- a/strix/config/settings.py +++ b/strix/config/settings.py @@ -115,8 +115,6 @@ class ShellSettings(BaseSettings): # 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" diff --git a/tests/test_agent_factory_shell.py b/tests/test_agent_factory_shell.py index 8cf356c8..1c8569be 100644 --- a/tests/test_agent_factory_shell.py +++ b/tests/test_agent_factory_shell.py @@ -151,16 +151,31 @@ async def test_wrap_exec_command_raises_default_yield_when_omitted() -> None: "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: +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_long_yield_ms + 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_long_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 -p- example.com", "yield_time_ms": 300_000}), + ) + + assert json.loads(captured["raw_input"])["yield_time_ms"] == 300_000 + + @pytest.mark.asyncio async def test_wrap_exec_command_preserves_explicit_yield() -> None: captured: dict[str, str] = {} @@ -174,7 +189,7 @@ async def test_wrap_exec_command_preserves_explicit_yield() -> None: @pytest.mark.asyncio -async def test_wrap_exec_command_unparsable_command_uses_plain_default() -> None: +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))