feat: resolve room-scoped MCP servers at native startup

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

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",
+]