Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 23 additions & 5 deletions strix/agents/factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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<session_id>\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",
Expand Down Expand Up @@ -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
Expand Down
33 changes: 33 additions & 0 deletions tests/test_agent_factory_shell.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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] = {}
Expand All @@ -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