Merge 4dcd543db7 into 717ffc8f4c

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)

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 <<EOF`) driven through
+  `write_stdin`; then run the file as a normal one-shot command.
+- Don't blind-`sleep` to wait for a background job. Poll it with
+  `write_stdin(session_id=<id>, chars="", yield_time_ms=<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

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",

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)

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