diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 977d01f19..83dc88280 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -6,7 +6,7 @@ import json import logging import re -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from agents.agent import ToolsToFinalOutputResult from agents.sandbox import SandboxAgent @@ -168,6 +168,11 @@ def _configure_chat_completions_filesystem_tools(toolset: Any) -> None: _CHARS_ESCAPE_RE = re.compile(r"\\(?:u[0-9a-fA-F]{4}|x[0-9a-fA-F]{2}|[0abtnvfr\\])") +_PTY_SESSION_NOT_FOUND_RE = re.compile(r"PTY session not found: (?P\d+)") +_WRITE_STDIN_DEAD_SESSION_RESULT_TEMPLATE = ( + "write_stdin: session {session_id} is already terminal/non-retryable " + "after PTY session not found." +) _CHARS_ESCAPE_MAP = { "\\\\": "\\", "\\n": "\n", @@ -234,19 +239,32 @@ async def invoke(ctx: Any, raw_input: str) -> Any: def _wrap_write_stdin(tool: FunctionTool) -> FunctionTool: invoke_tool = tool.on_invoke_tool + dead_pty_session_ids: set[int] = set() async def invoke(ctx: Any, raw_input: str) -> Any: + parsed: dict[str, Any] | None try: - parsed = json.loads(raw_input) + raw_parsed: Any = json.loads(raw_input) except json.JSONDecodeError: parsed = None - if isinstance(parsed, dict) and isinstance(parsed.get("chars"), str): - parsed["chars"] = _decode_chars_escape(parsed["chars"]) + else: + parsed = cast("dict[str, Any]", raw_parsed) if isinstance(raw_parsed, dict) else None + session_id = parsed.get("session_id") if isinstance(parsed, dict) else None + chars = parsed.get("chars") if isinstance(parsed, dict) else None + if isinstance(session_id, int) and chars == "" and session_id in dead_pty_session_ids: + return _WRITE_STDIN_DEAD_SESSION_RESULT_TEMPLATE.format(session_id=session_id) + if isinstance(parsed, dict) and isinstance(chars, str): + parsed["chars"] = _decode_chars_escape(chars) 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) + if isinstance(result, str): + match = _PTY_SESSION_NOT_FOUND_RE.search(result) + if match is not None: + dead_pty_session_ids.add(int(match.group("session_id"))) + 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 6de932190..7f2f2c551 100644 --- a/tests/test_agent_factory_shell.py +++ b/tests/test_agent_factory_shell.py @@ -11,6 +11,10 @@ from strix.agents import factory +_MISSING_PTY_SESSION_ID = 60418 +_MISSING_PTY_RESULT = f"write_stdin failed: PTY session not found: {_MISSING_PTY_SESSION_ID}" + + def _capturing_exec_tool(captured: dict[str, str]) -> FunctionTool: async def invoke(_ctx: Any, raw_input: str) -> str: captured["raw_input"] = raw_input @@ -24,6 +28,19 @@ async def invoke(_ctx: Any, raw_input: str) -> str: ) +def _capturing_write_stdin_tool(results: list[str], captured: list[str]) -> FunctionTool: + async def invoke(_ctx: Any, raw_input: str) -> str: + captured.append(raw_input) + return results.pop(0) + + 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] = {} @@ -49,3 +66,19 @@ async def test_wrap_exec_command_preserves_explicit_shell(shell: str) -> None: ) assert json.loads(captured["raw_input"])["shell"] == shell + + +@pytest.mark.asyncio +async def test_wrap_write_stdin_blocks_repeat_empty_poll_after_dead_pty() -> None: + captured: list[str] = [] + wrapped = factory._wrap_write_stdin( + _capturing_write_stdin_tool([_MISSING_PTY_RESULT], captured) + ) + raw_input = json.dumps({"session_id": _MISSING_PTY_SESSION_ID, "chars": ""}) + + first_result = await wrapped.on_invoke_tool(cast("Any", None), raw_input) + second_result = await wrapped.on_invoke_tool(cast("Any", None), raw_input) + + assert first_result == _MISSING_PTY_RESULT + assert "non-retryable" in second_result + assert len(captured) == 1