From 848129c84fd1499a62ff2c621baff18f5726c705 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Wed, 21 Oct 2026 20:43:39 +0300 Subject: [PATCH 01/35] fix(codex): let the Codex session exit on stdin EOF before killing its tree close() killed the child's session with SIGKILL before closing stdin, so Codex never saw the EOF it exits on and could do nothing at shutdown. Stdin now closes first and the session gets one second to end on its own; the tree kill remains the fallback. The session id stays reserved while any member lives, so a launcher shim's real CLI child that keeps the room directory is still reaped. --- src/band/integrations/codex/stdio_client.py | 32 ++++++++++++++++++- tests/integrations/codex/test_cleanup.py | 6 ++-- tests/integrations/codex/test_stdio_client.py | 14 ++++++++ 3 files changed, 48 insertions(+), 4 deletions(-) diff --git a/src/band/integrations/codex/stdio_client.py b/src/band/integrations/codex/stdio_client.py index c87961e5a..f5530b87a 100644 --- a/src/band/integrations/codex/stdio_client.py +++ b/src/band/integrations/codex/stdio_client.py @@ -30,6 +30,9 @@ # Long enough for the child watcher to record an exit, short enough that a # held stdio pipe cannot stall room cleanup. _EXIT_OBSERVE_S = 2.0 +# How long a closed stdin gets to end the whole Codex session before it is killed. +_SESSION_EXIT_S = 1.0 +_SESSION_POLL_S = 0.02 _STDOUT_LIMIT = 16 * 1024 * 1024 HELPER_PATH = Path(__file__).resolve().with_name("winmain.py") @@ -112,6 +115,30 @@ async def kill_process_tree(process: asyncio.subprocess.Process) -> None: return +async def _await_session_exit( + process: asyncio.subprocess.Process, timeout: float +) -> bool: + """True once the child and every other process in its session are gone. + + The session id stays reserved while any member lives, so probing the + group is exact; the launcher alone exiting is not enough, since its + real CLI child would keep the room's working directory. + """ + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while True: + try: + os.killpg(process.pid, 0) + except ProcessLookupError: + return True + except PermissionError: + pass + remaining = deadline - loop.time() + if remaining <= 0: + return False + await asyncio.sleep(min(_SESSION_POLL_S, remaining)) + + class CodexStdioClient(BaseJsonRpcClient): """Bidirectional JSON-RPC client with a single stdout read loop.""" @@ -258,8 +285,11 @@ async def _close_process(self) -> None: await self._windows_process.close(_EXIT_OBSERVE_S) self._windows_process = None elif (process := self._proc) is not None: - await kill_process_tree(process) + # Codex and its launcher exit on their own once stdin closes; the + # tree kill is the fallback for a session that does not. await close_stdin(process) + if not await _await_session_exit(process, _SESSION_EXIT_S): + await kill_process_tree(process) if not await await_returncode(process, _EXIT_OBSERVE_S): logger.warning( "Codex process %s still alive after tree kill", process.pid diff --git a/tests/integrations/codex/test_cleanup.py b/tests/integrations/codex/test_cleanup.py index b382d7644..70000a828 100644 --- a/tests/integrations/codex/test_cleanup.py +++ b/tests/integrations/codex/test_cleanup.py @@ -223,8 +223,8 @@ async def test_forced_close_stops_what_a_launcher_shim_spawned( async def test_close_reaps_a_worker_when_its_launcher_exits_on_stdin_eof( self, tmp_path: Path ) -> None: - """A launcher that exits on stdin EOF would otherwise leave its child - holding the workdir.""" + """A launcher gets to exit on stdin EOF, but its child still holding + the workdir is reaped by the tree kill that follows.""" workdir = tmp_path / "room" workdir.mkdir() client = CodexStdioClient( @@ -237,7 +237,7 @@ async def test_close_reaps_a_worker_when_its_launcher_exits_on_stdin_eof( await asyncio.wait_for(wait_for_heartbeat(heartbeat), timeout=CLOSE_BOUND_S) await asyncio.wait_for(client.close(), timeout=CLOSE_BOUND_S) await assert_heartbeat_stopped(heartbeat) - assert not (workdir / STDIN_EOF_MARKER).exists() + assert (workdir / STDIN_EOF_MARKER).exists() shutil.rmtree(workdir) finally: await asyncio.wait_for(client.close(), timeout=CLOSE_BOUND_S) diff --git a/tests/integrations/codex/test_stdio_client.py b/tests/integrations/codex/test_stdio_client.py index ff255cc58..332f425b1 100644 --- a/tests/integrations/codex/test_stdio_client.py +++ b/tests/integrations/codex/test_stdio_client.py @@ -408,3 +408,17 @@ async def test_stdio_client_eof_fails_pending_and_enqueues_disconnect( event = await client.recv_event(timeout_s=1.0) assert event.method == "transport/closed" await client.close() + + +@pytest.mark.asyncio +@pytest.mark.skipif(os.name == "nt", reason="Windows closes through the job object") +async def test_close_lets_codex_exit_on_stdin_eof_before_killing_its_tree( + fake_codex_server_script: Path, +) -> None: + """A harness that exits once its stdin closes ends on its own terms; the + tree kill is the fallback for one that does not.""" + client = await _build_client(fake_codex_server_script, scenario="basic") + process = client._proc + assert process is not None + await client.close() + assert process.returncode == 0 From 0b6c1e48ca06466e269b987d2be78d189056cdba Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Wed, 21 Oct 2026 15:49:12 +0300 Subject: [PATCH 02/35] feat(omp): add native RPC client and wire types The ACP bridge hides OMP's native lifecycle: sessions, model and thinking selection, host tool calls, permission dialogs and ask questions. This adds the transport a native adapter needs: OmpRpcClient (JSONL RPC over the omp process stdio, with process-tree reaping), typed wire frames, the native launch settings and SDK-owned MCP plugin projection in band.integrations.ompprocess, and the per-room session converter. The shared normalizer gains a namespace keyword so plugin-projected tool names map back to Band tools; omp_acp's spellings are unchanged. --- pyproject.toml | 4 + src/band/converters/__init__.py | 7 + src/band/converters/omp.py | 39 + src/band/integrations/omp.py | 21 +- src/band/integrations/ompprocess.py | 301 ++++++ src/band/integrations/omprpc.py | 958 ++++++++++++++++++ src/band/integrations/ompwire.py | 301 ++++++ tests/converters/test_omp.py | 86 ++ tests/framework_configs/converters.py | 7 +- tests/integrations/omppeer.py | 593 +++++++++++ tests/integrations/omppermissionpeer.py | 109 ++ tests/integrations/ompprocess/__init__.py | 0 tests/integrations/ompprocess/test_launch.py | 47 + .../ompprocess/test_mcp_plugin.py | 79 ++ .../ompprocess/test_mcp_startup.py | 107 ++ .../ompprocess/test_tool_names.py | 52 + tests/integrations/omprpc/__init__.py | 0 tests/integrations/omprpc/conftest.py | 65 ++ tests/integrations/omprpc/kit.py | 133 +++ tests/integrations/omprpc/test_callbacks.py | 179 ++++ tests/integrations/omprpc/test_commands.py | 40 + .../integrations/omprpc/test_interruption.py | 110 ++ .../integrations/omprpc/test_process_tree.py | 127 +++ .../omprpc/test_prompt_results.py | 142 +++ tests/integrations/omprpc/test_startup.py | 62 ++ .../omprpc/test_transport_failures.py | 82 ++ .../omprpc/test_turn_settlement.py | 95 ++ uv.lock | 36 +- 28 files changed, 3770 insertions(+), 12 deletions(-) create mode 100644 src/band/converters/omp.py create mode 100644 src/band/integrations/ompprocess.py create mode 100644 src/band/integrations/omprpc.py create mode 100644 src/band/integrations/ompwire.py create mode 100644 tests/converters/test_omp.py create mode 100644 tests/integrations/omppeer.py create mode 100644 tests/integrations/omppermissionpeer.py create mode 100644 tests/integrations/ompprocess/__init__.py create mode 100644 tests/integrations/ompprocess/test_launch.py create mode 100644 tests/integrations/ompprocess/test_mcp_plugin.py create mode 100644 tests/integrations/ompprocess/test_mcp_startup.py create mode 100644 tests/integrations/ompprocess/test_tool_names.py create mode 100644 tests/integrations/omprpc/__init__.py create mode 100644 tests/integrations/omprpc/conftest.py create mode 100644 tests/integrations/omprpc/kit.py create mode 100644 tests/integrations/omprpc/test_callbacks.py create mode 100644 tests/integrations/omprpc/test_commands.py create mode 100644 tests/integrations/omprpc/test_interruption.py create mode 100644 tests/integrations/omprpc/test_process_tree.py create mode 100644 tests/integrations/omprpc/test_prompt_results.py create mode 100644 tests/integrations/omprpc/test_startup.py create mode 100644 tests/integrations/omprpc/test_transport_failures.py create mode 100644 tests/integrations/omprpc/test_turn_settlement.py diff --git a/pyproject.toml b/pyproject.toml index 1272669c8..ff58e1638 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -139,6 +139,9 @@ acp = [ "starlette>=0.40.0", "uvicorn>=0.29", ] +omp = [ + "psutil>=6", +] slack = [ "slack-sdk>=3.27.0", # AsyncWebClient, signature verifier, Socket Mode "starlette>=0.40.0", # HTTP server framework @@ -206,6 +209,7 @@ dev = [ "pytest-rerunfailures>=14.0.0", "tenacity>=8.0.0", "pytest-timeout>=2.4.0", + "psutil>=6", # Native OMP process-tree cleanup tests "packaging>=23.0", # PEP 440 pre-release detection for the CI coverage-gate version guard "band-testing-python==0.1.4", "httpx>=0.24.0", # Already in main deps, for mocking diff --git a/src/band/converters/__init__.py b/src/band/converters/__init__.py index 431fb6876..bde4b094d 100644 --- a/src/band/converters/__init__.py +++ b/src/band/converters/__init__.py @@ -98,6 +98,12 @@ from band.converters.nooa import ( NooaTranscript as NooaTranscript, ) + from band.converters.omp import ( + OmpHistoryConverter as OmpHistoryConverter, + ) + from band.converters.omp import ( + OmpSessionState as OmpSessionState, + ) from band.converters.opencode import ( OpencodeHistoryConverter as OpencodeHistoryConverter, ) @@ -139,6 +145,7 @@ gemini=["GeminiHistoryConverter", "GeminiMessages"], google_adk=["GoogleADKHistoryConverter", "GoogleADKMessages"], opencode=["OpencodeHistoryConverter"], + omp=["OmpHistoryConverter", "OmpSessionState"], nooa=["NooaHistoryConverter", "NooaTranscript"], strands=["StrandsHistoryConverter", "StrandsMessages"], ) diff --git a/src/band/converters/omp.py b/src/band/converters/omp.py new file mode 100644 index 000000000..753a521b1 --- /dev/null +++ b/src/band/converters/omp.py @@ -0,0 +1,39 @@ +"""Native OMP resume candidates and read-only room transcript replay.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field + +from band.converters.helpers import build_replay_messages +from band.core.protocols import HistoryConverter + +# Task-event metadata keys the native adapter writes and this converter reads. +OMP_ROOM_METADATA_KEY = "omp_room_id" +OMP_SESSION_METADATA_KEY = "omp_session_file" + + +class OmpSessionState(BaseModel): + """The native session file each room may resume, plus the room's transcript.""" + + room_to_session: dict[str, str] = Field(default_factory=dict) + replay_messages: list[str] = Field(default_factory=list) + + +class OmpHistoryConverter(HistoryConverter[OmpSessionState]): + """Extract per-room native session files from task-event metadata.""" + + def convert(self, raw: list[dict[str, Any]]) -> OmpSessionState: + sessions: dict[str, str] = {} + for message in raw: + metadata = message.get("metadata") + if not isinstance(metadata, dict): + continue + room = metadata.get(OMP_ROOM_METADATA_KEY) + session = metadata.get(OMP_SESSION_METADATA_KEY) + if isinstance(room, str) and isinstance(session, str) and session: + sessions[room] = session + return OmpSessionState( + room_to_session=sessions, replay_messages=build_replay_messages(raw) + ) diff --git a/src/band/integrations/omp.py b/src/band/integrations/omp.py index a4370d697..69782d9c1 100644 --- a/src/band/integrations/omp.py +++ b/src/band/integrations/omp.py @@ -210,12 +210,12 @@ def is_omp_approve_deny_form(requested_schema: object) -> bool: return approve_deny_form_field(requested_schema) is not None -def _omp_mcp_remainder(name_or_path: str) -> str | None: - """Strip ``xd://`` and/or ``mcp__``; ``None`` when not an OMP MCP spelling.""" +def _omp_mcp_remainder(name_or_path: str, namespace: str) -> str | None: + """Strip ``xd://``, ``mcp__`` and ``namespace``; ``None`` when not an OMP MCP spelling.""" text = name_or_path.removeprefix(XD_URL_PREFIX) if not text.startswith(XD_MCP_PREFIX): return None - remainder = text.removeprefix(XD_MCP_PREFIX) + remainder = text.removeprefix(XD_MCP_PREFIX).removeprefix(namespace) return remainder or None @@ -243,10 +243,14 @@ def _omp_mcp_wire_candidates(remainder: str) -> list[str]: def normalize_omp_mcp_tool_name( - name_or_path: str, own_names: Collection[str] + name_or_path: str, own_names: Collection[str], *, namespace: str = "" ) -> str | None: - """Map an OMP ``mcp__…`` title or ``xd://mcp__…`` path to a Band tool name.""" - remainder = _omp_mcp_remainder(name_or_path) + """Map an OMP ``mcp__…`` title or ``xd://mcp__…`` path to a Band tool name. + + ``namespace`` is a server-name prefix OMP mints into every tool of a + plugin-projected server; it is removed before the remainder is matched. + """ + remainder = _omp_mcp_remainder(name_or_path, namespace) if remainder is None: return None for candidate in _omp_mcp_wire_candidates(remainder): @@ -262,6 +266,7 @@ def normalize_omp_mcp_device_call( own_names: Collection[str], *, kind: str | None = None, + namespace: str = "", ) -> tuple[str, dict[str, object]]: """Map OMP MCP device writes / ``mcp__`` titles to canonical Band tool names. @@ -281,7 +286,7 @@ def normalize_omp_mcp_device_call( break if path is not None: - canonical = normalize_omp_mcp_tool_name(path, own_names) + canonical = normalize_omp_mcp_tool_name(path, own_names, namespace=namespace) content = args.get("content") if canonical is None or kind == _ACP_KIND_READ: return name, args @@ -297,7 +302,7 @@ def normalize_omp_mcp_device_call( return canonical, args return canonical, parsed - canonical = normalize_omp_mcp_tool_name(name, own_names) + canonical = normalize_omp_mcp_tool_name(name, own_names, namespace=namespace) if canonical is None: return name, args return canonical, args diff --git a/src/band/integrations/ompprocess.py b/src/band/integrations/ompprocess.py new file mode 100644 index 000000000..655dc11a6 --- /dev/null +++ b/src/band/integrations/ompprocess.py @@ -0,0 +1,301 @@ +"""Native OMP (``omp --mode rpc``) launch settings and the SDK-owned MCP plugin.""" + +from __future__ import annotations + +import json +import os +import tempfile +from collections.abc import Collection, Iterator, Mapping, Sequence +from contextlib import contextmanager +from pathlib import Path +from typing import Literal, Self +from uuid import uuid4 + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from band.integrations.omp import ( + OMP_CWD_FLAG, + normalize_omp_mcp_device_call, + normalize_omp_mcp_tool_name, + omp_provider_env, +) +from band.integrations.omprpc import OmpRpcClient +from band.integrations.ompwire import ( + OmpEvent, + OmpHostToolCall, + OmpToolResult, + OmpUIRequest, + OmpUIResponse, +) + +DEFAULT_OMP_COMMAND: tuple[str, ...] = ("omp", "--mode", "rpc") + +# The native RPC ask dialog and MCP status events need a newer OMP than the +# ACP bridge's pin. +OMP_NATIVE_PINNED_PACKAGE = "@oh-my-pi/pi-coding-agent@18.8.9" + +# Plugin names prefix the MCP server names OMP mints, so the plugin root's +# basename must stay stable across starts; its sanitized form is the +# namespace every projected tool name carries. +OMP_PLUGIN_ROOT = "band-sdk" +OMP_PLUGIN_TOOL_NAMESPACE = f"{OMP_PLUGIN_ROOT.replace('-', '_')}_" + + +def omp_runtime_env( + *, + env: Mapping[str, str] | None, + model: str | None, + api_key: str | None, +) -> dict[str, str] | None: + """Build the provider environment, preserving explicitly supplied overrides.""" + credentials = ( + omp_provider_env(model=model, api_key=api_key) + if model and api_key is not None + else {} + ) + return {**credentials, **(env or {})} or None + + +def omp_command_in_workspace(command: Sequence[str], workspace: str) -> list[str]: + """Select OMP's room workspace without changing the subprocess launch cwd.""" + return [*command, f"{OMP_CWD_FLAG}={workspace}"] + + +class OmpMcpServer(BaseModel): + """An external MCP server projected into the SDK-owned OMP plugin.""" + + model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True) + + type: Literal["stdio", "http", "sse"] | None = None + command: str | None = None + args: tuple[str, ...] = () + env: dict[str, str] = Field(default_factory=dict, repr=False) + cwd: str | None = None + url: str | None = None + headers: dict[str, str] = Field(default_factory=dict, repr=False) + + @model_validator(mode="after") + def _require_one_transport(self) -> Self: + if bool(self.command) == bool(self.url): + raise ValueError("MCP server needs exactly one of command or url") + if self.command and self.type not in (None, "stdio"): + raise ValueError("a command requires stdio transport") + if self.url and self.type == "stdio": + raise ValueError("a URL requires http or sse transport") + return self + + +_MCP_STARTUP_EXTENSION = """ +import { writeFileSync } from "node:fs"; + +export default function (pi) { + const expected = new Set(config.servers); + const statuses = new Map(); + let finish; + const ready = new Promise(resolve => { finish = resolve; }); + const unsubscribe = pi.events.on("mcp:connection-status", event => { + if (!expected.has(event.serverName) || statuses.has(event.serverName)) return; + if (event.type !== "connected" && event.type !== "failed") return; + statuses.set(event.serverName, event); + if (event.type === "failed" || statuses.size === expected.size) { + unsubscribe(); + finish(); + } + }); + if (expected.size === 0) { + unsubscribe(); + finish(); + } + pi.registerCommand(config.command, { + description: "Verify the host-configured MCP connections", + handler: async (_args, ctx) => { + await ready; + const systemPromptReady = !config.systemPrompt + || ctx.getSystemPrompt().some(block => block.includes(config.systemPrompt.trim())); + writeFileSync(config.receipt, JSON.stringify({ systemPromptReady, servers: [...statuses.values()] }), { mode: 0o600 }); + }, + }); +} +""" + + +class OmpStartupError(RuntimeError): + """The required native connection context or MCP tools are not ready.""" + + +class _McpCommand(BaseModel): + name: str + + +class _McpCommands(BaseModel): + commands: list[_McpCommand] + + +class _ConnectedMcpServer(BaseModel): + type: Literal["connected"] + serverName: str + + +class _FailedMcpServer(BaseModel): + type: Literal["failed"] + serverName: str + error: str + + +class _McpStartupReceipt(BaseModel): + systemPromptReady: bool = Field(strict=True) + servers: list[_ConnectedMcpServer | _FailedMcpServer] + + +def _plugin_server_name(name: str) -> str: + return f"{OMP_PLUGIN_ROOT}:{name}" + + +def omp_mcp_ready_command(plugin: str) -> str: + """The slash command the plugin's startup extension registers for its receipt.""" + return f"band-sdk-mcp-{Path(plugin).parent.name}" + + +def omp_mcp_arguments(plugin: str) -> tuple[str, ...]: + """Enable native MCP status events without changing the user's persisted settings.""" + return ( + "--plugin-dir", + plugin, + "--mode=rpc-ui", + "--extension", + str(Path(plugin) / "startup.js"), + "--config", + str(Path(plugin) / "startup.json"), + ) + + +async def omp_mcp_ready( + runtime: OmpRpcClient, + plugin: str, + servers: Mapping[str, OmpMcpServer], +) -> None: + """Require the connection extension and configured MCP servers before inference.""" + command = omp_mcp_ready_command(plugin) + commands = _McpCommands.model_validate( + await runtime.request("get_available_commands") + ) + if not any(item.name == command for item in commands.commands): + raise OmpStartupError("OMP did not load the required SDK startup extension") + + async def on_event(event: OmpEvent) -> None: + if event.type == "agent_start": + raise OmpStartupError( + "OMP unexpectedly invoked an agent during SDK startup" + ) + + async def on_tool(_call: OmpHostToolCall) -> OmpToolResult: + raise OmpStartupError("OMP unexpectedly invoked a host tool during SDK startup") + + async def on_ui(_request: OmpUIRequest) -> OmpUIResponse: + return OmpUIResponse(cancelled=True) + + await runtime.prompt( + f"/{command}", + on_event=on_event, + on_tool=on_tool, + on_ui=on_ui, + ) + try: + receipt = _McpStartupReceipt.model_validate_json( + (Path(plugin) / "receipt.json").read_text(encoding="utf-8") + ) + except (OSError, ValueError) as error: + raise OmpStartupError("OMP returned no valid SDK startup receipt") from error + if not receipt.systemPromptReady: + raise OmpStartupError( + "OMP did not load the required Band connection instructions; " + "rules may be disabled or omitted by the system prompt template" + ) + expected = {_plugin_server_name(name) for name in servers} + for status in receipt.servers: + if status.serverName not in expected: + raise OmpStartupError("OMP returned an unexpected MCP startup receipt") + if isinstance(status, _FailedMcpServer): + raise OmpStartupError( + f"OMP MCP server {status.serverName} failed: {status.error}" + ) + if {status.serverName for status in receipt.servers} != expected or len( + receipt.servers + ) != len(expected): + raise OmpStartupError("OMP did not acknowledge every configured MCP server") + + +def _projected_server(server: OmpMcpServer, workspace: str) -> dict[str, object]: + payload = server.model_dump(mode="json", exclude_none=True, exclude_defaults=True) + if server.command: + # Plugin discovery otherwise rebases launch paths into the temporary overlay. + cwd = Path(os.path.abspath(Path(workspace) / (server.cwd or "."))) + payload["cwd"] = str(cwd) + if server.command.startswith(("./", "../", ".\\", "..\\")): + payload["command"] = os.path.abspath(cwd / server.command) + return payload + + +def _write_private(path: Path, content: str) -> None: + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as handle: + handle.write(content) + + +@contextmanager +def omp_mcp_plugin( + servers: Mapping[str, OmpMcpServer], *, workspace: str, system_prompt: str = "" +) -> Iterator[str | None]: + """Project connection instructions and servers without changing native user configs.""" + if not servers and not system_prompt: + yield None + return + projected = { + name: _projected_server(server, workspace) for name, server in servers.items() + } + with tempfile.TemporaryDirectory(prefix="band-omp-") as temporary: + root = Path(temporary) / OMP_PLUGIN_ROOT + root.mkdir(mode=0o700) + if system_prompt: + rules = root / "rules" + rules.mkdir(mode=0o700) + _write_private( + rules / f"band-connection-{uuid4().hex}.md", + f"---\nalwaysApply: true\n---\n{system_prompt}", + ) + startup = { + "servers": [_plugin_server_name(name) for name in servers], + "command": omp_mcp_ready_command(str(root)), + "receipt": str(root / "receipt.json"), + "systemPrompt": system_prompt, + } + _write_private(root / ".mcp.json", json.dumps({"mcpServers": projected})) + # Quiet startup suppresses the native status bus, not just UI rendering. + _write_private(root / "startup.json", json.dumps({"startup": {"quiet": False}})) + _write_private( + root / "startup.js", + f"const config = {json.dumps(startup)};\n{_MCP_STARTUP_EXTENSION}", + ) + yield str(root) + + +def normalize_native_mcp_tool_name( + name_or_path: str, own_names: Collection[str] +) -> str | None: + """Map a tool name OMP minted under the SDK plugin namespace to a Band tool.""" + return normalize_omp_mcp_tool_name( + name_or_path, own_names, namespace=OMP_PLUGIN_TOOL_NAMESPACE + ) + + +def normalize_native_mcp_device_call( + name: str, + arguments: Mapping[str, object], + own_names: Collection[str], + *, + kind: str | None = None, +) -> tuple[str, dict[str, object]]: + """Map a device call under the SDK plugin namespace to a canonical Band tool.""" + return normalize_omp_mcp_device_call( + name, arguments, own_names, kind=kind, namespace=OMP_PLUGIN_TOOL_NAMESPACE + ) diff --git a/src/band/integrations/omprpc.py b/src/band/integrations/omprpc.py new file mode 100644 index 000000000..16d04ce69 --- /dev/null +++ b/src/band/integrations/omprpc.py @@ -0,0 +1,958 @@ +"""Persistent subprocess transport for OMP's native JSONL RPC protocol.""" + +from __future__ import annotations + +import asyncio +import base64 +import logging +import os +from binascii import Error as Base64Error +from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence +from dataclasses import dataclass +from typing import cast + +import psutil +from pydantic import JsonValue, TypeAdapter, ValidationError + +from band.integrations.ompwire import ( + OMP_CHUNK_BYTES, + OMP_FRAME_BYTES, + OMP_REASSEMBLED_BYTES, + OmpChunk, + OmpCommand, + OmpEvent, + OmpEventType, + OmpFrame, + OmpFrameError, + OmpFrameType, + OmpHostToolCall, + OmpHostToolCancel, + OmpMessage, + OmpNegotiation, + OmpPromptAcceptance, + OmpPromptResult, + OmpReady, + OmpResponse, + OmpState, + OmpToolResult, + OmpUIRequest, + OmpUIResponse, +) + +logger = logging.getLogger(__name__) + +_DRAIN_BYTES = 64 * 1024 +_ABORT_TIMEOUT_S = 1.0 +_PROCESS_EXIT_TIMEOUT_S = 1.0 +_CALLBACK_EXIT_TIMEOUT_S = 1.0 +_PAYLOAD_ADAPTER: TypeAdapter[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]) +_EVENT_TYPES = frozenset(OmpEventType) +_UI_NOTIFICATIONS = frozenset( + {"notify", "setStatus", "setWidget", "setTitle", "set_editor_text", "open_url"} +) + + +class OmpRpcError(RuntimeError): + """An unsuccessful command response; the transport remains usable.""" + + def __init__(self, command: str, error: str, *, code: str | None = None) -> None: + self.command = command + self.error = error + self.code = code + super().__init__(f"OMP {command}: {error}") + + +class OmpProcessError(RuntimeError): + """The OMP process or its stdio transport is no longer usable.""" + + +class OmpProtocolError(OmpProcessError): + """An invalid or incorrectly correlated native RPC frame.""" + + +class OmpPromptError(RuntimeError): + """The prompt or its final background continuation failed or was aborted.""" + + def __init__(self, status: str, message: str) -> None: + self.status = status + super().__init__(message) + + +@dataclass +class PendingChunks: + chunk_id: str + count: int + byte_length: int + next_index: int + received_bytes: int + parts: list[bytes] + + +@dataclass(frozen=True) +class ReceivedResponse: + frame: OmpResponse + sequence: int + + +@dataclass(frozen=True) +class PendingRequest: + command: str + future: asyncio.Future[ReceivedResponse | OmpRpcError | OmpProcessError] + + +class PromptRun: + def __init__( + self, + request_id: str, + on_event: Callable[[OmpEvent], Awaitable[None]], + on_tool: Callable[[OmpHostToolCall], Awaitable[OmpToolResult]], + on_ui: Callable[[OmpUIRequest], Awaitable[OmpUIResponse]], + ) -> None: + loop = asyncio.get_running_loop() + self.id = request_id + self.on_event = on_event + self.on_tool = on_tool + self.on_ui = on_ui + self.events: asyncio.Queue[OmpEvent] = asyncio.Queue() + self.failure: asyncio.Future[BaseException] = loop.create_future() + self.result: asyncio.Future[OmpPromptResult] = loop.create_future() + self.changed = asyncio.Event() + self.last_activity_sequence = 0 + self.settled_sequence = 0 + self.last_assistant: OmpMessage | None = None + self.rejected = False + self.callbacks: dict[tuple[str, str], asyncio.Task[None]] = {} + self.callback_tasks: set[asyncio.Task[None]] = set() + self.event_worker: asyncio.Task[None] | None = None + + def fail(self, error: BaseException) -> None: + if not self.failure.done(): + self.failure.set_result(error) + + def observe(self, event: OmpEvent, sequence: int) -> None: + self.last_activity_sequence = sequence + if ( + event.type == OmpEventType.MESSAGE_END + and event.message is not None + and event.message.role == "assistant" + ): + self.last_assistant = event.message + if event.type == OmpEventType.AGENT_END: + for message in reversed(event.messages): + if message.role == "assistant": + self.last_assistant = message + break + self.events.put_nowait(event) + self.changed.set() + + +class OmpRpcClient: + """One OMP process, one persistent reader, and at most one active prompt.""" + + def __init__( + self, command: Sequence[str], *, env: Mapping[str, str] | None = None + ) -> None: + if not command: + raise ValueError("OMP command must not be empty") + self.command = tuple(command) + self.env = None if env is None else {**os.environ, **env} + self._proc: asyncio.subprocess.Process | None = None + self._process_identity: psutil.Process | None = None + self._descendants: set[psutil.Process] = set() + self._spawn_task: asyncio.Task[asyncio.subprocess.Process] | None = None + self._reader_task: asyncio.Task[None] | None = None + self._stderr_task: asyncio.Task[None] | None = None + self._close_task: asyncio.Task[None] | None = None + self._ready: asyncio.Future[OmpReady | OmpProcessError] | None = None + self._failure: OmpProcessError | None = None + self._closing = False + self._request_id = 0 + self._sequence = 0 + self._chunks: PendingChunks | None = None + self._max_frame_bytes = OMP_FRAME_BYTES + self._max_reassembled_bytes = OMP_REASSEMBLED_BYTES + self._pending: dict[str, PendingRequest] = {} + self._turn: PromptRun | None = None + self._start_lock = asyncio.Lock() + self._write_lock = asyncio.Lock() + self._prompt_lock = asyncio.Lock() + + async def start(self) -> None: + """Start the supplied command and require native RPC protocol v2.""" + async with self._start_lock: + if self._closing: + await self.close() + self._proc = None + self._spawn_task = None + self._process_identity = None + self._descendants.clear() + self._close_task = None + if self._proc is not None: + self._ensure_running() + return + self._failure = None + self._chunks = None + self._max_frame_bytes = OMP_FRAME_BYTES + self._max_reassembled_bytes = OMP_REASSEMBLED_BYTES + self._closing = False + self._ready = asyncio.get_running_loop().create_future() + self._spawn_task = asyncio.create_task(self._spawn_process()) + try: + self._proc = await asyncio.shield(self._spawn_task) + self._reader_task = asyncio.create_task(self._read_stdout(self._proc)) + self._stderr_task = asyncio.create_task(self._read_stderr(self._proc)) + ready = await asyncio.shield(self._ready) + if isinstance(ready, OmpProcessError): + raise ready + if 2 not in ready.supportedProtocolVersions: + raise OmpProtocolError("OMP RPC protocol v2 is required") + self._max_frame_bytes = ready.maxFrameBytes + self._max_reassembled_bytes = ready.maxReassembledFrameBytes + negotiated = await self._request_frame( + "negotiate_protocol", {"protocolVersion": 2} + ) + try: + OmpNegotiation.model_validate(self._response_data(negotiated.frame)) + except ValidationError as cause: + raise OmpProtocolError( + "OMP RPC protocol v2 negotiation failed" + ) from cause + except OSError as cause: + await self.close() + raise OmpProcessError("Unable to start OMP RPC process") from cause + except (Exception, asyncio.CancelledError): + await self.close() + raise + + async def _spawn_process(self) -> asyncio.subprocess.Process: + process = await asyncio.create_subprocess_exec( + *self.command, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=None, + env=self.env, + limit=OMP_FRAME_BYTES, + ) + self._proc = process + self._process_identity = await asyncio.to_thread( + self._capture_identity, process.pid + ) + return process + + @staticmethod + def _capture_identity(pid: int) -> psutil.Process | None: + try: + process = psutil.Process(pid) + # Process() can defer AccessDenied when reading the creation time. + process.create_time() + return process + except psutil.NoSuchProcess: + return None + + async def request( + self, command: str, **parameters: JsonValue + ) -> dict[str, JsonValue]: + """Correlate a native command response without consuming prompt events.""" + self._ensure_running() + response = await self._request_frame(command, parameters) + return self._response_data(response.frame) + + async def prompt( + self, + message: str, + *, + on_event: Callable[[OmpEvent], Awaitable[None]], + on_tool: Callable[[OmpHostToolCall], Awaitable[OmpToolResult]], + on_ui: Callable[[OmpUIRequest], Awaitable[OmpUIResponse]], + ) -> None: + """Deliver a turn through its correlated result and actual native settlement.""" + try: + async with self._prompt_lock: + self._ensure_running() + await self._execute_prompt(message, on_event, on_tool, on_ui) + except asyncio.CancelledError: + await self.close() + raise + + async def close(self) -> None: + """Finish shared bounded cleanup before propagating caller cancellation.""" + if self._proc is None and self._spawn_task is None: + return + self._begin_close() + closing = self._close_task + assert closing is not None + interrupted = False + while True: + try: + await asyncio.shield(closing) + break + except asyncio.CancelledError: + if closing.cancelled(): + raise + interrupted = True + if interrupted: + raise asyncio.CancelledError + + def _ensure_running(self) -> None: + if self._failure is not None: + raise self._failure + if self._closing or self._proc is None: + raise OmpProcessError("OMP RPC process is not running") + if self._proc.returncode is not None: + raise OmpProcessError("OMP RPC process exited") + + def _next_id(self) -> str: + self._request_id += 1 + return f"band-{self._request_id}" + + async def _request_frame( + self, + command: str, + parameters: Mapping[str, JsonValue], + *, + request_id: str | None = None, + allow_closing: bool = False, + ) -> ReceivedResponse: + if not allow_closing: + self._ensure_running() + if "id" in parameters or "type" in parameters: + raise ValueError("OMP command parameters must not override id or type") + request_id = request_id or self._next_id() + frame = OmpCommand.model_validate( + {"type": command, "id": request_id, **parameters} + ) + future: asyncio.Future[ReceivedResponse | OmpRpcError | OmpProcessError] = ( + asyncio.get_running_loop().create_future() + ) + self._pending[request_id] = PendingRequest(command, future) + try: + await self._write_line(frame.model_dump_json()) + response = await future + if isinstance(response, (OmpRpcError, OmpProcessError)): + raise response + return response + finally: + self._pending.pop(request_id, None) + + def _response_data(self, response: OmpResponse) -> dict[str, JsonValue]: + if response.data is None: + return {} + if not isinstance(response.data, dict): + error = OmpProtocolError("OMP command response data must be an object") + self._record_failure(error) + raise error + return response.data + + async def _execute_prompt( + self, + message: str, + on_event: Callable[[OmpEvent], Awaitable[None]], + on_tool: Callable[[OmpHostToolCall], Awaitable[OmpToolResult]], + on_ui: Callable[[OmpUIRequest], Awaitable[OmpUIResponse]], + ) -> None: + turn = PromptRun(self._next_id(), on_event, on_tool, on_ui) + self._turn = turn + event_worker = asyncio.create_task(self._deliver_events(turn)) + turn.event_worker = event_worker + work = asyncio.create_task(self._complete_prompt(turn, message)) + try: + done, _ = await asyncio.wait( + (work, turn.failure), return_when=asyncio.FIRST_COMPLETED + ) + if turn.failure in done: + raise turn.failure.result() + await work + except OmpRpcError: + if not turn.rejected: + await self.close() + raise + except BaseException: + await self.close() + raise + finally: + self._turn = None + await self._cancel_tasks([work, event_worker, *turn.callback_tasks]) + + async def _complete_prompt(self, turn: PromptRun, message: str) -> None: + try: + accepted = await self._request_frame( + "prompt", {"message": message}, request_id=turn.id + ) + except OmpRpcError: + turn.rejected = True + raise + try: + acceptance = OmpPromptAcceptance.model_validate( + self._response_data(accepted.frame) + ) + except ValidationError as cause: + raise OmpProtocolError( + "Invalid OMP prompt acceptance protocol data" + ) from cause + outcome: OmpPromptResult | None = None + # A synchronous local command explicitly owes no prompt_result. + if acceptance.agentInvoked is not False: + outcome = await turn.result + await self._wait_for_settlement(turn) + await turn.events.join() + if turn.callback_tasks: + _, pending = await asyncio.wait( + turn.callback_tasks, timeout=_CALLBACK_EXIT_TIMEOUT_S + ) + if pending: + raise OmpProcessError("OMP callbacks did not finish after settlement") + if outcome is not None and outcome.status != "completed": + message = ( + outcome.error.message + if outcome.error is not None + else f"OMP prompt {outcome.status}" + ) + raise OmpPromptError(outcome.status, message) + assistant = turn.last_assistant + if assistant is not None and assistant.stopReason in {"error", "aborted"}: + message = ( + assistant.errorMessage or f"OMP continuation {assistant.stopReason}" + ) + raise OmpPromptError(assistant.stopReason, message) + + async def _wait_for_settlement(self, turn: PromptRun) -> None: + snapshot = await self._request_frame("get_state", {}) + try: + state = OmpState.model_validate(self._response_data(snapshot.frame)) + except ValidationError as cause: + raise OmpProtocolError("Invalid OMP settlement protocol data") from cause + if state.isSettled and snapshot.sequence >= turn.last_activity_sequence: + return + # An old, uncorrelated session_settled cannot overrule a newer busy snapshot. + while ( + turn.settled_sequence <= snapshot.sequence + or turn.settled_sequence <= turn.last_activity_sequence + ): + turn.changed.clear() + await turn.changed.wait() + + async def _deliver_events(self, turn: PromptRun) -> None: + try: + while True: + event = await turn.events.get() + try: + await turn.on_event(event) + finally: + turn.events.task_done() + except asyncio.CancelledError: + raise + except Exception as error: # noqa: BLE001 -- arbitrary callback failures belong to the awaiting prompt + turn.fail(error) + + async def _write_line(self, line: str) -> None: + async with self._write_lock: + process = self._proc + if process is None or process.stdin is None: + raise OmpProcessError("OMP RPC stdin is unavailable") + try: + process.stdin.write((line + "\n").encode("utf-8")) + await process.stdin.drain() + except (BrokenPipeError, ConnectionError, OSError) as cause: + error = OmpProcessError("OMP RPC stdin closed") + error.__cause__ = cause + self._record_failure(error) + raise error from cause + + async def _write_frame(self, frame: dict[str, JsonValue]) -> None: + await self._write_line(_PAYLOAD_ADAPTER.dump_json(frame).decode("utf-8")) + + def _decode_payload( + self, payload: dict[str, JsonValue] + ) -> dict[str, JsonValue] | None: + frame = OmpFrame.model_validate(payload) + if frame.type != OmpFrameType.CHUNK: + if self._chunks is not None: + raise OmpProtocolError("OMP RPC protocol chunk sequence interrupted") + return payload + chunk = OmpChunk.model_validate(payload) + if chunk.byteLength > self._max_reassembled_bytes: + raise OmpProtocolError( + "OMP RPC protocol frame exceeds the negotiated limit" + ) + data = base64.b64decode(chunk.data, validate=True) + if ( + not data + or len(data) > OMP_CHUNK_BYTES + or base64.b64encode(data).decode("ascii") != chunk.data + ): + raise OmpProtocolError("Invalid OMP RPC protocol chunk data") + if self._chunks is None: + if chunk.index != 0: + raise OmpProtocolError( + "OMP RPC protocol chunks must start at index zero" + ) + self._chunks = PendingChunks( + chunk.chunkId, chunk.count, chunk.byteLength, 0, 0, [] + ) + pending = self._chunks + if ( + pending.chunk_id != chunk.chunkId + or pending.count != chunk.count + or pending.byte_length != chunk.byteLength + or pending.next_index != chunk.index + ): + raise OmpProtocolError("OMP RPC protocol chunk sequence mismatch") + pending.parts.append(data) + pending.received_bytes += len(data) + pending.next_index += 1 + if pending.received_bytes > pending.byte_length: + raise OmpProtocolError( + "OMP RPC protocol chunks exceed their declared length" + ) + if pending.next_index < pending.count: + return None + if pending.received_bytes != pending.byte_length: + raise OmpProtocolError("OMP RPC protocol chunk length mismatch") + self._chunks = None + decoded = _PAYLOAD_ADAPTER.validate_json(b"".join(pending.parts), strict=True) + if OmpFrame.model_validate(decoded).type == OmpFrameType.CHUNK: + raise OmpProtocolError("OMP RPC protocol chunks cannot nest") + return decoded + + async def _read_stdout(self, process: asyncio.subprocess.Process) -> None: + assert process.stdout is not None + try: + while True: + line = await process.stdout.readline() + if not line: + if self._chunks is not None: + raise OmpProtocolError( + "OMP RPC protocol chunks interrupted by stdout EOF" + ) + raise OmpProcessError("OMP RPC stdout EOF") + if len(line) > self._max_frame_bytes or not line.endswith(b"\n"): + raise OmpProtocolError("Invalid OMP RPC protocol frame length") + payload = self._decode_payload( + _PAYLOAD_ADAPTER.validate_json(line, strict=True) + ) + if payload is None: + continue + self._sequence += 1 + await self._dispatch_payload(payload, self._sequence) + except asyncio.CancelledError: + raise + except ValidationError as cause: + error = OmpProtocolError("Invalid OMP RPC protocol frame") + error.__cause__ = cause + self._record_failure(error) + except OmpProcessError as error: + self._record_failure(error) + except (ValueError, Base64Error) as cause: + error = OmpProtocolError("Invalid OMP RPC protocol frame") + error.__cause__ = cause + self._record_failure(error) + except Exception as cause: # noqa: BLE001 -- reader failure must unblock every pending request + error = OmpProcessError("OMP RPC stdout reader failed") + error.__cause__ = cause + self._record_failure(error) + + async def _dispatch_payload( + self, payload: dict[str, JsonValue], sequence: int + ) -> None: + frame = OmpFrame.model_validate(payload) + turn = self._turn + match frame.type: + case OmpFrameType.READY: + ready = OmpReady.model_validate(payload) + if self._ready is not None and not self._ready.done(): + self._ready.set_result(ready) + case OmpFrameType.FRAME_ERROR: + overflow = OmpFrameError.model_validate(payload) + raise OmpProtocolError( + f"OMP RPC protocol frame overflow: {overflow.error}" + ) + case OmpFrameType.RESPONSE: + self._dispatch_response(OmpResponse.model_validate(payload), sequence) + case OmpFrameType.PROMPT_RESULT: + result = OmpPromptResult.model_validate(payload) + if turn is not None and result.id == turn.id and not turn.result.done(): + turn.result.set_result(result) + case OmpFrameType.SESSION_SETTLED: + if turn is not None and ( + turn.last_activity_sequence > 0 or turn.result.done() + ): + turn.settled_sequence = sequence + turn.changed.set() + case OmpFrameType.HOST_TOOL_CALL: + call = OmpHostToolCall.model_validate(payload) + if turn is None: + raise OmpProtocolError( + "OMP host tool arrived without an active prompt" + ) + self._start_callback( + turn, "tool", call.id, self._handle_tool(turn, call) + ) + case OmpFrameType.HOST_TOOL_CANCEL: + cancellation = OmpHostToolCancel.model_validate(payload) + if turn is not None: + self._cancel_callback(turn, "tool", cancellation.targetId) + case OmpFrameType.UI_REQUEST: + request = OmpUIRequest.model_validate(payload) + if request.method == "cancel": + if request.targetId is None: + raise OmpProtocolError( + "OMP UI cancellation is missing targetId" + ) + if turn is not None: + self._cancel_callback(turn, "ui", request.targetId) + return + if request.method in _UI_NOTIFICATIONS: + return + if turn is None: + await self._write_frame( + { + "type": "extension_ui_response", + "id": request.id, + "cancelled": True, + } + ) + return + self._start_callback( + turn, "ui", request.id, self._handle_ui(turn, request) + ) + case event_type if event_type in _EVENT_TYPES: + event = OmpEvent.model_validate(payload) + if turn is not None: + turn.observe(event, sequence) + + def _dispatch_response(self, response: OmpResponse, sequence: int) -> None: + if response.id is None: + raise OmpProtocolError("OMP response is missing its command id") + pending = self._pending.get(response.id) + if pending is None: + if ( + not response.success + and self._turn is not None + and response.id == self._turn.id + ): + self._turn.fail( + OmpRpcError( + response.command, + response.error or "Command failed", + code=response.code, + ) + ) + return + if response.command != pending.command: + raise OmpProtocolError("OMP response command does not match its request") + if pending.future.done(): + return + if response.success: + pending.future.set_result(ReceivedResponse(response, sequence)) + else: + pending.future.set_result( + OmpRpcError( + response.command, + response.error or "Command failed", + code=response.code, + ) + ) + + def _start_callback( + self, + turn: PromptRun, + kind: str, + request_id: str, + callback: Coroutine[object, object, None], + ) -> None: + key = (kind, request_id) + task = asyncio.create_task(callback) + if key in turn.callbacks: + task.cancel() + raise OmpProtocolError("OMP callback request id was reused") + turn.callbacks[key] = task + turn.callback_tasks.add(task) + task.add_done_callback(lambda done: self._callback_finished(turn, key, done)) + + @staticmethod + def _cancel_callback(turn: PromptRun, kind: str, request_id: str) -> None: + task = turn.callbacks.pop((kind, request_id), None) + if task is not None: + task.cancel() + + @staticmethod + def _callback_finished( + turn: PromptRun, key: tuple[str, str], task: asyncio.Task[None] + ) -> None: + if turn.callbacks.get(key) is task: + turn.callbacks.pop(key) + turn.callback_tasks.discard(task) + if not task.cancelled() and (error := task.exception()) is not None: + turn.fail(error) + + async def _handle_tool(self, turn: PromptRun, call: OmpHostToolCall) -> None: + result = await turn.on_tool(call) + if ( + self._closing + or turn.callbacks.get(("tool", call.id)) is not asyncio.current_task() + ): + return + await self._write_frame( + { + "type": "host_tool_result", + "id": call.id, + "result": cast( + dict[str, JsonValue], + result.model_dump( + mode="json", exclude={"is_error"}, exclude_none=True + ), + ), + "isError": result.is_error, + } + ) + + async def _handle_ui(self, turn: PromptRun, request: OmpUIRequest) -> None: + response = await turn.on_ui(request) + if ( + self._closing + or turn.callbacks.get(("ui", request.id)) is not asyncio.current_task() + ): + return + frame: dict[str, JsonValue] = { + "type": "extension_ui_response", + "id": request.id, + } + if response.value is not None: + frame["value"] = response.value + elif response.confirmed is not None: + frame["confirmed"] = response.confirmed + else: + frame["cancelled"] = True + await self._write_frame(frame) + + async def _read_stderr(self, process: asyncio.subprocess.Process) -> None: + assert process.stderr is not None + while data := await process.stderr.read(_DRAIN_BYTES): + logger.debug( + "OMP stderr: %s", data.decode("utf-8", errors="replace").rstrip() + ) + + def _fail_pending(self, error: OmpProcessError) -> None: + for pending in self._pending.values(): + if not pending.future.done(): + pending.future.set_result(error) + if self._ready is not None and not self._ready.done(): + self._ready.set_result(error) + if self._turn is not None: + self._turn.fail(error) + + def _record_failure(self, error: OmpProcessError) -> None: + if self._failure is None: + self._failure = error + self._fail_pending(self._failure) + if not self._closing: + self._begin_close() + + def _begin_close(self) -> None: + self._closing = True + if self._close_task is None or ( + self._close_task.done() + and ( + self._close_task.cancelled() or self._close_task.exception() is not None + ) + ): + self._close_task = asyncio.create_task(self._close_process()) + self._close_task.add_done_callback(self._log_close_failure) + + @staticmethod + def _log_close_failure(task: asyncio.Task[None]) -> None: + if not task.cancelled() and (error := task.exception()) is not None: + logger.error("OMP process cleanup failed", exc_info=error) + + @staticmethod + def _consume_task(task: asyncio.Task[None]) -> None: + if not task.cancelled(): + task.exception() + + @staticmethod + async def _cancel_tasks(tasks: Sequence[asyncio.Task[None]]) -> None: + active = {task for task in tasks if not task.done()} + for task in active: + task.cancel() + if active: + _, pending = await asyncio.wait(active, timeout=_CALLBACK_EXIT_TIMEOUT_S) + for task in pending: + task.cancel() + logger.warning("OMP callback did not honor cancellation") + for task in tasks: + if task.done(): + OmpRpcClient._consume_task(task) + else: + task.add_done_callback(OmpRpcClient._consume_task) + + @staticmethod + async def _drain(stream: asyncio.StreamReader) -> None: + while await stream.read(_DRAIN_BYTES): + pass + + @staticmethod + async def _wait_for_exit(process: asyncio.subprocess.Process) -> bool: + try: + async with asyncio.timeout(_PROCESS_EXIT_TIMEOUT_S): + await process.wait() + except TimeoutError: + return False + return True + + @staticmethod + def _process_is_alive(process: psutil.Process) -> bool: + try: + return process.is_running() and process.status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + def _capture_descendants(self) -> None: + root = self._process_identity + if root is None or not self._process_is_alive(root): + return + try: + descendants = root.children(recursive=True) + except psutil.NoSuchProcess: + return + if not root.is_running(): + return + for process in descendants: + try: + process.create_time() + except psutil.NoSuchProcess: + continue + self._descendants.add(process) + + def _signal_root(self, *, kill: bool) -> None: + root = self._process_identity + if root is None: + if self._proc is not None and self._proc.returncode is None: + raise OmpProcessError("Cannot signal OMP without a captured identity") + return + try: + if kill: + root.kill() + else: + root.terminate() + except psutil.NoSuchProcess: + pass + + @classmethod + def _close_descendants(cls, processes: Sequence[psutil.Process]) -> None: + alive = list(processes) + errors: list[Exception] = [] + for kill in (False, True): + for process in alive: + try: + if kill: + process.kill() + else: + process.terminate() + except psutil.NoSuchProcess: + pass + except psutil.Error as cause: + errors.append(cause) + try: + _, alive = psutil.wait_procs(alive, timeout=_PROCESS_EXIT_TIMEOUT_S) + except psutil.Error as cause: + errors.append(cause) + remaining: list[psutil.Process] = [] + for process in alive: + try: + if cls._process_is_alive(process): + remaining.append(process) + except psutil.Error as cause: + errors.append(cause) + remaining.append(process) + alive = remaining + if not alive: + break + if alive: + pids = sorted(process.pid for process in alive) + errors.append(OmpProcessError(f"OMP descendants survived kill: {pids}")) + if len(errors) == 1: + raise errors[0] + if errors: + raise ExceptionGroup("OMP descendant cleanup failed", errors) + + async def _close_process(self) -> None: + process = self._proc + errors: list[Exception] = [] + if self._spawn_task is not None: + try: + process = await asyncio.shield(self._spawn_task) + except OSError: + return + except psutil.Error as cause: + errors.append(cause) + process = self._proc + if process is None: + return + # Abort and stdin EOF can orphan children before their parent is reaped. + try: + await asyncio.to_thread(self._capture_descendants) + except psutil.Error as cause: + errors.append(cause) + turn = self._turn + self._fail_pending(self._failure or OmpProcessError("OMP RPC client closed")) + if turn is not None: + workers = list(turn.callback_tasks) + if turn.event_worker is not None: + workers.append(turn.event_worker) + await self._cancel_tasks(workers) + if ( + self._failure is None + and process.returncode is None + and self._reader_task is not None + and not self._reader_task.done() + ): + try: + async with asyncio.timeout(_ABORT_TIMEOUT_S): + await self._request_frame("abort", {}, allow_closing=True) + except (TimeoutError, OmpRpcError, OmpProcessError): + pass + if self._reader_task is not None: + await self._cancel_tasks([self._reader_task]) + drains: list[asyncio.Task[None]] = [] + if process.stdout is not None: + drains.append(asyncio.create_task(self._drain(process.stdout))) + if ( + self._stderr_task is None or self._stderr_task.done() + ) and process.stderr is not None: + drains.append(asyncio.create_task(self._drain(process.stderr))) + try: + await asyncio.to_thread(self._capture_descendants) + except psutil.Error as cause: + errors.append(cause) + if process.stdin is not None: + process.stdin.close() + descendants = asyncio.create_task( + asyncio.to_thread(self._close_descendants, tuple(self._descendants)) + ) + try: + if not await self._wait_for_exit(process): + for kill in (False, True): + try: + await asyncio.to_thread(self._signal_root, kill=kill) + except (psutil.Error, OmpProcessError) as cause: + errors.append(cause) + if await self._wait_for_exit(process): + break + else: + errors.append( + OmpProcessError("OMP process did not exit after kill") + ) + finally: + try: + await descendants + except (psutil.Error, OmpProcessError, ExceptionGroup) as cause: + errors.append(cause) + finally: + if self._stderr_task is not None: + drains.append(self._stderr_task) + await self._cancel_tasks(drains) + if len(errors) == 1: + raise errors[0] + if errors: + raise ExceptionGroup("OMP process cleanup failed", errors) diff --git a/src/band/integrations/ompwire.py b/src/band/integrations/ompwire.py new file mode 100644 index 000000000..7e0042ce9 --- /dev/null +++ b/src/band/integrations/ompwire.py @@ -0,0 +1,301 @@ +"""Typed projections of OMP's native JSONL RPC frames.""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Literal + +from pydantic import ( + AliasChoices, + BaseModel, + ConfigDict, + Field, + JsonValue, + RootModel, + model_validator, +) + +OMP_FRAME_BYTES = 1024 * 1024 +OMP_REASSEMBLED_BYTES = 64 * 1024 * 1024 +OMP_CHUNK_BYTES = 256 * 1024 + + +class OmpFrameType(StrEnum): + READY = "ready" + RESPONSE = "response" + PROMPT_RESULT = "prompt_result" + SESSION_SETTLED = "session_settled" + HOST_TOOL_CALL = "host_tool_call" + HOST_TOOL_CANCEL = "host_tool_cancel" + UI_REQUEST = "extension_ui_request" + CHUNK = "rpc_chunk" + FRAME_ERROR = "rpc_frame_error" + + +class OmpEventType(StrEnum): + AGENT_START = "agent_start" + AGENT_END = "agent_end" + TURN_START = "turn_start" + TURN_END = "turn_end" + MESSAGE_START = "message_start" + MESSAGE_UPDATE = "message_update" + MESSAGE_END = "message_end" + TOOL_START = "tool_execution_start" + TOOL_UPDATE = "tool_execution_update" + TOOL_STREAM = "tool_stream_update" + TOOL_END = "tool_execution_end" + + +class OmpWireModel(BaseModel): + model_config = ConfigDict( + frozen=True, extra="ignore", strict=True, allow_inf_nan=False + ) + + +class OmpContent(OmpWireModel): + """Transcript content, including text, thoughts, images, and tool calls.""" + + type: str + text: str | None = None + thinking: str | None = None + data: str | None = None + mimeType: str | None = None + id: str | None = None + name: str | None = None + arguments: dict[str, JsonValue] | None = None + intent: str | None = None + + +class OmpUsageCost(OmpWireModel): + input: float = 0 + output: float = 0 + cacheRead: float = 0 + cacheWrite: float = 0 + total: float = 0 + + +class OmpUsage(OmpWireModel): + input: int = 0 + output: int = 0 + cacheRead: int = 0 + cacheWrite: int = 0 + totalTokens: int = 0 + contextTokens: int | None = None + reasoningTokens: int | None = None + cost: OmpUsageCost | None = None + + +class OmpMessage(OmpWireModel): + role: str + content: str | list[OmpContent] = Field(default_factory=list) + stopReason: str | None = None + errorMessage: str | None = None + usage: OmpUsage | None = None + model: str | None = None + provider: str | None = None + api: str | None = None + timestamp: float | None = None + toolCallId: str | None = None + toolName: str | None = None + isError: bool | None = None + + +class OmpToolResult(OmpWireModel): + content: list[OmpContent] + details: JsonValue = None + is_error: bool = Field( + default=False, validation_alias=AliasChoices("is_error", "isError") + ) + + @model_validator(mode="after") + def validate_content(self) -> OmpToolResult: + for content in self.content: + match content.type: + case "text" if content.text is not None: + continue + case "image" if ( + content.data is not None and content.mimeType is not None + ): + continue + case _: + raise ValueError("OMP tool content must be text or an inline image") + return self + + +class OmpAssistantEvent(OmpWireModel): + type: str + delta: str | None = None + contentIndex: int | None = None + content: JsonValue = None + + +class OmpEvent(OmpWireModel): + type: str + message: OmpMessage | None = None + messages: list[OmpMessage] = Field(default_factory=list) + messageId: str | None = None + assistantMessageEvent: OmpAssistantEvent | None = None + toolCallId: str | None = None + toolName: str | None = None + args: JsonValue = None + result: OmpToolResult | None = None + partialResult: JsonValue = None + update: JsonValue = None + isError: bool | None = None + intent: str | None = None + isTerminal: bool | None = None + yielded: bool | None = None + awaitingAsyncWork: bool | None = None + + @model_validator(mode="after") + def validate_event(self) -> OmpEvent: + match self.type: + case ( + OmpEventType.MESSAGE_START + | OmpEventType.MESSAGE_UPDATE + | OmpEventType.MESSAGE_END + | OmpEventType.TURN_END + ): + if self.message is None: + raise ValueError("OMP message event is missing its message") + case OmpEventType.AGENT_END: + if "messages" not in self.model_fields_set: + raise ValueError("OMP agent_end is missing its messages") + case ( + OmpEventType.TOOL_START + | OmpEventType.TOOL_UPDATE + | OmpEventType.TOOL_STREAM + | OmpEventType.TOOL_END + ): + if self.toolCallId is None or self.toolName is None: + raise ValueError("OMP tool event is missing its call identity") + return self + + +class OmpModel(OmpWireModel): + id: str + provider: str + name: str | None = None + reasoning: bool | None = None + contextWindow: int | None = None + + +class OmpState(OmpWireModel): + sessionId: str + sessionFile: str | None = None + isSettled: bool + model: OmpModel | None = None + thinkingLevel: str | None = None + isStreaming: bool = False + hasPendingAsyncWork: bool = False + queuedMessageCount: int = 0 + + +class OmpHostToolCall(OmpWireModel): + type: Literal["host_tool_call"] = "host_tool_call" + id: str + toolCallId: str + toolName: str + arguments: dict[str, JsonValue] + + +class OmpHostToolCancel(OmpWireModel): + type: Literal["host_tool_cancel"] = "host_tool_cancel" + id: str + targetId: str + + +class OmpUIRequest(OmpWireModel): + type: Literal["extension_ui_request"] = "extension_ui_request" + id: str + method: str + title: str | None = None + options: list[str] = Field(default_factory=list) + message: str | None = None + targetId: str | None = None + + +class OmpUIResponse(OmpWireModel): + value: str | None = None + confirmed: bool | None = None + cancelled: bool = False + + @model_validator(mode="after") + def validate_answer(self) -> OmpUIResponse: + if ( + sum((self.value is not None, self.confirmed is not None, self.cancelled)) + > 1 + ): + raise ValueError("OMP UI response must contain only one answer") + return self + + +class OmpReady(OmpWireModel): + type: Literal["ready"] + protocolVersion: Literal[1] + supportedProtocolVersions: list[int] + maxFrameBytes: int = Field(gt=0, le=OMP_FRAME_BYTES) + maxReassembledFrameBytes: int = Field(gt=0, le=OMP_REASSEMBLED_BYTES) + + +class OmpNegotiation(OmpWireModel): + protocolVersion: Literal[2] + + +class OmpChunk(OmpWireModel): + type: Literal["rpc_chunk"] + chunkId: str = Field(min_length=1, max_length=128) + index: int = Field(ge=0) + count: int = Field(ge=2, le=OMP_REASSEMBLED_BYTES // OMP_CHUNK_BYTES) + byteLength: int = Field(ge=OMP_FRAME_BYTES, le=OMP_REASSEMBLED_BYTES) + data: str = Field(min_length=1) + + @model_validator(mode="after") + def validate_index(self) -> OmpChunk: + if self.index >= self.count: + raise ValueError("OMP chunk index exceeds its count") + return self + + +class OmpFrameError(OmpWireModel): + type: Literal["rpc_frame_error"] + error: str + + +class OmpResponse(OmpWireModel): + type: Literal["response"] + id: str | None = None + command: str + success: bool + data: JsonValue = None + error: str | None = None + code: str | None = None + + +class OmpPromptAcceptance(OmpWireModel): + agentInvoked: bool | None = None + + +class OmpPromptFailure(OmpWireModel): + message: str + retryable: bool + provider: str | None = None + model: str | None = None + httpStatus: int | None = None + + +class OmpPromptResult(OmpWireModel): + type: Literal["prompt_result"] + id: str | None = None + agentInvoked: bool + status: Literal["completed", "error", "aborted"] + sessionSettled: bool + error: OmpPromptFailure | None = None + + +class OmpFrame(OmpWireModel): + type: str + + +class OmpCommand(RootModel[dict[str, JsonValue]]): + model_config = ConfigDict(strict=True, allow_inf_nan=False) diff --git a/tests/converters/test_omp.py b/tests/converters/test_omp.py new file mode 100644 index 000000000..cdb2abe50 --- /dev/null +++ b/tests/converters/test_omp.py @@ -0,0 +1,86 @@ +"""Tests for OmpHistoryConverter.""" + +from __future__ import annotations + +from band.converters.omp import ( + OMP_ROOM_METADATA_KEY, + OMP_SESSION_METADATA_KEY, + OmpHistoryConverter, +) + + +class TestOmpHistoryConverter: + def test_convert_empty_history(self) -> None: + state = OmpHistoryConverter().convert([]) + assert state.room_to_session == {} + assert state.replay_messages == [] + + def test_latest_session_file_wins_per_room(self) -> None: + state = OmpHistoryConverter().convert( + [ + { + "message_type": "task", + "metadata": { + OMP_ROOM_METADATA_KEY: "room-1", + OMP_SESSION_METADATA_KEY: "/sessions/old.jsonl", + }, + }, + { + "message_type": "task", + "metadata": { + OMP_ROOM_METADATA_KEY: "room-2", + OMP_SESSION_METADATA_KEY: "/sessions/other.jsonl", + }, + }, + { + "message_type": "task", + "metadata": { + OMP_ROOM_METADATA_KEY: "room-1", + OMP_SESSION_METADATA_KEY: "/sessions/new.jsonl", + }, + }, + ] + ) + assert state.room_to_session == { + "room-1": "/sessions/new.jsonl", + "room-2": "/sessions/other.jsonl", + } + + def test_unusable_metadata_is_ignored(self) -> None: + state = OmpHistoryConverter().convert( + [ + {"message_type": "task", "metadata": "not a mapping"}, + { + "message_type": "task", + "metadata": {OMP_ROOM_METADATA_KEY: "room-1"}, + }, + { + "message_type": "task", + "metadata": { + OMP_ROOM_METADATA_KEY: "room-1", + OMP_SESSION_METADATA_KEY: "", + }, + }, + { + "message_type": "task", + "metadata": { + OMP_ROOM_METADATA_KEY: 7, + OMP_SESSION_METADATA_KEY: "/sessions/x.jsonl", + }, + }, + ] + ) + assert state.room_to_session == {} + + def test_room_transcript_is_replayed(self) -> None: + state = OmpHistoryConverter().convert( + [ + { + "message_type": "text", + "content": "hello", + "sender": {"name": "Alice", "type": "User"}, + } + ] + ) + assert len(state.replay_messages) == 1 + assert "hello" in state.replay_messages[0] diff --git a/tests/framework_configs/converters.py b/tests/framework_configs/converters.py index 98b85b956..b930207df 100644 --- a/tests/framework_configs/converters.py +++ b/tests/framework_configs/converters.py @@ -319,9 +319,9 @@ def _build_gemini_config() -> ConverterConfig: # parsing and helpers are internal utility modules (shared helpers, not converters). # a2a / a2a_gateway use the A2A protocol which has a different message schema. # acp_client / acp_server use ACP protocol session updates, not standard convert(). -# codex, letta, and opencode are metadata-only converters that extract session state -# from task event metadata rather than converting message history. They don't implement -# the standard convert() -> framework-format contract that conformance tests validate. +# codex, omp, letta, and opencode extract session state from task-event metadata +# rather than converting message history. They don't implement the standard +# convert() -> framework-format contract that conformance tests validate. # crewai_flow is a metadata-only converter for orchestration state reconstructed from # task events; same exception as codex/letta/opencode. # slack is a metadata-only converter too: it recovers the Slack thread binding from the @@ -335,6 +335,7 @@ def _build_gemini_config() -> ConverterConfig: "acp_client", "acp_server", "codex", + "omp", "crewai_flow", "letta", "opencode", diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py new file mode 100644 index 000000000..4502c3614 --- /dev/null +++ b/tests/integrations/omppeer.py @@ -0,0 +1,593 @@ +"""Local OMP RPC peer with input-driven settlement and cancellation barriers.""" + +from __future__ import annotations + +import argparse +import base64 +import json +import os +import signal +import socket +import subprocess +import sys +import threading +from typing import cast + +import psutil +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +PROTOCOL_VERSION = 1 +FRAME_BYTES = 1024 * 1024 +CHUNK_BYTES = 256 * 1024 + + +class Frame(BaseModel): + model_config = ConfigDict(extra="allow") + __pydantic_extra__: dict[str, JsonValue] = Field(init=False) + + type: str + id: str | None = None + message: str | None = None + result: dict[str, JsonValue] | None = None + isError: bool | None = None + value: str | None = None + confirmed: bool | None = None + cancelled: bool | None = None + + +class ProcessReceipt(BaseModel): + pid: int + parent_pid: int + created: float + + +def start_descendant(scenario: str, port: int, stdio: str) -> subprocess.Popen[bytes]: + return subprocess.Popen( + ( + sys.executable, + "-u", + __file__, + "--scenario", + scenario, + "--gate-port", + str(port), + ), + # A parent's pending pipe read can block Python's stdin initialization on Windows. + # Descendants are controlled by the gate socket, not the RPC input pipe. + stdin=subprocess.DEVNULL, + stdout=None if stdio == "inherit" else subprocess.DEVNULL, + stderr=None if stdio == "inherit" else subprocess.DEVNULL, + ) + + +def hold_process_gate(port: int, *, spawn_child: bool) -> None: + if hasattr(signal, "SIGTERM"): + signal.signal(signal.SIGTERM, signal.SIG_IGN) + child = start_descendant("gated-child", port, "inherit") if spawn_child else None + with socket.create_connection(("127.0.0.1", port), timeout=5) as gate: + gate.settimeout(None) + receipt = ProcessReceipt( + pid=os.getpid(), + parent_pid=os.getppid(), + created=psutil.Process().create_time(), + ) + gate.sendall((receipt.model_dump_json() + "\n").encode()) + gate.recv(1) + if child is not None: + child.wait(timeout=5) + + +def send(frame: dict[str, JsonValue]) -> None: + payload = json.dumps(frame, separators=(",", ":")).encode("utf-8") + if PROTOCOL_VERSION == 2 and len(payload) + 1 > FRAME_BYTES: + count = (len(payload) + CHUNK_BYTES - 1) // CHUNK_BYTES + for index in range(count): + chunk: dict[str, JsonValue] = { + "type": "rpc_chunk", + "chunkId": "large-frame", + "index": index, + "count": count, + "byteLength": len(payload), + "data": base64.b64encode( + payload[index * CHUNK_BYTES : (index + 1) * CHUNK_BYTES] + ).decode("ascii"), + } + sys.stdout.write(json.dumps(chunk, separators=(",", ":")) + "\n") + else: + sys.stdout.write(payload.decode("utf-8") + "\n") + sys.stdout.flush() + + +def assistant(text: str, *, stop: str = "stop") -> dict[str, JsonValue]: + message: dict[str, JsonValue] = { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Considering the result"}, + {"type": "text", "text": text}, + ], + "api": "test-api", + "provider": "test-provider", + "model": "test-model", + "stopReason": stop, + "timestamp": 1, + "usage": { + "input": 7, + "output": 3, + "cacheRead": 2, + "cacheWrite": 1, + "totalTokens": 13, + "cost": { + "input": 0.1, + "output": 0.2, + "cacheRead": 0.01, + "cacheWrite": 0.02, + "total": 0.33, + }, + }, + } + if stop == "error": + message["errorMessage"] = "Resumed provider failed" + return message + + +class Peer: + def __init__(self, scenario: str) -> None: + self.scenario = scenario + self.settled = True + self.prompt_id: str | None = None + self.prompt_count = 0 + self.frames: list[dict[str, JsonValue]] = [] + self.tool_results: set[str] = set() + self.ui_answered = False + self.held: Frame | None = None + + def respond(self, frame: Frame, data: dict[str, JsonValue] | None = None) -> None: + result: dict[str, JsonValue] = { + "type": "response", + "id": frame.id, + "command": frame.type, + "success": True, + } + if data is not None: + result["data"] = data + send(result) + + def end( + self, + text: str, + *, + stop: str = "stop", + pending: bool = False, + compacted: bool = False, + ) -> None: + message = assistant(text, stop=stop) + send({"type": "message_end", "messageId": text, "message": message}) + send( + { + "type": "agent_end", + "messages": [] if compacted else [message], + "isTerminal": not pending, + "yielded": True, + "awaitingAsyncWork": pending, + } + ) + + def result(self, *, settled: bool, status: str = "completed") -> None: + result: dict[str, JsonValue] = { + "type": "prompt_result", + "id": self.prompt_id, + "agentInvoked": True, + "status": status, + "sessionSettled": settled, + } + if status == "error": + result["error"] = {"message": "Initial provider failed", "retryable": False} + send(result) + + def settle(self) -> None: + self.settled = True + send({"type": "session_settled"}) + + def finish(self, frame: Frame) -> None: + self.respond(frame) + send({"type": "agent_start"}) + self.end("complete") + self.result(settled=True) + self.settle() + + def invalid_chunk(self) -> None: + chunk: dict[str, JsonValue] = { + "type": "rpc_chunk", + "chunkId": "invalid-frame", + "index": 0, + "count": 2, + "byteLength": 1024 * 1024, + "data": "e30=", + } + match self.scenario: + case "chunk-overflow": + chunk["byteLength"] = 64 * 1024 * 1024 + 1 + send(chunk) + case "chunk-order": + chunk["index"] = 1 + send(chunk) + case "chunk-invalid-data": + chunk["data"] = "!" + send(chunk) + case "chunk-length": + send(chunk) + chunk["index"] = 1 + send(chunk) + case "chunk-eof": + send(chunk) + raise SystemExit(2) + case _: + send(chunk) + send({"type": "notice", "level": "info", "message": "Interrupted"}) + + def prompt(self, frame: Frame) -> None: + self.prompt_id = frame.id + self.prompt_count += 1 + self.settled = False + if self.prompt_count > 1 or frame.message == "second": + self.finish(frame) + return + match self.scenario: + case "rejected": + self.settled = True + send( + { + "type": "response", + "id": frame.id, + "command": "prompt", + "success": False, + "error": "Prompt rejected", + "code": "rejected", + } + ) + return + case "local": + self.settled = True + self.respond(frame, {"agentInvoked": False}) + return + case "bad-acceptance": + self.respond(frame, {"agentInvoked": "yes"}) + return + case "eof": + self.respond(frame) + send({"type": "agent_start"}) + raise SystemExit(2) + case "bad-json": + sys.stdout.write("{broken\n") + sys.stdout.flush() + return + case "bad-event": + self.respond(frame) + send({"type": "message_end", "message": {"role": 123}}) + return + case "wrong-command": + send( + { + "type": "response", + "id": frame.id, + "command": "get_state", + "success": True, + "data": {}, + } + ) + return + case "callbacks": + self.held = frame + send({"type": "agent_start"}) + for number in (1, 2): + send( + { + "type": "host_tool_call", + "id": f"host-{number}", + "toolCallId": f"call-{number}", + "toolName": "band_chat", + "arguments": {"number": number}, + } + ) + send( + { + "type": "extension_ui_request", + "id": "approval-1", + "method": "select", + "title": "Allow command?", + "options": ["Approve", "Deny"], + } + ) + return + case "stale": + send({"type": "session_settled"}) + send( + { + "type": "prompt_result", + "id": "old-prompt", + "agentInvoked": True, + "status": "completed", + "sessionSettled": True, + } + ) + case _: + pass + self.respond(frame) + send({"type": "agent_start"}) + match self.scenario: + case scenario if scenario.startswith("chunk-"): + self.invalid_chunk() + case "frame-overflow": + send( + { + "type": "rpc_frame_error", + "originalType": "message_end", + "error": "RPC frame exceeded the transport limit", + } + ) + case "large-image": + message = assistant("complete") + message["content"] = [ + { + "type": "image", + "mimeType": "image/png", + "data": "YWJj" * 300_000, + } + ] + send( + { + "type": "message_end", + "messageId": "large-image", + "message": message, + } + ) + send({"type": "agent_end", "messages": [message], "isTerminal": True}) + self.result(settled=True) + self.settle() + case "basic": + send({"type": "queue_update", "steering": [], "followUp": []}) + send({"type": "available_commands_update", "commands": []}) + send({"type": "model_changed"}) + send( + { + "type": "tool_execution_start", + "toolCallId": "native-call", + "toolName": "read", + "args": {"path": "note.txt"}, + } + ) + send( + { + "type": "tool_execution_end", + "toolCallId": "native-call", + "toolName": "read", + "result": { + "content": [{"type": "text", "text": "note"}], + "details": {"lines": 1}, + }, + "isError": False, + } + ) + self.end("complete") + self.result(settled=True) + self.settle() + case "settlement-before-result": + self.end("complete") + self.settle() + case "no-result": + self.end("complete") + self.settle() + case "host-cancel": + send( + { + "type": "host_tool_call", + "id": "host-cancelled", + "toolCallId": "call-cancelled", + "toolName": "band_chat", + "arguments": {}, + } + ) + case "ui-cancel" | "ui-error": + send( + { + "type": "extension_ui_request", + "id": "ui-cancelled", + "method": "select", + "title": "Allow command?", + "options": ["Approve", "Deny"], + } + ) + case "initial-error" | "initial-aborted": + stop = "error" if self.scenario == "initial-error" else "aborted" + self.end("failed", stop=stop) + self.result(settled=True, status=stop) + self.settle() + case _: + self.end("yielded", pending=True) + self.result(settled=self.scenario == "false-result-settled") + if self.scenario == "state-settles": + self.settled = True + + def handle(self, frame: Frame) -> None: + global PROTOCOL_VERSION + self.frames.append( + cast(dict[str, JsonValue], frame.model_dump(exclude_unset=True)) + ) + match frame.type: + case "negotiate_protocol": + if self.scenario == "negotiation-rejected": + send( + { + "type": "response", + "id": frame.id, + "command": frame.type, + "success": False, + "error": "Native protocol v2 is not supported", + } + ) + return + self.respond( + frame, + {"protocolVersion": 1 if self.scenario == "bad-negotiation" else 2}, + ) + PROTOCOL_VERSION = 2 + case "prompt": + self.prompt(frame) + case "get_state": + self.respond( + frame, + { + "sessionId": "session-local", + "sessionFile": "/local/session.jsonl", + "model": { + "id": "test-model", + "provider": "test-provider", + "name": "Test model", + }, + **( + {"isSettled": self.settled} + if self.scenario != "missing-state" + else {} + ), + "isStreaming": False, + "hasPendingAsyncWork": not self.settled, + "queuedMessageCount": 0, + }, + ) + if self.scenario == "settles-during-state" and not self.settled: + self.settle() + case "release": + if self.scenario == "settlement-before-result": + self.result(settled=False) + else: + send({"type": "agent_start"}) + stop = { + "resume-error": "error", + "resume-error-compacted": "error", + "resume-aborted": "aborted", + }.get(self.scenario, "stop") + self.end( + "background complete", + stop=stop, + compacted=self.scenario == "resume-error-compacted", + ) + self.settle() + self.respond(frame) + case "cancel-host": + send( + { + "type": "host_tool_cancel", + "id": "cancel-frame", + "targetId": "host-cancelled", + } + ) + self.end("host cancelled") + self.result(settled=True) + self.settle() + self.respond(frame) + case "cancel-ui": + send( + { + "type": "extension_ui_request", + "id": "ui-cancel-frame", + "method": "cancel", + "targetId": "ui-cancelled", + } + ) + self.end("ui cancelled") + self.result(settled=True) + self.settle() + self.respond(frame) + case "host_tool_result" | "extension_ui_response": + if frame.type == "host_tool_result" and frame.id is not None: + self.tool_results.add(frame.id) + if frame.type == "extension_ui_response": + self.ui_answered = True + if ( + self.scenario == "callbacks" + and self.tool_results == {"host-1", "host-2"} + and self.ui_answered + and self.held is not None + ): + self.respond(self.held) + self.held = None + self.end("callbacks complete") + self.result(settled=True) + self.settle() + case "checkpoint": + self.respond(frame, {"pid": os.getpid(), "frames": self.frames}) + case "abort": + if self.scenario != "stubborn": + self.respond(frame) + if self.prompt_id is not None: + self.end("aborted", stop="aborted") + self.settle() + case "first" | "hold": + self.held = frame + case "second": + self.respond(frame, {"value": "second"}) + if self.held is not None: + self.respond(self.held, {"value": "first"}) + self.held = None + case "error": + send( + { + "type": "response", + "id": frame.id, + "command": frame.type, + "success": False, + "error": "Command failed", + "code": "bad_command", + } + ) + case "crash": + raise SystemExit(3) + case "stderr": + sys.stderr.write("x" * (2 * 1024 * 1024)) + sys.stderr.flush() + self.respond(frame, {"drained": True}) + case _: + self.respond(frame) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--scenario", required=True) + parser.add_argument("--gate-port", type=int) + parser.add_argument( + "--child-stdio", choices=("inherit", "discard"), default="inherit" + ) + args = parser.parse_args() + if args.scenario in {"gated-parent", "gated-child"}: + if args.gate_port is None: + raise ValueError("A process gate port is required") + hold_process_gate(args.gate_port, spawn_child=args.scenario == "gated-parent") + return + if args.scenario in {"descendants", "no-ready-descendants"}: + if args.gate_port is None: + raise ValueError("A process gate port is required") + # The gate, not the peer's stdio EOF, owns these children's lifetime. + start_descendant("gated-parent", args.gate_port, args.child_stdio) + if args.scenario == "stubborn" and hasattr(signal, "SIGTERM"): + signal.signal(signal.SIGTERM, signal.SIG_IGN) + peer = Peer(args.scenario) + if args.scenario not in {"no-ready", "no-ready-descendants"}: + send( + { + "type": "ready", + "protocolVersion": 1, + "supportedProtocolVersions": [1] + if args.scenario == "no-v2" + else [1, 2], + "maxFrameBytes": 1024 * 1024, + "maxReassembledFrameBytes": 64 * 1024 * 1024, + } + ) + for line in sys.stdin: + peer.handle(Frame.model_validate_json(line, strict=True)) + if args.scenario == "stubborn": + threading.Event().wait() + + +if __name__ == "__main__": + main() diff --git a/tests/integrations/omppermissionpeer.py b/tests/integrations/omppermissionpeer.py new file mode 100644 index 000000000..6cd4c44d1 --- /dev/null +++ b/tests/integrations/omppermissionpeer.py @@ -0,0 +1,109 @@ +"""Script independent UI decisions and native events over the existing RPC peer.""" + +from __future__ import annotations + +import json +import os +import sys + +from tests.integrations.omppeer import Frame, Peer, send + + +class PermissionPeer(Peer): + def __init__(self, order: str) -> None: + super().__init__("permission-records") + self.order = order + + @staticmethod + def native_start() -> None: + send( + { + "type": "tool_execution_start", + "toolCallId": "ui-wire-id", + "toolName": os.environ.get("OMP_TEST_NATIVE_NAME", "read"), + "args": json.loads( + os.environ.get("OMP_TEST_NATIVE_ARGUMENTS", '{"path":"denied.txt"}') + ), + } + ) + + @staticmethod + def native_end() -> None: + send( + { + "type": "tool_execution_end", + "toolCallId": "ui-wire-id", + "toolName": os.environ.get("OMP_TEST_NATIVE_NAME", "read"), + "isError": os.environ.get("OMP_TEST_NATIVE_ERROR", "1") == "1", + "result": { + "content": [ + { + "type": "text", + "text": ( + "Native tool failed" + if os.environ.get("OMP_TEST_NATIVE_ERROR", "1") == "1" + else "Native tool finished" + ), + } + ] + }, + } + ) + + @staticmethod + def ask_ui() -> None: + send( + { + "type": "extension_ui_request", + "id": "ui-wire-id", + "method": "select", + "title": os.environ["OMP_TEST_PERMISSION_TITLE"], + "options": ["Approve", "Deny"], + } + ) + + def prompt(self, frame: Frame) -> None: + self.prompt_id = frame.id + self.settled = False + self.respond(frame) + send({"type": "agent_start"}) + if self.order == "native-first": + self.native_start() + else: + self.ask_ui() + + def handle(self, frame: Frame) -> None: + if frame.type == "permission-after-native": + self.respond(frame) + self.ask_ui() + return + if frame.type == "extension_ui_response": + if frame.id != "ui-wire-id" or not frame.cancelled: + raise AssertionError("The scripted UI decision must be cancelled") + if self.order == "permission-first": + self.native_start() + self.native_end() + self.end("") + self.result(settled=True) + self.settle() + return + super().handle(frame) + + +def main() -> None: + peer = PermissionPeer(sys.argv[1]) + send( + { + "type": "ready", + "protocolVersion": 1, + "supportedProtocolVersions": [1, 2], + "maxFrameBytes": 1024 * 1024, + "maxReassembledFrameBytes": 64 * 1024 * 1024, + } + ) + for line in sys.stdin: + peer.handle(Frame.model_validate_json(line, strict=True)) + + +if __name__ == "__main__": + main() diff --git a/tests/integrations/ompprocess/__init__.py b/tests/integrations/ompprocess/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/integrations/ompprocess/test_launch.py b/tests/integrations/ompprocess/test_launch.py new file mode 100644 index 000000000..4958592be --- /dev/null +++ b/tests/integrations/ompprocess/test_launch.py @@ -0,0 +1,47 @@ +"""Native OMP launch settings: command shape and provider environment.""" + +from __future__ import annotations + +from band.integrations.omp import ( + OMP_APPROVAL_MODE_ALWAYS_ASK, + OMP_APPROVAL_MODE_FLAG, + finalize_omp_command, +) +from band.integrations.ompprocess import ( + DEFAULT_OMP_COMMAND, + omp_command_in_workspace, + omp_runtime_env, +) + + +class TestNativeLaunchSettings: + def test_default_command_is_safe_after_finalize(self) -> None: + assert finalize_omp_command(DEFAULT_OMP_COMMAND)[-2:] == [ + OMP_APPROVAL_MODE_FLAG, + OMP_APPROVAL_MODE_ALWAYS_ASK, + ] + + def test_workspace_is_selected_without_a_subprocess_cwd(self) -> None: + assert omp_command_in_workspace(("bunx", *DEFAULT_OMP_COMMAND, "-v"), "/w") == [ + "bunx", + *DEFAULT_OMP_COMMAND, + "-v", + "--cwd=/w", + ] + + def test_runtime_env_routes_the_selected_provider_key(self) -> None: + assert omp_runtime_env(env=None, model="openai/custom", api_key="secret") == { + "OPENAI_API_KEY": "secret" + } + + def test_explicit_env_overrides_the_derived_credential(self) -> None: + env = omp_runtime_env( + env={"OPENAI_API_KEY": "override", "EXTRA": "1"}, + model="openai/custom", + api_key="secret", + ) + assert env == {"OPENAI_API_KEY": "override", "EXTRA": "1"} + + def test_no_credentials_and_no_env_means_inherit(self) -> None: + assert omp_runtime_env(env=None, model=None, api_key=None) is None + assert omp_runtime_env(env={}, model="openai/custom", api_key=None) is None diff --git a/tests/integrations/ompprocess/test_mcp_plugin.py b/tests/integrations/ompprocess/test_mcp_plugin.py new file mode 100644 index 000000000..c53f924a1 --- /dev/null +++ b/tests/integrations/ompprocess/test_mcp_plugin.py @@ -0,0 +1,79 @@ +"""The SDK-owned plugin projects MCP servers and connection rules for one process.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from band.integrations.ompprocess import ( + OmpMcpServer, + omp_mcp_arguments, + omp_mcp_plugin, + omp_mcp_ready_command, +) + + +class TestMcpPluginProjection: + def test_nothing_to_project_yields_no_plugin(self, tmp_path: Path) -> None: + with omp_mcp_plugin({}, workspace=str(tmp_path)) as plugin: + assert plugin is None + + def test_servers_land_in_the_plugin_mcp_config(self, tmp_path: Path) -> None: + servers = { + "local": OmpMcpServer(command="./server", args=("--port", "1")), + "remote": OmpMcpServer(type="sse", url="https://mcp.example/sse"), + } + with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: + assert plugin is not None + root = Path(plugin) + assert root.name == "band-sdk" + config = json.loads((root / ".mcp.json").read_text(encoding="utf-8")) + assert config["mcpServers"]["remote"] == { + "type": "sse", + "url": "https://mcp.example/sse", + } + local = config["mcpServers"]["local"] + assert local["args"] == ["--port", "1"] + assert Path(local["cwd"]) == tmp_path.resolve() + assert Path(local["command"]) == (tmp_path / "server").resolve() + startup = (root / "startup.js").read_text(encoding="utf-8") + assert omp_mcp_ready_command(plugin) in startup + assert '"band-sdk:local"' in startup + assert '"band-sdk:remote"' in startup + assert not Path(plugin).exists() + + def test_connection_instructions_become_an_always_applied_rule( + self, tmp_path: Path + ) -> None: + with omp_mcp_plugin( + {}, workspace=str(tmp_path), system_prompt="Reply through Band." + ) as plugin: + assert plugin is not None + rules = list((Path(plugin) / "rules").glob("band-connection-*.md")) + assert len(rules) == 1 + assert rules[0].read_text(encoding="utf-8") == ( + "---\nalwaysApply: true\n---\nReply through Band." + ) + + def test_arguments_load_the_plugin_and_its_startup_extension(self) -> None: + plugin = "/tmp/overlay/band-sdk" + arguments = omp_mcp_arguments(plugin) + assert arguments[:2] == ("--plugin-dir", plugin) + assert "--mode=rpc-ui" in arguments + assert str(Path(plugin) / "startup.js") in arguments + assert str(Path(plugin) / "startup.json") in arguments + + @pytest.mark.parametrize( + "fields", + [ + {}, + {"command": "server", "url": "https://mcp.example"}, + {"command": "server", "type": "http"}, + {"url": "https://mcp.example", "type": "stdio"}, + ], + ) + def test_a_server_needs_exactly_one_transport(self, fields: dict[str, str]) -> None: + with pytest.raises(ValueError): + OmpMcpServer(**fields) diff --git a/tests/integrations/ompprocess/test_mcp_startup.py b/tests/integrations/ompprocess/test_mcp_startup.py new file mode 100644 index 000000000..f25a91268 --- /dev/null +++ b/tests/integrations/ompprocess/test_mcp_startup.py @@ -0,0 +1,107 @@ +"""Required MCP connections cannot be inferred from an otherwise ready OMP process.""" + +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest + +from band.integrations.ompprocess import ( + OmpMcpServer, + OmpStartupError, + omp_mcp_ready, + omp_mcp_ready_command, +) +from band.integrations.omprpc import OmpRpcClient + + +class TestRequiredMcpStartup: + @pytest.mark.parametrize( + ("receipt", "reason"), + [ + (None, "no valid SDK startup receipt"), + ("not JSON", "no valid SDK startup receipt"), + ( + {"systemPromptReady": True, "servers": []}, + "did not acknowledge every configured MCP server", + ), + ( + { + "systemPromptReady": True, + "servers": [ + {"type": "connected", "serverName": "band-sdk:required"}, + {"type": "connected", "serverName": "band-sdk:required"}, + ], + }, + "did not acknowledge every configured MCP server", + ), + ( + { + "systemPromptReady": True, + "servers": [{"type": "connected", "serverName": "unrelated"}], + }, + "unexpected MCP startup receipt", + ), + ( + { + "systemPromptReady": True, + "servers": [{"type": "failed", "serverName": "band-sdk:required"}], + }, + "no valid SDK startup receipt", + ), + ( + { + "systemPromptReady": True, + "servers": [ + { + "type": "failed", + "serverName": "band-sdk:required", + "error": "initialize rejected", + } + ], + }, + "band-sdk:required failed: initialize rejected", + ), + ( + { + "systemPromptReady": False, + "servers": [ + {"type": "connected", "serverName": "band-sdk:required"} + ], + }, + "required Band connection instructions", + ), + ], + ) + async def test_rejects_incomplete_or_failed_receipts( + self, + tmp_path: Path, + receipt: str | dict[str, object] | None, + reason: str, + ) -> None: + plugin = tmp_path / "band-sdk" + plugin.mkdir() + if receipt is not None: + payload = json.dumps(receipt) if isinstance(receipt, dict) else receipt + (plugin / "receipt.json").write_text(payload, encoding="utf-8") + runtime = AsyncMock(spec=OmpRpcClient) + runtime.request.return_value = { + "commands": [{"name": omp_mcp_ready_command(str(plugin))}], + } + with pytest.raises(OmpStartupError, match=reason): + await omp_mcp_ready( + runtime, + str(plugin), + {"required": OmpMcpServer(command="server")}, + ) + + async def test_missing_startup_extension_cannot_fall_through_to_inference( + self, tmp_path: Path + ) -> None: + runtime = AsyncMock(spec=OmpRpcClient) + runtime.request.return_value = {"commands": []} + with pytest.raises(OmpStartupError, match="required SDK startup extension"): + await omp_mcp_ready(runtime, str(tmp_path), {}) + runtime.prompt.assert_not_awaited() diff --git a/tests/integrations/ompprocess/test_tool_names.py b/tests/integrations/ompprocess/test_tool_names.py new file mode 100644 index 000000000..ae8d3a52a --- /dev/null +++ b/tests/integrations/ompprocess/test_tool_names.py @@ -0,0 +1,52 @@ +"""Tool names OMP mints under the SDK plugin namespace map back to Band tools.""" + +from __future__ import annotations + +import pytest + +from band.integrations.ompprocess import ( + normalize_native_mcp_device_call, + normalize_native_mcp_tool_name, +) + +OWN = frozenset({"band_send_message"}) + + +class TestPluginNamespacedToolNames: + @pytest.mark.parametrize( + "path", + [ + "xd://mcp__band_send_message", + "xd://mcp__band_sdk_band_band_send_message", + ], + ) + def test_executed_device_write_maps_to_the_band_tool(self, path: str) -> None: + arguments = {"path": path, "content": '{"chat_id":"room-1","content":"hi"}'} + assert normalize_native_mcp_device_call( + "write", arguments, OWN, kind="execute" + ) == ("band_send_message", {"chat_id": "room-1", "content": "hi"}) + + @pytest.mark.parametrize( + "title", ["mcp__band_send_message", "mcp__band_sdk_band_band_send_message"] + ) + def test_mcp_title_maps_to_the_band_tool(self, title: str) -> None: + assert normalize_native_mcp_tool_name(title, OWN) == "band_send_message" + assert normalize_native_mcp_device_call(title, {"chat_id": "r1"}, OWN) == ( + "band_send_message", + {"chat_id": "r1"}, + ) + + def test_a_namespaced_read_is_never_the_band_tool(self) -> None: + arguments = { + "path": "xd://mcp__band_sdk_band_band_send_message", + "content": '{"content":"hi"}', + } + assert normalize_native_mcp_device_call( + "read", arguments, OWN, kind="read" + ) == ("read", arguments) + + def test_a_foreign_namespace_stays_unmapped(self) -> None: + assert ( + normalize_native_mcp_tool_name("mcp__other_sdk_band_send_message", OWN) + is None + ) diff --git a/tests/integrations/omprpc/__init__.py b/tests/integrations/omprpc/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/integrations/omprpc/conftest.py b/tests/integrations/omprpc/conftest.py new file mode 100644 index 000000000..c6b3f8724 --- /dev/null +++ b/tests/integrations/omprpc/conftest.py @@ -0,0 +1,65 @@ +"""Client lifecycle and process gate fixtures for OMP RPC tests.""" + +from __future__ import annotations + +import asyncio +import sys +from collections.abc import AsyncIterator + +import pytest +import pytest_asyncio + +pytest.importorskip("psutil", reason="psutil not installed (band-sdk[omp])") + +from band.integrations.omprpc import OmpRpcClient +from tests.integrations.omprpc.kit import ( + OMP_PEER_SCRIPT, + ClientFactory, + ProcessGate, + kill_gate_processes, +) + + +@pytest_asyncio.fixture(loop_scope="function") +async def client_factory() -> AsyncIterator[ClientFactory]: + clients: list[OmpRpcClient] = [] + + async def build(scenario: str) -> OmpRpcClient: + client = OmpRpcClient( + ( + sys.executable, + "-u", + OMP_PEER_SCRIPT, + "--scenario", + scenario, + ) + ) + clients.append(client) + await client.start() + return client + + try: + yield build + finally: + for client in clients: + await client.close() + + +@pytest_asyncio.fixture(loop_scope="function") +async def process_gate() -> AsyncIterator[ProcessGate]: + gate = ProcessGate() + server = await asyncio.start_server(gate.accept, "127.0.0.1", 0) + try: + sockets = server.sockets + assert sockets + gate.port = sockets[0].getsockname()[1] + yield gate + finally: + server.close() + for writer in gate.writers: + writer.close() + await asyncio.gather( + *(writer.wait_closed() for writer in gate.writers), return_exceptions=True + ) + await server.wait_closed() + await asyncio.to_thread(kill_gate_processes, gate.processes) diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py new file mode 100644 index 000000000..66253d247 --- /dev/null +++ b/tests/integrations/omprpc/kit.py @@ -0,0 +1,133 @@ +"""Shared observers and process gates for OMP RPC tests.""" + +from __future__ import annotations + +import asyncio +import time +from collections.abc import Awaitable, Callable +from contextlib import suppress +from dataclasses import dataclass +from typing import cast + +import psutil +from pydantic import JsonValue + +from band.integrations.omprpc import OmpRpcClient +from band.integrations.ompwire import ( + OmpContent, + OmpEvent, + OmpHostToolCall, + OmpToolResult, + OmpUIRequest, + OmpUIResponse, +) +from tests.integrations import omppeer +from tests.integrations.omppeer import ProcessReceipt + +ClientFactory = Callable[[str], Awaitable[OmpRpcClient]] + +# The fake OMP peer every RPC test launches as its subprocess. +OMP_PEER_SCRIPT = omppeer.__file__ + + +class Observer: + def __init__(self) -> None: + self.events: list[OmpEvent] = [] + self.queue: asyncio.Queue[OmpEvent] = asyncio.Queue() + + async def on_event(self, event: OmpEvent) -> None: + self.events.append(event) + self.queue.put_nowait(event) + + async def on_tool(self, call: OmpHostToolCall) -> OmpToolResult: + return OmpToolResult(content=[OmpContent(type="text", text=call.toolName)]) + + async def on_ui(self, request: OmpUIRequest) -> OmpUIResponse: + return OmpUIResponse(cancelled=True) + + async def wait_for(self, event_type: str) -> OmpEvent: + async with asyncio.timeout(5): + while True: + event = await self.queue.get() + if event.type == event_type: + return event + + +def prompt_task( + client: OmpRpcClient, observer: Observer, message: str = "first" +) -> asyncio.Task[None]: + return asyncio.create_task( + client.prompt( + message, + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=observer.on_ui, + ) + ) + + +async def second_prompt(client: OmpRpcClient) -> Observer: + observer = Observer() + async with asyncio.timeout(5): + await prompt_task(client, observer, "second") + assert any( + event.message is not None + and isinstance(event.message.content, list) + and any(content.text == "complete" for content in event.message.content) + for event in observer.events + ) + return observer + + +async def checkpoint_frames(client: OmpRpcClient) -> list[dict[str, JsonValue]]: + checkpoint = await client.request("checkpoint") + return cast(list[dict[str, JsonValue]], checkpoint["frames"]) + + +def process_is_alive(process: psutil.Process) -> bool: + try: + return process.is_running() and process.status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + +def kill_gate_processes(processes: list[psutil.Process]) -> None: + for process in processes: + with suppress(psutil.NoSuchProcess): + process.kill() + deadline = time.monotonic() + 5 + while alive := [process.pid for process in processes if process_is_alive(process)]: + if time.monotonic() >= deadline: + raise RuntimeError(f"Process gate cleanup left live processes: {alive}") + time.sleep(0.01) + + +@dataclass(frozen=True) +class GatedProcess: + receipt: ProcessReceipt + process: psutil.Process + + +class ProcessGate: + def __init__(self) -> None: + self.port = 0 + self.connections: asyncio.Queue[ + tuple[asyncio.StreamReader, asyncio.StreamWriter] + ] = asyncio.Queue() + self.writers: list[asyncio.StreamWriter] = [] + self.processes: list[psutil.Process] = [] + + def accept( + self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + self.writers.append(writer) + self.connections.put_nowait((reader, writer)) + + async def connection(self) -> GatedProcess: + async with asyncio.timeout(5): + reader, _ = await self.connections.get() + receipt = ProcessReceipt.model_validate_json(await reader.readline()) + process = await asyncio.to_thread(psutil.Process, receipt.pid) + self.processes.append(process) + assert await asyncio.to_thread(process.create_time) == receipt.created + return GatedProcess(receipt, process) diff --git a/tests/integrations/omprpc/test_callbacks.py b/tests/integrations/omprpc/test_callbacks.py new file mode 100644 index 000000000..f174a091c --- /dev/null +++ b/tests/integrations/omprpc/test_callbacks.py @@ -0,0 +1,179 @@ +"""Host tool and UI callbacks run concurrently and propagate cancellation or errors.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.ompwire import ( + OmpContent, + OmpHostToolCall, + OmpToolResult, + OmpUIRequest, + OmpUIResponse, +) +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + checkpoint_frames, + second_prompt, +) + + +class TestHostToolAndUICallbacks: + async def test_host_tools_and_ui_run_concurrently_even_before_prompt_acceptance( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("callbacks") + observer = Observer() + first_started = asyncio.Event() + release_first = asyncio.Event() + calls: list[OmpHostToolCall] = [] + approvals: list[OmpUIRequest] = [] + + async def on_tool(call: OmpHostToolCall) -> OmpToolResult: + calls.append(call) + if call.arguments["number"] == 1: + first_started.set() + await release_first.wait() + else: + await first_started.wait() + await client.request("checkpoint") + release_first.set() + return OmpToolResult( + content=[ + OmpContent(type="text", text=call.toolCallId), + OmpContent(type="image", data="aW1hZ2U=", mimeType="image/png"), + ], + details={"number": call.arguments["number"]}, + is_error=call.arguments["number"] == 2, + ) + + async def on_ui(request: OmpUIRequest) -> OmpUIResponse: + approvals.append(request) + return OmpUIResponse(value="Approve") + + async with asyncio.timeout(5): + await client.prompt( + "first", on_event=observer.on_event, on_tool=on_tool, on_ui=on_ui + ) + assert {call.id for call in calls} == {"host-1", "host-2"} + assert [(request.method, request.options) for request in approvals] == [ + ("select", ["Approve", "Deny"]) + ] + frames = await checkpoint_frames(client) + result = next(frame for frame in frames if frame.get("id") == "host-2") + assert result == { + "type": "host_tool_result", + "id": "host-2", + "result": { + "content": [ + {"type": "text", "text": "call-2"}, + {"type": "image", "data": "aW1hZ2U=", "mimeType": "image/png"}, + ], + "details": {"number": 2}, + }, + "isError": True, + } + ui_response = next(frame for frame in frames if frame.get("id") == "approval-1") + assert ui_response == { + "type": "extension_ui_response", + "id": "approval-1", + "value": "Approve", + } + await second_prompt(client) + + @pytest.mark.parametrize("kind", ["host", "ui"]) + async def test_native_cancel_targets_the_callback_and_keeps_the_next_turn_usable( + self, client_factory: ClientFactory, kind: str + ) -> None: + client = await client_factory(f"{kind}-cancel") + observer = Observer() + started = asyncio.Event() + cancelled = asyncio.Event() + forever = asyncio.Event() + + async def on_tool(call: OmpHostToolCall) -> OmpToolResult: + started.set() + try: + await forever.wait() + finally: + cancelled.set() + return OmpToolResult(content=[]) + + async def on_ui(request: OmpUIRequest) -> OmpUIResponse: + started.set() + try: + await forever.wait() + finally: + cancelled.set() + return OmpUIResponse(value="Approve") + + running = asyncio.create_task( + client.prompt( + "first", on_event=observer.on_event, on_tool=on_tool, on_ui=on_ui + ) + ) + try: + async with asyncio.timeout(5): + await started.wait() + await client.request(f"cancel-{kind}") + await cancelled.wait() + await running + frames = await checkpoint_frames(client) + target = "host-cancelled" if kind == "host" else "ui-cancelled" + assert not any(frame.get("id") == target for frame in frames) + await second_prompt(client) + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + async def test_ui_callback_failure_fails_closed_and_reaps_the_process( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("ui-error") + process = client._proc + assert process is not None + + async def on_ui(request: OmpUIRequest) -> OmpUIResponse: + raise ValueError("Unoffered permission choice") + + observer = Observer() + async with asyncio.timeout(6): + with pytest.raises(ValueError, match="Unoffered permission choice"): + await client.prompt( + "first", + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=on_ui, + ) + assert process.returncode is not None + + async def test_tool_callback_exception_propagates_unchanged_and_reaps_the_process( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("callbacks") + process = client._proc + assert process is not None + observer = Observer() + sentinel = ValueError("Band delivery failed") + + async def on_tool(call: OmpHostToolCall) -> OmpToolResult: + raise sentinel + + async with asyncio.timeout(6): + with pytest.raises(ValueError, match="Band delivery failed") as failed: + await client.prompt( + "first", + on_event=observer.on_event, + on_tool=on_tool, + on_ui=observer.on_ui, + ) + assert failed.value is sentinel + assert process.returncode is not None + await client.start() + await second_prompt(client) diff --git a/tests/integrations/omprpc/test_commands.py b/tests/integrations/omprpc/test_commands.py new file mode 100644 index 000000000..2c9efca4a --- /dev/null +++ b/tests/integrations/omprpc/test_commands.py @@ -0,0 +1,40 @@ +"""Command responses correlate independently and recover from errors or cancellation.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.omprpc import OmpRpcError +from tests.integrations.omprpc.kit import ClientFactory, second_prompt + + +class TestCommandCorrelation: + async def test_command_responses_are_correlated_out_of_order_and_errors_are_recoverable( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("basic") + async with asyncio.timeout(5): + first, second = await asyncio.gather( + client.request("first"), client.request("second") + ) + assert (first, second) == ({"value": "first"}, {"value": "second"}) + with pytest.raises(OmpRpcError, match="Command failed") as failed: + await client.request("error") + assert (failed.value.command, failed.value.code) == ("error", "bad_command") + await second_prompt(client) + + async def test_cancelled_command_drops_only_its_waiter( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("basic") + pending = asyncio.create_task(client.request("hold")) + await client.request("checkpoint") + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + assert await client.request("second") == {"value": "second"} + await second_prompt(client) diff --git a/tests/integrations/omprpc/test_interruption.py b/tests/integrations/omprpc/test_interruption.py new file mode 100644 index 000000000..74bd9244a --- /dev/null +++ b/tests/integrations/omprpc/test_interruption.py @@ -0,0 +1,110 @@ +"""Interrupted prompts and cancelled closes await callback cancellation and process reaping.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.ompwire import OmpHostToolCall, OmpToolResult +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + prompt_task, + second_prompt, +) + + +class TestInterruptionAndReaping: + @pytest.mark.parametrize("scenario", ["yield", "no-result", "stubborn"]) + @pytest.mark.parametrize("interruption", ["timeout", "cancel"]) + async def test_interrupted_prompt_aborts_reaps_and_can_restart( + self, client_factory: ClientFactory, scenario: str, interruption: str + ) -> None: + client = await client_factory(scenario) + process = client._proc + assert process is not None + observer = Observer() + running = prompt_task(client, observer) + try: + await observer.wait_for("agent_end") + if interruption == "cancel": + running.cancel() + async with asyncio.timeout(6): + with pytest.raises(asyncio.CancelledError): + await running + else: + with pytest.raises(TimeoutError): + async with asyncio.timeout(0.05): + await running + assert process.returncode is not None + async with asyncio.timeout(5): + await client.start() + await second_prompt(client) + assert client._proc is not process + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + async def test_repeatedly_cancelled_close_still_waits_for_reap( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("stubborn") + process = client._proc + assert process is not None + closing = asyncio.create_task(client.close()) + loop = asyncio.get_running_loop() + cancellations = [ + loop.call_later(delay, closing.cancel) for delay in (0.01, 0.05, 0.1) + ] + try: + async with asyncio.timeout(6): + with pytest.raises(asyncio.CancelledError): + await closing + assert process.returncode is not None + await client.close() + finally: + for cancellation in cancellations: + cancellation.cancel() + await asyncio.gather(closing, return_exceptions=True) + + async def test_cancelled_prompt_cancels_its_host_callback_before_reap( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("host-cancel") + process = client._proc + assert process is not None + observer = Observer() + started = asyncio.Event() + cancelled = asyncio.Event() + forever = asyncio.Event() + + async def on_tool(call: OmpHostToolCall) -> OmpToolResult: + started.set() + try: + await forever.wait() + finally: + cancelled.set() + return OmpToolResult(content=[]) + + running = asyncio.create_task( + client.prompt( + "first", + on_event=observer.on_event, + on_tool=on_tool, + on_ui=observer.on_ui, + ) + ) + try: + async with asyncio.timeout(6): + await started.wait() + running.cancel() + with pytest.raises(asyncio.CancelledError): + await running + assert cancelled.is_set() + assert process.returncode is not None + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) diff --git a/tests/integrations/omprpc/test_process_tree.py b/tests/integrations/omprpc/test_process_tree.py new file mode 100644 index 000000000..d1e881378 --- /dev/null +++ b/tests/integrations/omprpc/test_process_tree.py @@ -0,0 +1,127 @@ +"""Cancelled cleanup reaps descendants without affecting unrelated processes.""" + +from __future__ import annotations + +import asyncio +import sys +from contextlib import suppress + +import psutil +import pytest + +from band.integrations.omprpc import OmpRpcClient +from tests.integrations.omprpc.kit import ( + OMP_PEER_SCRIPT, + ProcessGate, + process_is_alive, +) + + +class TestProcessTreeCleanup: + @pytest.mark.parametrize("stdio", ["inherit", "discard"]) + @pytest.mark.parametrize("interruption", ["startup", "close"]) + async def test_cancelled_cleanup_reaps_descendants_but_not_unrelated_processes( + self, + process_gate: ProcessGate, + stdio: str, + interruption: str, + ) -> None: + peer = OMP_PEER_SCRIPT + port = str(process_gate.port) + unrelated = await asyncio.create_subprocess_exec( + sys.executable, + "-u", + peer, + "--scenario", + "gated-child", + "--gate-port", + port, + stdin=asyncio.subprocess.DEVNULL, + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + ) + client = OmpRpcClient( + ( + sys.executable, + "-u", + peer, + "--scenario", + "no-ready-descendants" if interruption == "startup" else "descendants", + "--gate-port", + port, + "--child-stdio", + stdio, + ) + ) + starting: asyncio.Task[None] | None = None + closing: asyncio.Task[None] | None = None + cancellations: list[asyncio.TimerHandle] = [] + try: + sibling = await process_gate.connection() + unrelated_process = await asyncio.to_thread(psutil.Process, unrelated.pid) + # Windows virtualenv launchers may remain between the parent and Python. + assert sibling.process == unrelated_process or unrelated_process in ( + await asyncio.to_thread(sibling.process.parents) + ) + starting = asyncio.create_task(client.start()) + try: + descendants = [ + await process_gate.connection(), + await process_gate.connection(), + ] + except TimeoutError: + if starting.done(): + await starting + process = client._proc + pytest.fail( + f"Peer did not connect its descendants: startup_done={starting.done()}, " + f"root_pid={process.pid if process else None}, " + f"returncode={process.returncode if process else None}" + ) + root = client._proc + assert root is not None + root_process = await asyncio.to_thread(psutil.Process, root.pid) + ancestors = [ + await asyncio.to_thread(process.process.parents) + for process in descendants + ] + assert all(root_process in parents for parents in ancestors) + assert ( + descendants[0].process in ancestors[1] + or descendants[1].process in ancestors[0] + ) + if interruption == "startup": + closing = starting + else: + async with asyncio.timeout(5): + await starting + closing = asyncio.create_task(client.close()) + await asyncio.sleep(0) + closing.cancel() + loop = asyncio.get_running_loop() + cancellations = [ + loop.call_later(delay, closing.cancel) for delay in (0.01, 0.05, 0.1) + ] + async with asyncio.timeout(10): + with pytest.raises(asyncio.CancelledError): + await closing + assert root.returncode is not None + for descendant in descendants: + assert not await asyncio.to_thread(process_is_alive, descendant.process) + assert await asyncio.to_thread(process_is_alive, sibling.process) + assert unrelated.returncode is None + await client.close() + finally: + for cancellation in cancellations: + cancellation.cancel() + tasks = {task for task in (starting, closing) if task is not None} + for task in tasks: + task.cancel() + with suppress(ProcessLookupError): + unrelated.kill() + async with asyncio.timeout(5): + await unrelated.wait() + if client._proc is not None and client._proc.returncode is None: + with suppress(ProcessLookupError): + client._proc.kill() + await asyncio.gather(*tasks, return_exceptions=True) diff --git a/tests/integrations/omprpc/test_prompt_results.py b/tests/integrations/omprpc/test_prompt_results.py new file mode 100644 index 000000000..787376864 --- /dev/null +++ b/tests/integrations/omprpc/test_prompt_results.py @@ -0,0 +1,142 @@ +"""Prompt narration and results preserve typed content and report failures.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.omprpc import OmpPromptError, OmpRpcError +from band.integrations.ompwire import OmpEvent, OmpState +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + checkpoint_frames, + prompt_task, + second_prompt, +) + + +class TestPromptNarrationAndResults: + async def test_prompt_delivers_typed_narration_and_ignores_other_events( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("basic") + observer = Observer() + async with asyncio.timeout(5): + await prompt_task(client, observer) + assert [event.type for event in observer.events] == [ + "agent_start", + "tool_execution_start", + "tool_execution_end", + "message_end", + "agent_end", + ] + tool_start, tool_end, message_end = observer.events[1:4] + assert (tool_start.toolCallId, tool_start.toolName, tool_start.args) == ( + "native-call", + "read", + {"path": "note.txt"}, + ) + assert tool_end.result is not None + assert tool_end.result.content[0].text == "note" + assert tool_end.result.details == {"lines": 1} + assert tool_end.isError is False + message = message_end.message + assert message is not None + assert message.role == "assistant" + assert isinstance(message.content, list) + assert [ + (content.type, content.text, content.thinking) + for content in message.content + ] == [ + ("thinking", None, "Considering the result"), + ("text", "complete", None), + ] + assert message.usage is not None + assert ( + message.usage.input, + message.usage.output, + message.usage.totalTokens, + ) == ( + 7, + 3, + 13, + ) + assert message.usage.cost is not None + assert message.usage.cost.total == 0.33 + state = OmpState.model_validate(await client.request("get_state")) + assert (state.sessionId, state.sessionFile, state.isSettled) == ( + "session-local", + "/local/session.jsonl", + True, + ) + assert state.model is not None + assert (state.model.provider, state.model.id) == ("test-provider", "test-model") + await second_prompt(client) + + async def test_rejected_prompt_does_not_wait_for_a_result_that_is_not_owed( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("rejected") + async with asyncio.timeout(5): + with pytest.raises(OmpRpcError, match="Prompt rejected") as rejected: + await prompt_task(client, Observer()) + assert (rejected.value.command, rejected.value.code) == ("prompt", "rejected") + await second_prompt(client) + + @pytest.mark.parametrize("scenario", ["initial-error", "initial-aborted"]) + async def test_failed_initial_result_is_reported( + self, client_factory: ClientFactory, scenario: str + ) -> None: + client = await client_factory(scenario) + async with asyncio.timeout(5): + with pytest.raises(OmpPromptError): + await prompt_task(client, Observer()) + + async def test_narration_failure_does_not_leave_the_turn_or_process_waiting( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("basic") + process = client._proc + assert process is not None + observer = Observer() + + async def on_event(event: OmpEvent) -> None: + raise ValueError("Narration failed") + + async with asyncio.timeout(6): + with pytest.raises(ValueError, match="Narration failed"): + await client.prompt( + "first", + on_event=on_event, + on_tool=observer.on_tool, + on_ui=observer.on_ui, + ) + assert process.returncode is not None + + async def test_native_v2_preserves_large_inline_image_frames( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("large-image") + observer = Observer() + async with asyncio.timeout(5): + await prompt_task(client, observer) + message = next( + event.message for event in observer.events if event.type == "message_end" + ) + assert message is not None and isinstance(message.content, list) + image = next(content for content in message.content if content.type == "image") + assert image.mimeType == "image/png" + assert image.data == "YWJj" * 300_000 + negotiations = [ + frame + for frame in await checkpoint_frames(client) + if frame["type"] == "negotiate_protocol" + ] + assert len(negotiations) == 1 + assert negotiations[0]["protocolVersion"] == 2 diff --git a/tests/integrations/omprpc/test_startup.py b/tests/integrations/omprpc/test_startup.py new file mode 100644 index 000000000..48d410111 --- /dev/null +++ b/tests/integrations/omprpc/test_startup.py @@ -0,0 +1,62 @@ +"""Startup requires native protocol negotiation and reaps failed launches.""" + +from __future__ import annotations + +import asyncio +import sys + +import pytest + +from band.integrations.omprpc import ( + OmpProcessError, + OmpProtocolError, + OmpRpcClient, + OmpRpcError, +) +from tests.integrations.omprpc.kit import OMP_PEER_SCRIPT, ClientFactory + + +class TestStartupNegotiation: + async def test_cancelled_startup_reaps_a_process_that_never_announces_ready( + self, + ) -> None: + client = OmpRpcClient( + ( + sys.executable, + "-u", + OMP_PEER_SCRIPT, + "--scenario", + "no-ready", + ) + ) + try: + with pytest.raises(TimeoutError): + async with asyncio.timeout(0.05): + await client.start() + assert client._proc is not None + assert client._proc.returncode is not None + finally: + await client.close() + + async def test_missing_executable_is_a_process_failure(self) -> None: + client = OmpRpcClient(("definitely-not-an-omp-executable-xyz",)) + try: + with pytest.raises(OmpProcessError, match="start"): + await client.start() + finally: + await client.close() + + @pytest.mark.parametrize( + ("scenario", "error"), + [ + ("no-v2", OmpProtocolError), + ("bad-negotiation", OmpProtocolError), + ("negotiation-rejected", OmpRpcError), + ], + ) + async def test_start_requires_successful_native_v2_negotiation( + self, client_factory: ClientFactory, scenario: str, error: type[RuntimeError] + ) -> None: + async with asyncio.timeout(6): + with pytest.raises(error, match="protocol"): + await client_factory(scenario) diff --git a/tests/integrations/omprpc/test_transport_failures.py b/tests/integrations/omprpc/test_transport_failures.py new file mode 100644 index 000000000..92e0dda3c --- /dev/null +++ b/tests/integrations/omprpc/test_transport_failures.py @@ -0,0 +1,82 @@ +"""Transport and protocol failures fail waiters and reap their subprocesses.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.omprpc import OmpProcessError, OmpProtocolError +from tests.integrations.omprpc.kit import ClientFactory, Observer, prompt_task + + +class TestTransportAndProtocolFailures: + @pytest.mark.parametrize( + ("scenario", "error", "match"), + [ + ("eof", OmpProcessError, "EOF"), + ("bad-json", OmpProtocolError, "protocol"), + ("bad-event", OmpProtocolError, "protocol"), + ("wrong-command", OmpProtocolError, "command"), + ("missing-state", OmpProtocolError, "protocol"), + ("bad-acceptance", OmpProtocolError, "protocol"), + ], + ) + async def test_eof_and_protocol_failure_cannot_leave_a_prompt_waiting( + self, + client_factory: ClientFactory, + scenario: str, + error: type[RuntimeError], + match: str, + ) -> None: + client = await client_factory(scenario) + process = client._proc + assert process is not None + async with asyncio.timeout(6): + with pytest.raises(error, match=match): + await prompt_task(client, Observer()) + assert process.returncode is not None + + async def test_eof_fails_a_command_waiter_and_closes_the_transport( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("basic") + process = client._proc + assert process is not None + async with asyncio.timeout(6): + with pytest.raises(OmpProcessError, match="EOF"): + await client.request("crash") + await client.close() + assert process.returncode is not None + + async def test_stderr_is_drained_without_newlines( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("basic") + async with asyncio.timeout(5): + assert await client.request("stderr") == {"drained": True} + + @pytest.mark.parametrize( + "scenario", + [ + "chunk-overflow", + "chunk-order", + "chunk-interrupted", + "chunk-invalid-data", + "chunk-length", + "chunk-eof", + "frame-overflow", + ], + ) + async def test_invalid_or_oversized_native_chunks_fail_and_reap( + self, client_factory: ClientFactory, scenario: str + ) -> None: + client = await client_factory(scenario) + process = client._proc + assert process is not None + async with asyncio.timeout(6): + with pytest.raises(OmpProtocolError, match="protocol"): + await prompt_task(client, Observer()) + assert process.returncode is not None diff --git a/tests/integrations/omprpc/test_turn_settlement.py b/tests/integrations/omprpc/test_turn_settlement.py new file mode 100644 index 000000000..af4dc67d9 --- /dev/null +++ b/tests/integrations/omprpc/test_turn_settlement.py @@ -0,0 +1,95 @@ +"""Turns await matching results and settlement across background work.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.omprpc import OmpPromptError +from band.integrations.ompwire import OmpState +from tests.integrations.omprpc.kit import ClientFactory, Observer, prompt_task + + +class TestTurnSettlement: + @pytest.mark.parametrize("scenario", ["yield", "stale", "false-result-settled"]) + async def test_initial_yield_never_finishes_a_background_turn( + self, client_factory: ClientFactory, scenario: str + ) -> None: + client = await client_factory(scenario) + observer = Observer() + running = prompt_task(client, observer) + try: + first_end = await observer.wait_for("agent_end") + assert first_end.awaitingAsyncWork is True + await client.request("checkpoint") + assert not running.done() + await client.request("release") + async with asyncio.timeout(5): + await running + messages = [ + content.text + for event in observer.events + if event.type == "message_end" + and event.message is not None + and isinstance(event.message.content, list) + for content in event.message.content + if content.type == "text" + ] + assert messages == ["yielded", "background complete"] + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + @pytest.mark.parametrize( + "scenario", ["resume-error", "resume-error-compacted", "resume-aborted"] + ) + async def test_resumed_failures_override_a_successful_initial_prompt_result( + self, client_factory: ClientFactory, scenario: str + ) -> None: + client = await client_factory(scenario) + observer = Observer() + running = prompt_task(client, observer) + try: + await observer.wait_for("agent_end") + await client.request("release") + expected = ( + "Resumed provider failed" if scenario != "resume-aborted" else "aborted" + ) + async with asyncio.timeout(5): + with pytest.raises(OmpPromptError, match=expected): + await running + assert sum(event.type == "agent_end" for event in observer.events) == 2 + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + async def test_settlement_before_matching_result_waits_for_both( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory("settlement-before-result") + observer = Observer() + running = prompt_task(client, observer) + try: + await observer.wait_for("agent_end") + await client.request("checkpoint") + assert not running.done() + await client.request("release") + async with asyncio.timeout(5): + await running + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + @pytest.mark.parametrize( + "scenario", ["state-settles", "settles-during-state", "local"] + ) + async def test_native_state_snapshot_and_subscribed_settlement_do_not_race( + self, client_factory: ClientFactory, scenario: str + ) -> None: + client = await client_factory(scenario) + observer = Observer() + async with asyncio.timeout(5): + await prompt_task(client, observer) + assert OmpState.model_validate(await client.request("get_state")).isSettled diff --git a/uv.lock b/uv.lock index 8684ecbef..12ff6aeb8 100644 --- a/uv.lock +++ b/uv.lock @@ -815,6 +815,7 @@ dev = [ { name = "opentelemetry-sdk" }, { name = "packaging" }, { name = "pre-commit" }, + { name = "psutil" }, { name = "pydantic-ai-slim", extra = ["anthropic", "openai"], marker = "extra == 'extra-8-band-sdk-dev' or (extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-dev-parlant') or (extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-dev-parlant') or (extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant') or (extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-dev-strands-nooa') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-nooa') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-strands') or (extra == 'extra-8-band-sdk-dev-strands-nooa' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-nooa' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-parlant' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-pydantic-ai' and extra == 'extra-8-band-sdk-strands')" }, { name = "pydantic-settings" }, { name = "pyrefly" }, @@ -955,6 +956,9 @@ nooa = [ { name = "nooa", marker = "(python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'win32') or (python_full_version < '3.12' and extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (python_full_version < '3.12' and extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant') or (python_full_version >= '3.14' and extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (python_full_version >= '3.14' and extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant') or (sys_platform == 'win32' and extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (sys_platform == 'win32' and extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant')" }, { name = "openai", version = "2.46.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'win32') or (python_full_version < '3.12' and extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (python_full_version < '3.12' and extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant') or (python_full_version >= '3.14' and extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (python_full_version >= '3.14' and extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant') or (sys_platform == 'win32' and extra == 'extra-8-band-sdk-crewai' and extra == 'extra-8-band-sdk-parlant') or (sys_platform == 'win32' and extra == 'extra-8-band-sdk-dev-crewai' and extra == 'extra-8-band-sdk-parlant')" }, ] +omp = [ + { name = "psutil" }, +] opencode = [ { name = "httpx" }, { name = "mcp", version = "1.28.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-8-band-sdk-crewai' or extra == 'extra-8-band-sdk-dev-crewai' or (extra == 'extra-8-band-sdk-dev' and extra == 'extra-8-band-sdk-dev-parlant') or (extra == 'extra-8-band-sdk-dev' and extra == 'extra-8-band-sdk-dev-strands-nooa') or (extra == 'extra-8-band-sdk-dev' and extra == 'extra-8-band-sdk-nooa') or (extra == 'extra-8-band-sdk-dev' and extra == 'extra-8-band-sdk-parlant') or (extra == 'extra-8-band-sdk-dev' and extra == 'extra-8-band-sdk-strands') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-dev-strands-nooa') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-nooa') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-dev-parlant' and extra == 'extra-8-band-sdk-strands') or (extra == 'extra-8-band-sdk-dev-strands-nooa' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-nooa' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-parlant' and extra == 'extra-8-band-sdk-pydantic-ai') or (extra == 'extra-8-band-sdk-pydantic-ai' and extra == 'extra-8-band-sdk-strands')" }, @@ -1104,6 +1108,8 @@ requires-dist = [ { name = "pillow", marker = "extra == 'crewai'", specifier = ">=12.1.1" }, { name = "pillow", marker = "extra == 'dev-crewai'", specifier = ">=12.1.1" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.0.0" }, + { name = "psutil", marker = "extra == 'dev'", specifier = ">=6" }, + { name = "psutil", marker = "extra == 'omp'", specifier = ">=6" }, { name = "pydantic", specifier = ">=2.0" }, { name = "pydantic-ai-slim", extras = ["anthropic", "openai"], marker = "extra == 'dev'", specifier = ">=2.55.0" }, { name = "pydantic-ai-slim", extras = ["anthropic", "openai"], marker = "extra == 'pydantic-ai'", specifier = ">=2.55.0" }, @@ -1191,7 +1197,7 @@ requires-dist = [ { name = "werkzeug", marker = "extra == 'dev-parlant'", specifier = ">=3.1.6" }, { name = "werkzeug", marker = "extra == 'parlant'", specifier = ">=3.1.6" }, ] -provides-extras = ["logging", "desktop", "codex", "opencode", "letta", "pydantic-ai", "anthropic", "langgraph", "claude-sdk", "copilot-sdk", "parlant", "crewai", "gemini", "a2a", "a2a-gateway", "a2a-gateway-demo", "acp", "slack", "bridge", "bridge-agentcore", "agentcore-runtime", "google-adk", "agno", "strands", "nooa", "dev", "dev-crewai", "dev-parlant", "dev-strands-nooa"] +provides-extras = ["logging", "desktop", "codex", "opencode", "letta", "pydantic-ai", "anthropic", "langgraph", "claude-sdk", "copilot-sdk", "parlant", "crewai", "gemini", "a2a", "a2a-gateway", "a2a-gateway-demo", "acp", "omp", "slack", "bridge", "bridge-agentcore", "agentcore-runtime", "google-adk", "agno", "strands", "nooa", "dev", "dev-crewai", "dev-parlant", "dev-strands-nooa"] [[package]] name = "band-sdk-core" @@ -6845,6 +6851,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c4/72/02445137af02769918a93807b2b7890047c32bfb9f90371cbc12688819eb/protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901", size = 170656, upload-time = "2026-03-18T19:04:59.826Z" }, ] +[[package]] +name = "psutil" +version = "7.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/08/510cbdb69c25a96f4ae523f733cdc963ae654904e8db864c07585ef99875/psutil-7.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b", size = 130595, upload-time = "2026-01-28T18:14:57.293Z" }, + { url = "https://files.pythonhosted.org/packages/d6/f5/97baea3fe7a5a9af7436301f85490905379b1c6f2dd51fe3ecf24b4c5fbf/psutil-7.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e78c8603dcd9a04c7364f1a3e670cea95d51ee865e4efb3556a3a63adef958ea", size = 131082, upload-time = "2026-01-28T18:14:59.732Z" }, + { url = "https://files.pythonhosted.org/packages/37/d6/246513fbf9fa174af531f28412297dd05241d97a75911ac8febefa1a53c6/psutil-7.2.2-cp313-cp313t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1a571f2330c966c62aeda00dd24620425d4b0cc86881c89861fbc04549e5dc63", size = 181476, upload-time = "2026-01-28T18:15:01.884Z" }, + { url = "https://files.pythonhosted.org/packages/b8/b5/9182c9af3836cca61696dabe4fd1304e17bc56cb62f17439e1154f225dd3/psutil-7.2.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:917e891983ca3c1887b4ef36447b1e0873e70c933afc831c6b6da078ba474312", size = 184062, upload-time = "2026-01-28T18:15:04.436Z" }, + { url = "https://files.pythonhosted.org/packages/16/ba/0756dca669f5a9300d0cbcbfae9a4c30e446dfc7440ffe43ded5724bfd93/psutil-7.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:ab486563df44c17f5173621c7b198955bd6b613fb87c71c161f827d3fb149a9b", size = 139893, upload-time = "2026-01-28T18:15:06.378Z" }, + { url = "https://files.pythonhosted.org/packages/1c/61/8fa0e26f33623b49949346de05ec1ddaad02ed8ba64af45f40a147dbfa97/psutil-7.2.2-cp313-cp313t-win_arm64.whl", hash = "sha256:ae0aefdd8796a7737eccea863f80f81e468a1e4cf14d926bd9b6f5f2d5f90ca9", size = 135589, upload-time = "2026-01-28T18:15:08.03Z" }, + { url = "https://files.pythonhosted.org/packages/81/69/ef179ab5ca24f32acc1dac0c247fd6a13b501fd5534dbae0e05a1c48b66d/psutil-7.2.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:eed63d3b4d62449571547b60578c5b2c4bcccc5387148db46e0c2313dad0ee00", size = 130664, upload-time = "2026-01-28T18:15:09.469Z" }, + { url = "https://files.pythonhosted.org/packages/7b/64/665248b557a236d3fa9efc378d60d95ef56dd0a490c2cd37dafc7660d4a9/psutil-7.2.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7b6d09433a10592ce39b13d7be5a54fbac1d1228ed29abc880fb23df7cb694c9", size = 131087, upload-time = "2026-01-28T18:15:11.724Z" }, + { url = "https://files.pythonhosted.org/packages/d5/2e/e6782744700d6759ebce3043dcfa661fb61e2fb752b91cdeae9af12c2178/psutil-7.2.2-cp314-cp314t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1fa4ecf83bcdf6e6c8f4449aff98eefb5d0604bf88cb883d7da3d8d2d909546a", size = 182383, upload-time = "2026-01-28T18:15:13.445Z" }, + { url = "https://files.pythonhosted.org/packages/57/49/0a41cefd10cb7505cdc04dab3eacf24c0c2cb158a998b8c7b1d27ee2c1f5/psutil-7.2.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e452c464a02e7dc7822a05d25db4cde564444a67e58539a00f929c51eddda0cf", size = 185210, upload-time = "2026-01-28T18:15:16.002Z" }, + { url = "https://files.pythonhosted.org/packages/dd/2c/ff9bfb544f283ba5f83ba725a3c5fec6d6b10b8f27ac1dc641c473dc390d/psutil-7.2.2-cp314-cp314t-win_amd64.whl", hash = "sha256:c7663d4e37f13e884d13994247449e9f8f574bc4655d509c3b95e9ec9e2b9dc1", size = 141228, upload-time = "2026-01-28T18:15:18.385Z" }, + { url = "https://files.pythonhosted.org/packages/f2/fc/f8d9c31db14fcec13748d373e668bc3bed94d9077dbc17fb0eebc073233c/psutil-7.2.2-cp314-cp314t-win_arm64.whl", hash = "sha256:11fe5a4f613759764e79c65cf11ebdf26e33d6dd34336f8a337aa2996d71c841", size = 136284, upload-time = "2026-01-28T18:15:19.912Z" }, + { url = "https://files.pythonhosted.org/packages/e7/36/5ee6e05c9bd427237b11b3937ad82bb8ad2752d72c6969314590dd0c2f6e/psutil-7.2.2-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486", size = 129090, upload-time = "2026-01-28T18:15:22.168Z" }, + { url = "https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979", size = 129859, upload-time = "2026-01-28T18:15:23.795Z" }, + { url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" }, + { url = "https://files.pythonhosted.org/packages/63/65/37648c0c158dc222aba51c089eb3bdfa238e621674dc42d48706e639204f/psutil-7.2.2-cp36-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e", size = 156997, upload-time = "2026-01-28T18:15:27.794Z" }, + { url = "https://files.pythonhosted.org/packages/8e/13/125093eadae863ce03c6ffdbae9929430d116a246ef69866dad94da3bfbc/psutil-7.2.2-cp36-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8", size = 148972, upload-time = "2026-01-28T18:15:29.342Z" }, + { url = "https://files.pythonhosted.org/packages/04/78/0acd37ca84ce3ddffaa92ef0f571e073faa6d8ff1f0559ab1272188ea2be/psutil-7.2.2-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc", size = 148266, upload-time = "2026-01-28T18:15:31.597Z" }, + { url = "https://files.pythonhosted.org/packages/b4/90/e2159492b5426be0c1fef7acba807a03511f97c5f86b3caeda6ad92351a7/psutil-7.2.2-cp37-abi3-win_amd64.whl", hash = "sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988", size = 137737, upload-time = "2026-01-28T18:15:33.849Z" }, + { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, +] + [[package]] name = "py-key-value-aio" version = "0.4.4" From 81c3a14a3f5bdc7129117b4950a58da3b5e75a14 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 07:13:33 +0300 Subject: [PATCH 03/35] fix: retain native process containment and prompt admission ownership --- src/band/integrations/codex/stdio_client.py | 59 +--- src/band/integrations/codex/winspawn.py | 49 ++++ src/band/integrations/omprpc.py | 266 ++++++++---------- tests/integrations/codex/conftest.py | 8 +- tests/integrations/codex/lifecycle.py | 3 +- tests/integrations/codex/lookup.py | 3 +- tests/integrations/codex/test_bootstrap.py | 5 +- tests/integrations/codex/test_cleanup.py | 12 +- tests/integrations/codex/test_lookup.py | 6 +- tests/integrations/codex/test_windows.py | 16 +- tests/integrations/omppeer.py | 6 +- tests/integrations/omprpc/conftest.py | 19 +- tests/integrations/omprpc/kit.py | 128 ++++++++- tests/integrations/omprpc/lifecycle.py | 252 +++++++++++++++++ .../integrations/omprpc/test_interruption.py | 34 +++ .../omprpc/test_lifecycle_model.py | 26 ++ .../integrations/omprpc/test_process_tree.py | 122 +++++++- 17 files changed, 765 insertions(+), 249 deletions(-) create mode 100644 tests/integrations/omprpc/lifecycle.py create mode 100644 tests/integrations/omprpc/test_lifecycle_model.py diff --git a/src/band/integrations/codex/stdio_client.py b/src/band/integrations/codex/stdio_client.py index c87961e5a..ec0636225 100644 --- a/src/band/integrations/codex/stdio_client.py +++ b/src/band/integrations/codex/stdio_client.py @@ -9,8 +9,6 @@ import shutil import signal from collections.abc import Awaitable, Callable, Mapping, Sequence -from contextlib import suppress -from pathlib import Path from typing import Any from band.core.exceptions import BandConnectionError @@ -21,8 +19,6 @@ WindowsProcess, await_returncode, close_stdin, - find_executable, - process_image, ) logger = logging.getLogger(__name__) @@ -31,31 +27,9 @@ # held stdio pipe cannot stall room cleanup. _EXIT_OBSERVE_S = 2.0 _STDOUT_LIMIT = 16 * 1024 * 1024 -HELPER_PATH = Path(__file__).resolve().with_name("winmain.py") -def _absolute_command(command: Sequence[str]) -> tuple[str, ...]: - """Resolve the executable in this process, not in the room directory. - - The helper's current directory is the room. ``CreateProcessW`` searches - that directory when the application name is relative, so a room-local - binary would run in place of the one this process would have started. - """ - head, *rest = command - if Path(head).is_absolute(): - return (head, *rest) - if os.name == "nt" and not os.path.dirname(head): - found = find_executable(head) - else: - found = shutil.which(head) - if found is None and Path(head).exists(): - found = str(Path(head).resolve()) - if found is None: - raise FileNotFoundError(head) - return (str(Path(found).resolve()), *rest) - - -def _session_pid_reused(process: asyncio.subprocess.Process) -> bool: +def session_pid_reused(process: asyncio.subprocess.Process) -> bool: """True when ``process.pid`` now names some other live process. After the leader is reaped, its session id stays reserved only while a @@ -70,12 +44,12 @@ def _session_pid_reused(process: asyncio.subprocess.Process) -> bool: return False except OSError: logger.warning( - "Codex process %s pid is not signalable; leaving its session", + "Process %s pid is not signalable; leaving its session", process.pid, ) return True logger.warning( - "Codex process %s pid was reused; not signaling that session", + "Process %s pid was reused; not signaling that session", process.pid, ) return True @@ -99,17 +73,18 @@ async def _wait_readers(readers: Sequence[asyncio.Task[None]]) -> None: async def kill_process_tree(process: asyncio.subprocess.Process) -> None: - """Force-stop a Unix process and everything in its session. + """Force-stop the owned Unix process group, including after root exit. The codex launcher is a shim whose real CLI is its child. On Unix the child leads a session, so ``killpg`` still reaches that group after the leader has exited. """ - if os.name != "nt": - if not _session_pid_reused(process): - with suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) + if os.name == "nt" or session_pid_reused(process): return + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass class CodexStdioClient(BaseJsonRpcClient): @@ -177,8 +152,8 @@ async def _spawn(self) -> asyncio.subprocess.Process: raise BandConnectionError("Codex process is already starting") self._windows_process = owner = WindowsProcess() try: - return await owner.start( - self._windows_command, + return await owner.start_command( + self.command, cwd=self.cwd, env=self.env, limit=_STDOUT_LIMIT, @@ -198,18 +173,6 @@ async def _spawn(self) -> asyncio.subprocess.Process: limit=_STDOUT_LIMIT, ) - def _windows_command(self, job_name: str, status_handle: int) -> list[str]: - return [ - str(process_image()), - "-I", - "-S", - str(HELPER_PATH), - job_name, - str(status_handle), - "--", - *_absolute_command(self.command), - ] - def _spawn_connection_error(self, failure: SpawnFailure) -> BandConnectionError: if failure.winerror in MISSING_EXECUTABLE: return BandConnectionError(_missing_cli_message(self.command[0])) diff --git a/src/band/integrations/codex/winspawn.py b/src/band/integrations/codex/winspawn.py index a7306ba9c..833ee8fe1 100644 --- a/src/band/integrations/codex/winspawn.py +++ b/src/band/integrations/codex/winspawn.py @@ -12,6 +12,7 @@ import functools import logging import os +import shutil import subprocess import sys import uuid @@ -25,6 +26,7 @@ # CreateProcess failed because the executable path does not exist. MISSING_EXECUTABLE = frozenset({2, 3}) +HELPER_PATH = Path(__file__).resolve().with_name("winmain.py") _READY = "ok" _STATUS_ERROR = "error" @@ -344,6 +346,37 @@ def find_executable(name: str) -> str | None: capacity = length + 1 +def absolute_command(command: Sequence[str]) -> tuple[str, ...]: + """Resolve the executable in this process, not in the command's directory.""" + head, *rest = command + if Path(head).is_absolute(): + return (head, *rest) + if os.name == "nt" and not os.path.dirname(head): + found = find_executable(head) + else: + found = shutil.which(head) + if found is None and Path(head).exists(): + found = str(Path(head).resolve()) + if found is None: + raise FileNotFoundError(head) + return (str(Path(found).resolve()), *rest) + + +def command_in_job( + command: Sequence[str], job_name: str, status_handle: int +) -> list[str]: + return [ + str(process_image()), + "-I", + "-S", + str(HELPER_PATH), + job_name, + str(status_handle), + "--", + *absolute_command(command), + ] + + def _winerror(exc: OSError) -> int: code = exc.winerror # type: ignore[attr-defined] if isinstance(code, int) and code: @@ -576,6 +609,22 @@ def __init__(self) -> None: self._status: asyncio.Task[str] | None = None self._read_fd: int | None = None + async def start_command( + self, + command: Sequence[str], + *, + cwd: str | None, + env: Mapping[str, str] | None, + limit: int, + ) -> asyncio.subprocess.Process: + """Start a command with job membership established before it can run.""" + return await self.start( + functools.partial(command_in_job, command), + cwd=cwd, + env=env, + limit=limit, + ) + async def start( self, command_for: Callable[[str, int], Sequence[str]], diff --git a/src/band/integrations/omprpc.py b/src/band/integrations/omprpc.py index 16d04ce69..fe57c64aa 100644 --- a/src/band/integrations/omprpc.py +++ b/src/band/integrations/omprpc.py @@ -14,6 +14,13 @@ import psutil from pydantic import JsonValue, TypeAdapter, ValidationError +from band.integrations.codex.stdio_client import kill_process_tree, session_pid_reused +from band.integrations.codex.winspawn import ( + SpawnFailure, + WindowsProcess, + await_returncode, + close_stdin, +) from band.integrations.ompwire import ( OMP_CHUNK_BYTES, OMP_FRAME_BYTES, @@ -157,8 +164,8 @@ def __init__( self.command = tuple(command) self.env = None if env is None else {**os.environ, **env} self._proc: asyncio.subprocess.Process | None = None - self._process_identity: psutil.Process | None = None - self._descendants: set[psutil.Process] = set() + self._windows_process: WindowsProcess | None = None + self._session_waiters: set[psutil.Process] = set() self._spawn_task: asyncio.Task[asyncio.subprocess.Process] | None = None self._reader_task: asyncio.Task[None] | None = None self._stderr_task: asyncio.Task[None] | None = None @@ -184,8 +191,6 @@ async def start(self) -> None: await self.close() self._proc = None self._spawn_task = None - self._process_identity = None - self._descendants.clear() self._close_task = None if self._proc is not None: self._ensure_running() @@ -198,7 +203,10 @@ async def start(self) -> None: self._ready = asyncio.get_running_loop().create_future() self._spawn_task = asyncio.create_task(self._spawn_process()) try: - self._proc = await asyncio.shield(self._spawn_task) + process = await asyncio.shield(self._spawn_task) + if self._closing: + raise OmpProcessError("OMP RPC client closed during startup") + self._proc = process self._reader_task = asyncio.create_task(self._read_stdout(self._proc)) self._stderr_task = asyncio.create_task(self._read_stderr(self._proc)) ready = await asyncio.shield(self._ready) @@ -217,7 +225,7 @@ async def start(self) -> None: raise OmpProtocolError( "OMP RPC protocol v2 negotiation failed" ) from cause - except OSError as cause: + except (OSError, SpawnFailure) as cause: await self.close() raise OmpProcessError("Unable to start OMP RPC process") from cause except (Exception, asyncio.CancelledError): @@ -225,31 +233,28 @@ async def start(self) -> None: raise async def _spawn_process(self) -> asyncio.subprocess.Process: - process = await asyncio.create_subprocess_exec( - *self.command, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=None, - env=self.env, - limit=OMP_FRAME_BYTES, - ) + if os.name == "nt": + self._windows_process = WindowsProcess() + process = await self._windows_process.start_command( + self.command, + cwd=None, + env=self.env, + limit=OMP_FRAME_BYTES, + ) + else: + process = await asyncio.create_subprocess_exec( + *self.command, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=None, + env=self.env, + start_new_session=True, + limit=OMP_FRAME_BYTES, + ) self._proc = process - self._process_identity = await asyncio.to_thread( - self._capture_identity, process.pid - ) return process - @staticmethod - def _capture_identity(pid: int) -> psutil.Process | None: - try: - process = psutil.Process(pid) - # Process() can defer AccessDenied when reading the creation time. - process.create_time() - return process - except psutil.NoSuchProcess: - return None - async def request( self, command: str, **parameters: JsonValue ) -> dict[str, JsonValue]: @@ -267,13 +272,9 @@ async def prompt( on_ui: Callable[[OmpUIRequest], Awaitable[OmpUIResponse]], ) -> None: """Deliver a turn through its correlated result and actual native settlement.""" - try: - async with self._prompt_lock: - self._ensure_running() - await self._execute_prompt(message, on_event, on_tool, on_ui) - except asyncio.CancelledError: - await self.close() - raise + async with self._prompt_lock: + self._ensure_running() + await self._execute_prompt(message, on_event, on_tool, on_ui) async def close(self) -> None: """Finish shared bounded cleanup before propagating caller cancellation.""" @@ -751,6 +752,8 @@ def _record_failure(self, error: OmpProcessError) -> None: def _begin_close(self) -> None: self._closing = True + if self._windows_process is not None: + self._windows_process.closing = True if self._close_task is None or ( self._close_task.done() and ( @@ -800,100 +803,54 @@ async def _wait_for_exit(process: asyncio.subprocess.Process) -> bool: return False return True - @staticmethod - def _process_is_alive(process: psutil.Process) -> bool: - try: - return process.is_running() and process.status() != psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return False - - def _capture_descendants(self) -> None: - root = self._process_identity - if root is None or not self._process_is_alive(root): - return - try: - descendants = root.children(recursive=True) - except psutil.NoSuchProcess: + def _capture_session_waiters(self, process: asyncio.subprocess.Process) -> None: + if session_pid_reused(process): return - if not root.is_running(): - return - for process in descendants: + members: list[psutil.Process] = [] + for pid in psutil.pids(): + # asyncio's child watcher remains the sole waiter for the root. + if pid == process.pid: + continue try: - process.create_time() - except psutil.NoSuchProcess: + if os.getpgid(pid) == process.pid: + member = psutil.Process(pid) + member.create_time() + members.append(member) + except (ProcessLookupError, psutil.NoSuchProcess): continue - self._descendants.add(process) + if not session_pid_reused(process): + self._session_waiters.update(members) - def _signal_root(self, *, kill: bool) -> None: - root = self._process_identity - if root is None: - if self._proc is not None and self._proc.returncode is None: - raise OmpProcessError("Cannot signal OMP without a captured identity") - return - try: - if kill: - root.kill() - else: - root.terminate() - except psutil.NoSuchProcess: - pass - - @classmethod - def _close_descendants(cls, processes: Sequence[psutil.Process]) -> None: - alive = list(processes) - errors: list[Exception] = [] - for kill in (False, True): - for process in alive: - try: - if kill: - process.kill() - else: - process.terminate() - except psutil.NoSuchProcess: - pass - except psutil.Error as cause: - errors.append(cause) + def _wait_for_session_exit(self) -> None: + _, alive = psutil.wait_procs( + self._session_waiters, timeout=_PROCESS_EXIT_TIMEOUT_S + ) + remaining: list[int] = [] + for member in alive: try: - _, alive = psutil.wait_procs(alive, timeout=_PROCESS_EXIT_TIMEOUT_S) - except psutil.Error as cause: - errors.append(cause) - remaining: list[psutil.Process] = [] - for process in alive: - try: - if cls._process_is_alive(process): - remaining.append(process) - except psutil.Error as cause: - errors.append(cause) - remaining.append(process) - alive = remaining - if not alive: - break - if alive: - pids = sorted(process.pid for process in alive) - errors.append(OmpProcessError(f"OMP descendants survived kill: {pids}")) - if len(errors) == 1: - raise errors[0] - if errors: - raise ExceptionGroup("OMP descendant cleanup failed", errors) + if member.is_running() and member.status() != psutil.STATUS_ZOMBIE: + remaining.append(member.pid) + except psutil.NoSuchProcess: + pass + if remaining: + raise OmpProcessError( + f"OMP session members survived containment kill: {sorted(remaining)}" + ) - async def _close_process(self) -> None: - process = self._proc - errors: list[Exception] = [] - if self._spawn_task is not None: - try: - process = await asyncio.shield(self._spawn_task) - except OSError: - return - except psutil.Error as cause: - errors.append(cause) - process = self._proc - if process is None: - return - # Abort and stdin EOF can orphan children before their parent is reaped. + async def _settle_spawn(self) -> asyncio.subprocess.Process | None: + spawning = self._spawn_task + if spawning is None: + return self._proc try: - await asyncio.to_thread(self._capture_descendants) - except psutil.Error as cause: - errors.append(cause) + return await asyncio.shield(spawning) + except (OSError, SpawnFailure): + return self._proc + except asyncio.CancelledError: + if not spawning.cancelled(): + raise + return self._proc + + async def _close_process(self) -> None: turn = self._turn self._fail_pending(self._failure or OmpProcessError("OMP RPC client closed")) if turn is not None: @@ -901,6 +858,20 @@ async def _close_process(self) -> None: if turn.event_worker is not None: workers.append(turn.event_worker) await self._cancel_tasks(workers) + + owner = self._windows_process + owner_closed = False + if owner is not None and self._proc is None: + # Stop incomplete helper creation/readiness before awaiting its spawn task. + try: + await owner.close(_PROCESS_EXIT_TIMEOUT_S) + owner_closed = True + finally: + self._proc = owner.process + process = await self._settle_spawn() + if process is None: + self._windows_process = None + return if ( self._failure is None and process.returncode is None @@ -922,37 +893,24 @@ async def _close_process(self) -> None: ) and process.stderr is not None: drains.append(asyncio.create_task(self._drain(process.stderr))) try: - await asyncio.to_thread(self._capture_descendants) - except psutil.Error as cause: - errors.append(cause) - if process.stdin is not None: - process.stdin.close() - descendants = asyncio.create_task( - asyncio.to_thread(self._close_descendants, tuple(self._descendants)) - ) - try: + if owner is None: + # EOF and signals can remove PGID visibility before exit is observable. + await asyncio.to_thread(self._capture_session_waiters, process) + if not owner_closed: + await close_stdin(process) + await await_returncode(process, _PROCESS_EXIT_TIMEOUT_S) + if owner is not None: + if not owner_closed: + await owner.close(_PROCESS_EXIT_TIMEOUT_S) + else: + await asyncio.to_thread(self._capture_session_waiters, process) + await kill_process_tree(process) + await asyncio.to_thread(self._wait_for_session_exit) if not await self._wait_for_exit(process): - for kill in (False, True): - try: - await asyncio.to_thread(self._signal_root, kill=kill) - except (psutil.Error, OmpProcessError) as cause: - errors.append(cause) - if await self._wait_for_exit(process): - break - else: - errors.append( - OmpProcessError("OMP process did not exit after kill") - ) + raise OmpProcessError("OMP process did not exit after containment kill") + self._windows_process = None + self._session_waiters.clear() finally: - try: - await descendants - except (psutil.Error, OmpProcessError, ExceptionGroup) as cause: - errors.append(cause) - finally: - if self._stderr_task is not None: - drains.append(self._stderr_task) - await self._cancel_tasks(drains) - if len(errors) == 1: - raise errors[0] - if errors: - raise ExceptionGroup("OMP process cleanup failed", errors) + if self._stderr_task is not None: + drains.append(self._stderr_task) + await self._cancel_tasks(drains) diff --git a/tests/integrations/codex/conftest.py b/tests/integrations/codex/conftest.py index 5d3fc862c..c3a6d2f99 100644 --- a/tests/integrations/codex/conftest.py +++ b/tests/integrations/codex/conftest.py @@ -21,8 +21,9 @@ def windows_lookup(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> WindowsLo peer = WindowsLookup(tmp_path) monkeypatch.chdir(peer.current) monkeypatch.setenv("PATH", "") - monkeypatch.setattr(stdio_client, "os", SimpleNamespace(name="nt", path=os.path)) - monkeypatch.setattr(winspawn, "os", SimpleNamespace(name="nt", pathsep=os.pathsep)) + monkeypatch.setattr( + winspawn, "os", SimpleNamespace(name="nt", path=os.path, pathsep=os.pathsep) + ) monkeypatch.setattr(winspawn, "_kernel32", lambda: peer.kernel) return peer @@ -40,9 +41,6 @@ async def windows_lifecycle( ) monkeypatch.setattr(asyncio, "create_subprocess_exec", peer.launch) monkeypatch.setattr(winspawn, "_read_status", peer.read_status) - monkeypatch.setattr( - peer.client, "_windows_command", lambda _job, _handle: ["unused"] - ) try: yield peer finally: diff --git a/tests/integrations/codex/lifecycle.py b/tests/integrations/codex/lifecycle.py index c6f2bb2ef..acd0fad2f 100644 --- a/tests/integrations/codex/lifecycle.py +++ b/tests/integrations/codex/lifecycle.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import sys from collections.abc import Mapping from typing import cast @@ -65,7 +66,7 @@ class WindowsLifecycle: def __init__(self) -> None: self.helper = HelperPeer() self.job = JobPeer(self.helper) - self.client = CodexStdioClient(command=["unused"]) + self.client = CodexStdioClient(command=[sys.executable]) self.creation_started = asyncio.Event() self.creation_allowed = asyncio.Event() self.creation_allowed.set() diff --git a/tests/integrations/codex/lookup.py b/tests/integrations/codex/lookup.py index 1aebc9ac1..b2820f664 100644 --- a/tests/integrations/codex/lookup.py +++ b/tests/integrations/codex/lookup.py @@ -9,6 +9,7 @@ from types import SimpleNamespace from band.integrations.codex.stdio_client import CodexStdioClient +from band.integrations.codex.winspawn import command_in_job async def launch_directly(client: CodexStdioClient) -> None: @@ -30,7 +31,7 @@ async def launch_directly(client: CodexStdioClient) -> None: def selected_executable(name: str) -> Path: - command = CodexStdioClient(command=[name])._windows_command("test-job", 1) + command = command_in_job([name], "test-job", 1) return Path(command[command.index("--") + 1]) diff --git a/tests/integrations/codex/test_bootstrap.py b/tests/integrations/codex/test_bootstrap.py index 41e47c799..0ad48891e 100644 --- a/tests/integrations/codex/test_bootstrap.py +++ b/tests/integrations/codex/test_bootstrap.py @@ -10,6 +10,7 @@ import pytest from band.integrations.codex.stdio_client import CodexStdioClient +from band.integrations.codex.winspawn import command_in_job class TestWindowsBootstrap: @@ -38,7 +39,7 @@ def test_room_pythonpath_cannot_supply_helper_code( cwd=str(workdir), env={"PYTHONPATH": pythonpath} if configured else None, ) - command = client._windows_command("job", 1) + command = command_in_job(client.command, "job", 1) completed = subprocess.run( command[:4], cwd=workdir, @@ -73,7 +74,7 @@ def test_child_path_does_not_change_parent_executable_selection( client = CodexStdioClient( command=[executable.name, "-c", "pass"], env={"PATH": str(child_path)} ) - command = client._windows_command("job", 1) + command = command_in_job(client.command, "job", 1) selected = command[command.index("--") + 1] assert Path(selected).samefile(executable) assert client.env is not None and client.env["PATH"] == str(child_path) diff --git a/tests/integrations/codex/test_cleanup.py b/tests/integrations/codex/test_cleanup.py index b382d7644..1b2f4e944 100644 --- a/tests/integrations/codex/test_cleanup.py +++ b/tests/integrations/codex/test_cleanup.py @@ -19,13 +19,15 @@ import pytest_asyncio from band.integrations.codex import CodexStdioClient, stdio_client -from band.integrations.codex.stdio_client import HELPER_PATH, kill_process_tree +from band.integrations.codex.stdio_client import kill_process_tree from band.integrations.codex.winspawn import ( + HELPER_PATH, Job, JobNotDrained, SpawnFailure, _read_status, _start_and_wait, + command_in_job, ) from tests.integrations.codex.cleanup import cleanup_snapshot from tests.paths import ROOM_PEER @@ -343,7 +345,7 @@ def test_windows_helper_does_not_import_a_room_package( command=[sys.executable, "-c", "pass"], cwd=str(workdir), ) - argv = client._windows_command("job", 1) + argv = command_in_job(client.command, "job", 1) completed = subprocess.run( argv[:4], cwd=workdir, @@ -379,11 +381,9 @@ def test_windows_command_does_not_resolve_a_binary_in_the_room( cwd=str(workdir), ) with pytest.raises(FileNotFoundError): - client._windows_command("job", 1) + command_in_job(client.command, "job", 1) named = Path(sys.executable).name - resolved = CodexStdioClient(command=[named, "-c", "pass"])._windows_command( - "job", 1 - ) + resolved = command_in_job([named, "-c", "pass"], "job", 1) assert Path(resolved[resolved.index("--") + 1]).is_absolute() @pytest.mark.skipif(os.name == "nt", reason="session-group teardown is Unix") diff --git a/tests/integrations/codex/test_lookup.py b/tests/integrations/codex/test_lookup.py index 932560ad4..963a36156 100644 --- a/tests/integrations/codex/test_lookup.py +++ b/tests/integrations/codex/test_lookup.py @@ -7,7 +7,7 @@ import pytest -from band.integrations.codex.stdio_client import CodexStdioClient +from band.integrations.codex.winspawn import command_in_job from tests.integrations.codex.lookup import WindowsLookup, selected_executable @@ -16,9 +16,7 @@ def test_helper_uses_the_running_image_instead_of_a_venv_launcher( self, windows_lookup: WindowsLookup ) -> None: peer = windows_lookup - command = CodexStdioClient(command=[str(peer.image)])._windows_command( - "test-job", 1 - ) + command = command_in_job([str(peer.image)], "test-job", 1) assert Path(command[0]).samefile(peer.image) @pytest.mark.parametrize( diff --git a/tests/integrations/codex/test_windows.py b/tests/integrations/codex/test_windows.py index 278f651ce..4a8f8cde5 100644 --- a/tests/integrations/codex/test_windows.py +++ b/tests/integrations/codex/test_windows.py @@ -4,13 +4,15 @@ import os import sys +from collections.abc import Sequence from pathlib import Path import pytest from band.core.exceptions import BandConnectionError -from band.integrations.codex.stdio_client import HELPER_PATH, CodexStdioClient -from band.integrations.codex.winspawn import WindowsProcess +from band.integrations.codex import winspawn +from band.integrations.codex.stdio_client import CodexStdioClient +from band.integrations.codex.winspawn import HELPER_PATH, WindowsProcess, command_in_job @pytest.mark.skipif(os.name != "nt", reason="requires native Windows job objects") @@ -57,12 +59,14 @@ async def test_rejected_spawn_releases_helper_and_tree_directory_handles( command=[sys.executable, "-c", "import sys; sys.stdin.buffer.read()"], cwd=str(workdir), ) - command_for = client._windows_command + command_for = command_in_job report = tmp_path / "helper-pid" owners: list[WindowsProcess] = [] - def rejected_command(job_name: str, status_handle: int) -> list[str]: - command = command_for(job_name, status_handle) + def rejected_command( + child_command: Sequence[str], job_name: str, status_handle: int + ) -> list[str]: + command = command_for(child_command, job_name, status_handle) assert client._windows_process is not None owners.append(client._windows_process) return [ @@ -74,7 +78,7 @@ def rejected_command(job_name: str, status_handle: int) -> list[str]: *command[4:], ] - monkeypatch.setattr(client, "_windows_command", rejected_command) + monkeypatch.setattr(winspawn, "command_in_job", rejected_command) try: with pytest.raises(BandConnectionError): await client.connect() diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index 4502c3614..e74c82945 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -229,7 +229,9 @@ def prompt(self, frame: Frame) -> None: self.prompt_id = frame.id self.prompt_count += 1 self.settled = False - if self.prompt_count > 1 or frame.message == "second": + if ( + self.prompt_count > 1 and self.scenario != "lifecycle" + ) or frame.message == "second": self.finish(frame) return match self.scenario: @@ -563,7 +565,7 @@ def main() -> None: raise ValueError("A process gate port is required") hold_process_gate(args.gate_port, spawn_child=args.scenario == "gated-parent") return - if args.scenario in {"descendants", "no-ready-descendants"}: + if args.scenario in {"descendants", "no-ready-descendants", "lifecycle"}: if args.gate_port is None: raise ValueError("A process gate port is required") # The gate, not the peer's stdio EOF, owns these children's lifetime. diff --git a/tests/integrations/omprpc/conftest.py b/tests/integrations/omprpc/conftest.py index c6b3f8724..a24840303 100644 --- a/tests/integrations/omprpc/conftest.py +++ b/tests/integrations/omprpc/conftest.py @@ -2,7 +2,6 @@ from __future__ import annotations -import asyncio import sys from collections.abc import AsyncIterator @@ -16,7 +15,7 @@ OMP_PEER_SCRIPT, ClientFactory, ProcessGate, - kill_gate_processes, + process_gate_session, ) @@ -47,19 +46,5 @@ async def build(scenario: str) -> OmpRpcClient: @pytest_asyncio.fixture(loop_scope="function") async def process_gate() -> AsyncIterator[ProcessGate]: - gate = ProcessGate() - server = await asyncio.start_server(gate.accept, "127.0.0.1", 0) - try: - sockets = server.sockets - assert sockets - gate.port = sockets[0].getsockname()[1] + async with process_gate_session() as gate: yield gate - finally: - server.close() - for writer in gate.writers: - writer.close() - await asyncio.gather( - *(writer.wait_closed() for writer in gate.writers), return_exceptions=True - ) - await server.wait_closed() - await asyncio.to_thread(kill_gate_processes, gate.processes) diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py index 66253d247..958c24448 100644 --- a/tests/integrations/omprpc/kit.py +++ b/tests/integrations/omprpc/kit.py @@ -3,9 +3,12 @@ from __future__ import annotations import asyncio +import os +import signal +import sys import time -from collections.abc import Awaitable, Callable -from contextlib import suppress +from collections.abc import AsyncIterator, Awaitable, Callable, Sequence +from contextlib import asynccontextmanager, suppress from dataclasses import dataclass from typing import cast @@ -35,6 +38,18 @@ def __init__(self) -> None: self.events: list[OmpEvent] = [] self.queue: asyncio.Queue[OmpEvent] = asyncio.Queue() + @property + def assistant_texts(self) -> list[str]: + return [ + part.text + for event in self.events + if event.message is not None + and event.message.role == "assistant" + and isinstance(event.message.content, list) + for part in event.message.content + if part.text is not None + ] + async def on_event(self, event: OmpEvent) -> None: self.events.append(event) self.queue.put_nowait(event) @@ -79,11 +94,62 @@ async def second_prompt(client: OmpRpcClient) -> Observer: return observer +async def announced_prompt( + client: OmpRpcClient, + observer: Observer, + entered: asyncio.Event, + message: str, +) -> None: + entered.set() + await client.prompt( + message, + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=observer.on_ui, + ) + + +async def queue_prompt( + client: OmpRpcClient, message: str = "queued" +) -> asyncio.Task[None]: + entered = asyncio.Event() + task = asyncio.create_task(announced_prompt(client, Observer(), entered, message)) + # The coroutine reaches the occupied admission lock before this waiter resumes. + await entered.wait() + return task + + +def descendant_client( + gate: ProcessGate, *, scenario: str = "descendants", stdio: str = "discard" +) -> OmpRpcClient: + return OmpRpcClient( + ( + sys.executable, + "-u", + OMP_PEER_SCRIPT, + "--scenario", + scenario, + "--gate-port", + str(gate.port), + "--child-stdio", + stdio, + ) + ) + + async def checkpoint_frames(client: OmpRpcClient) -> list[dict[str, JsonValue]]: checkpoint = await client.request("checkpoint") return cast(list[dict[str, JsonValue]], checkpoint["frames"]) +async def sent_prompt_messages(client: OmpRpcClient) -> list[str]: + return [ + cast(str, frame["message"]) + for frame in await checkpoint_frames(client) + if frame["type"] == "prompt" + ] + + def process_is_alive(process: psutil.Process) -> bool: try: return process.is_running() and process.status() != psutil.STATUS_ZOMBIE @@ -108,6 +174,13 @@ class GatedProcess: process: psutil.Process +async def descendant_liveness(descendants: Sequence[GatedProcess]) -> list[bool]: + return [ + await asyncio.to_thread(process_is_alive, child.process) + for child in descendants + ] + + class ProcessGate: def __init__(self) -> None: self.port = 0 @@ -131,3 +204,54 @@ async def connection(self) -> GatedProcess: self.processes.append(process) assert await asyncio.to_thread(process.create_time) == receipt.created return GatedProcess(receipt, process) + + +@asynccontextmanager +async def process_gate_session() -> AsyncIterator[ProcessGate]: + gate = ProcessGate() + server = await asyncio.start_server(gate.accept, "127.0.0.1", 0) + try: + sockets = server.sockets + assert sockets + gate.port = sockets[0].getsockname()[1] + yield gate + finally: + server.close() + for writer in gate.writers: + writer.close() + await asyncio.gather( + *(writer.wait_closed() for writer in gate.writers), return_exceptions=True + ) + await server.wait_closed() + await asyncio.to_thread(kill_gate_processes, gate.processes) + + +def refuse_process_signal(*args: object) -> None: + raise PermissionError("Process ownership is temporarily not signalable") + + +class DeferredSessionSignal: + """Keep real members alive while their post-signal group lookup disappears.""" + + def __init__( + self, + pgid: int, + members: Sequence[GatedProcess], + *, + hide_after_signal: bool, + ) -> None: + self.pgid = pgid + self.members = {member.process.pid for member in members} + self.hide_after_signal = hide_after_signal + self.signalled = False + self.original_group = os.getpgid + + def send(self, pgid: int, sig: int) -> None: + assert pgid == self.pgid + assert sig == signal.SIGKILL + self.signalled = True + + def group(self, pid: int) -> int: + if self.signalled and self.hide_after_signal and pid in self.members: + raise ProcessLookupError(pid) + return self.original_group(pid) diff --git a/tests/integrations/omprpc/lifecycle.py b/tests/integrations/omprpc/lifecycle.py new file mode 100644 index 000000000..e3d8093a4 --- /dev/null +++ b/tests/integrations/omprpc/lifecycle.py @@ -0,0 +1,252 @@ +"""Generated host operations against a real RPC peer and owned process tree.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Awaitable, Callable +from contextlib import asynccontextmanager, suppress +from enum import StrEnum +from unittest.mock import patch + +from hypothesis import strategies as st + +from band.integrations.omprpc import OmpProcessError +from tests.integrations.omprpc.kit import ( + GatedProcess, + Observer, + ProcessGate, + descendant_client, + descendant_liveness, + process_gate_session, + prompt_task, + queue_prompt, + second_prompt, +) + +STEP_GUARD_S = 15 + + +class Step(StrEnum): + PROMPT = "admit a held prompt" + CANCEL_QUEUED = "cancel a prompt waiting for admission" + SETTLE = "settle the admitted prompt" + CANCEL_OWNER = "cancel the admitted prompt" + CLOSE = "close the transport" + CRASH = "crash the root before closing" + CANCEL_START_RESTART = "cancel held spawn while a restart waits" + CANCEL_CLOSE_RESTART = "cancel close during held spawn while a restart waits" + + +sequences = st.lists(st.sampled_from(Step), min_size=1, max_size=6) + + +class SpawnBarrier: + """Hold publication of the first real subprocess, not its creation.""" + + def __init__( + self, spawn: Callable[..., Awaitable[asyncio.subprocess.Process]] + ) -> None: + self.spawn = spawn + self.created = asyncio.Event() + self.release = asyncio.Event() + self.process: asyncio.subprocess.Process | None = None + + async def create( + self, *command: str, **kwargs: object + ) -> asyncio.subprocess.Process: + process = await self.spawn(*command, **kwargs) + if self.process is None: + self.process = process + self.created.set() + await self.release.wait() + return process + + +async def announced_operation( + operation: Awaitable[None], entered: asyncio.Event +) -> None: + entered.set() + await operation + + +class RpcLifecycle: + def __init__(self, gate: ProcessGate) -> None: + self.gate = gate + self.client = descendant_client(gate, scenario="lifecycle") + self.live = False + self.active: asyncio.Task[None] | None = None + self.observer: Observer | None = None + self.tasks: list[asyncio.Task[None]] = [] + self.roots: list[asyncio.subprocess.Process] = [] + self.children: list[GatedProcess] = [] + + async def record_children(self) -> list[GatedProcess]: + children = [await self.gate.connection(), await self.gate.connection()] + self.children.extend(children) + return children + + async def open(self) -> None: + if self.live: + return + await self.client.start() + assert self.client._proc is not None + self.roots.append(self.client._proc) + await self.record_children() + self.live = True + + async def hold_prompt(self) -> None: + await self.open() + if self.active is not None: + return + self.observer = Observer() + self.active = prompt_task(self.client, self.observer) + self.tasks.append(self.active) + await self.observer.wait_for("agent_end") + assert not self.active.done() + assert (await self.client.request("get_state"))["isSettled"] is False + + async def assert_dead( + self, roots: list[asyncio.subprocess.Process], children: list[GatedProcess] + ) -> None: + assert all(root.returncode is not None for root in roots) + assert not any(await descendant_liveness(children)) + + async def close(self) -> None: + await self.client.close() + if self.active is not None: + [result] = await asyncio.gather(self.active, return_exceptions=True) + assert result is None or isinstance(result, OmpProcessError) + self.active = None + self.live = False + await self.assert_dead(self.roots, self.children) + + async def overlap_spawn_and_restart(self, *, cancel_start: bool) -> None: + await self.close() + barrier = SpawnBarrier(asyncio.create_subprocess_exec) + closing: asyncio.Task[None] | None = None + with patch.object(asyncio, "create_subprocess_exec", barrier.create): + starting = asyncio.create_task(self.client.start()) + self.tasks.append(starting) + try: + await barrier.created.wait() + assert barrier.process is not None + old_root = barrier.process + self.roots.append(old_root) + old_children = await self.record_children() + if cancel_start: + starting.cancel() + else: + close_entered = asyncio.Event() + closing = asyncio.create_task( + announced_operation(self.client.close(), close_entered) + ) + self.tasks.append(closing) + await close_entered.wait() + closing.cancel() + restart_entered = asyncio.Event() + restarting = asyncio.create_task( + announced_operation(self.client.start(), restart_entered) + ) + self.tasks.append(restarting) + await restart_entered.wait() + assert not restarting.done() + barrier.release.set() + [result] = await asyncio.gather(starting, return_exceptions=True) + assert isinstance( + result, + asyncio.CancelledError if cancel_start else OmpProcessError, + ) + if closing is not None: + [result] = await asyncio.gather(closing, return_exceptions=True) + assert isinstance(result, asyncio.CancelledError) + await restarting + assert self.client._proc is not None + assert self.client._proc is not old_root + self.roots.append(self.client._proc) + await self.record_children() + self.live = True + await self.assert_dead([old_root], old_children) + await second_prompt(self.client) + finally: + barrier.release.set() + + async def take(self, step: Step) -> None: + match step: + case Step.PROMPT: + await self.hold_prompt() + case Step.CANCEL_QUEUED: + await self.hold_prompt() + queued = await queue_prompt(self.client) + self.tasks.append(queued) + queued.cancel() + [result] = await asyncio.gather(queued, return_exceptions=True) + assert isinstance(result, asyncio.CancelledError) + assert self.active is not None and not self.active.done() + assert (await self.client.request("get_state"))["isSettled"] is False + case Step.SETTLE: + await self.hold_prompt() + await self.client.request("release") + assert self.active is not None and self.observer is not None + await self.active + assert "background complete" in self.observer.assistant_texts + self.active = None + assert (await self.client.request("get_state"))["isSettled"] is True + case Step.CANCEL_OWNER: + await self.hold_prompt() + assert self.active is not None + self.active.cancel() + [result] = await asyncio.gather(self.active, return_exceptions=True) + assert isinstance(result, asyncio.CancelledError) + self.active = None + await self.close() + case Step.CLOSE: + await self.close() + case Step.CRASH: + await self.open() + root = self.client._proc + assert root is not None + with suppress(OmpProcessError): + await self.client.request("crash") + await root.wait() + assert root.returncode == 3 + await self.close() + case Step.CANCEL_START_RESTART | Step.CANCEL_CLOSE_RESTART: + await self.overlap_spawn_and_restart( + cancel_start=step == Step.CANCEL_START_RESTART + ) + if self.live: + assert (await self.client.request("checkpoint"))["pid"] is not None + + async def finish(self) -> None: + if self.active is not None: + await self.take(Step.SETTLE) + await self.close() + assert all(task.done() for task in self.tasks) + await self.open() + await second_prompt(self.client) + await self.close() + + async def cleanup(self) -> None: + for task in self.tasks: + task.cancel() + await asyncio.gather(*self.tasks, return_exceptions=True) + await self.client.close() + + +@asynccontextmanager +async def rpc_lifecycle() -> AsyncIterator[RpcLifecycle]: + async with process_gate_session() as gate: + session = RpcLifecycle(gate) + try: + yield session + finally: + await session.cleanup() + + +async def play(steps: list[Step]) -> None: + async with rpc_lifecycle() as session: + for step in steps: + async with asyncio.timeout(STEP_GUARD_S): + await session.take(step) + async with asyncio.timeout(STEP_GUARD_S): + await session.finish() diff --git a/tests/integrations/omprpc/test_interruption.py b/tests/integrations/omprpc/test_interruption.py index 74bd9244a..302c05236 100644 --- a/tests/integrations/omprpc/test_interruption.py +++ b/tests/integrations/omprpc/test_interruption.py @@ -11,11 +11,45 @@ ClientFactory, Observer, prompt_task, + queue_prompt, second_prompt, + sent_prompt_messages, ) class TestInterruptionAndReaping: + async def test_cancelling_a_queued_prompt_preserves_the_active_and_next_prompt( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("yield") + process = client._proc + assert process is not None + observer = Observer() + active = prompt_task(client, observer) + queued: asyncio.Task[None] | None = None + try: + await observer.wait_for("agent_end") + queued = await queue_prompt(client) + queued.cancel() + with pytest.raises(asyncio.CancelledError): + await queued + + assert not active.done() + assert process.returncode is None + assert await sent_prompt_messages(client) == ["first"] + await client.request("release") + async with asyncio.timeout(5): + await active + await second_prompt(client) + assert client._proc is process + assert process.returncode is None + assert await sent_prompt_messages(client) == ["first", "second"] + finally: + tasks = [active, *([queued] if queued is not None else [])] + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + @pytest.mark.parametrize("scenario", ["yield", "no-result", "stubborn"]) @pytest.mark.parametrize("interruption", ["timeout", "cancel"]) async def test_interrupted_prompt_aborts_reaps_and_can_restart( diff --git a/tests/integrations/omprpc/test_lifecycle_model.py b/tests/integrations/omprpc/test_lifecycle_model.py new file mode 100644 index 000000000..06fdee3b7 --- /dev/null +++ b/tests/integrations/omprpc/test_lifecycle_model.py @@ -0,0 +1,26 @@ +"""Owned process trees and admitted turns survive generated host interleavings.""" + +from __future__ import annotations + +import asyncio + +from hypothesis import Phase, example, given, settings + +from tests.integrations.omprpc.lifecycle import Step, play, sequences + + +class TestRpcLifecycleModel: + @settings( + max_examples=20, + deadline=None, + phases=(Phase.explicit, Phase.reuse, Phase.generate), + ) + @given(steps=sequences) + @example(steps=[Step.PROMPT, Step.CANCEL_QUEUED, Step.SETTLE]) + @example(steps=[Step.PROMPT, Step.CRASH, Step.PROMPT, Step.CANCEL_OWNER]) + @example(steps=[Step.CANCEL_START_RESTART, Step.PROMPT, Step.CLOSE]) + @example(steps=[Step.CANCEL_CLOSE_RESTART, Step.PROMPT, Step.SETTLE]) + def test_generated_operations_preserve_custody_and_settlement( + self, steps: list[Step] + ) -> None: + asyncio.run(play(steps)) diff --git a/tests/integrations/omprpc/test_process_tree.py b/tests/integrations/omprpc/test_process_tree.py index d1e881378..80e12681f 100644 --- a/tests/integrations/omprpc/test_process_tree.py +++ b/tests/integrations/omprpc/test_process_tree.py @@ -3,21 +3,141 @@ from __future__ import annotations import asyncio +import os import sys from contextlib import suppress import psutil import pytest +from pydantic import JsonValue -from band.integrations.omprpc import OmpRpcClient +from band.integrations import omprpc +from band.integrations.codex.winspawn import Job +from band.integrations.omprpc import OmpProcessError, OmpRpcClient from tests.integrations.omprpc.kit import ( OMP_PEER_SCRIPT, + DeferredSessionSignal, ProcessGate, + descendant_client, + descendant_liveness, process_is_alive, + refuse_process_signal, ) class TestProcessTreeCleanup: + async def test_close_reaps_owned_descendants_after_the_root_has_crashed( + self, process_gate: ProcessGate + ) -> None: + client = descendant_client(process_gate) + crashing: asyncio.Task[dict[str, JsonValue]] | None = None + try: + await client.start() + descendants = [ + await process_gate.connection(), + await process_gate.connection(), + ] + root = client._proc + assert root is not None + assert await descendant_liveness(descendants) == [True, True] + + crashing = asyncio.create_task(client.request("crash")) + async with asyncio.timeout(5): + await root.wait() + assert root.returncode == 3 + with pytest.raises(OmpProcessError): + await crashing + await client.close() + + assert await descendant_liveness(descendants) == [False, False] + await client.close() + finally: + await client.close() + if crashing is not None: + await asyncio.gather(crashing, return_exceptions=True) + + async def test_failed_containment_cleanup_is_retried_before_restart( + self, process_gate: ProcessGate, monkeypatch: pytest.MonkeyPatch + ) -> None: + client = descendant_client(process_gate) + crashing: asyncio.Task[dict[str, JsonValue]] | None = None + try: + await client.start() + old_root = client._proc + assert old_root is not None + descendants = [ + await process_gate.connection(), + await process_gate.connection(), + ] + with monkeypatch.context() as patch: + if os.name == "nt": + patch.setattr(Job, "terminate", refuse_process_signal) + else: + patch.setattr(os, "killpg", refuse_process_signal) + crashing = asyncio.create_task(client.request("crash")) + async with asyncio.timeout(5): + await old_root.wait() + with pytest.raises(PermissionError, match="not signalable"): + await client.close() + assert await descendant_liveness(descendants) == [True, True] + + await client.start() + assert client._proc is not old_root + assert await descendant_liveness(descendants) == [False, False] + await process_gate.connection() + await process_gate.connection() + await client.close() + finally: + await client.close() + if crashing is not None: + await asyncio.gather(crashing, return_exceptions=True) + + @pytest.mark.skipif(os.name == "nt", reason="POSIX group settlement boundary") + @pytest.mark.parametrize( + "hide_after_signal", [False, True], ids=["visible", "hidden-after-signal"] + ) + async def test_unsettled_session_members_keep_close_retryable( + self, + process_gate: ProcessGate, + monkeypatch: pytest.MonkeyPatch, + hide_after_signal: bool, + ) -> None: + client = descendant_client(process_gate) + crashing: asyncio.Task[dict[str, JsonValue]] | None = None + try: + await client.start() + old_root = client._proc + assert old_root is not None + descendants = [ + await process_gate.connection(), + await process_gate.connection(), + ] + boundary = DeferredSessionSignal( + old_root.pid, descendants, hide_after_signal=hide_after_signal + ) + with monkeypatch.context() as patch: + patch.setattr(os, "killpg", boundary.send) + patch.setattr(os, "getpgid", boundary.group) + patch.setattr(omprpc, "_PROCESS_EXIT_TIMEOUT_S", 0) + crashing = asyncio.create_task(client.request("crash")) + async with asyncio.timeout(5): + await old_root.wait() + for _ in range(2): + with pytest.raises(OmpProcessError, match="session members"): + await client.close() + assert await descendant_liveness(descendants) == [True, True] + + await client.start() + assert client._proc is not old_root + assert await descendant_liveness(descendants) == [False, False] + await process_gate.connection() + await process_gate.connection() + await client.close() + finally: + await client.close() + if crashing is not None: + await asyncio.gather(crashing, return_exceptions=True) + @pytest.mark.parametrize("stdio", ["inherit", "discard"]) @pytest.mark.parametrize("interruption", ["startup", "close"]) async def test_cancelled_cleanup_reaps_descendants_but_not_unrelated_processes( From 259d1eb591ca2bc036c0b69fb892eab1de43a9bd Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 08:10:30 +0300 Subject: [PATCH 04/35] fix: reap discoverable native detached descendants --- src/band/integrations/omprpc.py | 148 +++++++++++++----- tests/integrations/omppeer.py | 28 +++- tests/integrations/omprpc/kit.py | 86 +++++++++- tests/integrations/omprpc/lifecycle.py | 12 +- .../omprpc/test_lifecycle_model.py | 19 ++- .../integrations/omprpc/test_process_tree.py | 84 +++++++++- 6 files changed, 318 insertions(+), 59 deletions(-) diff --git a/src/band/integrations/omprpc.py b/src/band/integrations/omprpc.py index fe57c64aa..73ed7d564 100644 --- a/src/band/integrations/omprpc.py +++ b/src/band/integrations/omprpc.py @@ -6,10 +6,12 @@ import base64 import logging import os +import time from binascii import Error as Base64Error from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass from typing import cast +from uuid import uuid4 import psutil from pydantic import JsonValue, TypeAdapter, ValidationError @@ -52,6 +54,7 @@ _ABORT_TIMEOUT_S = 1.0 _PROCESS_EXIT_TIMEOUT_S = 1.0 _CALLBACK_EXIT_TIMEOUT_S = 1.0 +_PROCESS_SCOPE_VALUE = "1" _PAYLOAD_ADAPTER: TypeAdapter[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]) _EVENT_TYPES = frozenset(OmpEventType) _UI_NOTIFICATIONS = frozenset( @@ -153,6 +156,88 @@ def observe(self, event: OmpEvent, sequence: int) -> None: self.changed.set() +class OmpProcessScope: + """Retain observable POSIX descendants for best-effort cleanup. + + The private PI_ marker survives native helper filters. Explicit environment + replacement and restricted process inspection can hide detached children. + """ + + def __init__(self) -> None: + self.key = f"PI_BAND_PROCESS_{uuid4().hex}" + self.processes: set[psutil.Process] = set() + + def environment(self, base: Mapping[str, str] | None) -> dict[str, str]: + return { + **(os.environ if base is None else base), + self.key: _PROCESS_SCOPE_VALUE, + } + + def capture(self, root: asyncio.subprocess.Process) -> None: + group_visible = not session_pid_reused(root) + group_members: list[psutil.Process] = [] + for pid in psutil.pids(): + # asyncio's child watcher remains the sole waiter for the root. + if pid == 0 or pid == root.pid: + continue + try: + member = psutil.Process(pid) + created = member.create_time() + marked = False + try: + marked = member.environ().get(self.key) == _PROCESS_SCOPE_VALUE + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass + in_group = False + if group_visible: + try: + in_group = os.getpgid(pid) == root.pid + except (ProcessLookupError, PermissionError): + pass + if not marked and not in_group: + continue + current = psutil.Process(pid) + if current.create_time() != created or current != member: + continue + if marked: + self.processes.add(member) + elif in_group: + group_members.append(member) + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + if group_visible and not session_pid_reused(root): + self.processes.update(group_members) + + @staticmethod + def _is_alive(process: psutil.Process) -> bool: + try: + return process.is_running() and process.status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + def signal(self) -> None: + for process in self.processes: + try: + process.kill() + except psutil.NoSuchProcess: + pass + + def close(self, root: asyncio.subprocess.Process, timeout: float) -> None: + deadline = time.monotonic() + timeout + self.capture(root) + alive = [process for process in self.processes if self._is_alive(process)] + while alive: + self.signal() + psutil.wait_procs(alive, timeout=max(0.0, deadline - time.monotonic())) + self.capture(root) + alive = [process for process in self.processes if self._is_alive(process)] + if alive and time.monotonic() >= deadline: + raise OmpProcessError( + "OMP owned session members survived cleanup: " + f"{sorted(process.pid for process in alive)}" + ) + + class OmpRpcClient: """One OMP process, one persistent reader, and at most one active prompt.""" @@ -165,7 +250,7 @@ def __init__( self.env = None if env is None else {**os.environ, **env} self._proc: asyncio.subprocess.Process | None = None self._windows_process: WindowsProcess | None = None - self._session_waiters: set[psutil.Process] = set() + self._process_scope: OmpProcessScope | None = None self._spawn_task: asyncio.Task[asyncio.subprocess.Process] | None = None self._reader_task: asyncio.Task[None] | None = None self._stderr_task: asyncio.Task[None] | None = None @@ -200,6 +285,7 @@ async def start(self) -> None: self._max_frame_bytes = OMP_FRAME_BYTES self._max_reassembled_bytes = OMP_REASSEMBLED_BYTES self._closing = False + self._process_scope = OmpProcessScope() if os.name != "nt" else None self._ready = asyncio.get_running_loop().create_future() self._spawn_task = asyncio.create_task(self._spawn_process()) try: @@ -242,13 +328,15 @@ async def _spawn_process(self) -> asyncio.subprocess.Process: limit=OMP_FRAME_BYTES, ) else: + scope = self._process_scope + assert scope is not None process = await asyncio.create_subprocess_exec( *self.command, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=None, - env=self.env, + env=scope.environment(self.env), start_new_session=True, limit=OMP_FRAME_BYTES, ) @@ -803,40 +891,6 @@ async def _wait_for_exit(process: asyncio.subprocess.Process) -> bool: return False return True - def _capture_session_waiters(self, process: asyncio.subprocess.Process) -> None: - if session_pid_reused(process): - return - members: list[psutil.Process] = [] - for pid in psutil.pids(): - # asyncio's child watcher remains the sole waiter for the root. - if pid == process.pid: - continue - try: - if os.getpgid(pid) == process.pid: - member = psutil.Process(pid) - member.create_time() - members.append(member) - except (ProcessLookupError, psutil.NoSuchProcess): - continue - if not session_pid_reused(process): - self._session_waiters.update(members) - - def _wait_for_session_exit(self) -> None: - _, alive = psutil.wait_procs( - self._session_waiters, timeout=_PROCESS_EXIT_TIMEOUT_S - ) - remaining: list[int] = [] - for member in alive: - try: - if member.is_running() and member.status() != psutil.STATUS_ZOMBIE: - remaining.append(member.pid) - except psutil.NoSuchProcess: - pass - if remaining: - raise OmpProcessError( - f"OMP session members survived containment kill: {sorted(remaining)}" - ) - async def _settle_spawn(self) -> asyncio.subprocess.Process | None: spawning = self._spawn_task if spawning is None: @@ -860,6 +914,7 @@ async def _close_process(self) -> None: await self._cancel_tasks(workers) owner = self._windows_process + scope = self._process_scope owner_closed = False if owner is not None and self._proc is None: # Stop incomplete helper creation/readiness before awaiting its spawn task. @@ -871,7 +926,10 @@ async def _close_process(self) -> None: process = await self._settle_spawn() if process is None: self._windows_process = None + self._process_scope = None return + if scope is not None: + await asyncio.to_thread(scope.capture, process) if ( self._failure is None and process.returncode is None @@ -893,9 +951,9 @@ async def _close_process(self) -> None: ) and process.stderr is not None: drains.append(asyncio.create_task(self._drain(process.stderr))) try: - if owner is None: - # EOF and signals can remove PGID visibility before exit is observable. - await asyncio.to_thread(self._capture_session_waiters, process) + if scope is not None: + # Ownership lookup can disappear before process exit is observable. + await asyncio.to_thread(scope.capture, process) if not owner_closed: await close_stdin(process) await await_returncode(process, _PROCESS_EXIT_TIMEOUT_S) @@ -903,13 +961,19 @@ async def _close_process(self) -> None: if not owner_closed: await owner.close(_PROCESS_EXIT_TIMEOUT_S) else: - await asyncio.to_thread(self._capture_session_waiters, process) + assert scope is not None + await asyncio.to_thread(scope.capture, process) await kill_process_tree(process) - await asyncio.to_thread(self._wait_for_session_exit) + await asyncio.to_thread(scope.signal) + if not await await_returncode(process, _PROCESS_EXIT_TIMEOUT_S): + raise OmpProcessError( + "OMP process did not exit after containment kill" + ) + await asyncio.to_thread(scope.close, process, _PROCESS_EXIT_TIMEOUT_S) if not await self._wait_for_exit(process): raise OmpProcessError("OMP process did not exit after containment kill") self._windows_process = None - self._session_waiters.clear() + self._process_scope = None finally: if self._stderr_task is not None: drains.append(self._stderr_task) diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index e74c82945..7cb594370 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -41,7 +41,9 @@ class ProcessReceipt(BaseModel): created: float -def start_descendant(scenario: str, port: int, stdio: str) -> subprocess.Popen[bytes]: +def start_descendant( + scenario: str, port: int, stdio: str, *, detached: bool = False +) -> subprocess.Popen[bytes]: return subprocess.Popen( ( sys.executable, @@ -51,19 +53,25 @@ def start_descendant(scenario: str, port: int, stdio: str) -> subprocess.Popen[b scenario, "--gate-port", str(port), + *(("--detached-children",) if detached else ()), ), # A parent's pending pipe read can block Python's stdin initialization on Windows. # Descendants are controlled by the gate socket, not the RPC input pipe. stdin=subprocess.DEVNULL, stdout=None if stdio == "inherit" else subprocess.DEVNULL, stderr=None if stdio == "inherit" else subprocess.DEVNULL, + start_new_session=detached and os.name != "nt", ) -def hold_process_gate(port: int, *, spawn_child: bool) -> None: +def hold_process_gate(port: int, *, spawn_child: bool, detached: bool = False) -> None: if hasattr(signal, "SIGTERM"): signal.signal(signal.SIGTERM, signal.SIG_IGN) - child = start_descendant("gated-child", port, "inherit") if spawn_child else None + child = ( + start_descendant("gated-child", port, "inherit", detached=detached) + if spawn_child + else None + ) with socket.create_connection(("127.0.0.1", port), timeout=5) as gate: gate.settimeout(None) receipt = ProcessReceipt( @@ -556,6 +564,7 @@ def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--scenario", required=True) parser.add_argument("--gate-port", type=int) + parser.add_argument("--detached-children", action="store_true") parser.add_argument( "--child-stdio", choices=("inherit", "discard"), default="inherit" ) @@ -563,13 +572,22 @@ def main() -> None: if args.scenario in {"gated-parent", "gated-child"}: if args.gate_port is None: raise ValueError("A process gate port is required") - hold_process_gate(args.gate_port, spawn_child=args.scenario == "gated-parent") + hold_process_gate( + args.gate_port, + spawn_child=args.scenario == "gated-parent", + detached=args.detached_children, + ) return if args.scenario in {"descendants", "no-ready-descendants", "lifecycle"}: if args.gate_port is None: raise ValueError("A process gate port is required") # The gate, not the peer's stdio EOF, owns these children's lifetime. - start_descendant("gated-parent", args.gate_port, args.child_stdio) + start_descendant( + "gated-parent", + args.gate_port, + args.child_stdio, + detached=args.detached_children, + ) if args.scenario == "stubborn" and hasattr(signal, "SIGTERM"): signal.signal(signal.SIGTERM, signal.SIG_IGN) peer = Peer(args.scenario) diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py index 958c24448..9dcb8de75 100644 --- a/tests/integrations/omprpc/kit.py +++ b/tests/integrations/omprpc/kit.py @@ -13,6 +13,7 @@ from typing import cast import psutil +import pytest from pydantic import JsonValue from band.integrations.omprpc import OmpRpcClient @@ -120,7 +121,11 @@ async def queue_prompt( def descendant_client( - gate: ProcessGate, *, scenario: str = "descendants", stdio: str = "discard" + gate: ProcessGate, + *, + scenario: str = "descendants", + stdio: str = "discard", + detached: bool = False, ) -> OmpRpcClient: return OmpRpcClient( ( @@ -133,6 +138,7 @@ def descendant_client( str(gate.port), "--child-stdio", stdio, + *(("--detached-children",) if detached else ()), ) ) @@ -231,7 +237,7 @@ def refuse_process_signal(*args: object) -> None: class DeferredSessionSignal: - """Keep real members alive while their post-signal group lookup disappears.""" + """Keep real members alive while post-signal ownership lookup disappears.""" def __init__( self, @@ -245,13 +251,89 @@ def __init__( self.hide_after_signal = hide_after_signal self.signalled = False self.original_group = os.getpgid + self.original_kill = os.kill + self.original_environment = psutil.Process.environ def send(self, pgid: int, sig: int) -> None: assert pgid == self.pgid assert sig == signal.SIGKILL self.signalled = True + def send_pid(self, pid: int, sig: int) -> None: + if sig == signal.SIGKILL: + assert pid in self.members + self.signalled = True + return + self.original_kill(pid, sig) + def group(self, pid: int) -> int: if self.signalled and self.hide_after_signal and pid in self.members: raise ProcessLookupError(pid) return self.original_group(pid) + + def environment(self, process: psutil.Process) -> dict[str, str]: + if self.signalled and self.hide_after_signal and process.pid in self.members: + raise psutil.NoSuchProcess(process.pid) + return self.original_environment(process) + + def install(self, patch: pytest.MonkeyPatch) -> None: + patch.setattr(os, "killpg", self.send) + patch.setattr(os, "kill", self.send_pid) + patch.setattr(os, "getpgid", self.group) + patch.setattr( + psutil.Process, "environ", lambda process: self.environment(process) + ) + + +@asynccontextmanager +async def unrelated_gate_child(gate: ProcessGate) -> AsyncIterator[GatedProcess]: + process = await asyncio.create_subprocess_exec( + sys.executable, + "-u", + OMP_PEER_SCRIPT, + "--scenario", + "gated-child", + "--gate-port", + str(gate.port), + stdin=asyncio.subprocess.DEVNULL, + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + start_new_session=os.name != "nt", + ) + try: + yield await gate.connection() + finally: + if process.returncode is None: + with suppress(ProcessLookupError): + process.kill() + async with asyncio.timeout(5): + await process.wait() + + +class RecycledEnvironment: + """Expose one stale matching environment while the process identity changes.""" + + def __init__(self, unrelated: GatedProcess, owner_pid: int) -> None: + self.pid = unrelated.process.pid + self.owner_environment = psutil.Process(owner_pid).environ() + self.original_environment = psutil.Process.environ + self.original_create_time = psutil.Process.create_time + self.recycled = False + + def environment(self, process: psutil.Process) -> dict[str, str]: + if process.pid == self.pid and not self.recycled: + self.recycled = True + return self.owner_environment + return self.original_environment(process) + + def create_time(self, process: psutil.Process) -> float: + created = self.original_create_time(process) + return created + 1 if self.recycled and process.pid == self.pid else created + + def install(self, patch: pytest.MonkeyPatch) -> None: + patch.setattr( + psutil.Process, "environ", lambda process: self.environment(process) + ) + patch.setattr( + psutil.Process, "create_time", lambda process: self.create_time(process) + ) diff --git a/tests/integrations/omprpc/lifecycle.py b/tests/integrations/omprpc/lifecycle.py index e3d8093a4..67ee5569a 100644 --- a/tests/integrations/omprpc/lifecycle.py +++ b/tests/integrations/omprpc/lifecycle.py @@ -70,9 +70,9 @@ async def announced_operation( class RpcLifecycle: - def __init__(self, gate: ProcessGate) -> None: + def __init__(self, gate: ProcessGate, *, detached: bool) -> None: self.gate = gate - self.client = descendant_client(gate, scenario="lifecycle") + self.client = descendant_client(gate, scenario="lifecycle", detached=detached) self.live = False self.active: asyncio.Task[None] | None = None self.observer: Observer | None = None @@ -234,17 +234,17 @@ async def cleanup(self) -> None: @asynccontextmanager -async def rpc_lifecycle() -> AsyncIterator[RpcLifecycle]: +async def rpc_lifecycle(*, detached: bool) -> AsyncIterator[RpcLifecycle]: async with process_gate_session() as gate: - session = RpcLifecycle(gate) + session = RpcLifecycle(gate, detached=detached) try: yield session finally: await session.cleanup() -async def play(steps: list[Step]) -> None: - async with rpc_lifecycle() as session: +async def play(steps: list[Step], *, detached: bool) -> None: + async with rpc_lifecycle(detached=detached) as session: for step in steps: async with asyncio.timeout(STEP_GUARD_S): await session.take(step) diff --git a/tests/integrations/omprpc/test_lifecycle_model.py b/tests/integrations/omprpc/test_lifecycle_model.py index 06fdee3b7..f34afb0a8 100644 --- a/tests/integrations/omprpc/test_lifecycle_model.py +++ b/tests/integrations/omprpc/test_lifecycle_model.py @@ -4,12 +4,16 @@ import asyncio +import pytest from hypothesis import Phase, example, given, settings from tests.integrations.omprpc.lifecycle import Step, play, sequences class TestRpcLifecycleModel: + @pytest.mark.parametrize("detached", [False, True], ids=["session", "detached"]) + # The generated workload launches hundreds of real peers; each step stays guarded. + @pytest.mark.timeout(120) @settings( max_examples=20, deadline=None, @@ -20,7 +24,18 @@ class TestRpcLifecycleModel: @example(steps=[Step.PROMPT, Step.CRASH, Step.PROMPT, Step.CANCEL_OWNER]) @example(steps=[Step.CANCEL_START_RESTART, Step.PROMPT, Step.CLOSE]) @example(steps=[Step.CANCEL_CLOSE_RESTART, Step.PROMPT, Step.SETTLE]) + @example(steps=[Step.PROMPT, Step.PROMPT, Step.CANCEL_OWNER, Step.CANCEL_OWNER]) + @example( + steps=[ + Step.CANCEL_CLOSE_RESTART, + Step.CANCEL_OWNER, + Step.CANCEL_CLOSE_RESTART, + Step.CANCEL_CLOSE_RESTART, + Step.CANCEL_OWNER, + Step.CANCEL_OWNER, + ] + ) def test_generated_operations_preserve_custody_and_settlement( - self, steps: list[Step] + self, steps: list[Step], detached: bool ) -> None: - asyncio.run(play(steps)) + asyncio.run(play(steps, detached=detached)) diff --git a/tests/integrations/omprpc/test_process_tree.py b/tests/integrations/omprpc/test_process_tree.py index 80e12681f..5f01ee6dc 100644 --- a/tests/integrations/omprpc/test_process_tree.py +++ b/tests/integrations/omprpc/test_process_tree.py @@ -18,14 +18,95 @@ OMP_PEER_SCRIPT, DeferredSessionSignal, ProcessGate, + RecycledEnvironment, descendant_client, descendant_liveness, process_is_alive, refuse_process_signal, + unrelated_gate_child, ) class TestProcessTreeCleanup: + @pytest.mark.skipif(os.name == "nt", reason="POSIX setsid ownership boundary") + @pytest.mark.parametrize("root_exit", ["close", "crash"]) + async def test_close_reaps_detached_descendants_without_affecting_other_owners( + self, process_gate: ProcessGate, root_exit: str + ) -> None: + client = descendant_client(process_gate, detached=True) + other = descendant_client(process_gate, detached=True) + crashing: asyncio.Task[dict[str, JsonValue]] | None = None + async with unrelated_gate_child(process_gate) as unrelated: + try: + await client.start() + descendants = [ + await process_gate.connection(), + await process_gate.connection(), + ] + root = client._proc + assert root is not None + assert all( + os.getsid(child.process.pid) == child.process.pid + for child in descendants + ) + assert all( + os.getsid(child.process.pid) != root.pid for child in descendants + ) + await other.start() + other_descendants = [ + await process_gate.connection(), + await process_gate.connection(), + ] + other_root = other._proc + assert other_root is not None + if root_exit == "crash": + crashing = asyncio.create_task(client.request("crash")) + async with asyncio.timeout(5): + await root.wait() + assert root.returncode == 3 + with pytest.raises(OmpProcessError): + await crashing + + await client.close() + assert root.returncode is not None + assert await descendant_liveness(descendants) == [False, False] + assert other_root.returncode is None + assert await descendant_liveness(other_descendants) == [True, True] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + assert (await other.request("get_state"))["isSettled"] is True + await other.close() + assert await descendant_liveness(other_descendants) == [False, False] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + finally: + await client.close() + await other.close() + if crashing is not None: + await asyncio.gather(crashing, return_exceptions=True) + + @pytest.mark.skipif(os.name == "nt", reason="POSIX inherited-scope boundary") + async def test_stale_matching_environment_cannot_adopt_a_reused_process_identity( + self, process_gate: ProcessGate, monkeypatch: pytest.MonkeyPatch + ) -> None: + client = descendant_client(process_gate, detached=True) + async with unrelated_gate_child(process_gate) as unrelated: + try: + await client.start() + descendants = [ + await process_gate.connection(), + await process_gate.connection(), + ] + root = client._proc + assert root is not None + boundary = RecycledEnvironment(unrelated, root.pid) + with monkeypatch.context() as patch: + boundary.install(patch) + await client.close() + assert boundary.recycled + assert await descendant_liveness(descendants) == [False, False] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + finally: + await client.close() + async def test_close_reaps_owned_descendants_after_the_root_has_crashed( self, process_gate: ProcessGate ) -> None: @@ -116,8 +197,7 @@ async def test_unsettled_session_members_keep_close_retryable( old_root.pid, descendants, hide_after_signal=hide_after_signal ) with monkeypatch.context() as patch: - patch.setattr(os, "killpg", boundary.send) - patch.setattr(os, "getpgid", boundary.group) + boundary.install(patch) patch.setattr(omprpc, "_PROCESS_EXIT_TIMEOUT_S", 0) crashing = asyncio.create_task(client.request("crash")) async with asyncio.timeout(5): From f696f09d2164fb0c3934c4ae849ddbea5fd53330 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 09:27:14 +0300 Subject: [PATCH 05/35] fix: bound native event delivery and complete wire contracts --- docs/acp.md | 12 ++ src/band/integrations/ompprocess.py | 22 +-- src/band/integrations/omprpc.py | 14 +- src/band/integrations/ompwire.py | 55 ++++-- tests/integrations/omppeer.py | 95 ++++++++- tests/integrations/omprpc/cancellation.py | 62 ++++++ .../omprpc/test_eof_descendants.py | 87 +++++++++ .../omprpc/test_event_backpressure.py | 150 +++++++++++++++ .../integrations/omprpc/test_interruption.py | 31 ++- .../integrations/omprpc/test_process_tree.py | 180 ++++++++---------- tests/integrations/test_ompwire.py | 74 +++++++ 11 files changed, 641 insertions(+), 141 deletions(-) create mode 100644 tests/integrations/omprpc/cancellation.py create mode 100644 tests/integrations/omprpc/test_eof_descendants.py create mode 100644 tests/integrations/omprpc/test_event_backpressure.py create mode 100644 tests/integrations/test_ompwire.py diff --git a/docs/acp.md b/docs/acp.md index b57c2f1d1..a1530358b 100644 --- a/docs/acp.md +++ b/docs/acp.md @@ -235,3 +235,15 @@ failure fails that room turn visibly instead of falling back. credential-redacted; values under credential-named fields are redacted in full. - The ACP Python client drains queued update handlers before surfacing a prompt error, so a slow client handler can still delay its own `prompt()`. + +## Native OMP transport + +`band.integrations.omprpc.OmpRpcClient` speaks OMP's native RPC v2 protocol, +including structured ask questions and answers. It is separate from the ACP +transport; adding it does not deprecate or remove the public `OmpACPAdapter`. + +Event callbacks run serially with a bounded backlog. Exhausting that backlog +fails the owning prompt with `OmpProcessError` and starts cleanup rather than +silently dropping events and reporting success. The wire reader does not wait +for callback capacity, so callbacks can still issue correlated commands and +answer native UI requests while delivery is held. diff --git a/src/band/integrations/ompprocess.py b/src/band/integrations/ompprocess.py index 655dc11a6..14f757726 100644 --- a/src/band/integrations/ompprocess.py +++ b/src/band/integrations/ompprocess.py @@ -123,28 +123,28 @@ class OmpStartupError(RuntimeError): """The required native connection context or MCP tools are not ready.""" -class _McpCommand(BaseModel): +class McpCommand(BaseModel): name: str -class _McpCommands(BaseModel): - commands: list[_McpCommand] +class McpCommands(BaseModel): + commands: list[McpCommand] -class _ConnectedMcpServer(BaseModel): +class ConnectedMcpServer(BaseModel): type: Literal["connected"] serverName: str -class _FailedMcpServer(BaseModel): +class FailedMcpServer(BaseModel): type: Literal["failed"] serverName: str error: str -class _McpStartupReceipt(BaseModel): +class McpStartupReceipt(BaseModel): systemPromptReady: bool = Field(strict=True) - servers: list[_ConnectedMcpServer | _FailedMcpServer] + servers: list[ConnectedMcpServer | FailedMcpServer] def _plugin_server_name(name: str) -> str: @@ -153,7 +153,7 @@ def _plugin_server_name(name: str) -> str: def omp_mcp_ready_command(plugin: str) -> str: """The slash command the plugin's startup extension registers for its receipt.""" - return f"band-sdk-mcp-{Path(plugin).parent.name}" + return f"{OMP_PLUGIN_ROOT}-mcp-{Path(plugin).parent.name}" def omp_mcp_arguments(plugin: str) -> tuple[str, ...]: @@ -176,7 +176,7 @@ async def omp_mcp_ready( ) -> None: """Require the connection extension and configured MCP servers before inference.""" command = omp_mcp_ready_command(plugin) - commands = _McpCommands.model_validate( + commands = McpCommands.model_validate( await runtime.request("get_available_commands") ) if not any(item.name == command for item in commands.commands): @@ -201,7 +201,7 @@ async def on_ui(_request: OmpUIRequest) -> OmpUIResponse: on_ui=on_ui, ) try: - receipt = _McpStartupReceipt.model_validate_json( + receipt = McpStartupReceipt.model_validate_json( (Path(plugin) / "receipt.json").read_text(encoding="utf-8") ) except (OSError, ValueError) as error: @@ -215,7 +215,7 @@ async def on_ui(_request: OmpUIRequest) -> OmpUIResponse: for status in receipt.servers: if status.serverName not in expected: raise OmpStartupError("OMP returned an unexpected MCP startup receipt") - if isinstance(status, _FailedMcpServer): + if isinstance(status, FailedMcpServer): raise OmpStartupError( f"OMP MCP server {status.serverName} failed: {status.error}" ) diff --git a/src/band/integrations/omprpc.py b/src/band/integrations/omprpc.py index 73ed7d564..595a7b3f3 100644 --- a/src/band/integrations/omprpc.py +++ b/src/band/integrations/omprpc.py @@ -51,6 +51,7 @@ logger = logging.getLogger(__name__) _DRAIN_BYTES = 64 * 1024 +_EVENT_BACKLOG_LIMIT = 256 _ABORT_TIMEOUT_S = 1.0 _PROCESS_EXIT_TIMEOUT_S = 1.0 _CALLBACK_EXIT_TIMEOUT_S = 1.0 @@ -123,7 +124,9 @@ def __init__( self.on_event = on_event self.on_tool = on_tool self.on_ui = on_ui - self.events: asyncio.Queue[OmpEvent] = asyncio.Queue() + self.events: asyncio.Queue[OmpEvent] = asyncio.Queue( + maxsize=_EVENT_BACKLOG_LIMIT + ) self.failure: asyncio.Future[BaseException] = loop.create_future() self.result: asyncio.Future[OmpPromptResult] = loop.create_future() self.changed = asyncio.Event() @@ -140,6 +143,14 @@ def fail(self, error: BaseException) -> None: self.failure.set_result(error) def observe(self, event: OmpEvent, sequence: int) -> None: + if self.failure.done(): + return + # Waiting for space would block the reader that services reentrant commands. + try: + self.events.put_nowait(event) + except asyncio.QueueFull: + self.fail(OmpProcessError("OMP event callback backlog exceeded")) + return self.last_activity_sequence = sequence if ( event.type == OmpEventType.MESSAGE_END @@ -152,7 +163,6 @@ def observe(self, event: OmpEvent, sequence: int) -> None: if message.role == "assistant": self.last_assistant = message break - self.events.put_nowait(event) self.changed.set() diff --git a/src/band/integrations/ompwire.py b/src/band/integrations/ompwire.py index 7e0042ce9..c6bacd8a4 100644 --- a/src/band/integrations/ompwire.py +++ b/src/band/integrations/ompwire.py @@ -192,7 +192,7 @@ class OmpState(OmpWireModel): class OmpHostToolCall(OmpWireModel): - type: Literal["host_tool_call"] = "host_tool_call" + type: Literal[OmpFrameType.HOST_TOOL_CALL] = OmpFrameType.HOST_TOOL_CALL id: str toolCallId: str toolName: str @@ -200,38 +200,69 @@ class OmpHostToolCall(OmpWireModel): class OmpHostToolCancel(OmpWireModel): - type: Literal["host_tool_cancel"] = "host_tool_cancel" + type: Literal[OmpFrameType.HOST_TOOL_CANCEL] = OmpFrameType.HOST_TOOL_CANCEL id: str targetId: str +class OmpAskOption(OmpWireModel): + label: str + description: str | None = None + preview: str | None = None + + +class OmpAskQuestion(OmpWireModel): + """A structured native ask question, enabled by ``set_ask_dialog``.""" + + id: str + question: str + header: str | None = None + options: list[OmpAskOption] = Field(default_factory=list) + multi: bool = False + # The tool schema accepts any number and forwards the recommendation unchanged. + recommended: float | None = None + + +class OmpAskAnswer(OmpWireModel): + """One answer per question, in question order, using offered option labels.""" + + id: str + selectedOptions: list[str] = Field(default_factory=list) + customInput: str | None = None + + class OmpUIRequest(OmpWireModel): - type: Literal["extension_ui_request"] = "extension_ui_request" + type: Literal[OmpFrameType.UI_REQUEST] = OmpFrameType.UI_REQUEST id: str method: str title: str | None = None options: list[str] = Field(default_factory=list) message: str | None = None targetId: str | None = None + questions: list[OmpAskQuestion] = Field(default_factory=list) class OmpUIResponse(OmpWireModel): value: str | None = None confirmed: bool | None = None + answers: list[OmpAskAnswer] | None = None cancelled: bool = False @model_validator(mode="after") def validate_answer(self) -> OmpUIResponse: - if ( - sum((self.value is not None, self.confirmed is not None, self.cancelled)) - > 1 - ): + given = ( + self.value is not None, + self.confirmed is not None, + self.answers is not None, + self.cancelled, + ) + if sum(given) > 1: raise ValueError("OMP UI response must contain only one answer") return self class OmpReady(OmpWireModel): - type: Literal["ready"] + type: Literal[OmpFrameType.READY] protocolVersion: Literal[1] supportedProtocolVersions: list[int] maxFrameBytes: int = Field(gt=0, le=OMP_FRAME_BYTES) @@ -243,7 +274,7 @@ class OmpNegotiation(OmpWireModel): class OmpChunk(OmpWireModel): - type: Literal["rpc_chunk"] + type: Literal[OmpFrameType.CHUNK] chunkId: str = Field(min_length=1, max_length=128) index: int = Field(ge=0) count: int = Field(ge=2, le=OMP_REASSEMBLED_BYTES // OMP_CHUNK_BYTES) @@ -258,12 +289,12 @@ def validate_index(self) -> OmpChunk: class OmpFrameError(OmpWireModel): - type: Literal["rpc_frame_error"] + type: Literal[OmpFrameType.FRAME_ERROR] error: str class OmpResponse(OmpWireModel): - type: Literal["response"] + type: Literal[OmpFrameType.RESPONSE] id: str | None = None command: str success: bool @@ -285,7 +316,7 @@ class OmpPromptFailure(OmpWireModel): class OmpPromptResult(OmpWireModel): - type: Literal["prompt_result"] + type: Literal[OmpFrameType.PROMPT_RESULT] id: str | None = None agentInvoked: bool status: Literal["completed", "error", "aborted"] diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index 7cb594370..3b663e04c 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -33,6 +33,7 @@ class Frame(BaseModel): value: str | None = None confirmed: bool | None = None cancelled: bool | None = None + count: int = 0 class ProcessReceipt(BaseModel): @@ -42,7 +43,12 @@ class ProcessReceipt(BaseModel): def start_descendant( - scenario: str, port: int, stdio: str, *, detached: bool = False + scenario: str, + port: int, + stdio: str, + *, + detached: bool = False, + spawn_on_eof: bool = False, ) -> subprocess.Popen[bytes]: return subprocess.Popen( ( @@ -54,24 +60,35 @@ def start_descendant( "--gate-port", str(port), *(("--detached-children",) if detached else ()), + *(("--spawn-on-eof",) if spawn_on_eof else ()), ), # A parent's pending pipe read can block Python's stdin initialization on Windows. # Descendants are controlled by the gate socket, not the RPC input pipe. - stdin=subprocess.DEVNULL, + stdin=subprocess.PIPE if spawn_on_eof else subprocess.DEVNULL, stdout=None if stdio == "inherit" else subprocess.DEVNULL, stderr=None if stdio == "inherit" else subprocess.DEVNULL, start_new_session=detached and os.name != "nt", ) -def hold_process_gate(port: int, *, spawn_child: bool, detached: bool = False) -> None: +def hold_process_gate( + port: int, + *, + spawn_child: bool, + detached: bool = False, + spawn_on_eof: bool = False, +) -> None: if hasattr(signal, "SIGTERM"): signal.signal(signal.SIGTERM, signal.SIG_IGN) - child = ( - start_descendant("gated-child", port, "inherit", detached=detached) - if spawn_child - else None - ) + + def launch_child() -> subprocess.Popen[bytes] | None: + return ( + start_descendant("gated-child", port, "inherit", detached=detached) + if spawn_child + else None + ) + + child = None if spawn_on_eof else launch_child() with socket.create_connection(("127.0.0.1", port), timeout=5) as gate: gate.settimeout(None) receipt = ProcessReceipt( @@ -80,6 +97,9 @@ def hold_process_gate(port: int, *, spawn_child: bool, detached: bool = False) - created=psutil.Process().create_time(), ) gate.sendall((receipt.model_dump_json() + "\n").encode()) + if spawn_on_eof: + sys.stdin.buffer.read() + child = launch_child() gate.recv(1) if child is not None: child.wait(timeout=5) @@ -327,6 +347,8 @@ def prompt(self, frame: Frame) -> None: match self.scenario: case scenario if scenario.startswith("chunk-"): self.invalid_chunk() + case "event-backpressure": + return case "frame-overflow": send( { @@ -466,6 +488,39 @@ def handle(self, frame: Frame) -> None: ) if self.scenario == "settles-during-state" and not self.settled: self.settle() + if ( + self.scenario == "event-backpressure" + and self.held is not None + and self.held.type == "event-result" + ): + self.respond(self.held) + self.held = None + case "event-burst": + for number in range(frame.count): + send( + { + "type": "message_update", + "messageId": f"burst-{number}", + "message": {"role": "assistant", "content": str(number)}, + } + ) + self.respond(frame) + case "event-result": + self.held = frame + self.end("complete", compacted=True) + self.result(settled=True) + self.settle() + case "event-ui": + self.held = frame + send( + { + "type": "extension_ui_request", + "id": "backlog-approval", + "method": "select", + "title": "Allow command?", + "options": ["Approve", "Deny"], + } + ) case "release": if self.scenario == "settlement-before-result": self.result(settled=False) @@ -513,6 +568,15 @@ def handle(self, frame: Frame) -> None: self.tool_results.add(frame.id) if frame.type == "extension_ui_response": self.ui_answered = True + if ( + self.scenario == "event-backpressure" + and frame.type == "extension_ui_response" + and frame.id == "backlog-approval" + and self.held is not None + and self.held.type == "event-ui" + ): + self.respond(self.held, {"value": frame.value}) + self.held = None if ( self.scenario == "callbacks" and self.tool_results == {"host-1", "host-2"} @@ -565,6 +629,7 @@ def main() -> None: parser.add_argument("--scenario", required=True) parser.add_argument("--gate-port", type=int) parser.add_argument("--detached-children", action="store_true") + parser.add_argument("--spawn-on-eof", action="store_true") parser.add_argument( "--child-stdio", choices=("inherit", "discard"), default="inherit" ) @@ -576,17 +641,25 @@ def main() -> None: args.gate_port, spawn_child=args.scenario == "gated-parent", detached=args.detached_children, + spawn_on_eof=args.spawn_on_eof, ) return - if args.scenario in {"descendants", "no-ready-descendants", "lifecycle"}: + managed_child: subprocess.Popen[bytes] | None = None + if args.scenario in { + "descendants", + "no-ready-descendants", + "lifecycle", + "eof-descendants", + }: if args.gate_port is None: raise ValueError("A process gate port is required") # The gate, not the peer's stdio EOF, owns these children's lifetime. - start_descendant( + managed_child = start_descendant( "gated-parent", args.gate_port, args.child_stdio, detached=args.detached_children, + spawn_on_eof=args.scenario == "eof-descendants", ) if args.scenario == "stubborn" and hasattr(signal, "SIGTERM"): signal.signal(signal.SIGTERM, signal.SIG_IGN) @@ -605,6 +678,8 @@ def main() -> None: ) for line in sys.stdin: peer.handle(Frame.model_validate_json(line, strict=True)) + if managed_child is not None and managed_child.stdin is not None: + managed_child.stdin.close() if args.scenario == "stubborn": threading.Event().wait() diff --git a/tests/integrations/omprpc/cancellation.py b/tests/integrations/omprpc/cancellation.py new file mode 100644 index 000000000..55d3aa71c --- /dev/null +++ b/tests/integrations/omprpc/cancellation.py @@ -0,0 +1,62 @@ +"""Hold real RPC cleanup and observe cancellation-safe joins without elapsed time.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable +from typing import TypeVar + +import pytest + +from band.integrations.omprpc import OmpRpcClient + +Result = TypeVar("Result") + + +class HeldCleanup: + def __init__(self, client: OmpRpcClient, monkeypatch: pytest.MonkeyPatch) -> None: + self.client = client + self.entered = asyncio.Event() + self.release = asyncio.Event() + self.finished = asyncio.Event() + self.task: asyncio.Task[object] | None = None + self.joiners: asyncio.Queue[asyncio.Task[object] | None] = asyncio.Queue() + close_process = client._close_process + shield = asyncio.shield + + async def hold_cleanup() -> None: + self.task = asyncio.current_task() + self.entered.set() + await self.release.wait() + await close_process() + self.finished.set() + + def observe_shield(operation: Awaitable[Result]) -> asyncio.Future[Result]: + waiter = shield(operation) + if operation is client._close_task: + # Queueing does not suspend the caller: it reaches the real shield + # wait before the observing test can resume. + self.joiners.put_nowait(asyncio.current_task()) + return waiter + + monkeypatch.setattr(client, "_close_process", hold_cleanup) + monkeypatch.setattr(asyncio, "shield", observe_shield) + + async def wait_for_join(self, caller: asyncio.Task[None]) -> None: + async with asyncio.timeout(5): + await self.entered.wait() + joining = await self.joiners.get() + assert joining is caller + assert self.task is not None + assert self.client._close_task is self.task + assert not self.task.done() + assert self.task.cancelling() == 0 + assert not caller.done() + + async def cancel_and_wait_for_rejoin(self, caller: asyncio.Task[None]) -> None: + cancellations = caller.cancelling() + assert caller.cancel() + # The next join is emitted only after close catches this cancellation + # and suspends on a fresh shield around the still-held cleanup task. + await self.wait_for_join(caller) + assert caller.cancelling() == cancellations + 1 diff --git a/tests/integrations/omprpc/test_eof_descendants.py b/tests/integrations/omprpc/test_eof_descendants.py new file mode 100644 index 000000000..76d77f2bf --- /dev/null +++ b/tests/integrations/omprpc/test_eof_descendants.py @@ -0,0 +1,87 @@ +"""Cleanup discovers descendants created by an owned helper after native stdin EOF.""" + +from __future__ import annotations + +import asyncio +import os +import threading + +import pytest + +from tests.integrations.omprpc.kit import ( + ProcessGate, + descendant_client, + descendant_liveness, + process_is_alive, + unrelated_gate_child, +) + + +class TestEofDescendants: + @pytest.mark.skipif(os.name == "nt", reason="POSIX detached-helper boundary") + @pytest.mark.parametrize("stdio", ["inherit", "discard"]) + async def test_close_reaps_a_detached_child_created_after_stdin_eof( + self, + process_gate: ProcessGate, + monkeypatch: pytest.MonkeyPatch, + stdio: str, + ) -> None: + client = descendant_client( + process_gate, scenario="eof-descendants", stdio=stdio, detached=True + ) + entered = asyncio.Event() + release = threading.Event() + loop = asyncio.get_running_loop() + closing: asyncio.Task[None] | None = None + + async with unrelated_gate_child(process_gate) as unrelated: + try: + await client.start() + helper = await process_gate.connection() + root = client._proc + scope = client._process_scope + assert root is not None and scope is not None + signal = scope.signal + + def hold_signal() -> None: + loop.call_soon_threadsafe(entered.set) + if not release.wait(timeout=10): + raise TimeoutError( + "Test did not release the containment reaper" + ) + signal() + + with monkeypatch.context() as patch: + patch.setattr(scope, "signal", hold_signal) + closing = asyncio.create_task(client.close()) + try: + async with asyncio.timeout(10): + await entered.wait() + late_child = await process_gate.connection() + assert root.returncode is not None + assert late_child.receipt.parent_pid == helper.process.pid + assert ( + os.getsid(late_child.process.pid) + == late_child.process.pid + ) + assert await descendant_liveness([helper, late_child]) == [ + True, + True, + ] + assert await asyncio.to_thread( + process_is_alive, unrelated.process + ) + assert not closing.done() + release.set() + await closing + finally: + release.set() + await asyncio.gather(closing, return_exceptions=True) + + assert await descendant_liveness([helper, late_child]) == [False, False] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + finally: + release.set() + if closing is not None: + await asyncio.gather(closing, return_exceptions=True) + await client.close() diff --git a/tests/integrations/omprpc/test_event_backpressure.py b/tests/integrations/omprpc/test_event_backpressure.py new file mode 100644 index 000000000..45684b18e --- /dev/null +++ b/tests/integrations/omprpc/test_event_backpressure.py @@ -0,0 +1,150 @@ +"""Callback overload fails closed without blocking native control correlation.""" + +from __future__ import annotations + +import asyncio + +import pytest +from pydantic import JsonValue + +from band.integrations import omprpc +from band.integrations.omprpc import OmpProcessError +from band.integrations.ompwire import ( + OmpEvent, + OmpEventType, + OmpUIRequest, + OmpUIResponse, +) +from tests.integrations.omprpc.kit import ClientFactory, Observer, second_prompt + + +class TestEventBackpressure: + async def test_full_backlog_keeps_reentrant_control_live_and_preserves_event_order( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("event-backpressure") + process = client._proc + assert process is not None + observer = Observer() + entered = asyncio.Event() + release = asyncio.Event() + delivered = asyncio.Event() + event_count = omprpc._EVENT_BACKLOG_LIMIT + + async def on_event(event: OmpEvent) -> None: + if event.type == OmpEventType.AGENT_START: + entered.set() + await release.wait() + await observer.on_event(event) + if event.messageId == f"burst-{event_count - 1}": + delivered.set() + + async def on_ui(request: OmpUIRequest) -> OmpUIResponse: + await client.request("checkpoint") + return OmpUIResponse(value="Approve") + + running = asyncio.create_task( + client.prompt( + "first", + on_event=on_event, + on_tool=observer.on_tool, + on_ui=on_ui, + ) + ) + try: + async with asyncio.timeout(8): + await entered.wait() + await client.request("event-burst", count=event_count) + assert await client.request("event-ui") == {"value": "Approve"} + assert observer.events == [] + assert not running.done() + assert process.returncode is None + + release.set() + await delivered.wait() + await client.request("event-result") + await running + + assert [(event.type, event.messageId) for event in observer.events] == [ + (OmpEventType.AGENT_START, None), + *[ + (OmpEventType.MESSAGE_UPDATE, f"burst-{number}") + for number in range(event_count) + ], + (OmpEventType.MESSAGE_END, "complete"), + (OmpEventType.AGENT_END, None), + ] + assert observer.assistant_texts == ["complete"] + await second_prompt(client) + assert client._proc is process + assert process.returncode is None + finally: + running.cancel() + release.set() + await asyncio.gather(running, return_exceptions=True) + + @pytest.mark.parametrize( + "result_seen", + [False, True], + ids=["before-result", "after-result-and-settlement"], + ) + async def test_overload_fails_the_prompt_cancels_delivery_reaps_and_can_restart( + self, client_factory: ClientFactory, result_seen: bool + ) -> None: + client = await client_factory("event-backpressure") + process = client._proc + assert process is not None + observer = Observer() + entered = asyncio.Event() + release = asyncio.Event() + stopped = asyncio.Event() + + async def on_event(event: OmpEvent) -> None: + entered.set() + try: + await release.wait() + finally: + stopped.set() + await observer.on_event(event) + + running = asyncio.create_task( + client.prompt( + "first", + on_event=on_event, + on_tool=observer.on_tool, + on_ui=observer.on_ui, + ) + ) + overflowing: asyncio.Task[dict[str, JsonValue]] | None = None + try: + async with asyncio.timeout(8): + await entered.wait() + if result_seen: + # The peer acknowledges only after the settled state snapshot. + await client.request("event-result") + await client.request( + "event-burst", + count=omprpc._EVENT_BACKLOG_LIMIT - (2 if result_seen else 0), + ) + assert not running.done() + assert not stopped.is_set() + + overflowing = asyncio.create_task( + client.request("event-burst", count=1) + ) + with pytest.raises(OmpProcessError, match="event callback backlog"): + await running + + assert stopped.is_set() + assert observer.events == [] + assert process.returncode is not None + await asyncio.gather(overflowing, return_exceptions=True) + await client.start() + await second_prompt(client) + assert client._proc is not process + finally: + tasks = [running, *([overflowing] if overflowing is not None else [])] + for task in tasks: + task.cancel() + release.set() + await asyncio.gather(*tasks, return_exceptions=True) diff --git a/tests/integrations/omprpc/test_interruption.py b/tests/integrations/omprpc/test_interruption.py index 302c05236..d5300f1f5 100644 --- a/tests/integrations/omprpc/test_interruption.py +++ b/tests/integrations/omprpc/test_interruption.py @@ -7,6 +7,7 @@ import pytest from band.integrations.ompwire import OmpHostToolCall, OmpToolResult +from tests.integrations.omprpc.cancellation import HeldCleanup from tests.integrations.omprpc.kit import ( ClientFactory, Observer, @@ -68,9 +69,12 @@ async def test_interrupted_prompt_aborts_reaps_and_can_restart( with pytest.raises(asyncio.CancelledError): await running else: - with pytest.raises(TimeoutError): - async with asyncio.timeout(0.05): - await running + async with asyncio.timeout(6): + with pytest.raises(TimeoutError): + async with asyncio.timeout(None) as deadline: + deadline.reschedule(asyncio.get_running_loop().time()) + await running + assert deadline.expired() assert process.returncode is not None async with asyncio.timeout(5): await client.start() @@ -83,25 +87,34 @@ async def test_interrupted_prompt_aborts_reaps_and_can_restart( async def test_repeatedly_cancelled_close_still_waits_for_reap( self, client_factory: ClientFactory, + monkeypatch: pytest.MonkeyPatch, ) -> None: client = await client_factory("stubborn") process = client._proc assert process is not None + cleanup = HeldCleanup(client, monkeypatch) closing = asyncio.create_task(client.close()) - loop = asyncio.get_running_loop() - cancellations = [ - loop.call_later(delay, closing.cancel) for delay in (0.01, 0.05, 0.1) - ] try: + await cleanup.wait_for_join(closing) + for _ in range(3): + await cleanup.cancel_and_wait_for_rejoin(closing) + assert client._proc is process + assert process.returncode is None + assert not cleanup.finished.is_set() + + cleanup.release.set() async with asyncio.timeout(6): with pytest.raises(asyncio.CancelledError): await closing + await process.wait() + assert cleanup.finished.is_set() assert process.returncode is not None await client.close() finally: - for cancellation in cancellations: - cancellation.cancel() + cleanup.release.set() + closing.cancel() await asyncio.gather(closing, return_exceptions=True) + await client.close() async def test_cancelled_prompt_cancels_its_host_callback_before_reap( self, diff --git a/tests/integrations/omprpc/test_process_tree.py b/tests/integrations/omprpc/test_process_tree.py index 5f01ee6dc..9521dad2b 100644 --- a/tests/integrations/omprpc/test_process_tree.py +++ b/tests/integrations/omprpc/test_process_tree.py @@ -4,8 +4,6 @@ import asyncio import os -import sys -from contextlib import suppress import psutil import pytest @@ -13,9 +11,9 @@ from band.integrations import omprpc from band.integrations.codex.winspawn import Job -from band.integrations.omprpc import OmpProcessError, OmpRpcClient +from band.integrations.omprpc import OmpProcessError +from tests.integrations.omprpc.cancellation import HeldCleanup from tests.integrations.omprpc.kit import ( - OMP_PEER_SCRIPT, DeferredSessionSignal, ProcessGate, RecycledEnvironment, @@ -223,105 +221,93 @@ async def test_unsettled_session_members_keep_close_retryable( async def test_cancelled_cleanup_reaps_descendants_but_not_unrelated_processes( self, process_gate: ProcessGate, + monkeypatch: pytest.MonkeyPatch, stdio: str, interruption: str, ) -> None: - peer = OMP_PEER_SCRIPT - port = str(process_gate.port) - unrelated = await asyncio.create_subprocess_exec( - sys.executable, - "-u", - peer, - "--scenario", - "gated-child", - "--gate-port", - port, - stdin=asyncio.subprocess.DEVNULL, - stdout=asyncio.subprocess.DEVNULL, - stderr=asyncio.subprocess.DEVNULL, - ) - client = OmpRpcClient( - ( - sys.executable, - "-u", - peer, - "--scenario", - "no-ready-descendants" if interruption == "startup" else "descendants", - "--gate-port", - port, - "--child-stdio", - stdio, - ) + client = descendant_client( + process_gate, + scenario=( + "no-ready-descendants" if interruption == "startup" else "descendants" + ), + stdio=stdio, ) + cleanup = HeldCleanup(client, monkeypatch) starting: asyncio.Task[None] | None = None closing: asyncio.Task[None] | None = None - cancellations: list[asyncio.TimerHandle] = [] - try: - sibling = await process_gate.connection() - unrelated_process = await asyncio.to_thread(psutil.Process, unrelated.pid) - # Windows virtualenv launchers may remain between the parent and Python. - assert sibling.process == unrelated_process or unrelated_process in ( - await asyncio.to_thread(sibling.process.parents) - ) - starting = asyncio.create_task(client.start()) + async with unrelated_gate_child(process_gate) as unrelated: try: - descendants = [ - await process_gate.connection(), - await process_gate.connection(), + starting = asyncio.create_task(client.start()) + try: + descendants = [ + await process_gate.connection(), + await process_gate.connection(), + ] + except TimeoutError: + if starting.done(): + await starting + process = client._proc + pytest.fail( + "Peer did not connect its descendants: " + f"startup_done={starting.done()}, " + f"root_pid={process.pid if process else None}, " + f"returncode={process.returncode if process else None}" + ) + spawning = client._spawn_task + assert spawning is not None + async with asyncio.timeout(5): + root = await asyncio.shield(spawning) + assert client._proc is root + root_process = await asyncio.to_thread(psutil.Process, root.pid) + ancestors = [ + await asyncio.to_thread(process.process.parents) + for process in descendants ] - except TimeoutError: - if starting.done(): - await starting - process = client._proc - pytest.fail( - f"Peer did not connect its descendants: startup_done={starting.done()}, " - f"root_pid={process.pid if process else None}, " - f"returncode={process.returncode if process else None}" + assert all(root_process in parents for parents in ancestors) + assert ( + descendants[0].process in ancestors[1] + or descendants[1].process in ancestors[0] ) - root = client._proc - assert root is not None - root_process = await asyncio.to_thread(psutil.Process, root.pid) - ancestors = [ - await asyncio.to_thread(process.process.parents) - for process in descendants - ] - assert all(root_process in parents for parents in ancestors) - assert ( - descendants[0].process in ancestors[1] - or descendants[1].process in ancestors[0] - ) - if interruption == "startup": - closing = starting - else: - async with asyncio.timeout(5): - await starting - closing = asyncio.create_task(client.close()) - await asyncio.sleep(0) - closing.cancel() - loop = asyncio.get_running_loop() - cancellations = [ - loop.call_later(delay, closing.cancel) for delay in (0.01, 0.05, 0.1) - ] - async with asyncio.timeout(10): - with pytest.raises(asyncio.CancelledError): - await closing - assert root.returncode is not None - for descendant in descendants: - assert not await asyncio.to_thread(process_is_alive, descendant.process) - assert await asyncio.to_thread(process_is_alive, sibling.process) - assert unrelated.returncode is None - await client.close() - finally: - for cancellation in cancellations: - cancellation.cancel() - tasks = {task for task in (starting, closing) if task is not None} - for task in tasks: - task.cancel() - with suppress(ProcessLookupError): - unrelated.kill() - async with asyncio.timeout(5): - await unrelated.wait() - if client._proc is not None and client._proc.returncode is None: - with suppress(ProcessLookupError): - client._proc.kill() - await asyncio.gather(*tasks, return_exceptions=True) + scope = client._process_scope + owner = client._windows_process + assert scope is not None or owner is not None + if interruption == "startup": + closing = starting + assert closing.cancel() + else: + async with asyncio.timeout(5): + await starting + closing = asyncio.create_task(client.close()) + await cleanup.wait_for_join(closing) + if interruption == "close": + await cleanup.cancel_and_wait_for_rejoin(closing) + + for _ in range(3): + await cleanup.cancel_and_wait_for_rejoin(closing) + assert client._proc is root + assert client._process_scope is scope + assert client._windows_process is owner + assert root.returncode is None + assert await asyncio.to_thread(process_is_alive, root_process) + assert await descendant_liveness(descendants) == [True, True] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + assert not cleanup.finished.is_set() + assert closing.cancelling() == 4 + + cleanup.release.set() + async with asyncio.timeout(10): + with pytest.raises(asyncio.CancelledError): + await closing + await root.wait() + assert cleanup.finished.is_set() + assert root.returncode is not None + assert await descendant_liveness(descendants) == [False, False] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + await client.close() + finally: + cleanup.release.set() + tasks = {task for task in (starting, closing) if task is not None} + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await client.close() diff --git a/tests/integrations/test_ompwire.py b/tests/integrations/test_ompwire.py new file mode 100644 index 000000000..012374774 --- /dev/null +++ b/tests/integrations/test_ompwire.py @@ -0,0 +1,74 @@ +"""Structured ask frames preserve native metadata and mutually exclusive answers.""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from band.integrations.ompwire import OmpUIRequest, OmpUIResponse + + +class TestOmpAskWire: + def test_structured_questions_preserve_native_metadata(self) -> None: + question = { + "id": "database-choice", + "question": "Which database should this workspace use?", + "header": "Database", + "options": [ + { + "label": "SQLite", + "description": "A local database", + "preview": "CREATE TABLE workspace (id INTEGER);", + } + ], + "multi": True, + "recommended": 0, + } + request = OmpUIRequest.model_validate( + {"id": "ask-1", "method": "ask", "questions": [question]} + ) + + assert request.model_dump(mode="json", exclude_none=True)["questions"] == [ + question + ] + + @pytest.mark.parametrize("recommended", [-1, 0.5, 8]) + def test_numeric_recommendations_do_not_break_the_reader( + self, recommended: float + ) -> None: + request = OmpUIRequest.model_validate( + { + "id": "ask-2", + "method": "ask", + "questions": [ + { + "id": "question-1", + "question": "Pick a database", + "options": [{"label": "SQLite", "preview": "schema preview"}], + "recommended": recommended, + } + ], + } + ) + assert request.questions[0].recommended == recommended + assert request.questions[0].options[0].preview == "schema preview" + + def test_structured_answers_preserve_an_empty_native_selection(self) -> None: + answers = [{"id": "question-2", "selectedOptions": []}] + response = OmpUIResponse.model_validate({"answers": answers}) + + assert response.model_dump(mode="json", exclude_none=True)["answers"] == answers + + @pytest.mark.parametrize( + "other_answer", [{"value": "SQLite"}, {"confirmed": True}, {"cancelled": True}] + ) + def test_structured_answers_cannot_be_mixed_with_another_ui_variant( + self, other_answer: dict[str, str | bool] + ) -> None: + with pytest.raises(ValidationError, match="only one answer"): + OmpUIResponse.model_validate( + { + "answers": [{"id": "question-3", "selectedOptions": []}], + **other_answer, + } + ) From 0f8e226fc0c2f8e4b594df5d99b0f29bd1ee7a9d Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 09:49:33 +0300 Subject: [PATCH 06/35] fix: serialize structured native answers in the transport --- src/band/integrations/omprpc.py | 4 ++ tests/integrations/omppeer.py | 32 ++++++++++ tests/integrations/omprpc/test_ui_answers.py | 61 ++++++++++++++++++++ 3 files changed, 97 insertions(+) create mode 100644 tests/integrations/omprpc/test_ui_answers.py diff --git a/src/band/integrations/omprpc.py b/src/band/integrations/omprpc.py index 595a7b3f3..59b46d0c6 100644 --- a/src/band/integrations/omprpc.py +++ b/src/band/integrations/omprpc.py @@ -821,6 +821,10 @@ async def _handle_ui(self, turn: PromptRun, request: OmpUIRequest) -> None: frame["value"] = response.value elif response.confirmed is not None: frame["confirmed"] = response.confirmed + elif response.answers is not None: + frame["answers"] = [ + answer.model_dump(exclude_none=True) for answer in response.answers + ] else: frame["cancelled"] = True await self._write_frame(frame) diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index 3b663e04c..d17dac217 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -521,6 +521,22 @@ def handle(self, frame: Frame) -> None: "options": ["Approve", "Deny"], } ) + case "ask-ui": + self.held = frame + send( + { + "type": "extension_ui_request", + "id": "structured-question", + "method": "ask", + "questions": [ + { + "id": "database-choice", + "question": "Which database?", + "options": [{"label": "SQLite"}], + } + ], + } + ) case "release": if self.scenario == "settlement-before-result": self.result(settled=False) @@ -577,6 +593,22 @@ def handle(self, frame: Frame) -> None: ): self.respond(self.held, {"value": frame.value}) self.held = None + if ( + frame.type == "extension_ui_response" + and frame.id == "structured-question" + and self.held is not None + and self.held.type == "ask-ui" + ): + self.respond( + self.held, + { + "answer": cast( + dict[str, JsonValue], + frame.model_dump(exclude_unset=True), + ) + }, + ) + self.held = None if ( self.scenario == "callbacks" and self.tool_results == {"host-1", "host-2"} diff --git a/tests/integrations/omprpc/test_ui_answers.py b/tests/integrations/omprpc/test_ui_answers.py new file mode 100644 index 000000000..110df012e --- /dev/null +++ b/tests/integrations/omprpc/test_ui_answers.py @@ -0,0 +1,61 @@ +"""Structured UI answers reach the native peer without becoming cancellation.""" + +from __future__ import annotations + +import asyncio + +import pytest +from pydantic import JsonValue + +from band.integrations.ompwire import OmpUIRequest, OmpUIResponse +from tests.integrations.omprpc.kit import ClientFactory, Observer + + +class TestStructuredUIAnswers: + @pytest.mark.parametrize( + "answer", + [ + {"id": "database-choice", "selectedOptions": []}, + { + "id": "database-choice", + "selectedOptions": [], + "customInput": "Use the existing database", + }, + ], + ids=["empty-selection", "custom-input"], + ) + async def test_answer_variant_survives_callback_to_wire( + self, client_factory: ClientFactory, answer: dict[str, JsonValue] + ) -> None: + client = await client_factory("event-backpressure") + observer = Observer() + + async def on_ui(request: OmpUIRequest) -> OmpUIResponse: + assert request.method == "ask" + assert request.questions[0].id == answer["id"] + return OmpUIResponse.model_validate({"answers": [answer]}) + + running = asyncio.create_task( + client.prompt( + "first", + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=on_ui, + ) + ) + try: + async with asyncio.timeout(8): + await observer.wait_for("agent_start") + receipt = await client.request("ask-ui") + assert receipt == { + "answer": { + "type": "extension_ui_response", + "id": "structured-question", + "answers": [answer], + } + } + await client.request("event-result") + await running + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) From 13ee3f3421d8f46e49d2527638ec2ba2138cc578 Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 12:38:05 +0300 Subject: [PATCH 07/35] fix: close native RPC review gaps in transport ownership and coverage Fail closed on unmatched host frames and MCP startup UI, log transport failures, and track Windows descendants with the same env-marker scope POSIX already used. Shared process helpers now live with the Windows job helpers instead of the Codex client. Co-authored-by: Cursor --- src/band/converters/omp.py | 2 + src/band/integrations/codex/stdio_client.py | 43 +------ src/band/integrations/codex/winspawn.py | 39 +++++++ src/band/integrations/ompprocess.py | 4 +- src/band/integrations/omprpc.py | 41 ++++--- src/band/integrations/ompwire.py | 5 +- tests/converters/test_omp.py | 17 +++ tests/integrations/codex/test_cleanup.py | 2 +- tests/integrations/omppeer.py | 43 +++++++ tests/integrations/omppermissionpeer.py | 109 ------------------ tests/integrations/ompprocess/test_launch.py | 6 +- .../ompprocess/test_mcp_startup.py | 58 ++++++++++ tests/integrations/omprpc/test_fail_closed.py | 83 +++++++++++++ .../omprpc/test_prompt_results.py | 4 +- 14 files changed, 284 insertions(+), 172 deletions(-) delete mode 100644 tests/integrations/omppermissionpeer.py create mode 100644 tests/integrations/omprpc/test_fail_closed.py diff --git a/src/band/converters/omp.py b/src/band/converters/omp.py index 753a521b1..0e45e54d3 100644 --- a/src/band/converters/omp.py +++ b/src/band/converters/omp.py @@ -27,6 +27,8 @@ class OmpHistoryConverter(HistoryConverter[OmpSessionState]): def convert(self, raw: list[dict[str, Any]]) -> OmpSessionState: sessions: dict[str, str] = {} for message in raw: + if message.get("message_type") != "task": + continue metadata = message.get("metadata") if not isinstance(metadata, dict): continue diff --git a/src/band/integrations/codex/stdio_client.py b/src/band/integrations/codex/stdio_client.py index ec0636225..da9d9be68 100644 --- a/src/band/integrations/codex/stdio_client.py +++ b/src/band/integrations/codex/stdio_client.py @@ -7,7 +7,6 @@ import logging import os import shutil -import signal from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import Any @@ -19,6 +18,7 @@ WindowsProcess, await_returncode, close_stdin, + kill_process_tree, ) logger = logging.getLogger(__name__) @@ -29,32 +29,6 @@ _STDOUT_LIMIT = 16 * 1024 * 1024 -def session_pid_reused(process: asyncio.subprocess.Process) -> bool: - """True when ``process.pid`` now names some other live process. - - After the leader is reaped, its session id stays reserved only while a - member of that session still exists. Once the id is free it can be given - to a new session, and ``killpg`` would signal that session. - """ - if process.returncode is None: - return False - try: - os.kill(process.pid, 0) - except ProcessLookupError: - return False - except OSError: - logger.warning( - "Process %s pid is not signalable; leaving its session", - process.pid, - ) - return True - logger.warning( - "Process %s pid was reused; not signaling that session", - process.pid, - ) - return True - - def _missing_cli_message(binary: str) -> str: return ( f"Codex CLI binary not found: '{binary}'. " @@ -72,21 +46,6 @@ async def _wait_readers(readers: Sequence[asyncio.Task[None]]) -> None: logger.warning("Codex stdio readers still pending after tree kill") -async def kill_process_tree(process: asyncio.subprocess.Process) -> None: - """Force-stop the owned Unix process group, including after root exit. - - The codex launcher is a shim whose real CLI is its child. On Unix the - child leads a session, so ``killpg`` still reaches that group after the - leader has exited. - """ - if os.name == "nt" or session_pid_reused(process): - return - try: - os.killpg(process.pid, signal.SIGKILL) - except ProcessLookupError: - pass - - class CodexStdioClient(BaseJsonRpcClient): """Bidirectional JSON-RPC client with a single stdout read loop.""" diff --git a/src/band/integrations/codex/winspawn.py b/src/band/integrations/codex/winspawn.py index 833ee8fe1..4943999ad 100644 --- a/src/band/integrations/codex/winspawn.py +++ b/src/band/integrations/codex/winspawn.py @@ -13,6 +13,7 @@ import logging import os import shutil +import signal import subprocess import sys import uuid @@ -562,6 +563,44 @@ async def await_returncode(process: asyncio.subprocess.Process, timeout: float) return True +def session_pid_reused(process: asyncio.subprocess.Process) -> bool: + """True when ``process.pid`` now names some other live process. + + After the leader is reaped, its session id stays reserved only while a + member of that session still exists. Once the id is free it can be given + to a new session, and ``killpg`` would signal that session. + """ + if process.returncode is None: + return False + try: + os.kill(process.pid, 0) + except ProcessLookupError: + return False + except OSError: + logger.warning( + "Process %s pid is not signalable; leaving its session", + process.pid, + ) + return True + logger.warning( + "Process %s pid was reused; not signaling that session", + process.pid, + ) + return True + + +async def kill_process_tree(process: asyncio.subprocess.Process) -> None: + """Force-stop the owned Unix process group, including after root exit. + + On Unix the launched child leads a session, so ``killpg`` still reaches + that group after the leader has exited. + """ + if os.name == "nt" or session_pid_reused(process): + return + with suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + + async def close_stdin(process: asyncio.subprocess.Process) -> None: """Drop the write end so a synchronous read in the child can finish. diff --git a/src/band/integrations/ompprocess.py b/src/band/integrations/ompprocess.py index 14f757726..958ccef96 100644 --- a/src/band/integrations/ompprocess.py +++ b/src/band/integrations/ompprocess.py @@ -56,7 +56,7 @@ def omp_runtime_env( return {**credentials, **(env or {})} or None -def omp_command_in_workspace(command: Sequence[str], workspace: str) -> list[str]: +def omp_rpc_command_in_workspace(command: Sequence[str], workspace: str) -> list[str]: """Select OMP's room workspace without changing the subprocess launch cwd.""" return [*command, f"{OMP_CWD_FLAG}={workspace}"] @@ -192,7 +192,7 @@ async def on_tool(_call: OmpHostToolCall) -> OmpToolResult: raise OmpStartupError("OMP unexpectedly invoked a host tool during SDK startup") async def on_ui(_request: OmpUIRequest) -> OmpUIResponse: - return OmpUIResponse(cancelled=True) + raise OmpStartupError("OMP unexpectedly requested UI during SDK startup") await runtime.prompt( f"/{command}", diff --git a/src/band/integrations/omprpc.py b/src/band/integrations/omprpc.py index 59b46d0c6..ececb3cfd 100644 --- a/src/band/integrations/omprpc.py +++ b/src/band/integrations/omprpc.py @@ -16,16 +16,18 @@ import psutil from pydantic import JsonValue, TypeAdapter, ValidationError -from band.integrations.codex.stdio_client import kill_process_tree, session_pid_reused from band.integrations.codex.winspawn import ( SpawnFailure, WindowsProcess, await_returncode, close_stdin, + kill_process_tree, + session_pid_reused, ) from band.integrations.ompwire import ( OMP_CHUNK_BYTES, OMP_FRAME_BYTES, + OMP_NATIVE_PROTOCOL_VERSION, OMP_REASSEMBLED_BYTES, OmpChunk, OmpCommand, @@ -167,10 +169,12 @@ def observe(self, event: OmpEvent, sequence: int) -> None: class OmpProcessScope: - """Retain observable POSIX descendants for best-effort cleanup. + """Retain observable descendants for best-effort cleanup. The private PI_ marker survives native helper filters. Explicit environment replacement and restricted process inspection can hide detached children. + Windows job membership is the primary containment; this marker catches + descendants the job does not observe. """ def __init__(self) -> None: @@ -199,9 +203,10 @@ def capture(self, root: asyncio.subprocess.Process) -> None: except (psutil.NoSuchProcess, psutil.AccessDenied): pass in_group = False - if group_visible: + getpgid = getattr(os, "getpgid", None) + if group_visible and getpgid is not None: try: - in_group = os.getpgid(pid) == root.pid + in_group = getpgid(pid) == root.pid except (ProcessLookupError, PermissionError): pass if not marked and not in_group: @@ -290,12 +295,16 @@ async def start(self) -> None: if self._proc is not None: self._ensure_running() return + if self._failure is not None: + logger.error( + "Discarding previous OMP RPC failure on restart: %s", self._failure + ) self._failure = None self._chunks = None self._max_frame_bytes = OMP_FRAME_BYTES self._max_reassembled_bytes = OMP_REASSEMBLED_BYTES self._closing = False - self._process_scope = OmpProcessScope() if os.name != "nt" else None + self._process_scope = OmpProcessScope() self._ready = asyncio.get_running_loop().create_future() self._spawn_task = asyncio.create_task(self._spawn_process()) try: @@ -308,12 +317,13 @@ async def start(self) -> None: ready = await asyncio.shield(self._ready) if isinstance(ready, OmpProcessError): raise ready - if 2 not in ready.supportedProtocolVersions: + if OMP_NATIVE_PROTOCOL_VERSION not in ready.supportedProtocolVersions: raise OmpProtocolError("OMP RPC protocol v2 is required") self._max_frame_bytes = ready.maxFrameBytes self._max_reassembled_bytes = ready.maxReassembledFrameBytes negotiated = await self._request_frame( - "negotiate_protocol", {"protocolVersion": 2} + "negotiate_protocol", + {"protocolVersion": OMP_NATIVE_PROTOCOL_VERSION}, ) try: OmpNegotiation.model_validate(self._response_data(negotiated.frame)) @@ -329,17 +339,17 @@ async def start(self) -> None: raise async def _spawn_process(self) -> asyncio.subprocess.Process: + scope = self._process_scope + assert scope is not None if os.name == "nt": self._windows_process = WindowsProcess() process = await self._windows_process.start_command( self.command, cwd=None, - env=self.env, + env=scope.environment(self.env), limit=OMP_FRAME_BYTES, ) else: - scope = self._process_scope - assert scope is not None process = await asyncio.create_subprocess_exec( *self.command, stdin=asyncio.subprocess.PIPE, @@ -706,7 +716,7 @@ async def _dispatch_payload( if turn is None: await self._write_frame( { - "type": "extension_ui_response", + "type": OmpFrameType.UI_RESPONSE, "id": request.id, "cancelled": True, } @@ -794,7 +804,7 @@ async def _handle_tool(self, turn: PromptRun, call: OmpHostToolCall) -> None: return await self._write_frame( { - "type": "host_tool_result", + "type": OmpFrameType.HOST_TOOL_RESULT, "id": call.id, "result": cast( dict[str, JsonValue], @@ -814,7 +824,7 @@ async def _handle_ui(self, turn: PromptRun, request: OmpUIRequest) -> None: ): return frame: dict[str, JsonValue] = { - "type": "extension_ui_response", + "type": OmpFrameType.UI_RESPONSE, "id": request.id, } if response.value is not None: @@ -848,6 +858,7 @@ def _fail_pending(self, error: OmpProcessError) -> None: def _record_failure(self, error: OmpProcessError) -> None: if self._failure is None: self._failure = error + logger.error("OMP RPC transport failed: %s", error, exc_info=error) self._fail_pending(self._failure) if not self._closing: self._begin_close() @@ -974,6 +985,10 @@ async def _close_process(self) -> None: if owner is not None: if not owner_closed: await owner.close(_PROCESS_EXIT_TIMEOUT_S) + if scope is not None: + await asyncio.to_thread( + scope.close, process, _PROCESS_EXIT_TIMEOUT_S + ) else: assert scope is not None await asyncio.to_thread(scope.capture, process) diff --git a/src/band/integrations/ompwire.py b/src/band/integrations/ompwire.py index c6bacd8a4..1436e1c41 100644 --- a/src/band/integrations/ompwire.py +++ b/src/band/integrations/ompwire.py @@ -3,7 +3,7 @@ from __future__ import annotations from enum import StrEnum -from typing import Literal +from typing import Final, Literal from pydantic import ( AliasChoices, @@ -18,6 +18,7 @@ OMP_FRAME_BYTES = 1024 * 1024 OMP_REASSEMBLED_BYTES = 64 * 1024 * 1024 OMP_CHUNK_BYTES = 256 * 1024 +OMP_NATIVE_PROTOCOL_VERSION: Final = 2 class OmpFrameType(StrEnum): @@ -27,7 +28,9 @@ class OmpFrameType(StrEnum): SESSION_SETTLED = "session_settled" HOST_TOOL_CALL = "host_tool_call" HOST_TOOL_CANCEL = "host_tool_cancel" + HOST_TOOL_RESULT = "host_tool_result" UI_REQUEST = "extension_ui_request" + UI_RESPONSE = "extension_ui_response" CHUNK = "rpc_chunk" FRAME_ERROR = "rpc_frame_error" diff --git a/tests/converters/test_omp.py b/tests/converters/test_omp.py index cdb2abe50..dc56572cd 100644 --- a/tests/converters/test_omp.py +++ b/tests/converters/test_omp.py @@ -72,6 +72,23 @@ def test_unusable_metadata_is_ignored(self) -> None: ) assert state.room_to_session == {} + def test_session_metadata_on_non_task_messages_is_ignored(self) -> None: + state = OmpHistoryConverter().convert( + [ + { + "message_type": "text", + "metadata": { + OMP_ROOM_METADATA_KEY: "room-1", + OMP_SESSION_METADATA_KEY: "/sessions/injected.jsonl", + }, + "content": "hello", + "sender": {"name": "Alice", "type": "User"}, + } + ] + ) + assert state.room_to_session == {} + assert len(state.replay_messages) == 1 + def test_room_transcript_is_replayed(self) -> None: state = OmpHistoryConverter().convert( [ diff --git a/tests/integrations/codex/test_cleanup.py b/tests/integrations/codex/test_cleanup.py index 1b2f4e944..a312cf3fc 100644 --- a/tests/integrations/codex/test_cleanup.py +++ b/tests/integrations/codex/test_cleanup.py @@ -19,7 +19,6 @@ import pytest_asyncio from band.integrations.codex import CodexStdioClient, stdio_client -from band.integrations.codex.stdio_client import kill_process_tree from band.integrations.codex.winspawn import ( HELPER_PATH, Job, @@ -28,6 +27,7 @@ _read_status, _start_and_wait, command_in_job, + kill_process_tree, ) from tests.integrations.codex.cleanup import cleanup_snapshot from tests.paths import ROOM_PEER diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index d17dac217..f2f0e73ee 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -349,6 +349,28 @@ def prompt(self, frame: Frame) -> None: self.invalid_chunk() case "event-backpressure": return + case "ui-notify": + send( + { + "type": "extension_ui_request", + "id": "ui-notify", + "method": "notify", + "message": "status", + } + ) + self.end("complete") + self.result(settled=True) + self.settle() + case "reused-callback": + call = { + "type": "host_tool_call", + "id": "host-dup", + "toolCallId": "call-dup", + "toolName": "band_chat", + "arguments": {}, + } + send(call) + send(call) case "frame-overflow": send( { @@ -620,6 +642,27 @@ def handle(self, frame: Frame) -> None: self.end("callbacks complete") self.result(settled=True) self.settle() + case "orphan-tool": + send( + { + "type": "host_tool_call", + "id": "host-orphan", + "toolCallId": "call-orphan", + "toolName": "band_chat", + "arguments": {}, + } + ) + case "orphan-ui": + send( + { + "type": "extension_ui_request", + "id": "ui-orphan", + "method": "select", + "title": "Allow command?", + "options": ["Approve", "Deny"], + } + ) + self.respond(frame) case "checkpoint": self.respond(frame, {"pid": os.getpid(), "frames": self.frames}) case "abort": diff --git a/tests/integrations/omppermissionpeer.py b/tests/integrations/omppermissionpeer.py deleted file mode 100644 index 6cd4c44d1..000000000 --- a/tests/integrations/omppermissionpeer.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Script independent UI decisions and native events over the existing RPC peer.""" - -from __future__ import annotations - -import json -import os -import sys - -from tests.integrations.omppeer import Frame, Peer, send - - -class PermissionPeer(Peer): - def __init__(self, order: str) -> None: - super().__init__("permission-records") - self.order = order - - @staticmethod - def native_start() -> None: - send( - { - "type": "tool_execution_start", - "toolCallId": "ui-wire-id", - "toolName": os.environ.get("OMP_TEST_NATIVE_NAME", "read"), - "args": json.loads( - os.environ.get("OMP_TEST_NATIVE_ARGUMENTS", '{"path":"denied.txt"}') - ), - } - ) - - @staticmethod - def native_end() -> None: - send( - { - "type": "tool_execution_end", - "toolCallId": "ui-wire-id", - "toolName": os.environ.get("OMP_TEST_NATIVE_NAME", "read"), - "isError": os.environ.get("OMP_TEST_NATIVE_ERROR", "1") == "1", - "result": { - "content": [ - { - "type": "text", - "text": ( - "Native tool failed" - if os.environ.get("OMP_TEST_NATIVE_ERROR", "1") == "1" - else "Native tool finished" - ), - } - ] - }, - } - ) - - @staticmethod - def ask_ui() -> None: - send( - { - "type": "extension_ui_request", - "id": "ui-wire-id", - "method": "select", - "title": os.environ["OMP_TEST_PERMISSION_TITLE"], - "options": ["Approve", "Deny"], - } - ) - - def prompt(self, frame: Frame) -> None: - self.prompt_id = frame.id - self.settled = False - self.respond(frame) - send({"type": "agent_start"}) - if self.order == "native-first": - self.native_start() - else: - self.ask_ui() - - def handle(self, frame: Frame) -> None: - if frame.type == "permission-after-native": - self.respond(frame) - self.ask_ui() - return - if frame.type == "extension_ui_response": - if frame.id != "ui-wire-id" or not frame.cancelled: - raise AssertionError("The scripted UI decision must be cancelled") - if self.order == "permission-first": - self.native_start() - self.native_end() - self.end("") - self.result(settled=True) - self.settle() - return - super().handle(frame) - - -def main() -> None: - peer = PermissionPeer(sys.argv[1]) - send( - { - "type": "ready", - "protocolVersion": 1, - "supportedProtocolVersions": [1, 2], - "maxFrameBytes": 1024 * 1024, - "maxReassembledFrameBytes": 64 * 1024 * 1024, - } - ) - for line in sys.stdin: - peer.handle(Frame.model_validate_json(line, strict=True)) - - -if __name__ == "__main__": - main() diff --git a/tests/integrations/ompprocess/test_launch.py b/tests/integrations/ompprocess/test_launch.py index 4958592be..8ddd687f1 100644 --- a/tests/integrations/ompprocess/test_launch.py +++ b/tests/integrations/ompprocess/test_launch.py @@ -9,7 +9,7 @@ ) from band.integrations.ompprocess import ( DEFAULT_OMP_COMMAND, - omp_command_in_workspace, + omp_rpc_command_in_workspace, omp_runtime_env, ) @@ -22,7 +22,9 @@ def test_default_command_is_safe_after_finalize(self) -> None: ] def test_workspace_is_selected_without_a_subprocess_cwd(self) -> None: - assert omp_command_in_workspace(("bunx", *DEFAULT_OMP_COMMAND, "-v"), "/w") == [ + assert omp_rpc_command_in_workspace( + ("bunx", *DEFAULT_OMP_COMMAND, "-v"), "/w" + ) == [ "bunx", *DEFAULT_OMP_COMMAND, "-v", diff --git a/tests/integrations/ompprocess/test_mcp_startup.py b/tests/integrations/ompprocess/test_mcp_startup.py index f25a91268..4d4d8ba1b 100644 --- a/tests/integrations/ompprocess/test_mcp_startup.py +++ b/tests/integrations/ompprocess/test_mcp_startup.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from collections.abc import Awaitable, Callable from pathlib import Path from unittest.mock import AsyncMock @@ -15,6 +16,13 @@ omp_mcp_ready_command, ) from band.integrations.omprpc import OmpRpcClient +from band.integrations.ompwire import ( + OmpEvent, + OmpHostToolCall, + OmpToolResult, + OmpUIRequest, + OmpUIResponse, +) class TestRequiredMcpStartup: @@ -105,3 +113,53 @@ async def test_missing_startup_extension_cannot_fall_through_to_inference( with pytest.raises(OmpStartupError, match="required SDK startup extension"): await omp_mcp_ready(runtime, str(tmp_path), {}) runtime.prompt.assert_not_awaited() + + async def test_complete_receipt_prompts_the_readiness_command( + self, tmp_path: Path + ) -> None: + plugin = tmp_path / "band-sdk" + plugin.mkdir() + command = omp_mcp_ready_command(str(plugin)) + (plugin / "receipt.json").write_text( + json.dumps( + { + "systemPromptReady": True, + "servers": [ + {"type": "connected", "serverName": "band-sdk:required"} + ], + } + ), + encoding="utf-8", + ) + runtime = AsyncMock(spec=OmpRpcClient) + runtime.request.return_value = {"commands": [{"name": command}]} + await omp_mcp_ready( + runtime, + str(plugin), + {"required": OmpMcpServer(command="server")}, + ) + runtime.prompt.assert_awaited_once() + assert runtime.prompt.await_args is not None + assert runtime.prompt.await_args.args[0] == f"/{command}" + + async def test_ui_request_during_startup_cannot_fall_through_to_inference( + self, tmp_path: Path + ) -> None: + runtime = AsyncMock(spec=OmpRpcClient) + runtime.request.return_value = { + "commands": [{"name": omp_mcp_ready_command(str(tmp_path))}], + } + + async def prompt( + message: str, + *, + on_event: Callable[[OmpEvent], Awaitable[None]], + on_tool: Callable[[OmpHostToolCall], Awaitable[OmpToolResult]], + on_ui: Callable[[OmpUIRequest], Awaitable[OmpUIResponse]], + ) -> None: + _ = (message, on_event, on_tool) + await on_ui(OmpUIRequest.model_validate({"id": "ask", "method": "select"})) + + runtime.prompt.side_effect = prompt + with pytest.raises(OmpStartupError, match="requested UI"): + await omp_mcp_ready(runtime, str(tmp_path), {}) diff --git a/tests/integrations/omprpc/test_fail_closed.py b/tests/integrations/omprpc/test_fail_closed.py new file mode 100644 index 000000000..9e20bb5c0 --- /dev/null +++ b/tests/integrations/omprpc/test_fail_closed.py @@ -0,0 +1,83 @@ +"""Dispatcher fail-closed paths keep OMP from hanging on unmatched host frames.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.omprpc import OmpProtocolError +from band.integrations.ompwire import OmpUIRequest, OmpUIResponse +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + checkpoint_frames, + prompt_task, +) + + +class TestFailClosedDispatch: + async def test_host_tool_without_an_active_prompt_fails_the_transport( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("basic") + async with asyncio.timeout(6): + with pytest.raises(OmpProtocolError, match="without an active prompt"): + await client.request("orphan-tool") + + async def test_interactive_ui_without_an_active_prompt_is_cancelled( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("basic") + async with asyncio.timeout(5): + await client.request("orphan-ui") + frames = await checkpoint_frames(client) + assert any( + frame.get("type") == "extension_ui_response" + and frame.get("id") == "ui-orphan" + and frame.get("cancelled") is True + for frame in frames + ) + + async def test_ui_notifications_do_not_start_a_callback( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("ui-notify") + observer = Observer() + seen: list[OmpUIRequest] = [] + + async def on_ui(request: OmpUIRequest) -> OmpUIResponse: + seen.append(request) + return await observer.on_ui(request) + + async with asyncio.timeout(5): + await client.prompt( + "first", + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=on_ui, + ) + assert seen == [] + + async def test_reused_callback_id_fails_the_transport( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("reused-callback") + observer = Observer() + async with asyncio.timeout(6): + with pytest.raises( + OmpProtocolError, match="callback request id was reused" + ): + await prompt_task(client, observer) + + async def test_request_rejects_id_and_type_overrides( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("basic") + with pytest.raises(ValueError, match="must not override"): + await client.request("checkpoint", id="stolen") + with pytest.raises(ValueError, match="must not override"): + await client.request("checkpoint", type="prompt") + async with asyncio.timeout(5): + checkpoint = await client.request("checkpoint") + assert "pid" in checkpoint diff --git a/tests/integrations/omprpc/test_prompt_results.py b/tests/integrations/omprpc/test_prompt_results.py index 787376864..b95aaaa5a 100644 --- a/tests/integrations/omprpc/test_prompt_results.py +++ b/tests/integrations/omprpc/test_prompt_results.py @@ -7,7 +7,7 @@ import pytest from band.integrations.omprpc import OmpPromptError, OmpRpcError -from band.integrations.ompwire import OmpEvent, OmpState +from band.integrations.ompwire import OMP_NATIVE_PROTOCOL_VERSION, OmpEvent, OmpState from tests.integrations.omprpc.kit import ( ClientFactory, Observer, @@ -139,4 +139,4 @@ async def test_native_v2_preserves_large_inline_image_frames( if frame["type"] == "negotiate_protocol" ] assert len(negotiations) == 1 - assert negotiations[0]["protocolVersion"] == 2 + assert negotiations[0]["protocolVersion"] == OMP_NATIVE_PROTOCOL_VERSION From 28ba70ae8b7f3a4b778c1dc0bc255a4e67aea85f Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 12:46:44 +0300 Subject: [PATCH 08/35] refactor: compare MCP startup events to OmpEventType Co-authored-by: Cursor --- src/band/integrations/ompprocess.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/band/integrations/ompprocess.py b/src/band/integrations/ompprocess.py index 958ccef96..77b0b26c0 100644 --- a/src/band/integrations/ompprocess.py +++ b/src/band/integrations/ompprocess.py @@ -22,6 +22,7 @@ from band.integrations.omprpc import OmpRpcClient from band.integrations.ompwire import ( OmpEvent, + OmpEventType, OmpHostToolCall, OmpToolResult, OmpUIRequest, @@ -183,7 +184,7 @@ async def omp_mcp_ready( raise OmpStartupError("OMP did not load the required SDK startup extension") async def on_event(event: OmpEvent) -> None: - if event.type == "agent_start": + if event.type == OmpEventType.AGENT_START: raise OmpStartupError( "OMP unexpectedly invoked an agent during SDK startup" ) From 73073c638f62978503d65f9ce3cf16fd9ffe9924 Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 13:04:12 +0300 Subject: [PATCH 09/35] refactor: share native RPC parse helpers and protocol constants Co-authored-by: Cursor --- src/band/integrations/omprpc.py | 131 +++++++++--------- tests/integrations/omppeer.py | 36 +++-- tests/integrations/ompprocess/test_launch.py | 3 +- tests/integrations/omprpc/kit.py | 3 +- tests/integrations/omprpc/lifecycle.py | 3 +- tests/integrations/omprpc/test_callbacks.py | 5 +- tests/integrations/omprpc/test_fail_closed.py | 4 +- .../integrations/omprpc/test_interruption.py | 6 +- .../omprpc/test_prompt_results.py | 21 ++- .../omprpc/test_turn_settlement.py | 24 ++-- tests/integrations/omprpc/test_ui_answers.py | 11 +- 11 files changed, 136 insertions(+), 111 deletions(-) diff --git a/src/band/integrations/omprpc.py b/src/band/integrations/omprpc.py index ececb3cfd..f89992469 100644 --- a/src/band/integrations/omprpc.py +++ b/src/band/integrations/omprpc.py @@ -10,7 +10,7 @@ from binascii import Error as Base64Error from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass -from typing import cast +from typing import TypeVar, cast from uuid import uuid4 import psutil @@ -48,8 +48,11 @@ OmpToolResult, OmpUIRequest, OmpUIResponse, + OmpWireModel, ) +TModel = TypeVar("TModel", bound=OmpWireModel) + logger = logging.getLogger(__name__) _DRAIN_BYTES = 64 * 1024 @@ -91,6 +94,14 @@ def __init__(self, status: str, message: str) -> None: super().__init__(message) +def _command_error(response: OmpResponse) -> OmpRpcError: + return OmpRpcError( + response.command, + response.error or "Command failed", + code=response.code, + ) + + @dataclass class PendingChunks: chunk_id: str @@ -154,17 +165,20 @@ def observe(self, event: OmpEvent, sequence: int) -> None: self.fail(OmpProcessError("OMP event callback backlog exceeded")) return self.last_activity_sequence = sequence - if ( - event.type == OmpEventType.MESSAGE_END - and event.message is not None - and event.message.role == "assistant" - ): - self.last_assistant = event.message - if event.type == OmpEventType.AGENT_END: - for message in reversed(event.messages): - if message.role == "assistant": - self.last_assistant = message - break + match event.type: + case OmpEventType.MESSAGE_END if ( + event.message is not None and event.message.role == "assistant" + ): + self.last_assistant = event.message + case OmpEventType.AGENT_END: + self.last_assistant = next( + ( + message + for message in reversed(event.messages) + if message.role == "assistant" + ), + self.last_assistant, + ) self.changed.set() @@ -190,6 +204,7 @@ def environment(self, base: Mapping[str, str] | None) -> dict[str, str]: def capture(self, root: asyncio.subprocess.Process) -> None: group_visible = not session_pid_reused(root) group_members: list[psutil.Process] = [] + getpgid = getattr(os, "getpgid", None) for pid in psutil.pids(): # asyncio's child watcher remains the sole waiter for the root. if pid == 0 or pid == root.pid: @@ -203,7 +218,6 @@ def capture(self, root: asyncio.subprocess.Process) -> None: except (psutil.NoSuchProcess, psutil.AccessDenied): pass in_group = False - getpgid = getattr(os, "getpgid", None) if group_visible and getpgid is not None: try: in_group = getpgid(pid) == root.pid @@ -325,12 +339,11 @@ async def start(self) -> None: "negotiate_protocol", {"protocolVersion": OMP_NATIVE_PROTOCOL_VERSION}, ) - try: - OmpNegotiation.model_validate(self._response_data(negotiated.frame)) - except ValidationError as cause: - raise OmpProtocolError( - "OMP RPC protocol v2 negotiation failed" - ) from cause + self._parse_response( + negotiated.frame, + OmpNegotiation, + "OMP RPC protocol v2 negotiation failed", + ) except (OSError, SpawnFailure) as cause: await self.close() raise OmpProcessError("Unable to start OMP RPC process") from cause @@ -453,6 +466,14 @@ def _response_data(self, response: OmpResponse) -> dict[str, JsonValue]: raise error return response.data + def _parse_response( + self, response: OmpResponse, model: type[TModel], error: str + ) -> TModel: + try: + return model.model_validate(self._response_data(response)) + except ValidationError as cause: + raise OmpProtocolError(error) from cause + async def _execute_prompt( self, message: str, @@ -491,14 +512,11 @@ async def _complete_prompt(self, turn: PromptRun, message: str) -> None: except OmpRpcError: turn.rejected = True raise - try: - acceptance = OmpPromptAcceptance.model_validate( - self._response_data(accepted.frame) - ) - except ValidationError as cause: - raise OmpProtocolError( - "Invalid OMP prompt acceptance protocol data" - ) from cause + acceptance = self._parse_response( + accepted.frame, + OmpPromptAcceptance, + "Invalid OMP prompt acceptance protocol data", + ) outcome: OmpPromptResult | None = None # A synchronous local command explicitly owes no prompt_result. if acceptance.agentInvoked is not False: @@ -512,25 +530,24 @@ async def _complete_prompt(self, turn: PromptRun, message: str) -> None: if pending: raise OmpProcessError("OMP callbacks did not finish after settlement") if outcome is not None and outcome.status != "completed": - message = ( + raise OmpPromptError( + outcome.status, outcome.error.message if outcome.error is not None - else f"OMP prompt {outcome.status}" + else f"OMP prompt {outcome.status}", ) - raise OmpPromptError(outcome.status, message) assistant = turn.last_assistant if assistant is not None and assistant.stopReason in {"error", "aborted"}: - message = ( - assistant.errorMessage or f"OMP continuation {assistant.stopReason}" + raise OmpPromptError( + assistant.stopReason, + assistant.errorMessage or f"OMP continuation {assistant.stopReason}", ) - raise OmpPromptError(assistant.stopReason, message) async def _wait_for_settlement(self, turn: PromptRun) -> None: snapshot = await self._request_frame("get_state", {}) - try: - state = OmpState.model_validate(self._response_data(snapshot.frame)) - except ValidationError as cause: - raise OmpProtocolError("Invalid OMP settlement protocol data") from cause + state = self._parse_response( + snapshot.frame, OmpState, "Invalid OMP settlement protocol data" + ) if state.isSettled and snapshot.sequence >= turn.last_activity_sequence: return # An old, uncorrelated session_settled cannot overrule a newer busy snapshot. @@ -646,13 +663,9 @@ async def _read_stdout(self, process: asyncio.subprocess.Process) -> None: await self._dispatch_payload(payload, self._sequence) except asyncio.CancelledError: raise - except ValidationError as cause: - error = OmpProtocolError("Invalid OMP RPC protocol frame") - error.__cause__ = cause - self._record_failure(error) except OmpProcessError as error: self._record_failure(error) - except (ValueError, Base64Error) as cause: + except (ValidationError, ValueError, Base64Error) as cause: error = OmpProtocolError("Invalid OMP RPC protocol frame") error.__cause__ = cause self._record_failure(error) @@ -740,13 +753,7 @@ def _dispatch_response(self, response: OmpResponse, sequence: int) -> None: and self._turn is not None and response.id == self._turn.id ): - self._turn.fail( - OmpRpcError( - response.command, - response.error or "Command failed", - code=response.code, - ) - ) + self._turn.fail(_command_error(response)) return if response.command != pending.command: raise OmpProtocolError("OMP response command does not match its request") @@ -755,13 +762,7 @@ def _dispatch_response(self, response: OmpResponse, sequence: int) -> None: if response.success: pending.future.set_result(ReceivedResponse(response, sequence)) else: - pending.future.set_result( - OmpRpcError( - response.command, - response.error or "Command failed", - code=response.code, - ) - ) + pending.future.set_result(_command_error(response)) def _start_callback( self, @@ -779,6 +780,14 @@ def _start_callback( turn.callback_tasks.add(task) task.add_done_callback(lambda done: self._callback_finished(turn, key, done)) + def _callback_still_active( + self, turn: PromptRun, kind: str, request_id: str + ) -> bool: + return ( + not self._closing + and turn.callbacks.get((kind, request_id)) is asyncio.current_task() + ) + @staticmethod def _cancel_callback(turn: PromptRun, kind: str, request_id: str) -> None: task = turn.callbacks.pop((kind, request_id), None) @@ -797,10 +806,7 @@ def _callback_finished( async def _handle_tool(self, turn: PromptRun, call: OmpHostToolCall) -> None: result = await turn.on_tool(call) - if ( - self._closing - or turn.callbacks.get(("tool", call.id)) is not asyncio.current_task() - ): + if not self._callback_still_active(turn, "tool", call.id): return await self._write_frame( { @@ -818,10 +824,7 @@ async def _handle_tool(self, turn: PromptRun, call: OmpHostToolCall) -> None: async def _handle_ui(self, turn: PromptRun, request: OmpUIRequest) -> None: response = await turn.on_ui(request) - if ( - self._closing - or turn.callbacks.get(("ui", request.id)) is not asyncio.current_task() - ): + if not self._callback_still_active(turn, "ui", request.id): return frame: dict[str, JsonValue] = { "type": OmpFrameType.UI_RESPONSE, diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index f2f0e73ee..54f3279ee 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -16,9 +16,14 @@ import psutil from pydantic import BaseModel, ConfigDict, Field, JsonValue +from band.integrations.ompwire import ( + OMP_CHUNK_BYTES, + OMP_FRAME_BYTES, + OMP_NATIVE_PROTOCOL_VERSION, + OMP_REASSEMBLED_BYTES, +) + PROTOCOL_VERSION = 1 -FRAME_BYTES = 1024 * 1024 -CHUNK_BYTES = 256 * 1024 class Frame(BaseModel): @@ -107,8 +112,11 @@ def launch_child() -> subprocess.Popen[bytes] | None: def send(frame: dict[str, JsonValue]) -> None: payload = json.dumps(frame, separators=(",", ":")).encode("utf-8") - if PROTOCOL_VERSION == 2 and len(payload) + 1 > FRAME_BYTES: - count = (len(payload) + CHUNK_BYTES - 1) // CHUNK_BYTES + if ( + PROTOCOL_VERSION == OMP_NATIVE_PROTOCOL_VERSION + and len(payload) + 1 > OMP_FRAME_BYTES + ): + count = (len(payload) + OMP_CHUNK_BYTES - 1) // OMP_CHUNK_BYTES for index in range(count): chunk: dict[str, JsonValue] = { "type": "rpc_chunk", @@ -117,7 +125,7 @@ def send(frame: dict[str, JsonValue]) -> None: "count": count, "byteLength": len(payload), "data": base64.b64encode( - payload[index * CHUNK_BYTES : (index + 1) * CHUNK_BYTES] + payload[index * OMP_CHUNK_BYTES : (index + 1) * OMP_CHUNK_BYTES] ).decode("ascii"), } sys.stdout.write(json.dumps(chunk, separators=(",", ":")) + "\n") @@ -229,12 +237,12 @@ def invalid_chunk(self) -> None: "chunkId": "invalid-frame", "index": 0, "count": 2, - "byteLength": 1024 * 1024, + "byteLength": OMP_FRAME_BYTES, "data": "e30=", } match self.scenario: case "chunk-overflow": - chunk["byteLength"] = 64 * 1024 * 1024 + 1 + chunk["byteLength"] = OMP_REASSEMBLED_BYTES + 1 send(chunk) case "chunk-order": chunk["index"] = 1 @@ -482,9 +490,13 @@ def handle(self, frame: Frame) -> None: return self.respond( frame, - {"protocolVersion": 1 if self.scenario == "bad-negotiation" else 2}, + { + "protocolVersion": 1 + if self.scenario == "bad-negotiation" + else OMP_NATIVE_PROTOCOL_VERSION + }, ) - PROTOCOL_VERSION = 2 + PROTOCOL_VERSION = OMP_NATIVE_PROTOCOL_VERSION case "prompt": self.prompt(frame) case "get_state": @@ -746,9 +758,9 @@ def main() -> None: "protocolVersion": 1, "supportedProtocolVersions": [1] if args.scenario == "no-v2" - else [1, 2], - "maxFrameBytes": 1024 * 1024, - "maxReassembledFrameBytes": 64 * 1024 * 1024, + else [1, OMP_NATIVE_PROTOCOL_VERSION], + "maxFrameBytes": OMP_FRAME_BYTES, + "maxReassembledFrameBytes": OMP_REASSEMBLED_BYTES, } ) for line in sys.stdin: diff --git a/tests/integrations/ompprocess/test_launch.py b/tests/integrations/ompprocess/test_launch.py index 8ddd687f1..cd247e231 100644 --- a/tests/integrations/ompprocess/test_launch.py +++ b/tests/integrations/ompprocess/test_launch.py @@ -5,6 +5,7 @@ from band.integrations.omp import ( OMP_APPROVAL_MODE_ALWAYS_ASK, OMP_APPROVAL_MODE_FLAG, + OMP_CWD_FLAG, finalize_omp_command, ) from band.integrations.ompprocess import ( @@ -28,7 +29,7 @@ def test_workspace_is_selected_without_a_subprocess_cwd(self) -> None: "bunx", *DEFAULT_OMP_COMMAND, "-v", - "--cwd=/w", + f"{OMP_CWD_FLAG}=/w", ] def test_runtime_env_routes_the_selected_provider_key(self) -> None: diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py index 9dcb8de75..7a08d736f 100644 --- a/tests/integrations/omprpc/kit.py +++ b/tests/integrations/omprpc/kit.py @@ -20,6 +20,7 @@ from band.integrations.ompwire import ( OmpContent, OmpEvent, + OmpEventType, OmpHostToolCall, OmpToolResult, OmpUIRequest, @@ -61,7 +62,7 @@ async def on_tool(self, call: OmpHostToolCall) -> OmpToolResult: async def on_ui(self, request: OmpUIRequest) -> OmpUIResponse: return OmpUIResponse(cancelled=True) - async def wait_for(self, event_type: str) -> OmpEvent: + async def wait_for(self, event_type: OmpEventType) -> OmpEvent: async with asyncio.timeout(5): while True: event = await self.queue.get() diff --git a/tests/integrations/omprpc/lifecycle.py b/tests/integrations/omprpc/lifecycle.py index 67ee5569a..27761b125 100644 --- a/tests/integrations/omprpc/lifecycle.py +++ b/tests/integrations/omprpc/lifecycle.py @@ -11,6 +11,7 @@ from hypothesis import strategies as st from band.integrations.omprpc import OmpProcessError +from band.integrations.ompwire import OmpEventType from tests.integrations.omprpc.kit import ( GatedProcess, Observer, @@ -101,7 +102,7 @@ async def hold_prompt(self) -> None: self.observer = Observer() self.active = prompt_task(self.client, self.observer) self.tasks.append(self.active) - await self.observer.wait_for("agent_end") + await self.observer.wait_for(OmpEventType.AGENT_END) assert not self.active.done() assert (await self.client.request("get_state"))["isSettled"] is False diff --git a/tests/integrations/omprpc/test_callbacks.py b/tests/integrations/omprpc/test_callbacks.py index f174a091c..96f8c3138 100644 --- a/tests/integrations/omprpc/test_callbacks.py +++ b/tests/integrations/omprpc/test_callbacks.py @@ -8,6 +8,7 @@ from band.integrations.ompwire import ( OmpContent, + OmpFrameType, OmpHostToolCall, OmpToolResult, OmpUIRequest, @@ -66,7 +67,7 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: frames = await checkpoint_frames(client) result = next(frame for frame in frames if frame.get("id") == "host-2") assert result == { - "type": "host_tool_result", + "type": OmpFrameType.HOST_TOOL_RESULT, "id": "host-2", "result": { "content": [ @@ -79,7 +80,7 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: } ui_response = next(frame for frame in frames if frame.get("id") == "approval-1") assert ui_response == { - "type": "extension_ui_response", + "type": OmpFrameType.UI_RESPONSE, "id": "approval-1", "value": "Approve", } diff --git a/tests/integrations/omprpc/test_fail_closed.py b/tests/integrations/omprpc/test_fail_closed.py index 9e20bb5c0..435b7d354 100644 --- a/tests/integrations/omprpc/test_fail_closed.py +++ b/tests/integrations/omprpc/test_fail_closed.py @@ -7,7 +7,7 @@ import pytest from band.integrations.omprpc import OmpProtocolError -from band.integrations.ompwire import OmpUIRequest, OmpUIResponse +from band.integrations.ompwire import OmpFrameType, OmpUIRequest, OmpUIResponse from tests.integrations.omprpc.kit import ( ClientFactory, Observer, @@ -33,7 +33,7 @@ async def test_interactive_ui_without_an_active_prompt_is_cancelled( await client.request("orphan-ui") frames = await checkpoint_frames(client) assert any( - frame.get("type") == "extension_ui_response" + frame.get("type") == OmpFrameType.UI_RESPONSE and frame.get("id") == "ui-orphan" and frame.get("cancelled") is True for frame in frames diff --git a/tests/integrations/omprpc/test_interruption.py b/tests/integrations/omprpc/test_interruption.py index d5300f1f5..6adda260c 100644 --- a/tests/integrations/omprpc/test_interruption.py +++ b/tests/integrations/omprpc/test_interruption.py @@ -6,7 +6,7 @@ import pytest -from band.integrations.ompwire import OmpHostToolCall, OmpToolResult +from band.integrations.ompwire import OmpEventType, OmpHostToolCall, OmpToolResult from tests.integrations.omprpc.cancellation import HeldCleanup from tests.integrations.omprpc.kit import ( ClientFactory, @@ -29,7 +29,7 @@ async def test_cancelling_a_queued_prompt_preserves_the_active_and_next_prompt( active = prompt_task(client, observer) queued: asyncio.Task[None] | None = None try: - await observer.wait_for("agent_end") + await observer.wait_for(OmpEventType.AGENT_END) queued = await queue_prompt(client) queued.cancel() with pytest.raises(asyncio.CancelledError): @@ -62,7 +62,7 @@ async def test_interrupted_prompt_aborts_reaps_and_can_restart( observer = Observer() running = prompt_task(client, observer) try: - await observer.wait_for("agent_end") + await observer.wait_for(OmpEventType.AGENT_END) if interruption == "cancel": running.cancel() async with asyncio.timeout(6): diff --git a/tests/integrations/omprpc/test_prompt_results.py b/tests/integrations/omprpc/test_prompt_results.py index b95aaaa5a..92f4983f0 100644 --- a/tests/integrations/omprpc/test_prompt_results.py +++ b/tests/integrations/omprpc/test_prompt_results.py @@ -7,7 +7,12 @@ import pytest from band.integrations.omprpc import OmpPromptError, OmpRpcError -from band.integrations.ompwire import OMP_NATIVE_PROTOCOL_VERSION, OmpEvent, OmpState +from band.integrations.ompwire import ( + OMP_NATIVE_PROTOCOL_VERSION, + OmpEvent, + OmpEventType, + OmpState, +) from tests.integrations.omprpc.kit import ( ClientFactory, Observer, @@ -27,11 +32,11 @@ async def test_prompt_delivers_typed_narration_and_ignores_other_events( async with asyncio.timeout(5): await prompt_task(client, observer) assert [event.type for event in observer.events] == [ - "agent_start", - "tool_execution_start", - "tool_execution_end", - "message_end", - "agent_end", + OmpEventType.AGENT_START, + OmpEventType.TOOL_START, + OmpEventType.TOOL_END, + OmpEventType.MESSAGE_END, + OmpEventType.AGENT_END, ] tool_start, tool_end, message_end = observer.events[1:4] assert (tool_start.toolCallId, tool_start.toolName, tool_start.args) == ( @@ -127,7 +132,9 @@ async def test_native_v2_preserves_large_inline_image_frames( async with asyncio.timeout(5): await prompt_task(client, observer) message = next( - event.message for event in observer.events if event.type == "message_end" + event.message + for event in observer.events + if event.type == OmpEventType.MESSAGE_END ) assert message is not None and isinstance(message.content, list) image = next(content for content in message.content if content.type == "image") diff --git a/tests/integrations/omprpc/test_turn_settlement.py b/tests/integrations/omprpc/test_turn_settlement.py index af4dc67d9..2529d72ee 100644 --- a/tests/integrations/omprpc/test_turn_settlement.py +++ b/tests/integrations/omprpc/test_turn_settlement.py @@ -7,7 +7,7 @@ import pytest from band.integrations.omprpc import OmpPromptError -from band.integrations.ompwire import OmpState +from band.integrations.ompwire import OmpEventType, OmpState from tests.integrations.omprpc.kit import ClientFactory, Observer, prompt_task @@ -20,23 +20,14 @@ async def test_initial_yield_never_finishes_a_background_turn( observer = Observer() running = prompt_task(client, observer) try: - first_end = await observer.wait_for("agent_end") + first_end = await observer.wait_for(OmpEventType.AGENT_END) assert first_end.awaitingAsyncWork is True await client.request("checkpoint") assert not running.done() await client.request("release") async with asyncio.timeout(5): await running - messages = [ - content.text - for event in observer.events - if event.type == "message_end" - and event.message is not None - and isinstance(event.message.content, list) - for content in event.message.content - if content.type == "text" - ] - assert messages == ["yielded", "background complete"] + assert observer.assistant_texts == ["yielded", "background complete"] finally: running.cancel() await asyncio.gather(running, return_exceptions=True) @@ -51,7 +42,7 @@ async def test_resumed_failures_override_a_successful_initial_prompt_result( observer = Observer() running = prompt_task(client, observer) try: - await observer.wait_for("agent_end") + await observer.wait_for(OmpEventType.AGENT_END) await client.request("release") expected = ( "Resumed provider failed" if scenario != "resume-aborted" else "aborted" @@ -59,7 +50,10 @@ async def test_resumed_failures_override_a_successful_initial_prompt_result( async with asyncio.timeout(5): with pytest.raises(OmpPromptError, match=expected): await running - assert sum(event.type == "agent_end" for event in observer.events) == 2 + assert ( + sum(event.type == OmpEventType.AGENT_END for event in observer.events) + == 2 + ) finally: running.cancel() await asyncio.gather(running, return_exceptions=True) @@ -72,7 +66,7 @@ async def test_settlement_before_matching_result_waits_for_both( observer = Observer() running = prompt_task(client, observer) try: - await observer.wait_for("agent_end") + await observer.wait_for(OmpEventType.AGENT_END) await client.request("checkpoint") assert not running.done() await client.request("release") diff --git a/tests/integrations/omprpc/test_ui_answers.py b/tests/integrations/omprpc/test_ui_answers.py index 110df012e..a5616ed70 100644 --- a/tests/integrations/omprpc/test_ui_answers.py +++ b/tests/integrations/omprpc/test_ui_answers.py @@ -7,7 +7,12 @@ import pytest from pydantic import JsonValue -from band.integrations.ompwire import OmpUIRequest, OmpUIResponse +from band.integrations.ompwire import ( + OmpEventType, + OmpFrameType, + OmpUIRequest, + OmpUIResponse, +) from tests.integrations.omprpc.kit import ClientFactory, Observer @@ -45,11 +50,11 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: ) try: async with asyncio.timeout(8): - await observer.wait_for("agent_start") + await observer.wait_for(OmpEventType.AGENT_START) receipt = await client.request("ask-ui") assert receipt == { "answer": { - "type": "extension_ui_response", + "type": OmpFrameType.UI_RESPONSE, "id": "structured-question", "answers": [answer], } From 292c805b46ddc9049c5a2831db325116ec3529a7 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 13:31:06 +0300 Subject: [PATCH 10/35] fix: respect platform EOF contracts in Codex cleanup tests --- docs/adapters/codex.md | 3 +++ tests/integrations/codex/test_cleanup.py | 6 +++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/docs/adapters/codex.md b/docs/adapters/codex.md index 41d6e5ce8..5f12b25ab 100644 --- a/docs/adapters/codex.md +++ b/docs/adapters/codex.md @@ -35,6 +35,9 @@ For existing integrations, follow the [SDK 4.x to 5.x migration](../acp.md#upgra close errors and close timeouts retain the claim for retry. A timed-out or cancelled caller does not abandon subprocess cleanup; retry room cleanup to finish releasing ownership. Shutdown attempts every room and reports failures. + On POSIX, close sends stdin EOF and gives the session a bounded chance to + exit before reaping remaining descendants. Windows shutdown does not promise + session-visible stdin EOF. On Windows, the client owns the helper and its process-tree job from startup through cleanup. It starts the interpreter image directly so a virtual-environment launcher cannot hide another helper process. Cleanup confirms the helper has diff --git a/tests/integrations/codex/test_cleanup.py b/tests/integrations/codex/test_cleanup.py index 70000a828..6765799cb 100644 --- a/tests/integrations/codex/test_cleanup.py +++ b/tests/integrations/codex/test_cleanup.py @@ -223,8 +223,7 @@ async def test_forced_close_stops_what_a_launcher_shim_spawned( async def test_close_reaps_a_worker_when_its_launcher_exits_on_stdin_eof( self, tmp_path: Path ) -> None: - """A launcher gets to exit on stdin EOF, but its child still holding - the workdir is reaped by the tree kill that follows.""" + """POSIX launchers observe stdin EOF; every platform reaps their worker.""" workdir = tmp_path / "room" workdir.mkdir() client = CodexStdioClient( @@ -237,7 +236,8 @@ async def test_close_reaps_a_worker_when_its_launcher_exits_on_stdin_eof( await asyncio.wait_for(wait_for_heartbeat(heartbeat), timeout=CLOSE_BOUND_S) await asyncio.wait_for(client.close(), timeout=CLOSE_BOUND_S) await assert_heartbeat_stopped(heartbeat) - assert (workdir / STDIN_EOF_MARKER).exists() + if sys.platform != "win32": + assert (workdir / STDIN_EOF_MARKER).exists() shutil.rmtree(workdir) finally: await asyncio.wait_for(client.close(), timeout=CLOSE_BOUND_S) From 19e3b43148042016901dbfaaa6a1486c1ab6f20f Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 13:40:15 +0300 Subject: [PATCH 11/35] fix: treat graceful RPC EOF as teardown, not a crash Co-authored-by: Cursor --- src/band/converters/omp.py | 3 +- src/band/integrations/ompprocess.py | 45 +++++++++---- src/band/integrations/omprpc.py | 64 +++++++++++-------- tests/converters/test_omp.py | 13 ++-- .../ompprocess/test_mcp_plugin.py | 17 +++-- .../ompprocess/test_mcp_startup.py | 38 +++++++++-- tests/integrations/omprpc/test_callbacks.py | 24 +++++++ tests/integrations/omprpc/test_fail_closed.py | 12 ++-- .../omprpc/test_prompt_results.py | 12 +++- tests/integrations/omprpc/test_startup.py | 19 +++++- 10 files changed, 180 insertions(+), 67 deletions(-) diff --git a/src/band/converters/omp.py b/src/band/converters/omp.py index 0e45e54d3..457564a82 100644 --- a/src/band/converters/omp.py +++ b/src/band/converters/omp.py @@ -8,6 +8,7 @@ from band.converters.helpers import build_replay_messages from band.core.protocols import HistoryConverter +from band.core.types import MessageType # Task-event metadata keys the native adapter writes and this converter reads. OMP_ROOM_METADATA_KEY = "omp_room_id" @@ -27,7 +28,7 @@ class OmpHistoryConverter(HistoryConverter[OmpSessionState]): def convert(self, raw: list[dict[str, Any]]) -> OmpSessionState: sessions: dict[str, str] = {} for message in raw: - if message.get("message_type") != "task": + if message.get("message_type") != MessageType.TASK: continue metadata = message.get("metadata") if not isinstance(metadata, dict): diff --git a/src/band/integrations/ompprocess.py b/src/band/integrations/ompprocess.py index 77b0b26c0..70051861c 100644 --- a/src/band/integrations/ompprocess.py +++ b/src/band/integrations/ompprocess.py @@ -7,6 +7,7 @@ import tempfile from collections.abc import Collection, Iterator, Mapping, Sequence from contextlib import contextmanager +from enum import StrEnum from pathlib import Path from typing import Literal, Self from uuid import uuid4 @@ -40,6 +41,15 @@ # namespace every projected tool name carries. OMP_PLUGIN_ROOT = "band-sdk" OMP_PLUGIN_TOOL_NAMESPACE = f"{OMP_PLUGIN_ROOT.replace('-', '_')}_" +OMP_MCP_STARTUP_JS = "startup.js" +OMP_MCP_STARTUP_CONFIG = "startup.json" +OMP_MCP_RECEIPT = "receipt.json" +OMP_MCP_SERVERS_CONFIG = ".mcp.json" + + +class McpConnectionStatus(StrEnum): + CONNECTED = "connected" + FAILED = "failed" def omp_runtime_env( @@ -96,9 +106,9 @@ def _require_one_transport(self) -> Self: const ready = new Promise(resolve => { finish = resolve; }); const unsubscribe = pi.events.on("mcp:connection-status", event => { if (!expected.has(event.serverName) || statuses.has(event.serverName)) return; - if (event.type !== "connected" && event.type !== "failed") return; + if (event.type !== CONNECTED && event.type !== FAILED) return; statuses.set(event.serverName, event); - if (event.type === "failed" || statuses.size === expected.size) { + if (event.type === FAILED || statuses.size === expected.size) { unsubscribe(); finish(); } @@ -133,12 +143,12 @@ class McpCommands(BaseModel): class ConnectedMcpServer(BaseModel): - type: Literal["connected"] + type: Literal[McpConnectionStatus.CONNECTED] serverName: str class FailedMcpServer(BaseModel): - type: Literal["failed"] + type: Literal[McpConnectionStatus.FAILED] serverName: str error: str @@ -164,9 +174,9 @@ def omp_mcp_arguments(plugin: str) -> tuple[str, ...]: plugin, "--mode=rpc-ui", "--extension", - str(Path(plugin) / "startup.js"), + str(Path(plugin) / OMP_MCP_STARTUP_JS), "--config", - str(Path(plugin) / "startup.json"), + str(Path(plugin) / OMP_MCP_STARTUP_CONFIG), ) @@ -203,7 +213,7 @@ async def on_ui(_request: OmpUIRequest) -> OmpUIResponse: ) try: receipt = McpStartupReceipt.model_validate_json( - (Path(plugin) / "receipt.json").read_text(encoding="utf-8") + (Path(plugin) / OMP_MCP_RECEIPT).read_text(encoding="utf-8") ) except (OSError, ValueError) as error: raise OmpStartupError("OMP returned no valid SDK startup receipt") from error @@ -267,15 +277,26 @@ def omp_mcp_plugin( startup = { "servers": [_plugin_server_name(name) for name in servers], "command": omp_mcp_ready_command(str(root)), - "receipt": str(root / "receipt.json"), + "receipt": str(root / OMP_MCP_RECEIPT), "systemPrompt": system_prompt, } - _write_private(root / ".mcp.json", json.dumps({"mcpServers": projected})) + _write_private( + root / OMP_MCP_SERVERS_CONFIG, json.dumps({"mcpServers": projected}) + ) # Quiet startup suppresses the native status bus, not just UI rendering. - _write_private(root / "startup.json", json.dumps({"startup": {"quiet": False}})) _write_private( - root / "startup.js", - f"const config = {json.dumps(startup)};\n{_MCP_STARTUP_EXTENSION}", + root / OMP_MCP_STARTUP_CONFIG, json.dumps({"startup": {"quiet": False}}) + ) + _write_private( + root / OMP_MCP_STARTUP_JS, + "\n".join( + ( + f"const config = {json.dumps(startup)};", + f"const CONNECTED = {json.dumps(McpConnectionStatus.CONNECTED)};", + f"const FAILED = {json.dumps(McpConnectionStatus.FAILED)};", + _MCP_STARTUP_EXTENSION, + ) + ), ) yield str(root) diff --git a/src/band/integrations/omprpc.py b/src/band/integrations/omprpc.py index f89992469..17ccd7210 100644 --- a/src/band/integrations/omprpc.py +++ b/src/band/integrations/omprpc.py @@ -10,6 +10,7 @@ from binascii import Error as Base64Error from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass +from enum import StrEnum from typing import TypeVar, cast from uuid import uuid4 @@ -68,6 +69,11 @@ ) +class OmpCallbackKind(StrEnum): + TOOL = "tool" + UI = "ui" + + class OmpRpcError(RuntimeError): """An unsuccessful command response; the transport remains usable.""" @@ -147,7 +153,7 @@ def __init__( self.settled_sequence = 0 self.last_assistant: OmpMessage | None = None self.rejected = False - self.callbacks: dict[tuple[str, str], asyncio.Task[None]] = {} + self.callbacks: dict[tuple[OmpCallbackKind, str], asyncio.Task[None]] = {} self.callback_tasks: set[asyncio.Task[None]] = set() self.event_worker: asyncio.Task[None] | None = None @@ -647,6 +653,8 @@ async def _read_stdout(self, process: asyncio.subprocess.Process) -> None: while True: line = await process.stdout.readline() if not line: + if self._closing: + return if self._chunks is not None: raise OmpProtocolError( "OMP RPC protocol chunks interrupted by stdout EOF" @@ -708,12 +716,14 @@ async def _dispatch_payload( "OMP host tool arrived without an active prompt" ) self._start_callback( - turn, "tool", call.id, self._handle_tool(turn, call) + turn, OmpCallbackKind.TOOL, call.id, self._handle_tool(turn, call) ) case OmpFrameType.HOST_TOOL_CANCEL: cancellation = OmpHostToolCancel.model_validate(payload) if turn is not None: - self._cancel_callback(turn, "tool", cancellation.targetId) + self._cancel_callback( + turn, OmpCallbackKind.TOOL, cancellation.targetId + ) case OmpFrameType.UI_REQUEST: request = OmpUIRequest.model_validate(payload) if request.method == "cancel": @@ -722,7 +732,9 @@ async def _dispatch_payload( "OMP UI cancellation is missing targetId" ) if turn is not None: - self._cancel_callback(turn, "ui", request.targetId) + self._cancel_callback( + turn, OmpCallbackKind.UI, request.targetId + ) return if request.method in _UI_NOTIFICATIONS: return @@ -736,7 +748,7 @@ async def _dispatch_payload( ) return self._start_callback( - turn, "ui", request.id, self._handle_ui(turn, request) + turn, OmpCallbackKind.UI, request.id, self._handle_ui(turn, request) ) case event_type if event_type in _EVENT_TYPES: event = OmpEvent.model_validate(payload) @@ -767,7 +779,7 @@ def _dispatch_response(self, response: OmpResponse, sequence: int) -> None: def _start_callback( self, turn: PromptRun, - kind: str, + kind: OmpCallbackKind, request_id: str, callback: Coroutine[object, object, None], ) -> None: @@ -781,7 +793,7 @@ def _start_callback( task.add_done_callback(lambda done: self._callback_finished(turn, key, done)) def _callback_still_active( - self, turn: PromptRun, kind: str, request_id: str + self, turn: PromptRun, kind: OmpCallbackKind, request_id: str ) -> bool: return ( not self._closing @@ -789,14 +801,16 @@ def _callback_still_active( ) @staticmethod - def _cancel_callback(turn: PromptRun, kind: str, request_id: str) -> None: + def _cancel_callback( + turn: PromptRun, kind: OmpCallbackKind, request_id: str + ) -> None: task = turn.callbacks.pop((kind, request_id), None) if task is not None: task.cancel() @staticmethod def _callback_finished( - turn: PromptRun, key: tuple[str, str], task: asyncio.Task[None] + turn: PromptRun, key: tuple[OmpCallbackKind, str], task: asyncio.Task[None] ) -> None: if turn.callbacks.get(key) is task: turn.callbacks.pop(key) @@ -806,7 +820,7 @@ def _callback_finished( async def _handle_tool(self, turn: PromptRun, call: OmpHostToolCall) -> None: result = await turn.on_tool(call) - if not self._callback_still_active(turn, "tool", call.id): + if not self._callback_still_active(turn, OmpCallbackKind.TOOL, call.id): return await self._write_frame( { @@ -824,7 +838,7 @@ async def _handle_tool(self, turn: PromptRun, call: OmpHostToolCall) -> None: async def _handle_ui(self, turn: PromptRun, request: OmpUIRequest) -> None: response = await turn.on_ui(request) - if not self._callback_still_active(turn, "ui", request.id): + if not self._callback_still_active(turn, OmpCallbackKind.UI, request.id): return frame: dict[str, JsonValue] = { "type": OmpFrameType.UI_RESPONSE, @@ -844,10 +858,15 @@ async def _handle_ui(self, turn: PromptRun, request: OmpUIRequest) -> None: async def _read_stderr(self, process: asyncio.subprocess.Process) -> None: assert process.stderr is not None - while data := await process.stderr.read(_DRAIN_BYTES): - logger.debug( - "OMP stderr: %s", data.decode("utf-8", errors="replace").rstrip() - ) + try: + while data := await process.stderr.read(_DRAIN_BYTES): + logger.debug( + "OMP stderr: %s", data.decode("utf-8", errors="replace").rstrip() + ) + except asyncio.CancelledError: + raise + except Exception: + logger.exception("OMP stderr reader failed") def _fail_pending(self, error: OmpProcessError) -> None: for pending in self._pending.values(): @@ -910,15 +929,6 @@ async def _drain(stream: asyncio.StreamReader) -> None: while await stream.read(_DRAIN_BYTES): pass - @staticmethod - async def _wait_for_exit(process: asyncio.subprocess.Process) -> bool: - try: - async with asyncio.timeout(_PROCESS_EXIT_TIMEOUT_S): - await process.wait() - except TimeoutError: - return False - return True - async def _settle_spawn(self) -> asyncio.subprocess.Process | None: spawning = self._spawn_task if spawning is None: @@ -967,8 +977,8 @@ async def _close_process(self) -> None: try: async with asyncio.timeout(_ABORT_TIMEOUT_S): await self._request_frame("abort", {}, allow_closing=True) - except (TimeoutError, OmpRpcError, OmpProcessError): - pass + except (TimeoutError, OmpRpcError, OmpProcessError) as error: + logger.warning("OMP graceful abort skipped: %s", error) if self._reader_task is not None: await self._cancel_tasks([self._reader_task]) drains: list[asyncio.Task[None]] = [] @@ -1002,7 +1012,7 @@ async def _close_process(self) -> None: "OMP process did not exit after containment kill" ) await asyncio.to_thread(scope.close, process, _PROCESS_EXIT_TIMEOUT_S) - if not await self._wait_for_exit(process): + if not await await_returncode(process, _PROCESS_EXIT_TIMEOUT_S): raise OmpProcessError("OMP process did not exit after containment kill") self._windows_process = None self._process_scope = None diff --git a/tests/converters/test_omp.py b/tests/converters/test_omp.py index dc56572cd..1b5fd6563 100644 --- a/tests/converters/test_omp.py +++ b/tests/converters/test_omp.py @@ -7,6 +7,7 @@ OMP_SESSION_METADATA_KEY, OmpHistoryConverter, ) +from band.core.types import MessageType class TestOmpHistoryConverter: @@ -19,21 +20,21 @@ def test_latest_session_file_wins_per_room(self) -> None: state = OmpHistoryConverter().convert( [ { - "message_type": "task", + "message_type": MessageType.TASK, "metadata": { OMP_ROOM_METADATA_KEY: "room-1", OMP_SESSION_METADATA_KEY: "/sessions/old.jsonl", }, }, { - "message_type": "task", + "message_type": MessageType.TASK, "metadata": { OMP_ROOM_METADATA_KEY: "room-2", OMP_SESSION_METADATA_KEY: "/sessions/other.jsonl", }, }, { - "message_type": "task", + "message_type": MessageType.TASK, "metadata": { OMP_ROOM_METADATA_KEY: "room-1", OMP_SESSION_METADATA_KEY: "/sessions/new.jsonl", @@ -51,18 +52,18 @@ def test_unusable_metadata_is_ignored(self) -> None: [ {"message_type": "task", "metadata": "not a mapping"}, { - "message_type": "task", + "message_type": MessageType.TASK, "metadata": {OMP_ROOM_METADATA_KEY: "room-1"}, }, { - "message_type": "task", + "message_type": MessageType.TASK, "metadata": { OMP_ROOM_METADATA_KEY: "room-1", OMP_SESSION_METADATA_KEY: "", }, }, { - "message_type": "task", + "message_type": MessageType.TASK, "metadata": { OMP_ROOM_METADATA_KEY: 7, OMP_SESSION_METADATA_KEY: "/sessions/x.jsonl", diff --git a/tests/integrations/ompprocess/test_mcp_plugin.py b/tests/integrations/ompprocess/test_mcp_plugin.py index c53f924a1..8fe326199 100644 --- a/tests/integrations/ompprocess/test_mcp_plugin.py +++ b/tests/integrations/ompprocess/test_mcp_plugin.py @@ -8,6 +8,9 @@ import pytest from band.integrations.ompprocess import ( + OMP_MCP_SERVERS_CONFIG, + OMP_MCP_STARTUP_CONFIG, + OMP_MCP_STARTUP_JS, OmpMcpServer, omp_mcp_arguments, omp_mcp_plugin, @@ -29,7 +32,13 @@ def test_servers_land_in_the_plugin_mcp_config(self, tmp_path: Path) -> None: assert plugin is not None root = Path(plugin) assert root.name == "band-sdk" - config = json.loads((root / ".mcp.json").read_text(encoding="utf-8")) + config = json.loads( + (root / OMP_MCP_SERVERS_CONFIG).read_text(encoding="utf-8") + ) + startup_config = json.loads( + (root / OMP_MCP_STARTUP_CONFIG).read_text(encoding="utf-8") + ) + assert startup_config == {"startup": {"quiet": False}} assert config["mcpServers"]["remote"] == { "type": "sse", "url": "https://mcp.example/sse", @@ -38,7 +47,7 @@ def test_servers_land_in_the_plugin_mcp_config(self, tmp_path: Path) -> None: assert local["args"] == ["--port", "1"] assert Path(local["cwd"]) == tmp_path.resolve() assert Path(local["command"]) == (tmp_path / "server").resolve() - startup = (root / "startup.js").read_text(encoding="utf-8") + startup = (root / OMP_MCP_STARTUP_JS).read_text(encoding="utf-8") assert omp_mcp_ready_command(plugin) in startup assert '"band-sdk:local"' in startup assert '"band-sdk:remote"' in startup @@ -62,8 +71,8 @@ def test_arguments_load_the_plugin_and_its_startup_extension(self) -> None: arguments = omp_mcp_arguments(plugin) assert arguments[:2] == ("--plugin-dir", plugin) assert "--mode=rpc-ui" in arguments - assert str(Path(plugin) / "startup.js") in arguments - assert str(Path(plugin) / "startup.json") in arguments + assert str(Path(plugin) / OMP_MCP_STARTUP_JS) in arguments + assert str(Path(plugin) / OMP_MCP_STARTUP_CONFIG) in arguments @pytest.mark.parametrize( "fields", diff --git a/tests/integrations/ompprocess/test_mcp_startup.py b/tests/integrations/ompprocess/test_mcp_startup.py index 4d4d8ba1b..38d424de9 100644 --- a/tests/integrations/ompprocess/test_mcp_startup.py +++ b/tests/integrations/ompprocess/test_mcp_startup.py @@ -18,6 +18,7 @@ from band.integrations.omprpc import OmpRpcClient from band.integrations.ompwire import ( OmpEvent, + OmpEventType, OmpHostToolCall, OmpToolResult, OmpUIRequest, @@ -142,8 +143,16 @@ async def test_complete_receipt_prompts_the_readiness_command( assert runtime.prompt.await_args is not None assert runtime.prompt.await_args.args[0] == f"/{command}" - async def test_ui_request_during_startup_cannot_fall_through_to_inference( - self, tmp_path: Path + @pytest.mark.parametrize( + ("invoke", "reason"), + [ + ("event", "invoked an agent"), + ("tool", "invoked a host tool"), + ("ui", "requested UI"), + ], + ) + async def test_startup_callbacks_cannot_fall_through_to_inference( + self, tmp_path: Path, invoke: str, reason: str ) -> None: runtime = AsyncMock(spec=OmpRpcClient) runtime.request.return_value = { @@ -157,9 +166,28 @@ async def prompt( on_tool: Callable[[OmpHostToolCall], Awaitable[OmpToolResult]], on_ui: Callable[[OmpUIRequest], Awaitable[OmpUIResponse]], ) -> None: - _ = (message, on_event, on_tool) - await on_ui(OmpUIRequest.model_validate({"id": "ask", "method": "select"})) + _ = message + match invoke: + case "event": + await on_event( + OmpEvent.model_validate({"type": OmpEventType.AGENT_START}) + ) + case "tool": + await on_tool( + OmpHostToolCall.model_validate( + { + "id": "host-1", + "toolCallId": "call-1", + "toolName": "band_chat", + "arguments": {}, + } + ) + ) + case "ui": + await on_ui( + OmpUIRequest.model_validate({"id": "ask", "method": "select"}) + ) runtime.prompt.side_effect = prompt - with pytest.raises(OmpStartupError, match="requested UI"): + with pytest.raises(OmpStartupError, match=reason): await omp_mcp_ready(runtime, str(tmp_path), {}) diff --git a/tests/integrations/omprpc/test_callbacks.py b/tests/integrations/omprpc/test_callbacks.py index 96f8c3138..e866bb780 100644 --- a/tests/integrations/omprpc/test_callbacks.py +++ b/tests/integrations/omprpc/test_callbacks.py @@ -86,6 +86,30 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: } await second_prompt(client) + async def test_confirmed_ui_reply_reaches_the_peer( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("callbacks") + observer = Observer() + + async def on_ui(_request: OmpUIRequest) -> OmpUIResponse: + return OmpUIResponse(confirmed=True) + + async with asyncio.timeout(5): + await client.prompt( + "first", + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=on_ui, + ) + frames = await checkpoint_frames(client) + ui_response = next(frame for frame in frames if frame.get("id") == "approval-1") + assert ui_response == { + "type": OmpFrameType.UI_RESPONSE, + "id": "approval-1", + "confirmed": True, + } + @pytest.mark.parametrize("kind", ["host", "ui"]) async def test_native_cancel_targets_the_callback_and_keeps_the_next_turn_usable( self, client_factory: ClientFactory, kind: str diff --git a/tests/integrations/omprpc/test_fail_closed.py b/tests/integrations/omprpc/test_fail_closed.py index 435b7d354..ccbc648ba 100644 --- a/tests/integrations/omprpc/test_fail_closed.py +++ b/tests/integrations/omprpc/test_fail_closed.py @@ -32,12 +32,12 @@ async def test_interactive_ui_without_an_active_prompt_is_cancelled( async with asyncio.timeout(5): await client.request("orphan-ui") frames = await checkpoint_frames(client) - assert any( - frame.get("type") == OmpFrameType.UI_RESPONSE - and frame.get("id") == "ui-orphan" - and frame.get("cancelled") is True - for frame in frames - ) + orphan = next(frame for frame in frames if frame.get("id") == "ui-orphan") + assert orphan == { + "type": OmpFrameType.UI_RESPONSE, + "id": "ui-orphan", + "cancelled": True, + } async def test_ui_notifications_do_not_start_a_callback( self, client_factory: ClientFactory diff --git a/tests/integrations/omprpc/test_prompt_results.py b/tests/integrations/omprpc/test_prompt_results.py index 92f4983f0..082ff8975 100644 --- a/tests/integrations/omprpc/test_prompt_results.py +++ b/tests/integrations/omprpc/test_prompt_results.py @@ -92,13 +92,19 @@ async def test_rejected_prompt_does_not_wait_for_a_result_that_is_not_owed( assert (rejected.value.command, rejected.value.code) == ("prompt", "rejected") await second_prompt(client) - @pytest.mark.parametrize("scenario", ["initial-error", "initial-aborted"]) + @pytest.mark.parametrize( + ("scenario", "match"), + [ + ("initial-error", "Initial provider failed"), + ("initial-aborted", "OMP prompt aborted"), + ], + ) async def test_failed_initial_result_is_reported( - self, client_factory: ClientFactory, scenario: str + self, client_factory: ClientFactory, scenario: str, match: str ) -> None: client = await client_factory(scenario) async with asyncio.timeout(5): - with pytest.raises(OmpPromptError): + with pytest.raises(OmpPromptError, match=match): await prompt_task(client, Observer()) async def test_narration_failure_does_not_leave_the_turn_or_process_waiting( diff --git a/tests/integrations/omprpc/test_startup.py b/tests/integrations/omprpc/test_startup.py index 48d410111..9542110dc 100644 --- a/tests/integrations/omprpc/test_startup.py +++ b/tests/integrations/omprpc/test_startup.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import logging import sys import pytest @@ -29,11 +30,15 @@ async def test_cancelled_startup_reaps_a_process_that_never_announces_ready( "no-ready", ) ) + starting = asyncio.create_task(client.start()) try: - with pytest.raises(TimeoutError): - async with asyncio.timeout(0.05): - await client.start() + async with asyncio.timeout(5): + while client._proc is None and not starting.done(): + await asyncio.sleep(0) assert client._proc is not None + assert starting.cancel() + with pytest.raises(asyncio.CancelledError): + await starting assert client._proc.returncode is not None finally: await client.close() @@ -60,3 +65,11 @@ async def test_start_requires_successful_native_v2_negotiation( async with asyncio.timeout(6): with pytest.raises(error, match="protocol"): await client_factory(scenario) + + async def test_graceful_close_does_not_log_stdout_eof_as_a_crash( + self, client_factory: ClientFactory, caplog: pytest.LogCaptureFixture + ) -> None: + client = await client_factory("basic") + with caplog.at_level(logging.ERROR, logger="band.integrations.omprpc"): + await client.close() + assert caplog.records == [] From 098b548b2604be9a20a1f2f2a57815c3a8687132 Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 14:19:45 +0300 Subject: [PATCH 12/35] fix: keep failed native prompts from reaping the RPC process Co-authored-by: Cursor --- src/band/integrations/omp.py | 8 +++- src/band/integrations/ompprocess.py | 4 +- src/band/integrations/omprpc.py | 42 +++++++++++++------ src/band/integrations/ompwire.py | 13 +++++- tests/integrations/omppeer.py | 13 ++++++ tests/integrations/ompprocess/test_launch.py | 4 +- .../ompprocess/test_mcp_plugin.py | 6 ++- .../ompprocess/test_mcp_startup.py | 21 +++++----- .../integrations/omprpc/test_process_tree.py | 18 ++++---- .../omprpc/test_prompt_results.py | 4 ++ .../omprpc/test_turn_settlement.py | 32 +++++++++++++- 11 files changed, 124 insertions(+), 41 deletions(-) diff --git a/src/band/integrations/omp.py b/src/band/integrations/omp.py index 69782d9c1..1e0e189da 100644 --- a/src/band/integrations/omp.py +++ b/src/band/integrations/omp.py @@ -24,6 +24,12 @@ # The acp subcommand's flag for the room workspace, passed as ``--cwd=``. OMP_CWD_FLAG = "--cwd" + +def omp_cwd_argument(workspace: str) -> str: + """Format OMP's workspace flag without changing launch cwd.""" + return f"{OMP_CWD_FLAG}={workspace}" + + DEFAULT_OMP_ACP_COMMAND: tuple[str, ...] = ( "omp", OMP_ACP_SUBCOMMAND, @@ -167,7 +173,7 @@ def omp_command_in_workspace(command: Sequence[str], workspace: str) -> list[str f"got {list(command)!r}" ) split = list(command).index(OMP_ACP_SUBCOMMAND) + 1 - return [*command[:split], f"{OMP_CWD_FLAG}={workspace}", *command[split:]] + return [*command[:split], omp_cwd_argument(workspace), *command[split:]] def omp_elicitation_call_id(session_id: str) -> str: diff --git a/src/band/integrations/ompprocess.py b/src/band/integrations/ompprocess.py index 70051861c..8059631e1 100644 --- a/src/band/integrations/ompprocess.py +++ b/src/band/integrations/ompprocess.py @@ -15,9 +15,9 @@ from pydantic import BaseModel, ConfigDict, Field, model_validator from band.integrations.omp import ( - OMP_CWD_FLAG, normalize_omp_mcp_device_call, normalize_omp_mcp_tool_name, + omp_cwd_argument, omp_provider_env, ) from band.integrations.omprpc import OmpRpcClient @@ -69,7 +69,7 @@ def omp_runtime_env( def omp_rpc_command_in_workspace(command: Sequence[str], workspace: str) -> list[str]: """Select OMP's room workspace without changing the subprocess launch cwd.""" - return [*command, f"{OMP_CWD_FLAG}={workspace}"] + return [*command, omp_cwd_argument(workspace)] class OmpMcpServer(BaseModel): diff --git a/src/band/integrations/omprpc.py b/src/band/integrations/omprpc.py index 17ccd7210..efa9ee165 100644 --- a/src/band/integrations/omprpc.py +++ b/src/band/integrations/omprpc.py @@ -26,6 +26,7 @@ session_pid_reused, ) from band.integrations.ompwire import ( + OMP_ASSISTANT_ROLE, OMP_CHUNK_BYTES, OMP_FRAME_BYTES, OMP_NATIVE_PROTOCOL_VERSION, @@ -43,6 +44,7 @@ OmpNegotiation, OmpPromptAcceptance, OmpPromptResult, + OmpPromptStatus, OmpReady, OmpResponse, OmpState, @@ -173,7 +175,7 @@ def observe(self, event: OmpEvent, sequence: int) -> None: self.last_activity_sequence = sequence match event.type: case OmpEventType.MESSAGE_END if ( - event.message is not None and event.message.role == "assistant" + event.message is not None and event.message.role == OMP_ASSISTANT_ROLE ): self.last_assistant = event.message case OmpEventType.AGENT_END: @@ -181,7 +183,7 @@ def observe(self, event: OmpEvent, sequence: int) -> None: ( message for message in reversed(event.messages) - if message.role == "assistant" + if message.role == OMP_ASSISTANT_ROLE ), self.last_assistant, ) @@ -303,6 +305,7 @@ def __init__( self._start_lock = asyncio.Lock() self._write_lock = asyncio.Lock() self._prompt_lock = asyncio.Lock() + self._write_tasks: set[asyncio.Task[None]] = set() async def start(self) -> None: """Start the supplied command and require native RPC protocol v2.""" @@ -338,7 +341,9 @@ async def start(self) -> None: if isinstance(ready, OmpProcessError): raise ready if OMP_NATIVE_PROTOCOL_VERSION not in ready.supportedProtocolVersions: - raise OmpProtocolError("OMP RPC protocol v2 is required") + raise OmpProtocolError( + f"OMP RPC protocol v{OMP_NATIVE_PROTOCOL_VERSION} is required" + ) self._max_frame_bytes = ready.maxFrameBytes self._max_reassembled_bytes = ready.maxReassembledFrameBytes negotiated = await self._request_frame( @@ -348,7 +353,7 @@ async def start(self) -> None: self._parse_response( negotiated.frame, OmpNegotiation, - "OMP RPC protocol v2 negotiation failed", + f"OMP RPC protocol v{OMP_NATIVE_PROTOCOL_VERSION} negotiation failed", ) except (OSError, SpawnFailure) as cause: await self.close() @@ -503,6 +508,8 @@ async def _execute_prompt( if not turn.rejected: await self.close() raise + except OmpPromptError: + raise except BaseException: await self.close() raise @@ -535,7 +542,7 @@ async def _complete_prompt(self, turn: PromptRun, message: str) -> None: ) if pending: raise OmpProcessError("OMP callbacks did not finish after settlement") - if outcome is not None and outcome.status != "completed": + if outcome is not None and outcome.status != OmpPromptStatus.COMPLETED: raise OmpPromptError( outcome.status, outcome.error.message @@ -543,7 +550,10 @@ async def _complete_prompt(self, turn: PromptRun, message: str) -> None: else f"OMP prompt {outcome.status}", ) assistant = turn.last_assistant - if assistant is not None and assistant.stopReason in {"error", "aborted"}: + if assistant is not None and assistant.stopReason in { + OmpPromptStatus.ERROR, + OmpPromptStatus.ABORTED, + }: raise OmpPromptError( assistant.stopReason, assistant.errorMessage or f"OMP continuation {assistant.stopReason}", @@ -739,13 +749,18 @@ async def _dispatch_payload( if request.method in _UI_NOTIFICATIONS: return if turn is None: - await self._write_frame( - { - "type": OmpFrameType.UI_RESPONSE, - "id": request.id, - "cancelled": True, - } + write = asyncio.create_task( + self._write_frame( + { + "type": OmpFrameType.UI_RESPONSE, + "id": request.id, + "cancelled": True, + } + ) ) + self._write_tasks.add(write) + write.add_done_callback(self._write_tasks.discard) + write.add_done_callback(self._consume_task) return self._start_callback( turn, OmpCallbackKind.UI, request.id, self._handle_ui(turn, request) @@ -783,6 +798,9 @@ def _start_callback( request_id: str, callback: Coroutine[object, object, None], ) -> None: + if self._closing or turn.failure.done(): + callback.close() + return key = (kind, request_id) task = asyncio.create_task(callback) if key in turn.callbacks: diff --git a/src/band/integrations/ompwire.py b/src/band/integrations/ompwire.py index 1436e1c41..b67f73eda 100644 --- a/src/band/integrations/ompwire.py +++ b/src/band/integrations/ompwire.py @@ -19,6 +19,13 @@ OMP_REASSEMBLED_BYTES = 64 * 1024 * 1024 OMP_CHUNK_BYTES = 256 * 1024 OMP_NATIVE_PROTOCOL_VERSION: Final = 2 +OMP_ASSISTANT_ROLE = "assistant" + + +class OmpPromptStatus(StrEnum): + COMPLETED = "completed" + ERROR = "error" + ABORTED = "aborted" class OmpFrameType(StrEnum): @@ -322,7 +329,11 @@ class OmpPromptResult(OmpWireModel): type: Literal[OmpFrameType.PROMPT_RESULT] id: str | None = None agentInvoked: bool - status: Literal["completed", "error", "aborted"] + status: Literal[ + OmpPromptStatus.COMPLETED, + OmpPromptStatus.ERROR, + OmpPromptStatus.ABORTED, + ] sessionSettled: bool error: OmpPromptFailure | None = None diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index 54f3279ee..cd284f3d6 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -439,6 +439,19 @@ def prompt(self, frame: Frame) -> None: case "no-result": self.end("complete") self.settle() + case "settled-during-callback": + send( + { + "type": "host_tool_call", + "id": "host-held", + "toolCallId": "call-held", + "toolName": "band_chat", + "arguments": {}, + } + ) + self.end("complete") + self.result(settled=True) + self.settle() case "host-cancel": send( { diff --git a/tests/integrations/ompprocess/test_launch.py b/tests/integrations/ompprocess/test_launch.py index cd247e231..5b6f40dad 100644 --- a/tests/integrations/ompprocess/test_launch.py +++ b/tests/integrations/ompprocess/test_launch.py @@ -5,8 +5,8 @@ from band.integrations.omp import ( OMP_APPROVAL_MODE_ALWAYS_ASK, OMP_APPROVAL_MODE_FLAG, - OMP_CWD_FLAG, finalize_omp_command, + omp_cwd_argument, ) from band.integrations.ompprocess import ( DEFAULT_OMP_COMMAND, @@ -29,7 +29,7 @@ def test_workspace_is_selected_without_a_subprocess_cwd(self) -> None: "bunx", *DEFAULT_OMP_COMMAND, "-v", - f"{OMP_CWD_FLAG}=/w", + omp_cwd_argument("/w"), ] def test_runtime_env_routes_the_selected_provider_key(self) -> None: diff --git a/tests/integrations/ompprocess/test_mcp_plugin.py b/tests/integrations/ompprocess/test_mcp_plugin.py index 8fe326199..bdda521d5 100644 --- a/tests/integrations/ompprocess/test_mcp_plugin.py +++ b/tests/integrations/ompprocess/test_mcp_plugin.py @@ -66,8 +66,10 @@ def test_connection_instructions_become_an_always_applied_rule( "---\nalwaysApply: true\n---\nReply through Band." ) - def test_arguments_load_the_plugin_and_its_startup_extension(self) -> None: - plugin = "/tmp/overlay/band-sdk" + def test_arguments_load_the_plugin_and_its_startup_extension( + self, tmp_path: Path + ) -> None: + plugin = str(tmp_path / "band-sdk") arguments = omp_mcp_arguments(plugin) assert arguments[:2] == ("--plugin-dir", plugin) assert "--mode=rpc-ui" in arguments diff --git a/tests/integrations/ompprocess/test_mcp_startup.py b/tests/integrations/ompprocess/test_mcp_startup.py index 38d424de9..567d4e5b1 100644 --- a/tests/integrations/ompprocess/test_mcp_startup.py +++ b/tests/integrations/ompprocess/test_mcp_startup.py @@ -121,19 +121,20 @@ async def test_complete_receipt_prompts_the_readiness_command( plugin = tmp_path / "band-sdk" plugin.mkdir() command = omp_mcp_ready_command(str(plugin)) - (plugin / "receipt.json").write_text( - json.dumps( - { - "systemPromptReady": True, - "servers": [ - {"type": "connected", "serverName": "band-sdk:required"} - ], - } - ), - encoding="utf-8", + receipt = json.dumps( + { + "systemPromptReady": True, + "servers": [{"type": "connected", "serverName": "band-sdk:required"}], + } ) runtime = AsyncMock(spec=OmpRpcClient) runtime.request.return_value = {"commands": [{"name": command}]} + + async def prompt(message: str, **_callbacks: object) -> None: + _ = message + (plugin / "receipt.json").write_text(receipt, encoding="utf-8") + + runtime.prompt.side_effect = prompt await omp_mcp_ready( runtime, str(plugin), diff --git a/tests/integrations/omprpc/test_process_tree.py b/tests/integrations/omprpc/test_process_tree.py index 9521dad2b..d5c78827f 100644 --- a/tests/integrations/omprpc/test_process_tree.py +++ b/tests/integrations/omprpc/test_process_tree.py @@ -26,7 +26,6 @@ class TestProcessTreeCleanup: - @pytest.mark.skipif(os.name == "nt", reason="POSIX setsid ownership boundary") @pytest.mark.parametrize("root_exit", ["close", "crash"]) async def test_close_reaps_detached_descendants_without_affecting_other_owners( self, process_gate: ProcessGate, root_exit: str @@ -43,13 +42,15 @@ async def test_close_reaps_detached_descendants_without_affecting_other_owners( ] root = client._proc assert root is not None - assert all( - os.getsid(child.process.pid) == child.process.pid - for child in descendants - ) - assert all( - os.getsid(child.process.pid) != root.pid for child in descendants - ) + if os.name != "nt": + assert all( + os.getsid(child.process.pid) == child.process.pid + for child in descendants + ) + assert all( + os.getsid(child.process.pid) != root.pid + for child in descendants + ) await other.start() other_descendants = [ await process_gate.connection(), @@ -81,7 +82,6 @@ async def test_close_reaps_detached_descendants_without_affecting_other_owners( if crashing is not None: await asyncio.gather(crashing, return_exceptions=True) - @pytest.mark.skipif(os.name == "nt", reason="POSIX inherited-scope boundary") async def test_stale_matching_environment_cannot_adopt_a_reused_process_identity( self, process_gate: ProcessGate, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/integrations/omprpc/test_prompt_results.py b/tests/integrations/omprpc/test_prompt_results.py index 082ff8975..2b202db65 100644 --- a/tests/integrations/omprpc/test_prompt_results.py +++ b/tests/integrations/omprpc/test_prompt_results.py @@ -103,9 +103,13 @@ async def test_failed_initial_result_is_reported( self, client_factory: ClientFactory, scenario: str, match: str ) -> None: client = await client_factory(scenario) + process = client._proc + assert process is not None async with asyncio.timeout(5): with pytest.raises(OmpPromptError, match=match): await prompt_task(client, Observer()) + assert process.returncode is None + await second_prompt(client) async def test_narration_failure_does_not_leave_the_turn_or_process_waiting( self, diff --git a/tests/integrations/omprpc/test_turn_settlement.py b/tests/integrations/omprpc/test_turn_settlement.py index 2529d72ee..398eb9495 100644 --- a/tests/integrations/omprpc/test_turn_settlement.py +++ b/tests/integrations/omprpc/test_turn_settlement.py @@ -6,8 +6,14 @@ import pytest -from band.integrations.omprpc import OmpPromptError -from band.integrations.ompwire import OmpEventType, OmpState +from band.integrations import omprpc +from band.integrations.omprpc import OmpProcessError, OmpPromptError +from band.integrations.ompwire import ( + OmpEventType, + OmpHostToolCall, + OmpState, + OmpToolResult, +) from tests.integrations.omprpc.kit import ClientFactory, Observer, prompt_task @@ -87,3 +93,25 @@ async def test_native_state_snapshot_and_subscribed_settlement_do_not_race( async with asyncio.timeout(5): await prompt_task(client, observer) assert OmpState.model_validate(await client.request("get_state")).isSettled + + async def test_settlement_does_not_succeed_while_a_callback_is_still_running( + self, client_factory: ClientFactory, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(omprpc, "_CALLBACK_EXIT_TIMEOUT_S", 0) + client = await client_factory("settled-during-callback") + held = asyncio.Event() + + async def on_tool(_call: OmpHostToolCall) -> OmpToolResult: + await held.wait() + return OmpToolResult(content=[]) + + observer = Observer() + with pytest.raises(OmpProcessError, match="callbacks did not finish"): + async with asyncio.timeout(5): + await client.prompt( + "first", + on_event=observer.on_event, + on_tool=on_tool, + on_ui=observer.on_ui, + ) + held.set() From 2ad483cc7bbb7d6e24a1121d1d4fb40e9b605b75 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 14:31:27 +0300 Subject: [PATCH 13/35] fix: await Windows job member kernel exit --- CHANGELOG.md | 6 + docs/adapters/codex.md | 8 +- src/band/integrations/codex/winspawn.py | 127 +++++++++++++++- tests/integrations/codex/conftest.py | 12 ++ tests/integrations/codex/job.py | 166 ++++++++++++++++++++ tests/integrations/codex/test_cleanup.py | 15 +- tests/integrations/codex/test_job.py | 183 +++++++++++++++++++++-- 7 files changed, 485 insertions(+), 32 deletions(-) create mode 100644 tests/integrations/codex/job.py diff --git a/CHANGELOG.md b/CHANGELOG.md index edac55b1f..5b3a784eb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,12 @@ The format is based on [Conventional Commits](https://www.conventionalcommits.or and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). This changelog is automatically generated by [Release Please](https://github.com/googleapis/release-please). +## Unreleased + +### Bug Fixes + +* Wait for Windows Job members' kernel-confirmed exit before releasing OMP or Codex subprocess ownership, preserving process handles across incomplete cleanup and preventing late child admission during termination. + ## [6.0.1](https://github.com/band-ai/band-sdk-python/compare/band-sdk-v6.0.0...band-sdk-v6.0.1) (2026-10-20) diff --git a/docs/adapters/codex.md b/docs/adapters/codex.md index 41d6e5ce8..8ec9cf13f 100644 --- a/docs/adapters/codex.md +++ b/docs/adapters/codex.md @@ -38,9 +38,11 @@ For existing integrations, follow the [SDK 4.x to 5.x migration](../acp.md#upgra On Windows, the client owns the helper and its process-tree job from startup through cleanup. It starts the interpreter image directly so a virtual-environment launcher cannot hide another helper process. Cleanup confirms the helper has - exited before terminating and draining the job; an empty job alone does not - establish helper exit. - Failed cleanup retains that same owner and job for retry. + exited, prevents new job members from being admitted, and pins the current + members' process identities before terminating the job. It waits for both empty + job accounting and kernel-confirmed member exit; an empty job alone does not + establish process exit. Failed cleanup retains that same owner, job and process + handles for retry. - **Leaving is final.** Leaving a room stops it admitting work at once. A leave that fails or runs over keeps the process, thread and workspace claim, and the room refuses work until a retried leave succeeds; a later rejoin diff --git a/src/band/integrations/codex/winspawn.py b/src/band/integrations/codex/winspawn.py index 4943999ad..08261ddf7 100644 --- a/src/band/integrations/codex/winspawn.py +++ b/src/band/integrations/codex/winspawn.py @@ -33,14 +33,19 @@ _STATUS_ERROR = "error" _SPAWN_ACK_S = 30.0 _JOB_POLL_S = 0.05 +_JOB_LIST_QUERIES = 2 _ERROR_ALREADY_EXISTS = 183 _ERROR_BROKEN_PIPE = 109 +_ERROR_INVALID_PARAMETER = 87 +_ERROR_MORE_DATA = 234 _STATUS_READ_SIZE = 256 _PATH_BUFFER_SIZE = 260 _JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x00002000 +_JOB_OBJECT_LIMIT_ACTIVE_PROCESS = 0x00000008 _JOB_OBJECT_EXTENDED_LIMIT_INFORMATION = 9 _JOB_OBJECT_BASIC_ACCOUNTING_INFORMATION = 1 +_JOB_OBJECT_BASIC_PROCESS_ID_LIST = 3 _JOB_OBJECT_ASSIGN_PROCESS = 0x0001 _CREATE_SUSPENDED = 0x00000004 @@ -49,6 +54,9 @@ _HANDLE_FLAG_INHERIT = 0x00000001 _PROC_THREAD_ATTRIBUTE_HANDLE_LIST = 0x00020002 _PROC_THREAD_ATTRIBUTE_JOB_LIST = 0x0002000D +_PROCESS_QUERY_LIMITED_INFORMATION = 0x1000 +_SYNCHRONIZE = 0x00100000 +_WAIT_OBJECT_0 = 0 _INFINITE = 0xFFFFFFFF _WAIT_FAILED = 0xFFFFFFFF @@ -119,6 +127,14 @@ class BasicAccountingInformation(ctypes.Structure): ) +class ProcessIdList(ctypes.Structure): + _fields_ = ( + ("NumberOfAssignedProcesses", wintypes.DWORD), + ("NumberOfProcessIdsInList", wintypes.DWORD), + ("ProcessIdList", ctypes.c_size_t * 1), + ) + + class StartupInfo(ctypes.Structure): _fields_ = ( ("cb", wintypes.DWORD), @@ -204,6 +220,14 @@ def _kernel32() -> Any: library.CloseHandle.restype = wintypes.BOOL library.OpenJobObjectW.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.LPCWSTR] library.OpenJobObjectW.restype = handle + library.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + library.OpenProcess.restype = handle + library.IsProcessInJob.argtypes = [ + handle, + handle, + ctypes.POINTER(wintypes.BOOL), + ] + library.IsProcessInJob.restype = wintypes.BOOL library.CreateProcessW.argtypes = [ wintypes.LPCWSTR, ctypes.POINTER(ctypes.c_wchar), @@ -392,13 +416,19 @@ def _fail(api: str, code: int | None = None) -> NoReturn: class Job: - """A job handle owned by the parent process.""" + """Retain the job and identity-pinned members until kernel exit.""" def __init__(self, handle: ctypes.c_void_p, name: str) -> None: self._handle: ctypes.c_void_p | None = handle self.name = name + self._process_handles: dict[int, ctypes.c_void_p] = {} def close(self) -> None: + while self._process_handles: + pid = next(iter(self._process_handles)) + if not _kernel32().CloseHandle(self._process_handles[pid]): + raise _win_os_error(_last_error()) + del self._process_handles[pid] handle = self._handle if handle is not None and not _kernel32().CloseHandle(handle): raise _win_os_error(_last_error()) @@ -428,12 +458,98 @@ def active_processes(self) -> int | None: return None return int(info.ActiveProcesses) + def _seal_admission(self) -> None: + limits = ExtendedLimitInformation() + handle = self._open_handle() + if not _kernel32().QueryInformationJobObject( + handle, + _JOB_OBJECT_EXTENDED_LIMIT_INFORMATION, + ctypes.byref(limits), + ctypes.sizeof(limits), + None, + ): + raise _win_os_error(_last_error()) + # No member may fork between enumeration and termination. + limits.BasicLimitInformation.LimitFlags |= _JOB_OBJECT_LIMIT_ACTIVE_PROCESS + limits.BasicLimitInformation.ActiveProcessLimit = 0 + if not _kernel32().SetInformationJobObject( + handle, + _JOB_OBJECT_EXTENDED_LIMIT_INFORMATION, + ctypes.byref(limits), + ctypes.sizeof(limits), + ): + raise _win_os_error(_last_error()) + + def _process_ids(self) -> list[int]: + capacity = 1 + # Sealed admission permits one size query and one complete snapshot. + for _attempt in range(_JOB_LIST_QUERIES): + buffer = ctypes.create_string_buffer( + ProcessIdList.ProcessIdList.offset + + capacity * ctypes.sizeof(ctypes.c_size_t) + ) + info = ctypes.cast(buffer, ctypes.POINTER(ProcessIdList)).contents + succeeded = _kernel32().QueryInformationJobObject( + self._open_handle(), + _JOB_OBJECT_BASIC_PROCESS_ID_LIST, + buffer, + ctypes.sizeof(buffer), + None, + ) + if not succeeded and (code := _last_error()) != _ERROR_MORE_DATA: + raise _win_os_error(code) + assigned = int(info.NumberOfAssignedProcesses) + count = int(info.NumberOfProcessIdsInList) + if succeeded and count == assigned and count <= capacity: + pids = (ctypes.c_size_t * count).from_buffer( + buffer, ProcessIdList.ProcessIdList.offset + ) + return list(pids) + capacity = max(capacity * 2, assigned) + raise _win_os_error(_ERROR_MORE_DATA) + + def _capture_process_handles(self) -> None: + for pid in self._process_ids(): + if pid in self._process_handles: + continue + handle = _kernel32().OpenProcess( + _SYNCHRONIZE | _PROCESS_QUERY_LIMITED_INFORMATION, False, pid + ) + if not handle: + code = _last_error() + if code == _ERROR_INVALID_PARAMETER: + continue + raise _win_os_error(code) + self._process_handles[pid] = ctypes.c_void_p(handle) + for pid, handle in tuple(self._process_handles.items()): + member = wintypes.BOOL() + if not _kernel32().IsProcessInJob( + handle, self._open_handle(), ctypes.byref(member) + ): + raise _win_os_error(_last_error()) + if not member.value: + # Opening a PID can race exit and reuse; never adopt that identity. + if not _kernel32().CloseHandle(handle): + raise _win_os_error(_last_error()) + del self._process_handles[pid] + + def _pending_process_exits(self) -> int: + pending = 0 + for handle in self._process_handles.values(): + state = _kernel32().WaitForSingleObject(handle, 0) + if state == _WAIT_FAILED: + raise _win_os_error(_last_error()) + pending += state != _WAIT_OBJECT_0 + return pending + def terminate(self) -> None: """Mark every member for termination. Members blocked in a synchronous pipe read stay inside that call until the other end closes. Wait only after that end is closed. """ + self._seal_admission() + self._capture_process_handles() if not _kernel32().TerminateJobObject(self._open_handle(), 1): raise _win_os_error(_last_error()) @@ -443,11 +559,16 @@ async def wait_until_idle(self, timeout: float) -> None: deadline = loop.time() + timeout while True: active = self.active_processes() - if active == 0: + pending = self._pending_process_exits() + if active == 0 and not pending: return remaining = deadline - loop.time() if active is None or remaining <= 0: - logger.warning("Codex job still has %s active processes", active) + logger.warning( + "Codex job has %s active processes and %s pending kernel exits", + active, + pending, + ) raise JobNotDrained(active) await asyncio.sleep(min(_JOB_POLL_S, remaining)) diff --git a/tests/integrations/codex/conftest.py b/tests/integrations/codex/conftest.py index c3a6d2f99..31eef2ca0 100644 --- a/tests/integrations/codex/conftest.py +++ b/tests/integrations/codex/conftest.py @@ -12,10 +12,22 @@ import pytest_asyncio from band.integrations.codex import stdio_client, winspawn +from tests.integrations.codex.job import JobKernel from tests.integrations.codex.lifecycle import WindowsLifecycle from tests.integrations.codex.lookup import WindowsLookup +@pytest.fixture +def job_kernel(monkeypatch: pytest.MonkeyPatch) -> JobKernel: + kernel = JobKernel() + monkeypatch.setattr(winspawn, "_kernel32", lambda: kernel) + monkeypatch.setattr(winspawn, "_last_error", lambda: kernel.last_error) + monkeypatch.setattr( + winspawn, "_win_os_error", lambda code: OSError(code, "native Job failure") + ) + return kernel + + @pytest.fixture def windows_lookup(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> WindowsLookup: peer = WindowsLookup(tmp_path) diff --git a/tests/integrations/codex/job.py b/tests/integrations/codex/job.py new file mode 100644 index 000000000..525a02868 --- /dev/null +++ b/tests/integrations/codex/job.py @@ -0,0 +1,166 @@ +"""Control Job accounting independently of native process-exit signals.""" + +from __future__ import annotations + +import ctypes +from ctypes import wintypes +from typing import Any + +from band.integrations.codex.winspawn import ( + BasicAccountingInformation, + ExtendedLimitInformation, + ProcessIdList, +) + +JOB_HANDLE = 1 +FIRST_MEMBER_PID = 21 +SECOND_MEMBER_PID = 22 +PROCESS_HANDLES = {FIRST_MEMBER_PID: 121, SECOND_MEMBER_PID: 122} +KILL_ON_JOB_CLOSE = 0x2000 +UNRELATED_HANDLE = 999 + + +def handle_value(handle: ctypes.c_void_p | int) -> int: + if isinstance(handle, int): + return handle + assert handle.value is not None + return handle.value + + +class JobKernel: + def __init__(self) -> None: + self.pids = list(PROCESS_HANDLES) + self.signaled: set[int] = set() + self.opened: list[int] = [] + self.closed: list[int] = [] + self.terminated: list[int] = [] + self.termination_attempts: list[int] = [] + self.close_failures: dict[int, int] = {} + self.termination_failures = 0 + self.wait_fails = False + self.last_error = 0 + self.admission_sealed = False + self.fork_during_open = False + self.late_fork_rejected = False + self.recycled_pid: int | None = None + self.missing_pid: int | None = None + self.inaccessible_pid: int | None = None + self.limits = ExtendedLimitInformation() + self.limits.BasicLimitInformation.LimitFlags = KILL_ON_JOB_CLOSE + + def QueryInformationJobObject( + self, + job: ctypes.c_void_p | int, + information_class: int, + information: Any, + size: int, + returned: object, + ) -> bool: + assert handle_value(job) == JOB_HANDLE + match information_class: + case 1: + accounting = ctypes.cast( + information, ctypes.POINTER(BasicAccountingInformation) + ).contents + accounting.ActiveProcesses = 0 if self.terminated else len(self.pids) + return True + case 3: + header = ctypes.cast( + information, ctypes.POINTER(ProcessIdList) + ).contents + pids = [] if self.terminated else self.pids + capacity = (size - ProcessIdList.ProcessIdList.offset) // ctypes.sizeof( + ctypes.c_size_t + ) + header.NumberOfAssignedProcesses = len(pids) + header.NumberOfProcessIdsInList = min(len(pids), capacity) + address = ctypes.addressof(header) + ProcessIdList.ProcessIdList.offset + buffer = (ctypes.c_size_t * capacity).from_address(address) + for index, pid in enumerate(pids[:capacity]): + buffer[index] = pid + if len(pids) > capacity: + self.last_error = 234 + return False + return True + case 9: + assert size == ctypes.sizeof(self.limits) + ctypes.memmove(information, ctypes.byref(self.limits), size) + return True + case _: + raise AssertionError(f"Unexpected Job query class: {information_class}") + + def SetInformationJobObject( + self, + job: ctypes.c_void_p | int, + information_class: int, + information: Any, + size: int, + ) -> bool: + assert handle_value(job) == JOB_HANDLE + assert information_class == 9 + assert size == ctypes.sizeof(self.limits) + ctypes.memmove(ctypes.byref(self.limits), information, size) + basic = self.limits.BasicLimitInformation + self.admission_sealed = ( + bool(basic.LimitFlags & 8) and basic.ActiveProcessLimit == 0 + ) + return True + + def OpenProcess(self, rights: int, inherit: bool, pid: int) -> int | None: + assert rights & 0x00100000, "process handles must support waiting" + assert rights & 0x1000, "process handles must support Job membership queries" + assert not inherit + if self.fork_during_open: + self.fork_during_open = False + self.late_fork_rejected = self.admission_sealed + if not self.admission_sealed: + self.pids.append(23) + if pid == self.missing_pid: + self.last_error = 87 + return None + if pid == self.inaccessible_pid: + self.last_error = 5 + return None + handle = UNRELATED_HANDLE if pid == self.recycled_pid else PROCESS_HANDLES[pid] + self.opened.append(handle) + return handle + + def IsProcessInJob( + self, process: ctypes.c_void_p | int, job: ctypes.c_void_p | int, result: Any + ) -> bool: + assert handle_value(job) == JOB_HANDLE + ctypes.cast(result, ctypes.POINTER(wintypes.BOOL)).contents.value = ( + handle_value(process) != UNRELATED_HANDLE + ) + return True + + def TerminateJobObject(self, job: ctypes.c_void_p | int, code: int) -> bool: + assert code == 1 + self.termination_attempts.append(handle_value(job)) + if self.termination_failures: + self.termination_failures -= 1 + self.last_error = 5 + return False + self.terminated.append(handle_value(job)) + return True + + def WaitForSingleObject(self, process: ctypes.c_void_p | int, timeout: int) -> int: + assert timeout == 0, "native probes must not block the event loop" + if self.wait_fails: + self.last_error = 6 + return 0xFFFFFFFF + return 0 if handle_value(process) in self.signaled else 258 + + def CloseHandle(self, handle: ctypes.c_void_p | int) -> bool: + value = handle_value(handle) + remaining = self.close_failures.get(value, 0) + if remaining: + self.close_failures[value] = remaining - 1 + self.last_error = 5 + return False + self.closed.append(value) + return True + + def assert_handles_released(self) -> None: + """Retirement releases every handle acquired across failed attempts.""" + assert sorted(self.closed) == sorted([JOB_HANDLE, *self.opened]) diff --git a/tests/integrations/codex/test_cleanup.py b/tests/integrations/codex/test_cleanup.py index a312cf3fc..a53c44aa0 100644 --- a/tests/integrations/codex/test_cleanup.py +++ b/tests/integrations/codex/test_cleanup.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio +import ctypes import logging import os import shutil @@ -313,27 +314,19 @@ def deny(_name: str) -> object: async def test_idle_wait_preserves_ownership_until_explicit_release( self, active: int | None, drained: bool ) -> None: - job = Job.__new__(Job) - job._handle = object() # type: ignore[assignment] - job.name = "test" - closed = False + handle = ctypes.c_void_p(1) + job = Job(handle, "test") def active_processes() -> int | None: return active - def close() -> None: - nonlocal closed - closed = True - job._handle = None - job.active_processes = active_processes # type: ignore[method-assign] - job.close = close # type: ignore[method-assign] if drained: await job.wait_until_idle(0) else: with pytest.raises(JobNotDrained): await job.wait_until_idle(0) - assert not closed + assert job._handle is handle def test_windows_helper_does_not_import_a_room_package( self, tmp_path: Path diff --git a/tests/integrations/codex/test_job.py b/tests/integrations/codex/test_job.py index c801264cf..7f5053fbc 100644 --- a/tests/integrations/codex/test_job.py +++ b/tests/integrations/codex/test_job.py @@ -2,12 +2,22 @@ from __future__ import annotations +import asyncio import ctypes from types import SimpleNamespace import pytest from band.integrations.codex import winspawn +from tests.integrations.codex.job import ( + FIRST_MEMBER_PID, + JOB_HANDLE, + KILL_ON_JOB_CLOSE, + PROCESS_HANDLES, + SECOND_MEMBER_PID, + UNRELATED_HANDLE, + JobKernel, +) class TestJobHandleOwnership: @@ -32,25 +42,168 @@ def close_handle(current: ctypes.c_void_p) -> bool: job.close() assert attempted == [handle, handle] - def test_failed_termination_is_reported_without_releasing_the_job( - self, monkeypatch: pytest.MonkeyPatch + async def test_failed_termination_is_reported_without_releasing_the_job( + self, job_kernel: JobKernel ) -> None: - handle = ctypes.c_void_p(1) + handle = ctypes.c_void_p(JOB_HANDLE) job = winspawn.Job(handle, "test") - terminated: list[ctypes.c_void_p] = [] + job_kernel.termination_failures = 1 + with pytest.raises(OSError): + job.terminate() + assert not job_kernel.closed + job.terminate() + assert job_kernel.termination_attempts == [JOB_HANDLE, JOB_HANDLE] + job_kernel.signaled.update(PROCESS_HANDLES.values()) + await job.wait_until_idle(0) + job.close() + assert sorted(job_kernel.closed) == [JOB_HANDLE, *PROCESS_HANDLES.values()] - def terminate(current: ctypes.c_void_p, _code: int) -> bool: - terminated.append(current) - return len(terminated) > 1 + async def test_empty_accounting_does_not_release_pending_kernel_processes( + self, job_kernel: JobKernel + ) -> None: + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + assert job.active_processes() == 0 + with pytest.raises(winspawn.JobNotDrained): + await job.wait_until_idle(0) + assert not job_kernel.closed - monkeypatch.setattr( - winspawn, - "_kernel32", - lambda: SimpleNamespace(TerminateJobObject=terminate), - ) - monkeypatch.setattr(winspawn, "_last_error", lambda: 5) - monkeypatch.setattr(winspawn, "_win_os_error", lambda code: OSError(code)) + job_kernel.signaled.update(PROCESS_HANDLES.values()) + await job.wait_until_idle(0) + job.close() + assert sorted(job_kernel.closed) == [JOB_HANDLE, *PROCESS_HANDLES.values()] + + def test_admission_is_sealed_before_a_member_can_fork_during_capture( + self, job_kernel: JobKernel + ) -> None: + job_kernel.fork_during_open = True + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + assert job_kernel.late_fork_rejected + assert job_kernel.limits.BasicLimitInformation.LimitFlags & KILL_ON_JOB_CLOSE + + async def test_recycled_pid_is_not_waited_as_an_owned_process( + self, job_kernel: JobKernel + ) -> None: + job_kernel.recycled_pid = FIRST_MEMBER_PID + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + job_kernel.signaled.add(PROCESS_HANDLES[SECOND_MEMBER_PID]) + await job.wait_until_idle(0) + job.close() + assert sorted(job_kernel.closed) == [ + JOB_HANDLE, + PROCESS_HANDLES[SECOND_MEMBER_PID], + UNRELATED_HANDLE, + ] + + async def test_process_that_exits_during_handle_capture_does_not_block_cleanup( + self, job_kernel: JobKernel + ) -> None: + job_kernel.missing_pid = FIRST_MEMBER_PID + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + job_kernel.signaled.add(PROCESS_HANDLES[SECOND_MEMBER_PID]) + await job.wait_until_idle(0) + job.close() + assert sorted(job_kernel.closed) == [ + JOB_HANDLE, + PROCESS_HANDLES[SECOND_MEMBER_PID], + ] + + async def test_native_wait_failure_retains_process_handles_for_retry( + self, job_kernel: JobKernel + ) -> None: + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + job_kernel.wait_fails = True + with pytest.raises(OSError): + await job.wait_until_idle(0) + assert not job_kernel.closed + job_kernel.wait_fails = False + job_kernel.signaled.update(PROCESS_HANDLES.values()) + await job.wait_until_idle(0) + job.close() + assert sorted(job_kernel.closed) == [JOB_HANDLE, *PROCESS_HANDLES.values()] + + @pytest.mark.parametrize("failing_pid", [FIRST_MEMBER_PID, SECOND_MEMBER_PID]) + async def test_failed_member_handle_close_keeps_the_job_owned_for_retry( + self, job_kernel: JobKernel, failing_pid: int + ) -> None: + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + job_kernel.signaled.update(PROCESS_HANDLES.values()) + await job.wait_until_idle(0) + job_kernel.close_failures[PROCESS_HANDLES[failing_pid]] = 1 + with pytest.raises(OSError): + job.close() + assert JOB_HANDLE not in job_kernel.closed + job.close() + assert sorted(job_kernel.closed) == [JOB_HANDLE, *PROCESS_HANDLES.values()] + + async def test_unrelated_identity_handle_close_failure_is_retained_for_retry( + self, job_kernel: JobKernel + ) -> None: + job_kernel.recycled_pid = FIRST_MEMBER_PID + job_kernel.close_failures[UNRELATED_HANDLE] = 1 + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") with pytest.raises(OSError): job.terminate() + assert not job_kernel.closed + assert not job_kernel.terminated job.terminate() - assert terminated == [handle, handle] + job_kernel.signaled.add(PROCESS_HANDLES[SECOND_MEMBER_PID]) + await job.wait_until_idle(0) + job.close() + job_kernel.assert_handles_released() + + async def test_partial_capture_failure_retains_the_handles_already_opened( + self, job_kernel: JobKernel + ) -> None: + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job_kernel.inaccessible_pid = SECOND_MEMBER_PID + with pytest.raises(OSError): + job.terminate() + assert not job_kernel.closed + assert not job_kernel.terminated + job_kernel.inaccessible_pid = None + job.terminate() + job_kernel.signaled.update(PROCESS_HANDLES.values()) + await job.wait_until_idle(0) + job.close() + job_kernel.assert_handles_released() + + async def test_cancelled_native_exit_wait_keeps_handles_for_explicit_retry( + self, job_kernel: JobKernel, monkeypatch: pytest.MonkeyPatch + ) -> None: + entered = asyncio.Event() + native_wait = job_kernel.WaitForSingleObject + + def observed_wait(handle: ctypes.c_void_p | int, timeout: int) -> int: + entered.set() + return native_wait(handle, timeout) + + monkeypatch.setattr(job_kernel, "WaitForSingleObject", observed_wait) + job = winspawn.Job(ctypes.c_void_p(JOB_HANDLE), "test") + job.terminate() + waiting = asyncio.create_task(job.wait_until_idle(30)) + observed = asyncio.create_task(entered.wait()) + try: + done, _pending = await asyncio.wait( + [waiting, observed], return_when=asyncio.FIRST_COMPLETED + ) + assert observed in done + assert not waiting.done() + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + assert not job_kernel.closed + job_kernel.signaled.update(PROCESS_HANDLES.values()) + job.terminate() + await job.wait_until_idle(0) + job.close() + assert sorted(job_kernel.closed) == [JOB_HANDLE, *PROCESS_HANDLES.values()] + finally: + waiting.cancel() + observed.cancel() + await asyncio.gather(waiting, observed, return_exceptions=True) From 1687240d1e46eb7840134a49d6e07edf9cccdf0c Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 15:02:49 +0300 Subject: [PATCH 14/35] fix: keep native RPC close from racing stdout and extra-gated imports Close kills the child before draining stdout so a cancelled Windows pipe read cannot share the stream. Launch helpers no longer import the psutil-gated client at module load. Co-authored-by: Cursor --- src/band/integrations/ompprocess.py | 6 +- src/band/integrations/omprpc.py | 70 +++++++++++-------- .../ompprocess/test_mcp_plugin.py | 32 +++++++++ .../ompprocess/test_mcp_startup.py | 2 + tests/integrations/omprpc/kit.py | 7 +- .../omprpc/test_turn_settlement.py | 12 +++- 6 files changed, 90 insertions(+), 39 deletions(-) diff --git a/src/band/integrations/ompprocess.py b/src/band/integrations/ompprocess.py index 8059631e1..69bf09649 100644 --- a/src/band/integrations/ompprocess.py +++ b/src/band/integrations/ompprocess.py @@ -9,7 +9,7 @@ from contextlib import contextmanager from enum import StrEnum from pathlib import Path -from typing import Literal, Self +from typing import TYPE_CHECKING, Literal, Self from uuid import uuid4 from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -20,7 +20,6 @@ omp_cwd_argument, omp_provider_env, ) -from band.integrations.omprpc import OmpRpcClient from band.integrations.ompwire import ( OmpEvent, OmpEventType, @@ -30,6 +29,9 @@ OmpUIResponse, ) +if TYPE_CHECKING: + from band.integrations.omprpc import OmpRpcClient + DEFAULT_OMP_COMMAND: tuple[str, ...] = ("omp", "--mode", "rpc") # The native RPC ask dialog and MCP status events need a newer OMP than the diff --git a/src/band/integrations/omprpc.py b/src/band/integrations/omprpc.py index efa9ee165..e2ab41cae 100644 --- a/src/band/integrations/omprpc.py +++ b/src/band/integrations/omprpc.py @@ -110,6 +110,26 @@ def _command_error(response: OmpResponse) -> OmpRpcError: ) +def _ui_response_frame( + request_id: str, response: OmpUIResponse +) -> dict[str, JsonValue]: + frame: dict[str, JsonValue] = { + "type": OmpFrameType.UI_RESPONSE, + "id": request_id, + } + if response.value is not None: + frame["value"] = response.value + elif response.confirmed is not None: + frame["confirmed"] = response.confirmed + elif response.answers is not None: + frame["answers"] = [ + answer.model_dump(exclude_none=True) for answer in response.answers + ] + else: + frame["cancelled"] = True + return frame + + @dataclass class PendingChunks: chunk_id: str @@ -751,11 +771,9 @@ async def _dispatch_payload( if turn is None: write = asyncio.create_task( self._write_frame( - { - "type": OmpFrameType.UI_RESPONSE, - "id": request.id, - "cancelled": True, - } + _ui_response_frame( + request.id, OmpUIResponse(cancelled=True) + ) ) ) self._write_tasks.add(write) @@ -858,21 +876,7 @@ async def _handle_ui(self, turn: PromptRun, request: OmpUIRequest) -> None: response = await turn.on_ui(request) if not self._callback_still_active(turn, OmpCallbackKind.UI, request.id): return - frame: dict[str, JsonValue] = { - "type": OmpFrameType.UI_RESPONSE, - "id": request.id, - } - if response.value is not None: - frame["value"] = response.value - elif response.confirmed is not None: - frame["confirmed"] = response.confirmed - elif response.answers is not None: - frame["answers"] = [ - answer.model_dump(exclude_none=True) for answer in response.answers - ] - else: - frame["cancelled"] = True - await self._write_frame(frame) + await self._write_frame(_ui_response_frame(request.id, response)) async def _read_stderr(self, process: asyncio.subprocess.Process) -> None: assert process.stderr is not None @@ -947,6 +951,11 @@ async def _drain(stream: asyncio.StreamReader) -> None: while await stream.read(_DRAIN_BYTES): pass + @staticmethod + async def _await_contained_exit(process: asyncio.subprocess.Process) -> None: + if not await await_returncode(process, _PROCESS_EXIT_TIMEOUT_S): + raise OmpProcessError("OMP process did not exit after containment kill") + async def _settle_spawn(self) -> asyncio.subprocess.Process | None: spawning = self._spawn_task if spawning is None: @@ -997,11 +1006,10 @@ async def _close_process(self) -> None: await self._request_frame("abort", {}, allow_closing=True) except (TimeoutError, OmpRpcError, OmpProcessError) as error: logger.warning("OMP graceful abort skipped: %s", error) - if self._reader_task is not None: - await self._cancel_tasks([self._reader_task]) + reader = self._reader_task + if reader is not None and not reader.done(): + reader.cancel() drains: list[asyncio.Task[None]] = [] - if process.stdout is not None: - drains.append(asyncio.create_task(self._drain(process.stdout))) if ( self._stderr_task is None or self._stderr_task.done() ) and process.stderr is not None: @@ -1025,16 +1033,18 @@ async def _close_process(self) -> None: await asyncio.to_thread(scope.capture, process) await kill_process_tree(process) await asyncio.to_thread(scope.signal) - if not await await_returncode(process, _PROCESS_EXIT_TIMEOUT_S): - raise OmpProcessError( - "OMP process did not exit after containment kill" - ) + await self._await_contained_exit(process) await asyncio.to_thread(scope.close, process, _PROCESS_EXIT_TIMEOUT_S) - if not await await_returncode(process, _PROCESS_EXIT_TIMEOUT_S): - raise OmpProcessError("OMP process did not exit after containment kill") + await self._await_contained_exit(process) self._windows_process = None self._process_scope = None finally: + # A cancelled pipe read on Windows stays pending until the child is + # killed; drain stdout only after that reader finishes. + if reader is not None: + await self._cancel_tasks([reader]) + if process.stdout is not None: + drains.append(asyncio.create_task(self._drain(process.stdout))) if self._stderr_task is not None: drains.append(self._stderr_task) await self._cancel_tasks(drains) diff --git a/tests/integrations/ompprocess/test_mcp_plugin.py b/tests/integrations/ompprocess/test_mcp_plugin.py index bdda521d5..4b3b51ebf 100644 --- a/tests/integrations/ompprocess/test_mcp_plugin.py +++ b/tests/integrations/ompprocess/test_mcp_plugin.py @@ -3,6 +3,8 @@ from __future__ import annotations import json +import os +import stat from pathlib import Path import pytest @@ -76,6 +78,36 @@ def test_arguments_load_the_plugin_and_its_startup_extension( assert str(Path(plugin) / OMP_MCP_STARTUP_JS) in arguments assert str(Path(plugin) / OMP_MCP_STARTUP_CONFIG) in arguments + @pytest.mark.parametrize( + "accepted_on, file_mode, dir_mode", + [ + (("posix",), 0o600, 0o700), + (("nt",), None, None), + ], + ) + def test_projected_plugin_files_are_owner_only( + self, + tmp_path: Path, + accepted_on: tuple[str, ...], + file_mode: int | None, + dir_mode: int | None, + ) -> None: + servers = {"local": OmpMcpServer(command="./server")} + with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: + assert plugin is not None + root = Path(plugin) + files = ( + root / OMP_MCP_SERVERS_CONFIG, + root / OMP_MCP_STARTUP_CONFIG, + root / OMP_MCP_STARTUP_JS, + ) + assert root.is_dir() + assert all(path.is_file() for path in files) + if os.name not in accepted_on or file_mode is None or dir_mode is None: + return + assert stat.S_IMODE(root.stat().st_mode) == dir_mode + assert all(stat.S_IMODE(path.stat().st_mode) == file_mode for path in files) + @pytest.mark.parametrize( "fields", [ diff --git a/tests/integrations/ompprocess/test_mcp_startup.py b/tests/integrations/ompprocess/test_mcp_startup.py index 567d4e5b1..480b11826 100644 --- a/tests/integrations/ompprocess/test_mcp_startup.py +++ b/tests/integrations/ompprocess/test_mcp_startup.py @@ -9,6 +9,8 @@ import pytest +pytest.importorskip("psutil", reason="psutil not installed (band-sdk[omp])") + from band.integrations.ompprocess import ( OmpMcpServer, OmpStartupError, diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py index 7a08d736f..7ed2fb3ae 100644 --- a/tests/integrations/omprpc/kit.py +++ b/tests/integrations/omprpc/kit.py @@ -87,12 +87,7 @@ async def second_prompt(client: OmpRpcClient) -> Observer: observer = Observer() async with asyncio.timeout(5): await prompt_task(client, observer, "second") - assert any( - event.message is not None - and isinstance(event.message.content, list) - and any(content.text == "complete" for content in event.message.content) - for event in observer.events - ) + assert "complete" in observer.assistant_texts return observer diff --git a/tests/integrations/omprpc/test_turn_settlement.py b/tests/integrations/omprpc/test_turn_settlement.py index 398eb9495..8324ca9cd 100644 --- a/tests/integrations/omprpc/test_turn_settlement.py +++ b/tests/integrations/omprpc/test_turn_settlement.py @@ -14,7 +14,12 @@ OmpState, OmpToolResult, ) -from tests.integrations.omprpc.kit import ClientFactory, Observer, prompt_task +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + prompt_task, + second_prompt, +) class TestTurnSettlement: @@ -99,6 +104,8 @@ async def test_settlement_does_not_succeed_while_a_callback_is_still_running( ) -> None: monkeypatch.setattr(omprpc, "_CALLBACK_EXIT_TIMEOUT_S", 0) client = await client_factory("settled-during-callback") + process = client._proc + assert process is not None held = asyncio.Event() async def on_tool(_call: OmpHostToolCall) -> OmpToolResult: @@ -115,3 +122,6 @@ async def on_tool(_call: OmpHostToolCall) -> OmpToolResult: on_ui=observer.on_ui, ) held.set() + assert process.returncode is not None + await client.start() + await second_prompt(client) From 76ff1f7db40abdd8fe0a99c6e6ab979ed708ee6c Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 13:31:07 +0300 Subject: [PATCH 15/35] fix: preserve accepted-turn finalizers during shutdown retries --- docs/turn-outcome.md | 19 ++- src/band/runtime/execution.py | 9 +- .../adaptercleanup/test_work_settlement.py | 59 ++++++++ tests/runtime/executioninterrupt/kit.py | 7 + .../executioninterrupt/test_room_shutdown.py | 133 +++++++++++++++--- 5 files changed, 198 insertions(+), 29 deletions(-) diff --git a/docs/turn-outcome.md b/docs/turn-outcome.md index cc9054e25..d92d33f44 100644 --- a/docs/turn-outcome.md +++ b/docs/turn-outcome.md @@ -389,13 +389,18 @@ originating turn, including a worker released for approval. Leave and shutdown clear reporting while released worker cleanup remains with adapter residency. Provider execution and recovery timeouts retain their adapter policy and coordinate with that same owner. -Unfinished tasks remain owned after grace expires. They cannot post late replies, -change a newer turn's activity, or overlap replacement provider work. Room cleanup -reports `CleanupIncompleteError` while the worker still needs its resources; -a later existing host cleanup retry completes retirement. Worker completion alone -does not readmit the room. Successful approval release retains its original tools -and judgment until the worker actually finishes; `Turn.detach()` is still only -an outcome-ledger operation. +Unfinished tasks remain owned after grace expires. The stop attempt that first +retains a task does not cancel it again or renew its grace. Retained work cannot +post late replies, change a newer turn's activity, or overlap replacement provider +work. Room cleanup reports `CleanupIncompleteError` while work remains unfinished, +and the host keeps dependent adapter resources open. + +A later explicit host cleanup retry may cancel previously retained work again +within a fresh one-second grace, so a task that swallowed its first cancellation +can still settle. Once work settles, a cleanup retry completes retirement. Worker +completion alone does not readmit the room. Successful approval release retains +its original tools and judgment until the worker actually finishes; +`Turn.detach()` is still only an outcome-ledger operation. Stopping the work an external backend runs for the turn is the adapter's half, owned by one `BackendTurns` (`band.core.turn_lifecycle`) per adapter. The task diff --git a/src/band/runtime/execution.py b/src/band/runtime/execution.py index 499b65d6e..918814ded 100644 --- a/src/band/runtime/execution.py +++ b/src/band/runtime/execution.py @@ -633,6 +633,8 @@ async def stop(self, timeout: float | None = None) -> bool: return await asyncio.shield(self._shutdown_task) async def _stop(self, timeout: float | None) -> bool: + # Only work retained before this attempt may receive another cancellation. + retry_owners = tuple(self._detached_cycle_tasks.values()) if self._process_loop_task is not None: self._stop_graceful = True if timeout is not None and self.is_processing: @@ -647,7 +649,7 @@ async def _stop(self, timeout: float | None) -> bool: self._process_loop_task = None self._activity_owner = None await self._join_activity_completion() - await self._stop_retained_tasks(terminal_deadline()) + await self._stop_retained_tasks(retry_owners) await self._working_reporter.stop() return self._stop_graceful @@ -869,9 +871,10 @@ def _aborted_outcome(scope: TurnScope) -> TurnOutcome: case _: return TurnOutcome.PAUSED - async def _stop_retained_tasks(self, deadline: float) -> None: + async def _stop_retained_tasks(self, retry_owners: tuple[TurnTask, ...]) -> None: owners = tuple(self._detached_cycle_tasks.values()) - for owner in owners: + deadline = terminal_deadline() + for owner in retry_owners: owner.retry(deadline=deadline) for owner in owners: await owner.drain() diff --git a/tests/core/adaptercleanup/test_work_settlement.py b/tests/core/adaptercleanup/test_work_settlement.py index df24fd089..85cd97d52 100644 --- a/tests/core/adaptercleanup/test_work_settlement.py +++ b/tests/core/adaptercleanup/test_work_settlement.py @@ -7,12 +7,19 @@ import pytest from band.agent import Agent +from band.core.exceptions import CleanupIncompleteError +from band.core.turn_lifecycle import CYCLE_CANCEL_GRACE_SECONDS +from band.runtime.execution import ExecutionContext from band.runtime.oneshot import OneShotInvoker from tests.adaptercleanup import HostKind, host_for, resource_peer +from tests.adapters.waiting import spin from tests.conftest import make_message_event from tests.core.adaptercleanup.work import SuppressingWorkingAdapter, WorkingAdapter from tests.runtime.helpers import ROOM_ID from tests.runtime.oneshot.kit import msg_body +from tests.toolkit.clock import binary_clock_resolution +from tests.toolkit.tasks import background +from tests.toolkit.timecontrol import controlled_clock class TestWorkSettlement: @@ -60,3 +67,55 @@ async def test_finalizer_keeps_resource_until_accepted_work_settles( await stopping assert adapter.finished.is_set() assert client.is_closed + + @pytest.mark.looptime(resolution=binary_clock_resolution()) + async def test_stop_retains_finalizer_whose_grace_expires_during_shutdown( + self, + ) -> None: + async with resource_peer() as client: + adapter = WorkingAdapter(client) + adapter.add_cleanup(client.aclose) + async with ( + controlled_clock() as clock, + host_for(adapter, HostKind.AGENT) as host, + ): + await host.start() + assert isinstance(host.owner, Agent) + runtime = host.owner._runtime.runtime + host.platform.add_message("work") + await runtime._on_room_joined(ROOM_ID, {}) + await runtime._on_room_event( + ROOM_ID, make_message_event(room_id=ROOM_ID, msg_id="work") + ) + await adapter.entered.wait() + execution = runtime.executions[ROOM_ID] + assert isinstance(execution, ExecutionContext) + scope = execution.current_scope + assert scope is not None and scope.handler is not None + work = scope.handler.task + execution.interrupt() + await adapter.finalizing.wait() + remaining_grace = CYCLE_CANCEL_GRACE_SECONDS / 2 + await clock.elapse(remaining_grace) + + async with background(host.stop()) as stopping: + try: + await spin() + assert scope.shutdown + await clock.elapse(remaining_grace) + with pytest.raises(ExceptionGroup) as raised: + await stopping + assert raised.value.subgroup(CleanupIncompleteError) is not None + assert work.cancelling() == 1 + assert not work.done() + assert not adapter.finished.is_set() + assert runtime.unsettled_rooms == [ROOM_ID] + assert not client.is_closed + assert (await client.get("/resource")).status_code == 200 + finally: + adapter.release.set() + await asyncio.gather(work, stopping, return_exceptions=True) + await host.stop() + + assert adapter.finished.is_set() + assert client.is_closed diff --git a/tests/runtime/executioninterrupt/kit.py b/tests/runtime/executioninterrupt/kit.py index 9ebd388f5..8b3686fa6 100644 --- a/tests/runtime/executioninterrupt/kit.py +++ b/tests/runtime/executioninterrupt/kit.py @@ -10,6 +10,7 @@ import asyncio from collections.abc import Awaitable, Callable +from enum import StrEnum from typing import Any import pytest @@ -24,6 +25,12 @@ from tests.runtime.helpers import AGENT_ID, ROOM_ID +class ShutdownStep(StrEnum): + STOP = "stop" + JOIN_STOP = "concurrent stop" + SETTLE = "settle accepted work" + + async def assert_fresh_cycle_still_propagates_shutdown_cancel( ctx: ExecutionContext, msg_id: str ) -> None: diff --git a/tests/runtime/executioninterrupt/test_room_shutdown.py b/tests/runtime/executioninterrupt/test_room_shutdown.py index d048dbcbd..26625ba30 100644 --- a/tests/runtime/executioninterrupt/test_room_shutdown.py +++ b/tests/runtime/executioninterrupt/test_room_shutdown.py @@ -7,14 +7,26 @@ from unittest.mock import AsyncMock import pytest +from hypothesis import example, given, settings +from hypothesis import strategies as st from band.core.exceptions import CleanupIncompleteError -from band.core.turn_lifecycle import TurnTask, bind_turn_task +from band.core.turn_lifecycle import ( + CYCLE_CANCEL_GRACE_SECONDS, + TurnTask, + bind_turn_task, +) from band.runtime.cycle import TurnScope from band.runtime.execution import ExecutionContext from band.runtime.tools.agent import AgentTools +from tests.adaptercleanup import CleanupPlatform +from tests.adapters.waiting import spin from tests.conftest import BlockingHandler, make_message_event -from tests.runtime.helpers import resist_cancellation_until +from tests.runtime.executioninterrupt.kit import ShutdownStep, auxiliary_context +from tests.runtime.helpers import resist_cancellation_until, rest_client_over +from tests.toolkit.clock import binary_clock_resolution +from tests.toolkit.tasks import background +from tests.toolkit.timecontrol import controlled_clock class TestRoomShutdown: @@ -83,33 +95,43 @@ async def observer(ctx: ExecutionContext, scope: TurnScope) -> None: await ctx.stop() assert cancelled.is_set() + @pytest.mark.looptime(resolution=binary_clock_resolution()) async def test_retried_stop_cancels_an_observer_that_swallowed_the_first( - self, mock_link, monkeypatch + self, mock_link ) -> None: - monkeypatch.setattr("band.core.turn_lifecycle.CYCLE_CANCEL_GRACE_SECONDS", 0.01) - entered, cancels = asyncio.Event(), 0 + entered, cancelled = asyncio.Event(), asyncio.Event() async def observer(ctx: ExecutionContext, scope: TurnScope) -> None: - nonlocal cancels + task = asyncio.current_task() + assert task is not None entered.set() while True: try: - await asyncio.Future() + await asyncio.Future[None]() except asyncio.CancelledError: - cancels += 1 - if cancels == 2: + cancelled.set() + if task.cancelling() == 2: raise - ctx = ExecutionContext( - "room-123", mock_link, AsyncMock(), on_platform_stop=observer - ) - await ctx._begin_scope() - ctx.observe_platform_stop(ctx.current_scope) - await entered.wait() - with pytest.raises(CleanupIncompleteError): - await ctx.stop() - await ctx.stop() - assert cancels == 2 + async with controlled_clock() as clock: + ctx = ExecutionContext( + "room-123", mock_link, AsyncMock(), on_platform_stop=observer + ) + scope = await ctx._begin_scope() + assert scope is not None + ctx.observe_platform_stop(scope) + assert scope.observer_task is not None + async with background(scope.observer_task) as observing: + await entered.wait() + async with background(ctx.stop()) as stopping: + await cancelled.wait() + await clock.elapse(CYCLE_CANCEL_GRACE_SECONDS) + with pytest.raises(CleanupIncompleteError): + await stopping + assert observing.cancelling() == 1 + await ctx.stop() + assert observing.cancelled() + assert observing.cancelling() == 2 async def test_interrupted_drain_still_keeps_the_stubborn_task_for_stop( self, mock_link @@ -213,3 +235,76 @@ async def report_activity( await started.wait() assert await ctx.stop() is True assert false_sent.is_set() + + @pytest.mark.looptime(resolution=binary_clock_resolution()) + @settings(max_examples=20, deadline=None) + @given(steps=st.lists(st.sampled_from(ShutdownStep), min_size=1, max_size=6)) + @example(steps=[ShutdownStep.STOP, ShutdownStep.STOP, ShutdownStep.SETTLE]) + @example(steps=[ShutdownStep.JOIN_STOP, ShutdownStep.SETTLE, ShutdownStep.STOP]) + async def test_generated_shutdown_steps_preserve_cancellation_ownership( + self, steps: list[ShutdownStep] + ) -> None: + started, finalizing, finish, finished = (asyncio.Event() for _ in range(4)) + + async def handler(ctx: ExecutionContext, event: Any) -> None: + started.set() + try: + await finish.wait() + finally: + finalizing.set() + while not finish.is_set(): + try: + await finish.wait() + except asyncio.CancelledError: + pass + finished.set() + + platform = CleanupPlatform() + platform.add_message("work") + async with ( + controlled_clock() as clock, + rest_client_over(platform.answer) as rest, + ): + ctx = auxiliary_context(rest, handler) + await ctx.start() + await ctx.on_event(make_message_event(msg_id="work")) + await started.wait() + scope = ctx.current_scope + assert scope is not None and scope.handler is not None + work = scope.handler.task + cancellations = 0 + settled = False + stops: list[asyncio.Task[bool]] = [] + try: + for step in steps: + if step is ShutdownStep.SETTLE: + finish.set() + await asyncio.gather(work, return_exceptions=True) + settled = True + assert finished.is_set() + continue + attempt = [asyncio.create_task(ctx.stop())] + if step is ShutdownStep.JOIN_STOP: + attempt.append(asyncio.create_task(ctx.stop())) + stops.extend(attempt) + if settled: + assert all(await asyncio.gather(*attempt)) + continue + if cancellations == 0: + await finalizing.wait() + else: + await spin() + cancellations += 1 + await clock.elapse(CYCLE_CANCEL_GRACE_SECONDS) + assert work.cancelling() == cancellations + for stopping in attempt: + with pytest.raises(CleanupIncompleteError): + await stopping + assert not work.done() + assert not finished.is_set() + assert not ctx.is_running + finally: + finish.set() + await asyncio.gather(work, *stops, return_exceptions=True) + await ctx.stop() + assert finished.is_set() From 10565fc9887780bb05a7a22258214d6081f7ffd7 Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 19:29:01 +0300 Subject: [PATCH 16/35] test: synchronize shutdown retries and release observer cleanup --- .../executioninterrupt/test_room_shutdown.py | 45 ++++++++++--------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/tests/runtime/executioninterrupt/test_room_shutdown.py b/tests/runtime/executioninterrupt/test_room_shutdown.py index 26625ba30..c21d6d2e0 100644 --- a/tests/runtime/executioninterrupt/test_room_shutdown.py +++ b/tests/runtime/executioninterrupt/test_room_shutdown.py @@ -20,7 +20,6 @@ from band.runtime.execution import ExecutionContext from band.runtime.tools.agent import AgentTools from tests.adaptercleanup import CleanupPlatform -from tests.adapters.waiting import spin from tests.conftest import BlockingHandler, make_message_event from tests.runtime.executioninterrupt.kit import ShutdownStep, auxiliary_context from tests.runtime.helpers import resist_cancellation_until, rest_client_over @@ -99,15 +98,15 @@ async def observer(ctx: ExecutionContext, scope: TurnScope) -> None: async def test_retried_stop_cancels_an_observer_that_swallowed_the_first( self, mock_link ) -> None: - entered, cancelled = asyncio.Event(), asyncio.Event() + entered, cancelled, release = (asyncio.Event() for _ in range(3)) async def observer(ctx: ExecutionContext, scope: TurnScope) -> None: task = asyncio.current_task() assert task is not None entered.set() - while True: + while not release.is_set(): try: - await asyncio.Future[None]() + await release.wait() except asyncio.CancelledError: cancelled.set() if task.cancelling() == 2: @@ -122,16 +121,20 @@ async def observer(ctx: ExecutionContext, scope: TurnScope) -> None: ctx.observe_platform_stop(scope) assert scope.observer_task is not None async with background(scope.observer_task) as observing: - await entered.wait() - async with background(ctx.stop()) as stopping: - await cancelled.wait() - await clock.elapse(CYCLE_CANCEL_GRACE_SECONDS) - with pytest.raises(CleanupIncompleteError): - await stopping - assert observing.cancelling() == 1 - await ctx.stop() - assert observing.cancelled() - assert observing.cancelling() == 2 + try: + await entered.wait() + async with background(ctx.stop()) as stopping: + await cancelled.wait() + await clock.elapse(CYCLE_CANCEL_GRACE_SECONDS) + with pytest.raises(CleanupIncompleteError): + await stopping + assert observing.cancelling() == 1 + await ctx.stop() + assert observing.cancelled() + assert observing.cancelling() == 2 + finally: + release.set() + await ctx.stop() async def test_interrupted_drain_still_keeps_the_stubborn_task_for_stop( self, mock_link @@ -244,19 +247,22 @@ async def report_activity( async def test_generated_shutdown_steps_preserve_cancellation_ownership( self, steps: list[ShutdownStep] ) -> None: - started, finalizing, finish, finished = (asyncio.Event() for _ in range(4)) + started, finish, finished = (asyncio.Event() for _ in range(3)) + cancellation_received: asyncio.Queue[None] = asyncio.Queue() async def handler(ctx: ExecutionContext, event: Any) -> None: started.set() try: await finish.wait() + except asyncio.CancelledError: + cancellation_received.put_nowait(None) + raise finally: - finalizing.set() while not finish.is_set(): try: await finish.wait() except asyncio.CancelledError: - pass + cancellation_received.put_nowait(None) finished.set() platform = CleanupPlatform() @@ -290,10 +296,7 @@ async def handler(ctx: ExecutionContext, event: Any) -> None: if settled: assert all(await asyncio.gather(*attempt)) continue - if cancellations == 0: - await finalizing.wait() - else: - await spin() + await cancellation_received.get() cancellations += 1 await clock.elapse(CYCLE_CANCEL_GRACE_SECONDS) assert work.cancelling() == cancellations From 67bbbd81094809a65c64417e0f2f1f4ba6415412 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 16:06:54 +0300 Subject: [PATCH 17/35] test: budget the complete native RPC lifecycle workload --- tests/integrations/omprpc/test_lifecycle_model.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/integrations/omprpc/test_lifecycle_model.py b/tests/integrations/omprpc/test_lifecycle_model.py index f34afb0a8..82abf3854 100644 --- a/tests/integrations/omprpc/test_lifecycle_model.py +++ b/tests/integrations/omprpc/test_lifecycle_model.py @@ -12,8 +12,9 @@ class TestRpcLifecycleModel: @pytest.mark.parametrize("detached", [False, True], ids=["session", "detached"]) - # The generated workload launches hundreds of real peers; each step stays guarded. - @pytest.mark.timeout(120) + # The full Hypothesis workload can exceed two minutes on Windows. + # Each generated step retains its separate, shorter hang guard. + @pytest.mark.timeout(600) @settings( max_examples=20, deadline=None, From 330e8ab88969ce34ff89821024ca2ba15f9716c5 Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 16:25:24 +0300 Subject: [PATCH 18/35] refactor: split native RPC into single-concern modules and intent-led tests Keep public imports stable while giving spawn and timeout public seams so tests assert outcomes instead of patching internals or waiting out grace. Co-authored-by: Cursor --- src/band/converters/omp.py | 29 +- src/band/integrations/omprpc/__init__.py | 19 + .../{omprpc.py => omprpc/client.py} | 390 +++--------------- src/band/integrations/omprpc/errors.py | 43 ++ src/band/integrations/omprpc/frames.py | 122 ++++++ src/band/integrations/omprpc/prompt.py | 85 ++++ src/band/integrations/omprpc/scope.py | 101 +++++ tests/converters/test_omp.py | 11 +- tests/integrations/omppeer.py | 65 +-- .../ompprocess/test_mcp_plugin.py | 27 +- tests/integrations/omprpc/kit.py | 24 +- tests/integrations/omprpc/lifecycle.py | 257 +++++++----- .../omprpc/test_event_backpressure.py | 196 +++++---- .../omprpc/test_lifecycle_model.py | 16 +- .../integrations/omprpc/test_process_tree.py | 10 +- .../omprpc/test_turn_settlement.py | 41 +- 16 files changed, 847 insertions(+), 589 deletions(-) create mode 100644 src/band/integrations/omprpc/__init__.py rename src/band/integrations/{omprpc.py => omprpc/client.py} (69%) create mode 100644 src/band/integrations/omprpc/errors.py create mode 100644 src/band/integrations/omprpc/frames.py create mode 100644 src/band/integrations/omprpc/prompt.py create mode 100644 src/band/integrations/omprpc/scope.py diff --git a/src/band/converters/omp.py b/src/band/converters/omp.py index 457564a82..331bedcf8 100644 --- a/src/band/converters/omp.py +++ b/src/band/converters/omp.py @@ -2,6 +2,7 @@ from __future__ import annotations +from enum import StrEnum from typing import Any from pydantic import BaseModel, Field @@ -10,9 +11,16 @@ from band.core.protocols import HistoryConverter from band.core.types import MessageType -# Task-event metadata keys the native adapter writes and this converter reads. -OMP_ROOM_METADATA_KEY = "omp_room_id" -OMP_SESSION_METADATA_KEY = "omp_session_file" + +class OmpMetadataKey(StrEnum): + """Task-event metadata keys the native adapter writes and this converter reads.""" + + ROOM = "omp_room_id" + SESSION = "omp_session_file" + + +OMP_ROOM_METADATA_KEY = OmpMetadataKey.ROOM +OMP_SESSION_METADATA_KEY = OmpMetadataKey.SESSION class OmpSessionState(BaseModel): @@ -28,15 +36,12 @@ class OmpHistoryConverter(HistoryConverter[OmpSessionState]): def convert(self, raw: list[dict[str, Any]]) -> OmpSessionState: sessions: dict[str, str] = {} for message in raw: - if message.get("message_type") != MessageType.TASK: - continue - metadata = message.get("metadata") - if not isinstance(metadata, dict): - continue - room = metadata.get(OMP_ROOM_METADATA_KEY) - session = metadata.get(OMP_SESSION_METADATA_KEY) - if isinstance(room, str) and isinstance(session, str) and session: - sessions[room] = session + match message.get("message_type"), message.get("metadata"): + case MessageType.TASK, { + OmpMetadataKey.ROOM: str(room), + OmpMetadataKey.SESSION: str(session), + } if session: + sessions[room] = session return OmpSessionState( room_to_session=sessions, replay_messages=build_replay_messages(raw) ) diff --git a/src/band/integrations/omprpc/__init__.py b/src/band/integrations/omprpc/__init__.py new file mode 100644 index 000000000..5df95e812 --- /dev/null +++ b/src/band/integrations/omprpc/__init__.py @@ -0,0 +1,19 @@ +"""Persistent subprocess transport for OMP's native JSONL RPC protocol.""" + +from __future__ import annotations + +from band.integrations.omprpc.client import OmpRpcClient +from band.integrations.omprpc.errors import ( + OmpProcessError, + OmpPromptError, + OmpProtocolError, + OmpRpcError, +) + +__all__ = [ + "OmpProcessError", + "OmpPromptError", + "OmpProtocolError", + "OmpRpcClient", + "OmpRpcError", +] diff --git a/src/band/integrations/omprpc.py b/src/band/integrations/omprpc/client.py similarity index 69% rename from src/band/integrations/omprpc.py rename to src/band/integrations/omprpc/client.py index e2ab41cae..f6633dc65 100644 --- a/src/band/integrations/omprpc.py +++ b/src/band/integrations/omprpc/client.py @@ -3,19 +3,14 @@ from __future__ import annotations import asyncio -import base64 import logging import os -import time from binascii import Error as Base64Error from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass -from enum import StrEnum from typing import TypeVar, cast -from uuid import uuid4 -import psutil -from pydantic import JsonValue, TypeAdapter, ValidationError +from pydantic import JsonValue, ValidationError from band.integrations.codex.winspawn import ( SpawnFailure, @@ -23,15 +18,27 @@ await_returncode, close_stdin, kill_process_tree, - session_pid_reused, ) +from band.integrations.omprpc.errors import ( + CALLBACKS_UNFINISHED, + OmpProcessError, + OmpPromptError, + OmpProtocolError, + OmpRpcError, + command_error, +) +from band.integrations.omprpc.frames import ( + ChunkAssembler, + parse_payload_line, + payload_json, + ui_response_frame, +) +from band.integrations.omprpc.prompt import OmpCallbackKind, PromptRun +from band.integrations.omprpc.scope import OmpProcessScope from band.integrations.ompwire import ( - OMP_ASSISTANT_ROLE, - OMP_CHUNK_BYTES, OMP_FRAME_BYTES, OMP_NATIVE_PROTOCOL_VERSION, OMP_REASSEMBLED_BYTES, - OmpChunk, OmpCommand, OmpEvent, OmpEventType, @@ -40,7 +47,6 @@ OmpFrameType, OmpHostToolCall, OmpHostToolCancel, - OmpMessage, OmpNegotiation, OmpPromptAcceptance, OmpPromptResult, @@ -55,91 +61,20 @@ ) TModel = TypeVar("TModel", bound=OmpWireModel) +CreateProcess = Callable[..., Awaitable[asyncio.subprocess.Process]] logger = logging.getLogger(__name__) _DRAIN_BYTES = 64 * 1024 -_EVENT_BACKLOG_LIMIT = 256 _ABORT_TIMEOUT_S = 1.0 -_PROCESS_EXIT_TIMEOUT_S = 1.0 -_CALLBACK_EXIT_TIMEOUT_S = 1.0 -_PROCESS_SCOPE_VALUE = "1" -_PAYLOAD_ADAPTER: TypeAdapter[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]) +PROCESS_EXIT_TIMEOUT_S = 1.0 +CALLBACK_EXIT_TIMEOUT_S = 1.0 _EVENT_TYPES = frozenset(OmpEventType) _UI_NOTIFICATIONS = frozenset( {"notify", "setStatus", "setWidget", "setTitle", "set_editor_text", "open_url"} ) -class OmpCallbackKind(StrEnum): - TOOL = "tool" - UI = "ui" - - -class OmpRpcError(RuntimeError): - """An unsuccessful command response; the transport remains usable.""" - - def __init__(self, command: str, error: str, *, code: str | None = None) -> None: - self.command = command - self.error = error - self.code = code - super().__init__(f"OMP {command}: {error}") - - -class OmpProcessError(RuntimeError): - """The OMP process or its stdio transport is no longer usable.""" - - -class OmpProtocolError(OmpProcessError): - """An invalid or incorrectly correlated native RPC frame.""" - - -class OmpPromptError(RuntimeError): - """The prompt or its final background continuation failed or was aborted.""" - - def __init__(self, status: str, message: str) -> None: - self.status = status - super().__init__(message) - - -def _command_error(response: OmpResponse) -> OmpRpcError: - return OmpRpcError( - response.command, - response.error or "Command failed", - code=response.code, - ) - - -def _ui_response_frame( - request_id: str, response: OmpUIResponse -) -> dict[str, JsonValue]: - frame: dict[str, JsonValue] = { - "type": OmpFrameType.UI_RESPONSE, - "id": request_id, - } - if response.value is not None: - frame["value"] = response.value - elif response.confirmed is not None: - frame["confirmed"] = response.confirmed - elif response.answers is not None: - frame["answers"] = [ - answer.model_dump(exclude_none=True) for answer in response.answers - ] - else: - frame["cancelled"] = True - return frame - - -@dataclass -class PendingChunks: - chunk_id: str - count: int - byte_length: int - next_index: int - received_bytes: int - parts: list[bytes] - - @dataclass(frozen=True) class ReceivedResponse: frame: OmpResponse @@ -152,159 +87,23 @@ class PendingRequest: future: asyncio.Future[ReceivedResponse | OmpRpcError | OmpProcessError] -class PromptRun: - def __init__( - self, - request_id: str, - on_event: Callable[[OmpEvent], Awaitable[None]], - on_tool: Callable[[OmpHostToolCall], Awaitable[OmpToolResult]], - on_ui: Callable[[OmpUIRequest], Awaitable[OmpUIResponse]], - ) -> None: - loop = asyncio.get_running_loop() - self.id = request_id - self.on_event = on_event - self.on_tool = on_tool - self.on_ui = on_ui - self.events: asyncio.Queue[OmpEvent] = asyncio.Queue( - maxsize=_EVENT_BACKLOG_LIMIT - ) - self.failure: asyncio.Future[BaseException] = loop.create_future() - self.result: asyncio.Future[OmpPromptResult] = loop.create_future() - self.changed = asyncio.Event() - self.last_activity_sequence = 0 - self.settled_sequence = 0 - self.last_assistant: OmpMessage | None = None - self.rejected = False - self.callbacks: dict[tuple[OmpCallbackKind, str], asyncio.Task[None]] = {} - self.callback_tasks: set[asyncio.Task[None]] = set() - self.event_worker: asyncio.Task[None] | None = None - - def fail(self, error: BaseException) -> None: - if not self.failure.done(): - self.failure.set_result(error) - - def observe(self, event: OmpEvent, sequence: int) -> None: - if self.failure.done(): - return - # Waiting for space would block the reader that services reentrant commands. - try: - self.events.put_nowait(event) - except asyncio.QueueFull: - self.fail(OmpProcessError("OMP event callback backlog exceeded")) - return - self.last_activity_sequence = sequence - match event.type: - case OmpEventType.MESSAGE_END if ( - event.message is not None and event.message.role == OMP_ASSISTANT_ROLE - ): - self.last_assistant = event.message - case OmpEventType.AGENT_END: - self.last_assistant = next( - ( - message - for message in reversed(event.messages) - if message.role == OMP_ASSISTANT_ROLE - ), - self.last_assistant, - ) - self.changed.set() - - -class OmpProcessScope: - """Retain observable descendants for best-effort cleanup. - - The private PI_ marker survives native helper filters. Explicit environment - replacement and restricted process inspection can hide detached children. - Windows job membership is the primary containment; this marker catches - descendants the job does not observe. - """ - - def __init__(self) -> None: - self.key = f"PI_BAND_PROCESS_{uuid4().hex}" - self.processes: set[psutil.Process] = set() - - def environment(self, base: Mapping[str, str] | None) -> dict[str, str]: - return { - **(os.environ if base is None else base), - self.key: _PROCESS_SCOPE_VALUE, - } - - def capture(self, root: asyncio.subprocess.Process) -> None: - group_visible = not session_pid_reused(root) - group_members: list[psutil.Process] = [] - getpgid = getattr(os, "getpgid", None) - for pid in psutil.pids(): - # asyncio's child watcher remains the sole waiter for the root. - if pid == 0 or pid == root.pid: - continue - try: - member = psutil.Process(pid) - created = member.create_time() - marked = False - try: - marked = member.environ().get(self.key) == _PROCESS_SCOPE_VALUE - except (psutil.NoSuchProcess, psutil.AccessDenied): - pass - in_group = False - if group_visible and getpgid is not None: - try: - in_group = getpgid(pid) == root.pid - except (ProcessLookupError, PermissionError): - pass - if not marked and not in_group: - continue - current = psutil.Process(pid) - if current.create_time() != created or current != member: - continue - if marked: - self.processes.add(member) - elif in_group: - group_members.append(member) - except (psutil.NoSuchProcess, psutil.AccessDenied): - continue - if group_visible and not session_pid_reused(root): - self.processes.update(group_members) - - @staticmethod - def _is_alive(process: psutil.Process) -> bool: - try: - return process.is_running() and process.status() != psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return False - - def signal(self) -> None: - for process in self.processes: - try: - process.kill() - except psutil.NoSuchProcess: - pass - - def close(self, root: asyncio.subprocess.Process, timeout: float) -> None: - deadline = time.monotonic() + timeout - self.capture(root) - alive = [process for process in self.processes if self._is_alive(process)] - while alive: - self.signal() - psutil.wait_procs(alive, timeout=max(0.0, deadline - time.monotonic())) - self.capture(root) - alive = [process for process in self.processes if self._is_alive(process)] - if alive and time.monotonic() >= deadline: - raise OmpProcessError( - "OMP owned session members survived cleanup: " - f"{sorted(process.pid for process in alive)}" - ) - - class OmpRpcClient: """One OMP process, one persistent reader, and at most one active prompt.""" def __init__( - self, command: Sequence[str], *, env: Mapping[str, str] | None = None + self, + command: Sequence[str], + *, + env: Mapping[str, str] | None = None, + create_process: CreateProcess | None = None, ) -> None: if not command: raise ValueError("OMP command must not be empty") self.command = tuple(command) self.env = None if env is None else {**os.environ, **env} + self.create_process = create_process + self.callback_exit_timeout_s = CALLBACK_EXIT_TIMEOUT_S + self.process_exit_timeout_s = PROCESS_EXIT_TIMEOUT_S self._proc: asyncio.subprocess.Process | None = None self._windows_process: WindowsProcess | None = None self._process_scope: OmpProcessScope | None = None @@ -317,9 +116,8 @@ def __init__( self._closing = False self._request_id = 0 self._sequence = 0 - self._chunks: PendingChunks | None = None + self._frames = ChunkAssembler() self._max_frame_bytes = OMP_FRAME_BYTES - self._max_reassembled_bytes = OMP_REASSEMBLED_BYTES self._pending: dict[str, PendingRequest] = {} self._turn: PromptRun | None = None self._start_lock = asyncio.Lock() @@ -343,9 +141,9 @@ async def start(self) -> None: "Discarding previous OMP RPC failure on restart: %s", self._failure ) self._failure = None - self._chunks = None + self._frames.reset() + self._frames.max_reassembled_bytes = OMP_REASSEMBLED_BYTES self._max_frame_bytes = OMP_FRAME_BYTES - self._max_reassembled_bytes = OMP_REASSEMBLED_BYTES self._closing = False self._process_scope = OmpProcessScope() self._ready = asyncio.get_running_loop().create_future() @@ -365,7 +163,7 @@ async def start(self) -> None: f"OMP RPC protocol v{OMP_NATIVE_PROTOCOL_VERSION} is required" ) self._max_frame_bytes = ready.maxFrameBytes - self._max_reassembled_bytes = ready.maxReassembledFrameBytes + self._frames.max_reassembled_bytes = ready.maxReassembledFrameBytes negotiated = await self._request_frame( "negotiate_protocol", {"protocolVersion": OMP_NATIVE_PROTOCOL_VERSION}, @@ -385,23 +183,22 @@ async def start(self) -> None: async def _spawn_process(self) -> asyncio.subprocess.Process: scope = self._process_scope assert scope is not None - if os.name == "nt": + env = scope.environment(self.env) + if os.name == "nt" and self.create_process is None: self._windows_process = WindowsProcess() process = await self._windows_process.start_command( - self.command, - cwd=None, - env=scope.environment(self.env), - limit=OMP_FRAME_BYTES, + self.command, cwd=None, env=env, limit=OMP_FRAME_BYTES ) else: - process = await asyncio.create_subprocess_exec( + create = self.create_process or asyncio.create_subprocess_exec + process = await create( *self.command, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=None, - env=scope.environment(self.env), - start_new_session=True, + env=env, + start_new_session=os.name != "nt", limit=OMP_FRAME_BYTES, ) self._proc = process @@ -535,7 +332,10 @@ async def _execute_prompt( raise finally: self._turn = None - await self._cancel_tasks([work, event_worker, *turn.callback_tasks]) + await self._cancel_tasks( + [work, event_worker, *turn.callback_tasks], + timeout=self.callback_exit_timeout_s, + ) async def _complete_prompt(self, turn: PromptRun, message: str) -> None: try: @@ -558,10 +358,10 @@ async def _complete_prompt(self, turn: PromptRun, message: str) -> None: await turn.events.join() if turn.callback_tasks: _, pending = await asyncio.wait( - turn.callback_tasks, timeout=_CALLBACK_EXIT_TIMEOUT_S + turn.callback_tasks, timeout=self.callback_exit_timeout_s ) if pending: - raise OmpProcessError("OMP callbacks did not finish after settlement") + raise OmpProcessError(CALLBACKS_UNFINISHED) if outcome is not None and outcome.status != OmpPromptStatus.COMPLETED: raise OmpPromptError( outcome.status, @@ -622,60 +422,7 @@ async def _write_line(self, line: str) -> None: raise error from cause async def _write_frame(self, frame: dict[str, JsonValue]) -> None: - await self._write_line(_PAYLOAD_ADAPTER.dump_json(frame).decode("utf-8")) - - def _decode_payload( - self, payload: dict[str, JsonValue] - ) -> dict[str, JsonValue] | None: - frame = OmpFrame.model_validate(payload) - if frame.type != OmpFrameType.CHUNK: - if self._chunks is not None: - raise OmpProtocolError("OMP RPC protocol chunk sequence interrupted") - return payload - chunk = OmpChunk.model_validate(payload) - if chunk.byteLength > self._max_reassembled_bytes: - raise OmpProtocolError( - "OMP RPC protocol frame exceeds the negotiated limit" - ) - data = base64.b64decode(chunk.data, validate=True) - if ( - not data - or len(data) > OMP_CHUNK_BYTES - or base64.b64encode(data).decode("ascii") != chunk.data - ): - raise OmpProtocolError("Invalid OMP RPC protocol chunk data") - if self._chunks is None: - if chunk.index != 0: - raise OmpProtocolError( - "OMP RPC protocol chunks must start at index zero" - ) - self._chunks = PendingChunks( - chunk.chunkId, chunk.count, chunk.byteLength, 0, 0, [] - ) - pending = self._chunks - if ( - pending.chunk_id != chunk.chunkId - or pending.count != chunk.count - or pending.byte_length != chunk.byteLength - or pending.next_index != chunk.index - ): - raise OmpProtocolError("OMP RPC protocol chunk sequence mismatch") - pending.parts.append(data) - pending.received_bytes += len(data) - pending.next_index += 1 - if pending.received_bytes > pending.byte_length: - raise OmpProtocolError( - "OMP RPC protocol chunks exceed their declared length" - ) - if pending.next_index < pending.count: - return None - if pending.received_bytes != pending.byte_length: - raise OmpProtocolError("OMP RPC protocol chunk length mismatch") - self._chunks = None - decoded = _PAYLOAD_ADAPTER.validate_json(b"".join(pending.parts), strict=True) - if OmpFrame.model_validate(decoded).type == OmpFrameType.CHUNK: - raise OmpProtocolError("OMP RPC protocol chunks cannot nest") - return decoded + await self._write_line(payload_json(frame)) async def _read_stdout(self, process: asyncio.subprocess.Process) -> None: assert process.stdout is not None @@ -685,16 +432,14 @@ async def _read_stdout(self, process: asyncio.subprocess.Process) -> None: if not line: if self._closing: return - if self._chunks is not None: + if self._frames.pending: raise OmpProtocolError( "OMP RPC protocol chunks interrupted by stdout EOF" ) raise OmpProcessError("OMP RPC stdout EOF") if len(line) > self._max_frame_bytes or not line.endswith(b"\n"): raise OmpProtocolError("Invalid OMP RPC protocol frame length") - payload = self._decode_payload( - _PAYLOAD_ADAPTER.validate_json(line, strict=True) - ) + payload = self._frames.decode(parse_payload_line(line)) if payload is None: continue self._sequence += 1 @@ -771,9 +516,7 @@ async def _dispatch_payload( if turn is None: write = asyncio.create_task( self._write_frame( - _ui_response_frame( - request.id, OmpUIResponse(cancelled=True) - ) + ui_response_frame(request.id, OmpUIResponse(cancelled=True)) ) ) self._write_tasks.add(write) @@ -798,7 +541,7 @@ def _dispatch_response(self, response: OmpResponse, sequence: int) -> None: and self._turn is not None and response.id == self._turn.id ): - self._turn.fail(_command_error(response)) + self._turn.fail(command_error(response)) return if response.command != pending.command: raise OmpProtocolError("OMP response command does not match its request") @@ -807,7 +550,7 @@ def _dispatch_response(self, response: OmpResponse, sequence: int) -> None: if response.success: pending.future.set_result(ReceivedResponse(response, sequence)) else: - pending.future.set_result(_command_error(response)) + pending.future.set_result(command_error(response)) def _start_callback( self, @@ -876,7 +619,7 @@ async def _handle_ui(self, turn: PromptRun, request: OmpUIRequest) -> None: response = await turn.on_ui(request) if not self._callback_still_active(turn, OmpCallbackKind.UI, request.id): return - await self._write_frame(_ui_response_frame(request.id, response)) + await self._write_frame(ui_response_frame(request.id, response)) async def _read_stderr(self, process: asyncio.subprocess.Process) -> None: assert process.stderr is not None @@ -931,12 +674,14 @@ def _consume_task(task: asyncio.Task[None]) -> None: task.exception() @staticmethod - async def _cancel_tasks(tasks: Sequence[asyncio.Task[None]]) -> None: + async def _cancel_tasks( + tasks: Sequence[asyncio.Task[None]], *, timeout: float + ) -> None: active = {task for task in tasks if not task.done()} for task in active: task.cancel() if active: - _, pending = await asyncio.wait(active, timeout=_CALLBACK_EXIT_TIMEOUT_S) + _, pending = await asyncio.wait(active, timeout=timeout) for task in pending: task.cancel() logger.warning("OMP callback did not honor cancellation") @@ -951,9 +696,8 @@ async def _drain(stream: asyncio.StreamReader) -> None: while await stream.read(_DRAIN_BYTES): pass - @staticmethod - async def _await_contained_exit(process: asyncio.subprocess.Process) -> None: - if not await await_returncode(process, _PROCESS_EXIT_TIMEOUT_S): + async def _await_contained_exit(self, process: asyncio.subprocess.Process) -> None: + if not await await_returncode(process, self.process_exit_timeout_s): raise OmpProcessError("OMP process did not exit after containment kill") async def _settle_spawn(self) -> asyncio.subprocess.Process | None: @@ -976,7 +720,7 @@ async def _close_process(self) -> None: workers = list(turn.callback_tasks) if turn.event_worker is not None: workers.append(turn.event_worker) - await self._cancel_tasks(workers) + await self._cancel_tasks(workers, timeout=self.callback_exit_timeout_s) owner = self._windows_process scope = self._process_scope @@ -984,7 +728,7 @@ async def _close_process(self) -> None: if owner is not None and self._proc is None: # Stop incomplete helper creation/readiness before awaiting its spawn task. try: - await owner.close(_PROCESS_EXIT_TIMEOUT_S) + await owner.close(self.process_exit_timeout_s) owner_closed = True finally: self._proc = owner.process @@ -1020,13 +764,13 @@ async def _close_process(self) -> None: await asyncio.to_thread(scope.capture, process) if not owner_closed: await close_stdin(process) - await await_returncode(process, _PROCESS_EXIT_TIMEOUT_S) + await await_returncode(process, self.process_exit_timeout_s) if owner is not None: if not owner_closed: - await owner.close(_PROCESS_EXIT_TIMEOUT_S) + await owner.close(self.process_exit_timeout_s) if scope is not None: await asyncio.to_thread( - scope.close, process, _PROCESS_EXIT_TIMEOUT_S + scope.close, process, self.process_exit_timeout_s ) else: assert scope is not None @@ -1034,7 +778,9 @@ async def _close_process(self) -> None: await kill_process_tree(process) await asyncio.to_thread(scope.signal) await self._await_contained_exit(process) - await asyncio.to_thread(scope.close, process, _PROCESS_EXIT_TIMEOUT_S) + await asyncio.to_thread( + scope.close, process, self.process_exit_timeout_s + ) await self._await_contained_exit(process) self._windows_process = None self._process_scope = None @@ -1042,9 +788,9 @@ async def _close_process(self) -> None: # A cancelled pipe read on Windows stays pending until the child is # killed; drain stdout only after that reader finishes. if reader is not None: - await self._cancel_tasks([reader]) + await self._cancel_tasks([reader], timeout=self.callback_exit_timeout_s) if process.stdout is not None: drains.append(asyncio.create_task(self._drain(process.stdout))) if self._stderr_task is not None: drains.append(self._stderr_task) - await self._cancel_tasks(drains) + await self._cancel_tasks(drains, timeout=self.callback_exit_timeout_s) diff --git a/src/band/integrations/omprpc/errors.py b/src/band/integrations/omprpc/errors.py new file mode 100644 index 000000000..6c5820128 --- /dev/null +++ b/src/band/integrations/omprpc/errors.py @@ -0,0 +1,43 @@ +"""Native OMP RPC failure types.""" + +from __future__ import annotations + +from band.integrations.ompwire import OmpResponse + + +class OmpRpcError(RuntimeError): + """An unsuccessful command response; the transport remains usable.""" + + def __init__(self, command: str, error: str, *, code: str | None = None) -> None: + self.command = command + self.error = error + self.code = code + super().__init__(f"OMP {command}: {error}") + + +class OmpProcessError(RuntimeError): + """The OMP process or its stdio transport is no longer usable.""" + + +BACKLOG_EXCEEDED = "OMP event callback backlog exceeded" +CALLBACKS_UNFINISHED = "OMP callbacks did not finish after settlement" + + +class OmpProtocolError(OmpProcessError): + """An invalid or incorrectly correlated native RPC frame.""" + + +class OmpPromptError(RuntimeError): + """The prompt or its final background continuation failed or was aborted.""" + + def __init__(self, status: str, message: str) -> None: + self.status = status + super().__init__(message) + + +def command_error(response: OmpResponse) -> OmpRpcError: + return OmpRpcError( + response.command, + response.error or "Command failed", + code=response.code, + ) diff --git a/src/band/integrations/omprpc/frames.py b/src/band/integrations/omprpc/frames.py new file mode 100644 index 000000000..0bb968029 --- /dev/null +++ b/src/band/integrations/omprpc/frames.py @@ -0,0 +1,122 @@ +"""Native OMP JSONL frame reassembly and host UI replies.""" + +from __future__ import annotations + +import base64 +from dataclasses import dataclass + +from pydantic import JsonValue, TypeAdapter + +from band.integrations.omprpc.errors import OmpProtocolError +from band.integrations.ompwire import ( + OMP_CHUNK_BYTES, + OMP_REASSEMBLED_BYTES, + OmpChunk, + OmpFrame, + OmpFrameType, + OmpUIResponse, +) + +_PAYLOAD_ADAPTER: TypeAdapter[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]) + + +@dataclass +class PendingChunks: + chunk_id: str + count: int + byte_length: int + next_index: int + received_bytes: int + parts: list[bytes] + + +class ChunkAssembler: + """Reassemble `rpc_chunk` sequences into one logical frame.""" + + def __init__(self) -> None: + self.max_reassembled_bytes = OMP_REASSEMBLED_BYTES + self._chunks: PendingChunks | None = None + + @property + def pending(self) -> bool: + return self._chunks is not None + + def reset(self) -> None: + self._chunks = None + + def decode(self, payload: dict[str, JsonValue]) -> dict[str, JsonValue] | None: + frame = OmpFrame.model_validate(payload) + if frame.type != OmpFrameType.CHUNK: + if self._chunks is not None: + raise OmpProtocolError("OMP RPC protocol chunk sequence interrupted") + return payload + chunk = OmpChunk.model_validate(payload) + if chunk.byteLength > self.max_reassembled_bytes: + raise OmpProtocolError( + "OMP RPC protocol frame exceeds the negotiated limit" + ) + data = base64.b64decode(chunk.data, validate=True) + if ( + not data + or len(data) > OMP_CHUNK_BYTES + or base64.b64encode(data).decode("ascii") != chunk.data + ): + raise OmpProtocolError("Invalid OMP RPC protocol chunk data") + if self._chunks is None: + if chunk.index != 0: + raise OmpProtocolError( + "OMP RPC protocol chunks must start at index zero" + ) + self._chunks = PendingChunks( + chunk.chunkId, chunk.count, chunk.byteLength, 0, 0, [] + ) + pending = self._chunks + if ( + pending.chunk_id != chunk.chunkId + or pending.count != chunk.count + or pending.byte_length != chunk.byteLength + or pending.next_index != chunk.index + ): + raise OmpProtocolError("OMP RPC protocol chunk sequence mismatch") + pending.parts.append(data) + pending.received_bytes += len(data) + pending.next_index += 1 + if pending.received_bytes > pending.byte_length: + raise OmpProtocolError( + "OMP RPC protocol chunks exceed their declared length" + ) + if pending.next_index < pending.count: + return None + if pending.received_bytes != pending.byte_length: + raise OmpProtocolError("OMP RPC protocol chunk length mismatch") + self._chunks = None + decoded = _PAYLOAD_ADAPTER.validate_json(b"".join(pending.parts), strict=True) + if OmpFrame.model_validate(decoded).type == OmpFrameType.CHUNK: + raise OmpProtocolError("OMP RPC protocol chunks cannot nest") + return decoded + + +def payload_json(frame: dict[str, JsonValue]) -> str: + return _PAYLOAD_ADAPTER.dump_json(frame).decode("utf-8") + + +def parse_payload_line(line: bytes) -> dict[str, JsonValue]: + return _PAYLOAD_ADAPTER.validate_json(line, strict=True) + + +def ui_response_frame(request_id: str, response: OmpUIResponse) -> dict[str, JsonValue]: + frame: dict[str, JsonValue] = { + "type": OmpFrameType.UI_RESPONSE, + "id": request_id, + } + if response.value is not None: + frame["value"] = response.value + elif response.confirmed is not None: + frame["confirmed"] = response.confirmed + elif response.answers is not None: + frame["answers"] = [ + answer.model_dump(exclude_none=True) for answer in response.answers + ] + else: + frame["cancelled"] = True + return frame diff --git a/src/band/integrations/omprpc/prompt.py b/src/band/integrations/omprpc/prompt.py new file mode 100644 index 000000000..dbb01756f --- /dev/null +++ b/src/band/integrations/omprpc/prompt.py @@ -0,0 +1,85 @@ +"""In-flight native OMP prompt state.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from enum import StrEnum + +from band.integrations.omprpc.errors import BACKLOG_EXCEEDED, OmpProcessError +from band.integrations.ompwire import ( + OMP_ASSISTANT_ROLE, + OmpEvent, + OmpEventType, + OmpHostToolCall, + OmpMessage, + OmpPromptResult, + OmpToolResult, + OmpUIRequest, + OmpUIResponse, +) + +EVENT_BACKLOG_LIMIT = 256 + + +class OmpCallbackKind(StrEnum): + TOOL = "tool" + UI = "ui" + + +class PromptRun: + def __init__( + self, + request_id: str, + on_event: Callable[[OmpEvent], Awaitable[None]], + on_tool: Callable[[OmpHostToolCall], Awaitable[OmpToolResult]], + on_ui: Callable[[OmpUIRequest], Awaitable[OmpUIResponse]], + ) -> None: + loop = asyncio.get_running_loop() + self.id = request_id + self.on_event = on_event + self.on_tool = on_tool + self.on_ui = on_ui + self.events: asyncio.Queue[OmpEvent] = asyncio.Queue( + maxsize=EVENT_BACKLOG_LIMIT + ) + self.failure: asyncio.Future[BaseException] = loop.create_future() + self.result: asyncio.Future[OmpPromptResult] = loop.create_future() + self.changed = asyncio.Event() + self.last_activity_sequence = 0 + self.settled_sequence = 0 + self.last_assistant: OmpMessage | None = None + self.rejected = False + self.callbacks: dict[tuple[OmpCallbackKind, str], asyncio.Task[None]] = {} + self.callback_tasks: set[asyncio.Task[None]] = set() + self.event_worker: asyncio.Task[None] | None = None + + def fail(self, error: BaseException) -> None: + if not self.failure.done(): + self.failure.set_result(error) + + def observe(self, event: OmpEvent, sequence: int) -> None: + if self.failure.done(): + return + # Waiting for space would block the reader that services reentrant commands. + try: + self.events.put_nowait(event) + except asyncio.QueueFull: + self.fail(OmpProcessError(BACKLOG_EXCEEDED)) + return + self.last_activity_sequence = sequence + match event.type: + case OmpEventType.MESSAGE_END if ( + event.message is not None and event.message.role == OMP_ASSISTANT_ROLE + ): + self.last_assistant = event.message + case OmpEventType.AGENT_END: + self.last_assistant = next( + ( + message + for message in reversed(event.messages) + if message.role == OMP_ASSISTANT_ROLE + ), + self.last_assistant, + ) + self.changed.set() diff --git a/src/band/integrations/omprpc/scope.py b/src/band/integrations/omprpc/scope.py new file mode 100644 index 000000000..6caf27aa5 --- /dev/null +++ b/src/band/integrations/omprpc/scope.py @@ -0,0 +1,101 @@ +"""Best-effort custody of an OMP RPC process tree.""" + +from __future__ import annotations + +import asyncio +import os +import time +from collections.abc import Mapping +from uuid import uuid4 + +import psutil + +from band.integrations.codex.winspawn import session_pid_reused +from band.integrations.omprpc.errors import OmpProcessError + +_PROCESS_SCOPE_VALUE = "1" + + +class OmpProcessScope: + """Retain observable descendants for best-effort cleanup. + + The private PI_ marker survives native helper filters. Explicit environment + replacement and restricted process inspection can hide detached children. + Windows job membership is the primary containment; this marker catches + descendants the job does not observe. + """ + + def __init__(self) -> None: + self.key = f"PI_BAND_PROCESS_{uuid4().hex}" + self.processes: set[psutil.Process] = set() + + def environment(self, base: Mapping[str, str] | None) -> dict[str, str]: + return { + **(os.environ if base is None else base), + self.key: _PROCESS_SCOPE_VALUE, + } + + def capture(self, root: asyncio.subprocess.Process) -> None: + group_visible = not session_pid_reused(root) + group_members: list[psutil.Process] = [] + getpgid = getattr(os, "getpgid", None) + for pid in psutil.pids(): + # asyncio's child watcher remains the sole waiter for the root. + if pid == 0 or pid == root.pid: + continue + try: + member = psutil.Process(pid) + created = member.create_time() + marked = False + try: + marked = member.environ().get(self.key) == _PROCESS_SCOPE_VALUE + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass + in_group = False + if group_visible and getpgid is not None: + try: + in_group = getpgid(pid) == root.pid + except (ProcessLookupError, PermissionError): + pass + if not marked and not in_group: + continue + current = psutil.Process(pid) + if current.create_time() != created or current != member: + continue + if marked: + self.processes.add(member) + elif in_group: + group_members.append(member) + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + if group_visible and not session_pid_reused(root): + self.processes.update(group_members) + + @staticmethod + def _is_alive(process: psutil.Process) -> bool: + try: + return process.is_running() and process.status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + def signal(self) -> None: + for process in self.processes: + try: + process.kill() + except psutil.NoSuchProcess: + pass + + def close(self, root: asyncio.subprocess.Process, timeout: float) -> None: + deadline = time.monotonic() + timeout + self.capture(root) + alive = [process for process in self.processes if self._is_alive(process)] + while alive: + self.signal() + psutil.wait_procs(alive, timeout=max(0.0, deadline - time.monotonic())) + self.capture(root) + alive = [process for process in self.processes if self._is_alive(process)] + if alive and time.monotonic() >= deadline: + raise OmpProcessError( + "OMP owned session members survived cleanup: " + f"{sorted(process.pid for process in alive)}" + ) diff --git a/tests/converters/test_omp.py b/tests/converters/test_omp.py index 1b5fd6563..cdccf9384 100644 --- a/tests/converters/test_omp.py +++ b/tests/converters/test_omp.py @@ -83,12 +83,13 @@ def test_session_metadata_on_non_task_messages_is_ignored(self) -> None: OMP_SESSION_METADATA_KEY: "/sessions/injected.jsonl", }, "content": "hello", - "sender": {"name": "Alice", "type": "User"}, + "sender_name": "Alice", + "sender_type": "User", } ] ) assert state.room_to_session == {} - assert len(state.replay_messages) == 1 + assert state.replay_messages == ["[Alice]: hello"] def test_room_transcript_is_replayed(self) -> None: state = OmpHistoryConverter().convert( @@ -96,9 +97,9 @@ def test_room_transcript_is_replayed(self) -> None: { "message_type": "text", "content": "hello", - "sender": {"name": "Alice", "type": "User"}, + "sender_name": "Alice", + "sender_type": "User", } ] ) - assert len(state.replay_messages) == 1 - assert "hello" in state.replay_messages[0] + assert state.replay_messages == ["[Alice]: hello"] diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index cd284f3d6..c77fe6f1b 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -21,9 +21,22 @@ OMP_FRAME_BYTES, OMP_NATIVE_PROTOCOL_VERSION, OMP_REASSEMBLED_BYTES, + OmpPromptStatus, ) PROTOCOL_VERSION = 1 +COMPLETE = "complete" +YIELDED = "yielded" +BACKGROUND_COMPLETE = "background complete" +RESUMED_PROVIDER_FAILED = "Resumed provider failed" +FOLLOW_UP_PROMPT = "second" +CRASH_EXIT = 3 +UI_APPROVE = "Approve" +UI_DENY = "Deny" + + +def burst_message_id(number: int) -> str: + return f"burst-{number}" class Frame(BaseModel): @@ -161,8 +174,8 @@ def assistant(text: str, *, stop: str = "stop") -> dict[str, JsonValue]: }, }, } - if stop == "error": - message["errorMessage"] = "Resumed provider failed" + if stop == OmpPromptStatus.ERROR: + message["errorMessage"] = RESUMED_PROVIDER_FAILED return message @@ -227,7 +240,7 @@ def settle(self) -> None: def finish(self, frame: Frame) -> None: self.respond(frame) send({"type": "agent_start"}) - self.end("complete") + self.end(COMPLETE) self.result(settled=True) self.settle() @@ -267,7 +280,7 @@ def prompt(self, frame: Frame) -> None: self.settled = False if ( self.prompt_count > 1 and self.scenario != "lifecycle" - ) or frame.message == "second": + ) or frame.message == FOLLOW_UP_PROMPT: self.finish(frame) return match self.scenario: @@ -333,7 +346,7 @@ def prompt(self, frame: Frame) -> None: "id": "approval-1", "method": "select", "title": "Allow command?", - "options": ["Approve", "Deny"], + "options": [UI_APPROVE, UI_DENY], } ) return @@ -366,7 +379,7 @@ def prompt(self, frame: Frame) -> None: "message": "status", } ) - self.end("complete") + self.end(COMPLETE) self.result(settled=True) self.settle() case "reused-callback": @@ -388,7 +401,7 @@ def prompt(self, frame: Frame) -> None: } ) case "large-image": - message = assistant("complete") + message = assistant(COMPLETE) message["content"] = [ { "type": "image", @@ -430,14 +443,14 @@ def prompt(self, frame: Frame) -> None: "isError": False, } ) - self.end("complete") + self.end(COMPLETE) self.result(settled=True) self.settle() case "settlement-before-result": - self.end("complete") + self.end(COMPLETE) self.settle() case "no-result": - self.end("complete") + self.end(COMPLETE) self.settle() case "settled-during-callback": send( @@ -449,7 +462,7 @@ def prompt(self, frame: Frame) -> None: "arguments": {}, } ) - self.end("complete") + self.end(COMPLETE) self.result(settled=True) self.settle() case "host-cancel": @@ -469,16 +482,20 @@ def prompt(self, frame: Frame) -> None: "id": "ui-cancelled", "method": "select", "title": "Allow command?", - "options": ["Approve", "Deny"], + "options": [UI_APPROVE, UI_DENY], } ) case "initial-error" | "initial-aborted": - stop = "error" if self.scenario == "initial-error" else "aborted" + stop = ( + OmpPromptStatus.ERROR + if self.scenario == "initial-error" + else OmpPromptStatus.ABORTED + ) self.end("failed", stop=stop) self.result(settled=True, status=stop) self.settle() case _: - self.end("yielded", pending=True) + self.end(YIELDED, pending=True) self.result(settled=self.scenario == "false-result-settled") if self.scenario == "state-settles": self.settled = True @@ -547,14 +564,14 @@ def handle(self, frame: Frame) -> None: send( { "type": "message_update", - "messageId": f"burst-{number}", + "messageId": burst_message_id(number), "message": {"role": "assistant", "content": str(number)}, } ) self.respond(frame) case "event-result": self.held = frame - self.end("complete", compacted=True) + self.end(COMPLETE, compacted=True) self.result(settled=True) self.settle() case "event-ui": @@ -565,7 +582,7 @@ def handle(self, frame: Frame) -> None: "id": "backlog-approval", "method": "select", "title": "Allow command?", - "options": ["Approve", "Deny"], + "options": [UI_APPROVE, UI_DENY], } ) case "ask-ui": @@ -590,12 +607,12 @@ def handle(self, frame: Frame) -> None: else: send({"type": "agent_start"}) stop = { - "resume-error": "error", - "resume-error-compacted": "error", - "resume-aborted": "aborted", + "resume-error": OmpPromptStatus.ERROR, + "resume-error-compacted": OmpPromptStatus.ERROR, + "resume-aborted": OmpPromptStatus.ABORTED, }.get(self.scenario, "stop") self.end( - "background complete", + BACKGROUND_COMPLETE, stop=stop, compacted=self.scenario == "resume-error-compacted", ) @@ -684,7 +701,7 @@ def handle(self, frame: Frame) -> None: "id": "ui-orphan", "method": "select", "title": "Allow command?", - "options": ["Approve", "Deny"], + "options": [UI_APPROVE, UI_DENY], } ) self.respond(frame) @@ -694,7 +711,7 @@ def handle(self, frame: Frame) -> None: if self.scenario != "stubborn": self.respond(frame) if self.prompt_id is not None: - self.end("aborted", stop="aborted") + self.end(OmpPromptStatus.ABORTED, stop=OmpPromptStatus.ABORTED) self.settle() case "first" | "hold": self.held = frame @@ -715,7 +732,7 @@ def handle(self, frame: Frame) -> None: } ) case "crash": - raise SystemExit(3) + raise SystemExit(CRASH_EXIT) case "stderr": sys.stderr.write("x" * (2 * 1024 * 1024)) sys.stderr.flush() diff --git a/tests/integrations/ompprocess/test_mcp_plugin.py b/tests/integrations/ompprocess/test_mcp_plugin.py index 4b3b51ebf..75cacbdd0 100644 --- a/tests/integrations/ompprocess/test_mcp_plugin.py +++ b/tests/integrations/ompprocess/test_mcp_plugin.py @@ -20,6 +20,12 @@ ) +def plugin_privacy(root: Path, files: tuple[Path, ...]) -> tuple[int, int | None]: + file_modes = {stat.S_IMODE(path.stat().st_mode) for path in files} + file_mode = next(iter(file_modes)) if len(file_modes) == 1 else None + return (stat.S_IMODE(root.stat().st_mode), file_mode) + + class TestMcpPluginProjection: def test_nothing_to_project_yields_no_plugin(self, tmp_path: Path) -> None: with omp_mcp_plugin({}, workspace=str(tmp_path)) as plugin: @@ -78,19 +84,9 @@ def test_arguments_load_the_plugin_and_its_startup_extension( assert str(Path(plugin) / OMP_MCP_STARTUP_JS) in arguments assert str(Path(plugin) / OMP_MCP_STARTUP_CONFIG) in arguments - @pytest.mark.parametrize( - "accepted_on, file_mode, dir_mode", - [ - (("posix",), 0o600, 0o700), - (("nt",), None, None), - ], - ) + @pytest.mark.parametrize("accepted_on", [("posix",)]) def test_projected_plugin_files_are_owner_only( - self, - tmp_path: Path, - accepted_on: tuple[str, ...], - file_mode: int | None, - dir_mode: int | None, + self, tmp_path: Path, accepted_on: tuple[str, ...] ) -> None: servers = {"local": OmpMcpServer(command="./server")} with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: @@ -103,10 +99,9 @@ def test_projected_plugin_files_are_owner_only( ) assert root.is_dir() assert all(path.is_file() for path in files) - if os.name not in accepted_on or file_mode is None or dir_mode is None: - return - assert stat.S_IMODE(root.stat().st_mode) == dir_mode - assert all(stat.S_IMODE(path.stat().st_mode) == file_mode for path in files) + assert (plugin_privacy(root, files) == (0o700, 0o600)) is ( + os.name in accepted_on + ) @pytest.mark.parametrize( "fields", diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py index 7ed2fb3ae..1712c6011 100644 --- a/tests/integrations/omprpc/kit.py +++ b/tests/integrations/omprpc/kit.py @@ -27,7 +27,7 @@ OmpUIResponse, ) from tests.integrations import omppeer -from tests.integrations.omppeer import ProcessReceipt +from tests.integrations.omppeer import COMPLETE, FOLLOW_UP_PROMPT, ProcessReceipt ClientFactory = Callable[[str], Awaitable[OmpRpcClient]] @@ -52,6 +52,18 @@ def assistant_texts(self) -> list[str]: if part.text is not None ] + @property + def outline(self) -> list[OmpEventType]: + return [event.type for event in self.events] + + @property + def update_ids(self) -> list[str | None]: + return [ + event.messageId + for event in self.events + if event.type == OmpEventType.MESSAGE_UPDATE + ] + async def on_event(self, event: OmpEvent) -> None: self.events.append(event) self.queue.put_nowait(event) @@ -83,11 +95,17 @@ def prompt_task( ) +def owned_rpc_process(client: OmpRpcClient) -> asyncio.subprocess.Process: + process = client._proc + assert process is not None + return process + + async def second_prompt(client: OmpRpcClient) -> Observer: observer = Observer() async with asyncio.timeout(5): - await prompt_task(client, observer, "second") - assert "complete" in observer.assistant_texts + await prompt_task(client, observer, FOLLOW_UP_PROMPT) + assert COMPLETE in observer.assistant_texts return observer diff --git a/tests/integrations/omprpc/lifecycle.py b/tests/integrations/omprpc/lifecycle.py index 27761b125..be8d9061c 100644 --- a/tests/integrations/omprpc/lifecycle.py +++ b/tests/integrations/omprpc/lifecycle.py @@ -1,33 +1,45 @@ -"""Generated host operations against a real RPC peer and owned process tree.""" +"""A live ``OmpRpcClient`` driven step by step against a real RPC peer. + +This is the model for ``test_lifecycle_model.py``, not a collected test. +Hypothesis picks a short sequence of ``Step`` values — the operations a host +can issue — and ``RpcLifecycle.take`` applies each one to one client and the +process tree it owns. After the sequence, the client must still start, finish +a turn, and reap descendants. +""" from __future__ import annotations import asyncio from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass from enum import StrEnum -from unittest.mock import patch from hypothesis import strategies as st from band.integrations.omprpc import OmpProcessError from band.integrations.ompwire import OmpEventType +from tests.integrations.omppeer import BACKGROUND_COMPLETE, CRASH_EXIT from tests.integrations.omprpc.kit import ( GatedProcess, Observer, ProcessGate, descendant_client, descendant_liveness, + owned_rpc_process, process_gate_session, prompt_task, queue_prompt, second_prompt, ) +# A healthy step takes milliseconds; a hung one fails its example within this. STEP_GUARD_S = 15 class Step(StrEnum): + """Host operations Hypothesis may interleave.""" + PROMPT = "admit a held prompt" CANCEL_QUEUED = "cancel a prompt waiting for admission" SETTLE = "settle the admitted prompt" @@ -42,7 +54,11 @@ class Step(StrEnum): class SpawnBarrier: - """Hold publication of the first real subprocess, not its creation.""" + """Create the first subprocess, then hold returning it until released. + + Descendants can still check in at the process gate; ``start()`` / ``close()`` + wait, so a host cancel can overlap spawn. + """ def __init__( self, spawn: Callable[..., Awaitable[asyncio.subprocess.Process]] @@ -63,6 +79,17 @@ async def create( return process +@dataclass +class HeldSpawn: + starting: asyncio.Task[None] + root: asyncio.subprocess.Process + children: list[GatedProcess] + barrier: SpawnBarrier + + def release(self) -> None: + self.barrier.release.set() + + async def announced_operation( operation: Awaitable[None], entered: asyncio.Event ) -> None: @@ -71,6 +98,8 @@ async def announced_operation( class RpcLifecycle: + """One RPC client, the process trees it has owned, and at most one admitted turn.""" + def __init__(self, gate: ProcessGate, *, detached: bool) -> None: self.gate = gate self.client = descendant_client(gate, scenario="lifecycle", detached=detached) @@ -90,8 +119,7 @@ async def open(self) -> None: if self.live: return await self.client.start() - assert self.client._proc is not None - self.roots.append(self.client._proc) + self.roots.append(owned_rpc_process(self.client)) await self.record_children() self.live = True @@ -106,7 +134,44 @@ async def hold_prompt(self) -> None: assert not self.active.done() assert (await self.client.request("get_state"))["isSettled"] is False - async def assert_dead( + async def cancel_queued(self) -> None: + await self.hold_prompt() + queued = await queue_prompt(self.client) + self.tasks.append(queued) + queued.cancel() + [result] = await asyncio.gather(queued, return_exceptions=True) + assert isinstance(result, asyncio.CancelledError) + assert self.active is not None and not self.active.done() + assert (await self.client.request("get_state"))["isSettled"] is False + + async def settle(self) -> None: + await self.hold_prompt() + await self.client.request("release") + assert self.active is not None and self.observer is not None + await self.active + assert BACKGROUND_COMPLETE in self.observer.assistant_texts + self.active = None + assert (await self.client.request("get_state"))["isSettled"] is True + + async def cancel_owner(self) -> None: + await self.hold_prompt() + assert self.active is not None + self.active.cancel() + [result] = await asyncio.gather(self.active, return_exceptions=True) + assert isinstance(result, asyncio.CancelledError) + self.active = None + await self.close() + + async def crash(self) -> None: + await self.open() + root = owned_rpc_process(self.client) + with suppress(OmpProcessError): + await self.client.request("crash") + await root.wait() + assert root.returncode == CRASH_EXIT + await self.close() + + async def assert_reaped( self, roots: list[asyncio.subprocess.Process], children: list[GatedProcess] ) -> None: assert all(root.returncode is not None for root in roots) @@ -119,106 +184,111 @@ async def close(self) -> None: assert result is None or isinstance(result, OmpProcessError) self.active = None self.live = False - await self.assert_dead(self.roots, self.children) + await self.assert_reaped(self.roots, self.children) + + def _track(self, task: asyncio.Task[None]) -> asyncio.Task[None]: + self.tasks.append(task) + return task + + async def _assert_finished( + self, task: asyncio.Task[None], error: type[BaseException] + ) -> None: + [result] = await asyncio.gather(task, return_exceptions=True) + assert isinstance(result, error) - async def overlap_spawn_and_restart(self, *, cancel_start: bool) -> None: + async def _queue_close(self) -> asyncio.Task[None]: + entered = asyncio.Event() + closing = self._track( + asyncio.create_task(announced_operation(self.client.close(), entered)) + ) + await entered.wait() + return closing + + async def _queue_restart(self) -> asyncio.Task[None]: + entered = asyncio.Event() + restarting = self._track( + asyncio.create_task(announced_operation(self.client.start(), entered)) + ) + await entered.wait() + assert not restarting.done() + return restarting + + async def _claim_restart( + self, + restarting: asyncio.Task[None], + old_root: asyncio.subprocess.Process, + old_children: list[GatedProcess], + ) -> None: + await restarting + process = owned_rpc_process(self.client) + assert process is not old_root + self.roots.append(process) + await self.record_children() + self.live = True + await self.assert_reaped([old_root], old_children) + await second_prompt(self.client) + + @asynccontextmanager + async def held_first_spawn(self) -> AsyncIterator[HeldSpawn]: await self.close() barrier = SpawnBarrier(asyncio.create_subprocess_exec) - closing: asyncio.Task[None] | None = None - with patch.object(asyncio, "create_subprocess_exec", barrier.create): - starting = asyncio.create_task(self.client.start()) - self.tasks.append(starting) - try: - await barrier.created.wait() - assert barrier.process is not None - old_root = barrier.process - self.roots.append(old_root) - old_children = await self.record_children() - if cancel_start: - starting.cancel() - else: - close_entered = asyncio.Event() - closing = asyncio.create_task( - announced_operation(self.client.close(), close_entered) - ) - self.tasks.append(closing) - await close_entered.wait() - closing.cancel() - restart_entered = asyncio.Event() - restarting = asyncio.create_task( - announced_operation(self.client.start(), restart_entered) - ) - self.tasks.append(restarting) - await restart_entered.wait() - assert not restarting.done() - barrier.release.set() - [result] = await asyncio.gather(starting, return_exceptions=True) - assert isinstance( - result, - asyncio.CancelledError if cancel_start else OmpProcessError, - ) - if closing is not None: - [result] = await asyncio.gather(closing, return_exceptions=True) - assert isinstance(result, asyncio.CancelledError) - await restarting - assert self.client._proc is not None - assert self.client._proc is not old_root - self.roots.append(self.client._proc) - await self.record_children() - self.live = True - await self.assert_dead([old_root], old_children) - await second_prompt(self.client) - finally: - barrier.release.set() + self.client.create_process = barrier.create + starting = self._track(asyncio.create_task(self.client.start())) + try: + await barrier.created.wait() + assert barrier.process is not None + self.roots.append(barrier.process) + yield HeldSpawn( + starting=starting, + root=barrier.process, + children=await self.record_children(), + barrier=barrier, + ) + finally: + self.client.create_process = None + barrier.release.set() + + async def cancel_held_start_then_restart(self) -> None: + async with self.held_first_spawn() as spawn: + spawn.starting.cancel() + restarting = await self._queue_restart() + spawn.release() + await self._assert_finished(spawn.starting, asyncio.CancelledError) + await self._claim_restart(restarting, spawn.root, spawn.children) + + async def cancel_held_close_then_restart(self) -> None: + async with self.held_first_spawn() as spawn: + closing = await self._queue_close() + closing.cancel() + restarting = await self._queue_restart() + spawn.release() + await self._assert_finished(spawn.starting, OmpProcessError) + await self._assert_finished(closing, asyncio.CancelledError) + await self._claim_restart(restarting, spawn.root, spawn.children) async def take(self, step: Step) -> None: match step: case Step.PROMPT: await self.hold_prompt() case Step.CANCEL_QUEUED: - await self.hold_prompt() - queued = await queue_prompt(self.client) - self.tasks.append(queued) - queued.cancel() - [result] = await asyncio.gather(queued, return_exceptions=True) - assert isinstance(result, asyncio.CancelledError) - assert self.active is not None and not self.active.done() - assert (await self.client.request("get_state"))["isSettled"] is False + await self.cancel_queued() case Step.SETTLE: - await self.hold_prompt() - await self.client.request("release") - assert self.active is not None and self.observer is not None - await self.active - assert "background complete" in self.observer.assistant_texts - self.active = None - assert (await self.client.request("get_state"))["isSettled"] is True + await self.settle() case Step.CANCEL_OWNER: - await self.hold_prompt() - assert self.active is not None - self.active.cancel() - [result] = await asyncio.gather(self.active, return_exceptions=True) - assert isinstance(result, asyncio.CancelledError) - self.active = None - await self.close() + await self.cancel_owner() case Step.CLOSE: await self.close() case Step.CRASH: - await self.open() - root = self.client._proc - assert root is not None - with suppress(OmpProcessError): - await self.client.request("crash") - await root.wait() - assert root.returncode == 3 - await self.close() - case Step.CANCEL_START_RESTART | Step.CANCEL_CLOSE_RESTART: - await self.overlap_spawn_and_restart( - cancel_start=step == Step.CANCEL_START_RESTART - ) + await self.crash() + case Step.CANCEL_START_RESTART: + await self.cancel_held_start_then_restart() + case Step.CANCEL_CLOSE_RESTART: + await self.cancel_held_close_then_restart() if self.live: assert (await self.client.request("checkpoint"))["pid"] is not None - async def finish(self) -> None: + async def assert_custody_survives(self) -> None: + """Settle leftover work, reap the tree, and prove a fresh turn still works.""" if self.active is not None: await self.take(Step.SETTLE) await self.close() @@ -242,12 +312,3 @@ async def rpc_lifecycle(*, detached: bool) -> AsyncIterator[RpcLifecycle]: yield session finally: await session.cleanup() - - -async def play(steps: list[Step], *, detached: bool) -> None: - async with rpc_lifecycle(detached=detached) as session: - for step in steps: - async with asyncio.timeout(STEP_GUARD_S): - await session.take(step) - async with asyncio.timeout(STEP_GUARD_S): - await session.finish() diff --git a/tests/integrations/omprpc/test_event_backpressure.py b/tests/integrations/omprpc/test_event_backpressure.py index 45684b18e..57580d733 100644 --- a/tests/integrations/omprpc/test_event_backpressure.py +++ b/tests/integrations/omprpc/test_event_backpressure.py @@ -3,19 +3,79 @@ from __future__ import annotations import asyncio +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager import pytest -from pydantic import JsonValue -from band.integrations import omprpc -from band.integrations.omprpc import OmpProcessError +from band.integrations.omprpc import OmpProcessError, OmpRpcClient +from band.integrations.omprpc.errors import BACKLOG_EXCEEDED +from band.integrations.omprpc.prompt import EVENT_BACKLOG_LIMIT from band.integrations.ompwire import ( OmpEvent, OmpEventType, OmpUIRequest, OmpUIResponse, ) -from tests.integrations.omprpc.kit import ClientFactory, Observer, second_prompt +from tests.integrations.omppeer import COMPLETE, UI_APPROVE, burst_message_id +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + owned_rpc_process, + second_prompt, +) + + +@asynccontextmanager +async def held_event_prompt( + client: OmpRpcClient, +) -> AsyncIterator[ + tuple[asyncio.Task[None], asyncio.subprocess.Process, asyncio.Event] +]: + observer = Observer() + entered = asyncio.Event() + release = asyncio.Event() + stopped = asyncio.Event() + + async def on_event(event: OmpEvent) -> None: + entered.set() + try: + await release.wait() + finally: + stopped.set() + await observer.on_event(event) + + running = asyncio.create_task( + client.prompt( + "first", + on_event=on_event, + on_tool=observer.on_tool, + on_ui=observer.on_ui, + ) + ) + try: + await entered.wait() + yield running, owned_rpc_process(client), stopped + finally: + running.cancel() + release.set() + await asyncio.gather(running, return_exceptions=True) + + +async def overflow_reaps_and_restarts( + client: OmpRpcClient, + running: asyncio.Task[None], + process: asyncio.subprocess.Process, + stopped: asyncio.Event, +) -> None: + await client.request("event-burst", count=EVENT_BACKLOG_LIMIT + 1) + with pytest.raises(OmpProcessError, match=BACKLOG_EXCEEDED): + await running + assert stopped.is_set() + assert process.returncode is not None + await client.start() + await second_prompt(client) + assert owned_rpc_process(client) is not process class TestEventBackpressure: @@ -23,25 +83,26 @@ async def test_full_backlog_keeps_reentrant_control_live_and_preserves_event_ord self, client_factory: ClientFactory ) -> None: client = await client_factory("event-backpressure") - process = client._proc - assert process is not None + process = owned_rpc_process(client) observer = Observer() entered = asyncio.Event() release = asyncio.Event() delivered = asyncio.Event() - event_count = omprpc._EVENT_BACKLOG_LIMIT + burst_ids = [burst_message_id(number) for number in range(EVENT_BACKLOG_LIMIT)] + last_burst = burst_ids[-1] async def on_event(event: OmpEvent) -> None: - if event.type == OmpEventType.AGENT_START: - entered.set() - await release.wait() + match event.type: + case OmpEventType.AGENT_START: + entered.set() + await release.wait() await observer.on_event(event) - if event.messageId == f"burst-{event_count - 1}": + if event.messageId == last_burst: delivered.set() async def on_ui(request: OmpUIRequest) -> OmpUIResponse: await client.request("checkpoint") - return OmpUIResponse(value="Approve") + return OmpUIResponse(value=UI_APPROVE) running = asyncio.create_task( client.prompt( @@ -54,8 +115,8 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: try: async with asyncio.timeout(8): await entered.wait() - await client.request("event-burst", count=event_count) - assert await client.request("event-ui") == {"value": "Approve"} + await client.request("event-burst", count=EVENT_BACKLOG_LIMIT) + assert await client.request("event-ui") == {"value": UI_APPROVE} assert observer.events == [] assert not running.done() assert process.returncode is None @@ -65,86 +126,47 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: await client.request("event-result") await running - assert [(event.type, event.messageId) for event in observer.events] == [ - (OmpEventType.AGENT_START, None), - *[ - (OmpEventType.MESSAGE_UPDATE, f"burst-{number}") - for number in range(event_count) - ], - (OmpEventType.MESSAGE_END, "complete"), - (OmpEventType.AGENT_END, None), + assert observer.outline == [ + OmpEventType.AGENT_START, + *([OmpEventType.MESSAGE_UPDATE] * EVENT_BACKLOG_LIMIT), + OmpEventType.MESSAGE_END, + OmpEventType.AGENT_END, ] - assert observer.assistant_texts == ["complete"] + assert observer.update_ids == burst_ids + assert observer.assistant_texts == [COMPLETE] await second_prompt(client) - assert client._proc is process + assert owned_rpc_process(client) is process assert process.returncode is None finally: running.cancel() release.set() await asyncio.gather(running, return_exceptions=True) - @pytest.mark.parametrize( - "result_seen", - [False, True], - ids=["before-result", "after-result-and-settlement"], - ) - async def test_overload_fails_the_prompt_cancels_delivery_reaps_and_can_restart( - self, client_factory: ClientFactory, result_seen: bool + async def test_overload_before_result_fails_closed_and_restarts( + self, client_factory: ClientFactory ) -> None: client = await client_factory("event-backpressure") - process = client._proc - assert process is not None - observer = Observer() - entered = asyncio.Event() - release = asyncio.Event() - stopped = asyncio.Event() - - async def on_event(event: OmpEvent) -> None: - entered.set() - try: - await release.wait() - finally: - stopped.set() - await observer.on_event(event) - - running = asyncio.create_task( - client.prompt( - "first", - on_event=on_event, - on_tool=observer.on_tool, - on_ui=observer.on_ui, - ) - ) - overflowing: asyncio.Task[dict[str, JsonValue]] | None = None - try: - async with asyncio.timeout(8): - await entered.wait() - if result_seen: - # The peer acknowledges only after the settled state snapshot. - await client.request("event-result") - await client.request( - "event-burst", - count=omprpc._EVENT_BACKLOG_LIMIT - (2 if result_seen else 0), - ) - assert not running.done() - assert not stopped.is_set() - - overflowing = asyncio.create_task( - client.request("event-burst", count=1) - ) - with pytest.raises(OmpProcessError, match="event callback backlog"): - await running - - assert stopped.is_set() - assert observer.events == [] - assert process.returncode is not None - await asyncio.gather(overflowing, return_exceptions=True) - await client.start() - await second_prompt(client) - assert client._proc is not process - finally: - tasks = [running, *([overflowing] if overflowing is not None else [])] - for task in tasks: - task.cancel() - release.set() - await asyncio.gather(*tasks, return_exceptions=True) + async with ( + asyncio.timeout(8), + held_event_prompt(client) as ( + running, + process, + stopped, + ), + ): + await overflow_reaps_and_restarts(client, running, process, stopped) + + async def test_overload_after_settlement_fails_closed_and_restarts( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("event-backpressure") + async with ( + asyncio.timeout(8), + held_event_prompt(client) as ( + running, + process, + stopped, + ), + ): + await client.request("event-result") + await overflow_reaps_and_restarts(client, running, process, stopped) diff --git a/tests/integrations/omprpc/test_lifecycle_model.py b/tests/integrations/omprpc/test_lifecycle_model.py index 82abf3854..937296b39 100644 --- a/tests/integrations/omprpc/test_lifecycle_model.py +++ b/tests/integrations/omprpc/test_lifecycle_model.py @@ -7,7 +7,21 @@ import pytest from hypothesis import Phase, example, given, settings -from tests.integrations.omprpc.lifecycle import Step, play, sequences +from tests.integrations.omprpc.lifecycle import ( + STEP_GUARD_S, + Step, + rpc_lifecycle, + sequences, +) + + +async def play(steps: list[Step], *, detached: bool) -> None: + async with rpc_lifecycle(detached=detached) as session: + for step in steps: + async with asyncio.timeout(STEP_GUARD_S): + await session.take(step) + async with asyncio.timeout(STEP_GUARD_S): + await session.assert_custody_survives() class TestRpcLifecycleModel: diff --git a/tests/integrations/omprpc/test_process_tree.py b/tests/integrations/omprpc/test_process_tree.py index d5c78827f..718a17688 100644 --- a/tests/integrations/omprpc/test_process_tree.py +++ b/tests/integrations/omprpc/test_process_tree.py @@ -9,9 +9,10 @@ import pytest from pydantic import JsonValue -from band.integrations import omprpc from band.integrations.codex.winspawn import Job from band.integrations.omprpc import OmpProcessError +from band.integrations.omprpc.client import PROCESS_EXIT_TIMEOUT_S +from tests.integrations.omppeer import CRASH_EXIT from tests.integrations.omprpc.cancellation import HeldCleanup from tests.integrations.omprpc.kit import ( DeferredSessionSignal, @@ -62,7 +63,7 @@ async def test_close_reaps_detached_descendants_without_affecting_other_owners( crashing = asyncio.create_task(client.request("crash")) async with asyncio.timeout(5): await root.wait() - assert root.returncode == 3 + assert root.returncode == CRASH_EXIT with pytest.raises(OmpProcessError): await crashing @@ -123,7 +124,7 @@ async def test_close_reaps_owned_descendants_after_the_root_has_crashed( crashing = asyncio.create_task(client.request("crash")) async with asyncio.timeout(5): await root.wait() - assert root.returncode == 3 + assert root.returncode == CRASH_EXIT with pytest.raises(OmpProcessError): await crashing await client.close() @@ -196,7 +197,7 @@ async def test_unsettled_session_members_keep_close_retryable( ) with monkeypatch.context() as patch: boundary.install(patch) - patch.setattr(omprpc, "_PROCESS_EXIT_TIMEOUT_S", 0) + client.process_exit_timeout_s = 0 crashing = asyncio.create_task(client.request("crash")) async with asyncio.timeout(5): await old_root.wait() @@ -204,6 +205,7 @@ async def test_unsettled_session_members_keep_close_retryable( with pytest.raises(OmpProcessError, match="session members"): await client.close() assert await descendant_liveness(descendants) == [True, True] + client.process_exit_timeout_s = PROCESS_EXIT_TIMEOUT_S await client.start() assert client._proc is not old_root diff --git a/tests/integrations/omprpc/test_turn_settlement.py b/tests/integrations/omprpc/test_turn_settlement.py index 8324ca9cd..64c17d21b 100644 --- a/tests/integrations/omprpc/test_turn_settlement.py +++ b/tests/integrations/omprpc/test_turn_settlement.py @@ -6,17 +6,24 @@ import pytest -from band.integrations import omprpc from band.integrations.omprpc import OmpProcessError, OmpPromptError +from band.integrations.omprpc.errors import CALLBACKS_UNFINISHED from band.integrations.ompwire import ( OmpEventType, OmpHostToolCall, + OmpPromptStatus, OmpState, OmpToolResult, ) +from tests.integrations.omppeer import ( + BACKGROUND_COMPLETE, + RESUMED_PROVIDER_FAILED, + YIELDED, +) from tests.integrations.omprpc.kit import ( ClientFactory, Observer, + owned_rpc_process, prompt_task, second_prompt, ) @@ -38,16 +45,24 @@ async def test_initial_yield_never_finishes_a_background_turn( await client.request("release") async with asyncio.timeout(5): await running - assert observer.assistant_texts == ["yielded", "background complete"] + assert observer.assistant_texts == [YIELDED, BACKGROUND_COMPLETE] finally: running.cancel() await asyncio.gather(running, return_exceptions=True) @pytest.mark.parametrize( - "scenario", ["resume-error", "resume-error-compacted", "resume-aborted"] + ("scenario", "error"), + [ + ("resume-error", RESUMED_PROVIDER_FAILED), + ("resume-error-compacted", RESUMED_PROVIDER_FAILED), + ( + "resume-aborted", + f"OMP continuation {OmpPromptStatus.ABORTED}", + ), + ], ) async def test_resumed_failures_override_a_successful_initial_prompt_result( - self, client_factory: ClientFactory, scenario: str + self, client_factory: ClientFactory, scenario: str, error: str ) -> None: client = await client_factory(scenario) observer = Observer() @@ -55,16 +70,9 @@ async def test_resumed_failures_override_a_successful_initial_prompt_result( try: await observer.wait_for(OmpEventType.AGENT_END) await client.request("release") - expected = ( - "Resumed provider failed" if scenario != "resume-aborted" else "aborted" - ) async with asyncio.timeout(5): - with pytest.raises(OmpPromptError, match=expected): + with pytest.raises(OmpPromptError, match=error): await running - assert ( - sum(event.type == OmpEventType.AGENT_END for event in observer.events) - == 2 - ) finally: running.cancel() await asyncio.gather(running, return_exceptions=True) @@ -100,12 +108,11 @@ async def test_native_state_snapshot_and_subscribed_settlement_do_not_race( assert OmpState.model_validate(await client.request("get_state")).isSettled async def test_settlement_does_not_succeed_while_a_callback_is_still_running( - self, client_factory: ClientFactory, monkeypatch: pytest.MonkeyPatch + self, client_factory: ClientFactory ) -> None: - monkeypatch.setattr(omprpc, "_CALLBACK_EXIT_TIMEOUT_S", 0) client = await client_factory("settled-during-callback") - process = client._proc - assert process is not None + client.callback_exit_timeout_s = 0 + process = owned_rpc_process(client) held = asyncio.Event() async def on_tool(_call: OmpHostToolCall) -> OmpToolResult: @@ -113,7 +120,7 @@ async def on_tool(_call: OmpHostToolCall) -> OmpToolResult: return OmpToolResult(content=[]) observer = Observer() - with pytest.raises(OmpProcessError, match="callbacks did not finish"): + with pytest.raises(OmpProcessError, match=CALLBACKS_UNFINISHED): async with asyncio.timeout(5): await client.prompt( "first", From 9ac09f4495e5f272fbc0ac46b0d0745487ddb2a1 Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 16:25:32 +0300 Subject: [PATCH 19/35] docs: make test conventions and review catch intent, seams, and branches The previous guidance named coverage in the abstract, so reviews treated plumbing assertions, private patches, and wall-clock grace as style nits instead of blocking test-quality issues. Co-authored-by: Cursor --- .../skills/band-code-review/angle-groups.md | 37 +++- .claude/skills/band-here-is-a-review/SKILL.md | 4 +- .claude/skills/band-plan/plan-template.md | 12 +- .../skills/band-python-conventions/SKILL.md | 19 +- .../band-python-test-conventions/SKILL.md | 200 +++++++++++++----- .claude/skills/band-write-e2e-test/SKILL.md | 6 +- AGENTS.md | 7 +- 7 files changed, 209 insertions(+), 76 deletions(-) diff --git a/.claude/skills/band-code-review/angle-groups.md b/.claude/skills/band-code-review/angle-groups.md index 1bc52381a..72664715b 100644 --- a/.claude/skills/band-code-review/angle-groups.md +++ b/.claude/skills/band-code-review/angle-groups.md @@ -37,15 +37,34 @@ and consequences differ. contents, secrets). **D — Test Quality** -Test code gets every angle above, more strictly, plus: -- Apply `band-python-test-conventions` "Coverage and realistic flows" to the - changed behavior and relevant existing tests, even when no test file changed. - Map scenarios to meaningful assertions, identify gaps, and name the concrete - regression each proposed test would catch. Check complete complex flows at - a boundary that preserves their real interactions; passing tests or high - line coverage alone are not sufficient evidence. -- A skipped, disabled, or silently-ignored failing test is always a blocking - issue (`band-python-test-conventions`). +Test code belongs here alone. Read the whole `band-python-test-conventions` +skill, not only "Coverage is branches, not lines". Apply A–C more strictly +than on production code (SSoT, intent-named helpers, no mixed-concern files, +no magic wire strings). Each of the following is an `issue:`, not a nit: + +- Coverage of complementary branches, not line counts. Map each changed + `match`/`if`/timeout/OS decision to the test that proves each reachable + outcome. Proving fail-closed does not prove the success side; a + `skipif(os.name)` is a missing branch on the skipped OS. Name the + regression each proposed test would catch. Passing tests or high line + coverage alone are not evidence. +- Assertions of plumbing or assumptions: occupancy arithmetic, private + fields, mock call counts, event-type tallies as a proxy, empty collections + the test does not own, or tautologies. The assertion must be a named + projection of the decision. +- A test function that branches on scenario/kind (`if scenario != ...`) or + re-types a wire literal the peer or product already names. Parametrize + `(scenario, expected)` so each row *is* the claim. +- The interaction under test is patched away (`unittest.mock.patch` of + private names, stdlib spawn, or timeout constants) instead of a public + seam or a real peer. `AsyncMock` of the unit under test is the same smell. +- Wall-clock sleep or waiting out a grace period used as proof. + `asyncio.timeout` is a hung-test guard only; a public timeout of `0` is + "check now", and it still needs the complementary "finished before wait" + sequence. +- A skipped, disabled, or silently-ignored failing test. +- A Hypothesis model or kit file that looks like a collected test (a `Test*` + class, or a missing role docstring, in a non-`test_*.py` module). **E — Edge-Case Flow Analysis** Owns failures that need a *sequence* of state transitions or interactions. A diff --git a/.claude/skills/band-here-is-a-review/SKILL.md b/.claude/skills/band-here-is-a-review/SKILL.md index f427e4cb1..5e23385e0 100644 --- a/.claude/skills/band-here-is-a-review/SKILL.md +++ b/.claude/skills/band-here-is-a-review/SKILL.md @@ -109,7 +109,9 @@ Only CONFIRMED and PARTIALLY-RIGHT findings that **matter** proceed to fixing. "Matters" means: user-visible breakage, data corruption, security exposure, a real maintenance trap, or a violated project invariant or rule (`AGENTS.md` and the guideline skills `band-python-conventions` and `band-python-test-conventions` — e.g. duplicated logic, -magic literals, narrating comments). Filter out: +magic literals, narrating comments, tests that assert plumbing or assumptions, +patch internals, wait out wall-clock grace, branch on scenario inside the test, +or treat line coverage as branch coverage). Filter out: - Pedantry and taste conflicts with the project's established conventions (the codebase's existing idiom wins over the reviewer's preference). diff --git a/.claude/skills/band-plan/plan-template.md b/.claude/skills/band-plan/plan-template.md index 599249381..03c386044 100644 --- a/.claude/skills/band-plan/plan-template.md +++ b/.claude/skills/band-plan/plan-template.md @@ -52,12 +52,20 @@ Test rules: - Test behavior that can actually break, not defaults or implementation trivia (see the `band-python-test-conventions` skill). A test that would still pass under a symptom-level patch doesn't cover the root cause. + Cover complementary branches of each owned decision, not line execution; + assert a named projection of the outcome, not internals or unclaimed + empties. Parametrize `(scenario, expected)` — do not `if` on scenario + inside the test. - When a material scenario needs end-to-end coverage to prove the behavior, reuse or extend the `tests/e2e/baseline` toolkit (see the `band-write-e2e-test` skill) instead of building new plumbing. Otherwise use the lowest boundary that preserves the behavior being tested. -- A missing helper, fixture, or local peer is its own step, before the step - that uses it — never an empty mock in place of a real boundary. +- A missing helper, fixture, constructor seam, or local peer is its own + step, before the step that uses it — never `unittest.mock.patch` of a + private name, stdlib spawn, or timeout constant, and never an empty mock + in place of a real boundary. Timeouts are proven with a public `0` ("check + now") plus the complementary "finished before wait" sequence, not by + sleeping the grace period. Scope rules: diff --git a/.claude/skills/band-python-conventions/SKILL.md b/.claude/skills/band-python-conventions/SKILL.md index 88dfe5e5f..bf209eb4a 100644 --- a/.claude/skills/band-python-conventions/SKILL.md +++ b/.claude/skills/band-python-conventions/SKILL.md @@ -1,6 +1,6 @@ --- name: band-python-conventions -description: Band Python SDK conventions for Python source, comments, docstrings, and error handling. Read before editing Python in this repository. +description: Band Python SDK conventions for Python source, comments, docstrings, and error handling. Read before editing Python in this repository. Test files also follow band-python-test-conventions; this skill still owns SSoT, intent-named helpers, and public test seams. paths: - "**/*.py" --- @@ -44,7 +44,12 @@ paths: - Use `%s` placeholders in log messages for lazy evaluation - Use Pydantic v2 for data models; use `model_dump()` not `dict()` - Target Python 3.11+; use `list[str]` not `List[str]`, `str | None` not `Optional[str]` -- Use async/await everywhere in async codebases; use `AsyncMock` for testing async methods +- Use async/await everywhere in async codebases. `AsyncMock` is a test double + at an external boundary, not a stand-in for the unit whose interaction a + test claims to prove. Public constructor seams (a `create_process` factory, + an instance timeout of `0` meaning "check now") beat `unittest.mock.patch` + of private names or stdlib spawn. Tests follow + `band-python-test-conventions`. - Catch `pydantic.ValidationError` separately from generic `Exception` - Use `raise ValueError(...)` for missing required config, not `logger.error()` + `sys.exit()` - Never put issue-tracker references in code — no Linear issue IDs (e.g. `INT-123`), Linear URLs, or ticket numbers in comments, docstrings, or strings. Explain the *why* in plain terms instead. (Branch names, commit messages, and PR descriptions may reference issues.) @@ -62,10 +67,12 @@ paths: - Prefer a single source of truth for a value or closed vocabulary consumed in more than one place: give it one definition — a constant, a `StrEnum`, or a small helper — that every site references, rather than re-typing the same magic literal in a - producer and the consumer that reads it (a typo then fails silently). Keep genuinely - distinct vocabularies separate, though — don't merge two sets that only happen to - share some values today (e.g. the ACP `ChunkType` a chunk carries vs. the platform - `message_type` an event is posted under). + producer and the consumer that reads it (a typo then fails silently). This includes + tests: a wire status, prompt text, or error message typed in a peer and again in an + assertion is two copies. Keep genuinely distinct vocabularies separate, though — + don't merge two sets that only happen to share some values today (e.g. the ACP + `ChunkType` a chunk carries vs. the platform `message_type` an event is posted + under). - Before writing new logic, search the codebase for an existing helper that does or almost does it; extending it beats a parallel implementation. - Use idiomatic constructs (`match`, decorators, context managers, recursion) diff --git a/.claude/skills/band-python-test-conventions/SKILL.md b/.claude/skills/band-python-test-conventions/SKILL.md index 7287c9350..aa4bdcf88 100644 --- a/.claude/skills/band-python-test-conventions/SKILL.md +++ b/.claude/skills/band-python-test-conventions/SKILL.md @@ -1,6 +1,6 @@ --- name: band-python-test-conventions -description: Band Python SDK test conventions and coverage assessment. Use when writing or reviewing tests, assessing whether changed behavior is adequately covered, or checking that tests validate meaningful outcomes and realistic complex flows. +description: Band Python SDK test conventions and coverage assessment. Use when writing or reviewing tests, assessing whether changed behavior is adequately covered, or checking that tests validate meaningful outcomes and realistic complex flows. Also use when a test asserts internals, patches private names, waits on wall-clock grace, branches on scenario inside the test, or reports line coverage instead of complementary branches. paths: - "tests/**/*.py" - "src/**/*.py" @@ -8,71 +8,165 @@ paths: # Test Conventions -## Coverage and realistic flows +`band-python-conventions` still applies: single source of truth, intent-named +helpers, and no mixed-concern files. This skill is the extra bar for tests. + +## Coverage is branches, not lines Assess the relevant existing tests whenever behavior changes or coverage is reviewed, even if the diff adds no tests. A green suite, a test count, or a line-coverage percentage alone does not establish adequate coverage. -- Map user-visible scenarios and important invariants to the tests and - assertions that prove them. Identify missing or weak coverage, prioritizing - consequential behavior, changed branches, and regressions over a numeric target. -- For each test, name the plausible bug it would catch. Assert the observable - outcome and relevant side effects; a test that only checks that a mock was - called or that nothing raised may leave the real behavior unverified. -- Exercise complex flows as realistic sequences, not just isolated steps. - Where relevant, cover follow-up turns, failure and recovery, cancellation, - restart/rehydration, concurrent rooms or agents, permissions, and cleanup. - Select scenarios the actual feature can reach; do not add every case to - every change. -- Drive the real entry point and interacting components at the lowest level - that preserves the behavior under test. Use realistic peers or fakes at - external boundaries; do not mock away the interaction whose correctness - the test claims to prove. Use the baseline toolkit for needed live E2E - coverage, not as a mandatory layer for every test. -- Extend useful existing tests before duplicating them. Add or strengthen - tests for meaningful in-scope gaps; report remaining gaps and unverified - scenarios explicitly. Follow `AGENTS.md` for synchronization and validation - scope, without silently expanding the task or launching unrelated suites. - -## Test implementation +Map each production decision the change owns — a `match`/`case`, an `if` that +chooses an outcome, an OS spawn path, a timeout pending-vs-done wait — onto +the tests that prove it. For every such decision: + +1. Name the reachable outcomes (succeed, fail-closed, ignore, retry, OS-A vs + OS-B). +2. Name the test that proves each as an observable sequence, not as a line + that executed. +3. Treat complementary pairs as separate scenarios. Proving one side does not + prove the other: backlog live-at-limit vs overflow; a callback that + finishes before the wait vs one still running; settlement before the + matching result vs result before settlement; POSIX process-group vs + Windows job. +4. Report remaining unexamined outcomes explicitly. A Hypothesis model covers + generated sequences of its Step alphabet, not branches those steps never + take. + +Prioritize consequential behavior, changed branches, and regressions over a +numeric target. Do not add every case to every change; select scenarios the +actual feature can reach. + +For each test, name the plausible bug it would catch. Assert the observable +outcome and relevant side effects; a test that only checks that a mock was +called or that nothing raised may leave the real behavior unverified. + +Exercise complex flows as realistic sequences, not just isolated steps. +Where relevant, cover follow-up turns, failure and recovery, cancellation, +restart/rehydration, concurrent rooms or agents, permissions, and cleanup. + +Extend useful existing tests before duplicating them. Follow `AGENTS.md` for +synchronization and validation scope, without silently expanding the task or +launching unrelated suites. + +## Assertions show intent, not plumbing or assumptions + +Tests must be declarative and intent-revealing, not a transcript of the +implementation. Assert on a readable projection of the observable outcome — +the thing the test is actually about — never on raw internals or on a side +effect that merely implies the real answer. + +- `assert reply.outline == ["tool_call (permission)", "message", ...]` over a + hand-rolled comprehension pulling `message_type` out of each event dict. +- `assert record.levelno == logging.DEBUG` over inferring a log level + indirectly from whether two separate capture windows came back empty. +- A named helper (`owned_rpc_process`, `outline`) over poking `_proc` or + counting `AGENT_END` as a proxy for "the turn finished twice". +- Overload of `EVENT_BACKLOG_LIMIT + 1` over occupancy arithmetic + (`limit - 2`) that restates the queue's internals. + +If writing the assertion requires re-deriving *how* the code decided +something, the test is checking the wrong thing — assert the decision +itself. + +Do not assert an assumption the test does not own. Drop empty-collection +checks on a mapping test that is about session keys, tautologies such as +Windows `(None, None)`, and "observer stayed empty" when the claim is +fail-closed-and-restarted. If the assertion would still pass after inverting +the behavior on an unclaimed axis, it is noise, not coverage. + +Test what really matters — behavior that can break. Don't write tests that +restate definitions (asserting dataclass defaults equal themselves, echoing a +constant) or otherwise cannot fail for a real reason. + +Decide what correct looks like before running the code, then assert that; +never bend an assertion to match whatever the code currently produces. + +## One concern per test, no in-test branching + +Each `test_*.py` holds exactly one `class Test...` named for the behavior it +covers — no module-level test functions and no helper classes beside it. +When a suite outgrows one file, make it a package: `conftest.py` for +fixtures, a plain helper module (`kit.py`, a Hypothesis `lifecycle.py`) for +shared doubles, and one test file per concern. + +A helper module in that package is not a collected test. Its module +docstring names the role (kit, peer, lifecycle model). It does not define a +`Test*` class. The collected test that *plays* a Hypothesis model lives in +`test_lifecycle_model.py` (see `tests/platform/linkclose/test_lifecycle_model.py`). + +One test function, one claim. Parametrize `(scenario, error)` so each row's +assertion *is* the claim. A long test that dispatches with `if scenario != +"aborted"` or `if os.name == "nt": return` is several tests hiding in one +function; split them. + +Wire statuses, prompt texts, and error messages belong to one named constant +or enum that the peer, the product, and the assertion all reference. A +hard-coded `"complete"` / `"Approve"` in the test is a second copy of a +closed vocabulary. + +## Drive the real boundary; inject, don't patch + +Drive the real entry point and interacting components at the lowest level +that preserves the behavior under test. Use realistic peers or fakes at +external boundaries; do not mock away the interaction whose correctness the +test claims to prove. + +- A constructor seam (`create_process=`) or a public instance timeout of `0` + ("check now / no grace") is the test hook. `unittest.mock.patch` of + `asyncio.create_subprocess_exec`, a private `_close_process`, or a + `_CALLBACK_EXIT_TIMEOUT_S` constant is not. +- `AsyncMock` is a double at an external boundary (a client the function + *takes*), not a stand-in for the unit under test. A receipt-table test may + mock the RPC client; it then does not count as coverage of the client's + prompt path. +- Use the baseline toolkit for needed live E2E coverage, not as a mandatory + layer for every test. + +## Time is a hung-test guard, not proof + +Ordering, concurrency, and timeout behavior is proven with explicit +synchronization and a controlled clock (`AGENTS.md` "Test Synchronization"): +events, gates, barriers, or observable state transitions. Never use elapsed +wall time, `asyncio.sleep` of a grace period, or "it finished within a +duration" as proof. + +- `async with asyncio.timeout(N)` around a test is a hung-test guard. It + must not be the mechanism that establishes the outcome. +- To prove a timeout fail-closed path, set the public timeout to `0` and + assert the failure. Do not wait out the production grace period. +- `timeout = 0` proves "still pending at the wait". It does not prove the + complementary branch: work that finishes before the wait succeeds. Cover + that with a sequence that completes, not with a longer sleep. + +Real subprocess and network readiness stay independent of a controlled +clock. Cleanup must account for resources created during failed or cancelled +startup and await their confirmed release. + +## Tests run on Ubuntu and Windows CI + +Test across OSes; don't skip one. Build host paths with `tmp_path` or +`tests.paths.host_absolute_path`, never a hard-coded `"/opt/..."` (no drive, +so relative on Windows); check absoluteness with +`band.workspaces.is_host_absolute`. + +Where behavior differs per OS, write one parametrize table that runs +everywhere, each row naming the `os.name` values that accept it. +`skipif(os.name ...)` is only for a boundary that cannot run on that OS, +never for a path, separator, or "this assertion is awkward on Windows". A +skip is a missing branch on the skipped OS. + +## Paths, bugs, skips - In tests, never derive repository-anchored paths with per-file `Path(__file__).parents[N]` arithmetic — import the anchors from `tests/paths.py` (`REPO_ROOT`, `SRC_ROOT`, `EXAMPLES_ROOT`, `KIT_DIR`, `ENV_TEST_FILE`). Only genuinely package-relative paths (a fixture file next to its test) stay relative to their own `__file__`. -- Test what really matters — behavior that can break. Don't write tests that - restate definitions (asserting dataclass defaults equal themselves, echoing a - constant) or otherwise cannot fail for a real reason; they add maintenance - cost without protection. -- **Tests must be declarative and intent-revealing, not a transcript of the - implementation.** Assert on a readable projection of the observable outcome - — the thing the test is actually about — never on raw internals or on a - side effect that merely implies the real answer. Concretely: - - `assert reply.outline == ["tool_call (permission)", "message", ...]` over - a hand-rolled comprehension pulling `message_type` out of each event dict. - - `assert record.levelno == logging.DEBUG` over inferring a log level - indirectly from whether two separate capture windows came back empty. - If writing the assertion requires re-deriving *how* the code decided - something, the test is checking the wrong thing — assert the decision - itself. -- **Tests run on Ubuntu and Windows CI — test across OSes, don't skip one.** - Build host paths with `tmp_path` or `tests.paths.host_absolute_path`, never a - hard-coded `"/opt/..."` (no drive, so relative on Windows); check absoluteness - with `band.workspaces.is_host_absolute`. Where behavior differs per OS, write - one parametrize table that runs everywhere, each row naming the `os.name` - values that accept it. `skipif(os.name ...)` is only for a boundary that - cannot run on that OS, never for a path or separator assumption. -- Decide what correct looks like before running the code, then assert that; - never bend an assertion to match whatever the code currently produces. - A skipped, disabled, or silently ignored failing test is always a blocking issue. - A fix for a confirmed bug comes with a test that fails before the fix and passes after it. - For a stateful lifecycle component, add a Hypothesis model-based test over - generated step sequences (see `tests/platform/linkclose/test_lifecycle_model.py`). -- Ordering, concurrency and timeout behavior is proven with explicit - synchronization and a controlled clock: `AGENTS.md` "Test Synchronization". -- File layout (one test class per file, packages for larger suites): - `AGENTS.md` "Testing Structure". + generated step sequences. diff --git a/.claude/skills/band-write-e2e-test/SKILL.md b/.claude/skills/band-write-e2e-test/SKILL.md index ec98bb44f..54cf6a00b 100644 --- a/.claude/skills/band-write-e2e-test/SKILL.md +++ b/.claude/skills/band-write-e2e-test/SKILL.md @@ -79,9 +79,9 @@ reuse, target location, precedents, and any coverage overlap). ## Phase 2 — Plan the test Apply [test conventions](../band-python-test-conventions/SKILL.md) -"Coverage and realistic flows" when mapping existing coverage and selecting -the scenario. Show which meaningful gap this test closes and which complete -interaction it proves. +"Coverage is branches, not lines" when mapping existing coverage and selecting +the scenario. Show which complementary-branch gap this test closes and which +complete interaction it proves. Write a concrete plan (not code yet). It must answer: diff --git a/AGENTS.md b/AGENTS.md index ffb1ee70d..55a5541b4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -80,7 +80,7 @@ explicit user approval. | Skill | Use when | |---|---| | [band-python-conventions](.claude/skills/band-python-conventions/SKILL.md) | Python source, comments, docstrings, and error handling | -| [band-python-test-conventions](.claude/skills/band-python-test-conventions/SKILL.md) | writing/reviewing tests or assessing meaningful coverage and realistic flows | +| [band-python-test-conventions](.claude/skills/band-python-test-conventions/SKILL.md) | writing/reviewing tests; complementary branches not line counts; intent assertions; seams not patches; no wall-clock proof | | Skill | Use when | |---|---| @@ -447,7 +447,10 @@ install` it and exercise the real call in this repo's venv. For Python changes, before selecting checks, re-read your own diff once against the `band-python-conventions` skill — including code you just wrote this session, -not only code you started from. The two rules that get skipped under time pressure: +not only code you started from. Test diffs also re-read +`band-python-test-conventions` (complementary branches, intent assertions, +seams not patches, no wall-clock proof). The two rules that get skipped under +time pressure: - **Single source of truth**: a literal, magic string, or multi-line block re-typed in more than one place (a second copy you just wrote counts) instead From d867375c6ede2613f40fa9d775f35c1c8b82f987 Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 17:27:39 +0300 Subject: [PATCH 20/35] fix: apply native RPC ready limits before the next frame The reader must honor negotiated frame limits immediately, stdin close must wait for in-flight writes, and fail-closed parse errors must log through the transport. Tests now inject public close/join seams and cover the complementary settlement and cancel branches. Co-authored-by: Cursor --- src/band/integrations/ompprocess.py | 11 ++- src/band/integrations/omprpc/client.py | 67 ++++++++++++++----- src/band/integrations/omprpc/scope.py | 13 +++- tests/converters/test_omp.py | 4 +- tests/integrations/omppeer.py | 8 +++ .../ompprocess/test_mcp_plugin.py | 19 +++++- tests/integrations/omprpc/cancellation.py | 32 +++------ tests/integrations/omprpc/kit.py | 15 +++++ tests/integrations/omprpc/test_callbacks.py | 14 ++-- .../omprpc/test_eof_descendants.py | 15 ++--- tests/integrations/omprpc/test_fail_closed.py | 8 +++ .../integrations/omprpc/test_interruption.py | 65 ++++++++++-------- .../integrations/omprpc/test_process_tree.py | 58 +++++----------- .../omprpc/test_prompt_results.py | 9 ++- .../omprpc/test_turn_settlement.py | 33 +++++++++ 15 files changed, 235 insertions(+), 136 deletions(-) diff --git a/src/band/integrations/ompprocess.py b/src/band/integrations/ompprocess.py index 69bf09649..700c79f90 100644 --- a/src/band/integrations/ompprocess.py +++ b/src/band/integrations/ompprocess.py @@ -32,7 +32,11 @@ if TYPE_CHECKING: from band.integrations.omprpc import OmpRpcClient -DEFAULT_OMP_COMMAND: tuple[str, ...] = ("omp", "--mode", "rpc") +OMP_MODE_FLAG = "--mode" +OMP_RPC_MODE = "rpc" +OMP_RPC_UI_MODE = "rpc-ui" +DEFAULT_OMP_COMMAND: tuple[str, ...] = ("omp", OMP_MODE_FLAG, OMP_RPC_MODE) +RELATIVE_MCP_COMMAND_PREFIXES = ("./", "../", ".\\", "..\\") # The native RPC ask dialog and MCP status events need a newer OMP than the # ACP bridge's pin. @@ -174,7 +178,8 @@ def omp_mcp_arguments(plugin: str) -> tuple[str, ...]: return ( "--plugin-dir", plugin, - "--mode=rpc-ui", + OMP_MODE_FLAG, + OMP_RPC_UI_MODE, "--extension", str(Path(plugin) / OMP_MCP_STARTUP_JS), "--config", @@ -244,7 +249,7 @@ def _projected_server(server: OmpMcpServer, workspace: str) -> dict[str, object] # Plugin discovery otherwise rebases launch paths into the temporary overlay. cwd = Path(os.path.abspath(Path(workspace) / (server.cwd or "."))) payload["cwd"] = str(cwd) - if server.command.startswith(("./", "../", ".\\", "..\\")): + if server.command.startswith(RELATIVE_MCP_COMMAND_PREFIXES): payload["command"] = os.path.abspath(cwd / server.command) return payload diff --git a/src/band/integrations/omprpc/client.py b/src/band/integrations/omprpc/client.py index f6633dc65..076c65708 100644 --- a/src/band/integrations/omprpc/client.py +++ b/src/band/integrations/omprpc/client.py @@ -8,7 +8,7 @@ from binascii import Error as Base64Error from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass -from typing import TypeVar, cast +from typing import Any, TypeVar, cast from pydantic import JsonValue, ValidationError @@ -62,6 +62,8 @@ TModel = TypeVar("TModel", bound=OmpWireModel) CreateProcess = Callable[..., Awaitable[asyncio.subprocess.Process]] +CloseProcess = Callable[[], Coroutine[Any, Any, None]] +JoinClose = Callable[[Awaitable[None]], Awaitable[None]] logger = logging.getLogger(__name__) @@ -102,6 +104,8 @@ def __init__( self.command = tuple(command) self.env = None if env is None else {**os.environ, **env} self.create_process = create_process + self.close_process: CloseProcess = self._close_process + self.join_close: JoinClose = asyncio.shield self.callback_exit_timeout_s = CALLBACK_EXIT_TIMEOUT_S self.process_exit_timeout_s = PROCESS_EXIT_TIMEOUT_S self._proc: asyncio.subprocess.Process | None = None @@ -126,7 +130,7 @@ def __init__( self._write_tasks: set[asyncio.Task[None]] = set() async def start(self) -> None: - """Start the supplied command and require native RPC protocol v2.""" + """Start the supplied command and require the native RPC protocol.""" async with self._start_lock: if self._closing: await self.close() @@ -159,11 +163,12 @@ async def start(self) -> None: if isinstance(ready, OmpProcessError): raise ready if OMP_NATIVE_PROTOCOL_VERSION not in ready.supportedProtocolVersions: - raise OmpProtocolError( + error = OmpProtocolError( f"OMP RPC protocol v{OMP_NATIVE_PROTOCOL_VERSION} is required" ) - self._max_frame_bytes = ready.maxFrameBytes - self._frames.max_reassembled_bytes = ready.maxReassembledFrameBytes + self._record_failure(error) + raise error + self._apply_ready_limits(ready) negotiated = await self._request_frame( "negotiate_protocol", {"protocolVersion": OMP_NATIVE_PROTOCOL_VERSION}, @@ -235,7 +240,7 @@ async def close(self) -> None: interrupted = False while True: try: - await asyncio.shield(closing) + await self.join_close(closing) break except asyncio.CancelledError: if closing.cancelled(): @@ -300,7 +305,10 @@ def _parse_response( try: return model.model_validate(self._response_data(response)) except ValidationError as cause: - raise OmpProtocolError(error) from cause + failure = OmpProtocolError(error) + failure.__cause__ = cause + self._record_failure(failure) + raise failure from cause async def _execute_prompt( self, @@ -335,6 +343,7 @@ async def _execute_prompt( await self._cancel_tasks( [work, event_worker, *turn.callback_tasks], timeout=self.callback_exit_timeout_s, + label="callback", ) async def _complete_prompt(self, turn: PromptRun, message: str) -> None: @@ -361,7 +370,9 @@ async def _complete_prompt(self, turn: PromptRun, message: str) -> None: turn.callback_tasks, timeout=self.callback_exit_timeout_s ) if pending: - raise OmpProcessError(CALLBACKS_UNFINISHED) + error = OmpProcessError(CALLBACKS_UNFINISHED) + self._record_failure(error) + raise error if outcome is not None and outcome.status != OmpPromptStatus.COMPLETED: raise OmpPromptError( outcome.status, @@ -465,8 +476,11 @@ async def _dispatch_payload( match frame.type: case OmpFrameType.READY: ready = OmpReady.model_validate(payload) + self._apply_ready_limits(ready) if self._ready is not None and not self._ready.done(): self._ready.set_result(ready) + else: + logger.debug("OMP RPC ignored a duplicate ready frame") case OmpFrameType.FRAME_ERROR: overflow = OmpFrameError.model_validate(payload) raise OmpProtocolError( @@ -530,6 +544,8 @@ async def _dispatch_payload( event = OmpEvent.model_validate(payload) if turn is not None: turn.observe(event, sequence) + else: + logger.debug("OMP RPC ignored an event without an active prompt") def _dispatch_response(self, response: OmpResponse, sequence: int) -> None: if response.id is None: @@ -542,6 +558,8 @@ def _dispatch_response(self, response: OmpResponse, sequence: int) -> None: and response.id == self._turn.id ): self._turn.fail(command_error(response)) + return + logger.debug("OMP RPC ignored an unmatched command response") return if response.command != pending.command: raise OmpProtocolError("OMP response command does not match its request") @@ -642,6 +660,10 @@ def _fail_pending(self, error: OmpProcessError) -> None: if self._turn is not None: self._turn.fail(error) + def _apply_ready_limits(self, ready: OmpReady) -> None: + self._max_frame_bytes = ready.maxFrameBytes + self._frames.max_reassembled_bytes = ready.maxReassembledFrameBytes + def _record_failure(self, error: OmpProcessError) -> None: if self._failure is None: self._failure = error @@ -660,7 +682,7 @@ def _begin_close(self) -> None: self._close_task.cancelled() or self._close_task.exception() is not None ) ): - self._close_task = asyncio.create_task(self._close_process()) + self._close_task = asyncio.create_task(self.close_process()) self._close_task.add_done_callback(self._log_close_failure) @staticmethod @@ -675,7 +697,10 @@ def _consume_task(task: asyncio.Task[None]) -> None: @staticmethod async def _cancel_tasks( - tasks: Sequence[asyncio.Task[None]], *, timeout: float + tasks: Sequence[asyncio.Task[None]], + *, + timeout: float, + label: str = "callback", ) -> None: active = {task for task in tasks if not task.done()} for task in active: @@ -684,7 +709,7 @@ async def _cancel_tasks( _, pending = await asyncio.wait(active, timeout=timeout) for task in pending: task.cancel() - logger.warning("OMP callback did not honor cancellation") + logger.warning("OMP %s did not honor cancellation", label) for task in tasks: if task.done(): OmpRpcClient._consume_task(task) @@ -720,7 +745,9 @@ async def _close_process(self) -> None: workers = list(turn.callback_tasks) if turn.event_worker is not None: workers.append(turn.event_worker) - await self._cancel_tasks(workers, timeout=self.callback_exit_timeout_s) + await self._cancel_tasks( + workers, timeout=self.callback_exit_timeout_s, label="callback" + ) owner = self._windows_process scope = self._process_scope @@ -762,8 +789,14 @@ async def _close_process(self) -> None: if scope is not None: # Ownership lookup can disappear before process exit is observable. await asyncio.to_thread(scope.capture, process) + await self._cancel_tasks( + list(self._write_tasks), + timeout=self.callback_exit_timeout_s, + label="write", + ) if not owner_closed: - await close_stdin(process) + async with self._write_lock: + await close_stdin(process) await await_returncode(process, self.process_exit_timeout_s) if owner is not None: if not owner_closed: @@ -788,9 +821,13 @@ async def _close_process(self) -> None: # A cancelled pipe read on Windows stays pending until the child is # killed; drain stdout only after that reader finishes. if reader is not None: - await self._cancel_tasks([reader], timeout=self.callback_exit_timeout_s) + await self._cancel_tasks( + [reader], timeout=self.callback_exit_timeout_s, label="reader" + ) if process.stdout is not None: drains.append(asyncio.create_task(self._drain(process.stdout))) if self._stderr_task is not None: drains.append(self._stderr_task) - await self._cancel_tasks(drains, timeout=self.callback_exit_timeout_s) + await self._cancel_tasks( + drains, timeout=self.callback_exit_timeout_s, label="drain" + ) diff --git a/src/band/integrations/omprpc/scope.py b/src/band/integrations/omprpc/scope.py index 6caf27aa5..c3502609a 100644 --- a/src/band/integrations/omprpc/scope.py +++ b/src/band/integrations/omprpc/scope.py @@ -14,6 +14,8 @@ from band.integrations.omprpc.errors import OmpProcessError _PROCESS_SCOPE_VALUE = "1" +# pid 0 is the kernel idle process; scanning it is wasted work. +KERNEL_IDLE_PID = 0 class OmpProcessScope: @@ -41,7 +43,7 @@ def capture(self, root: asyncio.subprocess.Process) -> None: getpgid = getattr(os, "getpgid", None) for pid in psutil.pids(): # asyncio's child watcher remains the sole waiter for the root. - if pid == 0 or pid == root.pid: + if pid == KERNEL_IDLE_PID or pid == root.pid: continue try: member = psutil.Process(pid) @@ -49,7 +51,12 @@ def capture(self, root: asyncio.subprocess.Process) -> None: marked = False try: marked = member.environ().get(self.key) == _PROCESS_SCOPE_VALUE - except (psutil.NoSuchProcess, psutil.AccessDenied): + except ( + psutil.NoSuchProcess, + psutil.AccessDenied, + PermissionError, + SystemError, + ): pass in_group = False if group_visible and getpgid is not None: @@ -68,7 +75,7 @@ def capture(self, root: asyncio.subprocess.Process) -> None: group_members.append(member) except (psutil.NoSuchProcess, psutil.AccessDenied): continue - if group_visible and not session_pid_reused(root): + if group_visible: self.processes.update(group_members) @staticmethod diff --git a/tests/converters/test_omp.py b/tests/converters/test_omp.py index cdccf9384..eaa7f5244 100644 --- a/tests/converters/test_omp.py +++ b/tests/converters/test_omp.py @@ -77,7 +77,7 @@ def test_session_metadata_on_non_task_messages_is_ignored(self) -> None: state = OmpHistoryConverter().convert( [ { - "message_type": "text", + "message_type": MessageType.TEXT, "metadata": { OMP_ROOM_METADATA_KEY: "room-1", OMP_SESSION_METADATA_KEY: "/sessions/injected.jsonl", @@ -95,7 +95,7 @@ def test_room_transcript_is_replayed(self) -> None: state = OmpHistoryConverter().convert( [ { - "message_type": "text", + "message_type": MessageType.TEXT, "content": "hello", "sender_name": "Alice", "sender_type": "User", diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index c77fe6f1b..bbeacce95 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -705,6 +705,14 @@ def handle(self, frame: Frame) -> None: } ) self.respond(frame) + case "orphan-ui-cancel": + send( + { + "type": "extension_ui_request", + "id": "ui-cancel-orphan", + "method": "cancel", + } + ) case "checkpoint": self.respond(frame, {"pid": os.getpid(), "frames": self.frames}) case "abort": diff --git a/tests/integrations/ompprocess/test_mcp_plugin.py b/tests/integrations/ompprocess/test_mcp_plugin.py index 75cacbdd0..f9de08028 100644 --- a/tests/integrations/ompprocess/test_mcp_plugin.py +++ b/tests/integrations/ompprocess/test_mcp_plugin.py @@ -13,6 +13,9 @@ OMP_MCP_SERVERS_CONFIG, OMP_MCP_STARTUP_CONFIG, OMP_MCP_STARTUP_JS, + OMP_MODE_FLAG, + OMP_RPC_UI_MODE, + RELATIVE_MCP_COMMAND_PREFIXES, OmpMcpServer, omp_mcp_arguments, omp_mcp_plugin, @@ -80,7 +83,7 @@ def test_arguments_load_the_plugin_and_its_startup_extension( plugin = str(tmp_path / "band-sdk") arguments = omp_mcp_arguments(plugin) assert arguments[:2] == ("--plugin-dir", plugin) - assert "--mode=rpc-ui" in arguments + assert arguments[arguments.index(OMP_MODE_FLAG) + 1] == OMP_RPC_UI_MODE assert str(Path(plugin) / OMP_MCP_STARTUP_JS) in arguments assert str(Path(plugin) / OMP_MCP_STARTUP_CONFIG) in arguments @@ -88,6 +91,8 @@ def test_arguments_load_the_plugin_and_its_startup_extension( def test_projected_plugin_files_are_owner_only( self, tmp_path: Path, accepted_on: tuple[str, ...] ) -> None: + if os.name not in accepted_on: + pytest.skip("Unix permission bits are not the Windows ACL boundary") servers = {"local": OmpMcpServer(command="./server")} with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: assert plugin is not None @@ -99,9 +104,17 @@ def test_projected_plugin_files_are_owner_only( ) assert root.is_dir() assert all(path.is_file() for path in files) - assert (plugin_privacy(root, files) == (0o700, 0o600)) is ( - os.name in accepted_on + assert plugin_privacy(root, files) == (0o700, 0o600) + + @pytest.mark.parametrize("prefix", RELATIVE_MCP_COMMAND_PREFIXES) + def test_relative_command_is_absolutized(self, tmp_path: Path, prefix: str) -> None: + servers = {"local": OmpMcpServer(command=f"{prefix}server")} + with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: + assert plugin is not None + config = json.loads( + (Path(plugin) / OMP_MCP_SERVERS_CONFIG).read_text(encoding="utf-8") ) + assert Path(config["mcpServers"]["local"]["command"]).is_absolute() @pytest.mark.parametrize( "fields", diff --git a/tests/integrations/omprpc/cancellation.py b/tests/integrations/omprpc/cancellation.py index 55d3aa71c..82621a3eb 100644 --- a/tests/integrations/omprpc/cancellation.py +++ b/tests/integrations/omprpc/cancellation.py @@ -4,43 +4,35 @@ import asyncio from collections.abc import Awaitable -from typing import TypeVar - -import pytest from band.integrations.omprpc import OmpRpcClient -Result = TypeVar("Result") - class HeldCleanup: - def __init__(self, client: OmpRpcClient, monkeypatch: pytest.MonkeyPatch) -> None: + def __init__(self, client: OmpRpcClient) -> None: self.client = client self.entered = asyncio.Event() self.release = asyncio.Event() self.finished = asyncio.Event() self.task: asyncio.Task[object] | None = None self.joiners: asyncio.Queue[asyncio.Task[object] | None] = asyncio.Queue() - close_process = client._close_process - shield = asyncio.shield + closer = client.close_process + join_close = client.join_close async def hold_cleanup() -> None: self.task = asyncio.current_task() self.entered.set() await self.release.wait() - await close_process() + await closer() self.finished.set() - def observe_shield(operation: Awaitable[Result]) -> asyncio.Future[Result]: - waiter = shield(operation) - if operation is client._close_task: - # Queueing does not suspend the caller: it reaches the real shield - # wait before the observing test can resume. - self.joiners.put_nowait(asyncio.current_task()) + def observe_join(operation: Awaitable[None]) -> Awaitable[None]: + waiter = join_close(operation) + self.joiners.put_nowait(asyncio.current_task()) return waiter - monkeypatch.setattr(client, "_close_process", hold_cleanup) - monkeypatch.setattr(asyncio, "shield", observe_shield) + client.close_process = hold_cleanup + client.join_close = observe_join async def wait_for_join(self, caller: asyncio.Task[None]) -> None: async with asyncio.timeout(5): @@ -48,15 +40,9 @@ async def wait_for_join(self, caller: asyncio.Task[None]) -> None: joining = await self.joiners.get() assert joining is caller assert self.task is not None - assert self.client._close_task is self.task assert not self.task.done() - assert self.task.cancelling() == 0 assert not caller.done() async def cancel_and_wait_for_rejoin(self, caller: asyncio.Task[None]) -> None: - cancellations = caller.cancelling() assert caller.cancel() - # The next join is emitted only after close catches this cancellation - # and suspends on a fresh shield around the still-held cleanup task. await self.wait_for_join(caller) - assert caller.cancelling() == cancellations + 1 diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py index 1712c6011..4baa1f807 100644 --- a/tests/integrations/omprpc/kit.py +++ b/tests/integrations/omprpc/kit.py @@ -16,6 +16,7 @@ import pytest from pydantic import JsonValue +from band.integrations.codex.winspawn import Job from band.integrations.omprpc import OmpRpcClient from band.integrations.ompwire import ( OmpContent, @@ -250,6 +251,20 @@ def refuse_process_signal(*args: object) -> None: raise PermissionError("Process ownership is temporarily not signalable") +def refuse_containment(monkeypatch: pytest.MonkeyPatch) -> None: + if os.name == "nt": + monkeypatch.setattr(Job, "terminate", refuse_process_signal) + return + monkeypatch.setattr(os, "killpg", refuse_process_signal) + + +def session_ids(pids: Sequence[int]) -> list[int] | None: + getsid = getattr(os, "getsid", None) + if getsid is None: + return None + return [getsid(pid) for pid in pids] + + class DeferredSessionSignal: """Keep real members alive while post-signal ownership lookup disappears.""" diff --git a/tests/integrations/omprpc/test_callbacks.py b/tests/integrations/omprpc/test_callbacks.py index e866bb780..039926534 100644 --- a/tests/integrations/omprpc/test_callbacks.py +++ b/tests/integrations/omprpc/test_callbacks.py @@ -14,10 +14,12 @@ OmpUIRequest, OmpUIResponse, ) +from tests.integrations.omppeer import UI_APPROVE, UI_DENY from tests.integrations.omprpc.kit import ( ClientFactory, Observer, checkpoint_frames, + owned_rpc_process, second_prompt, ) @@ -54,7 +56,7 @@ async def on_tool(call: OmpHostToolCall) -> OmpToolResult: async def on_ui(request: OmpUIRequest) -> OmpUIResponse: approvals.append(request) - return OmpUIResponse(value="Approve") + return OmpUIResponse(value=UI_APPROVE) async with asyncio.timeout(5): await client.prompt( @@ -62,7 +64,7 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: ) assert {call.id for call in calls} == {"host-1", "host-2"} assert [(request.method, request.options) for request in approvals] == [ - ("select", ["Approve", "Deny"]) + ("select", [UI_APPROVE, UI_DENY]) ] frames = await checkpoint_frames(client) result = next(frame for frame in frames if frame.get("id") == "host-2") @@ -134,7 +136,7 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: await forever.wait() finally: cancelled.set() - return OmpUIResponse(value="Approve") + return OmpUIResponse(value=UI_APPROVE) running = asyncio.create_task( client.prompt( @@ -160,8 +162,7 @@ async def test_ui_callback_failure_fails_closed_and_reaps_the_process( client_factory: ClientFactory, ) -> None: client = await client_factory("ui-error") - process = client._proc - assert process is not None + process = owned_rpc_process(client) async def on_ui(request: OmpUIRequest) -> OmpUIResponse: raise ValueError("Unoffered permission choice") @@ -182,8 +183,7 @@ async def test_tool_callback_exception_propagates_unchanged_and_reaps_the_proces client_factory: ClientFactory, ) -> None: client = await client_factory("callbacks") - process = client._proc - assert process is not None + process = owned_rpc_process(client) observer = Observer() sentinel = ValueError("Band delivery failed") diff --git a/tests/integrations/omprpc/test_eof_descendants.py b/tests/integrations/omprpc/test_eof_descendants.py index 76d77f2bf..8b7a5f4c2 100644 --- a/tests/integrations/omprpc/test_eof_descendants.py +++ b/tests/integrations/omprpc/test_eof_descendants.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio -import os import threading import pytest @@ -12,13 +11,14 @@ ProcessGate, descendant_client, descendant_liveness, + owned_rpc_process, process_is_alive, + session_ids, unrelated_gate_child, ) class TestEofDescendants: - @pytest.mark.skipif(os.name == "nt", reason="POSIX detached-helper boundary") @pytest.mark.parametrize("stdio", ["inherit", "discard"]) async def test_close_reaps_a_detached_child_created_after_stdin_eof( self, @@ -38,9 +38,9 @@ async def test_close_reaps_a_detached_child_created_after_stdin_eof( try: await client.start() helper = await process_gate.connection() - root = client._proc + root = owned_rpc_process(client) scope = client._process_scope - assert root is not None and scope is not None + assert scope is not None signal = scope.signal def hold_signal() -> None: @@ -60,10 +60,9 @@ def hold_signal() -> None: late_child = await process_gate.connection() assert root.returncode is not None assert late_child.receipt.parent_pid == helper.process.pid - assert ( - os.getsid(late_child.process.pid) - == late_child.process.pid - ) + ids = session_ids([late_child.process.pid]) + if ids is not None: + assert ids == [late_child.process.pid] assert await descendant_liveness([helper, late_child]) == [ True, True, diff --git a/tests/integrations/omprpc/test_fail_closed.py b/tests/integrations/omprpc/test_fail_closed.py index ccbc648ba..341a25161 100644 --- a/tests/integrations/omprpc/test_fail_closed.py +++ b/tests/integrations/omprpc/test_fail_closed.py @@ -25,6 +25,14 @@ async def test_host_tool_without_an_active_prompt_fails_the_transport( with pytest.raises(OmpProtocolError, match="without an active prompt"): await client.request("orphan-tool") + async def test_ui_cancel_without_target_id_fails_the_transport( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("basic") + async with asyncio.timeout(6): + with pytest.raises(OmpProtocolError, match="missing targetId"): + await client.request("orphan-ui-cancel") + async def test_interactive_ui_without_an_active_prompt_is_cancelled( self, client_factory: ClientFactory ) -> None: diff --git a/tests/integrations/omprpc/test_interruption.py b/tests/integrations/omprpc/test_interruption.py index 6adda260c..7031d1f18 100644 --- a/tests/integrations/omprpc/test_interruption.py +++ b/tests/integrations/omprpc/test_interruption.py @@ -11,6 +11,7 @@ from tests.integrations.omprpc.kit import ( ClientFactory, Observer, + owned_rpc_process, prompt_task, queue_prompt, second_prompt, @@ -23,8 +24,7 @@ async def test_cancelling_a_queued_prompt_preserves_the_active_and_next_prompt( self, client_factory: ClientFactory ) -> None: client = await client_factory("yield") - process = client._proc - assert process is not None + process = owned_rpc_process(client) observer = Observer() active = prompt_task(client, observer) queued: asyncio.Task[None] | None = None @@ -42,7 +42,7 @@ async def test_cancelling_a_queued_prompt_preserves_the_active_and_next_prompt( async with asyncio.timeout(5): await active await second_prompt(client) - assert client._proc is process + assert owned_rpc_process(client) is process assert process.returncode is None assert await sent_prompt_messages(client) == ["first", "second"] finally: @@ -52,34 +52,49 @@ async def test_cancelling_a_queued_prompt_preserves_the_active_and_next_prompt( await asyncio.gather(*tasks, return_exceptions=True) @pytest.mark.parametrize("scenario", ["yield", "no-result", "stubborn"]) - @pytest.mark.parametrize("interruption", ["timeout", "cancel"]) - async def test_interrupted_prompt_aborts_reaps_and_can_restart( - self, client_factory: ClientFactory, scenario: str, interruption: str + async def test_cancelled_prompt_aborts_reaps_and_can_restart( + self, client_factory: ClientFactory, scenario: str ) -> None: client = await client_factory(scenario) - process = client._proc - assert process is not None + process = owned_rpc_process(client) observer = Observer() running = prompt_task(client, observer) try: await observer.wait_for(OmpEventType.AGENT_END) - if interruption == "cancel": - running.cancel() - async with asyncio.timeout(6): - with pytest.raises(asyncio.CancelledError): + running.cancel() + async with asyncio.timeout(6): + with pytest.raises(asyncio.CancelledError): + await running + assert process.returncode is not None + async with asyncio.timeout(5): + await client.start() + await second_prompt(client) + assert owned_rpc_process(client) is not process + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + @pytest.mark.parametrize("scenario", ["yield", "no-result", "stubborn"]) + async def test_timed_out_prompt_aborts_reaps_and_can_restart( + self, client_factory: ClientFactory, scenario: str + ) -> None: + client = await client_factory(scenario) + process = owned_rpc_process(client) + observer = Observer() + running = prompt_task(client, observer) + try: + await observer.wait_for(OmpEventType.AGENT_END) + async with asyncio.timeout(6): + with pytest.raises(TimeoutError): + async with asyncio.timeout(None) as deadline: + deadline.reschedule(asyncio.get_running_loop().time()) await running - else: - async with asyncio.timeout(6): - with pytest.raises(TimeoutError): - async with asyncio.timeout(None) as deadline: - deadline.reschedule(asyncio.get_running_loop().time()) - await running - assert deadline.expired() + assert deadline.expired() assert process.returncode is not None async with asyncio.timeout(5): await client.start() await second_prompt(client) - assert client._proc is not process + assert owned_rpc_process(client) is not process finally: running.cancel() await asyncio.gather(running, return_exceptions=True) @@ -87,18 +102,15 @@ async def test_interrupted_prompt_aborts_reaps_and_can_restart( async def test_repeatedly_cancelled_close_still_waits_for_reap( self, client_factory: ClientFactory, - monkeypatch: pytest.MonkeyPatch, ) -> None: client = await client_factory("stubborn") - process = client._proc - assert process is not None - cleanup = HeldCleanup(client, monkeypatch) + process = owned_rpc_process(client) + cleanup = HeldCleanup(client) closing = asyncio.create_task(client.close()) try: await cleanup.wait_for_join(closing) for _ in range(3): await cleanup.cancel_and_wait_for_rejoin(closing) - assert client._proc is process assert process.returncode is None assert not cleanup.finished.is_set() @@ -121,8 +133,7 @@ async def test_cancelled_prompt_cancels_its_host_callback_before_reap( client_factory: ClientFactory, ) -> None: client = await client_factory("host-cancel") - process = client._proc - assert process is not None + process = owned_rpc_process(client) observer = Observer() started = asyncio.Event() cancelled = asyncio.Event() diff --git a/tests/integrations/omprpc/test_process_tree.py b/tests/integrations/omprpc/test_process_tree.py index 718a17688..6a8b9a2b6 100644 --- a/tests/integrations/omprpc/test_process_tree.py +++ b/tests/integrations/omprpc/test_process_tree.py @@ -9,7 +9,6 @@ import pytest from pydantic import JsonValue -from band.integrations.codex.winspawn import Job from band.integrations.omprpc import OmpProcessError from band.integrations.omprpc.client import PROCESS_EXIT_TIMEOUT_S from tests.integrations.omppeer import CRASH_EXIT @@ -20,8 +19,10 @@ RecycledEnvironment, descendant_client, descendant_liveness, + owned_rpc_process, process_is_alive, - refuse_process_signal, + refuse_containment, + session_ids, unrelated_gate_child, ) @@ -41,24 +42,17 @@ async def test_close_reaps_detached_descendants_without_affecting_other_owners( await process_gate.connection(), await process_gate.connection(), ] - root = client._proc - assert root is not None - if os.name != "nt": - assert all( - os.getsid(child.process.pid) == child.process.pid - for child in descendants - ) - assert all( - os.getsid(child.process.pid) != root.pid - for child in descendants - ) + root = owned_rpc_process(client) + ids = session_ids([child.process.pid for child in descendants]) + if ids is not None: + assert ids == [child.process.pid for child in descendants] + assert all(session != root.pid for session in ids) await other.start() other_descendants = [ await process_gate.connection(), await process_gate.connection(), ] - other_root = other._proc - assert other_root is not None + other_root = owned_rpc_process(other) if root_exit == "crash": crashing = asyncio.create_task(client.request("crash")) async with asyncio.timeout(5): @@ -94,8 +88,7 @@ async def test_stale_matching_environment_cannot_adopt_a_reused_process_identity await process_gate.connection(), await process_gate.connection(), ] - root = client._proc - assert root is not None + root = owned_rpc_process(client) boundary = RecycledEnvironment(unrelated, root.pid) with monkeypatch.context() as patch: boundary.install(patch) @@ -117,8 +110,7 @@ async def test_close_reaps_owned_descendants_after_the_root_has_crashed( await process_gate.connection(), await process_gate.connection(), ] - root = client._proc - assert root is not None + root = owned_rpc_process(client) assert await descendant_liveness(descendants) == [True, True] crashing = asyncio.create_task(client.request("crash")) @@ -143,17 +135,13 @@ async def test_failed_containment_cleanup_is_retried_before_restart( crashing: asyncio.Task[dict[str, JsonValue]] | None = None try: await client.start() - old_root = client._proc - assert old_root is not None + old_root = owned_rpc_process(client) descendants = [ await process_gate.connection(), await process_gate.connection(), ] with monkeypatch.context() as patch: - if os.name == "nt": - patch.setattr(Job, "terminate", refuse_process_signal) - else: - patch.setattr(os, "killpg", refuse_process_signal) + refuse_containment(patch) crashing = asyncio.create_task(client.request("crash")) async with asyncio.timeout(5): await old_root.wait() @@ -162,7 +150,7 @@ async def test_failed_containment_cleanup_is_retried_before_restart( assert await descendant_liveness(descendants) == [True, True] await client.start() - assert client._proc is not old_root + assert owned_rpc_process(client) is not old_root assert await descendant_liveness(descendants) == [False, False] await process_gate.connection() await process_gate.connection() @@ -186,8 +174,7 @@ async def test_unsettled_session_members_keep_close_retryable( crashing: asyncio.Task[dict[str, JsonValue]] | None = None try: await client.start() - old_root = client._proc - assert old_root is not None + old_root = owned_rpc_process(client) descendants = [ await process_gate.connection(), await process_gate.connection(), @@ -208,7 +195,7 @@ async def test_unsettled_session_members_keep_close_retryable( client.process_exit_timeout_s = PROCESS_EXIT_TIMEOUT_S await client.start() - assert client._proc is not old_root + assert owned_rpc_process(client) is not old_root assert await descendant_liveness(descendants) == [False, False] await process_gate.connection() await process_gate.connection() @@ -223,7 +210,6 @@ async def test_unsettled_session_members_keep_close_retryable( async def test_cancelled_cleanup_reaps_descendants_but_not_unrelated_processes( self, process_gate: ProcessGate, - monkeypatch: pytest.MonkeyPatch, stdio: str, interruption: str, ) -> None: @@ -234,7 +220,7 @@ async def test_cancelled_cleanup_reaps_descendants_but_not_unrelated_processes( ), stdio=stdio, ) - cleanup = HeldCleanup(client, monkeypatch) + cleanup = HeldCleanup(client) starting: asyncio.Task[None] | None = None closing: asyncio.Task[None] | None = None async with unrelated_gate_child(process_gate) as unrelated: @@ -248,7 +234,7 @@ async def test_cancelled_cleanup_reaps_descendants_but_not_unrelated_processes( except TimeoutError: if starting.done(): await starting - process = client._proc + process = getattr(client, "_proc", None) pytest.fail( "Peer did not connect its descendants: " f"startup_done={starting.done()}, " @@ -259,7 +245,6 @@ async def test_cancelled_cleanup_reaps_descendants_but_not_unrelated_processes( assert spawning is not None async with asyncio.timeout(5): root = await asyncio.shield(spawning) - assert client._proc is root root_process = await asyncio.to_thread(psutil.Process, root.pid) ancestors = [ await asyncio.to_thread(process.process.parents) @@ -270,9 +255,6 @@ async def test_cancelled_cleanup_reaps_descendants_but_not_unrelated_processes( descendants[0].process in ancestors[1] or descendants[1].process in ancestors[0] ) - scope = client._process_scope - owner = client._windows_process - assert scope is not None or owner is not None if interruption == "startup": closing = starting assert closing.cancel() @@ -286,15 +268,11 @@ async def test_cancelled_cleanup_reaps_descendants_but_not_unrelated_processes( for _ in range(3): await cleanup.cancel_and_wait_for_rejoin(closing) - assert client._proc is root - assert client._process_scope is scope - assert client._windows_process is owner assert root.returncode is None assert await asyncio.to_thread(process_is_alive, root_process) assert await descendant_liveness(descendants) == [True, True] assert await asyncio.to_thread(process_is_alive, unrelated.process) assert not cleanup.finished.is_set() - assert closing.cancelling() == 4 cleanup.release.set() async with asyncio.timeout(10): diff --git a/tests/integrations/omprpc/test_prompt_results.py b/tests/integrations/omprpc/test_prompt_results.py index 2b202db65..cdf0c5970 100644 --- a/tests/integrations/omprpc/test_prompt_results.py +++ b/tests/integrations/omprpc/test_prompt_results.py @@ -17,6 +17,7 @@ ClientFactory, Observer, checkpoint_frames, + owned_rpc_process, prompt_task, second_prompt, ) @@ -31,7 +32,7 @@ async def test_prompt_delivers_typed_narration_and_ignores_other_events( observer = Observer() async with asyncio.timeout(5): await prompt_task(client, observer) - assert [event.type for event in observer.events] == [ + assert observer.outline == [ OmpEventType.AGENT_START, OmpEventType.TOOL_START, OmpEventType.TOOL_END, @@ -103,8 +104,7 @@ async def test_failed_initial_result_is_reported( self, client_factory: ClientFactory, scenario: str, match: str ) -> None: client = await client_factory(scenario) - process = client._proc - assert process is not None + process = owned_rpc_process(client) async with asyncio.timeout(5): with pytest.raises(OmpPromptError, match=match): await prompt_task(client, Observer()) @@ -116,8 +116,7 @@ async def test_narration_failure_does_not_leave_the_turn_or_process_waiting( client_factory: ClientFactory, ) -> None: client = await client_factory("basic") - process = client._proc - assert process is not None + process = owned_rpc_process(client) observer = Observer() async def on_event(event: OmpEvent) -> None: diff --git a/tests/integrations/omprpc/test_turn_settlement.py b/tests/integrations/omprpc/test_turn_settlement.py index 64c17d21b..4fc9f8968 100644 --- a/tests/integrations/omprpc/test_turn_settlement.py +++ b/tests/integrations/omprpc/test_turn_settlement.py @@ -132,3 +132,36 @@ async def on_tool(_call: OmpHostToolCall) -> OmpToolResult: assert process.returncode is not None await client.start() await second_prompt(client) + + async def test_settlement_succeeds_after_the_running_callback_finishes( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("settled-during-callback") + started = asyncio.Event() + release = asyncio.Event() + + async def on_tool(_call: OmpHostToolCall) -> OmpToolResult: + started.set() + await release.wait() + return OmpToolResult(content=[]) + + observer = Observer() + running = asyncio.create_task( + client.prompt( + "first", + on_event=observer.on_event, + on_tool=on_tool, + on_ui=observer.on_ui, + ) + ) + try: + await started.wait() + release.set() + async with asyncio.timeout(5): + await running + assert owned_rpc_process(client).returncode is None + await second_prompt(client) + finally: + running.cancel() + release.set() + await asyncio.gather(running, return_exceptions=True) From d9609a098f1f25368dc14f07f9b998978b175bc1 Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 18:44:06 +0300 Subject: [PATCH 21/35] test: cover complementary native RPC branches without private polls Group D kept finding leftover OS skips, in-test ifs, and private seams as a family. Drive those scenarios from public hooks, OS tables, and shared peer constants so each branch is a named claim. Co-authored-by: Cursor --- src/band/integrations/omprpc/client.py | 7 + tests/converters/test_omp.py | 2 +- tests/integrations/omppeer.py | 22 +- .../ompprocess/test_mcp_plugin.py | 22 +- tests/integrations/omprpc/kit.py | 114 ++++++- tests/integrations/omprpc/lifecycle.py | 29 +- tests/integrations/omprpc/test_callbacks.py | 19 +- .../omprpc/test_eof_descendants.py | 102 +++---- .../integrations/omprpc/test_process_tree.py | 277 +++++++++--------- .../omprpc/test_prompt_results.py | 3 +- tests/integrations/omprpc/test_startup.py | 23 +- .../omprpc/test_transport_failures.py | 16 +- tests/ostable.py | 12 + 13 files changed, 396 insertions(+), 252 deletions(-) create mode 100644 tests/ostable.py diff --git a/src/band/integrations/omprpc/client.py b/src/band/integrations/omprpc/client.py index 076c65708..2fa853105 100644 --- a/src/band/integrations/omprpc/client.py +++ b/src/band/integrations/omprpc/client.py @@ -67,6 +67,11 @@ logger = logging.getLogger(__name__) + +def _noop() -> None: + return None + + _DRAIN_BYTES = 64 * 1024 _ABORT_TIMEOUT_S = 1.0 PROCESS_EXIT_TIMEOUT_S = 1.0 @@ -106,6 +111,7 @@ def __init__( self.create_process = create_process self.close_process: CloseProcess = self._close_process self.join_close: JoinClose = asyncio.shield + self.after_stdin_closed: Callable[[], None] = _noop self.callback_exit_timeout_s = CALLBACK_EXIT_TIMEOUT_S self.process_exit_timeout_s = PROCESS_EXIT_TIMEOUT_S self._proc: asyncio.subprocess.Process | None = None @@ -798,6 +804,7 @@ async def _close_process(self) -> None: async with self._write_lock: await close_stdin(process) await await_returncode(process, self.process_exit_timeout_s) + await asyncio.to_thread(self.after_stdin_closed) if owner is not None: if not owner_closed: await owner.close(self.process_exit_timeout_s) diff --git a/tests/converters/test_omp.py b/tests/converters/test_omp.py index eaa7f5244..51fc19ff8 100644 --- a/tests/converters/test_omp.py +++ b/tests/converters/test_omp.py @@ -50,7 +50,7 @@ def test_latest_session_file_wins_per_room(self) -> None: def test_unusable_metadata_is_ignored(self) -> None: state = OmpHistoryConverter().convert( [ - {"message_type": "task", "metadata": "not a mapping"}, + {"message_type": MessageType.TASK, "metadata": "not a mapping"}, { "message_type": MessageType.TASK, "metadata": {OMP_ROOM_METADATA_KEY: "room-1"}, diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index bbeacce95..c0d04cb33 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -33,6 +33,10 @@ CRASH_EXIT = 3 UI_APPROVE = "Approve" UI_DENY = "Deny" +HOST_CANCELLED_ID = "host-cancelled" +UI_CANCELLED_ID = "ui-cancelled" +OVERSIZE_AFTER_READY = "oversize-after-ready" +READY_LIMIT_BYTES = 1024 def burst_message_id(number: int) -> str: @@ -469,7 +473,7 @@ def prompt(self, frame: Frame) -> None: send( { "type": "host_tool_call", - "id": "host-cancelled", + "id": HOST_CANCELLED_ID, "toolCallId": "call-cancelled", "toolName": "band_chat", "arguments": {}, @@ -479,7 +483,7 @@ def prompt(self, frame: Frame) -> None: send( { "type": "extension_ui_request", - "id": "ui-cancelled", + "id": UI_CANCELLED_ID, "method": "select", "title": "Allow command?", "options": [UI_APPROVE, UI_DENY], @@ -623,7 +627,7 @@ def handle(self, frame: Frame) -> None: { "type": "host_tool_cancel", "id": "cancel-frame", - "targetId": "host-cancelled", + "targetId": HOST_CANCELLED_ID, } ) self.end("host cancelled") @@ -636,7 +640,7 @@ def handle(self, frame: Frame) -> None: "type": "extension_ui_request", "id": "ui-cancel-frame", "method": "cancel", - "targetId": "ui-cancelled", + "targetId": UI_CANCELLED_ID, } ) self.end("ui cancelled") @@ -790,6 +794,11 @@ def main() -> None: signal.signal(signal.SIGTERM, signal.SIG_IGN) peer = Peer(args.scenario) if args.scenario not in {"no-ready", "no-ready-descendants"}: + limit = ( + READY_LIMIT_BYTES + if args.scenario == OVERSIZE_AFTER_READY + else OMP_FRAME_BYTES + ) send( { "type": "ready", @@ -797,10 +806,13 @@ def main() -> None: "supportedProtocolVersions": [1] if args.scenario == "no-v2" else [1, OMP_NATIVE_PROTOCOL_VERSION], - "maxFrameBytes": OMP_FRAME_BYTES, + "maxFrameBytes": limit, "maxReassembledFrameBytes": OMP_REASSEMBLED_BYTES, } ) + if args.scenario == OVERSIZE_AFTER_READY: + sys.stdout.write("{" + ("x" * (limit + 8)) + "}\n") + sys.stdout.flush() for line in sys.stdin: peer.handle(Frame.model_validate_json(line, strict=True)) if managed_child is not None and managed_child.stdin is not None: diff --git a/tests/integrations/ompprocess/test_mcp_plugin.py b/tests/integrations/ompprocess/test_mcp_plugin.py index f9de08028..3aa2c2c99 100644 --- a/tests/integrations/ompprocess/test_mcp_plugin.py +++ b/tests/integrations/ompprocess/test_mcp_plugin.py @@ -21,6 +21,7 @@ omp_mcp_plugin, omp_mcp_ready_command, ) +from tests.ostable import skip_unless_os def plugin_privacy(root: Path, files: tuple[Path, ...]) -> tuple[int, int | None]: @@ -91,8 +92,7 @@ def test_arguments_load_the_plugin_and_its_startup_extension( def test_projected_plugin_files_are_owner_only( self, tmp_path: Path, accepted_on: tuple[str, ...] ) -> None: - if os.name not in accepted_on: - pytest.skip("Unix permission bits are not the Windows ACL boundary") + skip_unless_os(accepted_on) servers = {"local": OmpMcpServer(command="./server")} with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: assert plugin is not None @@ -106,6 +106,24 @@ def test_projected_plugin_files_are_owner_only( assert all(path.is_file() for path in files) assert plugin_privacy(root, files) == (0o700, 0o600) + @pytest.mark.parametrize("accepted_on", [("nt",)]) + def test_projected_plugin_files_are_readable( + self, tmp_path: Path, accepted_on: tuple[str, ...] + ) -> None: + skip_unless_os(accepted_on) + servers = {"local": OmpMcpServer(command="./server")} + with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: + assert plugin is not None + root = Path(plugin) + files = ( + root / OMP_MCP_SERVERS_CONFIG, + root / OMP_MCP_STARTUP_CONFIG, + root / OMP_MCP_STARTUP_JS, + ) + assert root.is_dir() + assert all(path.is_file() for path in files) + assert all(os.access(path, os.R_OK) for path in files) + @pytest.mark.parametrize("prefix", RELATIVE_MCP_COMMAND_PREFIXES) def test_relative_command_is_absolutized(self, tmp_path: Path, prefix: str) -> None: servers = {"local": OmpMcpServer(command=f"{prefix}server")} diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py index 4baa1f807..e9daba5cd 100644 --- a/tests/integrations/omprpc/kit.py +++ b/tests/integrations/omprpc/kit.py @@ -6,6 +6,7 @@ import os import signal import sys +import threading import time from collections.abc import AsyncIterator, Awaitable, Callable, Sequence from contextlib import asynccontextmanager, suppress @@ -17,7 +18,7 @@ from pydantic import JsonValue from band.integrations.codex.winspawn import Job -from band.integrations.omprpc import OmpRpcClient +from band.integrations.omprpc import OmpProcessError, OmpRpcClient from band.integrations.ompwire import ( OmpContent, OmpEvent, @@ -29,6 +30,7 @@ ) from tests.integrations import omppeer from tests.integrations.omppeer import COMPLETE, FOLLOW_UP_PROMPT, ProcessReceipt +from tests.integrations.omprpc.cancellation import HeldCleanup ClientFactory = Callable[[str], Awaitable[OmpRpcClient]] @@ -36,6 +38,28 @@ OMP_PEER_SCRIPT = omppeer.__file__ +class SpawnBarrier: + """Create the first subprocess, then hold returning it until released.""" + + def __init__( + self, spawn: Callable[..., Awaitable[asyncio.subprocess.Process]] + ) -> None: + self.spawn = spawn + self.created = asyncio.Event() + self.release = asyncio.Event() + self.process: asyncio.subprocess.Process | None = None + + async def create( + self, *command: str, **kwargs: object + ) -> asyncio.subprocess.Process: + process = await self.spawn(*command, **kwargs) + if self.process is None: + self.process = process + self.created.set() + await self.release.wait() + return process + + class Observer: def __init__(self) -> None: self.events: list[OmpEvent] = [] @@ -258,13 +282,95 @@ def refuse_containment(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(os, "killpg", refuse_process_signal) -def session_ids(pids: Sequence[int]) -> list[int] | None: +def session_ids(pids: Sequence[int]) -> list[int]: getsid = getattr(os, "getsid", None) - if getsid is None: - return None + assert getsid is not None return [getsid(pid) for pid in pids] +def job_member_pids(client: OmpRpcClient) -> set[int]: + owner = client._windows_process + assert owner is not None and owner.job is not None + return set(owner.job._process_ids()) + + +async def wait_for_descendants(gate: ProcessGate, count: int = 2) -> list[GatedProcess]: + return [await gate.connection() for _ in range(count)] + + +async def crash_rpc_root(client: OmpRpcClient) -> None: + root = owned_rpc_process(client) + crashing = asyncio.create_task(client.request("crash")) + async with asyncio.timeout(5): + await root.wait() + assert root.returncode == omppeer.CRASH_EXIT + with pytest.raises(OmpProcessError): + await crashing + + +@dataclass +class DescendantTree: + client: OmpRpcClient + cleanup: HeldCleanup + starting: asyncio.Task[None] + root: asyncio.subprocess.Process + descendants: list[GatedProcess] + + +@asynccontextmanager +async def hold_after_stdin_closed( + client: OmpRpcClient, +) -> AsyncIterator[asyncio.Task[None]]: + entered = asyncio.Event() + release = threading.Event() + loop = asyncio.get_running_loop() + + def hold() -> None: + loop.call_soon_threadsafe(entered.set) + if not release.wait(timeout=10): + raise TimeoutError("Test did not release after_stdin_closed") + + client.after_stdin_closed = hold + closing = asyncio.create_task(client.close()) + try: + async with asyncio.timeout(10): + await entered.wait() + yield closing + if not closing.done(): + release.set() + await closing + finally: + release.set() + await asyncio.gather(closing, return_exceptions=True) + + +async def start_eof_client( + gate: ProcessGate, *, stdio: str +) -> tuple[OmpRpcClient, GatedProcess, asyncio.subprocess.Process]: + client = descendant_client( + gate, scenario="eof-descendants", stdio=stdio, detached=True + ) + await client.start() + helper = await gate.connection() + return client, helper, owned_rpc_process(client) + + +async def connect_descendant_tree( + gate: ProcessGate, *, scenario: str, stdio: str +) -> DescendantTree: + client = descendant_client(gate, scenario=scenario, stdio=stdio) + cleanup = HeldCleanup(client) + starting = asyncio.create_task(client.start()) + descendants = await wait_for_descendants(gate) + return DescendantTree( + client=client, + cleanup=cleanup, + starting=starting, + root=owned_rpc_process(client), + descendants=descendants, + ) + + class DeferredSessionSignal: """Keep real members alive while post-signal ownership lookup disappears.""" diff --git a/tests/integrations/omprpc/lifecycle.py b/tests/integrations/omprpc/lifecycle.py index be8d9061c..79356c098 100644 --- a/tests/integrations/omprpc/lifecycle.py +++ b/tests/integrations/omprpc/lifecycle.py @@ -10,7 +10,7 @@ from __future__ import annotations import asyncio -from collections.abc import AsyncIterator, Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable from contextlib import asynccontextmanager, suppress from dataclasses import dataclass from enum import StrEnum @@ -24,6 +24,7 @@ GatedProcess, Observer, ProcessGate, + SpawnBarrier, descendant_client, descendant_liveness, owned_rpc_process, @@ -53,32 +54,6 @@ class Step(StrEnum): sequences = st.lists(st.sampled_from(Step), min_size=1, max_size=6) -class SpawnBarrier: - """Create the first subprocess, then hold returning it until released. - - Descendants can still check in at the process gate; ``start()`` / ``close()`` - wait, so a host cancel can overlap spawn. - """ - - def __init__( - self, spawn: Callable[..., Awaitable[asyncio.subprocess.Process]] - ) -> None: - self.spawn = spawn - self.created = asyncio.Event() - self.release = asyncio.Event() - self.process: asyncio.subprocess.Process | None = None - - async def create( - self, *command: str, **kwargs: object - ) -> asyncio.subprocess.Process: - process = await self.spawn(*command, **kwargs) - if self.process is None: - self.process = process - self.created.set() - await self.release.wait() - return process - - @dataclass class HeldSpawn: starting: asyncio.Task[None] diff --git a/tests/integrations/omprpc/test_callbacks.py b/tests/integrations/omprpc/test_callbacks.py index 039926534..fd0304bf6 100644 --- a/tests/integrations/omprpc/test_callbacks.py +++ b/tests/integrations/omprpc/test_callbacks.py @@ -14,7 +14,12 @@ OmpUIRequest, OmpUIResponse, ) -from tests.integrations.omppeer import UI_APPROVE, UI_DENY +from tests.integrations.omppeer import ( + HOST_CANCELLED_ID, + UI_APPROVE, + UI_CANCELLED_ID, + UI_DENY, +) from tests.integrations.omprpc.kit import ( ClientFactory, Observer, @@ -84,7 +89,7 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: assert ui_response == { "type": OmpFrameType.UI_RESPONSE, "id": "approval-1", - "value": "Approve", + "value": UI_APPROVE, } await second_prompt(client) @@ -112,9 +117,12 @@ async def on_ui(_request: OmpUIRequest) -> OmpUIResponse: "confirmed": True, } - @pytest.mark.parametrize("kind", ["host", "ui"]) + @pytest.mark.parametrize( + ("kind", "target"), + [("host", HOST_CANCELLED_ID), ("ui", UI_CANCELLED_ID)], + ) async def test_native_cancel_targets_the_callback_and_keeps_the_next_turn_usable( - self, client_factory: ClientFactory, kind: str + self, client_factory: ClientFactory, kind: str, target: str ) -> None: client = await client_factory(f"{kind}-cancel") observer = Observer() @@ -150,7 +158,6 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: await cancelled.wait() await running frames = await checkpoint_frames(client) - target = "host-cancelled" if kind == "host" else "ui-cancelled" assert not any(frame.get("id") == target for frame in frames) await second_prompt(client) finally: @@ -177,6 +184,8 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: on_ui=on_ui, ) assert process.returncode is not None + await client.start() + await second_prompt(client) async def test_tool_callback_exception_propagates_unchanged_and_reaps_the_process( self, diff --git a/tests/integrations/omprpc/test_eof_descendants.py b/tests/integrations/omprpc/test_eof_descendants.py index 8b7a5f4c2..c8f6dcae5 100644 --- a/tests/integrations/omprpc/test_eof_descendants.py +++ b/tests/integrations/omprpc/test_eof_descendants.py @@ -3,19 +3,20 @@ from __future__ import annotations import asyncio -import threading import pytest from tests.integrations.omprpc.kit import ( ProcessGate, - descendant_client, descendant_liveness, - owned_rpc_process, + hold_after_stdin_closed, + job_member_pids, process_is_alive, session_ids, + start_eof_client, unrelated_gate_child, ) +from tests.ostable import skip_unless_os class TestEofDescendants: @@ -23,64 +24,53 @@ class TestEofDescendants: async def test_close_reaps_a_detached_child_created_after_stdin_eof( self, process_gate: ProcessGate, - monkeypatch: pytest.MonkeyPatch, stdio: str, ) -> None: - client = descendant_client( - process_gate, scenario="eof-descendants", stdio=stdio, detached=True - ) - entered = asyncio.Event() - release = threading.Event() - loop = asyncio.get_running_loop() - closing: asyncio.Task[None] | None = None - + client, helper, root = await start_eof_client(process_gate, stdio=stdio) async with unrelated_gate_child(process_gate) as unrelated: try: - await client.start() - helper = await process_gate.connection() - root = owned_rpc_process(client) - scope = client._process_scope - assert scope is not None - signal = scope.signal - - def hold_signal() -> None: - loop.call_soon_threadsafe(entered.set) - if not release.wait(timeout=10): - raise TimeoutError( - "Test did not release the containment reaper" - ) - signal() - - with monkeypatch.context() as patch: - patch.setattr(scope, "signal", hold_signal) - closing = asyncio.create_task(client.close()) - try: - async with asyncio.timeout(10): - await entered.wait() - late_child = await process_gate.connection() - assert root.returncode is not None - assert late_child.receipt.parent_pid == helper.process.pid - ids = session_ids([late_child.process.pid]) - if ids is not None: - assert ids == [late_child.process.pid] - assert await descendant_liveness([helper, late_child]) == [ - True, - True, - ] - assert await asyncio.to_thread( - process_is_alive, unrelated.process - ) - assert not closing.done() - release.set() - await closing - finally: - release.set() - await asyncio.gather(closing, return_exceptions=True) - + async with hold_after_stdin_closed(client) as closing: + late_child = await process_gate.connection() + assert root.returncode is not None + assert late_child.receipt.parent_pid == helper.process.pid + assert await descendant_liveness([helper, late_child]) == [ + True, + True, + ] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + assert not closing.done() assert await descendant_liveness([helper, late_child]) == [False, False] assert await asyncio.to_thread(process_is_alive, unrelated.process) finally: - release.set() - if closing is not None: - await asyncio.gather(closing, return_exceptions=True) await client.close() + + @pytest.mark.parametrize("accepted_on", [("posix",)]) + async def test_eof_late_child_is_its_own_session( + self, process_gate: ProcessGate, accepted_on: tuple[str, ...] + ) -> None: + skip_unless_os(accepted_on) + client, helper, root = await start_eof_client(process_gate, stdio="discard") + try: + async with hold_after_stdin_closed(client): + late_child = await process_gate.connection() + ids = session_ids([late_child.process.pid]) + assert ids == [late_child.process.pid] + assert ids[0] != root.pid + assert helper.process.pid != late_child.process.pid + finally: + await client.close() + + @pytest.mark.parametrize("accepted_on", [("nt",)]) + async def test_eof_late_child_is_a_job_member( + self, process_gate: ProcessGate, accepted_on: tuple[str, ...] + ) -> None: + skip_unless_os(accepted_on) + client, helper, _root = await start_eof_client(process_gate, stdio="discard") + try: + async with hold_after_stdin_closed(client): + late_child = await process_gate.connection() + members = job_member_pids(client) + assert helper.process.pid in members + assert late_child.process.pid in members + finally: + await client.close() diff --git a/tests/integrations/omprpc/test_process_tree.py b/tests/integrations/omprpc/test_process_tree.py index 6a8b9a2b6..e137ae072 100644 --- a/tests/integrations/omprpc/test_process_tree.py +++ b/tests/integrations/omprpc/test_process_tree.py @@ -3,63 +3,45 @@ from __future__ import annotations import asyncio -import os import psutil import pytest -from pydantic import JsonValue from band.integrations.omprpc import OmpProcessError from band.integrations.omprpc.client import PROCESS_EXIT_TIMEOUT_S -from tests.integrations.omppeer import CRASH_EXIT -from tests.integrations.omprpc.cancellation import HeldCleanup from tests.integrations.omprpc.kit import ( DeferredSessionSignal, ProcessGate, RecycledEnvironment, + connect_descendant_tree, + crash_rpc_root, descendant_client, descendant_liveness, + job_member_pids, owned_rpc_process, process_is_alive, refuse_containment, session_ids, unrelated_gate_child, + wait_for_descendants, ) +from tests.ostable import skip_unless_os class TestProcessTreeCleanup: - @pytest.mark.parametrize("root_exit", ["close", "crash"]) async def test_close_reaps_detached_descendants_without_affecting_other_owners( - self, process_gate: ProcessGate, root_exit: str + self, process_gate: ProcessGate ) -> None: client = descendant_client(process_gate, detached=True) other = descendant_client(process_gate, detached=True) - crashing: asyncio.Task[dict[str, JsonValue]] | None = None async with unrelated_gate_child(process_gate) as unrelated: try: await client.start() - descendants = [ - await process_gate.connection(), - await process_gate.connection(), - ] + descendants = await wait_for_descendants(process_gate) root = owned_rpc_process(client) - ids = session_ids([child.process.pid for child in descendants]) - if ids is not None: - assert ids == [child.process.pid for child in descendants] - assert all(session != root.pid for session in ids) await other.start() - other_descendants = [ - await process_gate.connection(), - await process_gate.connection(), - ] + other_descendants = await wait_for_descendants(process_gate) other_root = owned_rpc_process(other) - if root_exit == "crash": - crashing = asyncio.create_task(client.request("crash")) - async with asyncio.timeout(5): - await root.wait() - assert root.returncode == CRASH_EXIT - with pytest.raises(OmpProcessError): - await crashing await client.close() assert root.returncode is not None @@ -74,8 +56,64 @@ async def test_close_reaps_detached_descendants_without_affecting_other_owners( finally: await client.close() await other.close() - if crashing is not None: - await asyncio.gather(crashing, return_exceptions=True) + + async def test_crash_then_close_reaps_detached_descendants_without_affecting_other_owners( + self, process_gate: ProcessGate + ) -> None: + client = descendant_client(process_gate, detached=True) + other = descendant_client(process_gate, detached=True) + async with unrelated_gate_child(process_gate) as unrelated: + try: + await client.start() + descendants = await wait_for_descendants(process_gate) + root = owned_rpc_process(client) + await other.start() + other_descendants = await wait_for_descendants(process_gate) + other_root = owned_rpc_process(other) + await crash_rpc_root(client) + + await client.close() + assert root.returncode is not None + assert await descendant_liveness(descendants) == [False, False] + assert other_root.returncode is None + assert await descendant_liveness(other_descendants) == [True, True] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + await other.close() + assert await descendant_liveness(other_descendants) == [False, False] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + finally: + await client.close() + await other.close() + + @pytest.mark.parametrize("accepted_on", [("posix",)]) + async def test_detached_descendants_are_their_own_sessions( + self, process_gate: ProcessGate, accepted_on: tuple[str, ...] + ) -> None: + skip_unless_os(accepted_on) + client = descendant_client(process_gate, detached=True) + try: + await client.start() + descendants = await wait_for_descendants(process_gate) + root = owned_rpc_process(client) + ids = session_ids([child.process.pid for child in descendants]) + assert ids == [child.process.pid for child in descendants] + assert all(session != root.pid for session in ids) + finally: + await client.close() + + @pytest.mark.parametrize("accepted_on", [("nt",)]) + async def test_detached_descendants_are_job_members( + self, process_gate: ProcessGate, accepted_on: tuple[str, ...] + ) -> None: + skip_unless_os(accepted_on) + client = descendant_client(process_gate, detached=True) + try: + await client.start() + descendants = await wait_for_descendants(process_gate) + members = job_member_pids(client) + assert {child.process.pid for child in descendants} <= members + finally: + await client.close() async def test_stale_matching_environment_cannot_adopt_a_reused_process_identity( self, process_gate: ProcessGate, monkeypatch: pytest.MonkeyPatch @@ -84,10 +122,7 @@ async def test_stale_matching_environment_cannot_adopt_a_reused_process_identity async with unrelated_gate_child(process_gate) as unrelated: try: await client.start() - descendants = [ - await process_gate.connection(), - await process_gate.connection(), - ] + descendants = await wait_for_descendants(process_gate) root = owned_rpc_process(client) boundary = RecycledEnvironment(unrelated, root.pid) with monkeypatch.context() as patch: @@ -103,48 +138,27 @@ async def test_close_reaps_owned_descendants_after_the_root_has_crashed( self, process_gate: ProcessGate ) -> None: client = descendant_client(process_gate) - crashing: asyncio.Task[dict[str, JsonValue]] | None = None try: await client.start() - descendants = [ - await process_gate.connection(), - await process_gate.connection(), - ] - root = owned_rpc_process(client) + descendants = await wait_for_descendants(process_gate) assert await descendant_liveness(descendants) == [True, True] - - crashing = asyncio.create_task(client.request("crash")) - async with asyncio.timeout(5): - await root.wait() - assert root.returncode == CRASH_EXIT - with pytest.raises(OmpProcessError): - await crashing + await crash_rpc_root(client) await client.close() - assert await descendant_liveness(descendants) == [False, False] - await client.close() finally: await client.close() - if crashing is not None: - await asyncio.gather(crashing, return_exceptions=True) async def test_failed_containment_cleanup_is_retried_before_restart( self, process_gate: ProcessGate, monkeypatch: pytest.MonkeyPatch ) -> None: client = descendant_client(process_gate) - crashing: asyncio.Task[dict[str, JsonValue]] | None = None try: await client.start() old_root = owned_rpc_process(client) - descendants = [ - await process_gate.connection(), - await process_gate.connection(), - ] + descendants = await wait_for_descendants(process_gate) with monkeypatch.context() as patch: refuse_containment(patch) - crashing = asyncio.create_task(client.request("crash")) - async with asyncio.timeout(5): - await old_root.wait() + await crash_rpc_root(client) with pytest.raises(PermissionError, match="not signalable"): await client.close() assert await descendant_liveness(descendants) == [True, True] @@ -152,15 +166,12 @@ async def test_failed_containment_cleanup_is_retried_before_restart( await client.start() assert owned_rpc_process(client) is not old_root assert await descendant_liveness(descendants) == [False, False] - await process_gate.connection() - await process_gate.connection() + await wait_for_descendants(process_gate) await client.close() finally: await client.close() - if crashing is not None: - await asyncio.gather(crashing, return_exceptions=True) - @pytest.mark.skipif(os.name == "nt", reason="POSIX group settlement boundary") + @pytest.mark.parametrize("accepted_on", [("posix",)]) @pytest.mark.parametrize( "hide_after_signal", [False, True], ids=["visible", "hidden-after-signal"] ) @@ -168,26 +179,22 @@ async def test_unsettled_session_members_keep_close_retryable( self, process_gate: ProcessGate, monkeypatch: pytest.MonkeyPatch, + accepted_on: tuple[str, ...], hide_after_signal: bool, ) -> None: + skip_unless_os(accepted_on) client = descendant_client(process_gate) - crashing: asyncio.Task[dict[str, JsonValue]] | None = None try: await client.start() old_root = owned_rpc_process(client) - descendants = [ - await process_gate.connection(), - await process_gate.connection(), - ] + descendants = await wait_for_descendants(process_gate) boundary = DeferredSessionSignal( old_root.pid, descendants, hide_after_signal=hide_after_signal ) with monkeypatch.context() as patch: boundary.install(patch) client.process_exit_timeout_s = 0 - crashing = asyncio.create_task(client.request("crash")) - async with asyncio.timeout(5): - await old_root.wait() + await crash_rpc_root(client) for _ in range(2): with pytest.raises(OmpProcessError, match="session members"): await client.close() @@ -197,97 +204,91 @@ async def test_unsettled_session_members_keep_close_retryable( await client.start() assert owned_rpc_process(client) is not old_root assert await descendant_liveness(descendants) == [False, False] - await process_gate.connection() - await process_gate.connection() + await wait_for_descendants(process_gate) await client.close() finally: await client.close() - if crashing is not None: - await asyncio.gather(crashing, return_exceptions=True) @pytest.mark.parametrize("stdio", ["inherit", "discard"]) - @pytest.mark.parametrize("interruption", ["startup", "close"]) - async def test_cancelled_cleanup_reaps_descendants_but_not_unrelated_processes( + async def test_cancelled_startup_cleanup_reaps_descendants_but_not_unrelated_processes( + self, + process_gate: ProcessGate, + stdio: str, + ) -> None: + tree = await connect_descendant_tree( + process_gate, scenario="no-ready-descendants", stdio=stdio + ) + async with unrelated_gate_child(process_gate) as unrelated: + try: + root_process = await asyncio.to_thread(psutil.Process, tree.root.pid) + closing = tree.starting + assert closing.cancel() + await tree.cleanup.wait_for_join(closing) + for _ in range(3): + await tree.cleanup.cancel_and_wait_for_rejoin(closing) + assert tree.root.returncode is None + assert await asyncio.to_thread(process_is_alive, root_process) + assert await descendant_liveness(tree.descendants) == [True, True] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + assert not tree.cleanup.finished.is_set() + + tree.cleanup.release.set() + async with asyncio.timeout(10): + with pytest.raises(asyncio.CancelledError): + await closing + await tree.root.wait() + assert tree.cleanup.finished.is_set() + assert tree.root.returncode is not None + assert await descendant_liveness(tree.descendants) == [False, False] + assert await asyncio.to_thread(process_is_alive, unrelated.process) + await tree.client.close() + finally: + tree.cleanup.release.set() + tree.starting.cancel() + await asyncio.gather(tree.starting, return_exceptions=True) + await tree.client.close() + + @pytest.mark.parametrize("stdio", ["inherit", "discard"]) + async def test_cancelled_close_cleanup_reaps_descendants_but_not_unrelated_processes( self, process_gate: ProcessGate, stdio: str, - interruption: str, ) -> None: - client = descendant_client( - process_gate, - scenario=( - "no-ready-descendants" if interruption == "startup" else "descendants" - ), - stdio=stdio, + tree = await connect_descendant_tree( + process_gate, scenario="descendants", stdio=stdio ) - cleanup = HeldCleanup(client) - starting: asyncio.Task[None] | None = None closing: asyncio.Task[None] | None = None async with unrelated_gate_child(process_gate) as unrelated: try: - starting = asyncio.create_task(client.start()) - try: - descendants = [ - await process_gate.connection(), - await process_gate.connection(), - ] - except TimeoutError: - if starting.done(): - await starting - process = getattr(client, "_proc", None) - pytest.fail( - "Peer did not connect its descendants: " - f"startup_done={starting.done()}, " - f"root_pid={process.pid if process else None}, " - f"returncode={process.returncode if process else None}" - ) - spawning = client._spawn_task - assert spawning is not None async with asyncio.timeout(5): - root = await asyncio.shield(spawning) - root_process = await asyncio.to_thread(psutil.Process, root.pid) - ancestors = [ - await asyncio.to_thread(process.process.parents) - for process in descendants - ] - assert all(root_process in parents for parents in ancestors) - assert ( - descendants[0].process in ancestors[1] - or descendants[1].process in ancestors[0] - ) - if interruption == "startup": - closing = starting - assert closing.cancel() - else: - async with asyncio.timeout(5): - await starting - closing = asyncio.create_task(client.close()) - await cleanup.wait_for_join(closing) - if interruption == "close": - await cleanup.cancel_and_wait_for_rejoin(closing) - + await tree.starting + root_process = await asyncio.to_thread(psutil.Process, tree.root.pid) + closing = asyncio.create_task(tree.client.close()) + await tree.cleanup.wait_for_join(closing) + await tree.cleanup.cancel_and_wait_for_rejoin(closing) for _ in range(3): - await cleanup.cancel_and_wait_for_rejoin(closing) - assert root.returncode is None + await tree.cleanup.cancel_and_wait_for_rejoin(closing) + assert tree.root.returncode is None assert await asyncio.to_thread(process_is_alive, root_process) - assert await descendant_liveness(descendants) == [True, True] + assert await descendant_liveness(tree.descendants) == [True, True] assert await asyncio.to_thread(process_is_alive, unrelated.process) - assert not cleanup.finished.is_set() + assert not tree.cleanup.finished.is_set() - cleanup.release.set() + tree.cleanup.release.set() async with asyncio.timeout(10): with pytest.raises(asyncio.CancelledError): await closing - await root.wait() - assert cleanup.finished.is_set() - assert root.returncode is not None - assert await descendant_liveness(descendants) == [False, False] + await tree.root.wait() + assert tree.cleanup.finished.is_set() + assert tree.root.returncode is not None + assert await descendant_liveness(tree.descendants) == [False, False] assert await asyncio.to_thread(process_is_alive, unrelated.process) - await client.close() + await tree.client.close() finally: - cleanup.release.set() - tasks = {task for task in (starting, closing) if task is not None} + tree.cleanup.release.set() + tree.starting.cancel() + tasks = [tree.starting, *([closing] if closing is not None else [])] for task in tasks: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) - await client.close() + await tree.client.close() diff --git a/tests/integrations/omprpc/test_prompt_results.py b/tests/integrations/omprpc/test_prompt_results.py index cdf0c5970..3c50b1c1d 100644 --- a/tests/integrations/omprpc/test_prompt_results.py +++ b/tests/integrations/omprpc/test_prompt_results.py @@ -13,6 +13,7 @@ OmpEventType, OmpState, ) +from tests.integrations.omppeer import COMPLETE from tests.integrations.omprpc.kit import ( ClientFactory, Observer, @@ -58,7 +59,7 @@ async def test_prompt_delivers_typed_narration_and_ignores_other_events( for content in message.content ] == [ ("thinking", None, "Considering the result"), - ("text", "complete", None), + ("text", COMPLETE, None), ] assert message.usage is not None assert ( diff --git a/tests/integrations/omprpc/test_startup.py b/tests/integrations/omprpc/test_startup.py index 9542110dc..8db912538 100644 --- a/tests/integrations/omprpc/test_startup.py +++ b/tests/integrations/omprpc/test_startup.py @@ -14,13 +14,15 @@ OmpRpcClient, OmpRpcError, ) -from tests.integrations.omprpc.kit import OMP_PEER_SCRIPT, ClientFactory +from tests.integrations.omppeer import OVERSIZE_AFTER_READY +from tests.integrations.omprpc.kit import OMP_PEER_SCRIPT, ClientFactory, SpawnBarrier class TestStartupNegotiation: async def test_cancelled_startup_reaps_a_process_that_never_announces_ready( self, ) -> None: + barrier = SpawnBarrier(asyncio.create_subprocess_exec) client = OmpRpcClient( ( sys.executable, @@ -28,19 +30,21 @@ async def test_cancelled_startup_reaps_a_process_that_never_announces_ready( OMP_PEER_SCRIPT, "--scenario", "no-ready", - ) + ), + create_process=barrier.create, ) starting = asyncio.create_task(client.start()) try: async with asyncio.timeout(5): - while client._proc is None and not starting.done(): - await asyncio.sleep(0) - assert client._proc is not None + await barrier.created.wait() + assert barrier.process is not None assert starting.cancel() + barrier.release.set() with pytest.raises(asyncio.CancelledError): await starting - assert client._proc.returncode is not None + assert barrier.process.returncode is not None finally: + barrier.release.set() await client.close() async def test_missing_executable_is_a_process_failure(self) -> None: @@ -66,6 +70,13 @@ async def test_start_requires_successful_native_v2_negotiation( with pytest.raises(error, match="protocol"): await client_factory(scenario) + async def test_negotiated_frame_limit_applies_before_the_next_line( + self, client_factory: ClientFactory + ) -> None: + async with asyncio.timeout(6): + with pytest.raises(OmpProtocolError, match="frame length"): + await client_factory(OVERSIZE_AFTER_READY) + async def test_graceful_close_does_not_log_stdout_eof_as_a_crash( self, client_factory: ClientFactory, caplog: pytest.LogCaptureFixture ) -> None: diff --git a/tests/integrations/omprpc/test_transport_failures.py b/tests/integrations/omprpc/test_transport_failures.py index 92e0dda3c..90d787746 100644 --- a/tests/integrations/omprpc/test_transport_failures.py +++ b/tests/integrations/omprpc/test_transport_failures.py @@ -7,7 +7,12 @@ import pytest from band.integrations.omprpc import OmpProcessError, OmpProtocolError -from tests.integrations.omprpc.kit import ClientFactory, Observer, prompt_task +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + owned_rpc_process, + prompt_task, +) class TestTransportAndProtocolFailures: @@ -30,8 +35,7 @@ async def test_eof_and_protocol_failure_cannot_leave_a_prompt_waiting( match: str, ) -> None: client = await client_factory(scenario) - process = client._proc - assert process is not None + process = owned_rpc_process(client) async with asyncio.timeout(6): with pytest.raises(error, match=match): await prompt_task(client, Observer()) @@ -42,8 +46,7 @@ async def test_eof_fails_a_command_waiter_and_closes_the_transport( client_factory: ClientFactory, ) -> None: client = await client_factory("basic") - process = client._proc - assert process is not None + process = owned_rpc_process(client) async with asyncio.timeout(6): with pytest.raises(OmpProcessError, match="EOF"): await client.request("crash") @@ -74,8 +77,7 @@ async def test_invalid_or_oversized_native_chunks_fail_and_reap( self, client_factory: ClientFactory, scenario: str ) -> None: client = await client_factory(scenario) - process = client._proc - assert process is not None + process = owned_rpc_process(client) async with asyncio.timeout(6): with pytest.raises(OmpProtocolError, match="protocol"): await prompt_task(client, Observer()) diff --git a/tests/ostable.py b/tests/ostable.py new file mode 100644 index 000000000..ae5860324 --- /dev/null +++ b/tests/ostable.py @@ -0,0 +1,12 @@ +"""OS-table rows that run everywhere and skip only on the named platforms.""" + +from __future__ import annotations + +import os + +import pytest + + +def skip_unless_os(accepted_on: tuple[str, ...]) -> None: + if os.name not in accepted_on: + pytest.skip(f"this row runs on {', '.join(accepted_on)}") From e1adbd661ec3d7af3d2e87e273f8fd3b294acb24 Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 18:51:52 +0300 Subject: [PATCH 22/35] docs: record Windows job containment as closed unless close leaks Nested jobs and breakaway look like missing features on this path. They are out of scope while close() drains the existing kill-on-close job. Co-authored-by: Cursor --- src/band/integrations/codex/winspawn.py | 6 ++++++ src/band/integrations/omprpc/client.py | 4 ++++ src/band/integrations/omprpc/scope.py | 4 +++- tests/integrations/omprpc/test_process_tree.py | 1 + 4 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/band/integrations/codex/winspawn.py b/src/band/integrations/codex/winspawn.py index 08261ddf7..3f61a5687 100644 --- a/src/band/integrations/codex/winspawn.py +++ b/src/band/integrations/codex/winspawn.py @@ -3,6 +3,11 @@ The parent retains the job handle. The isolated helper creates the real command in that job before it can run, then waits. Close terminates the job so descendants die even after their launcher exits. + +This is a single kill-on-close job with no breakaway and no nested jobs. +Children inherit membership at create; they must not inherit the handle +(or kill-on-close never fires). Nested jobs or BREAKAWAY_OK belong here +only when close still leaves members alive. """ from __future__ import annotations @@ -583,6 +588,7 @@ def create_job() -> Job: _kernel32().CloseHandle(handle) _fail("CreateJobObject", _ERROR_ALREADY_EXISTS) limits = ExtendedLimitInformation() + # Breakaway flags stay unset so descendants cannot leave this job. limits.BasicLimitInformation.LimitFlags = _JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE if not _kernel32().SetInformationJobObject( handle, diff --git a/src/band/integrations/omprpc/client.py b/src/band/integrations/omprpc/client.py index 2fa853105..277c1c0f4 100644 --- a/src/band/integrations/omprpc/client.py +++ b/src/band/integrations/omprpc/client.py @@ -195,6 +195,10 @@ async def _spawn_process(self) -> asyncio.subprocess.Process: scope = self._process_scope assert scope is not None env = scope.environment(self.env) + # Windows has no process group. The job is the whole tree (membership + # at create, no breakaway). A create_process factory skips it, so that + # path is not a job test. Nested jobs belong here only if close() + # still leaves live members. if os.name == "nt" and self.create_process is None: self._windows_process = WindowsProcess() process = await self._windows_process.start_command( diff --git a/src/band/integrations/omprpc/scope.py b/src/band/integrations/omprpc/scope.py index c3502609a..2fcccf7aa 100644 --- a/src/band/integrations/omprpc/scope.py +++ b/src/band/integrations/omprpc/scope.py @@ -24,7 +24,9 @@ class OmpProcessScope: The private PI_ marker survives native helper filters. Explicit environment replacement and restricted process inspection can hide detached children. Windows job membership is the primary containment; this marker catches - descendants the job does not observe. + descendants the job does not observe. The job does not allow breakaway or + nest other jobs — a second job is not a substitute for this scan unless + close() still reports live members. """ def __init__(self) -> None: diff --git a/tests/integrations/omprpc/test_process_tree.py b/tests/integrations/omprpc/test_process_tree.py index e137ae072..0a7374495 100644 --- a/tests/integrations/omprpc/test_process_tree.py +++ b/tests/integrations/omprpc/test_process_tree.py @@ -105,6 +105,7 @@ async def test_detached_descendants_are_their_own_sessions( async def test_detached_descendants_are_job_members( self, process_gate: ProcessGate, accepted_on: tuple[str, ...] ) -> None: + """Windows identity complementary to POSIX session ids, not a nested-job claim.""" skip_unless_os(accepted_on) client = descendant_client(process_gate, detached=True) try: From 1abe091b99eff4dcbb40b410d30d7711ca7cb961 Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 19:04:32 +0300 Subject: [PATCH 23/35] fix: fail-close native RPC overflow, logs, and blocked-drain close Overflow now records the transport ERROR, protocol failures omit frame payloads from logs, and close no longer waits on a drain that holds the write lock. Co-authored-by: Cursor --- src/band/integrations/omprpc/client.py | 10 ++++--- src/band/integrations/omprpc/prompt.py | 10 ++++--- tests/integrations/omppeer.py | 3 +- .../omprpc/test_event_backpressure.py | 25 ++++++++++++----- .../integrations/omprpc/test_interruption.py | 28 +++++++++++++++++++ .../omprpc/test_transport_failures.py | 15 ++++++++++ 6 files changed, 75 insertions(+), 16 deletions(-) diff --git a/src/band/integrations/omprpc/client.py b/src/band/integrations/omprpc/client.py index 277c1c0f4..37eb3d15b 100644 --- a/src/band/integrations/omprpc/client.py +++ b/src/band/integrations/omprpc/client.py @@ -553,7 +553,8 @@ async def _dispatch_payload( case event_type if event_type in _EVENT_TYPES: event = OmpEvent.model_validate(payload) if turn is not None: - turn.observe(event, sequence) + if overflow := turn.observe(event, sequence): + self._record_failure(overflow) else: logger.debug("OMP RPC ignored an event without an active prompt") @@ -677,7 +678,7 @@ def _apply_ready_limits(self, ready: OmpReady) -> None: def _record_failure(self, error: OmpProcessError) -> None: if self._failure is None: self._failure = error - logger.error("OMP RPC transport failed: %s", error, exc_info=error) + logger.error("OMP RPC transport failed: %s", error) self._fail_pending(self._failure) if not self._closing: self._begin_close() @@ -805,8 +806,9 @@ async def _close_process(self) -> None: label="write", ) if not owner_closed: - async with self._write_lock: - await close_stdin(process) + # A drain blocked on a full pipe keeps this lock; close must + # not wait for it. Writers are already cancelled. + await close_stdin(process) await await_returncode(process, self.process_exit_timeout_s) await asyncio.to_thread(self.after_stdin_closed) if owner is not None: diff --git a/src/band/integrations/omprpc/prompt.py b/src/band/integrations/omprpc/prompt.py index dbb01756f..88ab9590f 100644 --- a/src/band/integrations/omprpc/prompt.py +++ b/src/band/integrations/omprpc/prompt.py @@ -58,15 +58,16 @@ def fail(self, error: BaseException) -> None: if not self.failure.done(): self.failure.set_result(error) - def observe(self, event: OmpEvent, sequence: int) -> None: + def observe(self, event: OmpEvent, sequence: int) -> OmpProcessError | None: if self.failure.done(): - return + return None # Waiting for space would block the reader that services reentrant commands. try: self.events.put_nowait(event) except asyncio.QueueFull: - self.fail(OmpProcessError(BACKLOG_EXCEEDED)) - return + error = OmpProcessError(BACKLOG_EXCEEDED) + self.fail(error) + return error self.last_activity_sequence = sequence match event.type: case OmpEventType.MESSAGE_END if ( @@ -83,3 +84,4 @@ def observe(self, event: OmpEvent, sequence: int) -> None: self.last_assistant, ) self.changed.set() + return None diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index c0d04cb33..ea80144a8 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -31,6 +31,7 @@ RESUMED_PROVIDER_FAILED = "Resumed provider failed" FOLLOW_UP_PROMPT = "second" CRASH_EXIT = 3 +FRAME_PAYLOAD_MARK = "omp-frame-payload-must-not-be-logged" UI_APPROVE = "Approve" UI_DENY = "Deny" HOST_CANCELLED_ID = "host-cancelled" @@ -313,7 +314,7 @@ def prompt(self, frame: Frame) -> None: send({"type": "agent_start"}) raise SystemExit(2) case "bad-json": - sys.stdout.write("{broken\n") + sys.stdout.write("{" + FRAME_PAYLOAD_MARK + "\n") sys.stdout.flush() return case "bad-event": diff --git a/tests/integrations/omprpc/test_event_backpressure.py b/tests/integrations/omprpc/test_event_backpressure.py index 57580d733..c1288c7a7 100644 --- a/tests/integrations/omprpc/test_event_backpressure.py +++ b/tests/integrations/omprpc/test_event_backpressure.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import logging from collections.abc import AsyncIterator from contextlib import asynccontextmanager @@ -67,10 +68,20 @@ async def overflow_reaps_and_restarts( running: asyncio.Task[None], process: asyncio.subprocess.Process, stopped: asyncio.Event, + caplog: pytest.LogCaptureFixture, ) -> None: - await client.request("event-burst", count=EVENT_BACKLOG_LIMIT + 1) - with pytest.raises(OmpProcessError, match=BACKLOG_EXCEEDED): - await running + caplog.clear() + with caplog.at_level(logging.ERROR, logger="band.integrations.omprpc.client"): + burst = asyncio.create_task( + client.request("event-burst", count=EVENT_BACKLOG_LIMIT + 1) + ) + with pytest.raises(OmpProcessError, match=BACKLOG_EXCEEDED): + await running + with pytest.raises(OmpProcessError, match=BACKLOG_EXCEEDED): + await burst + errors = [record for record in caplog.records if record.levelno == logging.ERROR] + assert [record.levelno for record in errors] == [logging.ERROR] + assert BACKLOG_EXCEEDED in errors[0].getMessage() assert stopped.is_set() assert process.returncode is not None await client.start() @@ -143,7 +154,7 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: await asyncio.gather(running, return_exceptions=True) async def test_overload_before_result_fails_closed_and_restarts( - self, client_factory: ClientFactory + self, client_factory: ClientFactory, caplog: pytest.LogCaptureFixture ) -> None: client = await client_factory("event-backpressure") async with ( @@ -154,10 +165,10 @@ async def test_overload_before_result_fails_closed_and_restarts( stopped, ), ): - await overflow_reaps_and_restarts(client, running, process, stopped) + await overflow_reaps_and_restarts(client, running, process, stopped, caplog) async def test_overload_after_settlement_fails_closed_and_restarts( - self, client_factory: ClientFactory + self, client_factory: ClientFactory, caplog: pytest.LogCaptureFixture ) -> None: client = await client_factory("event-backpressure") async with ( @@ -169,4 +180,4 @@ async def test_overload_after_settlement_fails_closed_and_restarts( ), ): await client.request("event-result") - await overflow_reaps_and_restarts(client, running, process, stopped) + await overflow_reaps_and_restarts(client, running, process, stopped, caplog) diff --git a/tests/integrations/omprpc/test_interruption.py b/tests/integrations/omprpc/test_interruption.py index 7031d1f18..5e99b9c53 100644 --- a/tests/integrations/omprpc/test_interruption.py +++ b/tests/integrations/omprpc/test_interruption.py @@ -166,3 +166,31 @@ async def on_tool(call: OmpHostToolCall) -> OmpToolResult: finally: running.cancel() await asyncio.gather(running, return_exceptions=True) + + async def test_close_finishes_while_a_stdin_drain_is_held( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory("basic") + process = owned_rpc_process(client) + stdin = process.stdin + assert stdin is not None + entered = asyncio.Event() + release = asyncio.Event() + drain = stdin.drain + + async def hold_drain() -> None: + entered.set() + await release.wait() + await drain() + + stdin.drain = hold_drain + writing = asyncio.create_task(client.request("checkpoint")) + try: + async with asyncio.timeout(5): + await entered.wait() + await client.close() + assert process.returncode is not None + finally: + release.set() + writing.cancel() + await asyncio.gather(writing, return_exceptions=True) diff --git a/tests/integrations/omprpc/test_transport_failures.py b/tests/integrations/omprpc/test_transport_failures.py index 90d787746..4c70fff5c 100644 --- a/tests/integrations/omprpc/test_transport_failures.py +++ b/tests/integrations/omprpc/test_transport_failures.py @@ -3,10 +3,12 @@ from __future__ import annotations import asyncio +import logging import pytest from band.integrations.omprpc import OmpProcessError, OmpProtocolError +from tests.integrations.omppeer import FRAME_PAYLOAD_MARK from tests.integrations.omprpc.kit import ( ClientFactory, Observer, @@ -41,6 +43,19 @@ async def test_eof_and_protocol_failure_cannot_leave_a_prompt_waiting( await prompt_task(client, Observer()) assert process.returncode is not None + async def test_protocol_failure_logs_omit_the_frame_payload( + self, + client_factory: ClientFactory, + caplog: pytest.LogCaptureFixture, + ) -> None: + client = await client_factory("bad-json") + caplog.clear() + with caplog.at_level(logging.ERROR, logger="band.integrations.omprpc.client"): + async with asyncio.timeout(6): + with pytest.raises(OmpProtocolError, match="protocol"): + await prompt_task(client, Observer()) + assert FRAME_PAYLOAD_MARK not in caplog.text + async def test_eof_fails_a_command_waiter_and_closes_the_transport( self, client_factory: ClientFactory, From d621cb40acc498a06a268dcd693874740b6cb998 Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 19:05:21 +0300 Subject: [PATCH 24/35] refactor: assert the backlog ERROR itself instead of a filtered tautology Co-authored-by: Cursor --- tests/integrations/omprpc/test_event_backpressure.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/integrations/omprpc/test_event_backpressure.py b/tests/integrations/omprpc/test_event_backpressure.py index c1288c7a7..014c961f0 100644 --- a/tests/integrations/omprpc/test_event_backpressure.py +++ b/tests/integrations/omprpc/test_event_backpressure.py @@ -79,9 +79,10 @@ async def overflow_reaps_and_restarts( await running with pytest.raises(OmpProcessError, match=BACKLOG_EXCEEDED): await burst - errors = [record for record in caplog.records if record.levelno == logging.ERROR] - assert [record.levelno for record in errors] == [logging.ERROR] - assert BACKLOG_EXCEEDED in errors[0].getMessage() + assert any( + record.levelno == logging.ERROR and BACKLOG_EXCEEDED in record.getMessage() + for record in caplog.records + ) assert stopped.is_set() assert process.returncode is not None await client.start() From 724389037604b847cf2251e244571f6e7d799bb8 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 16:52:46 +0300 Subject: [PATCH 25/35] fix: preserve Windows ownership for injected native process creation --- docs/acp.md | 5 ++ src/band/integrations/codex/winspawn.py | 11 ++-- src/band/integrations/omprpc/client.py | 12 ++--- .../omprpc/test_injected_spawn.py | 51 +++++++++++++++++++ 4 files changed, 69 insertions(+), 10 deletions(-) create mode 100644 tests/integrations/omprpc/test_injected_spawn.py diff --git a/docs/acp.md b/docs/acp.md index 2a095eb23..8bed2b578 100644 --- a/docs/acp.md +++ b/docs/acp.md @@ -293,3 +293,8 @@ fails the owning prompt with `OmpProcessError` and starts cleanup rather than silently dropping events and reporting success. The wire reader does not wait for callback capacity, so callbacks can still issue correlated commands and answer native UI requests while delivery is held. + +A `create_process` factory customizes creation without bypassing containment. +On Windows it launches the existing helper; the client still owns the Job +containing the native process and its descendants, including during cancelled +or incomplete startup. diff --git a/src/band/integrations/codex/winspawn.py b/src/band/integrations/codex/winspawn.py index 3f61a5687..e85b4c62f 100644 --- a/src/band/integrations/codex/winspawn.py +++ b/src/band/integrations/codex/winspawn.py @@ -22,12 +22,14 @@ import subprocess import sys import uuid -from collections.abc import Callable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from contextlib import contextmanager, suppress from ctypes import wintypes from pathlib import Path from typing import Any, NoReturn +CreateProcess = Callable[..., Awaitable[asyncio.subprocess.Process]] + logger = logging.getLogger(__name__) # CreateProcess failed because the executable path does not exist. @@ -749,10 +751,11 @@ async def _launch_helper( cwd: str | None, env: Mapping[str, str] | None, limit: int, + create_process: CreateProcess, ) -> asyncio.subprocess.Process: startup = subprocess.STARTUPINFO() # type: ignore[attr-defined] startup.lpAttributeList = {"handle_list": [status_handle]} - return await asyncio.create_subprocess_exec( + return await create_process( *command, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, @@ -767,7 +770,8 @@ async def _launch_helper( class WindowsProcess: """One client-owned helper and job, including incomplete startup.""" - def __init__(self) -> None: + def __init__(self, *, create_process: CreateProcess | None = None) -> None: + self.create_process = create_process or asyncio.create_subprocess_exec self.job: Job | None = create_job() self.process: asyncio.subprocess.Process | None = None self.closing = False @@ -837,6 +841,7 @@ async def _create( cwd=cwd, env=env, limit=limit, + create_process=self.create_process, ) finally: # Only the helper may keep the writer; its exit must produce EOF. diff --git a/src/band/integrations/omprpc/client.py b/src/band/integrations/omprpc/client.py index 37eb3d15b..27cd53623 100644 --- a/src/band/integrations/omprpc/client.py +++ b/src/band/integrations/omprpc/client.py @@ -13,6 +13,7 @@ from pydantic import JsonValue, ValidationError from band.integrations.codex.winspawn import ( + CreateProcess, SpawnFailure, WindowsProcess, await_returncode, @@ -61,7 +62,6 @@ ) TModel = TypeVar("TModel", bound=OmpWireModel) -CreateProcess = Callable[..., Awaitable[asyncio.subprocess.Process]] CloseProcess = Callable[[], Coroutine[Any, Any, None]] JoinClose = Callable[[Awaitable[None]], Awaitable[None]] @@ -195,12 +195,10 @@ async def _spawn_process(self) -> asyncio.subprocess.Process: scope = self._process_scope assert scope is not None env = scope.environment(self.env) - # Windows has no process group. The job is the whole tree (membership - # at create, no breakaway). A create_process factory skips it, so that - # path is not a job test. Nested jobs belong here only if close() - # still leaves live members. - if os.name == "nt" and self.create_process is None: - self._windows_process = WindowsProcess() + # Windows has no process group. Descendants inherit job membership at + # creation and cannot break away. + if os.name == "nt": + self._windows_process = WindowsProcess(create_process=self.create_process) process = await self._windows_process.start_command( self.command, cwd=None, env=env, limit=OMP_FRAME_BYTES ) diff --git a/tests/integrations/omprpc/test_injected_spawn.py b/tests/integrations/omprpc/test_injected_spawn.py new file mode 100644 index 000000000..98e43b454 --- /dev/null +++ b/tests/integrations/omprpc/test_injected_spawn.py @@ -0,0 +1,51 @@ +"""Injected creation preserves forced retirement and a usable replacement.""" + +from __future__ import annotations + +import asyncio +import sys + +import pytest + +from band.integrations.omprpc import OmpProcessError, OmpRpcClient +from band.integrations.ompwire import OmpEventType +from tests.integrations.omprpc.kit import ( + OMP_PEER_SCRIPT, + Observer, + owned_rpc_process, + prompt_task, + second_prompt, +) + + +class TestInjectedSpawnOwnership: + async def test_stubborn_injected_process_is_reaped_before_restart(self) -> None: + client = OmpRpcClient( + (sys.executable, "-u", OMP_PEER_SCRIPT, "--scenario", "stubborn"), + create_process=asyncio.create_subprocess_exec, + ) + process: asyncio.subprocess.Process | None = None + running: asyncio.Task[None] | None = None + try: + await client.start() + process = owned_rpc_process(client) + observer = Observer() + running = prompt_task(client, observer) + await observer.wait_for(OmpEventType.AGENT_END) + async with asyncio.timeout(10): + await client.close() + with pytest.raises(OmpProcessError): + await running + assert process.returncode is not None + await client.start() + await second_prompt(client) + assert owned_rpc_process(client) is not process + finally: + # Release an uncontained direct child even when the regression fails. + if process is not None and process.returncode is None: + process.kill() + await process.wait() + if running is not None: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + await client.close() From e893e4ced65387bd0c85629f90ef21ab9043148d Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 16:52:46 +0300 Subject: [PATCH 26/35] test: allow backlog overflow to retire before acknowledgement --- tests/integrations/omprpc/test_event_backpressure.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/integrations/omprpc/test_event_backpressure.py b/tests/integrations/omprpc/test_event_backpressure.py index 014c961f0..12df978c9 100644 --- a/tests/integrations/omprpc/test_event_backpressure.py +++ b/tests/integrations/omprpc/test_event_backpressure.py @@ -75,10 +75,14 @@ async def overflow_reaps_and_restarts( burst = asyncio.create_task( client.request("event-burst", count=EVENT_BACKLOG_LIMIT + 1) ) - with pytest.raises(OmpProcessError, match=BACKLOG_EXCEEDED): - await running - with pytest.raises(OmpProcessError, match=BACKLOG_EXCEEDED): - await burst + try: + with pytest.raises(OmpProcessError, match=BACKLOG_EXCEEDED): + await running + with pytest.raises(OmpProcessError, match=BACKLOG_EXCEEDED): + await burst + finally: + burst.cancel() + await asyncio.gather(burst, return_exceptions=True) assert any( record.levelno == logging.ERROR and BACKLOG_EXCEEDED in record.getMessage() for record in caplog.records From cbba792228c26747e17130230c6e7b1fe7a554dc Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 19:04:24 +0300 Subject: [PATCH 27/35] fix: recheck native process group identity before claiming custody --- docs/acp.md | 6 +++++ src/band/integrations/omprpc/scope.py | 7 ++--- tests/integrations/omprpc/kit.py | 26 +++++++++++++++++++ .../integrations/omprpc/test_process_tree.py | 26 +++++++++++++++++++ 4 files changed, 62 insertions(+), 3 deletions(-) diff --git a/docs/acp.md b/docs/acp.md index 8bed2b578..ea42e546e 100644 --- a/docs/acp.md +++ b/docs/acp.md @@ -298,3 +298,9 @@ A `create_process` factory customizes creation without bypassing containment. On Windows it launches the existing helper; the client still owns the Job containing the native process and its descendants, including during cancelled or incomplete startup. + +On POSIX, process-group capture rechecks a reaped leader's identity after +enumerating members, before retaining that group. A reused leader PID does not +transfer custody to an unrelated session. Platforms without process-group support +use their existing Job and environment-marker ownership instead of probing that +PID as a group identity. diff --git a/src/band/integrations/omprpc/scope.py b/src/band/integrations/omprpc/scope.py index 2fcccf7aa..4873e50ef 100644 --- a/src/band/integrations/omprpc/scope.py +++ b/src/band/integrations/omprpc/scope.py @@ -40,9 +40,9 @@ def environment(self, base: Mapping[str, str] | None) -> dict[str, str]: } def capture(self, root: asyncio.subprocess.Process) -> None: - group_visible = not session_pid_reused(root) - group_members: list[psutil.Process] = [] getpgid = getattr(os, "getpgid", None) + group_visible = getpgid is not None and not session_pid_reused(root) + group_members: list[psutil.Process] = [] for pid in psutil.pids(): # asyncio's child watcher remains the sole waiter for the root. if pid == KERNEL_IDLE_PID or pid == root.pid: @@ -77,7 +77,8 @@ def capture(self, root: asyncio.subprocess.Process) -> None: group_members.append(member) except (psutil.NoSuchProcess, psutil.AccessDenied): continue - if group_visible: + # A reaped leader's id can be reused while membership is enumerated. + if group_visible and not session_pid_reused(root): self.processes.update(group_members) @staticmethod diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py index e9daba5cd..d7ba38f8c 100644 --- a/tests/integrations/omprpc/kit.py +++ b/tests/integrations/omprpc/kit.py @@ -445,6 +445,32 @@ async def unrelated_gate_child(gate: ProcessGate) -> AsyncIterator[GatedProcess] await process.wait() +class RecycledGroup: + """Reuse a reaped leader's group id during membership enumeration.""" + + def __init__( + self, root: asyncio.subprocess.Process, unrelated: GatedProcess + ) -> None: + self.root = root + self.unrelated = unrelated + self.reused = False + + def check_leader(self, pid: int, sig: int) -> None: + assert (pid, sig) == (self.root.pid, 0) + if not self.reused: + raise ProcessLookupError + + def group(self, pid: int) -> int: + assert pid == self.unrelated.process.pid + self.reused = True + return self.root.pid + + def install(self, patch: pytest.MonkeyPatch) -> None: + patch.setattr(psutil, "pids", lambda: [self.unrelated.process.pid]) + patch.setattr(os, "kill", self.check_leader) + patch.setattr(os, "getpgid", self.group) + + class RecycledEnvironment: """Expose one stale matching environment while the process identity changes.""" diff --git a/tests/integrations/omprpc/test_process_tree.py b/tests/integrations/omprpc/test_process_tree.py index 0a7374495..e80f7d270 100644 --- a/tests/integrations/omprpc/test_process_tree.py +++ b/tests/integrations/omprpc/test_process_tree.py @@ -9,10 +9,13 @@ from band.integrations.omprpc import OmpProcessError from band.integrations.omprpc.client import PROCESS_EXIT_TIMEOUT_S +from band.integrations.omprpc.scope import OmpProcessScope from tests.integrations.omprpc.kit import ( + ClientFactory, DeferredSessionSignal, ProcessGate, RecycledEnvironment, + RecycledGroup, connect_descendant_tree, crash_rpc_root, descendant_client, @@ -29,6 +32,29 @@ class TestProcessTreeCleanup: + @pytest.mark.parametrize("accepted_on", [("posix",)]) + async def test_reused_group_during_capture_cannot_adopt_an_unrelated_process( + self, + client_factory: ClientFactory, + process_gate: ProcessGate, + monkeypatch: pytest.MonkeyPatch, + accepted_on: tuple[str, ...], + ) -> None: + skip_unless_os(accepted_on) + client = await client_factory("basic") + root = owned_rpc_process(client) + await client.close() + assert root.returncode is not None + scope = OmpProcessScope() + async with unrelated_gate_child(process_gate) as unrelated: + boundary = RecycledGroup(root, unrelated) + with monkeypatch.context() as patch: + boundary.install(patch) + scope.capture(root) + assert boundary.reused + assert unrelated.process not in scope.processes + assert await asyncio.to_thread(process_is_alive, unrelated.process) + async def test_close_reaps_detached_descendants_without_affecting_other_owners( self, process_gate: ProcessGate ) -> None: From f73c2be8f8b94213ed39c3439805c7c4062b5f77 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 19:20:57 +0300 Subject: [PATCH 28/35] test: cover immediate ready limits and platform-neutral EOF custody --- tests/integrations/omppeer.py | 46 +++++++++++++---------- tests/integrations/omprpc/test_startup.py | 23 ++++++++---- 2 files changed, 43 insertions(+), 26 deletions(-) diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index ea80144a8..afdc44fd0 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -754,6 +754,32 @@ def handle(self, frame: Frame) -> None: self.respond(frame) +def announce_ready(scenario: str) -> None: + ready: dict[str, JsonValue] = { + "type": "ready", + "protocolVersion": 1, + "supportedProtocolVersions": [1] + if scenario == "no-v2" + else [1, OMP_NATIVE_PROTOCOL_VERSION], + "maxFrameBytes": OMP_FRAME_BYTES, + "maxReassembledFrameBytes": OMP_REASSEMBLED_BYTES, + } + if scenario == OVERSIZE_AFTER_READY: + ready["maxFrameBytes"] = READY_LIMIT_BYTES + frames = ( + ready, + {"type": "session_settled", "padding": "x" * READY_LIMIT_BYTES}, + ) + # One write keeps the oversized frame buffered behind its ready limits. + sys.stdout.write( + "\n".join(json.dumps(frame, separators=(",", ":")) for frame in frames) + + "\n" + ) + sys.stdout.flush() + return + send(ready) + + def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--scenario", required=True) @@ -795,25 +821,7 @@ def main() -> None: signal.signal(signal.SIGTERM, signal.SIG_IGN) peer = Peer(args.scenario) if args.scenario not in {"no-ready", "no-ready-descendants"}: - limit = ( - READY_LIMIT_BYTES - if args.scenario == OVERSIZE_AFTER_READY - else OMP_FRAME_BYTES - ) - send( - { - "type": "ready", - "protocolVersion": 1, - "supportedProtocolVersions": [1] - if args.scenario == "no-v2" - else [1, OMP_NATIVE_PROTOCOL_VERSION], - "maxFrameBytes": limit, - "maxReassembledFrameBytes": OMP_REASSEMBLED_BYTES, - } - ) - if args.scenario == OVERSIZE_AFTER_READY: - sys.stdout.write("{" + ("x" * (limit + 8)) + "}\n") - sys.stdout.flush() + announce_ready(args.scenario) for line in sys.stdin: peer.handle(Frame.model_validate_json(line, strict=True)) if managed_child is not None and managed_child.stdin is not None: diff --git a/tests/integrations/omprpc/test_startup.py b/tests/integrations/omprpc/test_startup.py index 8db912538..6233c80be 100644 --- a/tests/integrations/omprpc/test_startup.py +++ b/tests/integrations/omprpc/test_startup.py @@ -15,7 +15,12 @@ OmpRpcError, ) from tests.integrations.omppeer import OVERSIZE_AFTER_READY -from tests.integrations.omprpc.kit import OMP_PEER_SCRIPT, ClientFactory, SpawnBarrier +from tests.integrations.omprpc.kit import ( + OMP_PEER_SCRIPT, + ClientFactory, + SpawnBarrier, + owned_rpc_process, +) class TestStartupNegotiation: @@ -70,12 +75,16 @@ async def test_start_requires_successful_native_v2_negotiation( with pytest.raises(error, match="protocol"): await client_factory(scenario) - async def test_negotiated_frame_limit_applies_before_the_next_line( - self, client_factory: ClientFactory - ) -> None: - async with asyncio.timeout(6): - with pytest.raises(OmpProtocolError, match="frame length"): - await client_factory(OVERSIZE_AFTER_READY) + async def test_negotiated_frame_limit_applies_before_the_next_line(self) -> None: + client = OmpRpcClient( + (sys.executable, "-u", OMP_PEER_SCRIPT, "--scenario", OVERSIZE_AFTER_READY) + ) + try: + with pytest.raises(OmpProtocolError): + await client.start() + assert owned_rpc_process(client).returncode is not None + finally: + await client.close() async def test_graceful_close_does_not_log_stdout_eof_as_a_crash( self, client_factory: ClientFactory, caplog: pytest.LogCaptureFixture From 8756ce8cb37a91517e3640931b2d0c4122b52678 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 21:21:12 +0300 Subject: [PATCH 29/35] fix: bound native callback retirement and preserve response identity --- CHANGELOG.md | 2 + docs/acp.md | 11 ++ src/band/converters/omp.py | 4 - src/band/integrations/ompprocess.py | 23 +++- src/band/integrations/omprpc/client.py | 73 +++++------ src/band/integrations/omprpc/frames.py | 33 +++-- src/band/integrations/ompwire.py | 12 ++ tests/converters/test_omp.py | 29 ++--- tests/integrations/omppeer.py | 66 +++++++++- tests/integrations/ompprocess/test_launch.py | 28 +--- .../ompprocess/test_mcp_plugin.py | 105 +++++++++++---- .../ompprocess/test_mcp_startup.py | 122 ++++++++---------- .../ompprocess/test_tool_names.py | 24 +++- tests/integrations/omprpc/backpressure.py | 35 +++++ .../omprpc/test_callback_retirement.py | 121 +++++++++++++++++ .../integrations/omprpc/test_interruption.py | 62 +++++---- .../omprpc/test_prompt_results.py | 43 +++++- .../omprpc/test_transport_failures.py | 53 +++++++- .../omprpc/test_turn_settlement.py | 5 +- .../integrations/omprpc/test_ui_responses.py | 58 +++++++++ tests/toolkit/timecontrol.py | 63 ++++++++- 21 files changed, 729 insertions(+), 243 deletions(-) create mode 100644 tests/integrations/omprpc/backpressure.py create mode 100644 tests/integrations/omprpc/test_callback_retirement.py create mode 100644 tests/integrations/omprpc/test_ui_responses.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5b3a784eb..3cb737b6a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,8 @@ This changelog is automatically generated by [Release Please](https://github.com ### Bug Fixes * Wait for Windows Job members' kernel-confirmed exit before releasing OMP or Codex subprocess ownership, preserving process handles across incomplete cleanup and preventing late child admission during termination. +* Retire native OMP sessions on an observed root exit, preserve usable sessions after post-admission command errors, and keep callbacks admitted during settlement within the original exit deadline. +* Keep native OMP stderr payloads out of SDK debug logs. ## [6.0.1](https://github.com/band-ai/band-sdk-python/compare/band-sdk-v6.0.0...band-sdk-v6.0.1) (2026-10-20) diff --git a/docs/acp.md b/docs/acp.md index ea42e546e..3073f256d 100644 --- a/docs/acp.md +++ b/docs/acp.md @@ -294,6 +294,17 @@ silently dropping events and reporting success. The wire reader does not wait for callback capacity, so callbacks can still issue correlated commands and answer native UI requests while delivery is held. +An observed root-process exit also retires the transport when descendants still +hold stdout open. Command errors, including an error reported after prompt +admission, raise `OmpRpcError` without retiring an otherwise usable process. +Host-tool and UI callbacks admitted while a settled turn drains share its +original callback-exit deadline; an unfinished callback fails the turn and starts +cleanup. SDK debug logs report stderr byte counts, never the child's raw stderr. + +Explicit MCP transports use `band.integrations.ompprocess.OmpMcpTransport`. +Native session metadata uses `band.converters.omp.OmpMetadataKey.ROOM` and +`OmpMetadataKey.SESSION`. + A `create_process` factory customizes creation without bypassing containment. On Windows it launches the existing helper; the client still owns the Job containing the native process and its descendants, including during cancelled diff --git a/src/band/converters/omp.py b/src/band/converters/omp.py index 331bedcf8..e83f643c6 100644 --- a/src/band/converters/omp.py +++ b/src/band/converters/omp.py @@ -19,10 +19,6 @@ class OmpMetadataKey(StrEnum): SESSION = "omp_session_file" -OMP_ROOM_METADATA_KEY = OmpMetadataKey.ROOM -OMP_SESSION_METADATA_KEY = OmpMetadataKey.SESSION - - class OmpSessionState(BaseModel): """The native session file each room may resume, plus the room's transcript.""" diff --git a/src/band/integrations/ompprocess.py b/src/band/integrations/ompprocess.py index 700c79f90..9549a4834 100644 --- a/src/band/integrations/ompprocess.py +++ b/src/band/integrations/ompprocess.py @@ -21,6 +21,7 @@ omp_provider_env, ) from band.integrations.ompwire import ( + OmpCommandName, OmpEvent, OmpEventType, OmpHostToolCall, @@ -78,12 +79,20 @@ def omp_rpc_command_in_workspace(command: Sequence[str], workspace: str) -> list return [*command, omp_cwd_argument(workspace)] +class OmpMcpTransport(StrEnum): + """The transports OMP accepts for external MCP servers.""" + + STDIO = "stdio" + HTTP = "http" + SSE = "sse" + + class OmpMcpServer(BaseModel): """An external MCP server projected into the SDK-owned OMP plugin.""" model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True) - type: Literal["stdio", "http", "sse"] | None = None + type: OmpMcpTransport | None = None command: str | None = None args: tuple[str, ...] = () env: dict[str, str] = Field(default_factory=dict, repr=False) @@ -95,10 +104,12 @@ class OmpMcpServer(BaseModel): def _require_one_transport(self) -> Self: if bool(self.command) == bool(self.url): raise ValueError("MCP server needs exactly one of command or url") - if self.command and self.type not in (None, "stdio"): - raise ValueError("a command requires stdio transport") - if self.url and self.type == "stdio": - raise ValueError("a URL requires http or sse transport") + if self.command and self.type not in (None, OmpMcpTransport.STDIO): + raise ValueError(f"a command requires {OmpMcpTransport.STDIO} transport") + if self.url and self.type == OmpMcpTransport.STDIO: + raise ValueError( + f"a URL requires {OmpMcpTransport.HTTP} or {OmpMcpTransport.SSE} transport" + ) return self @@ -195,7 +206,7 @@ async def omp_mcp_ready( """Require the connection extension and configured MCP servers before inference.""" command = omp_mcp_ready_command(plugin) commands = McpCommands.model_validate( - await runtime.request("get_available_commands") + await runtime.request(OmpCommandName.GET_AVAILABLE_COMMANDS) ) if not any(item.name == command for item in commands.commands): raise OmpStartupError("OMP did not load the required SDK startup extension") diff --git a/src/band/integrations/omprpc/client.py b/src/band/integrations/omprpc/client.py index 27cd53623..e9bcd951b 100644 --- a/src/band/integrations/omprpc/client.py +++ b/src/band/integrations/omprpc/client.py @@ -8,7 +8,7 @@ from binascii import Error as Base64Error from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass -from typing import Any, TypeVar, cast +from typing import Any, TypeVar from pydantic import JsonValue, ValidationError @@ -30,6 +30,7 @@ ) from band.integrations.omprpc.frames import ( ChunkAssembler, + host_tool_result_frame, parse_payload_line, payload_json, ui_response_frame, @@ -41,6 +42,7 @@ OMP_NATIVE_PROTOCOL_VERSION, OMP_REASSEMBLED_BYTES, OmpCommand, + OmpCommandName, OmpEvent, OmpEventType, OmpFrame, @@ -176,7 +178,7 @@ async def start(self) -> None: raise error self._apply_ready_limits(ready) negotiated = await self._request_frame( - "negotiate_protocol", + OmpCommandName.NEGOTIATE_PROTOCOL, {"protocolVersion": OMP_NATIVE_PROTOCOL_VERSION}, ) self._parse_response( @@ -263,7 +265,9 @@ def _ensure_running(self) -> None: if self._closing or self._proc is None: raise OmpProcessError("OMP RPC process is not running") if self._proc.returncode is not None: - raise OmpProcessError("OMP RPC process exited") + error = OmpProcessError("OMP RPC process exited") + self._record_failure(error) + raise error def _next_id(self) -> str: self._request_id += 1 @@ -290,7 +294,7 @@ async def _request_frame( ) self._pending[request_id] = PendingRequest(command, future) try: - await self._write_line(frame.model_dump_json()) + await self._write_frame(frame.root) response = await future if isinstance(response, (OmpRpcError, OmpProcessError)): raise response @@ -357,7 +361,7 @@ async def _execute_prompt( async def _complete_prompt(self, turn: PromptRun, message: str) -> None: try: accepted = await self._request_frame( - "prompt", {"message": message}, request_id=turn.id + OmpCommandName.PROMPT, {"message": message}, request_id=turn.id ) except OmpRpcError: turn.rejected = True @@ -374,13 +378,14 @@ async def _complete_prompt(self, turn: PromptRun, message: str) -> None: await self._wait_for_settlement(turn) await turn.events.join() if turn.callback_tasks: - _, pending = await asyncio.wait( - turn.callback_tasks, timeout=self.callback_exit_timeout_s - ) - if pending: + try: + async with asyncio.timeout(self.callback_exit_timeout_s): + while any(not task.done() for task in turn.callback_tasks): + await asyncio.wait(turn.callback_tasks) + except TimeoutError: error = OmpProcessError(CALLBACKS_UNFINISHED) self._record_failure(error) - raise error + raise error from None if outcome is not None and outcome.status != OmpPromptStatus.COMPLETED: raise OmpPromptError( outcome.status, @@ -399,7 +404,7 @@ async def _complete_prompt(self, turn: PromptRun, message: str) -> None: ) async def _wait_for_settlement(self, turn: PromptRun) -> None: - snapshot = await self._request_frame("get_state", {}) + snapshot = await self._request_frame(OmpCommandName.GET_STATE, {}) state = self._parse_response( snapshot.frame, OmpState, "Invalid OMP settlement protocol data" ) @@ -521,6 +526,10 @@ async def _dispatch_payload( self._cancel_callback( turn, OmpCallbackKind.TOOL, cancellation.targetId ) + else: + logger.debug( + "OMP RPC ignored a host-tool cancellation without an active prompt" + ) case OmpFrameType.UI_REQUEST: request = OmpUIRequest.model_validate(payload) if request.method == "cancel": @@ -555,23 +564,27 @@ async def _dispatch_payload( self._record_failure(overflow) else: logger.debug("OMP RPC ignored an event without an active prompt") + case _: + logger.debug("OMP RPC ignored a frame with an unknown type") def _dispatch_response(self, response: OmpResponse, sequence: int) -> None: if response.id is None: raise OmpProtocolError("OMP response is missing its command id") pending = self._pending.get(response.id) + turn = self._turn + active_prompt = turn if turn is not None and response.id == turn.id else None + expected_command = pending.command if pending is not None else None + if expected_command is None and active_prompt is not None: + expected_command = OmpCommandName.PROMPT + if expected_command is not None and response.command != expected_command: + raise OmpProtocolError("OMP response command does not match its request") if pending is None: - if ( - not response.success - and self._turn is not None - and response.id == self._turn.id - ): - self._turn.fail(command_error(response)) + if not response.success and active_prompt is not None: + active_prompt.rejected = True + active_prompt.fail(command_error(response)) return logger.debug("OMP RPC ignored an unmatched command response") return - if response.command != pending.command: - raise OmpProtocolError("OMP response command does not match its request") if pending.future.done(): return if response.success: @@ -628,19 +641,7 @@ async def _handle_tool(self, turn: PromptRun, call: OmpHostToolCall) -> None: result = await turn.on_tool(call) if not self._callback_still_active(turn, OmpCallbackKind.TOOL, call.id): return - await self._write_frame( - { - "type": OmpFrameType.HOST_TOOL_RESULT, - "id": call.id, - "result": cast( - dict[str, JsonValue], - result.model_dump( - mode="json", exclude={"is_error"}, exclude_none=True - ), - ), - "isError": result.is_error, - } - ) + await self._write_frame(host_tool_result_frame(call.id, result)) async def _handle_ui(self, turn: PromptRun, request: OmpUIRequest) -> None: response = await turn.on_ui(request) @@ -652,9 +653,7 @@ async def _read_stderr(self, process: asyncio.subprocess.Process) -> None: assert process.stderr is not None try: while data := await process.stderr.read(_DRAIN_BYTES): - logger.debug( - "OMP stderr: %s", data.decode("utf-8", errors="replace").rstrip() - ) + logger.debug("OMP stderr drained: %s bytes", len(data)) except asyncio.CancelledError: raise except Exception: @@ -783,7 +782,9 @@ async def _close_process(self) -> None: ): try: async with asyncio.timeout(_ABORT_TIMEOUT_S): - await self._request_frame("abort", {}, allow_closing=True) + await self._request_frame( + OmpCommandName.ABORT, {}, allow_closing=True + ) except (TimeoutError, OmpRpcError, OmpProcessError) as error: logger.warning("OMP graceful abort skipped: %s", error) reader = self._reader_task diff --git a/src/band/integrations/omprpc/frames.py b/src/band/integrations/omprpc/frames.py index 0bb968029..7cfb3d25c 100644 --- a/src/band/integrations/omprpc/frames.py +++ b/src/band/integrations/omprpc/frames.py @@ -1,9 +1,10 @@ -"""Native OMP JSONL frame reassembly and host UI replies.""" +"""Native OMP JSONL frame reassembly and host replies.""" from __future__ import annotations import base64 from dataclasses import dataclass +from typing import cast from pydantic import JsonValue, TypeAdapter @@ -14,6 +15,7 @@ OmpChunk, OmpFrame, OmpFrameType, + OmpToolResult, OmpUIResponse, ) @@ -104,19 +106,24 @@ def parse_payload_line(line: bytes) -> dict[str, JsonValue]: return _PAYLOAD_ADAPTER.validate_json(line, strict=True) +def host_tool_result_frame( + request_id: str, result: OmpToolResult +) -> dict[str, JsonValue]: + return { + "type": OmpFrameType.HOST_TOOL_RESULT, + "id": request_id, + "result": cast( + dict[str, JsonValue], + result.model_dump(mode="json", exclude={"is_error"}, exclude_none=True), + ), + "isError": result.is_error, + } + + def ui_response_frame(request_id: str, response: OmpUIResponse) -> dict[str, JsonValue]: - frame: dict[str, JsonValue] = { + answer = response.model_dump(mode="json", exclude={"cancelled"}, exclude_none=True) + return { "type": OmpFrameType.UI_RESPONSE, "id": request_id, + **(answer or {"cancelled": True}), } - if response.value is not None: - frame["value"] = response.value - elif response.confirmed is not None: - frame["confirmed"] = response.confirmed - elif response.answers is not None: - frame["answers"] = [ - answer.model_dump(exclude_none=True) for answer in response.answers - ] - else: - frame["cancelled"] = True - return frame diff --git a/src/band/integrations/ompwire.py b/src/band/integrations/ompwire.py index b67f73eda..f8624e328 100644 --- a/src/band/integrations/ompwire.py +++ b/src/band/integrations/ompwire.py @@ -22,6 +22,16 @@ OMP_ASSISTANT_ROLE = "assistant" +class OmpCommandName(StrEnum): + """Name SDK-issued commands without restricting arbitrary native requests.""" + + NEGOTIATE_PROTOCOL = "negotiate_protocol" + PROMPT = "prompt" + GET_STATE = "get_state" + GET_AVAILABLE_COMMANDS = "get_available_commands" + ABORT = "abort" + + class OmpPromptStatus(StrEnum): COMPLETED = "completed" ERROR = "error" @@ -253,6 +263,8 @@ class OmpUIRequest(OmpWireModel): class OmpUIResponse(OmpWireModel): + """Answer one dialog, or cancel it when no answer variant is supplied.""" + value: str | None = None confirmed: bool | None = None answers: list[OmpAskAnswer] | None = None diff --git a/tests/converters/test_omp.py b/tests/converters/test_omp.py index 51fc19ff8..ef46b0c42 100644 --- a/tests/converters/test_omp.py +++ b/tests/converters/test_omp.py @@ -3,9 +3,8 @@ from __future__ import annotations from band.converters.omp import ( - OMP_ROOM_METADATA_KEY, - OMP_SESSION_METADATA_KEY, OmpHistoryConverter, + OmpMetadataKey, ) from band.core.types import MessageType @@ -22,22 +21,22 @@ def test_latest_session_file_wins_per_room(self) -> None: { "message_type": MessageType.TASK, "metadata": { - OMP_ROOM_METADATA_KEY: "room-1", - OMP_SESSION_METADATA_KEY: "/sessions/old.jsonl", + OmpMetadataKey.ROOM: "room-1", + OmpMetadataKey.SESSION: "/sessions/old.jsonl", }, }, { "message_type": MessageType.TASK, "metadata": { - OMP_ROOM_METADATA_KEY: "room-2", - OMP_SESSION_METADATA_KEY: "/sessions/other.jsonl", + OmpMetadataKey.ROOM: "room-2", + OmpMetadataKey.SESSION: "/sessions/other.jsonl", }, }, { "message_type": MessageType.TASK, "metadata": { - OMP_ROOM_METADATA_KEY: "room-1", - OMP_SESSION_METADATA_KEY: "/sessions/new.jsonl", + OmpMetadataKey.ROOM: "room-1", + OmpMetadataKey.SESSION: "/sessions/new.jsonl", }, }, ] @@ -53,20 +52,20 @@ def test_unusable_metadata_is_ignored(self) -> None: {"message_type": MessageType.TASK, "metadata": "not a mapping"}, { "message_type": MessageType.TASK, - "metadata": {OMP_ROOM_METADATA_KEY: "room-1"}, + "metadata": {OmpMetadataKey.ROOM: "room-1"}, }, { "message_type": MessageType.TASK, "metadata": { - OMP_ROOM_METADATA_KEY: "room-1", - OMP_SESSION_METADATA_KEY: "", + OmpMetadataKey.ROOM: "room-1", + OmpMetadataKey.SESSION: "", }, }, { "message_type": MessageType.TASK, "metadata": { - OMP_ROOM_METADATA_KEY: 7, - OMP_SESSION_METADATA_KEY: "/sessions/x.jsonl", + OmpMetadataKey.ROOM: 7, + OmpMetadataKey.SESSION: "/sessions/x.jsonl", }, }, ] @@ -79,8 +78,8 @@ def test_session_metadata_on_non_task_messages_is_ignored(self) -> None: { "message_type": MessageType.TEXT, "metadata": { - OMP_ROOM_METADATA_KEY: "room-1", - OMP_SESSION_METADATA_KEY: "/sessions/injected.jsonl", + OmpMetadataKey.ROOM: "room-1", + OmpMetadataKey.SESSION: "/sessions/injected.jsonl", }, "content": "hello", "sender_name": "Alice", diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index afdc44fd0..906e635b8 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -21,6 +21,7 @@ OMP_FRAME_BYTES, OMP_NATIVE_PROTOCOL_VERSION, OMP_REASSEMBLED_BYTES, + OmpCommandName, OmpPromptStatus, ) @@ -38,6 +39,11 @@ UI_CANCELLED_ID = "ui-cancelled" OVERSIZE_AFTER_READY = "oversize-after-ready" READY_LIMIT_BYTES = 1024 +LATE_PROMPT_ERROR = "late-prompt-error" +SETTLED_CALLBACK_SCENARIO = "settled-during-callback" +SETTLED_CALLBACK_ID = "host-held" +LATE_CALLBACK_ID = "host-late" +STDIN_BACKPRESSURE = "stdin-backpressure" def burst_message_id(number: int) -> str: @@ -51,6 +57,7 @@ class Frame(BaseModel): type: str id: str | None = None message: str | None = None + failed_command: str | None = None result: dict[str, JsonValue] | None = None isError: bool | None = None value: str | None = None @@ -188,6 +195,7 @@ class Peer: def __init__(self, scenario: str) -> None: self.scenario = scenario self.settled = True + self.state_requested = False self.prompt_id: str | None = None self.prompt_count = 0 self.frames: list[dict[str, JsonValue]] = [] @@ -206,6 +214,16 @@ def respond(self, frame: Frame, data: dict[str, JsonValue] | None = None) -> Non result["data"] = data send(result) + def checkpoint(self, frame: Frame) -> None: + if ( + self.scenario == SETTLED_CALLBACK_SCENARIO + and self.prompt_id is not None + and not self.state_requested + ): + self.held = frame + return + self.respond(frame, {"pid": os.getpid(), "frames": self.frames}) + def end( self, text: str, @@ -283,6 +301,7 @@ def prompt(self, frame: Frame) -> None: self.prompt_id = frame.id self.prompt_count += 1 self.settled = False + self.state_requested = False if ( self.prompt_count > 1 and self.scenario != "lifecycle" ) or frame.message == FOLLOW_UP_PROMPT: @@ -457,11 +476,11 @@ def prompt(self, frame: Frame) -> None: case "no-result": self.end(COMPLETE) self.settle() - case "settled-during-callback": + case scenario if scenario == SETTLED_CALLBACK_SCENARIO: send( { "type": "host_tool_call", - "id": "host-held", + "id": SETTLED_CALLBACK_ID, "toolCallId": "call-held", "toolName": "band_chat", "arguments": {}, @@ -535,6 +554,7 @@ def handle(self, frame: Frame) -> None: case "prompt": self.prompt(frame) case "get_state": + self.state_requested = True self.respond( frame, { @@ -555,6 +575,13 @@ def handle(self, frame: Frame) -> None: "queuedMessageCount": 0, }, ) + if ( + self.scenario == SETTLED_CALLBACK_SCENARIO + and self.held is not None + and self.held.type == "checkpoint" + ): + self.checkpoint(self.held) + self.held = None if self.scenario == "settles-during-state" and not self.settled: self.settle() if ( @@ -607,6 +634,30 @@ def handle(self, frame: Frame) -> None: } ) case "release": + if self.scenario == SETTLED_CALLBACK_SCENARIO: + send( + { + "type": "host_tool_call", + "id": LATE_CALLBACK_ID, + "toolCallId": f"call-{LATE_CALLBACK_ID}", + "toolName": "band_chat", + "arguments": {}, + } + ) + self.respond(frame) + return + if self.scenario == LATE_PROMPT_ERROR: + send( + { + "type": "response", + "id": self.prompt_id, + "command": frame.failed_command or OmpCommandName.PROMPT, + "success": False, + "error": "Prompt command failed after admission", + } + ) + self.respond(frame) + return if self.scenario == "settlement-before-result": self.result(settled=False) else: @@ -719,7 +770,7 @@ def handle(self, frame: Frame) -> None: } ) case "checkpoint": - self.respond(frame, {"pid": os.getpid(), "frames": self.frames}) + self.checkpoint(frame) case "abort": if self.scenario != "stubborn": self.respond(frame) @@ -727,6 +778,10 @@ def handle(self, frame: Frame) -> None: self.end(OmpPromptStatus.ABORTED, stop=OmpPromptStatus.ABORTED) self.settle() case "first" | "hold": + if self.scenario == STDIN_BACKPRESSURE: + self.respond(frame) + threading.Event().wait() + return self.held = frame case "second": self.respond(frame, {"value": "second"}) @@ -747,7 +802,10 @@ def handle(self, frame: Frame) -> None: case "crash": raise SystemExit(CRASH_EXIT) case "stderr": - sys.stderr.write("x" * (2 * 1024 * 1024)) + sys.stderr.write( + FRAME_PAYLOAD_MARK + * (2 * OMP_FRAME_BYTES // len(FRAME_PAYLOAD_MARK)) + ) sys.stderr.flush() self.respond(frame, {"drained": True}) case _: diff --git a/tests/integrations/ompprocess/test_launch.py b/tests/integrations/ompprocess/test_launch.py index 5b6f40dad..629a5c1ba 100644 --- a/tests/integrations/ompprocess/test_launch.py +++ b/tests/integrations/ompprocess/test_launch.py @@ -2,36 +2,10 @@ from __future__ import annotations -from band.integrations.omp import ( - OMP_APPROVAL_MODE_ALWAYS_ASK, - OMP_APPROVAL_MODE_FLAG, - finalize_omp_command, - omp_cwd_argument, -) -from band.integrations.ompprocess import ( - DEFAULT_OMP_COMMAND, - omp_rpc_command_in_workspace, - omp_runtime_env, -) +from band.integrations.ompprocess import omp_runtime_env class TestNativeLaunchSettings: - def test_default_command_is_safe_after_finalize(self) -> None: - assert finalize_omp_command(DEFAULT_OMP_COMMAND)[-2:] == [ - OMP_APPROVAL_MODE_FLAG, - OMP_APPROVAL_MODE_ALWAYS_ASK, - ] - - def test_workspace_is_selected_without_a_subprocess_cwd(self) -> None: - assert omp_rpc_command_in_workspace( - ("bunx", *DEFAULT_OMP_COMMAND, "-v"), "/w" - ) == [ - "bunx", - *DEFAULT_OMP_COMMAND, - "-v", - omp_cwd_argument("/w"), - ] - def test_runtime_env_routes_the_selected_provider_key(self) -> None: assert omp_runtime_env(env=None, model="openai/custom", api_key="secret") == { "OPENAI_API_KEY": "secret" diff --git a/tests/integrations/ompprocess/test_mcp_plugin.py b/tests/integrations/ompprocess/test_mcp_plugin.py index 3aa2c2c99..f7e826afd 100644 --- a/tests/integrations/ompprocess/test_mcp_plugin.py +++ b/tests/integrations/ompprocess/test_mcp_plugin.py @@ -8,18 +8,16 @@ from pathlib import Path import pytest +from pydantic import ValidationError from band.integrations.ompprocess import ( OMP_MCP_SERVERS_CONFIG, OMP_MCP_STARTUP_CONFIG, OMP_MCP_STARTUP_JS, - OMP_MODE_FLAG, - OMP_RPC_UI_MODE, RELATIVE_MCP_COMMAND_PREFIXES, OmpMcpServer, - omp_mcp_arguments, + OmpMcpTransport, omp_mcp_plugin, - omp_mcp_ready_command, ) from tests.ostable import skip_unless_os @@ -38,12 +36,13 @@ def test_nothing_to_project_yields_no_plugin(self, tmp_path: Path) -> None: def test_servers_land_in_the_plugin_mcp_config(self, tmp_path: Path) -> None: servers = { "local": OmpMcpServer(command="./server", args=("--port", "1")), - "remote": OmpMcpServer(type="sse", url="https://mcp.example/sse"), + "remote": OmpMcpServer( + type=OmpMcpTransport.SSE, url="https://mcp.example/sse" + ), } with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: assert plugin is not None root = Path(plugin) - assert root.name == "band-sdk" config = json.loads( (root / OMP_MCP_SERVERS_CONFIG).read_text(encoding="utf-8") ) @@ -52,17 +51,13 @@ def test_servers_land_in_the_plugin_mcp_config(self, tmp_path: Path) -> None: ) assert startup_config == {"startup": {"quiet": False}} assert config["mcpServers"]["remote"] == { - "type": "sse", + "type": OmpMcpTransport.SSE, "url": "https://mcp.example/sse", } local = config["mcpServers"]["local"] assert local["args"] == ["--port", "1"] assert Path(local["cwd"]) == tmp_path.resolve() assert Path(local["command"]) == (tmp_path / "server").resolve() - startup = (root / OMP_MCP_STARTUP_JS).read_text(encoding="utf-8") - assert omp_mcp_ready_command(plugin) in startup - assert '"band-sdk:local"' in startup - assert '"band-sdk:remote"' in startup assert not Path(plugin).exists() def test_connection_instructions_become_an_always_applied_rule( @@ -78,16 +73,6 @@ def test_connection_instructions_become_an_always_applied_rule( "---\nalwaysApply: true\n---\nReply through Band." ) - def test_arguments_load_the_plugin_and_its_startup_extension( - self, tmp_path: Path - ) -> None: - plugin = str(tmp_path / "band-sdk") - arguments = omp_mcp_arguments(plugin) - assert arguments[:2] == ("--plugin-dir", plugin) - assert arguments[arguments.index(OMP_MODE_FLAG) + 1] == OMP_RPC_UI_MODE - assert str(Path(plugin) / OMP_MCP_STARTUP_JS) in arguments - assert str(Path(plugin) / OMP_MCP_STARTUP_CONFIG) in arguments - @pytest.mark.parametrize("accepted_on", [("posix",)]) def test_projected_plugin_files_are_owner_only( self, tmp_path: Path, accepted_on: tuple[str, ...] @@ -134,15 +119,79 @@ def test_relative_command_is_absolutized(self, tmp_path: Path, prefix: str) -> N ) assert Path(config["mcpServers"]["local"]["command"]).is_absolute() + @pytest.mark.parametrize("transport", [None, OmpMcpTransport.STDIO]) + def test_command_transport_projects_a_launchable_server( + self, tmp_path: Path, transport: OmpMcpTransport | None + ) -> None: + servers = { + "local": OmpMcpServer( + type=transport, command="server", args=("--port", "1") + ) + } + with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: + assert plugin is not None + config = json.loads( + (Path(plugin) / OMP_MCP_SERVERS_CONFIG).read_text(encoding="utf-8") + ) + local = config["mcpServers"]["local"] + assert local.get("type") == transport + assert local["command"] == "server" + assert local["args"] == ["--port", "1"] + assert Path(local["cwd"]) == tmp_path.resolve() + + @pytest.mark.parametrize( + "transport", [None, OmpMcpTransport.HTTP, OmpMcpTransport.SSE] + ) + def test_url_transport_preserves_the_remote_endpoint( + self, tmp_path: Path, transport: OmpMcpTransport | None + ) -> None: + servers = { + "remote": OmpMcpServer( + type=transport, + url="https://mcp.example/tools", + headers={"Authorization": "Bearer token"}, + ) + } + with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: + assert plugin is not None + config = json.loads( + (Path(plugin) / OMP_MCP_SERVERS_CONFIG).read_text(encoding="utf-8") + ) + remote = config["mcpServers"]["remote"] + assert remote.get("type") == transport + assert remote["url"] == "https://mcp.example/tools" + assert remote["headers"] == {"Authorization": "Bearer token"} + assert "command" not in remote + assert "cwd" not in remote + @pytest.mark.parametrize( "fields", [ - {}, - {"command": "server", "url": "https://mcp.example"}, - {"command": "server", "type": "http"}, - {"url": "https://mcp.example", "type": "stdio"}, + pytest.param({}, id="no-endpoint"), + pytest.param( + {"command": "server", "url": "https://mcp.example"}, + id="two-endpoints", + ), + pytest.param( + {"command": "server", "type": OmpMcpTransport.HTTP}, + id="command-http", + ), + pytest.param( + {"command": "server", "type": OmpMcpTransport.SSE}, + id="command-sse", + ), + pytest.param( + {"url": "https://mcp.example", "type": OmpMcpTransport.STDIO}, + id="url-stdio", + ), + pytest.param( + {"command": "server", "type": "websocket"}, + id="unknown-transport", + ), ], ) - def test_a_server_needs_exactly_one_transport(self, fields: dict[str, str]) -> None: - with pytest.raises(ValueError): - OmpMcpServer(**fields) + def test_conflicting_transport_settings_are_rejected( + self, fields: dict[str, str] + ) -> None: + with pytest.raises(ValidationError): + OmpMcpServer.model_validate(fields) diff --git a/tests/integrations/ompprocess/test_mcp_startup.py b/tests/integrations/ompprocess/test_mcp_startup.py index 480b11826..01f34b39c 100644 --- a/tests/integrations/ompprocess/test_mcp_startup.py +++ b/tests/integrations/ompprocess/test_mcp_startup.py @@ -12,6 +12,9 @@ pytest.importorskip("psutil", reason="psutil not installed (band-sdk[omp])") from band.integrations.ompprocess import ( + OMP_MCP_RECEIPT, + OMP_PLUGIN_ROOT, + McpConnectionStatus, OmpMcpServer, OmpStartupError, omp_mcp_ready, @@ -30,59 +33,78 @@ class TestRequiredMcpStartup: @pytest.mark.parametrize( - ("receipt", "reason"), + "receipt", [ - (None, "no valid SDK startup receipt"), - ("not JSON", "no valid SDK startup receipt"), - ( + pytest.param(None, id="missing-receipt"), + pytest.param("not JSON", id="malformed-receipt"), + pytest.param( {"systemPromptReady": True, "servers": []}, - "did not acknowledge every configured MCP server", + id="missing-server", ), - ( + pytest.param( { "systemPromptReady": True, "servers": [ - {"type": "connected", "serverName": "band-sdk:required"}, - {"type": "connected", "serverName": "band-sdk:required"}, + { + "type": McpConnectionStatus.CONNECTED, + "serverName": f"{OMP_PLUGIN_ROOT}:required", + }, + { + "type": McpConnectionStatus.CONNECTED, + "serverName": f"{OMP_PLUGIN_ROOT}:required", + }, ], }, - "did not acknowledge every configured MCP server", + id="duplicate-server", ), - ( + pytest.param( { "systemPromptReady": True, - "servers": [{"type": "connected", "serverName": "unrelated"}], + "servers": [ + { + "type": McpConnectionStatus.CONNECTED, + "serverName": "unrelated", + } + ], }, - "unexpected MCP startup receipt", + id="unexpected-server", ), - ( + pytest.param( { "systemPromptReady": True, - "servers": [{"type": "failed", "serverName": "band-sdk:required"}], + "servers": [ + { + "type": McpConnectionStatus.FAILED, + "serverName": f"{OMP_PLUGIN_ROOT}:required", + } + ], }, - "no valid SDK startup receipt", + id="missing-failure-detail", ), - ( + pytest.param( { "systemPromptReady": True, "servers": [ { - "type": "failed", - "serverName": "band-sdk:required", + "type": McpConnectionStatus.FAILED, + "serverName": f"{OMP_PLUGIN_ROOT}:required", "error": "initialize rejected", } ], }, - "band-sdk:required failed: initialize rejected", + id="failed-server", ), - ( + pytest.param( { "systemPromptReady": False, "servers": [ - {"type": "connected", "serverName": "band-sdk:required"} + { + "type": McpConnectionStatus.CONNECTED, + "serverName": f"{OMP_PLUGIN_ROOT}:required", + } ], }, - "required Band connection instructions", + id="unloaded-instructions", ), ], ) @@ -90,72 +112,32 @@ async def test_rejects_incomplete_or_failed_receipts( self, tmp_path: Path, receipt: str | dict[str, object] | None, - reason: str, ) -> None: - plugin = tmp_path / "band-sdk" + plugin = tmp_path / OMP_PLUGIN_ROOT plugin.mkdir() if receipt is not None: payload = json.dumps(receipt) if isinstance(receipt, dict) else receipt - (plugin / "receipt.json").write_text(payload, encoding="utf-8") + (plugin / OMP_MCP_RECEIPT).write_text(payload, encoding="utf-8") runtime = AsyncMock(spec=OmpRpcClient) runtime.request.return_value = { "commands": [{"name": omp_mcp_ready_command(str(plugin))}], } - with pytest.raises(OmpStartupError, match=reason): + with pytest.raises(OmpStartupError): await omp_mcp_ready( runtime, str(plugin), {"required": OmpMcpServer(command="server")}, ) - async def test_missing_startup_extension_cannot_fall_through_to_inference( - self, tmp_path: Path - ) -> None: + async def test_missing_startup_extension_is_rejected(self, tmp_path: Path) -> None: runtime = AsyncMock(spec=OmpRpcClient) runtime.request.return_value = {"commands": []} - with pytest.raises(OmpStartupError, match="required SDK startup extension"): + with pytest.raises(OmpStartupError): await omp_mcp_ready(runtime, str(tmp_path), {}) - runtime.prompt.assert_not_awaited() - async def test_complete_receipt_prompts_the_readiness_command( - self, tmp_path: Path - ) -> None: - plugin = tmp_path / "band-sdk" - plugin.mkdir() - command = omp_mcp_ready_command(str(plugin)) - receipt = json.dumps( - { - "systemPromptReady": True, - "servers": [{"type": "connected", "serverName": "band-sdk:required"}], - } - ) - runtime = AsyncMock(spec=OmpRpcClient) - runtime.request.return_value = {"commands": [{"name": command}]} - - async def prompt(message: str, **_callbacks: object) -> None: - _ = message - (plugin / "receipt.json").write_text(receipt, encoding="utf-8") - - runtime.prompt.side_effect = prompt - await omp_mcp_ready( - runtime, - str(plugin), - {"required": OmpMcpServer(command="server")}, - ) - runtime.prompt.assert_awaited_once() - assert runtime.prompt.await_args is not None - assert runtime.prompt.await_args.args[0] == f"/{command}" - - @pytest.mark.parametrize( - ("invoke", "reason"), - [ - ("event", "invoked an agent"), - ("tool", "invoked a host tool"), - ("ui", "requested UI"), - ], - ) + @pytest.mark.parametrize("invoke", ["event", "tool", "ui"]) async def test_startup_callbacks_cannot_fall_through_to_inference( - self, tmp_path: Path, invoke: str, reason: str + self, tmp_path: Path, invoke: str ) -> None: runtime = AsyncMock(spec=OmpRpcClient) runtime.request.return_value = { @@ -192,5 +174,5 @@ async def prompt( ) runtime.prompt.side_effect = prompt - with pytest.raises(OmpStartupError, match=reason): + with pytest.raises(OmpStartupError): await omp_mcp_ready(runtime, str(tmp_path), {}) diff --git a/tests/integrations/ompprocess/test_tool_names.py b/tests/integrations/ompprocess/test_tool_names.py index ae8d3a52a..1ea17aeeb 100644 --- a/tests/integrations/ompprocess/test_tool_names.py +++ b/tests/integrations/ompprocess/test_tool_names.py @@ -4,7 +4,9 @@ import pytest +from band.integrations.omp import XD_MCP_PREFIX, XD_URL_PREFIX from band.integrations.ompprocess import ( + OMP_PLUGIN_TOOL_NAMESPACE, normalize_native_mcp_device_call, normalize_native_mcp_tool_name, ) @@ -16,8 +18,11 @@ class TestPluginNamespacedToolNames: @pytest.mark.parametrize( "path", [ - "xd://mcp__band_send_message", - "xd://mcp__band_sdk_band_band_send_message", + f"{XD_URL_PREFIX}{XD_MCP_PREFIX}band_send_message", + ( + f"{XD_URL_PREFIX}{XD_MCP_PREFIX}" + f"{OMP_PLUGIN_TOOL_NAMESPACE}band_band_send_message" + ), ], ) def test_executed_device_write_maps_to_the_band_tool(self, path: str) -> None: @@ -27,7 +32,11 @@ def test_executed_device_write_maps_to_the_band_tool(self, path: str) -> None: ) == ("band_send_message", {"chat_id": "room-1", "content": "hi"}) @pytest.mark.parametrize( - "title", ["mcp__band_send_message", "mcp__band_sdk_band_band_send_message"] + "title", + [ + f"{XD_MCP_PREFIX}band_send_message", + f"{XD_MCP_PREFIX}{OMP_PLUGIN_TOOL_NAMESPACE}band_band_send_message", + ], ) def test_mcp_title_maps_to_the_band_tool(self, title: str) -> None: assert normalize_native_mcp_tool_name(title, OWN) == "band_send_message" @@ -38,7 +47,10 @@ def test_mcp_title_maps_to_the_band_tool(self, title: str) -> None: def test_a_namespaced_read_is_never_the_band_tool(self) -> None: arguments = { - "path": "xd://mcp__band_sdk_band_band_send_message", + "path": ( + f"{XD_URL_PREFIX}{XD_MCP_PREFIX}" + f"{OMP_PLUGIN_TOOL_NAMESPACE}band_band_send_message" + ), "content": '{"content":"hi"}', } assert normalize_native_mcp_device_call( @@ -47,6 +59,8 @@ def test_a_namespaced_read_is_never_the_band_tool(self) -> None: def test_a_foreign_namespace_stays_unmapped(self) -> None: assert ( - normalize_native_mcp_tool_name("mcp__other_sdk_band_send_message", OWN) + normalize_native_mcp_tool_name( + f"{XD_MCP_PREFIX}other_sdk_band_send_message", OWN + ) is None ) diff --git a/tests/integrations/omprpc/backpressure.py b/tests/integrations/omprpc/backpressure.py new file mode 100644 index 000000000..1597758a0 --- /dev/null +++ b/tests/integrations/omprpc/backpressure.py @@ -0,0 +1,35 @@ +"""Observe real subprocess pipe saturation without replacing writes or drains.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +from band.integrations.ompwire import OMP_FRAME_BYTES + +PRESSURE_WRITE_COUNT = 8 + + +class BackpressureProtocol(asyncio.subprocess.SubprocessStreamProtocol): + def __init__( + self, limit: int, loop: asyncio.AbstractEventLoop, paused: asyncio.Event + ) -> None: + super().__init__(limit, loop) + self.paused = paused + + def pause_writing(self) -> None: + super().pause_writing() + self.paused.set() + + +class BackpressureSpawn: + def __init__(self) -> None: + self.paused = asyncio.Event() + + async def create( + self, *command: str, limit: int = OMP_FRAME_BYTES, **kwargs: Any + ) -> asyncio.subprocess.Process: + loop = asyncio.get_running_loop() + protocol = BackpressureProtocol(limit, loop, self.paused) + transport, _ = await loop.subprocess_exec(lambda: protocol, *command, **kwargs) + return asyncio.subprocess.Process(transport, protocol, loop) diff --git a/tests/integrations/omprpc/test_callback_retirement.py b/tests/integrations/omprpc/test_callback_retirement.py new file mode 100644 index 000000000..d0cfb6c07 --- /dev/null +++ b/tests/integrations/omprpc/test_callback_retirement.py @@ -0,0 +1,121 @@ +"""Callbacks admitted during settlement share the original retirement deadline.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from band.integrations.omprpc import OmpProcessError +from band.integrations.ompwire import ( + OmpCommandName, + OmpEventType, + OmpHostToolCall, + OmpToolResult, +) +from tests.integrations.omppeer import ( + LATE_CALLBACK_ID, + SETTLED_CALLBACK_ID, + SETTLED_CALLBACK_SCENARIO, +) +from tests.integrations.omprpc.cancellation import HeldCleanup +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + checkpoint_frames, + owned_rpc_process, +) +from tests.toolkit.clock import binary_clock_resolution +from tests.toolkit.timecontrol import controlled_clock + + +class TestCallbackRetirement: + @pytest.mark.parametrize( + "retire_second_on_first", + [False, True], + ids=["held-late-callback", "callback-turnover"], + ) + @pytest.mark.looptime(resolution=binary_clock_resolution()) + async def test_a_callback_admitted_during_the_join_cannot_escape_the_deadline( + self, client_factory: ClientFactory, retire_second_on_first: bool + ) -> None: + client = await client_factory(SETTLED_CALLBACK_SCENARIO) + grace = 1.0 + client.callback_exit_timeout_s = grace + process = owned_rpc_process(client) + observer = Observer() + cleanup = HeldCleanup(client) + joining: asyncio.Task[None] | None = None + first_started, first_retired, second_started = ( + asyncio.Event() for _ in range(3) + ) + release_first, release_second = (asyncio.Event() for _ in range(2)) + second_cancelled = asyncio.Event() + + async def on_tool(call: OmpHostToolCall) -> OmpToolResult: + if call.id == SETTLED_CALLBACK_ID: + task = asyncio.current_task() + assert task is not None + task.add_done_callback(lambda _done: first_retired.set()) + first_started.set() + await release_first.wait() + else: + assert call.id == LATE_CALLBACK_ID + second_started.set() + try: + await ( + first_retired.wait() + if retire_second_on_first + else release_second.wait() + ) + except asyncio.CancelledError: + second_cancelled.set() + raise + return OmpToolResult(content=[]) + + async with ( + controlled_clock() as clock, + clock.background( + client.prompt( + "first", + on_event=observer.on_event, + on_tool=on_tool, + on_ui=observer.on_ui, + ) + ) as running, + ): + try: + await first_started.wait() + await observer.wait_for(OmpEventType.AGENT_END) + frames = await checkpoint_frames(client) + assert any( + frame["type"] == OmpCommandName.GET_STATE for frame in frames + ) + await client.request("release") + await second_started.wait() + if retire_second_on_first: + clock.advance(grace) + release_first.set() + await first_retired.wait() + if not retire_second_on_first: + clock.advance(grace) + joining = asyncio.create_task(cleanup.wait_for_join(running)) + done, _ = await asyncio.wait( + (running, joining), return_when=asyncio.FIRST_COMPLETED + ) + assert joining in done + await joining + assert process.returncode is None + cleanup.release.set() + with pytest.raises(OmpProcessError): + await clock.run_io(running) + if not retire_second_on_first: + assert second_cancelled.is_set() + assert process.returncode is not None + finally: + release_first.set() + release_second.set() + cleanup.release.set() + if joining is not None: + joining.cancel() + await asyncio.gather(joining, return_exceptions=True) diff --git a/tests/integrations/omprpc/test_interruption.py b/tests/integrations/omprpc/test_interruption.py index 5e99b9c53..ccd7f0b52 100644 --- a/tests/integrations/omprpc/test_interruption.py +++ b/tests/integrations/omprpc/test_interruption.py @@ -3,12 +3,26 @@ from __future__ import annotations import asyncio +import sys import pytest - -from band.integrations.ompwire import OmpEventType, OmpHostToolCall, OmpToolResult +from pydantic import JsonValue + +from band.integrations.omprpc import OmpProcessError, OmpRpcClient +from band.integrations.ompwire import ( + OMP_FRAME_BYTES, + OmpEventType, + OmpHostToolCall, + OmpToolResult, +) +from tests.integrations.omppeer import STDIN_BACKPRESSURE +from tests.integrations.omprpc.backpressure import ( + PRESSURE_WRITE_COUNT, + BackpressureSpawn, +) from tests.integrations.omprpc.cancellation import HeldCleanup from tests.integrations.omprpc.kit import ( + OMP_PEER_SCRIPT, ClientFactory, Observer, owned_rpc_process, @@ -167,30 +181,30 @@ async def on_tool(call: OmpHostToolCall) -> OmpToolResult: running.cancel() await asyncio.gather(running, return_exceptions=True) - async def test_close_finishes_while_a_stdin_drain_is_held( - self, client_factory: ClientFactory - ) -> None: - client = await client_factory("basic") - process = owned_rpc_process(client) - stdin = process.stdin - assert stdin is not None - entered = asyncio.Event() - release = asyncio.Event() - drain = stdin.drain - - async def hold_drain() -> None: - entered.set() - await release.wait() - await drain() - - stdin.drain = hold_drain - writing = asyncio.create_task(client.request("checkpoint")) + async def test_close_finishes_while_a_stdin_drain_is_held(self) -> None: + spawn = BackpressureSpawn() + client = OmpRpcClient( + (sys.executable, "-u", OMP_PEER_SCRIPT, "--scenario", STDIN_BACKPRESSURE), + create_process=spawn.create, + ) + writes: list[asyncio.Task[dict[str, JsonValue]]] = [] try: + await client.start() + process = owned_rpc_process(client) + await client.request("hold") + payload = "x" * (OMP_FRAME_BYTES // 2) + writes = [ + asyncio.create_task(client.request("pressure", data=payload)) + for _ in range(PRESSURE_WRITE_COUNT) + ] async with asyncio.timeout(5): - await entered.wait() + await spawn.paused.wait() await client.close() assert process.returncode is not None + outcomes = await asyncio.gather(*writes, return_exceptions=True) + assert all(isinstance(outcome, OmpProcessError) for outcome in outcomes) finally: - release.set() - writing.cancel() - await asyncio.gather(writing, return_exceptions=True) + for writing in writes: + writing.cancel() + await asyncio.gather(*writes, return_exceptions=True) + await client.close() diff --git a/tests/integrations/omprpc/test_prompt_results.py b/tests/integrations/omprpc/test_prompt_results.py index 3c50b1c1d..1c0a8561d 100644 --- a/tests/integrations/omprpc/test_prompt_results.py +++ b/tests/integrations/omprpc/test_prompt_results.py @@ -6,14 +6,15 @@ import pytest -from band.integrations.omprpc import OmpPromptError, OmpRpcError +from band.integrations.omprpc import OmpPromptError, OmpProtocolError, OmpRpcError from band.integrations.ompwire import ( OMP_NATIVE_PROTOCOL_VERSION, + OmpCommandName, OmpEvent, OmpEventType, OmpState, ) -from tests.integrations.omppeer import COMPLETE +from tests.integrations.omppeer import COMPLETE, LATE_PROMPT_ERROR from tests.integrations.omprpc.kit import ( ClientFactory, Observer, @@ -94,6 +95,44 @@ async def test_rejected_prompt_does_not_wait_for_a_result_that_is_not_owed( assert (rejected.value.command, rejected.value.code) == ("prompt", "rejected") await second_prompt(client) + async def test_command_error_after_admission_keeps_the_process_usable( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory(LATE_PROMPT_ERROR) + process = owned_rpc_process(client) + observer = Observer() + running = prompt_task(client, observer) + try: + await observer.wait_for(OmpEventType.AGENT_END) + await client.request("release") + with pytest.raises(OmpRpcError) as failed: + await running + assert failed.value.command == OmpCommandName.PROMPT + assert process.returncode is None + await second_prompt(client) + assert owned_rpc_process(client) is process + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + + async def test_a_late_response_for_the_wrong_command_retires_the_transport( + self, client_factory: ClientFactory + ) -> None: + client = await client_factory(LATE_PROMPT_ERROR) + process = owned_rpc_process(client) + observer = Observer() + running = prompt_task(client, observer) + try: + await observer.wait_for(OmpEventType.AGENT_END) + with pytest.raises(OmpProtocolError): + await client.request("release", failed_command=OmpCommandName.GET_STATE) + with pytest.raises(OmpProtocolError): + await running + assert process.returncode is not None + finally: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + @pytest.mark.parametrize( ("scenario", "match"), [ diff --git a/tests/integrations/omprpc/test_transport_failures.py b/tests/integrations/omprpc/test_transport_failures.py index 4c70fff5c..cfc1065a9 100644 --- a/tests/integrations/omprpc/test_transport_failures.py +++ b/tests/integrations/omprpc/test_transport_failures.py @@ -6,14 +6,22 @@ import logging import pytest +from pydantic import JsonValue +from band.integrations.codex.winspawn import await_returncode from band.integrations.omprpc import OmpProcessError, OmpProtocolError +from band.integrations.ompwire import OmpEventType from tests.integrations.omppeer import FRAME_PAYLOAD_MARK from tests.integrations.omprpc.kit import ( ClientFactory, Observer, + ProcessGate, + descendant_client, + descendant_liveness, owned_rpc_process, prompt_task, + unrelated_gate_child, + wait_for_descendants, ) @@ -68,13 +76,54 @@ async def test_eof_fails_a_command_waiter_and_closes_the_transport( await client.close() assert process.returncode is not None + async def test_observed_root_exit_retires_an_inflight_prompt_and_owned_children( + self, process_gate: ProcessGate + ) -> None: + client = descendant_client(process_gate, stdio="inherit", detached=True) + observer = Observer() + running: asyncio.Task[None] | None = None + crashing: asyncio.Task[dict[str, JsonValue]] | None = None + async with unrelated_gate_child(process_gate) as unrelated: + try: + await client.start() + descendants = await wait_for_descendants(process_gate) + process = owned_rpc_process(client) + running = prompt_task(client, observer) + await observer.wait_for(OmpEventType.AGENT_END) + crashing = asyncio.create_task(client.request("crash")) + assert await await_returncode(process, float("inf")) + assert await descendant_liveness(descendants) == [True, True] + + with pytest.raises(OmpProcessError): + await client.request("get_state") + async with asyncio.timeout(5): + with pytest.raises(OmpProcessError): + await running + with pytest.raises(OmpProcessError): + await crashing + assert await descendant_liveness(descendants) == [False, False] + assert await descendant_liveness([unrelated]) == [True] + finally: + await client.close() + tasks = [task for task in (running, crashing) if task is not None] + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + async def test_stderr_is_drained_without_newlines( self, client_factory: ClientFactory, + caplog: pytest.LogCaptureFixture, ) -> None: client = await client_factory("basic") - async with asyncio.timeout(5): - assert await client.request("stderr") == {"drained": True} + with caplog.at_level(logging.DEBUG, logger="band.integrations.omprpc.client"): + async with asyncio.timeout(5): + assert await client.request("stderr") == {"drained": True} + leaked = any( + FRAME_PAYLOAD_MARK in record.getMessage() for record in caplog.records + ) + caplog.clear() + assert not leaked @pytest.mark.parametrize( "scenario", diff --git a/tests/integrations/omprpc/test_turn_settlement.py b/tests/integrations/omprpc/test_turn_settlement.py index 4fc9f8968..d9545c57c 100644 --- a/tests/integrations/omprpc/test_turn_settlement.py +++ b/tests/integrations/omprpc/test_turn_settlement.py @@ -18,6 +18,7 @@ from tests.integrations.omppeer import ( BACKGROUND_COMPLETE, RESUMED_PROVIDER_FAILED, + SETTLED_CALLBACK_SCENARIO, YIELDED, ) from tests.integrations.omprpc.kit import ( @@ -110,7 +111,7 @@ async def test_native_state_snapshot_and_subscribed_settlement_do_not_race( async def test_settlement_does_not_succeed_while_a_callback_is_still_running( self, client_factory: ClientFactory ) -> None: - client = await client_factory("settled-during-callback") + client = await client_factory(SETTLED_CALLBACK_SCENARIO) client.callback_exit_timeout_s = 0 process = owned_rpc_process(client) held = asyncio.Event() @@ -136,7 +137,7 @@ async def on_tool(_call: OmpHostToolCall) -> OmpToolResult: async def test_settlement_succeeds_after_the_running_callback_finishes( self, client_factory: ClientFactory ) -> None: - client = await client_factory("settled-during-callback") + client = await client_factory(SETTLED_CALLBACK_SCENARIO) started = asyncio.Event() release = asyncio.Event() diff --git a/tests/integrations/omprpc/test_ui_responses.py b/tests/integrations/omprpc/test_ui_responses.py new file mode 100644 index 000000000..eecdadbdd --- /dev/null +++ b/tests/integrations/omprpc/test_ui_responses.py @@ -0,0 +1,58 @@ +"""Preserve native UI answer variants through the public prompt callback.""" + +from __future__ import annotations + +import asyncio + +import pytest +from pydantic import JsonValue + +from band.integrations.ompwire import OmpFrameType, OmpUIRequest, OmpUIResponse +from tests.integrations.omprpc.kit import ( + ClientFactory, + Observer, + checkpoint_frames, + second_prompt, +) + + +class TestUIResponseVariants: + @pytest.mark.parametrize( + ("response", "answer"), + [ + pytest.param(OmpUIResponse(value=""), {"value": ""}, id="empty-text"), + pytest.param( + OmpUIResponse(confirmed=False), + {"confirmed": False}, + id="declined-confirmation", + ), + pytest.param(OmpUIResponse(), {"cancelled": True}, id="unanswered-dialog"), + ], + ) + async def test_callback_answer_reaches_the_peer_without_other_variants( + self, + client_factory: ClientFactory, + response: OmpUIResponse, + answer: dict[str, JsonValue], + ) -> None: + client = await client_factory("callbacks") + observer = Observer() + + async def on_ui(_request: OmpUIRequest) -> OmpUIResponse: + return response + + async with asyncio.timeout(5): + await client.prompt( + "first", + on_event=observer.on_event, + on_tool=observer.on_tool, + on_ui=on_ui, + ) + frames = await checkpoint_frames(client) + reply = next(frame for frame in frames if frame.get("id") == "approval-1") + assert reply == { + "type": OmpFrameType.UI_RESPONSE, + "id": "approval-1", + **answer, + } + await second_prompt(client) diff --git a/tests/toolkit/timecontrol.py b/tests/toolkit/timecontrol.py index 8e7fa4c6d..3212063cd 100644 --- a/tests/toolkit/timecontrol.py +++ b/tests/toolkit/timecontrol.py @@ -3,8 +3,10 @@ from __future__ import annotations import asyncio -from collections.abc import AsyncIterator -from contextlib import asynccontextmanager +import threading +from collections.abc import AsyncIterator, Awaitable, Coroutine +from contextlib import AsyncExitStack, asynccontextmanager +from typing import Any, TypeVar from looptime import LoopTimeEventLoop @@ -12,12 +14,17 @@ from tests.toolkit.clock import binary_clock_resolution from tests.toolkit.tasks import background +Result = TypeVar("Result") + +REAL_IO_STEP_S = 0.01 + class ControlledClock: """Keep time stationary until the test crosses a deliberate boundary.""" - def __init__(self, loop: LoopTimeEventLoop) -> None: + def __init__(self, loop: LoopTimeEventLoop, stationary: asyncio.Task[None]) -> None: self.loop = loop + self._stationary = stationary def advance(self, seconds: float) -> None: self.loop.setup_looptime( @@ -29,6 +36,48 @@ async def elapse(self, seconds: float) -> None: self.advance(seconds) await spin() + async def _pause(self) -> None: + self._stationary.cancel() + await asyncio.gather(self._stationary, return_exceptions=True) + + async def run_io(self, operation: Awaitable[Result]) -> Result: + """Let real subprocess readiness and retirement make timed progress.""" + # Edge-triggered I/O cannot block on a second infinite selector poll + # before the loop services the completion returned by the first one. + self.loop.setup_looptime( + start=self.loop.time(), + resolution=binary_clock_resolution(), + idle_step=REAL_IO_STEP_S, + ) + running = asyncio.ensure_future(operation) + release = threading.Event() + # Looptime follows real time while an executor operation is pending. + # Hold that mode until the whole I/O operation, including cleanup, ends. + waiting = asyncio.create_task(asyncio.to_thread(release.wait)) + try: + await self._pause() + return await running + finally: + try: + running.cancel() + await asyncio.gather(running, return_exceptions=True) + finally: + release.set() + await waiting + self._stationary = asyncio.create_task(keep_time_stationary()) + + @asynccontextmanager + async def background( + self, operation: Coroutine[Any, Any, Result] | asyncio.Task[Result] + ) -> AsyncIterator[asyncio.Task[Result]]: + """Join a native operation with real I/O progress on exceptional exits.""" + owner = AsyncExitStack() + try: + task = await owner.enter_async_context(background(operation)) + yield task + finally: + await self.run_io(owner.aclose()) + async def keep_time_stationary() -> None: # A ready task prevents looptime from jumping to a timer during socket I/O. @@ -40,5 +89,9 @@ async def keep_time_stationary() -> None: async def controlled_clock() -> AsyncIterator[ControlledClock]: loop = asyncio.get_running_loop() assert isinstance(loop, LoopTimeEventLoop) and loop.looptime_on - async with background(keep_time_stationary()): - yield ControlledClock(loop) + async with background(keep_time_stationary()) as stationary: + clock = ControlledClock(loop, stationary) + try: + yield clock + finally: + await clock._pause() From 0720167358abccceb65be410378e871783089c6d Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 19:48:44 +0300 Subject: [PATCH 30/35] fix: keep native RPC usable after a late prompt error Record root-exit as a transport failure, join callbacks that start during the wait, and stop treating an empty UI reply as cancelled. Tighten plugin and converter tests onto product names. Co-authored-by: Cursor --- CHANGELOG.md | 1 + src/band/integrations/ompprocess.py | 14 ++++---- src/band/integrations/omprpc/client.py | 9 +++++ src/band/integrations/omprpc/frames.py | 8 +++-- src/band/integrations/ompwire.py | 2 +- tests/converters/test_omp.py | 5 +-- tests/integrations/omppeer.py | 14 ++++++++ .../ompprocess/test_mcp_plugin.py | 11 +++++-- .../ompprocess/test_mcp_startup.py | 12 ++++--- .../ompprocess/test_tool_names.py | 33 ++++++++----------- .../omprpc/test_prompt_results.py | 22 ++++++++++++- .../integrations/omprpc/test_ui_responses.py | 9 ++++- 12 files changed, 98 insertions(+), 42 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3cb737b6a..8a0117a0c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ This changelog is automatically generated by [Release Please](https://github.com * Wait for Windows Job members' kernel-confirmed exit before releasing OMP or Codex subprocess ownership, preserving process handles across incomplete cleanup and preventing late child admission during termination. * Retire native OMP sessions on an observed root exit, preserve usable sessions after post-admission command errors, and keep callbacks admitted during settlement within the original exit deadline. * Keep native OMP stderr payloads out of SDK debug logs. +* Preserve false and empty native OMP UI answers without adding a cancellation flag, and ignore whitespace-only connection instructions. ## [6.0.1](https://github.com/band-ai/band-sdk-python/compare/band-sdk-v6.0.0...band-sdk-v6.0.1) (2026-10-20) diff --git a/src/band/integrations/ompprocess.py b/src/band/integrations/ompprocess.py index 9549a4834..3e3f73515 100644 --- a/src/band/integrations/ompprocess.py +++ b/src/band/integrations/ompprocess.py @@ -138,8 +138,9 @@ def _require_one_transport(self) -> Self: description: "Verify the host-configured MCP connections", handler: async (_args, ctx) => { await ready; - const systemPromptReady = !config.systemPrompt - || ctx.getSystemPrompt().some(block => block.includes(config.systemPrompt.trim())); + const trimmed = (config.systemPrompt || "").trim(); + const systemPromptReady = !trimmed + || ctx.getSystemPrompt().some(block => block.includes(trimmed)); writeFileSync(config.receipt, JSON.stringify({ systemPromptReady, servers: [...statuses.values()] }), { mode: 0o600 }); }, }); @@ -276,7 +277,8 @@ def omp_mcp_plugin( servers: Mapping[str, OmpMcpServer], *, workspace: str, system_prompt: str = "" ) -> Iterator[str | None]: """Project connection instructions and servers without changing native user configs.""" - if not servers and not system_prompt: + trimmed_prompt = system_prompt.strip() + if not servers and not trimmed_prompt: yield None return projected = { @@ -285,18 +287,18 @@ def omp_mcp_plugin( with tempfile.TemporaryDirectory(prefix="band-omp-") as temporary: root = Path(temporary) / OMP_PLUGIN_ROOT root.mkdir(mode=0o700) - if system_prompt: + if trimmed_prompt: rules = root / "rules" rules.mkdir(mode=0o700) _write_private( rules / f"band-connection-{uuid4().hex}.md", - f"---\nalwaysApply: true\n---\n{system_prompt}", + f"---\nalwaysApply: true\n---\n{trimmed_prompt}", ) startup = { "servers": [_plugin_server_name(name) for name in servers], "command": omp_mcp_ready_command(str(root)), "receipt": str(root / OMP_MCP_RECEIPT), - "systemPrompt": system_prompt, + "systemPrompt": trimmed_prompt, } _write_private( root / OMP_MCP_SERVERS_CONFIG, json.dumps({"mcpServers": projected}) diff --git a/src/band/integrations/omprpc/client.py b/src/band/integrations/omprpc/client.py index e9bcd951b..a30685e31 100644 --- a/src/band/integrations/omprpc/client.py +++ b/src/band/integrations/omprpc/client.py @@ -505,6 +505,8 @@ async def _dispatch_payload( result = OmpPromptResult.model_validate(payload) if turn is not None and result.id == turn.id and not turn.result.done(): turn.result.set_result(result) + else: + logger.debug("OMP RPC ignored an unmatched prompt_result") case OmpFrameType.SESSION_SETTLED: if turn is not None and ( turn.last_activity_sequence > 0 or turn.result.done() @@ -586,6 +588,13 @@ def _dispatch_response(self, response: OmpResponse, sequence: int) -> None: logger.debug("OMP RPC ignored an unmatched command response") return if pending.future.done(): + if ( + not response.success + and self._turn is not None + and response.id == self._turn.id + ): + self._turn.rejected = True + self._turn.fail(command_error(response)) return if response.success: pending.future.set_result(ReceivedResponse(response, sequence)) diff --git a/src/band/integrations/omprpc/frames.py b/src/band/integrations/omprpc/frames.py index 7cfb3d25c..b1e70b6a3 100644 --- a/src/band/integrations/omprpc/frames.py +++ b/src/band/integrations/omprpc/frames.py @@ -121,9 +121,13 @@ def host_tool_result_frame( def ui_response_frame(request_id: str, response: OmpUIResponse) -> dict[str, JsonValue]: - answer = response.model_dump(mode="json", exclude={"cancelled"}, exclude_none=True) + answer = response.model_dump( + mode="json", + exclude_none=True, + exclude=None if response.cancelled else {"cancelled"}, + ) return { "type": OmpFrameType.UI_RESPONSE, "id": request_id, - **(answer or {"cancelled": True}), + **answer, } diff --git a/src/band/integrations/ompwire.py b/src/band/integrations/ompwire.py index f8624e328..e9fe80256 100644 --- a/src/band/integrations/ompwire.py +++ b/src/band/integrations/ompwire.py @@ -263,7 +263,7 @@ class OmpUIRequest(OmpWireModel): class OmpUIResponse(OmpWireModel): - """Answer one dialog, or cancel it when no answer variant is supplied.""" + """Answer or explicitly cancel one dialog.""" value: str | None = None confirmed: bool | None = None diff --git a/tests/converters/test_omp.py b/tests/converters/test_omp.py index ef46b0c42..f59b1c758 100644 --- a/tests/converters/test_omp.py +++ b/tests/converters/test_omp.py @@ -2,10 +2,7 @@ from __future__ import annotations -from band.converters.omp import ( - OmpHistoryConverter, - OmpMetadataKey, -) +from band.converters.omp import OmpHistoryConverter, OmpMetadataKey from band.core.types import MessageType diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index 906e635b8..f8e3974aa 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -40,6 +40,7 @@ OVERSIZE_AFTER_READY = "oversize-after-ready" READY_LIMIT_BYTES = 1024 LATE_PROMPT_ERROR = "late-prompt-error" +ADMITTED_PROMPT_ERROR = "late-error" SETTLED_CALLBACK_SCENARIO = "settled-during-callback" SETTLED_CALLBACK_ID = "host-held" LATE_CALLBACK_ID = "host-late" @@ -321,6 +322,19 @@ def prompt(self, frame: Frame) -> None: } ) return + case scenario if scenario == ADMITTED_PROMPT_ERROR: + self.respond(frame) + send( + { + "type": "response", + "id": frame.id, + "command": OmpCommandName.PROMPT, + "success": False, + "error": "Admitted prompt failed later", + "code": "post_admit", + } + ) + return case "local": self.settled = True self.respond(frame, {"agentInvoked": False}) diff --git a/tests/integrations/ompprocess/test_mcp_plugin.py b/tests/integrations/ompprocess/test_mcp_plugin.py index f7e826afd..f800d03af 100644 --- a/tests/integrations/ompprocess/test_mcp_plugin.py +++ b/tests/integrations/ompprocess/test_mcp_plugin.py @@ -29,8 +29,13 @@ def plugin_privacy(root: Path, files: tuple[Path, ...]) -> tuple[int, int | None class TestMcpPluginProjection: - def test_nothing_to_project_yields_no_plugin(self, tmp_path: Path) -> None: - with omp_mcp_plugin({}, workspace=str(tmp_path)) as plugin: + @pytest.mark.parametrize("system_prompt", ["", " "]) + def test_nothing_to_project_yields_no_plugin( + self, tmp_path: Path, system_prompt: str + ) -> None: + with omp_mcp_plugin( + {}, workspace=str(tmp_path), system_prompt=system_prompt + ) as plugin: assert plugin is None def test_servers_land_in_the_plugin_mcp_config(self, tmp_path: Path) -> None: @@ -43,6 +48,7 @@ def test_servers_land_in_the_plugin_mcp_config(self, tmp_path: Path) -> None: with omp_mcp_plugin(servers, workspace=str(tmp_path)) as plugin: assert plugin is not None root = Path(plugin) + config = json.loads( (root / OMP_MCP_SERVERS_CONFIG).read_text(encoding="utf-8") ) @@ -58,6 +64,7 @@ def test_servers_land_in_the_plugin_mcp_config(self, tmp_path: Path) -> None: assert local["args"] == ["--port", "1"] assert Path(local["cwd"]) == tmp_path.resolve() assert Path(local["command"]) == (tmp_path / "server").resolve() + assert not Path(plugin).exists() def test_connection_instructions_become_an_always_applied_rule( diff --git a/tests/integrations/ompprocess/test_mcp_startup.py b/tests/integrations/ompprocess/test_mcp_startup.py index 01f34b39c..df5f6ae1e 100644 --- a/tests/integrations/ompprocess/test_mcp_startup.py +++ b/tests/integrations/ompprocess/test_mcp_startup.py @@ -30,6 +30,8 @@ OmpUIResponse, ) +REQUIRED_SERVER = f"{OMP_PLUGIN_ROOT}:required" + class TestRequiredMcpStartup: @pytest.mark.parametrize( @@ -47,11 +49,11 @@ class TestRequiredMcpStartup: "servers": [ { "type": McpConnectionStatus.CONNECTED, - "serverName": f"{OMP_PLUGIN_ROOT}:required", + "serverName": REQUIRED_SERVER, }, { "type": McpConnectionStatus.CONNECTED, - "serverName": f"{OMP_PLUGIN_ROOT}:required", + "serverName": REQUIRED_SERVER, }, ], }, @@ -75,7 +77,7 @@ class TestRequiredMcpStartup: "servers": [ { "type": McpConnectionStatus.FAILED, - "serverName": f"{OMP_PLUGIN_ROOT}:required", + "serverName": REQUIRED_SERVER, } ], }, @@ -87,7 +89,7 @@ class TestRequiredMcpStartup: "servers": [ { "type": McpConnectionStatus.FAILED, - "serverName": f"{OMP_PLUGIN_ROOT}:required", + "serverName": REQUIRED_SERVER, "error": "initialize rejected", } ], @@ -100,7 +102,7 @@ class TestRequiredMcpStartup: "servers": [ { "type": McpConnectionStatus.CONNECTED, - "serverName": f"{OMP_PLUGIN_ROOT}:required", + "serverName": REQUIRED_SERVER, } ], }, diff --git a/tests/integrations/ompprocess/test_tool_names.py b/tests/integrations/ompprocess/test_tool_names.py index 1ea17aeeb..ddddc7f5e 100644 --- a/tests/integrations/ompprocess/test_tool_names.py +++ b/tests/integrations/ompprocess/test_tool_names.py @@ -11,46 +11,39 @@ normalize_native_mcp_tool_name, ) -OWN = frozenset({"band_send_message"}) +OWN_TOOL = "band_send_message" +OWN = frozenset({OWN_TOOL}) +BARE_TITLE = f"{XD_MCP_PREFIX}{OWN_TOOL}" +NAMESPACED_TITLE = f"{XD_MCP_PREFIX}{OMP_PLUGIN_TOOL_NAMESPACE}band_{OWN_TOOL}" +BARE_PATH = f"{XD_URL_PREFIX}{BARE_TITLE}" +NAMESPACED_PATH = f"{XD_URL_PREFIX}{NAMESPACED_TITLE}" class TestPluginNamespacedToolNames: @pytest.mark.parametrize( "path", [ - f"{XD_URL_PREFIX}{XD_MCP_PREFIX}band_send_message", - ( - f"{XD_URL_PREFIX}{XD_MCP_PREFIX}" - f"{OMP_PLUGIN_TOOL_NAMESPACE}band_band_send_message" - ), + BARE_PATH, + NAMESPACED_PATH, ], ) def test_executed_device_write_maps_to_the_band_tool(self, path: str) -> None: arguments = {"path": path, "content": '{"chat_id":"room-1","content":"hi"}'} assert normalize_native_mcp_device_call( "write", arguments, OWN, kind="execute" - ) == ("band_send_message", {"chat_id": "room-1", "content": "hi"}) + ) == (OWN_TOOL, {"chat_id": "room-1", "content": "hi"}) - @pytest.mark.parametrize( - "title", - [ - f"{XD_MCP_PREFIX}band_send_message", - f"{XD_MCP_PREFIX}{OMP_PLUGIN_TOOL_NAMESPACE}band_band_send_message", - ], - ) + @pytest.mark.parametrize("title", [BARE_TITLE, NAMESPACED_TITLE]) def test_mcp_title_maps_to_the_band_tool(self, title: str) -> None: - assert normalize_native_mcp_tool_name(title, OWN) == "band_send_message" + assert normalize_native_mcp_tool_name(title, OWN) == OWN_TOOL assert normalize_native_mcp_device_call(title, {"chat_id": "r1"}, OWN) == ( - "band_send_message", + OWN_TOOL, {"chat_id": "r1"}, ) def test_a_namespaced_read_is_never_the_band_tool(self) -> None: arguments = { - "path": ( - f"{XD_URL_PREFIX}{XD_MCP_PREFIX}" - f"{OMP_PLUGIN_TOOL_NAMESPACE}band_band_send_message" - ), + "path": NAMESPACED_PATH, "content": '{"content":"hi"}', } assert normalize_native_mcp_device_call( diff --git a/tests/integrations/omprpc/test_prompt_results.py b/tests/integrations/omprpc/test_prompt_results.py index 1c0a8561d..f6ef6399c 100644 --- a/tests/integrations/omprpc/test_prompt_results.py +++ b/tests/integrations/omprpc/test_prompt_results.py @@ -14,7 +14,11 @@ OmpEventType, OmpState, ) -from tests.integrations.omppeer import COMPLETE, LATE_PROMPT_ERROR +from tests.integrations.omppeer import ( + ADMITTED_PROMPT_ERROR, + COMPLETE, + LATE_PROMPT_ERROR, +) from tests.integrations.omprpc.kit import ( ClientFactory, Observer, @@ -133,6 +137,22 @@ async def test_a_late_response_for_the_wrong_command_retires_the_transport( running.cancel() await asyncio.gather(running, return_exceptions=True) + async def test_a_late_prompt_error_does_not_reap_the_process( + self, + client_factory: ClientFactory, + ) -> None: + client = await client_factory(ADMITTED_PROMPT_ERROR) + process = owned_rpc_process(client) + async with asyncio.timeout(5): + with pytest.raises(OmpRpcError) as late: + await prompt_task(client, Observer()) + assert (late.value.command, late.value.code) == ( + OmpCommandName.PROMPT, + "post_admit", + ) + assert process.returncode is None + await second_prompt(client) + @pytest.mark.parametrize( ("scenario", "match"), [ diff --git a/tests/integrations/omprpc/test_ui_responses.py b/tests/integrations/omprpc/test_ui_responses.py index eecdadbdd..2def5fd95 100644 --- a/tests/integrations/omprpc/test_ui_responses.py +++ b/tests/integrations/omprpc/test_ui_responses.py @@ -26,7 +26,14 @@ class TestUIResponseVariants: {"confirmed": False}, id="declined-confirmation", ), - pytest.param(OmpUIResponse(), {"cancelled": True}, id="unanswered-dialog"), + pytest.param( + OmpUIResponse(cancelled=False), {}, id="unanswered-not-cancelled" + ), + pytest.param( + OmpUIResponse(cancelled=True), + {"cancelled": True}, + id="explicit-cancellation", + ), ], ) async def test_callback_answer_reaches_the_peer_without_other_variants( From c4ae7d844cc981205474ae6135ec5c5ce7505f3f Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 23:22:47 +0300 Subject: [PATCH 31/35] test: initialize callback retirement clock before native startup --- .../omprpc/test_callback_retirement.py | 142 +++++++++--------- 1 file changed, 70 insertions(+), 72 deletions(-) diff --git a/tests/integrations/omprpc/test_callback_retirement.py b/tests/integrations/omprpc/test_callback_retirement.py index d0cfb6c07..04dac3164 100644 --- a/tests/integrations/omprpc/test_callback_retirement.py +++ b/tests/integrations/omprpc/test_callback_retirement.py @@ -39,83 +39,81 @@ class TestCallbackRetirement: async def test_a_callback_admitted_during_the_join_cannot_escape_the_deadline( self, client_factory: ClientFactory, retire_second_on_first: bool ) -> None: - client = await client_factory(SETTLED_CALLBACK_SCENARIO) - grace = 1.0 - client.callback_exit_timeout_s = grace - process = owned_rpc_process(client) - observer = Observer() - cleanup = HeldCleanup(client) - joining: asyncio.Task[None] | None = None - first_started, first_retired, second_started = ( - asyncio.Event() for _ in range(3) - ) - release_first, release_second = (asyncio.Event() for _ in range(2)) - second_cancelled = asyncio.Event() + async with controlled_clock() as clock: + client = await clock.run_io(client_factory(SETTLED_CALLBACK_SCENARIO)) + grace = 1.0 + client.callback_exit_timeout_s = grace + process = owned_rpc_process(client) + observer = Observer() + cleanup = HeldCleanup(client) + joining: asyncio.Task[None] | None = None + first_started, first_retired, second_started = ( + asyncio.Event() for _ in range(3) + ) + release_first, release_second = (asyncio.Event() for _ in range(2)) + second_cancelled = asyncio.Event() - async def on_tool(call: OmpHostToolCall) -> OmpToolResult: - if call.id == SETTLED_CALLBACK_ID: - task = asyncio.current_task() - assert task is not None - task.add_done_callback(lambda _done: first_retired.set()) - first_started.set() - await release_first.wait() - else: - assert call.id == LATE_CALLBACK_ID - second_started.set() - try: - await ( - first_retired.wait() - if retire_second_on_first - else release_second.wait() - ) - except asyncio.CancelledError: - second_cancelled.set() - raise - return OmpToolResult(content=[]) + async def on_tool(call: OmpHostToolCall) -> OmpToolResult: + if call.id == SETTLED_CALLBACK_ID: + task = asyncio.current_task() + assert task is not None + task.add_done_callback(lambda _done: first_retired.set()) + first_started.set() + await release_first.wait() + else: + assert call.id == LATE_CALLBACK_ID + second_started.set() + try: + await ( + first_retired.wait() + if retire_second_on_first + else release_second.wait() + ) + except asyncio.CancelledError: + second_cancelled.set() + raise + return OmpToolResult(content=[]) - async with ( - controlled_clock() as clock, - clock.background( + async with clock.background( client.prompt( "first", on_event=observer.on_event, on_tool=on_tool, on_ui=observer.on_ui, ) - ) as running, - ): - try: - await first_started.wait() - await observer.wait_for(OmpEventType.AGENT_END) - frames = await checkpoint_frames(client) - assert any( - frame["type"] == OmpCommandName.GET_STATE for frame in frames - ) - await client.request("release") - await second_started.wait() - if retire_second_on_first: - clock.advance(grace) - release_first.set() - await first_retired.wait() - if not retire_second_on_first: - clock.advance(grace) - joining = asyncio.create_task(cleanup.wait_for_join(running)) - done, _ = await asyncio.wait( - (running, joining), return_when=asyncio.FIRST_COMPLETED - ) - assert joining in done - await joining - assert process.returncode is None - cleanup.release.set() - with pytest.raises(OmpProcessError): - await clock.run_io(running) - if not retire_second_on_first: - assert second_cancelled.is_set() - assert process.returncode is not None - finally: - release_first.set() - release_second.set() - cleanup.release.set() - if joining is not None: - joining.cancel() - await asyncio.gather(joining, return_exceptions=True) + ) as running: + try: + await first_started.wait() + await observer.wait_for(OmpEventType.AGENT_END) + frames = await checkpoint_frames(client) + assert any( + frame["type"] == OmpCommandName.GET_STATE for frame in frames + ) + await client.request("release") + await second_started.wait() + if retire_second_on_first: + clock.advance(grace) + release_first.set() + await first_retired.wait() + if not retire_second_on_first: + clock.advance(grace) + joining = asyncio.create_task(cleanup.wait_for_join(running)) + done, _ = await asyncio.wait( + (running, joining), return_when=asyncio.FIRST_COMPLETED + ) + assert joining in done + await joining + assert process.returncode is None + cleanup.release.set() + with pytest.raises(OmpProcessError): + await clock.run_io(running) + if not retire_second_on_first: + assert second_cancelled.is_set() + assert process.returncode is not None + finally: + release_first.set() + release_second.set() + cleanup.release.set() + if joining is not None: + joining.cancel() + await asyncio.gather(joining, return_exceptions=True) From a65aea233c3a7971935e6c6cfea38fd4b5e966cd Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 21:35:43 +0300 Subject: [PATCH 32/35] fix: keep Windows native RPC tests on the owned job spawn create_process wrappers that used create_subprocess_exec skipped the Windows job, so generated restart sequences could not reap descendants. Co-authored-by: Cursor --- docs/acp.md | 4 ++ src/band/integrations/omprpc/client.py | 51 ++++++++++++------- tests/integrations/omprpc/kit.py | 39 ++++++++++++++ tests/integrations/omprpc/lifecycle.py | 35 ++++++------- .../omprpc/test_eof_descendants.py | 3 +- .../omprpc/test_injected_spawn.py | 33 ++++++++++++ 6 files changed, 128 insertions(+), 37 deletions(-) diff --git a/docs/acp.md b/docs/acp.md index 3073f256d..3b120f95f 100644 --- a/docs/acp.md +++ b/docs/acp.md @@ -309,6 +309,10 @@ A `create_process` factory customizes creation without bypassing containment. On Windows it launches the existing helper; the client still owns the Job containing the native process and its descendants, including during cancelled or incomplete startup. +To hold or observe the owned launch, wrap `spawn_owned_process` and delegate to +its bound original (or override it and call `super()`). Do not pass that wrapper +as `create_process`: that factory runs below the ownership boundary. The owned +method retains the process before returning, so wrapper failure still reaps it. On POSIX, process-group capture rechecks a reaped leader's identity after enumerating members, before retaining that group. A reused leader PID does not diff --git a/src/band/integrations/omprpc/client.py b/src/band/integrations/omprpc/client.py index a30685e31..323896452 100644 --- a/src/band/integrations/omprpc/client.py +++ b/src/band/integrations/omprpc/client.py @@ -8,7 +8,7 @@ from binascii import Error as Base64Error from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass -from typing import Any, TypeVar +from typing import Any, TypeVar, cast from pydantic import JsonValue, ValidationError @@ -193,32 +193,45 @@ async def start(self) -> None: await self.close() raise - async def _spawn_process(self) -> asyncio.subprocess.Process: - scope = self._process_scope - assert scope is not None - env = scope.environment(self.env) - # Windows has no process group. Descendants inherit job membership at - # creation and cannot break away. + async def spawn_owned_process( + self, *command: str, **kwargs: Any + ) -> asyncio.subprocess.Process: + """Start and retain the command within this client's process containment. + + ``create_process`` is the low-level launcher, including for the Windows + helper. Wrappers override this method and delegate to its bound original + or ``super()`` so creation owns exactly one helper and Job. + """ if os.name == "nt": self._windows_process = WindowsProcess(create_process=self.create_process) process = await self._windows_process.start_command( - self.command, cwd=None, env=env, limit=OMP_FRAME_BYTES + command, + cwd=cast(str | None, kwargs.get("cwd")), + env=cast(Mapping[str, str] | None, kwargs.get("env")), + limit=cast(int, kwargs.get("limit", OMP_FRAME_BYTES)), ) else: create = self.create_process or asyncio.create_subprocess_exec - process = await create( - *self.command, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=None, - env=env, - start_new_session=os.name != "nt", - limit=OMP_FRAME_BYTES, - ) + process = await create(*command, **kwargs) + # A wrapper may hold or fail after creation; cleanup still owns the child. self._proc = process return process + async def _spawn_process(self) -> asyncio.subprocess.Process: + scope = self._process_scope + assert scope is not None + env = scope.environment(self.env) + return await self.spawn_owned_process( + *self.command, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=None, + env=env, + start_new_session=os.name != "nt", + limit=OMP_FRAME_BYTES, + ) + async def request( self, command: str, **parameters: JsonValue ) -> dict[str, JsonValue]: @@ -748,7 +761,7 @@ async def _settle_spawn(self) -> asyncio.subprocess.Process | None: return self._proc try: return await asyncio.shield(spawning) - except (OSError, SpawnFailure): + except Exception: # noqa: BLE001 -- startup propagates the error; cleanup retains its child. return self._proc except asyncio.CancelledError: if not spawning.cancelled(): diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py index d7ba38f8c..4ff9d83e9 100644 --- a/tests/integrations/omprpc/kit.py +++ b/tests/integrations/omprpc/kit.py @@ -202,6 +202,33 @@ def process_is_alive(process: psutil.Process) -> bool: return False +def is_descendant(child: psutil.Process, ancestor: psutil.Process) -> bool: + """True when ``ancestor`` is a living parent of ``child``. + + Windows Python launchers insert an extra process, so the gate child's + ``parent_pid`` is not always the owned helper. + """ + try: + ancestor_pid = ancestor.pid + ancestor_created = ancestor.create_time() + current = child.parent() + except (psutil.NoSuchProcess, psutil.AccessDenied): + return False + seen: set[int] = set() + while current is not None: + try: + pid = current.pid + if pid == ancestor_pid and current.create_time() == ancestor_created: + return True + if pid in seen: + return False + seen.add(pid) + current = current.parent() + except (psutil.NoSuchProcess, psutil.AccessDenied): + return False + return False + + def kill_gate_processes(processes: list[psutil.Process]) -> None: for process in processes: with suppress(psutil.NoSuchProcess): @@ -298,6 +325,18 @@ async def wait_for_descendants(gate: ProcessGate, count: int = 2) -> list[GatedP return [await gate.connection() for _ in range(count)] +async def descendants_started_with( + gate: ProcessGate, root: asyncio.subprocess.Process, count: int = 2 +) -> list[GatedProcess]: + created = await asyncio.to_thread(psutil.Process(root.pid).create_time) + children: list[GatedProcess] = [] + while len(children) < count: + child = await gate.connection() + if child.receipt.created >= created: + children.append(child) + return children + + async def crash_rpc_root(client: OmpRpcClient) -> None: root = owned_rpc_process(client) crashing = asyncio.create_task(client.request("crash")) diff --git a/tests/integrations/omprpc/lifecycle.py b/tests/integrations/omprpc/lifecycle.py index 79356c098..1f2fb0cad 100644 --- a/tests/integrations/omprpc/lifecycle.py +++ b/tests/integrations/omprpc/lifecycle.py @@ -14,6 +14,7 @@ from contextlib import asynccontextmanager, suppress from dataclasses import dataclass from enum import StrEnum +from unittest.mock import patch from hypothesis import strategies as st @@ -27,6 +28,7 @@ SpawnBarrier, descendant_client, descendant_liveness, + descendants_started_with, owned_rpc_process, process_gate_session, prompt_task, @@ -86,7 +88,7 @@ def __init__(self, gate: ProcessGate, *, detached: bool) -> None: self.children: list[GatedProcess] = [] async def record_children(self) -> list[GatedProcess]: - children = [await self.gate.connection(), await self.gate.connection()] + children = await descendants_started_with(self.gate, self.roots[-1]) self.children.extend(children) return children @@ -206,22 +208,21 @@ async def _claim_restart( @asynccontextmanager async def held_first_spawn(self) -> AsyncIterator[HeldSpawn]: await self.close() - barrier = SpawnBarrier(asyncio.create_subprocess_exec) - self.client.create_process = barrier.create - starting = self._track(asyncio.create_task(self.client.start())) - try: - await barrier.created.wait() - assert barrier.process is not None - self.roots.append(barrier.process) - yield HeldSpawn( - starting=starting, - root=barrier.process, - children=await self.record_children(), - barrier=barrier, - ) - finally: - self.client.create_process = None - barrier.release.set() + barrier = SpawnBarrier(self.client.spawn_owned_process) + with patch.object(self.client, "spawn_owned_process", barrier.create): + starting = self._track(asyncio.create_task(self.client.start())) + try: + await barrier.created.wait() + assert barrier.process is not None + self.roots.append(barrier.process) + yield HeldSpawn( + starting=starting, + root=barrier.process, + children=await self.record_children(), + barrier=barrier, + ) + finally: + barrier.release.set() async def cancel_held_start_then_restart(self) -> None: async with self.held_first_spawn() as spawn: diff --git a/tests/integrations/omprpc/test_eof_descendants.py b/tests/integrations/omprpc/test_eof_descendants.py index c8f6dcae5..4d7072eca 100644 --- a/tests/integrations/omprpc/test_eof_descendants.py +++ b/tests/integrations/omprpc/test_eof_descendants.py @@ -10,6 +10,7 @@ ProcessGate, descendant_liveness, hold_after_stdin_closed, + is_descendant, job_member_pids, process_is_alive, session_ids, @@ -32,7 +33,7 @@ async def test_close_reaps_a_detached_child_created_after_stdin_eof( async with hold_after_stdin_closed(client) as closing: late_child = await process_gate.connection() assert root.returncode is not None - assert late_child.receipt.parent_pid == helper.process.pid + assert is_descendant(late_child.process, helper.process) assert await descendant_liveness([helper, late_child]) == [ True, True, diff --git a/tests/integrations/omprpc/test_injected_spawn.py b/tests/integrations/omprpc/test_injected_spawn.py index 98e43b454..e28da13ae 100644 --- a/tests/integrations/omprpc/test_injected_spawn.py +++ b/tests/integrations/omprpc/test_injected_spawn.py @@ -4,6 +4,7 @@ import asyncio import sys +from typing import Any import pytest @@ -18,6 +19,18 @@ ) +class FailingOwnedClient(OmpRpcClient): + fail_spawn = True + + async def spawn_owned_process( + self, *command: str, **kwargs: Any + ) -> asyncio.subprocess.Process: + process = await super().spawn_owned_process(*command, **kwargs) + if self.fail_spawn: + raise ValueError("Owned spawn wrapper failed") + return process + + class TestInjectedSpawnOwnership: async def test_stubborn_injected_process_is_reaped_before_restart(self) -> None: client = OmpRpcClient( @@ -49,3 +62,23 @@ async def test_stubborn_injected_process_is_reaped_before_restart(self) -> None: running.cancel() await asyncio.gather(running, return_exceptions=True) await client.close() + + async def test_failed_owned_wrapper_reaps_before_restart(self) -> None: + client = FailingOwnedClient( + (sys.executable, "-u", OMP_PEER_SCRIPT, "--scenario", "basic"), + create_process=asyncio.create_subprocess_exec, + ) + process: asyncio.subprocess.Process | None = None + try: + with pytest.raises(ValueError, match="Owned spawn wrapper failed"): + await client.start() + process = owned_rpc_process(client) + assert process.returncode is not None + client.fail_spawn = False + await client.start() + await second_prompt(client) + finally: + if process is not None and process.returncode is None: + process.kill() + await process.wait() + await client.close() From 7b4857a1d71d1f8730a66f63b79f394ff4a79023 Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 22:14:01 +0300 Subject: [PATCH 33/35] refactor: type native RPC command names and host-tool results Keep the wire vocabulary in one enum and send commands through the same frame encoder as host-tool replies. Co-authored-by: Cursor --- src/band/integrations/ompprocess.py | 2 +- src/band/integrations/omprpc/client.py | 8 +++++--- src/band/integrations/omprpc/frames.py | 16 --------------- src/band/integrations/ompwire.py | 20 ++++++++++++++++++- tests/integrations/omppeer.py | 12 +++++------ tests/integrations/omprpc/kit.py | 3 ++- .../omprpc/test_prompt_results.py | 2 +- 7 files changed, 34 insertions(+), 29 deletions(-) diff --git a/src/band/integrations/ompprocess.py b/src/band/integrations/ompprocess.py index 3e3f73515..eb0ed34ae 100644 --- a/src/band/integrations/ompprocess.py +++ b/src/band/integrations/ompprocess.py @@ -106,7 +106,7 @@ def _require_one_transport(self) -> Self: raise ValueError("MCP server needs exactly one of command or url") if self.command and self.type not in (None, OmpMcpTransport.STDIO): raise ValueError(f"a command requires {OmpMcpTransport.STDIO} transport") - if self.url and self.type == OmpMcpTransport.STDIO: + if self.url and self.type is OmpMcpTransport.STDIO: raise ValueError( f"a URL requires {OmpMcpTransport.HTTP} or {OmpMcpTransport.SSE} transport" ) diff --git a/src/band/integrations/omprpc/client.py b/src/band/integrations/omprpc/client.py index 323896452..dec072fab 100644 --- a/src/band/integrations/omprpc/client.py +++ b/src/band/integrations/omprpc/client.py @@ -30,7 +30,6 @@ ) from band.integrations.omprpc.frames import ( ChunkAssembler, - host_tool_result_frame, parse_payload_line, payload_json, ui_response_frame, @@ -50,6 +49,7 @@ OmpFrameType, OmpHostToolCall, OmpHostToolCancel, + OmpHostToolResult, OmpNegotiation, OmpPromptAcceptance, OmpPromptResult, @@ -300,7 +300,7 @@ async def _request_frame( raise ValueError("OMP command parameters must not override id or type") request_id = request_id or self._next_id() frame = OmpCommand.model_validate( - {"type": command, "id": request_id, **parameters} + {"type": str(command), "id": request_id, **parameters} ) future: asyncio.Future[ReceivedResponse | OmpRpcError | OmpProcessError] = ( asyncio.get_running_loop().create_future() @@ -663,7 +663,9 @@ async def _handle_tool(self, turn: PromptRun, call: OmpHostToolCall) -> None: result = await turn.on_tool(call) if not self._callback_still_active(turn, OmpCallbackKind.TOOL, call.id): return - await self._write_frame(host_tool_result_frame(call.id, result)) + await self._write_frame( + OmpHostToolResult.from_tool(call.id, result).model_dump(mode="json") + ) async def _handle_ui(self, turn: PromptRun, request: OmpUIRequest) -> None: response = await turn.on_ui(request) diff --git a/src/band/integrations/omprpc/frames.py b/src/band/integrations/omprpc/frames.py index b1e70b6a3..dbcd71dad 100644 --- a/src/band/integrations/omprpc/frames.py +++ b/src/band/integrations/omprpc/frames.py @@ -4,7 +4,6 @@ import base64 from dataclasses import dataclass -from typing import cast from pydantic import JsonValue, TypeAdapter @@ -15,7 +14,6 @@ OmpChunk, OmpFrame, OmpFrameType, - OmpToolResult, OmpUIResponse, ) @@ -106,20 +104,6 @@ def parse_payload_line(line: bytes) -> dict[str, JsonValue]: return _PAYLOAD_ADAPTER.validate_json(line, strict=True) -def host_tool_result_frame( - request_id: str, result: OmpToolResult -) -> dict[str, JsonValue]: - return { - "type": OmpFrameType.HOST_TOOL_RESULT, - "id": request_id, - "result": cast( - dict[str, JsonValue], - result.model_dump(mode="json", exclude={"is_error"}, exclude_none=True), - ), - "isError": result.is_error, - } - - def ui_response_frame(request_id: str, response: OmpUIResponse) -> dict[str, JsonValue]: answer = response.model_dump( mode="json", diff --git a/src/band/integrations/ompwire.py b/src/band/integrations/ompwire.py index e9fe80256..a1c14f4eb 100644 --- a/src/band/integrations/ompwire.py +++ b/src/band/integrations/ompwire.py @@ -3,7 +3,7 @@ from __future__ import annotations from enum import StrEnum -from typing import Final, Literal +from typing import Final, Literal, Self, cast from pydantic import ( AliasChoices, @@ -225,6 +225,24 @@ class OmpHostToolCancel(OmpWireModel): targetId: str +class OmpHostToolResult(OmpWireModel): + type: Literal[OmpFrameType.HOST_TOOL_RESULT] = OmpFrameType.HOST_TOOL_RESULT + id: str + result: dict[str, JsonValue] + isError: bool = False + + @classmethod + def from_tool(cls, call_id: str, tool: OmpToolResult) -> Self: + return cls( + id=call_id, + result=cast( + dict[str, JsonValue], + tool.model_dump(mode="json", exclude={"is_error"}, exclude_none=True), + ), + isError=tool.is_error, + ) + + class OmpAskOption(OmpWireModel): label: str description: str | None = None diff --git a/tests/integrations/omppeer.py b/tests/integrations/omppeer.py index f8e3974aa..76c797ee7 100644 --- a/tests/integrations/omppeer.py +++ b/tests/integrations/omppeer.py @@ -315,7 +315,7 @@ def prompt(self, frame: Frame) -> None: { "type": "response", "id": frame.id, - "command": "prompt", + "command": OmpCommandName.PROMPT, "success": False, "error": "Prompt rejected", "code": "rejected", @@ -359,7 +359,7 @@ def prompt(self, frame: Frame) -> None: { "type": "response", "id": frame.id, - "command": "get_state", + "command": OmpCommandName.GET_STATE, "success": True, "data": {}, } @@ -544,7 +544,7 @@ def handle(self, frame: Frame) -> None: cast(dict[str, JsonValue], frame.model_dump(exclude_unset=True)) ) match frame.type: - case "negotiate_protocol": + case OmpCommandName.NEGOTIATE_PROTOCOL: if self.scenario == "negotiation-rejected": send( { @@ -565,9 +565,9 @@ def handle(self, frame: Frame) -> None: }, ) PROTOCOL_VERSION = OMP_NATIVE_PROTOCOL_VERSION - case "prompt": + case OmpCommandName.PROMPT: self.prompt(frame) - case "get_state": + case OmpCommandName.GET_STATE: self.state_requested = True self.respond( frame, @@ -785,7 +785,7 @@ def handle(self, frame: Frame) -> None: ) case "checkpoint": self.checkpoint(frame) - case "abort": + case OmpCommandName.ABORT: if self.scenario != "stubborn": self.respond(frame) if self.prompt_id is not None: diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py index 4ff9d83e9..7867fbec5 100644 --- a/tests/integrations/omprpc/kit.py +++ b/tests/integrations/omprpc/kit.py @@ -20,6 +20,7 @@ from band.integrations.codex.winspawn import Job from band.integrations.omprpc import OmpProcessError, OmpRpcClient from band.integrations.ompwire import ( + OmpCommandName, OmpContent, OmpEvent, OmpEventType, @@ -191,7 +192,7 @@ async def sent_prompt_messages(client: OmpRpcClient) -> list[str]: return [ cast(str, frame["message"]) for frame in await checkpoint_frames(client) - if frame["type"] == "prompt" + if frame["type"] == OmpCommandName.PROMPT ] diff --git a/tests/integrations/omprpc/test_prompt_results.py b/tests/integrations/omprpc/test_prompt_results.py index f6ef6399c..d1e43d34f 100644 --- a/tests/integrations/omprpc/test_prompt_results.py +++ b/tests/integrations/omprpc/test_prompt_results.py @@ -212,7 +212,7 @@ async def test_native_v2_preserves_large_inline_image_frames( negotiations = [ frame for frame in await checkpoint_frames(client) - if frame["type"] == "negotiate_protocol" + if frame["type"] == OmpCommandName.NEGOTIATE_PROTOCOL ] assert len(negotiations) == 1 assert negotiations[0]["protocolVersion"] == OMP_NATIVE_PROTOCOL_VERSION From b39b41ad6f24bfbacc99fcbba988db71d06cec6a Mon Sep 17 00:00:00 2001 From: Alexander Zaikman Date: Thu, 22 Oct 2026 22:34:42 +0300 Subject: [PATCH 34/35] fix: wait for owned descendants on the gate event, not a clock A five-second connection timeout was proving Windows children had arrived. Identify them as descendants of the current root instead. Co-authored-by: Cursor --- tests/integrations/omprpc/kit.py | 9 ++++----- tests/integrations/omprpc/test_lifecycle_model.py | 7 +++++++ 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py index 7867fbec5..55aa263d7 100644 --- a/tests/integrations/omprpc/kit.py +++ b/tests/integrations/omprpc/kit.py @@ -270,9 +270,8 @@ def accept( self.connections.put_nowait((reader, writer)) async def connection(self) -> GatedProcess: - async with asyncio.timeout(5): - reader, _ = await self.connections.get() - receipt = ProcessReceipt.model_validate_json(await reader.readline()) + reader, _ = await self.connections.get() + receipt = ProcessReceipt.model_validate_json(await reader.readline()) process = await asyncio.to_thread(psutil.Process, receipt.pid) self.processes.append(process) assert await asyncio.to_thread(process.create_time) == receipt.created @@ -329,11 +328,11 @@ async def wait_for_descendants(gate: ProcessGate, count: int = 2) -> list[GatedP async def descendants_started_with( gate: ProcessGate, root: asyncio.subprocess.Process, count: int = 2 ) -> list[GatedProcess]: - created = await asyncio.to_thread(psutil.Process(root.pid).create_time) + ancestor = await asyncio.to_thread(psutil.Process, root.pid) children: list[GatedProcess] = [] while len(children) < count: child = await gate.connection() - if child.receipt.created >= created: + if is_descendant(child.process, ancestor): children.append(child) return children diff --git a/tests/integrations/omprpc/test_lifecycle_model.py b/tests/integrations/omprpc/test_lifecycle_model.py index 937296b39..0a851c2eb 100644 --- a/tests/integrations/omprpc/test_lifecycle_model.py +++ b/tests/integrations/omprpc/test_lifecycle_model.py @@ -39,6 +39,13 @@ class TestRpcLifecycleModel: @example(steps=[Step.PROMPT, Step.CRASH, Step.PROMPT, Step.CANCEL_OWNER]) @example(steps=[Step.CANCEL_START_RESTART, Step.PROMPT, Step.CLOSE]) @example(steps=[Step.CANCEL_CLOSE_RESTART, Step.PROMPT, Step.SETTLE]) + @example( + steps=[ + Step.CANCEL_CLOSE_RESTART, + Step.CANCEL_OWNER, + Step.CANCEL_CLOSE_RESTART, + ] + ) @example(steps=[Step.PROMPT, Step.PROMPT, Step.CANCEL_OWNER, Step.CANCEL_OWNER]) @example( steps=[ From 0b0620e0b57727b1f774387f7c76a2abbbfd204d Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 23:46:37 +0300 Subject: [PATCH 35/35] test: keep owned launch doubles in the shared RPC kit --- tests/integrations/omprpc/kit.py | 14 +++++++++++++- tests/integrations/omprpc/test_injected_spawn.py | 16 ++-------------- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/tests/integrations/omprpc/kit.py b/tests/integrations/omprpc/kit.py index 55aa263d7..6512ddf62 100644 --- a/tests/integrations/omprpc/kit.py +++ b/tests/integrations/omprpc/kit.py @@ -11,7 +11,7 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Sequence from contextlib import asynccontextmanager, suppress from dataclasses import dataclass -from typing import cast +from typing import Any, cast import psutil import pytest @@ -39,6 +39,18 @@ OMP_PEER_SCRIPT = omppeer.__file__ +class FailingOwnedClient(OmpRpcClient): + fail_spawn = True + + async def spawn_owned_process( + self, *command: str, **kwargs: Any + ) -> asyncio.subprocess.Process: + process = await super().spawn_owned_process(*command, **kwargs) + if self.fail_spawn: + raise ValueError("Owned spawn wrapper failed") + return process + + class SpawnBarrier: """Create the first subprocess, then hold returning it until released.""" diff --git a/tests/integrations/omprpc/test_injected_spawn.py b/tests/integrations/omprpc/test_injected_spawn.py index e28da13ae..2831aea53 100644 --- a/tests/integrations/omprpc/test_injected_spawn.py +++ b/tests/integrations/omprpc/test_injected_spawn.py @@ -4,7 +4,6 @@ import asyncio import sys -from typing import Any import pytest @@ -12,6 +11,7 @@ from band.integrations.ompwire import OmpEventType from tests.integrations.omprpc.kit import ( OMP_PEER_SCRIPT, + FailingOwnedClient, Observer, owned_rpc_process, prompt_task, @@ -19,18 +19,6 @@ ) -class FailingOwnedClient(OmpRpcClient): - fail_spawn = True - - async def spawn_owned_process( - self, *command: str, **kwargs: Any - ) -> asyncio.subprocess.Process: - process = await super().spawn_owned_process(*command, **kwargs) - if self.fail_spawn: - raise ValueError("Owned spawn wrapper failed") - return process - - class TestInjectedSpawnOwnership: async def test_stubborn_injected_process_is_reaped_before_restart(self) -> None: client = OmpRpcClient( @@ -70,7 +58,7 @@ async def test_failed_owned_wrapper_reaps_before_restart(self) -> None: ) process: asyncio.subprocess.Process | None = None try: - with pytest.raises(ValueError, match="Owned spawn wrapper failed"): + with pytest.raises(ValueError): await client.start() process = owned_rpc_process(client) assert process.returncode is not None