diff --git a/docs/acp.md b/docs/acp.md index 2ec237dba..268b17c11 100644 --- a/docs/acp.md +++ b/docs/acp.md @@ -160,8 +160,9 @@ assert slow.retention.transition_timeout_s == 120 than twice on the same one. Only a bootstrap resume candidate is closed when its configuration is refused. - **Closure is confirmed, not assumed.** The process handle stays on the room until - `__aexit__` succeeds and the agent's stdout reaches EOF, so a descendant that - inherited it keeps the room closing until it exits, with a warning naming the + `__aexit__` succeeds, its child is reaped and the agent's stdout reaches EOF. + A child left running by a successful context exit is killed and waited for; + a descendant that inherited stdout keeps the room closing until it exits, with a warning naming the agent's pid. A failed or cancelled exit leaves the room owned, and a retry kills and waits for the same process. `cleanup_all()` and its alias `stop()` share one `transition_timeout_s` budget across the rooms and the shared Band MCP server, @@ -176,6 +177,39 @@ assert slow.retention.transition_timeout_s == 120 permission or extension requests carry no prompt id, so a call arriving after its turn is served with the room's current tools. +## Per-room MCP servers + +`resolve_mcp_servers(workspace)`, a keyword-only constructor argument of +`ACPClientAdapter`, `CopilotACPAdapter` and `CursorACPAdapter`, returns native +`acp.schema` descriptors for that room's workspace. +It runs once per session setup after the agent initializes. A reused live session +keeps its servers; a restored session and its fresh-session fallback share one +deep-copied snapshot. Recovery and a deliberate wake resolve anew on the new +connection, as does a new session replacing a stale Band MCP endpoint. + +The result replaces `config.mcp_servers`, including when empty, except Band routing: +a configured `band` entry and the injected Band server remain after the resolved +external servers. Callback failure never falls back to configured external servers. + +- `McpServerStdio` is always accepted. `HttpMcpServer` and `SseMcpServer` require the + agent to explicitly advertise that transport in `initialize`. `AcpMcpServer` is + refused. Names must be unique and may not be `band`. +- A refused descriptor fails the turn with `ACPMcpServerResolutionError`, naming + only its server and transport, never its URL, headers, command or environment. + Resolution failures recover the room's process through its retained owner; the + next turn resolves again on a fresh connection. +- Queued turns share the room's initializer. One cancelled waiter leaves it for + another; the last waiter cancels it once without waiting for callback settlement. + The initializer's retained lease holds the process and workspace until it settles. + A callback that swallows cancellation or requests teardown cannot publish a session + or pass its late result to the agent. +- A callback requesting `on_cleanup`, `cleanup_all` or `stop` gets + `CleanupIncompleteError` rather than waiting on itself. The retained owner continues + teardown after the callback unwinds, preserving the process, workspace and Band + MCP server until closure is confirmed. +- There is no resolver timeout: a callback that never returns holds that room's + setup, while other rooms continue independently. + ## Session config `resolve_session_config` runs after each new or restored session and before its first diff --git a/src/band/adapters/copilot_acp.py b/src/band/adapters/copilot_acp.py index f1557fb68..0535ab142 100644 --- a/src/band/adapters/copilot_acp.py +++ b/src/band/adapters/copilot_acp.py @@ -34,6 +34,7 @@ from band.integrations.acp.client_adapter import ( ACPClientAdapter, ACPClientAdapterConfig, + McpServerResolver, PermissionResolver, ) from band.integrations.acp.session_config import SessionConfigResolver @@ -80,6 +81,7 @@ def __init__( workspace_for_room: WorkspaceResolver | None = None, resolve_session_config: SessionConfigResolver | None = None, resolve_permission: PermissionResolver | None = None, + resolve_mcp_servers: McpServerResolver | None = None, **features: Unpack[FeatureKwargs], ) -> None: """Bridge Band rooms to ``copilot --acp``. @@ -92,6 +94,8 @@ def __init__( resolve_session_config: Picks session config options; exclusive with ``config.model`` and ``config.reasoning_effort``. resolve_permission: Chooses a permission option per tool call. + resolve_mcp_servers: Returns native MCP server descriptors once per + room session setup; see :class:`ACPClientAdapter`. """ config = config or CopilotACPAdapterConfig() super().__init__( @@ -100,6 +104,7 @@ def __init__( workspace_for_room=workspace_for_room, resolve_session_config=resolve_session_config, resolve_permission=resolve_permission, + resolve_mcp_servers=resolve_mcp_servers, **features, ) diff --git a/src/band/adapters/cursor_acp.py b/src/band/adapters/cursor_acp.py index bcf3e249a..862a16847 100644 --- a/src/band/adapters/cursor_acp.py +++ b/src/band/adapters/cursor_acp.py @@ -28,6 +28,7 @@ ACPClientAdapter, ACPClientAdapterConfig, ACPPermissionRequest, + McpServerResolver, ) from band.integrations.acp.client_profiles import ( CURSOR_ASK_QUESTION_METHOD, @@ -207,6 +208,7 @@ def __init__( additional_tools: list[CustomToolDef] | None = None, workspace_for_room: WorkspaceResolver | None = None, resolve_session_config: SessionConfigResolver | None = None, + resolve_mcp_servers: McpServerResolver | None = None, **features: Unpack[FeatureKwargs], ) -> None: """Bridge Band rooms to Cursor's ``agent acp``. @@ -217,6 +219,8 @@ def __init__( workspace_for_room: Maps a room id to its absolute workspace; exclusive with ``config.cwd``. resolve_session_config: Picks session config options. + resolve_mcp_servers: Returns native MCP server descriptors once per + room session setup; see :class:`ACPClientAdapter`. """ config = config or CursorACPAdapterConfig() # Each room's one accepted turn, from acceptance until it actually @@ -231,6 +235,7 @@ def __init__( additional_tools=additional_tools, workspace_for_room=workspace_for_room, resolve_session_config=resolve_session_config, + resolve_mcp_servers=resolve_mcp_servers, resolve_permission=self._resolve_cursor_permission, **features, ) diff --git a/src/band/integrations/acp/client_adapter.py b/src/band/integrations/acp/client_adapter.py index 24e6dc98f..e84e8a37c 100644 --- a/src/band/integrations/acp/client_adapter.py +++ b/src/band/integrations/acp/client_adapter.py @@ -179,6 +179,13 @@ class ACPTurnTimeoutError(TimeoutError): LocalMcpServerConfig = HttpMcpServer | SseMcpServer # What ACP's session/new takes; YAML/JSON entries validate into these. SessionMcpServer = HttpMcpServer | SseMcpServer | AcpMcpServer | McpServerStdio +ResolvedMcpServer: TypeAlias = HttpMcpServer | SseMcpServer | McpServerStdio +McpServerResolver: TypeAlias = Callable[[str], Awaitable[Sequence[ResolvedMcpServer]]] + + +class ACPMcpServerResolutionError(ValueError): + """The room's resolved MCP servers cannot be used by its ACP agent.""" + # Prefixes the change-triggered roster/contacts updates injected into a # prompt, so the model reads them as platform state, not as the requester @@ -408,6 +415,7 @@ def __init__( resolve_session_config: SessionConfigResolver | None = None, resolve_permission: PermissionResolver | None = None, client_capabilities: ClientCapabilities | None = None, + resolve_mcp_servers: McpServerResolver | None = None, spawn_process: SpawnProcess | None = None, **features: Unpack[FeatureKwargs], ) -> None: @@ -425,6 +433,10 @@ def __init__( resolve_permission: Chooses a permission option per tool call; ``None`` approves with the agent's allow option. client_capabilities: Capabilities advertised at ``initialize``. + resolve_mcp_servers: Returns native MCP descriptors for the room's + workspace once per session setup, after initialize. Replaces + configured external servers, keeping Band routing. Stdio is + always accepted; HTTP/SSE require explicit agent support. spawn_process: Rejected: a custom transport cannot guarantee one process per room. """ @@ -450,6 +462,7 @@ def __init__( self._resolve_session_config = resolve_session_config self._resolve_permission = resolve_permission self._client_capabilities = client_capabilities + self._resolve_mcp_servers = resolve_mcp_servers self._workspaces = RoomWorkspaces(self._workspace_for_room) self._retention = RoomRetention( config.retention, settle_grace_s=config.turn_settle_timeout_s @@ -748,15 +761,34 @@ async def bind_tools() -> None: room.tools = tools room.retained.require_admission() - async with room.reservation, AsyncExitStack() as lease: - try: - await lease.enter_async_context(room.retained.use(prepare=bind_tools)) - except TimeoutError as error: - await self._report_not_ready(room, tools, error) - raise - if room.runtime is None: - raise RuntimeError(f"Room {room.room_id} has no ACP process") - yield room.runtime + room.session_waiters += 1 + try: + async with room.reservation, AsyncExitStack() as lease: + try: + await lease.enter_async_context( + room.retained.use(prepare=bind_tools) + ) + except TimeoutError as error: + await self._report_not_ready(room, tools, error) + raise + if room.runtime is None: + raise RuntimeError(f"Room {room.room_id} has no ACP process") + yield room.runtime + finally: + room.session_waiters -= 1 + setup = room.session_setup + if ( + not room.session_waiters + and not room.retained.closing + and setup is not None + and not setup.done() + and setup is not room.abandoned_setup + ): + # A setup already settling its own cancellation is abandoned + # too: the callback's timeout may clear that request later. + room.abandoned_setup = setup + if not setup.cancelling(): + setup.cancel() @staticmethod async def _report_not_ready( @@ -1261,27 +1293,80 @@ async def _get_or_create_session( is reloaded when the agent still has it. A fresh session holds no conversation context; the caller owes it a transcript replay. """ + self._require_current_runtime(room, runtime) mcp = await self._session_mcp_servers(room, runtime) + self._require_current_runtime(room, runtime) session = room.session - restore_id: str | None = None if session is not None and session.live: if session.band_url == mcp.band_url: return session.session_id - # Never restore: the retained id is the one being retired. room.session = None + history = None self._retire_stale_session(room, runtime, session) - elif session is not None: - restore_id = session.session_id - elif history is not None: - restore_id = history.room_to_session.get(room.room_id) - - restored = None - if restore_id is not None: - restored = await self._restore_session( - room, runtime, restore_id, mcp, retained=session is not None + setup = room.session_setup + if setup is not None and (setup is room.abandoned_setup or setup.cancelling()): + # A callback may still be settling an abandoned setup. Its lease + # holds the process, and no successor resolves beside it. + await asyncio.wait((setup,)) + if room.abandoned_setup is setup: + room.abandoned_setup = None + self._require_current_runtime(room, runtime) + # A setup that was only settling its own cancellation may still + # have published a session for the waiters it kept serving. + joined = room.session + if joined is not None and joined.live and joined.band_url == mcp.band_url: + return joined.session_id + setup = None + if setup is None or setup.done(): + setup = room.retained.start_owned( + partial(self._initialize_session, room, runtime, history, mcp) ) - room.session = restored or await self._create_session(room, runtime, mcp) - return room.session.session_id + room.session_setup = setup + return await asyncio.shield(setup) + + def _require_current_runtime(self, room: ACPRoom, runtime: ACPRuntime) -> None: + task = asyncio.current_task() + if ( + self._rooms.get(room.room_id) is not room + or room.runtime is not runtime + or runtime.closed + or not room.retained.has_authority(room.generation) + or ( + task is not None and (task.cancelling() or task is room.abandoned_setup) + ) + ): + raise asyncio.CancelledError + + async def _initialize_session( + self, + room: ACPRoom, + runtime: ACPRuntime, + history: ACPClientSessionState | None, + mcp: SessionMcpServers, + ) -> str: + # The setup has its own lease: losing a waiter must not close the + # process under another waiter or an abandoned callback's settlement. + async with room.retained.use(): + self._require_current_runtime(room, runtime) + mcp = await self._with_resolved_mcp_servers(room, runtime, mcp) + self._require_current_runtime(room, runtime) + session = room.session + restore_id = ( + session.session_id + if session is not None + else None + if history is None + else history.room_to_session.get(room.room_id) + ) + restored = None + if restore_id is not None: + restored = await self._restore_session( + room, runtime, restore_id, mcp, retained=session is not None + ) + initialized = restored or await self._create_session(room, runtime, mcp) + self._require_current_runtime(room, runtime) + room.session = initialized + return initialized.session_id async def _restore_session( self, @@ -1316,6 +1401,7 @@ async def _restore_session( try: await self._configure_session(runtime, room.room_id, session_id) + self._require_setup_result(room) except BaseException as error: if retained: room.retained.recover( @@ -1343,6 +1429,8 @@ async def _restore_retained_session( back to a new session or a transcript replay. """ mcp = await self._session_mcp_servers(room, runtime) + mcp = await self._with_resolved_mcp_servers(room, runtime, mcp) + self._require_setup_result(room) await runtime.restore_session_response( cwd=room.workspace, session_id=session.session_id, @@ -1351,6 +1439,7 @@ async def _restore_retained_session( await self._configure_session( runtime, room.room_id, session.session_id, keep_refused_selection=True ) + self._require_setup_result(room) session.band_url = mcp.band_url session.live = True @@ -1360,6 +1449,7 @@ async def _create_session( """Create and configure a session for one room.""" async with self._fresh_session(room, runtime, mcp.servers) as session: await self._configure_session(runtime, room.room_id, session.session_id) + self._require_setup_result(room) logger.info( "Created ACP session %s for room %s (mcp_servers=%d)", @@ -1454,6 +1544,76 @@ async def _session_mcp_servers( servers=[*servers, band_server], band_url=band_server.url ) + @staticmethod + def _require_setup_result(room: ACPRoom) -> None: + # Callbacks may swallow cancellation or request their own teardown. + # Neither result may reach the agent or publish a session afterwards. + task = asyncio.current_task() + if room.retained.closing or ( + task is not None and (task.cancelling() or task is room.abandoned_setup) + ): + raise asyncio.CancelledError + + async def _with_resolved_mcp_servers( + self, room: ACPRoom, runtime: ACPRuntime, mcp: SessionMcpServers + ) -> SessionMcpServers: + resolver = self._resolve_mcp_servers + if resolver is None: + return mcp + resolved = await resolver(room.workspace) + self._require_setup_result(room) + if isinstance(resolved, str | bytes) or not isinstance(resolved, Sequence): + raise TypeError("resolve_mcp_servers must return a sequence of MCP servers") + reserved = {BAND_MCP_SERVER_NAME} + external: list[object] = [] + for value in resolved: + server = self._resolved_mcp_server(runtime, value, reserved) + reserved.add(server.name) + # A restore and its new-session fallback share this snapshot even + # when the callback's owner mutates its original nested values. + external.append(server.model_copy(deep=True)) + band = [ + server + for server in mcp.servers + if getattr(server, "name", None) == BAND_MCP_SERVER_NAME + ] + return SessionMcpServers(servers=[*external, *band], band_url=mcp.band_url) + + @staticmethod + def _resolved_mcp_server( + runtime: ACPRuntime, server: object, reserved: set[str] + ) -> ResolvedMcpServer: + match server: + case McpServerStdio(): + transport, supported = "stdio", True + case HttpMcpServer(): + transport, supported = "http", runtime.agent_mcp_http + case SseMcpServer(): + transport, supported = "sse", runtime.agent_mcp_sse + case AcpMcpServer(): + raise ACPMcpServerResolutionError( + f"MCP server {server.name!r}: resolve_mcp_servers does not " + "accept the acp transport" + ) + case _: + raise TypeError( + "resolve_mcp_servers returned a " + f"{type(server).__name__}, not an MCP server descriptor" + ) + # Descriptor values can contain credentials; diagnostics disclose + # only the server's name and the refused transport. + if not supported: + raise ACPMcpServerResolutionError( + f"MCP server {server.name!r}: the ACP agent does not advertise " + f"{transport} MCP servers" + ) + if server.name in reserved: + raise ACPMcpServerResolutionError( + f"MCP server {server.name!r} ({transport}) duplicates a server " + "name already in this session" + ) + return server + async def _configure_session( self, runtime: ACPRuntime, diff --git a/src/band/integrations/acp/client_runtime.py b/src/band/integrations/acp/client_runtime.py index acae4704e..0a4450069 100644 --- a/src/band/integrations/acp/client_runtime.py +++ b/src/band/integrations/acp/client_runtime.py @@ -114,6 +114,8 @@ def __init__( self._exit_attempted = False self._closure: asyncio.Task[None] | None = None self._stderr = ACPStderrDrain(logger) + self._agent_mcp_http = False + self._agent_mcp_sse = False async def start(self) -> None: """Spawn the ACP agent subprocess and initialize its connection.""" @@ -167,6 +169,8 @@ async def _initialize_connection(self, conn: ACPConnectionProtocol) -> None: init_kwargs["client_capabilities"] = self._client_capabilities init_response = await cast(Any, conn).initialize(**init_kwargs) self._agent_mcp_transport = self._select_mcp_transport(init_response) + self._agent_mcp_http = self._advertised_mcp(init_response, "http") + self._agent_mcp_sse = self._advertised_mcp(init_response, "sse") self._agent_supports_session_load = self._select_session_load(init_response) self._agent_supports_session_close = self._select_session_close(init_response) if self._auth_method: @@ -174,7 +178,7 @@ async def _initialize_connection(self, conn: ACPConnectionProtocol) -> None: logger.info("Authenticated with method: %s", self._auth_method) async def ensure_connection(self) -> ACPConnectionProtocol: - if self._conn is None: + if self._conn is None or self._closure is not None: raise RuntimeError("ACP client connection is not open") return self._conn @@ -231,7 +235,9 @@ async def _exit_transport(self) -> None: "ACP agent subprocess %s failed to exit", _pid(process) ) raise - elif isinstance(process, asyncio.subprocess.Process): + # Even a context that returned successfully can leave a child + # running; its process handle, not that return, confirms exit. + if isinstance(process, asyncio.subprocess.Process): await _kill(process) if isinstance(process, asyncio.subprocess.Process) and process.stdout: await _read_to_eof(process.stdout) @@ -245,6 +251,8 @@ def _forget_connection(self) -> None: self._closure = None self._agent_supports_session_load = False self._agent_supports_session_close = False + self._agent_mcp_http = False + self._agent_mcp_sse = False async def _close_failed_start(self) -> None: """Close what a failed start opened; a failed close keeps it for @@ -271,6 +279,22 @@ def _select_mcp_transport(init_response: object) -> BandMCPTransport: return BandMCPTransport.HTTP + @property + def agent_mcp_http(self) -> bool: + """The connected agent explicitly advertised HTTP MCP servers.""" + return self._agent_mcp_http + + @property + def agent_mcp_sse(self) -> bool: + """The connected agent explicitly advertised SSE MCP servers.""" + return self._agent_mcp_sse + + @staticmethod + def _advertised_mcp(init_response: object, transport: str) -> bool: + capabilities = getattr(init_response, "agent_capabilities", None) + mcp_capabilities = getattr(capabilities, "mcp_capabilities", None) + return getattr(mcp_capabilities, transport, None) is True + @staticmethod def _select_session_load(init_response: object) -> bool: capabilities = getattr(init_response, "agent_capabilities", None) diff --git a/src/band/integrations/acp/client_types.py b/src/band/integrations/acp/client_types.py index fa8b30311..75805bb7b 100644 --- a/src/band/integrations/acp/client_types.py +++ b/src/band/integrations/acp/client_types.py @@ -72,6 +72,12 @@ class ACPRoom: from session setup through their actual end. ``session_closes`` are best-effort closes of retired sessions, which finish before the process is released. + ``session_setup`` is the shared, room-owned initializer; its lease keeps + cancellation settlement ahead of process closure. ``abandoned_setup`` is + the initializer its last waiter left: a callback can clear its task's + cancellation with ``uncancel()``, so abandonment is recorded here, not + read back from the task. ``session_waiters`` counts prompts running or + queued on the room's reservation. """ room_id: str @@ -82,6 +88,9 @@ class ACPRoom: tools: AgentToolsProtocol | None = None reservation: asyncio.Lock = field(default_factory=asyncio.Lock) session_closes: set[asyncio.Task[None]] = field(default_factory=set) + session_setup: asyncio.Task[str] | None = None + abandoned_setup: asyncio.Task[str] | None = None + session_waiters: int = 0 @property def generation(self) -> int: diff --git a/tests/integrations/acp/acp_toolkit/mcpresolver.py b/tests/integrations/acp/acp_toolkit/mcpresolver.py new file mode 100644 index 000000000..50b959038 --- /dev/null +++ b/tests/integrations/acp/acp_toolkit/mcpresolver.py @@ -0,0 +1,118 @@ +"""Native MCP descriptors and callback barriers for ACP resolver tests.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable, Sequence +from typing import Any + +from acp.schema import ( + EnvVariable, + HttpHeader, + HttpMcpServer, + McpServerStdio, + SseMcpServer, +) + +SECRET = "s3cret-token" + + +def http_server( + name: str = "repo", url: str = "http://127.0.0.1:9/mcp" +) -> HttpMcpServer: + return HttpMcpServer( + type="http", + name=name, + url=url, + headers=[HttpHeader(name="Authorization", value=f"Bearer {SECRET}")], + ) + + +def sse_server(name: str = "events") -> SseMcpServer: + return SseMcpServer( + type="sse", + name=name, + url="http://127.0.0.1:9/sse", + headers=[HttpHeader(name="Authorization", value=f"Bearer {SECRET}")], + ) + + +def stdio_server(name: str = "tools") -> McpServerStdio: + return McpServerStdio( + name=name, + command="mcp-tools", + args=["mcp"], + env=[EnvVariable(name="TOKEN", value=SECRET)], + ) + + +def server_names(servers: Sequence[Any]) -> list[str]: + return [server.name for server in servers] + + +class Resolver: + """Record the workspace of each resolution and return scripted servers.""" + + def __init__(self, *servers: Any) -> None: + self.servers: list[Any] = list(servers) + self.workspaces: list[str] = [] + + async def __call__(self, workspace: str) -> Sequence[Any]: + self.workspaces.append(workspace) + return self.servers + + +class HeldResolver(Resolver): + """Hold the first resolution, then finish cancellation only when released. + + ``swallow`` clears the task's cancellation request, as a callback that + catches ``CancelledError`` and calls ``uncancel()`` would.""" + + def __init__( + self, + *, + self_cancel: bool = False, + wait_for_cancel: bool = False, + swallow: bool = False, + request: Callable[[], Awaitable[None]] | None = None, + ) -> None: + super().__init__(http_server()) + self.entered = asyncio.Event() + self.settling = asyncio.Event() + self.settle = asyncio.Event() + self.cancellations = 0 + self.log: list[str] = [] + self.self_cancel = self_cancel + self.wait_for_cancel = wait_for_cancel + self.request = request + self.swallow = swallow + + async def __call__(self, workspace: str) -> Sequence[Any]: + self.workspaces.append(workspace) + if len(self.workspaces) > 1: + self.log.append("second call") + return self.servers + task = asyncio.current_task() + assert task is not None + if self.self_cancel: + task.cancel() + self.entered.set() + if self.request is None or self.self_cancel or self.wait_for_cancel: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancellations += 1 + if self.swallow: + task.uncancel() + if self.request is not None: + await self.request() + self.log.append("request returned") + self.settling.set() + try: + await self.settle.wait() + except asyncio.CancelledError: + self.cancellations += 1 + self.log.append("settlement interrupted") + raise + self.log.append("settled") + return self.servers diff --git a/tests/integrations/acp/acp_toolkit/processreaping.py b/tests/integrations/acp/acp_toolkit/processreaping.py new file mode 100644 index 000000000..c1fd9eea4 --- /dev/null +++ b/tests/integrations/acp/acp_toolkit/processreaping.py @@ -0,0 +1,46 @@ +"""A spawn context that leaves a real child for the runtime to reap.""" + +from __future__ import annotations + +import asyncio +import sys +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager, suppress +from typing import Any + +from band.integrations.acp.client_runtime import ACPRuntime +from tests.integrations.acp.acp_toolkit.harness import make_acp_connection + + +@asynccontextmanager +async def leaky_runtime( + failure: BaseException | None = None, +) -> AsyncIterator[tuple[ACPRuntime, asyncio.subprocess.Process]]: + process = await asyncio.create_subprocess_exec( + sys.executable, + "-c", + "import time; time.sleep(60)", + stdin=asyncio.subprocess.DEVNULL, + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.PIPE, + ) + kill, wait = process.kill, process.wait + connection = make_acp_connection() + if failure is not None: + connection.initialize.side_effect = failure + + @asynccontextmanager + async def spawn(*args: Any, **kwargs: Any) -> AsyncIterator[tuple[Any, Any]]: + yield connection, process + + runtime = ACPRuntime(command=["agent"], spawn_process=spawn) + try: + yield runtime, process + finally: + process.kill = kill # type: ignore[method-assign] + process.wait = wait # type: ignore[method-assign] + if process.returncode is None: + with suppress(ProcessLookupError): + kill() + await wait() + await runtime.stop() diff --git a/tests/integrations/acp/clientadapter/test_mcp_capabilities.py b/tests/integrations/acp/clientadapter/test_mcp_capabilities.py new file mode 100644 index 000000000..021ffb812 --- /dev/null +++ b/tests/integrations/acp/clientadapter/test_mcp_capabilities.py @@ -0,0 +1,43 @@ +"""Remote MCP support comes only from explicit initialize advertisements.""" + +from __future__ import annotations + +import pytest +from acp.schema import AgentCapabilities, InitializeResponse, McpCapabilities + +from band.integrations.acp.client_runtime import ACPRuntime +from tests.integrations.acp.acp_toolkit.harness import FakeSpawn + + +class TestMcpCapabilities: + """Band's preferred transport is not an agent's advertised support.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("mcp", "http", "sse"), + [ + (None, False, False), + (McpCapabilities(), False, False), + (McpCapabilities(http=None, sse=None), False, False), + (McpCapabilities(http=False, sse=False), False, False), + (McpCapabilities(http=True), True, False), + (McpCapabilities(sse=True), False, True), + (McpCapabilities(http=True, sse=True), True, True), + ], + ) + async def test_support_is_only_what_the_agent_advertises( + self, mcp: McpCapabilities | None, http: bool, sse: bool + ) -> None: + spawn = FakeSpawn() + spawn.conn.initialize.return_value = InitializeResponse( + protocol_version=1, + agent_capabilities=AgentCapabilities(mcp_capabilities=mcp), + ) + runtime = ACPRuntime(command=["agent"], spawn_process=spawn) + try: + await runtime.start() + assert (runtime.agent_mcp_http, runtime.agent_mcp_sse) == (http, sse) + finally: + await runtime.stop() + + assert (runtime.agent_mcp_http, runtime.agent_mcp_sse) == (False, False) diff --git a/tests/integrations/acp/clientadapter/test_mcp_resolution.py b/tests/integrations/acp/clientadapter/test_mcp_resolution.py new file mode 100644 index 000000000..28da2fe13 --- /dev/null +++ b/tests/integrations/acp/clientadapter/test_mcp_resolution.py @@ -0,0 +1,211 @@ +"""Resolved native MCP descriptors replace external configuration safely.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from acp.schema import AcpMcpServer + +from band.adapters.copilot_acp import CopilotACPAdapter, CopilotACPAdapterConfig +from band.adapters.cursor_acp import CursorACPAdapter, CursorACPAdapterConfig +from band.integrations.acp.client_adapter import ACPMcpServerResolutionError +from band.runtime.tools import BAND_MCP_SERVER_NAME +from tests.integrations.acp.acp_toolkit import ( + FakeACPAgent, + acp_adapter, + fake_agent_config, + started_acp_adapter, +) +from tests.integrations.acp.acp_toolkit.mcpresolver import ( + SECRET, + Resolver, + http_server, + server_names, + sse_server, + stdio_server, +) +from tests.mcpbackends import backends_created_by + + +class TestMcpResolution: + """Only validated external descriptors replace the configured servers.""" + + @pytest.mark.asyncio + async def test_empty_result_is_authoritative(self) -> None: + agent = FakeACPAgent(http=True) + config = fake_agent_config(mcp_servers=[http_server(name="static")]) + + async with acp_adapter(agent, config, resolve_mcp_servers=Resolver()) as acp: + await acp.send("hello") + + assert agent.sessions[0]["mcp_servers"] == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize("resolved", [[], [stdio_server()]], ids=["empty", "one"]) + async def test_configured_band_routing_is_kept(self, resolved: list[Any]) -> None: + agent = FakeACPAgent(http=True) + band = http_server(name=BAND_MCP_SERVER_NAME, url="http://127.0.0.1:9/band") + config = fake_agent_config(mcp_servers=[http_server(name="static"), band]) + + async with acp_adapter( + agent, config, resolve_mcp_servers=Resolver(*resolved) + ) as acp: + await acp.send("hello") + + sent = agent.sessions[0]["mcp_servers"] + assert server_names(sent) == [*server_names(resolved), BAND_MCP_SERVER_NAME] + assert sent[-1].url == "http://127.0.0.1:9/band" + assert sent[-1].headers[0].value == f"Bearer {SECRET}" + + @pytest.mark.asyncio + async def test_injected_band_server_follows_resolved_servers(self) -> None: + agent = FakeACPAgent(http=True) + with backends_created_by(): + async with acp_adapter( + agent, + fake_agent_config(inject_band_tools=True), + resolve_mcp_servers=Resolver(http_server()), + ) as acp: + await acp.send("hello") + + assert server_names(agent.sessions[0]["mcp_servers"]) == [ + "repo", + BAND_MCP_SERVER_NAME, + ] + + @pytest.mark.asyncio + async def test_no_resolver_preserves_static_configuration(self) -> None: + agent = FakeACPAgent(http=False, sse=False).will_say("ok") + config = fake_agent_config(mcp_servers=[http_server(name="static")]) + + async with acp_adapter(agent, config) as acp: + reply = await acp.send("hello") + + assert reply.thoughts == ["ok"] + assert server_names(agent.sessions[0]["mcp_servers"]) == ["static"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("http", "sse", "server", "transport"), + [ + (False, False, http_server(), "http"), + (False, True, http_server(), "http"), + (True, False, sse_server(), "sse"), + ], + ) + async def test_unadvertised_transport_is_refused_without_secrets( + self, + http: bool, + sse: bool, + server: Any, + transport: str, + caplog: pytest.LogCaptureFixture, + ) -> None: + agent = FakeACPAgent(http=http, sse=sse) + async with acp_adapter(agent, resolve_mcp_servers=Resolver(server)) as acp: + with pytest.raises(ACPMcpServerResolutionError) as caught: + await acp.send("hello") + reported = acp.last_reply.errors + + message = str(caught.value) + assert server.name in message and transport in message + assert SECRET not in message and server.url not in message + assert SECRET not in caplog.text and server.url not in caplog.text + assert all( + SECRET not in error and server.url not in error for error in reported + ) + assert agent.sessions == [] + + @pytest.mark.asyncio + async def test_advertised_transports_and_stdio_are_accepted(self) -> None: + agent = FakeACPAgent(http=True, sse=True) + resolver = Resolver(http_server(), sse_server(), stdio_server()) + + async with acp_adapter(agent, resolve_mcp_servers=resolver) as acp: + await acp.send("hello") + + assert server_names(agent.sessions[0]["mcp_servers"]) == [ + "repo", + "events", + "tools", + ] + + @pytest.mark.asyncio + async def test_stdio_needs_no_remote_transport_advertisement(self) -> None: + agent = FakeACPAgent(http=False, sse=False) + async with acp_adapter( + agent, resolve_mcp_servers=Resolver(stdio_server()) + ) as acp: + await acp.send("hello") + + assert server_names(agent.sessions[0]["mcp_servers"]) == ["tools"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "servers", + [ + [http_server(name="repo"), stdio_server(name="repo")], + [http_server(name=BAND_MCP_SERVER_NAME)], + [AcpMcpServer(type="acp", name="peer", server_id="peer-1")], + ], + ids=["duplicate", "band", "acp-transport"], + ) + async def test_unacceptable_servers_fail_setup(self, servers: list[Any]) -> None: + agent = FakeACPAgent(http=True) + config = fake_agent_config(mcp_servers=[http_server(name="static")]) + + async with acp_adapter( + agent, config, resolve_mcp_servers=Resolver(*servers) + ) as acp: + with pytest.raises(ACPMcpServerResolutionError): + await acp.send("hello") + + assert agent.sessions == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "returned", ["repo", b"repo", [{"name": "repo"}], None, iter(())] + ) + async def test_non_descriptor_results_are_type_errors(self, returned: Any) -> None: + async def resolver(workspace: str) -> Any: + return returned + + agent = FakeACPAgent(http=True) + async with acp_adapter(agent, resolve_mcp_servers=resolver) as acp: + with pytest.raises(TypeError): + await acp.send("hello") + + assert agent.sessions == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize("adapter", ["copilot", "cursor"]) + async def test_backend_constructors_use_the_shared_resolver( + self, adapter: str + ) -> None: + resolver = Resolver(http_server()) + agent = FakeACPAgent(http=True) + match adapter: + case "copilot": + client = CopilotACPAdapter( + CopilotACPAdapterConfig( + command="fake-agent", inject_band_tools=False + ), + resolve_mcp_servers=resolver, + ) + case "cursor": + client = CursorACPAdapter( + CursorACPAdapterConfig( + command="fake-agent", inject_band_tools=False + ), + resolve_mcp_servers=resolver, + ) + case _: + raise AssertionError(adapter) + + async with started_acp_adapter(client, agent) as acp: + await acp.send("hello") + workspace = acp.room().workspace + + assert resolver.workspaces == [workspace] + assert server_names(agent.sessions[0]["mcp_servers"]) == ["repo"] diff --git a/tests/integrations/acp/clientadapter/test_mcp_snapshot.py b/tests/integrations/acp/clientadapter/test_mcp_snapshot.py new file mode 100644 index 000000000..b1bcc01c9 --- /dev/null +++ b/tests/integrations/acp/clientadapter/test_mcp_snapshot.py @@ -0,0 +1,164 @@ +"""Each ACP session setup consumes one owned MCP resolution snapshot.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from acp.schema import LoadSessionResponse + +from band.integrations.acp.client_types import ACPClientSessionState +from tests.integrations.acp.acp_toolkit import ( + DEFAULT_ROOM, + FakeACPAgent, + acp_adapter, + fake_agent_config, +) +from tests.integrations.acp.acp_toolkit.mcpresolver import ( + SECRET, + Resolver, + http_server, + server_names, + stdio_server, +) +from tests.integrations.acp.roomownership.kit import recover, woken +from tests.mcpclient import crash_backend + + +class TestMcpSnapshot: + """Reused sessions keep their servers; new setups resolve a fresh snapshot.""" + + @pytest.mark.asyncio + async def test_resolved_servers_replace_configured_ones_once_per_session( + self, + ) -> None: + agent = FakeACPAgent(http=True) + resolver = Resolver(http_server(), stdio_server()) + config = fake_agent_config(mcp_servers=[http_server(name="static")]) + + async with acp_adapter(agent, config, resolve_mcp_servers=resolver) as acp: + await acp.send("one") + await acp.send("two") + workspace = acp.room().workspace + + assert resolver.workspaces == [workspace] + assert len(agent.sessions) == 1 + sent = agent.sessions[0]["mcp_servers"] + assert server_names(sent) == ["repo", "tools"] + assert sent[0].headers[0].value == f"Bearer {SECRET}" + assert sent[1].env[0].value == SECRET + + @pytest.mark.asyncio + async def test_each_fresh_connection_resolves_anew(self) -> None: + agent = FakeACPAgent(http=True) + resolver = Resolver(http_server(url="http://127.0.0.1:9/grant-a")) + + async with acp_adapter(agent, resolve_mcp_servers=resolver) as acp: + await acp.send("one") + await acp.adapter.on_cleanup(DEFAULT_ROOM) + resolver.servers = [http_server(url="http://127.0.0.1:9/grant-b")] + await acp.send("two") + await acp.send("three") + + assert len(resolver.workspaces) == 2 + assert [s["mcp_servers"][0].url for s in agent.sessions] == [ + "http://127.0.0.1:9/grant-a", + "http://127.0.0.1:9/grant-b", + ] + + @pytest.mark.asyncio + async def test_load_and_fresh_fallback_share_one_deep_snapshot(self) -> None: + agent = FakeACPAgent(http=True, supports_session_load=True) + remote = http_server(url="http://127.0.0.1:9/original") + local = stdio_server() + resolver = Resolver(remote, local) + loaded_with: list[Any] = [] + load = agent.load_session + + async def load_session( + cwd: str, session_id: str, mcp_servers: Any = None, **kwargs: Any + ) -> LoadSessionResponse: + loaded_with.extend(mcp_servers or []) + remote.url = "http://127.0.0.1:9/changed" + remote.headers[0].value = "Bearer rotated" + local.args.append("--changed") + local.env[0].value = "rotated" + return await load(cwd, session_id, mcp_servers, **kwargs) + + agent.load_session = load_session # type: ignore[method-assign] + history = ACPClientSessionState(room_to_session={DEFAULT_ROOM: "gone"}) + + async with acp_adapter(agent, resolve_mcp_servers=resolver) as acp: + await acp.send("hello", history=history, bootstrap=True) + + assert len(resolver.workspaces) == 1 + loaded_remote, loaded_local = loaded_with + assert loaded_remote.url == "http://127.0.0.1:9/original" + assert loaded_remote.headers[0].value == f"Bearer {SECRET}" + assert loaded_local.args == ["mcp"] + assert loaded_local.env[0].value == SECRET + assert agent.sessions[0]["mcp_servers"] == loaded_with + + @pytest.mark.asyncio + @pytest.mark.parametrize("wake", [False, True], ids=["recovery", "deliberate-wake"]) + async def test_a_reloaded_retained_session_resolves_once_on_its_new_connection( + self, wake: bool + ) -> None: + agent = FakeACPAgent(http=True, supports_session_load=True) + resolver = Resolver(http_server(url="http://127.0.0.1:9/grant-a")) + loaded_with: list[Any] = [] + load = agent.load_session + + async def load_session( + cwd: str, session_id: str, mcp_servers: Any = None, **kwargs: Any + ) -> LoadSessionResponse: + loaded_with.extend(mcp_servers or []) + return await load(cwd, session_id, mcp_servers, **kwargs) + + agent.load_session = load_session # type: ignore[method-assign] + async with acp_adapter(agent, resolve_mcp_servers=resolver) as acp: + await acp.send("one") + retained_id = acp.session_id(DEFAULT_ROOM) + agent.knows_session(retained_id) + room = await recover(acp) + resolver.servers = [http_server(url="http://127.0.0.1:9/grant-b")] + if wake: + async with woken(acp.adapter, room): + assert room.session is not None and room.session.live + assert len(resolver.workspaces) == 2 + else: + await acp.send("two") + await acp.send("three") + assert acp.session_id(DEFAULT_ROOM) == retained_id + + assert len(resolver.workspaces) == 2 + assert loaded_with[0].url == "http://127.0.0.1:9/grant-b" + assert agent.session_load_requests == [retained_id] + assert len(agent.sessions) == 1 + + @pytest.mark.asyncio + async def test_a_replaced_band_endpoint_resolves_a_new_session_snapshot( + self, + ) -> None: + agent = FakeACPAgent(http=True) + resolver = Resolver(http_server(url="http://127.0.0.1:9/grant-a")) + async with acp_adapter( + agent, + fake_agent_config(inject_band_tools=True), + resolve_mcp_servers=resolver, + ) as acp: + await acp.send("one") + resolver.servers = [http_server(url="http://127.0.0.1:9/grant-b")] + await crash_backend(acp.adapter._mcp) + await acp.send("two") + await acp.send("three") + + assert len(resolver.workspaces) == 2 + assert [s["mcp_servers"][0].url for s in agent.sessions] == [ + "http://127.0.0.1:9/grant-a", + "http://127.0.0.1:9/grant-b", + ] + assert ( + agent.sessions[0]["mcp_servers"][-1].url + != agent.sessions[1]["mcp_servers"][-1].url + ) diff --git a/tests/integrations/acp/clientadapter/test_process_reaping.py b/tests/integrations/acp/clientadapter/test_process_reaping.py new file mode 100644 index 000000000..54ee3b182 --- /dev/null +++ b/tests/integrations/acp/clientadapter/test_process_reaping.py @@ -0,0 +1,98 @@ +"""Successful spawn-context exit cannot abandon a live ACP child.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from tests.integrations.acp.acp_toolkit.processreaping import leaky_runtime + + +class TestProcessReaping: + """The runtime confirms exit itself when the spawn context leaves a child.""" + + @pytest.mark.asyncio + async def test_stop_reaps_a_process_its_context_left_running(self) -> None: + async with leaky_runtime() as (runtime, process): + await runtime.start() + await runtime.stop() + + assert process.returncode is not None + assert runtime.closed + + @pytest.mark.asyncio + @pytest.mark.parametrize("confirmation", ["kill", "wait"]) + async def test_failed_confirmation_keeps_the_child_and_refuses_its_connection( + self, confirmation: str + ) -> None: + async with leaky_runtime() as (runtime, process): + await runtime.start() + kill, wait = process.kill, process.wait + + def refused_kill() -> None: + raise PermissionError("kill refused") + + async def refused_wait() -> int: + raise OSError("wait failed") + + if confirmation == "kill": + process.kill = refused_kill # type: ignore[method-assign] + else: + process.wait = refused_wait # type: ignore[method-assign] + with pytest.raises(OSError): + await runtime.stop() + + assert runtime._transport is process + assert not runtime.closed + with pytest.raises(RuntimeError, match="connection is not open"): + await runtime.ensure_connection() + with pytest.raises(RuntimeError, match="not yet closed"): + await runtime.start() + + process.kill = kill # type: ignore[method-assign] + process.wait = wait # type: ignore[method-assign] + await runtime.stop() + + assert process.returncode is not None + assert runtime.closed + + @pytest.mark.asyncio + @pytest.mark.parametrize("cancelled", [False, True], ids=["failed", "cancelled"]) + async def test_failed_start_reaps_its_child_without_changing_the_start_outcome( + self, cancelled: bool + ) -> None: + error = asyncio.CancelledError() if cancelled else LookupError("init failed") + async with leaky_runtime(error) as (runtime, process): + with pytest.raises(type(error)): + await runtime.start() + + assert process.returncode is not None + assert runtime.closed + + @pytest.mark.asyncio + @pytest.mark.parametrize("cancelled", [False, True], ids=["failed", "cancelled"]) + async def test_failed_start_keeps_its_outcome_when_the_child_refuses_to_close( + self, cancelled: bool + ) -> None: + error = asyncio.CancelledError() if cancelled else LookupError("init failed") + async with leaky_runtime(error) as (runtime, process): + kill = process.kill + + def refused_kill() -> None: + raise PermissionError("kill refused") + + process.kill = refused_kill # type: ignore[method-assign] + with pytest.raises(type(error)): + await runtime.start() + + assert process.returncode is None + assert runtime._transport is process + assert not runtime.closed + with pytest.raises(RuntimeError, match="connection is not open"): + await runtime.ensure_connection() + process.kill = kill # type: ignore[method-assign] + await runtime.stop() + + assert process.returncode is not None + assert runtime.closed diff --git a/tests/integrations/acp/roomownership/test_mcp_setup.py b/tests/integrations/acp/roomownership/test_mcp_setup.py new file mode 100644 index 000000000..72453863e --- /dev/null +++ b/tests/integrations/acp/roomownership/test_mcp_setup.py @@ -0,0 +1,504 @@ +"""MCP initialization stays owned by the room through cancellation and teardown.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest +from acp import RequestError + +from band.core.exceptions import CleanupIncompleteError, RoomClosingError +from band.core.retention import Residency +from tests.adapters.waiting import spin, until +from tests.integrations.acp.acp_toolkit import ( + DEFAULT_ROOM, + FakeACPAgent, + acp_adapter, + fake_agent_config, +) +from tests.integrations.acp.acp_toolkit.agent import AGENT_EXIT_FAILURE +from tests.integrations.acp.acp_toolkit.mcpresolver import ( + HeldResolver, + Resolver, + http_server, + server_names, +) +from tests.mcpbackends import backends_created_by + + +class TestMcpSetup: + """One setup serves waiting turns while retained custody owns its actual end.""" + + @pytest.mark.asyncio + async def test_concurrent_turns_share_resolution_and_rooms_are_independent( + self, + ) -> None: + agent = FakeACPAgent(http=True) + entered, release = asyncio.Event(), asyncio.Event() + calls: list[str] = [] + + async def resolver(workspace: str) -> list[Any]: + calls.append(workspace) + if workspace.endswith(DEFAULT_ROOM): + entered.set() + await release.wait() + return [http_server()] + + async with acp_adapter(agent, resolve_mcp_servers=resolver) as acp: + first = asyncio.create_task(acp.send("a")) + second = asyncio.create_task(acp.send("b")) + try: + await entered.wait() + await acp.send("other room", room="room-2") + assert not first.done() and not second.done() + finally: + release.set() + await asyncio.gather(first, second) + + assert len(calls) == 2 + assert len(set(calls)) == 2 + assert len(agent.sessions) == 2 + + @pytest.mark.asyncio + async def test_one_waiter_leaving_does_not_cancel_the_shared_setup(self) -> None: + agent = FakeACPAgent(http=True) + entered, release = asyncio.Event(), asyncio.Event() + calls = 0 + + async def resolver(workspace: str) -> list[Any]: + nonlocal calls + calls += 1 + entered.set() + await release.wait() + return [http_server()] + + async with acp_adapter(agent, resolve_mcp_servers=resolver) as acp: + leaving = asyncio.create_task(acp.send("a")) + staying = asyncio.create_task(acp.send("b")) + try: + await entered.wait() + await spin() + leaving.cancel() + with pytest.raises(asyncio.CancelledError): + await leaving + finally: + release.set() + await staying + + assert calls == 1 + assert len(agent.sessions) == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize("swallow", [False, True], ids=["cancelled", "swallowed"]) + async def test_last_waiter_leaving_cancels_setup_before_any_session( + self, swallow: bool + ) -> None: + agent = FakeACPAgent(http=True) + callback = HeldResolver(swallow=swallow) + async with acp_adapter(agent, resolve_mcp_servers=callback) as acp: + turn = asyncio.create_task(acp.send("a")) + try: + await callback.entered.wait() + turn.cancel() + with pytest.raises(asyncio.CancelledError): + await turn + await callback.settling.wait() + assert agent.sessions == [] + assert acp.room().runtime is not None + finally: + callback.settle.set() + await spin() + await until( + lambda: acp.room().session_setup.done(), + failure="the abandoned setup never settled", + ) + assert acp.room().session is None + + assert callback.cancellations == 1 + assert agent.sessions == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize("swallow", [False, True], ids=["cancelled", "swallowed"]) + async def test_abandoned_setup_holds_the_room_until_it_settles( + self, swallow: bool + ) -> None: + agent = FakeACPAgent(http=True).will_say("ok") + callback = HeldResolver(swallow=swallow) + async with acp_adapter(agent, resolve_mcp_servers=callback) as acp: + abandoned = asyncio.create_task(acp.send("abandoned")) + try: + await callback.entered.wait() + abandoned.cancel() + with pytest.raises(asyncio.CancelledError): + await abandoned + await callback.settling.wait() + follow_up = asyncio.create_task(acp.send("next")) + await spin() + assert len(callback.workspaces) == 1 + assert not follow_up.done() + finally: + callback.settle.set() + reply = await follow_up + + assert reply.thoughts == ["ok"] + assert callback.log == ["settled", "second call"] + assert callback.cancellations == 1 + assert len(agent.sessions) == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize("self_cancel", [False, True]) + async def test_cleanup_joins_a_cancelling_setup_without_cancelling_again( + self, self_cancel: bool + ) -> None: + agent = FakeACPAgent(http=True) + callback = HeldResolver(self_cancel=self_cancel) + config = fake_agent_config(turn_settle_timeout_s=0) + async with acp_adapter(agent, config, resolve_mcp_servers=callback) as acp: + turn = asyncio.create_task(acp.send("a")) + try: + await callback.entered.wait() + if not self_cancel: + turn.cancel() + with pytest.raises(asyncio.CancelledError): + await turn + await callback.settling.wait() + exiting = agent.exits_slowly() + cleanups = [ + asyncio.create_task(acp.adapter.on_cleanup(DEFAULT_ROOM)) + for _ in range(2) + ] + await spin() + assert not exiting.received.is_set() + assert acp.room().runtime is not None + assert DEFAULT_ROOM in acp.adapter._workspaces.rooms + assert callback.cancellations == 1 + finally: + callback.settle.set() + await exiting.received.wait() + exiting.release.set() + await asyncio.gather(*cleanups) + with pytest.raises(asyncio.CancelledError): + await turn + assert DEFAULT_ROOM not in acp.adapter._rooms + assert DEFAULT_ROOM not in acp.adapter._workspaces.rooms + + assert callback.log == ["settled"] + assert callback.cancellations == 1 + assert agent.sessions == [] + + @pytest.mark.asyncio + async def test_cleanup_waits_for_resolver_settlement_and_refuses_its_late_result( + self, + ) -> None: + agent = FakeACPAgent(http=True) + callback = HeldResolver() + config = fake_agent_config(turn_settle_timeout_s=0) + async with acp_adapter(agent, config, resolve_mcp_servers=callback) as acp: + turn = asyncio.create_task(acp.send("a")) + try: + await callback.entered.wait() + cleanup = asyncio.create_task(acp.adapter.on_cleanup(DEFAULT_ROOM)) + await callback.settling.wait() + assert not cleanup.done() + assert acp.room().runtime is not None + assert DEFAULT_ROOM in acp.adapter._workspaces.rooms + finally: + callback.settle.set() + await cleanup + with pytest.raises(asyncio.CancelledError): + await turn + + assert callback.cancellations == 1 + assert agent.sessions == [] + + @pytest.mark.asyncio + async def test_a_setup_cancelled_by_itself_never_reaches_the_agent(self) -> None: + agent = FakeACPAgent(http=True) + + async def resolver(workspace: str) -> list[Any]: + task = asyncio.current_task() + assert task is not None + task.cancel() + try: + await asyncio.sleep(0) + except asyncio.CancelledError: + pass + return [http_server()] + + async with acp_adapter(agent, resolve_mcp_servers=resolver) as acp: + with pytest.raises(asyncio.CancelledError): + await acp.send("hello") + + assert agent.sessions == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("cleanup", "cancellation"), + [ + ("leave", "none"), + ("leave", "self"), + ("stop", "none"), + ("stop", "self"), + ("cleanup_all", "none"), + ("cleanup_all", "self"), + ("cleanup_all", "shutdown"), + ], + ) + async def test_setup_requesting_teardown_does_not_wait_on_itself( + self, cleanup: str, cancellation: str + ) -> None: + agent = FakeACPAgent(http=True).will_say("ok") + config = fake_agent_config(inject_band_tools=True, turn_settle_timeout_s=0) + callback = HeldResolver( + self_cancel=cancellation == "self", + wait_for_cancel=cancellation == "shutdown", + ) + with backends_created_by(): + async with acp_adapter(agent, config, resolve_mcp_servers=callback) as acp: + + async def request() -> None: + with pytest.raises(CleanupIncompleteError, match="calling work"): + match cleanup: + case "leave": + await acp.adapter.on_cleanup(DEFAULT_ROOM) + case "stop": + await acp.adapter.stop() + case "cleanup_all": + await acp.adapter.cleanup_all() + + callback.request = request + turn = asyncio.create_task(acp.send("a")) + outside = None + if cancellation == "shutdown": + await callback.entered.wait() + outside = asyncio.create_task(acp.adapter.cleanup_all()) + try: + async with asyncio.timeout(5): + await callback.settling.wait() + assert callback.log == ["request returned"] + assert acp.room().runtime is not None + assert DEFAULT_ROOM in acp.adapter._workspaces.rooms + assert acp.adapter._mcp.current is not None + finally: + callback.settle.set() + with pytest.raises(asyncio.CancelledError): + await turn + if outside is not None: + await outside + if cleanup == "leave": + await acp.adapter.on_cleanup(DEFAULT_ROOM) + else: + await acp.adapter.cleanup_all() + with pytest.raises(RoomClosingError): + await acp.send("refused", room="room-2") + await acp.adapter.on_started("Fake Agent", "restarted") + reply = await acp.send("after", room="room-2") + + assert reply.thoughts == ["ok"] + assert callback.log == ["request returned", "settled", "second call"] + assert callback.cancellations == int(cancellation != "none") + assert len(agent.sessions) == 1 + + @pytest.mark.asyncio + async def test_callback_failure_never_falls_back_and_next_turn_recovers( + self, + ) -> None: + agent = FakeACPAgent(http=True).will_say("ok") + config = fake_agent_config(mcp_servers=[http_server(name="static")]) + resolver = Resolver(http_server(name="fixed")) + fail = True + + async def flaky(workspace: str) -> list[Any]: + if fail: + raise LookupError("grant unavailable") + return list(await resolver(workspace)) + + async with acp_adapter(agent, config, resolve_mcp_servers=flaky) as acp: + with pytest.raises(LookupError, match="grant unavailable"): + await acp.send("first") + failed_runtime = acp.room().runtime + assert agent.sessions == [] + fail = False + await acp.send("second") + assert acp.runtime() is not failed_runtime + + assert [server_names(s["mcp_servers"]) for s in agent.sessions] == [["fixed"]] + + @pytest.mark.asyncio + @pytest.mark.parametrize("failure", ["resolver", "session"]) + async def test_setup_failure_survives_failed_close_and_next_turn_recovers( + self, failure: str + ) -> None: + agent = FakeACPAgent(http=True).will_say("recovered") + if failure == "session": + agent.fails_next_new_session() + resolver = Resolver(http_server()) + + async def callback(workspace: str) -> list[Any]: + result = list(await resolver(workspace)) + if len(resolver.workspaces) == 1: + agent.fails_next_exit() + if failure == "resolver": + raise LookupError("resolver failed") + return result + + async with acp_adapter(agent, resolve_mcp_servers=callback) as acp: + with pytest.raises(LookupError if failure == "resolver" else RequestError): + await acp.send("first") + room = acp.room() + await until( + lambda: ( + room.retained.snapshot.error is not None + and AGENT_EXIT_FAILURE in room.retained.snapshot.error + ), + failure="failed close was not retained", + ) + failed_runtime = room.runtime + assert failed_runtime is not None and not failed_runtime.closed + assert room.retained.snapshot.state == Residency.CLEANUP_PENDING + assert DEFAULT_ROOM in acp.adapter._workspaces.rooms + + reply = await acp.send("second") + assert reply.thoughts == ["recovered"] + assert acp.runtime() is not failed_runtime + assert failed_runtime.closed + + assert len(resolver.workspaces) == 2 + + @pytest.mark.asyncio + @pytest.mark.parametrize("failed_close", [False, True]) + async def test_a_captured_runtime_after_cleanup_starts_no_server_or_resolution( + self, failed_close: bool + ) -> None: + agent = FakeACPAgent(http=True).will_say("ok") + resolver = Resolver(http_server()) + config = fake_agent_config(inject_band_tools=True) + with backends_created_by(): + async with acp_adapter(agent, config, resolve_mcp_servers=resolver) as acp: + await acp.send("before") + room, captured = acp.room(), acp.runtime() + if failed_close: + agent.fails_next_exit() + with pytest.RaisesGroup( + pytest.RaisesExc(RuntimeError, match=AGENT_EXIT_FAILURE) + ): + await acp.adapter.cleanup_all() + assert room.runtime is captured + assert DEFAULT_ROOM in acp.adapter._workspaces.rooms + assert acp.adapter._mcp.current is not None + else: + await acp.adapter.cleanup_all() + assert acp.adapter._mcp.current is None + with pytest.raises(asyncio.CancelledError): + await acp.adapter._get_or_create_session(room, captured, None) + assert len(resolver.workspaces) == 1 + assert len(agent.sessions) == 1 + await acp.adapter.cleanup_all() + assert acp.adapter._mcp.current is None + await acp.adapter.on_started("Fake Agent", "restarted") + await acp.send("after") + + assert len(resolver.workspaces) == 2 + + @pytest.mark.asyncio + async def test_setup_abandoned_while_its_own_timeout_settles_is_not_reused( + self, + ) -> None: + """The last waiter can leave while the callback is already cancelling + itself; once the callback's timeout clears that request and it returns, + its result must still not reach the agent or serve the next turn.""" + settling, release = asyncio.Event(), asyncio.Event() + calls: list[str] = [] + + async def resolver(workspace: str) -> list[Any]: + calls.append(workspace) + if len(calls) == 1: + try: + async with asyncio.timeout(0): + try: + await asyncio.Event().wait() + finally: + settling.set() + await release.wait() + except TimeoutError: + pass + return [http_server(name=f"attempt-{len(calls)}")] + + agent = FakeACPAgent(http=True).will_say("ok") + async with acp_adapter(agent, resolve_mcp_servers=resolver) as acp: + abandoned = asyncio.create_task(acp.send("abandoned")) + try: + await settling.wait() + setup = acp.room().session_setup + assert setup is not None and setup.cancelling() + abandoned.cancel() + with pytest.raises(asyncio.CancelledError): + await abandoned + release.set() + await asyncio.wait({setup}) + reply = await acp.send("next") + finally: + release.set() + await asyncio.gather(abandoned, return_exceptions=True) + + assert reply.thoughts == ["ok"] + assert len(calls) == 2 + assert [server_names(s["mcp_servers"])[0] for s in agent.sessions] == [ + "attempt-2" + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("settling", [False, True], ids=["held", "own-timeout"]) + async def test_a_waiter_that_stays_reuses_the_shared_setup_result( + self, settling: bool + ) -> None: + """When one of two waiters leaves, the shared setup is still the other's: + even if the callback was cancelling its own timeout meanwhile, its result + serves the remaining waiter instead of a second resolution and session.""" + entered, release = asyncio.Event(), asyncio.Event() + calls: list[str] = [] + + async def resolver(workspace: str) -> list[Any]: + calls.append(workspace) + if len(calls) == 1: + if settling: + try: + async with asyncio.timeout(0): + try: + await asyncio.Event().wait() + finally: + entered.set() + await release.wait() + except TimeoutError: + pass + else: + entered.set() + await release.wait() + return [http_server(name=f"attempt-{len(calls)}")] + + agent = FakeACPAgent(http=True).will_say("ok") + async with acp_adapter(agent, resolve_mcp_servers=resolver) as acp: + first = asyncio.create_task(acp.send("first")) + successor: asyncio.Task[Any] | None = None + try: + await entered.wait() + room = acp.room() + successor = asyncio.create_task(acp.send("successor")) + await until( + lambda: room.session_waiters == 2, failure="successor never queued" + ) + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + release.set() + reply = await successor + finally: + release.set() + if successor is not None: + await asyncio.gather(successor, return_exceptions=True) + await asyncio.gather(first, return_exceptions=True) + + assert reply.thoughts == ["ok"] + assert len(calls) == 1 + assert len(agent.sessions) == 1