From a6aec4e712bb5dcab8f0bbd4fea84d0fd471f39a Mon Sep 17 00:00:00 2001 From: Ousama Ben Younes Date: Thu, 23 Jul 2026 23:11:11 +0000 Subject: [PATCH 1/2] fix(agents): stop repeat polls of dead PTY sessions --- strix/agents/factory.py | 38 ++++++++++++++-- tests/test_agent_factory_shell.py | 76 +++++++++++++++++++++++++++++++ 2 files changed, 109 insertions(+), 5 deletions(-) diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 1b8742ba4..a6e1dc3ae 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -7,7 +7,7 @@ import inspect 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 @@ -395,6 +395,13 @@ def _make_filesystem_configurator( _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"write_stdin failed: 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", @@ -515,21 +522,42 @@ def _wrap_exec_command(tool: FunctionTool) -> FunctionTool: 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 + 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 + # An empty poll sends no input: `chars` is "" or omitted entirely (None). + # Both must hit the dead-session guard, otherwise dropping `chars` reopens + # the exact repeat-poll storm this wrapper exists to stop. + is_empty_poll = chars is None or chars == "" + if isinstance(session_id, int) and is_empty_poll and session_id in dead_pty_session_ids: + return _WRITE_STDIN_DEAD_SESSION_RESULT_TEMPLATE.format(session_id=session_id) if isinstance(parsed, dict): - if isinstance(parsed.get("chars"), str): - parsed["chars"] = _decode_chars_escape(parsed["chars"]) + if isinstance(chars, str): + parsed["chars"] = _decode_chars_escape(chars) _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) + if isinstance(result, str): + # Only a result that is *exactly* the terminal failure marks the + # session dead. A live poll whose captured output merely contains + # this text must never poison the session (cf. the running-marker + # guard in _annotate_running_non_tty). + match = _PTY_SESSION_NOT_FOUND_RE.fullmatch(result.strip()) + 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 c924af191..804d80a15 100644 --- a/tests/test_agent_factory_shell.py +++ b/tests/test_agent_factory_shell.py @@ -13,6 +13,10 @@ from strix.agents import factory from strix.config import load_settings +_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 @@ -26,6 +30,19 @@ def _capturing_exec_tool(captured: dict[str, str]) -> FunctionTool: ) +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] = {} @@ -215,3 +232,62 @@ def test_view_image_is_left_out_for_text_only_models( view_image = next(tool for tool in filesystem.tools() if tool.name == "view_image") assert cast("FunctionTool", view_image).is_enabled is has_view_image + + +@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 + + +@pytest.mark.asyncio +async def test_wrap_write_stdin_does_not_mark_session_dead_from_echoed_output() -> None: + # A live poll whose captured output merely echoes the not-found text must + # not poison the session; later empty polls must still reach the tool. + live_output = ( + f"Process running with session ID {_MISSING_PTY_SESSION_ID}\nOutput:\n" + f"{_MISSING_PTY_RESULT}\n" + ) + captured: list[str] = [] + wrapped = factory._wrap_write_stdin( + _capturing_write_stdin_tool([live_output, live_output], 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 == live_output + assert second_result == live_output + assert len(captured) == 2 + + +@pytest.mark.asyncio +async def test_wrap_write_stdin_blocks_repeat_poll_when_chars_is_omitted() -> None: + # An empty poll may omit `chars` entirely; a known-dead session must still + # be short-circuited rather than polled again. + captured: list[str] = [] + wrapped = factory._wrap_write_stdin( + _capturing_write_stdin_tool([_MISSING_PTY_RESULT, _MISSING_PTY_RESULT], captured) + ) + + await wrapped.on_invoke_tool( + cast("Any", None), + json.dumps({"session_id": _MISSING_PTY_SESSION_ID, "chars": ""}), + ) + second_result = await wrapped.on_invoke_tool( + cast("Any", None), json.dumps({"session_id": _MISSING_PTY_SESSION_ID}) + ) + + assert "non-retryable" in second_result + assert len(captured) == 1 From 5f7cfe7dc3a7e2a964cf852115882180ed9f4a89 Mon Sep 17 00:00:00 2001 From: Ben Younes <2910651+ousamabenyounes@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:49:02 +0000 Subject: [PATCH 2/2] refactor(agents): isolate shell adapters --- strix/agents/factory.py | 175 +-------------------------------- strix/agents/shell_tools.py | 186 ++++++++++++++++++++++++++++++++++++ 2 files changed, 188 insertions(+), 173 deletions(-) create mode 100644 strix/agents/shell_tools.py diff --git a/strix/agents/factory.py b/strix/agents/factory.py index a6e1dc3ae..6cf39986c 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -7,17 +7,15 @@ import inspect import json import logging import re -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any 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 from strix.agents.prompt import render_system_prompt +from strix.agents.shell_tools import _wrap_exec_command, _wrap_write_stdin from strix.config import load_settings from strix.config.tool_call_arguments import describe_malformed_arguments from strix.tools.agents_graph.tools import ( @@ -394,175 +392,6 @@ def _make_filesystem_configurator( return configure -_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"write_stdin failed: 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", - "\\t": "\t", - "\\r": "\r", - "\\0": "\x00", - "\\a": "\x07", - "\\b": "\x08", - "\\v": "\x0b", - "\\f": "\x0c", -} - - -def _decode_chars_escape(s: str) -> str: - if "\\" not in s: - return s - - def sub(match: re.Match[str]) -> str: - token = match.group(0) - if token in _CHARS_ESCAPE_MAP: - return _CHARS_ESCAPE_MAP[token] - if token.startswith(("\\u", "\\x")): - return chr(int(token[2:], 16)) - return token - - return _CHARS_ESCAPE_RE.sub(sub, s) - - -def _format_validation_error(tool_name: str, exc: ValidationError) -> str: - parts: list[str] = [] - for err in exc.errors(): - loc = ".".join(str(x) for x in err.get("loc", ())) - msg = err.get("msg", "invalid") - parts.append(f"{loc}: {msg}" if loc else msg) - return f"{tool_name}: invalid arguments — " + "; ".join(parts) - - -def _apply_shell_output_cap(parsed: dict[str, Any]) -> None: - """Clamp the SDK shell tools' ``max_output_tokens`` to the configured - ceiling; a smaller explicit value is respected.""" - ceiling = load_settings().context.tool_output_max_tokens - requested = parsed.get("max_output_tokens") - parsed["max_output_tokens"] = ( - ceiling if not isinstance(requested, int) or requested > ceiling else requested - ) - - -_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) -_OUTPUT_HEADER_RE = re.compile(r"^Output:\n", 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 - header = _OUTPUT_HEADER_RE.search(result) - head_end = len(result) if header is None else header.start() - if _PROCESS_RUNNING_RE.search(result, 0, head_end) 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: ") - ) - if header is None: - return f"{result}\n{notice}" - return f"{result[: header.end()]}{notice}\n{result[header.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: - parsed = json.loads(raw_input) - except (json.JSONDecodeError, TypeError): - parsed = None - if isinstance(parsed, dict): - if "shell" not in parsed: - parsed["shell"] = "bash" - _apply_shell_output_cap(parsed) - raw_input = json.dumps(parsed) - try: - result = await invoke_tool(ctx, raw_input) - except ValidationError as exc: - return _format_validation_error(tool.name, exc) - except InvalidManifestPathError as exc: - rel = exc.context.get("rel", "?") - return ( - "exec_command: workdir must be a path inside /workspace " - "(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 - - -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: - raw_parsed: Any = json.loads(raw_input) - except json.JSONDecodeError: - parsed = None - 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 - # An empty poll sends no input: `chars` is "" or omitted entirely (None). - # Both must hit the dead-session guard, otherwise dropping `chars` reopens - # the exact repeat-poll storm this wrapper exists to stop. - is_empty_poll = chars is None or chars == "" - if isinstance(session_id, int) and is_empty_poll and session_id in dead_pty_session_ids: - return _WRITE_STDIN_DEAD_SESSION_RESULT_TEMPLATE.format(session_id=session_id) - if isinstance(parsed, dict): - if isinstance(chars, str): - parsed["chars"] = _decode_chars_escape(chars) - _apply_shell_output_cap(parsed) - raw_input = json.dumps(parsed) - try: - result = await invoke_tool(ctx, raw_input) - except ValidationError as exc: - return _format_validation_error(tool.name, exc) - if isinstance(result, str): - # Only a result that is *exactly* the terminal failure marks the - # session dead. A live poll whose captured output merely contains - # this text must never poison the session (cf. the running-marker - # guard in _annotate_running_non_tty). - match = _PTY_SESSION_NOT_FOUND_RE.fullmatch(result.strip()) - if match is not None: - dead_pty_session_ids.add(int(match.group("session_id"))) - return result - - tool.on_invoke_tool = invoke - return tool - - def _configure_shell_tools( toolset: Any, *, chat_completions: bool, strict_schemas: bool = True ) -> None: diff --git a/strix/agents/shell_tools.py b/strix/agents/shell_tools.py new file mode 100644 index 000000000..effebedb1 --- /dev/null +++ b/strix/agents/shell_tools.py @@ -0,0 +1,186 @@ +"""Adapt SDK shell tools: output limits, argument decoding and PTY session state.""" + +from __future__ import annotations + +import json +import re +from typing import TYPE_CHECKING, Any, cast + +from agents.sandbox.errors import InvalidManifestPathError +from agents.sandbox.session.pty_types import PTY_YIELD_TIME_MS_MAX +from pydantic import ValidationError + +from strix.config import load_settings + + +if TYPE_CHECKING: + from agents.tool import FunctionTool + + +_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"write_stdin failed: 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", + "\\t": "\t", + "\\r": "\r", + "\\0": "\x00", + "\\a": "\x07", + "\\b": "\x08", + "\\v": "\x0b", + "\\f": "\x0c", +} + + +def _decode_chars_escape(s: str) -> str: + if "\\" not in s: + return s + + def sub(match: re.Match[str]) -> str: + token = match.group(0) + if token in _CHARS_ESCAPE_MAP: + return _CHARS_ESCAPE_MAP[token] + if token.startswith(("\\u", "\\x")): + return chr(int(token[2:], 16)) + return token + + return _CHARS_ESCAPE_RE.sub(sub, s) + + +def _format_validation_error(tool_name: str, exc: ValidationError) -> str: + parts: list[str] = [] + for err in exc.errors(): + loc = ".".join(str(x) for x in err.get("loc", ())) + msg = err.get("msg", "invalid") + parts.append(f"{loc}: {msg}" if loc else msg) + return f"{tool_name}: invalid arguments — " + "; ".join(parts) + + +def _apply_shell_output_cap(parsed: dict[str, Any]) -> None: + """Clamp the SDK shell tools' ``max_output_tokens`` to the configured + ceiling; a smaller explicit value is respected.""" + ceiling = load_settings().context.tool_output_max_tokens + requested = parsed.get("max_output_tokens") + parsed["max_output_tokens"] = ( + ceiling if not isinstance(requested, int) or requested > ceiling else requested + ) + + +_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) +_OUTPUT_HEADER_RE = re.compile(r"^Output:\n", 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 + header = _OUTPUT_HEADER_RE.search(result) + head_end = len(result) if header is None else header.start() + if _PROCESS_RUNNING_RE.search(result, 0, head_end) 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: ") + ) + if header is None: + return f"{result}\n{notice}" + return f"{result[: header.end()]}{notice}\n{result[header.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: + parsed = json.loads(raw_input) + except (json.JSONDecodeError, TypeError): + parsed = None + if isinstance(parsed, dict): + if "shell" not in parsed: + parsed["shell"] = "bash" + _apply_shell_output_cap(parsed) + raw_input = json.dumps(parsed) + try: + result = await invoke_tool(ctx, raw_input) + except ValidationError as exc: + return _format_validation_error(tool.name, exc) + except InvalidManifestPathError as exc: + rel = exc.context.get("rel", "?") + return ( + "exec_command: workdir must be a path inside /workspace " + "(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 + + +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: + raw_parsed: Any = json.loads(raw_input) + except json.JSONDecodeError: + parsed = None + 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 + # An empty poll sends no input: `chars` is "" or omitted entirely (None). + # Both must hit the dead-session guard, otherwise dropping `chars` reopens + # the exact repeat-poll storm this wrapper exists to stop. + is_empty_poll = chars is None or chars == "" + if isinstance(session_id, int) and is_empty_poll and session_id in dead_pty_session_ids: + return _WRITE_STDIN_DEAD_SESSION_RESULT_TEMPLATE.format(session_id=session_id) + if isinstance(parsed, dict): + if isinstance(chars, str): + parsed["chars"] = _decode_chars_escape(chars) + _apply_shell_output_cap(parsed) + raw_input = json.dumps(parsed) + try: + result = await invoke_tool(ctx, raw_input) + except ValidationError as exc: + return _format_validation_error(tool.name, exc) + if isinstance(result, str): + # Only a result that is *exactly* the terminal failure marks the + # session dead. A live poll whose captured output merely contains + # this text must never poison the session (cf. the running-marker + # guard in _annotate_running_non_tty). + match = _PTY_SESSION_NOT_FOUND_RE.fullmatch(result.strip()) + if match is not None: + dead_pty_session_ids.add(int(match.group("session_id"))) + return result + + tool.on_invoke_tool = invoke + return tool