diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 82b633a2..8e3790e5 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -13,6 +13,7 @@ from agents.agent import ToolsToFinalOutputResult from agents.sandbox import SandboxAgent from agents.sandbox.capabilities import Filesystem, Shell from agents.sandbox.errors import InvalidManifestPathError +from agents.sandbox.session.pty_types import PTY_YIELD_TIME_MS_MAX from agents.tool import CustomTool, FunctionTool, Tool from pydantic import ValidationError @@ -441,8 +442,43 @@ def _apply_shell_output_cap(parsed: dict[str, Any]) -> None: ) +_EXEC_YIELD_CAP_S = PTY_YIELD_TIME_MS_MAX // 1000 + +_EXEC_COMMAND_DESCRIPTION_NOTE = ( + f"The wait for output is capped at {_EXEC_YIELD_CAP_S} s whatever yield_time_ms says. " + "A non-TTY command still running at the yield cannot be read later (write_stdin fails " + "on it and its output is lost): run anything longer in the background with its output " + "redirected to a file (`cmd > /tmp/out.log 2>&1 &`, then read the log), or start it with " + 'tty=true and poll with write_stdin(session_id=, chars="").' +) + +_RUNNING_NON_TTY_NOTICE = ( + "Note: this non-TTY process cannot be polled and its later output is lost. Rerun it in " + "the background (`cmd > /tmp/out.log 2>&1 &`, then read the log) or with tty=true." +) + +_PROCESS_RUNNING_RE = re.compile(r"^Process running with session ID \d+$", re.MULTILINE) + + +def _annotate_running_non_tty(parsed: dict[str, Any], result: Any) -> Any: + if parsed.get("tty") or not isinstance(result, str): + return result + match = _PROCESS_RUNNING_RE.search(result) + if match is None: + return result + notice = _RUNNING_NON_TTY_NOTICE + requested = parsed.get("yield_time_ms") + if isinstance(requested, int | float) and requested > PTY_YIELD_TIME_MS_MAX: + notice = ( + f"Note: yield_time_ms={int(requested)} was capped at {PTY_YIELD_TIME_MS_MAX} " + f"({_EXEC_YIELD_CAP_S} s). " + notice.removeprefix("Note: ") + ) + return f"{result[: match.end()]}\n{notice}{result[match.end() :]}" + + def _wrap_exec_command(tool: FunctionTool) -> FunctionTool: invoke_tool = tool.on_invoke_tool + tool.description = f"{tool.description.rstrip()} {_EXEC_COMMAND_DESCRIPTION_NOTE}" async def invoke(ctx: Any, raw_input: str) -> Any: try: @@ -455,7 +491,7 @@ def _wrap_exec_command(tool: FunctionTool) -> FunctionTool: _apply_shell_output_cap(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: @@ -465,6 +501,9 @@ def _wrap_exec_command(tool: FunctionTool) -> FunctionTool: "(or omitted to use the turn's cwd). " f"Got: {rel!r}." ) + if isinstance(parsed, dict): + return _annotate_running_non_tty(parsed, result) + return result tool.on_invoke_tool = invoke return tool diff --git a/tests/test_agent_factory_shell.py b/tests/test_agent_factory_shell.py index 0109bf52..a663c9a5 100644 --- a/tests/test_agent_factory_shell.py +++ b/tests/test_agent_factory_shell.py @@ -80,6 +80,79 @@ async def test_wrap_exec_command_preserves_explicit_shell(shell: str) -> None: assert json.loads(captured["raw_input"])["shell"] == shell +def _running_exec_tool(result: str) -> FunctionTool: + async def invoke(_ctx: Any, _raw_input: str) -> str: + return result + + return FunctionTool( + name="exec_command", + description="Runs a command.", + params_json_schema={"type": "object", "properties": {}}, + on_invoke_tool=invoke, + ) + + +_RUNNING_RESULT = ( + "Chunk ID: abc123\nWall time: 30.0012 seconds\nProcess running with session ID 7\nOutput:\n" +) + + +def test_wrap_exec_command_describes_the_yield_cap() -> None: + wrapped = factory._wrap_exec_command(_running_exec_tool("ok")) + + assert wrapped.description.startswith("Runs a command. The wait for output is capped at 30 s") + assert "cannot be read later" in wrapped.description + assert "tty=true" in wrapped.description + + +@pytest.mark.asyncio +async def test_wrap_exec_command_flags_a_running_non_tty_process() -> None: + wrapped = factory._wrap_exec_command(_running_exec_tool(_RUNNING_RESULT)) + + result = await wrapped.on_invoke_tool( + cast("Any", None), json.dumps({"cmd": "sleep 60; curl x", "yield_time_ms": 90000}) + ) + + assert result == ( + "Chunk ID: abc123\nWall time: 30.0012 seconds\n" + "Process running with session ID 7\n" + "Note: yield_time_ms=90000 was capped at 30000 (30 s). this non-TTY process cannot be " + "polled and its later output is lost. Rerun it in the background " + "(`cmd > /tmp/out.log 2>&1 &`, then read the log) or with tty=true.\nOutput:\n" + ) + + +@pytest.mark.asyncio +async def test_wrap_exec_command_flags_a_running_non_tty_process_within_the_cap() -> None: + wrapped = factory._wrap_exec_command(_running_exec_tool(_RUNNING_RESULT)) + + result = await wrapped.on_invoke_tool(cast("Any", None), json.dumps({"cmd": "sleep 60"})) + + assert "Note: this non-TTY process cannot be polled" in result + assert "was capped" not in result + + +@pytest.mark.asyncio +async def test_wrap_exec_command_leaves_a_running_tty_process_alone() -> None: + wrapped = factory._wrap_exec_command(_running_exec_tool(_RUNNING_RESULT)) + + result = await wrapped.on_invoke_tool( + cast("Any", None), json.dumps({"cmd": "python3 -m http.server", "tty": True}) + ) + + assert result == _RUNNING_RESULT + + +@pytest.mark.asyncio +async def test_wrap_exec_command_leaves_an_exited_process_alone() -> None: + exited = "Chunk ID: abc123\nWall time: 0.1 seconds\nProcess exited with code 0\nOutput:\nhi" + wrapped = factory._wrap_exec_command(_running_exec_tool(exited)) + + result = await wrapped.on_invoke_tool(cast("Any", None), json.dumps({"cmd": "echo hi"})) + + assert result == exited + + @pytest.mark.asyncio async def test_responses_filesystem_custom_tool_output_is_bounded() -> None: async def invoke(_ctx: Any, _inp: str) -> str: