From 338086b181dbaec30f43c2ba368d5d068339aa23 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Wed, 21 Oct 2026 16:17:54 +0300 Subject: [PATCH 1/3] feat(omp): resolve room-scoped MCP servers at native startup A host can now pass resolve_mcp_servers, an async callable from the claimed room workspace to the room's complete external MCP server set, called once per new native process before the overlay or process exists so the overlay and the startup receipt check share one snapshot. The configured Band endpoint is kept beside the resolved set and a resolver may not claim its name. Room retirement joins overlapping cleanup calls, cancels only the room's active turn, and keeps a room whose native close failed owned until a later close succeeds. --- docs/acp.md | 40 ++ src/band/adapters/omp.py | 218 ++++++++++- tests/adapters/omp/kit.py | 222 ++++++++++- tests/adapters/omp/test_cancellation_entry.py | 64 ++++ .../adapters/omp/test_deadline_settlement.py | 202 ++++++++++ tests/adapters/omp/test_mcp_resolution.py | 360 ++++++++++++++++++ tests/adapters/omp/test_room_retirement.py | 323 ++++++++++++++++ 7 files changed, 1412 insertions(+), 17 deletions(-) create mode 100644 tests/adapters/omp/test_cancellation_entry.py create mode 100644 tests/adapters/omp/test_deadline_settlement.py create mode 100644 tests/adapters/omp/test_mcp_resolution.py create mode 100644 tests/adapters/omp/test_room_retirement.py diff --git a/docs/acp.md b/docs/acp.md index 59cec7c1d..1bfc03fac 100644 --- a/docs/acp.md +++ b/docs/acp.md @@ -330,6 +330,46 @@ negotiate RPC protocol v2 to preserve structured text/image tool results. enable those receipts without rewriting user settings. Missing or failed receipts fail startup rather than allowing a model to answer without its required tools. The enclosing turn timeout still bounds startup. +- **Per-connection MCP servers.** Pass the constructor callable + `resolve_mcp_servers`, an async `(workspace: str) -> Mapping[str, OmpMcpServer]`, + when external servers (or their credentials) can change while the adapter runs. + It is called once for each new room process, with the room's claimed workspace, + before the overlay or process exists; the same result feeds both the overlay and + the startup receipt check. Its result is the room's complete **external** server + set: it replaces every external entry in `mcp_servers` (an empty mapping means + none), and an exception fails that startup with no fallback to them. A `band` + entry in `mcp_servers` is the SDK's own Band endpoint, not an external server: it + is kept beside the resolved set, so Band tools still route through it with + `inject_band_tools=False`. The resolver itself may not return `band`, because OMP + would mint that server's tools under Band tool names. A reused room process never + calls it again; a retired room (cancellation, failure, or cleanup) resolves afresh + on its next start. The callable is runtime-only and never part of + `OmpAdapterConfig`. +- **Room retirement.** `on_cleanup` starts one retirement per room, and overlapping + `on_cleanup` and `cleanup_all` calls join it. The room's active turn is the + current `on_message` call, not the task that made it. A cancellation already + requested when that call starts is delivered at its first step, before anything + is resolved or launched; one its caller had already received and handled does + not stop it. Retirement marks the active turn retired, so nothing it resolves + afterwards can start a process. While the turn is still running, retirement + becomes the cause of its end: it disarms the turn deadline if that has not fired, + so the deadline cannot cancel the turn's settlement again, and cancels the turn + once, or joins a cancellation it is already handling. A turn whose deadline has + fired, or that is already handling its own outcome, keeps its timeout or other + result and its failure report. A fired deadline counts as that cancellation even + if the resolver consumed it with `uncancel()`: cleanup joins it rather than + cancelling again, and a server set the resolver returns afterwards is refused, so + the turn ends with `TimeoutError` and nothing is launched. A turn that has already + finished is never + cancelled, so its caller's later work is untouched. Retirement then waits for + that call's own finalization (not for the rest of its caller's work), and only + then closes the process and releases the workspace. A turn may request cleanup of + its own room, for example from its resolver: that call returns at once, the turn + is not cancelled again but its unfired deadline is disarmed all the same, it stops + before launching, and teardown follows its settlement. Cancelling a cleanup + call stops only that caller's wait. Until the retirement ends, turns already + waiting for the room and turns that arrive during it are cancelled rather than + starting a new process. Once cleanup succeeds, a later turn can start normally. - **Permissions.** `approval_mode="yolo"` bypasses native approval checks without changing Band platform permissions. Otherwise `resolve_permission` receives `OmpPermissionRequest(room_id, session_id, title, options)` and returns an diff --git a/src/band/adapters/omp.py b/src/band/adapters/omp.py index ce9b7b0eb..35a80971e 100644 --- a/src/band/adapters/omp.py +++ b/src/band/adapters/omp.py @@ -3,12 +3,14 @@ from __future__ import annotations import asyncio +import collections.abc import json import logging from collections.abc import Awaitable, Callable, Iterable from contextlib import ExitStack from dataclasses import dataclass, field from pathlib import Path +from types import MappingProxyType from typing import ClassVar, Literal, Self, cast from uuid import uuid4 @@ -83,6 +85,7 @@ render_system_prompt, ) from band.runtime.tools import ( + BAND_MCP_SERVER_NAME, get_band_tool_category, get_tool_description, is_image_passthrough_result, @@ -136,6 +139,9 @@ class OmpPermissionRequest(BaseModel): PermissionResolver = Callable[[OmpPermissionRequest], Awaitable[str | None]] +McpServerResolver = Callable[ + [str], Awaitable[collections.abc.Mapping[str, OmpMcpServer]] +] class _ThinkingLevels(BaseModel): @@ -164,6 +170,64 @@ class OmpTurn: closing_text: list[str] = field(default_factory=list) +@dataclass +class _ActiveTurn: + """One ``on_message`` invocation that owns its room. + + ``task`` is what retirement cancels; ``settled`` is set by the invocation's own + finalization, so retirement never waits on the rest of its caller's lifetime. + """ + + task: asyncio.Task[object] + settled: asyncio.Event = field(default_factory=asyncio.Event) + retired: bool = False # retirement claimed the room; nothing may launch now + ending: bool = False # already handling its outcome; must not be interrupted + # The turn deadline while its context is entered; None before and after. + deadline: asyncio.Timeout | None = None + # A caller task may carry cancellations it handled earlier; only newer ones + # belong to this invocation. Fixed once the admission checkpoint has passed. + baseline: int = field(init=False) + + def __post_init__(self) -> None: + self.baseline = self.task.cancelling() + + def cancel_requested(self) -> bool: + """Whether this invocation is already being ended: a cancellation newer + than its baseline, or its own deadline having fired. The fired deadline + stays authoritative even after a callback consumed the request count with + ``uncancel()``, so takeover joins it and its resolver result is refused.""" + timer = self.deadline + return self.task.cancelling() > self.baseline or ( + timer is not None and timer.expired() + ) + + def running(self) -> bool: + """Neither settled nor already handling its own outcome.""" + return not self.settled.is_set() and not self.ending + + def retire_deadline(self) -> None: + """Disarm the deadline of a running invocation whose end retirement now + owns, so it cannot cancel the settlement a second time. A fired deadline, + or one whose context has exited, is never touched.""" + timer = self.deadline + if self.running() and timer is not None and not timer.expired(): + timer.reschedule(None) + + def take_over(self, *, cancel: bool) -> None: + """Make retirement the cause of this still-running invocation's end. + + A settled or outcome-handling invocation is left alone, and so is a fired + deadline. The live deadline is disarmed for every retirement; the task is + cancelled only when ``cancel`` allows it (never for a turn that requested + its own retirement) and no cancellation is already being handled. + """ + if not self.running(): + return + self.retire_deadline() + if cancel and not self.cancel_requested(): + self.task.cancel() + + class OmpAdapter(SimpleAdapter[OmpSessionState]): """Keep Band's turn owned until OMP reports native session settlement. @@ -185,6 +249,7 @@ def __init__( additional_tools: list[CustomToolDef] | None = None, workspace_for_room: WorkspaceResolver | None = None, resolve_permission: PermissionResolver | None = None, + resolve_mcp_servers: McpServerResolver | None = None, **features: Unpack[FeatureKwargs], ) -> None: super().__init__(history_converter=OmpHistoryConverter(), **features) @@ -196,8 +261,11 @@ def __init__( ) self._rooms: dict[str, OmpRoom] = {} self._locks: dict[str, asyncio.Lock] = {} - self._owners: dict[str, asyncio.Task[object]] = {} + self._owners: dict[str, _ActiveTurn] = {} self._room_selections: dict[str, ModelSelection] = {} + self._retirements: dict[str, int] = {} + self._retiring: dict[str, asyncio.Task[None]] = {} + self._resolve_mcp_servers = resolve_mcp_servers self._host_schemas = self._build_host_tools() def apply_effective_features(self, features: AdapterFeatures) -> None: @@ -266,13 +334,46 @@ def _spawn_command(self, workspace: str | None) -> list[str]: ) return omp_command_in_workspace(command, workspace) if workspace else command + async def _connection_mcp_servers( + self, room_id: str, workspace: str + ) -> collections.abc.Mapping[str, OmpMcpServer]: + """The external MCP servers for one new native connection. + + Without a resolver this is the static config. With one, its result is the + complete external set for this connection, used for both the overlay and + native readiness; a failure propagates rather than falling back to config. + """ + if self._resolve_mcp_servers is None: + return self.config.mcp_servers + servers = dict(await self._resolve_mcp_servers(workspace)) + active = self._owners.get(room_id) + if active is not None and (active.retired or active.cancel_requested()): + # Retirement or cancellation arrived while the resolver ran and the + # resolver returned anyway; nothing may be launched from its result. + raise asyncio.CancelledError + for name, server in servers.items(): + if not isinstance(name, str) or not isinstance(server, OmpMcpServer): + raise TypeError("resolve_mcp_servers must map names to OmpMcpServer") + if name == BAND_MCP_SERVER_NAME: + # OMP mints this server's tools under Band's own tool names. + raise ValueError( + f"MCP server name {BAND_MCP_SERVER_NAME!r} is reserved for Band tools" + ) + protected = self.config.mcp_servers.get(BAND_MCP_SERVER_NAME) + if protected is not None: + # The configured Band endpoint is the SDK's own channel, not an external + # server, so it is kept beside whatever external set was resolved. + servers = {BAND_MCP_SERVER_NAME: protected, **servers} + return MappingProxyType(servers) + async def _open_room(self, room_id: str, history: OmpSessionState) -> OmpRoom: workspace = self._workspaces.claim(room_id) resources = ExitStack() try: + servers = await self._connection_mcp_servers(room_id, workspace) plugin = resources.enter_context( omp_mcp_plugin( - self.config.mcp_servers, + servers, workspace=workspace, system_prompt=self._build_system_context(), ) @@ -289,7 +390,7 @@ async def _open_room(self, room_id: str, history: OmpSessionState) -> OmpRoom: self._rooms[room_id] = room await room.runtime.start() if plugin: - await omp_mcp_ready(room.runtime, plugin, self.config.mcp_servers) + await omp_mcp_ready(room.runtime, plugin, servers) candidate = history.room_to_session.get(room_id) if ( candidate @@ -441,12 +542,29 @@ async def on_message( is_session_bootstrap: bool, room_id: str, ) -> None: + if room_id in self._retiring: + raise asyncio.CancelledError # the room is being retired + retirement = self._retirements.get(room_id, 0) async with self._locks.setdefault(room_id, asyncio.Lock()): - owner = asyncio.current_task() - if owner is not None: - self._owners[room_id] = owner + if room_id in self._retiring or ( + self._retirements.get(room_id, 0) != retirement + ): + # The room was retired, or is being retired, while this turn waited. + raise asyncio.CancelledError + task = asyncio.current_task() + assert task is not None, "OMP turns run inside an asyncio task" + active = _ActiveTurn(task) + self._owners[room_id] = active deadline = asyncio.timeout(self.config.turn_timeout_s) try: + # A cancellation requested before this call is delivered here, inside + # its cleanup ownership and before any resolver or launch; only a + # count that survives this checkpoint was handled earlier. + await asyncio.sleep(0) + if active.retired: + raise asyncio.CancelledError + active.baseline = task.cancelling() + active.deadline = deadline async with deadline: room = self._rooms.get(room_id) if room is not None and room.retiring: @@ -501,14 +619,18 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: OMP_SESSION_METADATA_KEY: room.session_file, }, ) + active.deadline = None except asyncio.CancelledError: - await self._close_room(room_id) + active.ending = True + await self._retire_after_failure(room_id) raise except DeliveryFailedError as error: - await self._close_room(room_id) + active.ending = True + await self._retire_after_failure(room_id) reraise_delivery_cause(error) except Exception: - await self._close_room(room_id) + active.ending = True + await self._retire_after_failure(room_id) logger.exception("OMP turn failed in room %s", room_id) failure = ( AgentFailure( @@ -522,7 +644,17 @@ async def on_ui(request: OmpUIRequest) -> OmpUIResponse: await tools.send_failure(failure) raise finally: - self._owners.pop(room_id, None) + if self._owners.get(room_id) is active: + del self._owners[room_id] + active.settled.set() + + async def _retire_after_failure(self, room_id: str) -> None: + try: + await self._close_room(room_id) + except Exception: + logger.exception( + "OMP room %s could not be closed; it stays owned until cleanup", room_id + ) async def _execute_tool( self, call: OmpHostToolCall, tools: AgentToolsProtocol @@ -812,10 +944,57 @@ async def _close_room(self, room_id: str) -> None: self._workspaces.release(room_id) async def on_cleanup(self, room_id: str) -> None: - owner = self._owners.get(room_id) - if owner is not None and owner is not asyncio.current_task(): - owner.cancel() - await asyncio.gather(owner, return_exceptions=True) + """Retire a room, joining a retirement already in progress for it. + + Cancelling this call stops only the wait: the shared retirement keeps + running, keeps the room closed to new turns, and still owns teardown. + """ + operation = self._retiring.get(room_id) + if operation is None: + operation = self._begin_retirement(room_id) + active = self._owners.get(room_id) + if active is not None and active.task is asyncio.current_task(): + # The room's own turn cannot wait for a retirement that waits for it. + # It is already marked retired, so it stops before launching anything, + # and teardown follows its settlement. + return + await asyncio.shield(operation) + + def _begin_retirement(self, room_id: str) -> asyncio.Task[None]: + # Everything here is synchronous, so admission closes before any turn can + # run: queued turns see the new generation, arriving ones the retirement, + # and the active turn can no longer launch from whatever it resolves. + self._retirements[room_id] = self._retirements.get(room_id, 0) + 1 + active = self._owners.get(room_id) + if active is not None: + active.retired = True + # Retirement owns the turn's end from here, including when the turn + # requested it itself: its unfired deadline must not cancel again. + active.retire_deadline() + operation = asyncio.create_task( + self._retire(room_id, active, asyncio.current_task()) + ) + self._retiring[room_id] = operation + operation.add_done_callback(lambda task: self._end_retirement(room_id, task)) + return operation + + def _end_retirement(self, room_id: str, operation: asyncio.Task[None]) -> None: + if self._retiring.get(room_id) is operation: + del self._retiring[room_id] + if not operation.cancelled() and operation.exception() is not None: + logger.warning("OMP room %s retirement did not complete", room_id) + + async def _retire( + self, + room_id: str, + active: _ActiveTurn | None, + requester: asyncio.Task[object] | None, + ) -> None: + if active is not None: + if self._owners.get(room_id) is active: + active.take_over(cancel=active.task is not requester) + # The invocation's own finalization, not its caller's whole lifetime. + await active.settled.wait() async with self._locks.setdefault(room_id, asyncio.Lock()): room = self._rooms.get(room_id) retrying = room is not None and room.retiring @@ -824,10 +1003,17 @@ async def on_cleanup(self, room_id: str) -> None: self._room_selections.pop(room_id, None) def _cleanup_rooms(self) -> Iterable[Awaitable[None]]: - return (self.on_cleanup(room_id) for room_id in tuple(self._rooms)) + # Resolving and retiring rooms may own work without an installed process. + rooms = dict.fromkeys((*self._rooms, *self._owners, *self._retiring)) + return (self.on_cleanup(room_id) for room_id in rooms) async def cleanup_all(self) -> None: await asyncio.gather(*self._cleanup_rooms()) -__all__ = ["OmpAdapter", "OmpAdapterConfig", "OmpPermissionRequest"] +__all__ = [ + "McpServerResolver", + "OmpAdapter", + "OmpAdapterConfig", + "OmpPermissionRequest", +] diff --git a/tests/adapters/omp/kit.py b/tests/adapters/omp/kit.py index 6819a498f..2b39b4abd 100644 --- a/tests/adapters/omp/kit.py +++ b/tests/adapters/omp/kit.py @@ -12,6 +12,7 @@ from band.adapters.omp import OmpAdapter, OmpAdapterConfig from band.converters.omp import OmpSessionState +from band.integrations.ompprocess import OmpMcpServer from band.integrations.omprpc import OmpRpcClient from band.integrations.ompwire import ( OmpContent, @@ -21,7 +22,7 @@ OmpUIRequest, ) from band.runtime.tools import ToolCallOutcome -from band.testing import FakeAgentTools +from band.testing import FakeAgentTools, reported_failures from tests.integrations.acp.conftest import make_platform_message from tests.integrations.omprpc.kit import process_is_alive from tests.paths import REPO_ROOT @@ -284,3 +285,222 @@ def launches(monkeypatch: pytest.MonkeyPatch) -> Launches: record = Launches() monkeypatch.setattr("band.adapters.omp.OmpRpcClient", record) return record + + +def remote(token: str) -> OmpMcpServer: + return OmpMcpServer( + type="http", + url="http://127.0.0.1:9/mcp", + headers={"Authorization": f"Bearer {token}"}, + ) + + +class BandMcpReplyRuntime(PendingRuntime): + """Replies through the Band MCP server mounted in the overlay, not a host tool.""" + + async def prompt(self, _message, *, on_event, on_tool, on_ui): + name = "mcp__band_sdk_band_send_message" + args = {"content": "done"} + await on_event( + OmpEvent( + type="tool_execution_start", toolCallId="c1", toolName=name, args=args + ) + ) + await on_event( + OmpEvent( + type="tool_execution_end", toolCallId="c1", toolName=name, isError=False + ) + ) + + +class SignallingLock(asyncio.Lock): + """A room lock that reports when a second turn starts waiting for it.""" + + def __init__(self) -> None: + super().__init__() + self.contended = asyncio.Event() + + async def acquire(self) -> bool: + if self.locked(): + self.contended.set() + return await super().acquire() + + +class SettlingResolver: + """Blocks the first resolution; when cancelled it settles durable work before + re-raising, and counts every cancellation it receives along the way.""" + + def __init__(self) -> None: + self.calls = 0 + self.entered = asyncio.Event() + self.cancel_seen = asyncio.Event() + self.settle = asyncio.Event() + self.settled = False + self.cancellations = 0 + + async def __call__(self, _workspace: str): + self.calls += 1 + if self.calls > 1: + return {} + self.entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancellations += 1 + self.cancel_seen.set() + try: + await self.settle.wait() + except asyncio.CancelledError: + self.cancellations += 1 # settlement interrupted by a second cancel + raise + self.settled = True + raise + return {} + + +def settling_adapter(tmp_path: Path, resolver, **config): + adapter = OmpAdapter( + OmpAdapterConfig(cwd=str(tmp_path), **config), + resolve_mcp_servers=resolver, + emit=(), + ) + releases: list[bool] = [] # whether settlement had finished at each release + real_release = adapter._workspaces.release + + def release(room_id, workspace=None): + releases.append(resolver.settled) + return real_release(room_id, workspace) + + adapter._workspaces.release = release # type: ignore[method-assign] + return adapter, releases + + +async def start_blocked_turn(adapter, resolver: SettlingResolver): + first = asyncio.create_task(run_turn(adapter, FakeAgentTools())) + await asyncio.wait_for(resolver.entered.wait(), 1) + return first + + +async def checkpoint() -> None: + resumed = asyncio.Event() + asyncio.get_running_loop().call_soon(resumed.set) + await resumed.wait() + + +async def start_cleanup(adapter: OmpAdapter, *, shutdown: bool = False): + entered = asyncio.Event() + original = adapter.on_cleanup + + async def cleanup(room_id): + entered.set() + await original(room_id) + + adapter.on_cleanup = cleanup # type: ignore[method-assign] + try: + operation = ( + adapter.cleanup_all() if shutdown else adapter.on_cleanup("room-123") + ) + task = asyncio.create_task(operation) + await asyncio.wait_for(entered.wait(), 1) + finally: + adapter.on_cleanup = original # type: ignore[method-assign] + return task + + +async def cross_deadline(clock, seconds: float) -> None: + crossed = asyncio.Event() + handle = asyncio.get_running_loop().call_later(seconds, crossed.set) + try: + clock.advance(seconds) + await asyncio.wait_for(crossed.wait(), 1) + finally: + handle.cancel() + + +class SwallowingResolver: + """Catches a cancellation at its first await and returns a server anyway.""" + + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self, _workspace: str): + self.calls += 1 + try: + await checkpoint() + except asyncio.CancelledError: + pass + return {"linear": remote("late")} + + +class SelfRetiringResolver(SettlingResolver): + """On its first cancellation it asks for its own room's cleanup, then holds + its settlement until released.""" + + adapter: OmpAdapter + + async def __call__(self, workspace: str): + self.calls += 1 + if self.calls > 1: + return {} + self.entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancellations += 1 + await self.adapter.on_cleanup("room-123") # returns at once for self + self.cancel_seen.set() + try: + await self.settle.wait() + except asyncio.CancelledError: + self.cancellations += 1 + raise + self.settled = True + raise + return {} + + +class UncancellingResolver(SettlingResolver): + """Consumes the deadline's cancellation with the public ``uncancel()``, then + settles; ``returns`` decides whether it re-raises or returns a server.""" + + def __init__(self, *, returns: bool) -> None: + super().__init__() + self.returns = returns + + async def __call__(self, _workspace: str): + self.calls += 1 + if self.calls > 1: + return {} + self.entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + asyncio.current_task().uncancel() # the request count is now consumed + self.cancellations += 1 + self.cancel_seen.set() + try: + await self.settle.wait() + except asyncio.CancelledError: + self.cancellations += 1 + raise + self.settled = True + if not self.returns: + raise + return {"linear": remote("late")} + + +def timeout_reports(tools) -> list[tuple[str, str | None]]: + return [(f["provider"], f.get("code")) for f in reported_failures(tools)] + + +def failing_runtime(closing: asyncio.Event, finish_close: asyncio.Event): + class FailingRuntime(PendingRuntime): + async def prompt(self, *_args, **_kwargs): + raise RuntimeError("native turn failed") + + async def close(self): + closing.set() + await finish_close.wait() + self.closed = True + + return FailingRuntime() diff --git a/tests/adapters/omp/test_cancellation_entry.py b/tests/adapters/omp/test_cancellation_entry.py new file mode 100644 index 000000000..cb9b0b068 --- /dev/null +++ b/tests/adapters/omp/test_cancellation_entry.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest + +from band.adapters.omp import OmpAdapter, OmpAdapterConfig +from band.testing import FakeAgentTools +from tests.adapters.omp.kit import ( + Launches, + SwallowingResolver, + checkpoint, + launches, + run_turn, +) + +__all__ = ["launches"] + + +class TestCancellationEntry: + @pytest.mark.asyncio + async def test_a_cancellation_pending_at_entry_launches_nothing( + self, launches: Launches, tmp_path: Path + ): + resolver = SwallowingResolver() + adapter = OmpAdapter( + OmpAdapterConfig(cwd=str(tmp_path)), resolve_mcp_servers=resolver, emit=() + ) + + async def caller(): + asyncio.current_task().cancel() # requested, not yet delivered + await run_turn(adapter, FakeAgentTools()) + + with pytest.raises(asyncio.CancelledError): + await asyncio.create_task(caller()) + + assert resolver.calls == 0 and launches.runtimes == [] + assert not adapter._workspaces.rooms and not adapter._owners + await run_turn(adapter, FakeAgentTools()) # explicit recovery + assert len(launches.runtimes) == 1 + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_a_cancellation_handled_before_entry_does_not_stop_the_turn( + self, launches: Launches, tmp_path: Path + ): + resolver = SwallowingResolver() + adapter = OmpAdapter( + OmpAdapterConfig(cwd=str(tmp_path)), resolve_mcp_servers=resolver, emit=() + ) + + async def caller(): + asyncio.current_task().cancel() + try: + await checkpoint() # delivered and handled by the caller itself + except asyncio.CancelledError: + pass + await run_turn(adapter, FakeAgentTools()) + + await asyncio.create_task(caller()) + + assert resolver.calls == 1 and len(launches.runtimes) == 1 + await adapter.cleanup_all() diff --git a/tests/adapters/omp/test_deadline_settlement.py b/tests/adapters/omp/test_deadline_settlement.py new file mode 100644 index 000000000..ab09882a0 --- /dev/null +++ b/tests/adapters/omp/test_deadline_settlement.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest + +from band.testing import FakeAgentTools, reported_failures +from tests.adapters.omp.kit import ( + Launches, + SelfRetiringResolver, + SettlingResolver, + UncancellingResolver, + cross_deadline, + launches, + run_turn, + settling_adapter, + start_cleanup, + timeout_reports, +) +from tests.toolkit.clock import binary_clock_resolution +from tests.toolkit.timecontrol import controlled_clock + +__all__ = ["launches"] + + +@pytest.mark.looptime(resolution=binary_clock_resolution()) +class TestDeadlineSettlement: + @pytest.mark.asyncio + async def test_cleanup_during_a_deadline_settlement_keeps_the_timeout_report( + self, launches: Launches, tmp_path: Path + ): + async with controlled_clock() as deadline_clock: + resolver = SettlingResolver() + adapter, releases = settling_adapter( + tmp_path, resolver, turn_timeout_s=0.05 + ) + tools = FakeAgentTools() + first = asyncio.create_task(run_turn(adapter, tools)) + await asyncio.wait_for(resolver.entered.wait(), 1) + deadline_clock.advance(0.1) + await asyncio.wait_for(resolver.cancel_seen.wait(), 1) # the deadline fired + + cleanup = await start_cleanup(adapter) + + assert not cleanup.done() + resolver.settle.set() + await asyncio.wait_for(cleanup, 1) + with pytest.raises(TimeoutError): + await first + assert resolver.cancellations == 1 and releases and all(releases) + assert [ + (failure["provider"], failure.get("code")) + for failure in reported_failures(tools) + ] == [("omp", "timeout")] + assert launches.runtimes == [] + + @pytest.mark.asyncio + async def test_cleanup_before_the_deadline_owns_the_turn_through_slow_settlement( + self, launches: Launches, tmp_path: Path + ): + async with controlled_clock() as deadline_clock: + resolver = SettlingResolver() + adapter, releases = settling_adapter(tmp_path, resolver, turn_timeout_s=0.2) + tools = FakeAgentTools() + first = asyncio.create_task(run_turn(adapter, tools)) + await asyncio.wait_for(resolver.entered.wait(), 1) + cleanup = asyncio.create_task(adapter.on_cleanup("room-123")) + await asyncio.wait_for(resolver.cancel_seen.wait(), 1) + + # Virtual time crosses the original deadline while settlement is held. + await cross_deadline(deadline_clock, 0.4) + + assert resolver.cancellations == 1 and not cleanup.done() + assert "room-123" in adapter._workspaces.rooms + resolver.settle.set() + await asyncio.wait_for(cleanup, 1) + with pytest.raises(asyncio.CancelledError): + await first + assert resolver.cancellations == 1 and releases and all(releases) + assert not reported_failures(tools) and launches.runtimes == [] + await run_turn(adapter, FakeAgentTools()) # explicit recovery + assert len(launches.runtimes) == 1 + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_self_cleanup_after_a_caller_cancellation_disarms_the_deadline( + self, launches: Launches, tmp_path: Path + ): + async with controlled_clock() as deadline_clock: + resolver = SelfRetiringResolver() + adapter, releases = settling_adapter(tmp_path, resolver, turn_timeout_s=0.2) + resolver.adapter = adapter + tools = FakeAgentTools() + first = asyncio.create_task(run_turn(adapter, tools)) + await asyncio.wait_for(resolver.entered.wait(), 1) + first.cancel() # the caller cancels; the resolver then retires its own room + await asyncio.wait_for(resolver.cancel_seen.wait(), 1) + shutdown = await start_cleanup(adapter, shutdown=True) + + # Virtual time crosses the original deadline while settlement is held. + await cross_deadline(deadline_clock, 0.4) + + assert ( + resolver.cancellations == 1 and not first.done() and not shutdown.done() + ) + assert "room-123" in adapter._workspaces.rooms + resolver.settle.set() + await asyncio.wait_for(shutdown, 1) + with pytest.raises(asyncio.CancelledError): + await first + assert resolver.cancellations == 1 and releases and all(releases) + assert not reported_failures(tools) and launches.runtimes == [] + await run_turn(adapter, FakeAgentTools()) # explicit recovery + assert len(launches.runtimes) == 1 + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_self_cleanup_after_the_deadline_keeps_the_timeout_report( + self, launches: Launches, tmp_path: Path + ): + async with controlled_clock() as deadline_clock: + resolver = SelfRetiringResolver() + adapter, releases = settling_adapter( + tmp_path, resolver, turn_timeout_s=0.05 + ) + resolver.adapter = adapter + tools = FakeAgentTools() + first = asyncio.create_task(run_turn(adapter, tools)) + await asyncio.wait_for(resolver.entered.wait(), 1) + deadline_clock.advance(0.1) + await asyncio.wait_for( + resolver.cancel_seen.wait(), 1 + ) # the deadline fired first + + resolver.settle.set() + with pytest.raises(TimeoutError): + await asyncio.wait_for(first, 1) + await asyncio.wait_for(adapter.cleanup_all(), 1) + + assert resolver.cancellations == 1 and resolver.settled + assert releases and all(releases) + assert [ + (failure["provider"], failure.get("code")) + for failure in reported_failures(tools) + ] == [("omp", "timeout")] + assert launches.runtimes == [] + + @pytest.mark.asyncio + async def test_cleanup_joins_a_fired_deadline_whose_cancellation_was_consumed( + self, launches: Launches, tmp_path: Path + ): + async with controlled_clock() as deadline_clock: + resolver = UncancellingResolver(returns=False) + adapter, releases = settling_adapter( + tmp_path, resolver, turn_timeout_s=0.05 + ) + tools = FakeAgentTools() + first = asyncio.create_task(run_turn(adapter, tools)) + await asyncio.wait_for(resolver.entered.wait(), 1) + deadline_clock.advance(0.1) + await asyncio.wait_for(resolver.cancel_seen.wait(), 1) # the deadline fired + + cleanup = await start_cleanup(adapter) + + assert not cleanup.done() + assert ( + resolver.cancellations == 1 and "room-123" in adapter._workspaces.rooms + ) + resolver.settle.set() + await asyncio.wait_for(cleanup, 1) + with pytest.raises(TimeoutError): + await first + assert resolver.cancellations == 1 and releases and all(releases) + assert timeout_reports(tools) == [("omp", "timeout")] + assert launches.runtimes == [] + + @pytest.mark.asyncio + async def test_a_result_returned_after_the_deadline_fired_launches_nothing( + self, launches: Launches, tmp_path: Path + ): + async with controlled_clock() as deadline_clock: + resolver = UncancellingResolver(returns=True) + resolver.settle.set() # settlement finishes and the resolver returns a server + adapter, releases = settling_adapter( + tmp_path, resolver, turn_timeout_s=0.05 + ) + tools = FakeAgentTools() + turn = asyncio.create_task(run_turn(adapter, tools)) + await asyncio.wait_for(resolver.entered.wait(), 1) + deadline_clock.advance(0.1) + + with pytest.raises(TimeoutError): + await asyncio.wait_for(turn, 1) + + assert resolver.cancellations == 1 and resolver.settled + assert launches.runtimes == [] and launches.plugins == [] + assert timeout_reports(tools) == [("omp", "timeout")] + assert not adapter._workspaces.rooms and releases and all(releases) + await run_turn(adapter, FakeAgentTools()) # explicit recovery + assert len(launches.runtimes) == 1 + await adapter.cleanup_all() diff --git a/tests/adapters/omp/test_mcp_resolution.py b/tests/adapters/omp/test_mcp_resolution.py new file mode 100644 index 000000000..d57c8f50d --- /dev/null +++ b/tests/adapters/omp/test_mcp_resolution.py @@ -0,0 +1,360 @@ +from __future__ import annotations + +import asyncio +import json +import os +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest + +import band.integrations.ompprocess as omp_module +from band.adapters.omp import OmpAdapter, OmpAdapterConfig +from band.converters.omp import OmpSessionState +from band.integrations.ompprocess import OmpMcpServer +from band.testing import FakeAgentTools, failure_reports +from tests.adapters.omp.kit import ( + BandMcpReplyRuntime, + Launches, + PendingRuntime, + checkpoint, + launches, + remote, + run_turn, +) +from tests.integrations.acp.conftest import make_platform_message + +__all__ = ["launches"] + + +class TestMcpResolution: + @pytest.mark.asyncio + async def test_each_new_connection_uses_its_own_resolved_snapshot( + self, launches: Launches, tmp_path: Path + ): + current = {"linear": remote("grant-A")} + seen: list[str] = [] + + async def resolve(workspace: str): + seen.append(workspace) + return dict(current) + + adapter = OmpAdapter( + OmpAdapterConfig(cwd=str(tmp_path)), resolve_mcp_servers=resolve, emit=() + ) + await run_turn(adapter, FakeAgentTools()) + await adapter.on_cleanup("room-123") # the native connection is retired + current["linear"] = remote("grant-B") + await run_turn(adapter, FakeAgentTools()) + + assert [s["linear"]["headers"]["Authorization"] for s in launches.servers] == [ + "Bearer grant-A", + "Bearer grant-B", + ] + claimed = os.path.realpath(tmp_path / "room-123") + assert seen == [claimed, claimed] # the claimed workspace, each time + assert "grant" not in repr(adapter.config) + assert adapter.config.mcp_servers == {} + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_reusing_a_live_connection_does_not_resolve_again( + self, launches: Launches, tmp_path: Path + ): + resolve = AsyncMock(return_value={"linear": remote("grant-A")}) + adapter = OmpAdapter( + OmpAdapterConfig(cwd=str(tmp_path)), resolve_mcp_servers=resolve, emit=() + ) + + await run_turn(adapter, FakeAgentTools()) + await run_turn(adapter, FakeAgentTools()) + + assert resolve.await_count == 1 + assert len(launches.runtimes) == 1 + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_a_resolved_set_replaces_static_servers_even_when_empty( + self, launches: Launches, tmp_path: Path, monkeypatch + ): + ready = AsyncMock() + monkeypatch.setattr("band.adapters.omp.omp_mcp_ready", ready) + adapter = OmpAdapter( + OmpAdapterConfig(cwd=str(tmp_path), mcp_servers={"stale": remote("old")}), + resolve_mcp_servers=AsyncMock(return_value={}), + emit=(), + ) + + await run_turn(adapter, FakeAgentTools()) + + assert launches.servers == [{}] + assert ready.await_args.args[2] == {} # readiness expects the same empty set + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_readiness_checks_the_same_snapshot_the_overlay_carries( + self, launches: Launches, tmp_path: Path, monkeypatch + ): + ready = AsyncMock() + monkeypatch.setattr("band.adapters.omp.omp_mcp_ready", ready) + snapshots = iter([{"linear": remote("grant-A")}, {"linear": remote("grant-B")}]) + + async def resolve(_workspace: str): + return next(snapshots) + + adapter = OmpAdapter( + OmpAdapterConfig(cwd=str(tmp_path)), resolve_mcp_servers=resolve, emit=() + ) + await run_turn(adapter, FakeAgentTools()) + + assert ready.await_count == 1 + assert ready.await_args.args[2] == {"linear": remote("grant-A")} + assert ( + launches.servers[0]["linear"]["headers"]["Authorization"] + == "Bearer grant-A" + ) + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_a_failed_resolution_launches_nothing_and_a_later_one_recovers( + self, launches: Launches, tmp_path: Path + ): + resolve = AsyncMock( + side_effect=[ + RuntimeError("grant unavailable"), + {"linear": remote("grant-B")}, + ] + ) + adapter = OmpAdapter( + OmpAdapterConfig( + cwd=str(tmp_path), mcp_servers={"linear": remote("static")} + ), + resolve_mcp_servers=resolve, + emit=(), + ) + failed = FakeAgentTools() + + with pytest.raises(RuntimeError, match="grant unavailable"): + await run_turn(adapter, failed) + + assert launches.runtimes == [] # no static fallback, no process + assert failure_reports(failed) + assert not adapter._workspaces.rooms # the claim was released + await run_turn(adapter, FakeAgentTools()) + assert [s["linear"]["headers"]["Authorization"] for s in launches.servers] == [ + "Bearer grant-B" + ] + await adapter.cleanup_all() + + @pytest.mark.asyncio + @pytest.mark.parametrize("resolver_settles", [False, True]) + async def test_shutdown_during_resolution_leaves_no_overlay_or_process( + self, launches: Launches, tmp_path: Path, monkeypatch, resolver_settles: bool + ): + entered = asyncio.Event() + overlays: list[str] = [] + real_plugin = omp_module.omp_mcp_plugin + + def recording_plugin(servers, **kwargs): + overlays.append(kwargs["workspace"]) + return real_plugin(servers, **kwargs) + + monkeypatch.setattr("band.adapters.omp.omp_mcp_plugin", recording_plugin) + + async def resolve(_workspace: str): + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + if not resolver_settles: + raise + # A resolver may settle durable work and return instead of re-raising. + await checkpoint() + return {"linear": remote("late")} + + adapter = OmpAdapter( + OmpAdapterConfig(cwd=str(tmp_path)), resolve_mcp_servers=resolve, emit=() + ) + turn = asyncio.create_task(run_turn(adapter, FakeAgentTools())) + await asyncio.wait_for(entered.wait(), 1) + + await asyncio.wait_for(adapter.cleanup_all(), 1) + + assert turn.done() + assert overlays == [] and launches.runtimes == [] + assert not adapter._workspaces.rooms + with pytest.raises(asyncio.CancelledError): + await turn + + @pytest.mark.asyncio + async def test_a_blocked_resolution_does_not_hold_another_room( + self, launches: Launches, tmp_path: Path + ): + release = asyncio.Event() + slow_entered = asyncio.Event() + + async def resolve(workspace: str): + if workspace.endswith("slow"): + slow_entered.set() + await release.wait() + return {} + + adapter = OmpAdapter( + resolve_mcp_servers=resolve, + workspace_for_room=lambda room: str(tmp_path / room), + emit=(), + ) + for room in ("slow", "fast"): + (tmp_path / room).mkdir() + + async def turn(room: str): + await adapter.run_judged_turn( + msg=make_platform_message("hello"), + tools=FakeAgentTools(), + history=OmpSessionState(), + participants_msg=None, + contacts_msg=None, + is_session_bootstrap=True, + room_id=room, + ) + + slow = asyncio.create_task(turn("slow")) + await asyncio.wait_for( + slow_entered.wait(), 1 + ) # slow is now inside its resolver + await asyncio.wait_for(turn("fast"), 1) + + assert ( + not release.is_set() and not slow.done() + ) # fast finished while slow blocked + assert len(launches.runtimes) == 1 + release.set() + await asyncio.wait_for(slow, 1) + assert len(launches.runtimes) == 2 + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_resolved_stdio_paths_are_anchored_in_the_claimed_workspace( + self, launches: Launches, tmp_path: Path + ): + server = OmpMcpServer( + command="./bin/tool", cwd="tools", env={"TOKEN": "s3cret"} + ) + adapter = OmpAdapter( + OmpAdapterConfig(cwd=str(tmp_path)), + resolve_mcp_servers=AsyncMock(return_value={"tool": server}), + emit=(), + ) + + await run_turn(adapter, FakeAgentTools()) + + (servers,) = launches.servers + workspace = Path(os.path.realpath(tmp_path)) / "room-123" + assert servers["tool"]["cwd"] == str(workspace / "tools") + assert servers["tool"]["command"] == str(workspace / "tools" / "bin" / "tool") + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_a_resolved_server_cannot_take_the_band_tool_namespace( + self, launches: Launches, tmp_path: Path + ): + adapter = OmpAdapter( + OmpAdapterConfig(cwd=str(tmp_path)), + resolve_mcp_servers=AsyncMock(return_value={"band": remote("x")}), + emit=(), + ) + tools = FakeAgentTools() + + with pytest.raises(ValueError, match="reserved"): + await run_turn(adapter, tools) + + assert launches.runtimes == [] + assert failure_reports(tools) + + @pytest.mark.asyncio + async def test_a_restored_session_survives_a_resolved_connection( + self, launches: Launches, tmp_path: Path + ): + session_dir = ( + Path(os.path.realpath(tmp_path)) / "room-123" / ".band-omp-sessions" + ) + session_dir.mkdir(parents=True) + saved = session_dir / "saved.jsonl" + saved.write_text("{}", encoding="utf-8") + requests: list[str] = [] + adapter = OmpAdapter( + OmpAdapterConfig(cwd=str(tmp_path)), + resolve_mcp_servers=AsyncMock(return_value={"linear": remote("grant-A")}), + emit=(), + ) + + async def bootstrap_turn(): + await adapter.run_judged_turn( + msg=make_platform_message("hello"), + tools=FakeAgentTools(), + history=OmpSessionState(room_to_session={"room-123": str(saved)}), + participants_msg=None, + contacts_msg=None, + is_session_bootstrap=True, + room_id="room-123", + ) + + original = PendingRuntime.request + + async def recording(self, command, **parameters): + requests.append(command) + if command == "switch_session": + return {"cancelled": False} + return await original(self, command, **parameters) + + PendingRuntime.request = recording # type: ignore[method-assign] + try: + await bootstrap_turn() + finally: + PendingRuntime.request = original # type: ignore[method-assign] + + assert "switch_session" in requests + await adapter.cleanup_all() + + @pytest.mark.asyncio + @pytest.mark.parametrize("external", [{}, {"linear": remote("grant-A")}]) + async def test_the_configured_band_endpoint_survives_any_resolved_external_set( + self, tmp_path: Path, monkeypatch, external + ): + overlays: list[dict[str, object]] = [] + + def spawn(command, **_kwargs): + plugin = Path(command[command.index("--plugin-dir") + 1]) + overlays.append( + json.loads((plugin / ".mcp.json").read_text(encoding="utf-8")) + ) + return BandMcpReplyRuntime() + + ready = AsyncMock() + monkeypatch.setattr("band.adapters.omp.OmpRpcClient", spawn) + monkeypatch.setattr("band.adapters.omp.omp_mcp_ready", ready) + band = OmpMcpServer(type="http", url="http://127.0.0.1:9/band") + adapter = OmpAdapter( + OmpAdapterConfig( + cwd=str(tmp_path), + inject_band_tools=False, + mcp_servers={"band": band, "stale": remote("old")}, + ), + resolve_mcp_servers=AsyncMock(return_value=external), + emit=(), + ) + tools = FakeAgentTools() + + await run_turn(adapter, tools) + + expected = {"band", *external} + (overlay,) = overlays + assert ( + set(overlay["mcpServers"]) == expected + ) # the cached external entry is gone + assert overlay["mcpServers"]["band"]["url"] == "http://127.0.0.1:9/band" + assert set(ready.await_args.args[2]) == expected + assert ( + tools.turn.replied + ) # Band's reply tool still routes through its own server + await adapter.cleanup_all() diff --git a/tests/adapters/omp/test_room_retirement.py b/tests/adapters/omp/test_room_retirement.py new file mode 100644 index 000000000..a603652a1 --- /dev/null +++ b/tests/adapters/omp/test_room_retirement.py @@ -0,0 +1,323 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest + +from band.adapters.omp import OmpAdapter, OmpAdapterConfig +from band.testing import FakeAgentTools, reported_failures +from tests.adapters.omp.kit import ( + Launches, + PendingRuntime, + SettlingResolver, + SignallingLock, + failing_runtime, + launches, + remote, + run_turn, + settling_adapter, + start_blocked_turn, + start_cleanup, +) + +__all__ = ["launches"] + + +class TestRoomRetirement: + @pytest.mark.asyncio + async def test_cleanup_settles_a_turn_queued_behind_a_blocked_resolution( + self, launches: Launches, tmp_path: Path + ): + entered = asyncio.Event() + calls = 0 + blocking = True + + async def resolve(_workspace: str): + nonlocal calls + calls += 1 + entered.set() + if blocking: + await asyncio.Event().wait() + return {} + + adapter = OmpAdapter( + OmpAdapterConfig(cwd=str(tmp_path)), resolve_mcp_servers=resolve, emit=() + ) + room_lock = SignallingLock() + adapter._locks["room-123"] = room_lock + first = asyncio.create_task(run_turn(adapter, FakeAgentTools())) + await asyncio.wait_for(entered.wait(), 1) + blocking = False # anything that resolves from here on would launch + queued = asyncio.create_task(run_turn(adapter, FakeAgentTools())) + await asyncio.wait_for( + room_lock.contended.wait(), 1 + ) # queued waits on the lock + + await asyncio.wait_for(adapter.cleanup_all(), 1) + + assert first.done() and queued.done() + assert calls == 1 and launches.runtimes == [] + assert not adapter._workspaces.rooms + for task in (first, queued): + with pytest.raises(asyncio.CancelledError): + await task + await run_turn(adapter, FakeAgentTools()) # a later explicit turn recovers + assert calls == 2 and len(launches.runtimes) == 1 + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_a_turn_arriving_during_retirement_is_refused( + self, launches: Launches, tmp_path: Path + ): + resolver = SettlingResolver() + adapter, releases = settling_adapter(tmp_path, resolver) + first = await start_blocked_turn(adapter, resolver) + cleanup = asyncio.create_task(adapter.on_cleanup("room-123")) + await asyncio.wait_for(resolver.cancel_seen.wait(), 1) + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(run_turn(adapter, FakeAgentTools()), 1) + assert resolver.calls == 1 and launches.runtimes == [] + assert ( + "room-123" in adapter._workspaces.rooms + ) # settlement still owns the claim + + resolver.settle.set() + await asyncio.wait_for(cleanup, 1) + with pytest.raises(asyncio.CancelledError): + await first + assert resolver.cancellations == 1 and releases and all(releases) + await run_turn(adapter, FakeAgentTools()) # admission reopens after retirement + assert resolver.calls == 2 and len(launches.runtimes) == 1 + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_overlapping_cleanups_cancel_the_settling_turn_once( + self, launches: Launches, tmp_path: Path + ): + resolver = SettlingResolver() + adapter, releases = settling_adapter(tmp_path, resolver) + first = await start_blocked_turn(adapter, resolver) + room_cleanup = asyncio.create_task(adapter.on_cleanup("room-123")) + await asyncio.wait_for(resolver.cancel_seen.wait(), 1) + shutdown = await start_cleanup(adapter, shutdown=True) + # Shutdown stays joined to the running retirement, not finishing or re-cancelling. + + assert not shutdown.done() and resolver.cancellations == 1 + assert "room-123" in adapter._workspaces.rooms + resolver.settle.set() + await asyncio.wait_for(asyncio.gather(room_cleanup, shutdown), 1) + + assert resolver.cancellations == 1 and resolver.settled + assert releases and all(releases) and not adapter._workspaces.rooms + assert launches.runtimes == [] and first.cancelled() + + @pytest.mark.asyncio + async def test_cancelling_a_cleanup_waiter_leaves_retirement_running( + self, launches: Launches, tmp_path: Path + ): + resolver = SettlingResolver() + adapter, releases = settling_adapter(tmp_path, resolver) + first = await start_blocked_turn(adapter, resolver) + waiter = asyncio.create_task(adapter.on_cleanup("room-123")) + await asyncio.wait_for(resolver.cancel_seen.wait(), 1) + + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + + assert resolver.cancellations == 1 + assert "room-123" in adapter._workspaces.rooms + with pytest.raises(asyncio.CancelledError): # admission stays closed + await asyncio.wait_for(run_turn(adapter, FakeAgentTools()), 1) + resolver.settle.set() + await asyncio.wait_for(adapter.cleanup_all(), 1) # joins the same retirement + assert resolver.cancellations == 1 and releases and all(releases) + assert first.cancelled() and launches.runtimes == [] + + @pytest.mark.asyncio + async def test_cleanup_all_joins_a_retirement_blocked_in_native_close( + self, launches: Launches, tmp_path: Path + ): + adapter = OmpAdapter(OmpAdapterConfig(cwd=str(tmp_path)), emit=()) + await run_turn(adapter, FakeAgentTools()) # a settled turn leaves a live room + (runtime,) = launches.runtimes + closing = asyncio.Event() + finish_close = asyncio.Event() + closes = 0 + + async def close(): + nonlocal closes + closes += 1 + closing.set() + await finish_close.wait() + runtime.closed = True + + runtime.close = close # type: ignore[method-assign] + room_cleanup = asyncio.create_task(adapter.on_cleanup("room-123")) + await asyncio.wait_for(closing.wait(), 1) + shutdown = asyncio.create_task(adapter.cleanup_all()) + with pytest.raises(asyncio.CancelledError): # no new turn during teardown + await asyncio.wait_for(run_turn(adapter, FakeAgentTools()), 1) + + assert not shutdown.done() and "room-123" in adapter._workspaces.rooms + finish_close.set() + await asyncio.wait_for(asyncio.gather(room_cleanup, shutdown), 1) + + assert closes == 1 and runtime.closed + assert len(launches.runtimes) == 1 and not adapter._workspaces.rooms + await run_turn(adapter, FakeAgentTools()) + assert len(launches.runtimes) == 2 + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_cleanup_joins_a_turn_already_settling_an_external_cancellation( + self, launches: Launches, tmp_path: Path + ): + resolver = SettlingResolver() + adapter, releases = settling_adapter(tmp_path, resolver) + first = await start_blocked_turn(adapter, resolver) + first.cancel() + await asyncio.wait_for(resolver.cancel_seen.wait(), 1) + + cleanup = await start_cleanup(adapter) + + assert not cleanup.done() + assert resolver.cancellations == 1 and "room-123" in adapter._workspaces.rooms + resolver.settle.set() + await asyncio.wait_for(cleanup, 1) + assert resolver.cancellations == 1 and releases and all(releases) + assert first.cancelled() and launches.runtimes == [] + await run_turn(adapter, FakeAgentTools()) # explicit recovery + assert len(launches.runtimes) == 1 + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_cleanup_requested_by_the_resolving_turn_launches_nothing( + self, launches: Launches, tmp_path: Path + ): + adapter: OmpAdapter + + async def resolve(_workspace: str): + if not launches.runtimes and resolve.calls == 0: + resolve.calls += 1 + await adapter.on_cleanup("room-123") # the room's own turn retires it + return {"linear": remote("late")} + + resolve.calls = 0 # type: ignore[attr-defined] + adapter = OmpAdapter( + OmpAdapterConfig(cwd=str(tmp_path)), resolve_mcp_servers=resolve, emit=() + ) + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(run_turn(adapter, FakeAgentTools()), 1) + await asyncio.wait_for(adapter.cleanup_all(), 1) + + assert launches.runtimes == [] and not adapter._workspaces.rooms + await run_turn(adapter, FakeAgentTools()) # explicit recovery + assert len(launches.runtimes) == 1 + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_retirement_waits_for_the_turn_not_for_its_callers_other_work( + self, launches: Launches, tmp_path: Path + ): + resolver = SettlingResolver() + resolver.settle.set() # cancellation settles immediately + adapter, releases = settling_adapter(tmp_path, resolver) + turn_returned = asyncio.Event() + + async def caller(): + try: + await run_turn(adapter, FakeAgentTools()) + except asyncio.CancelledError: + turn_returned.set() + await asyncio.Event().wait() # unrelated work keeps the caller alive + + wrapper = asyncio.create_task(caller()) + await asyncio.wait_for(resolver.entered.wait(), 1) + + await asyncio.wait_for(adapter.on_cleanup("room-123"), 1) + + assert turn_returned.is_set() and not wrapper.done() + assert resolver.cancellations == 1 and releases and all(releases) + assert not adapter._workspaces.rooms and launches.runtimes == [] + wrapper.cancel() + await run_turn(adapter, FakeAgentTools()) # explicit recovery + assert len(launches.runtimes) == 1 + await adapter.cleanup_all() + + @pytest.mark.asyncio + async def test_cleanup_all_joins_a_retirement_with_no_room_or_turn_left( + self, launches: Launches, tmp_path: Path + ): + adapter = OmpAdapter(OmpAdapterConfig(cwd=str(tmp_path)), emit=()) + room_lock = SignallingLock() + adapter._locks["room-123"] = room_lock + await room_lock.acquire() # another holder keeps the retirement from closing + retirement = asyncio.create_task(adapter.on_cleanup("room-123")) + await asyncio.wait_for(room_lock.contended.wait(), 1) + assert not adapter._rooms and not adapter._owners + + shutdown = await start_cleanup(adapter, shutdown=True) + + assert not shutdown.done() # it joined the in-flight retirement + room_lock.release() + await asyncio.wait_for(asyncio.gather(retirement, shutdown), 1) + assert launches.runtimes == [] + + @pytest.mark.asyncio + async def test_retirement_never_cancels_work_after_the_captured_turn_finished( + self, tmp_path: Path, monkeypatch + ): + adapter = OmpAdapter(OmpAdapterConfig(cwd=str(tmp_path)), emit=()) + unrelated = asyncio.Event() + continued = asyncio.Event() + runtime = PendingRuntime() + monkeypatch.setattr("band.adapters.omp.OmpRpcClient", lambda *a, **k: runtime) + + async def caller(): + await run_turn(adapter, FakeAgentTools()) + await unrelated.wait() # the caller's own work after its turn + continued.set() + + wrapper = asyncio.create_task(caller()) + await asyncio.wait_for(runtime.yielded.wait(), 1) + # Cleanup captures the live turn; the turn then finishes before retirement runs. + cleanup = asyncio.create_task(adapter.on_cleanup("room-123")) + runtime.finish.set() + await asyncio.wait_for(cleanup, 1) + + assert not wrapper.cancelled() and not wrapper.done() + assert runtime.closed and not adapter._workspaces.rooms + unrelated.set() + await asyncio.wait_for(wrapper, 1) + assert continued.is_set() + + @pytest.mark.asyncio + async def test_cleanup_joins_outcome_handling_after_the_deadline_context_exited( + self, launches: Launches, tmp_path: Path, monkeypatch + ): + closing = asyncio.Event() + finish_close = asyncio.Event() + + runtime = failing_runtime(closing, finish_close) + monkeypatch.setattr("band.adapters.omp.OmpRpcClient", lambda *a, **k: runtime) + adapter = OmpAdapter(OmpAdapterConfig(cwd=str(tmp_path)), emit=()) + tools = FakeAgentTools() + turn = asyncio.create_task(run_turn(adapter, tools)) + await asyncio.wait_for( + closing.wait(), 1 + ) # the turn is handling its own failure + + cleanup = await start_cleanup(adapter) + assert not cleanup.done() + finish_close.set() + await asyncio.wait_for(cleanup, 1) + + with pytest.raises(RuntimeError, match="native turn failed"): + await turn + assert [failure["provider"] for failure in reported_failures(tools)] == ["omp"] + assert runtime.closed and not adapter._workspaces.rooms From 34c1d7efa8b18f8d90f4b4f2c5286411bf766230 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 06:58:32 +0300 Subject: [PATCH 2/3] test: model native room retirement and recovery --- tests/adapters/omp/lifecycle.py | 178 +++++++++++++++++++++ tests/adapters/omp/test_lifecycle_model.py | 25 +++ 2 files changed, 203 insertions(+) create mode 100644 tests/adapters/omp/lifecycle.py create mode 100644 tests/adapters/omp/test_lifecycle_model.py diff --git a/tests/adapters/omp/lifecycle.py b/tests/adapters/omp/lifecycle.py new file mode 100644 index 000000000..4320d4ccb --- /dev/null +++ b/tests/adapters/omp/lifecycle.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from enum import StrEnum, auto +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import AsyncMock + +from pytest import MonkeyPatch, raises + +from band.adapters.omp import OmpAdapter +from band.testing import FakeAgentTools, reported_failures +from tests.adapters.omp.kit import Launches, PendingRuntime, SignallingLock, run_turn + +ROOM = "room-123" + + +class Step(StrEnum): + TURN = auto() + CLEANUP = auto() + FAILED_CLOSE = auto() + RETIRE_QUEUED = auto() + CANCEL_CLEANUP_WAITER = auto() + + +class HeldClose: + def __init__(self, close: Callable[[], Awaitable[None]]) -> None: + self.close = close + self.entered = asyncio.Event() + self.release = asyncio.Event() + + async def __call__(self) -> None: + self.entered.set() + await self.release.wait() + await self.close() + + +class RoomLifecycle: + def __init__(self, workspace: Path, launches: Launches) -> None: + self.launches = launches + self.adapter = OmpAdapter( + workspace_for_room=lambda _room: str(workspace), emit=() + ) + self.lock = SignallingLock() + self.adapter._locks[ROOM] = self.lock + self.expected: PendingRuntime | None = None + self.retiring = False + + async def take(self, step: Step) -> None: + match step: + case Step.TURN: + await self.turn() + case Step.CLEANUP: + await self.cleanup() + case Step.FAILED_CLOSE: + await self.failed_close() + case Step.RETIRE_QUEUED | Step.CANCEL_CLEANUP_WAITER: + await self.retire_queued( + cancel_waiter=step == Step.CANCEL_CLEANUP_WAITER + ) + self.assert_custody() + + async def turn(self) -> None: + previous = self.expected + fresh = previous is None or self.retiring + tools = FakeAgentTools() + await run_turn(self.adapter, tools) + assert tools.turn.replied + assert not reported_failures(tools) + current = self.launches.runtimes[-1] + if fresh: + assert current is not previous + if previous is not None: + assert previous.closed + else: + assert current is previous + self.expected = current + self.retiring = False + + async def cleanup(self) -> None: + await self.adapter.cleanup_all() + self.expected = None + self.retiring = False + + async def failed_close(self) -> None: + if self.expected is None or self.retiring: + await self.turn() + runtime = self.expected + assert runtime is not None + close = runtime.close + runtime.close = AsyncMock(side_effect=OSError("process still alive")) + try: + with raises(OSError, match="still alive"): + await self.adapter.on_cleanup(ROOM) + finally: + runtime.close = close + self.retiring = True + + async def retire_queued(self, *, cancel_waiter: bool) -> None: + if self.expected is None or self.retiring: + await self.turn() + runtime = self.expected + assert runtime is not None + runtime.finish.clear() + runtime.yielded.clear() + self.lock.contended.clear() + close = runtime.close + held = HeldClose(close) + runtime.close = held # type: ignore[method-assign] + tasks: list[asyncio.Task[None]] = [] + try: + active = asyncio.create_task(run_turn(self.adapter, FakeAgentTools())) + tasks.append(active) + await runtime.yielded.wait() + queued_tools = FakeAgentTools() + queued = asyncio.create_task(run_turn(self.adapter, queued_tools)) + tasks.append(queued) + await self.lock.contended.wait() + cleanup = asyncio.create_task(self.adapter.on_cleanup(ROOM)) + tasks.append(cleanup) + await held.entered.wait() + assert not runtime.closed + assert self.launches.plugins[-1].exists() + if cancel_waiter: + cleanup.cancel() + with raises(asyncio.CancelledError): + await cleanup + assert not runtime.closed + assert self.launches.plugins[-1].exists() + held.release.set() + for turn in (active, queued): + with raises(asyncio.CancelledError): + await turn + assert not queued_tools.turn.replied + if not cancel_waiter: + await cleanup + await self.cleanup() + finally: + held.release.set() + runtime.close = close + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await self.adapter.cleanup_all() + + def assert_custody(self) -> None: + if self.expected is None: + assert not self.adapter._rooms + assert not self.adapter._workspaces.rooms + assert all(runtime.closed for runtime in self.launches.runtimes) + assert all(not overlay.exists() for overlay in self.launches.plugins) + return + assert self.adapter._rooms[ROOM].runtime is self.expected + assert self.launches.plugins[-1].exists() + with raises(ValueError): + self.adapter._workspaces.claim("another-room") + assert all(runtime.closed for runtime in self.launches.runtimes[:-1]) + assert all(not overlay.exists() for overlay in self.launches.plugins[:-1]) + + +async def play(steps: list[Step]) -> None: + with ( + TemporaryDirectory(prefix="band-omp-model-") as directory, + MonkeyPatch.context() as patch, + ): + launches = Launches() + patch.setattr("band.adapters.omp.OmpRpcClient", launches) + model = RoomLifecycle(Path(directory), launches) + try: + async with asyncio.timeout(15): + for step in steps: + await model.take(step) + await model.cleanup() + model.assert_custody() + finally: + await model.adapter.cleanup_all() diff --git a/tests/adapters/omp/test_lifecycle_model.py b/tests/adapters/omp/test_lifecycle_model.py new file mode 100644 index 000000000..74ca1cc69 --- /dev/null +++ b/tests/adapters/omp/test_lifecycle_model.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +import asyncio + +from hypothesis import Phase, example, given, settings +from hypothesis import strategies as st + +from tests.adapters.omp.lifecycle import Step, play + + +class TestRoomLifecycleModel: + @settings( + max_examples=20, + deadline=None, + phases=(Phase.explicit, Phase.reuse, Phase.generate), + ) + @given(steps=st.lists(st.sampled_from(Step), min_size=1, max_size=10)) + @example(steps=[Step.TURN, Step.FAILED_CLOSE, Step.TURN, Step.CLEANUP]) + @example( + steps=[Step.RETIRE_QUEUED, Step.TURN, Step.CANCEL_CLEANUP_WAITER, Step.TURN] + ) + def test_sequences_preserve_runtime_and_workspace_custody( + self, steps: list[Step] + ) -> None: + asyncio.run(play(steps)) From 01e20db29188c7d919168547c52063d58195a61d Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Thu, 22 Oct 2026 08:25:49 +0300 Subject: [PATCH 3/3] test: remove obsolete room retirement mock import --- tests/adapters/omp/test_room_retirement.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/adapters/omp/test_room_retirement.py b/tests/adapters/omp/test_room_retirement.py index a603652a1..84f2581d5 100644 --- a/tests/adapters/omp/test_room_retirement.py +++ b/tests/adapters/omp/test_room_retirement.py @@ -2,7 +2,6 @@ import asyncio from pathlib import Path -from unittest.mock import AsyncMock import pytest