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/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_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_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)) 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..84f2581d5 --- /dev/null +++ b/tests/adapters/omp/test_room_retirement.py @@ -0,0 +1,322 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path + +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