From 3301dbb30db0538fb6127c6b0d7df94e8ae88650 Mon Sep 17 00:00:00 2001 From: Natan Zalkin Date: Thu, 15 Oct 2026 05:34:49 +0300 Subject: [PATCH 01/13] feat(codex): resolve ambient servers inside owned room startup --- docs/adapters/codex.md | 65 ++ src/band/adapters/codex.py | 362 ++++++++-- src/band/core/retention/room.py | 19 +- src/band/core/turn_lifecycle.py | 55 +- src/band/integrations/codex/mcp.py | 512 ++++++++++++++ src/band/integrations/codex/rpc_base.py | 34 +- src/band/integrations/codex/stdio_client.py | 84 ++- tests/adapters/codex/kit.py | 2 +- tests/adapters/codex/lifecycleprofile.py | 4 +- tests/adapters/codex/mcpkit.py | 130 ++++ tests/adapters/codex/test_mcp_acquisition.py | 38 + tests/adapters/codex/test_mcp_cancellation.py | 204 ++++++ tests/adapters/codex/test_mcp_launch.py | 659 ++++++++++++++++++ .../codex/test_mcp_owner_lifecycle.py | 150 ++++ tests/adapters/codex/test_mcp_recovery.py | 367 ++++++++++ tests/adapters/codex/test_mcp_startup.py | 649 +++++++++++++++++ tests/adapters/codexturns.py | 2 +- tests/adapters/opencode/helpers.py | 8 +- tests/adapters/opencode/test_waiting.py | 44 ++ tests/core/test_turn_lifecycle.py | 96 +++ tests/integrations/codex/test_cleanup.py | 61 +- .../integrations/codex/test_mcp_redaction.py | 252 +++++++ .../codex/test_process_group_cleanup.py | 510 ++++++++++++++ 23 files changed, 4195 insertions(+), 112 deletions(-) create mode 100644 src/band/integrations/codex/mcp.py create mode 100644 tests/adapters/codex/mcpkit.py create mode 100644 tests/adapters/codex/test_mcp_acquisition.py create mode 100644 tests/adapters/codex/test_mcp_cancellation.py create mode 100644 tests/adapters/codex/test_mcp_launch.py create mode 100644 tests/adapters/codex/test_mcp_owner_lifecycle.py create mode 100644 tests/adapters/codex/test_mcp_recovery.py create mode 100644 tests/adapters/codex/test_mcp_startup.py create mode 100644 tests/adapters/opencode/test_waiting.py create mode 100644 tests/integrations/codex/test_mcp_redaction.py create mode 100644 tests/integrations/codex/test_process_group_cleanup.py diff --git a/docs/adapters/codex.md b/docs/adapters/codex.md index 41d6e5ce8..2b0c1f0cc 100644 --- a/docs/adapters/codex.md +++ b/docs/adapters/codex.md @@ -41,6 +41,13 @@ For existing integrations, follow the [SDK 4.x to 5.x migration](../acp.md#upgra exited before terminating and draining the job; an empty job alone does not establish helper exit. Failed cleanup retains that same owner and job for retry. +- **POSIX cleanup settles the original owned process group.** Close allows the + leader one second to exit on stdin EOF, then sends TERM to its original group, + waits up to one second, and escalates to KILL with one more second to confirm + `ESRCH`. Unconfirmed cleanup retains the client and workspace for retry. + A reused leader PID is never signaled. POSIX groups lack generation-bound + handles: reuse between a probe and signal remains a residual race. + Descendants that leave the group are outside this guarantee. - **Leaving is final.** Leaving a room stops it admitting work at once. A leave that fails or runs over keeps the process, thread and workspace claim, and the room refuses work until a retried leave succeeds; a later rejoin @@ -111,3 +118,61 @@ Approval commands use the token displayed in the room prompt. Native Codex request and approval IDs stay internal; a fresh request gets a new chat token after completion or room cleanup, even when Codex reuses its native ID. A redelivery of a still-pending request keeps its displayed token. + +## Runtime MCP servers + +Pass the keyword-only `resolve_mcp_servers` callback to `CodexAdapter`, not +`CodexAdapterConfig`. It receives the room's absolute workspace and returns a +`Mapping[str, CodexMcpServer]`. Import the descriptors and startup error from +`band.integrations.codex.mcp`. + +- `CodexMcpStdioServer(command, args=(), env={}, cwd=None)` freezes arguments + and child-only environment values. Relative commands and arguments use the + server's `cwd`. On POSIX, nonempty `env` uses the SDK's interpreter to restore + generated aliases before replacing itself with the server process; native + protected environment names are refused. On Windows, nonempty stdio `env` + fails before client construction with stage `unsupported-platform`: + `secret-bearing stdio MCP servers require POSIX exec`. Env-free Windows stdio + and all HTTP servers remain supported; no Windows supervisor is included. +- `CodexMcpHttpServer(url, bearer_token=None, headers={})` describes streamable + HTTP. Other descriptor types, including SSE, are refused. Names must match + `[A-Za-z0-9_-]+`, be unique, and not be the reserved native name `codex_apps`. +- Resolution runs inside the retained room's acquisition, once per new process + after any predecessor is confirmed closed. Reusing a process does not resolve + again. To rotate a running grant, clean up the room first. An empty result + adds no servers or readiness requests; callback failures never fall back. +- The snapshot is frozen before launch without changing configuration or + `codex_env`. Validated TOML overrides are appended after existing arguments, + including a relay's `--` boundary. Supplied servers are enabled, required and + use connection readiness; other native servers are unchanged. +- Bearer tokens, headers and stdio environment values use generated child-env + names, never secret values in generated TOML or argv. Names reserve inherited + env names too. Secret fields are excluded from repr, equality and supported + Pydantic serialization; these objects are runtime-only, not persistence data. +- Before the process becomes resident or restores a thread, readiness pages + `mcpServerStatus/list` under one 30-second deadline. Every supplied server must + report information and capabilities, no discovery error, and null pre-thread + or `connected` runtime status. Tool-less servers are supported. Auth settings + or an empty catalog alone are not readiness proof. Malformed pages, duplicate + names, repeated cursors, negative statuses and consumed timeout cancellation + fail startup. +- Cancellation and retirement fence each stage and every discovery request. + A resolver may request its own room cleanup without self-waiting. Cleanup + disarms readiness deadlines and joins an already-cancelling acquisition + without a second cancellation. Settlement beyond the reporting bound retains + acquisition, child, workspace and diagnostic scope until cleanup succeeds. + Failed registration unwinds the gated worker before resolution can start. + Invocation-only cancellation timers end with their owner, never later work. +- Thread start/resume errors with adopted servers trigger bounded readiness + rather than error-text classification. Failed readiness refuses history + fallback and retires the process; the next acquisition resolves a fresh grant. + Positive readiness preserves genuine missing-history fallback and unrelated + errors. Recovery still strictly restores retained threads as described above. + +`CodexMcpStartupError.stage` is a secret-free startup or validation category. +Transport diagnostics omit RPC methods, wire contents and process-stderr text. +Unknown frames report transport, type and character length; stderr reports byte +length; full queues report only the dropped event kind; cleanup reports only +exception type. Forwarded RPC errors remain scoped-redacted and are frozen +before teardown clears the secret scope. + diff --git a/src/band/adapters/codex.py b/src/band/adapters/codex.py index 3a41024fd..5643b5982 100644 --- a/src/band/adapters/codex.py +++ b/src/band/adapters/codex.py @@ -21,6 +21,7 @@ from pydantic import ( AliasChoices, BaseModel, + ConfigDict, Field, ValidationError, field_validator, @@ -43,7 +44,7 @@ deliver_notice, reraise_delivery_cause, ) -from band.core.exceptions import InputNotAccepted +from band.core.exceptions import CleanupIncompleteError, InputNotAccepted from band.core.protocols import ( FAILURE_CODE_TIMEOUT, GENERIC_PROVIDER_FAILURE_MESSAGE, @@ -69,7 +70,6 @@ BackendTurns, TurnSettleTimeout, TurnTask, - bind_turn_task, start_releasable_turn, ) from band.core.types import ( @@ -88,6 +88,13 @@ CodexStdioClient, RpcEvent, ) +from band.integrations.codex.mcp import ( + CodexMcpLaunch, + CodexMcpResolver, + CodexMcpSnapshot, + CodexMcpStartupError, + freeze_mcp_servers, +) from band.integrations.codex.types import ( CODEX_APPROVAL_METHODS, CODEX_PROVIDER, @@ -627,6 +634,35 @@ def changes_in(self, restored: ThreadSettings) -> list[str]: ] +class McpServerInfo(BaseModel): + model_config = ConfigDict(strict=True) + + name: str + version: str + + +class McpServerStatus(BaseModel): + model_config = ConfigDict(strict=True) + + name: str = Field(min_length=1) + server_info: McpServerInfo | None = Field(default=None, alias="serverInfo") + server_capabilities: dict[str, Any] | None = Field( + default=None, alias="serverCapabilities" + ) + tools_error: str | None = Field(default=None, alias="toolsError") + runtime_status: str | None = Field(default=None, alias="runtimeStatus") + + +class McpStatusPage(BaseModel): + model_config = ConfigDict(strict=True) + + data: list[McpServerStatus] + next_cursor: str | None = Field(default=None, alias="nextCursor") + + +_MCP_READINESS_TIMEOUT_S = 30.0 + + @dataclass class RoomCodexClient: """One Band room's Codex process and what its thread was started with. @@ -643,6 +679,8 @@ class RoomCodexClient: client: CodexClientProtocol | None = None reader: asyncio.Task[None] | None = None turn: NativeTurn | None = None + mcp_snapshot: CodexMcpSnapshot | None = dataclass_field(default=None, repr=False) + startup_owner: TurnTask | None = dataclass_field(default=None, repr=False) selected_model: str | None = None reasoning_effort: str | None = None reasoning_summary: str | None = None @@ -919,6 +957,7 @@ def __init__( *, additional_tools: list[CustomToolDef] | None = None, history_converter: CodexHistoryConverter | None = None, + resolve_mcp_servers: CodexMcpResolver | None = None, client_factory: Callable[[CodexAdapterConfig], CodexClientProtocol] | None = None, **features: Unpack[FeatureKwargs], @@ -945,6 +984,7 @@ def __init__( raise ValueError( "custom Codex clients cannot guarantee room process isolation" ) + self._resolve_mcp_servers = resolve_mcp_servers self._room_clients: dict[str, RoomCodexClient] = {} # Final disposal gives a room's turn ``turn_settle_timeout_s`` to wind # down before cancelling it. @@ -1197,16 +1237,23 @@ async def on_message( is_session_bootstrap: bool, room_id: str, ) -> None: + task = asyncio.current_task() + if task is not None and task.cancelling(): + raise asyncio.CancelledError await self._admit_turn(room_id, "Codex") room = self._room_client(room_id) self._active_room.set(room_id) command = self._extract_local_command(msg.content) release: asyncio.Future[None] = asyncio.get_running_loop().create_future() + owner_ready: asyncio.Future[TurnTask] = ( + asyncio.get_running_loop().create_future() + ) turn = start_releasable_turn( room.retained, partial( self._run_turn, + owner_ready=owner_ready, room=room, msg=msg, tools=tools, @@ -1220,15 +1267,35 @@ async def on_message( room_id=room_id, ) self._turn_release[room_id] = release - self._turn_tasks[room_id] = bind_turn_task(tools, turn) + owner = self._turn_tasks[room_id] = TurnTask(turn) turn.add_done_callback(lambda task: self._forget_turn(room_id, task)) + try: + register = getattr(tools, "register_turn_task", None) + if callable(register): + register(owner) + if owner.task.cancelling(): + raise asyncio.CancelledError + owner_ready.set_result(owner) + except BaseException: + owner_ready.cancel() + owner.terminate() + try: + await owner.drain() + finally: + owner.close() + self._forget_turn(room_id, turn) + raise try: await asyncio.wait({release, turn}, return_when=asyncio.FIRST_COMPLETED) except asyncio.CancelledError: await self._cancel_turn(room_id) raise finally: - if self._turn_release.get(room_id) is release: + if ( + turn.done() + and room.startup_owner is None + and self._turn_release.get(room_id) is release + ): del self._turn_release[room_id] if turn.done(): # No human was asked: the turn's outcome is on_message's own. @@ -1237,6 +1304,7 @@ async def on_message( async def _run_turn( self, *, + owner_ready: asyncio.Future[TurnTask], room: RoomCodexClient, msg: PlatformMessage, tools: AgentToolsProtocol, @@ -1247,6 +1315,7 @@ async def _run_turn( command: tuple[CodexCommand, str] | None, ) -> None: """Run one turn to completion, leasing its room's process throughout.""" + await owner_ready async with room.rpc_lock: try: if command is not None and await self._run_local_command( @@ -1269,12 +1338,13 @@ async def _run_turn( # A lease refused by a closing room is no failed turn: the # host keeps the message. raise - except (CodexJsonRpcError, CodexContinuityError) as e: - # A structured RPC error from the app-server (e.g. "model not - # available") or a refused restore is safe, curated text -- - # unlike an arbitrary caught exception, worth showing verbatim. - await tools.send_failure(AgentFailure(CODEX_PROVIDER, str(e))) + except (CodexMcpStartupError, CodexContinuityError) as error: + await tools.send_failure(AgentFailure(CODEX_PROVIDER, str(error))) raise + except CodexJsonRpcError as error: + safe = self._safe_rpc_error(room, error) + await tools.send_failure(AgentFailure(CODEX_PROVIDER, str(safe))) + raise safe from None except TimeoutError as timeout: # Only waiting on the room's process start or stop raises # this; an abandoned start stays owned until recovery closes it. @@ -1501,6 +1571,9 @@ def _forget_turn(self, room_id: str, task: asyncio.Task[None]) -> None: owner = self._turn_tasks.get(room_id) if owner is not None and owner.task is task: del self._turn_tasks[room_id] + room = self._room_clients.get(room_id) + if room is None or room.startup_owner is None: + self._turn_release.pop(room_id, None) async def _emit_failed_turn_outcome( self, @@ -1888,7 +1961,20 @@ async def on_cleanup(self, room_id: str) -> None: if room is None: return self._clear_pending_approvals_for_room(room_id) - await room.retained.dispose() + if room.startup_owner is not None: + room.startup_owner.terminate( + deadline=asyncio.get_running_loop().time() + + self.config.turn_settle_timeout_s, + cooperative=True, + ) + try: + await room.retained.dispose() + except CleanupIncompleteError: + if ( + room.startup_owner is None + or room.startup_owner.task is not asyncio.current_task() + ): + raise async def cleanup_all(self) -> None: """Stop the residency owner: admission closes at once, every room's @@ -1900,39 +1986,168 @@ async def cleanup_all(self) -> None: workspace. Success stays sealed until ``on_started``. """ for room_id in list(self._room_clients): + room = self._room_clients[room_id] + if room.startup_owner is not None: + room.startup_owner.terminate( + deadline=asyncio.get_running_loop().time() + + self.config.turn_settle_timeout_s, + cooperative=True, + ) self._clear_pending_approvals_for_room(room_id) await self._retention.stop() await self.backend_turns.drain() + @staticmethod + def _check_startup_current(room: RoomCodexClient, generation: int) -> None: + task = asyncio.current_task() + if task is not None and task.cancelling(): + raise asyncio.CancelledError + room.retained.require_admission() + if ( + not room.retained.has_authority(generation) + or (room.startup_owner is not None and room.startup_owner.task.cancelling()) + or ( + room.startup_owner is not None + and room.startup_owner.deadline is not None + ) + ): + raise asyncio.CancelledError + async def _acquire( self, room: RoomCodexClient, generation: int, resume: bool ) -> None: - """Start the room's process and strictly restore the thread it retains. - - A crash keeps the thread id and the room's model, so ``resume`` - changes nothing: a retained thread is always restored on its model. - The reader runs before ``thread/resume``, which can start native work. - """ - room.client = client = self._build_client(room.workspace) - await client.connect() - room.reader = spawn_owned( - self._read_native_events(room, client, generation), - label=f"Room {room.room_id} Codex event reader", - ) - await client.initialize( - client_name=self.config.client_name, - client_title=self.config.client_title, - client_version=self.config.client_version, - experimental_api=self.config.experimental_api, - ) - await self._register_skill_roots(client) - if room.selected_model is None: - room.selected_model = await self._select_model(client) - logger.info( - "Room %s: Codex model selected: %s", room.room_id, room.selected_model + """Resolve and ready one process under the retained acquisition owner.""" + task = asyncio.current_task() + assert task is not None + owner = room.startup_owner = TurnTask(task) + stage = "resolution" + try: + self._check_startup_current(room, generation) + launch = None + if self._resolve_mcp_servers is not None: + resolved = await self._resolve_mcp_servers(room.workspace) + self._check_startup_current(room, generation) + stage = "snapshot" + room.mcp_snapshot = freeze_mcp_servers(resolved) + self._check_startup_current(room, generation) + stage = "launch" + launch = room.mcp_snapshot.launch( + self.config.codex_command + or CodexStdioClient._default_codex_command(), + self.config.codex_env, + ) + self._check_startup_current(room, generation) + room.client = client = self._build_client(room, launch) + stage = "connect" + await client.connect() + self._check_startup_current(room, generation) + room.reader = spawn_owned( + self._read_native_events(room, client, generation), + label=f"Room {room.room_id} Codex event reader", + ) + stage = "initialize" + await client.initialize( + client_name=self.config.client_name, + client_title=self.config.client_title, + client_version=self.config.client_version, + experimental_api=self.config.experimental_api, ) - if (thread_id := self._room_threads.get(room.room_id)) is not None: - await self._restore_thread(room, client, thread_id) + self._check_startup_current(room, generation) + stage = "skills" + await self._register_skill_roots(client) + self._check_startup_current(room, generation) + stage = "model" + selected_model = room.selected_model or await self._select_model(client) + self._check_startup_current(room, generation) + stage = "readiness" + await self._check_mcp_readiness(client, room, generation, owner) + self._check_startup_current(room, generation) + room.selected_model = selected_model + if (thread_id := self._room_threads.get(room.room_id)) is not None: + stage = "thread/resume" + await self._restore_thread(room, client, thread_id) + self._check_startup_current(room, generation) + except asyncio.CancelledError: + raise asyncio.CancelledError from None + except InputNotAccepted: + raise + except Exception as error: + if self._resolve_mcp_servers is None: + raise + if ( + stage == "snapshot" + and type(error) is CodexMcpStartupError + and error.stage + in { + "unsupported-transport", + "unsupported-platform", + "invalid-name", + "duplicate-name", + "reserved-name", + "invalid-descriptor", + } + ): + stage = error.stage + raise CodexMcpStartupError(stage) from None + finally: + owner.close() + room.startup_owner = None + if not self._turn_running(room.room_id): + self._turn_release.pop(room.room_id, None) + + @staticmethod + def _check_mcp_deadline(readiness: asyncio.Timeout) -> None: + deadline = readiness.when() + if readiness.expired() or ( + deadline is not None and asyncio.get_running_loop().time() >= deadline + ): + raise CodexMcpStartupError("readiness") + + async def _check_mcp_readiness( + self, + client: CodexClientProtocol, + room: RoomCodexClient, + generation: int, + owner: TurnTask, + ) -> None: + snapshot = room.mcp_snapshot + if snapshot is None or not snapshot.expected_names: + return + seen_names: set[str] = set() + cursors: set[str] = set() + cursor = None + async with owner.timed(_MCP_READINESS_TIMEOUT_S) as readiness: + while True: + self._check_startup_current(room, generation) + raw_page = await client.request( + "mcpServerStatus/list", + {"cursor": cursor, "limit": 100, "detail": "full"}, + ) + self._check_mcp_deadline(readiness) + self._check_startup_current(room, generation) + page = McpStatusPage.model_validate(raw_page) + for row in page.data: + if row.name in seen_names: + raise CodexMcpStartupError("readiness") + seen_names.add(row.name) + if row.name in snapshot.expected_names and ( + row.server_info is None + or row.server_capabilities is None + or "tools_error" not in row.model_fields_set + or row.tools_error is not None + or "runtime_status" not in row.model_fields_set + or row.runtime_status not in (None, "connected") + ): + raise CodexMcpStartupError("readiness") + cursor = page.next_cursor + if cursor is None: + break + if not cursor or cursor in cursors: + raise CodexMcpStartupError("readiness") + cursors.add(cursor) + if not snapshot.expected_names <= seen_names: + raise CodexMcpStartupError("readiness") + self._check_mcp_deadline(readiness) def _resume_params(self, room: RoomCodexClient, thread_id: str) -> dict[str, Any]: """``thread/resume`` overrides pinning the thread to the room's workspace. @@ -2154,7 +2369,14 @@ async def _release( f"Codex process for room {room.room_id} did not close " f"within {timeout}s" ) from error + except Exception: + if room.mcp_snapshot is not None: + raise CodexMcpStartupError("retirement") from None + raise room.client = None + if room.mcp_snapshot is not None: + room.mcp_snapshot.redactor.clear() + room.mcp_snapshot = None room.restored = False return ReleaseOutcome.RELEASED @@ -2192,11 +2414,16 @@ async def _register_skill_roots(self, client: CodexClientProtocol) -> None: except CodexJsonRpcError as exc: raise RuntimeError(f"Codex rejected skill_roots {roots}: {exc}") from exc - def _build_client(self, workspace: str) -> CodexClientProtocol: + def _build_client( + self, room: RoomCodexClient, launch: CodexMcpLaunch | None + ) -> CodexClientProtocol: return CodexStdioClient( - command=self.config.codex_command, - cwd=workspace, - env=self.config.codex_env, + command=launch.command if launch is not None else self.config.codex_command, + cwd=room.workspace, + env=launch.env if launch is not None else self.config.codex_env, + redactor=room.mcp_snapshot.redactor + if room.mcp_snapshot is not None + else None, ) async def _select_model(self, client: CodexClientProtocol) -> str: @@ -2218,6 +2445,52 @@ async def _select_model(self, client: CodexClientProtocol) -> str: logger.warning("model/list has no visible models; using default Codex model") return OPENAI_MODEL + @staticmethod + def _safe_rpc_error( + room: RoomCodexClient, error: CodexJsonRpcError + ) -> CodexJsonRpcError: + if room.mcp_snapshot is None: + return error + redactor = room.mcp_snapshot.redactor + return CodexJsonRpcError( + code=error.code, + message=redactor.text(error.message), + data=redactor.value(error.data), + ) + + async def _verify_mcp_after_thread_error( + self, room: RoomCodexClient, stage: str + ) -> None: + if room.mcp_snapshot is None or not room.mcp_snapshot.expected_names: + return + task = asyncio.current_task() + assert task is not None + owner = self._turn_tasks.get(room.room_id) + temporary = owner is None + if owner is None: + owner = TurnTask(task) + try: + await self._check_mcp_readiness( + room.require_client(), room, room.generation, owner + ) + except asyncio.CancelledError: + room.retained.recover( + generation=room.generation, reason="MCP readiness cancelled" + ) + raise asyncio.CancelledError from None + except Exception as error: # noqa: BLE001 -- discovery errors may carry credentials + logger.debug( + "Codex MCP readiness failed after a thread error (type=%s)", + type(error).__name__, + ) + room.retained.recover( + generation=room.generation, reason="MCP readiness failed" + ) + raise CodexMcpStartupError(stage) from None + finally: + if temporary: + owner.close() + async def _ensure_thread( self, *, @@ -2280,11 +2553,12 @@ async def _resume_from_history( self._resume_params(room, named), ) except CodexJsonRpcError as exc: + await self._verify_mcp_after_thread_error(room, "thread/resume") logger.warning( "thread/resume failed for room %s thread %s: %s", room_id, named, - exc, + self._safe_rpc_error(room, exc), ) if self.config.inject_history_on_resume_failure: self._needs_history_injection.add(room_id) @@ -2332,7 +2606,13 @@ async def _start_thread( } self._apply_thread_sandbox(start_params, room_id=room_id) - started = await client.request(CodexRequestMethod.THREAD_START, start_params) + try: + started = await client.request( + CodexRequestMethod.THREAD_START, start_params + ) + except CodexJsonRpcError as error: + await self._verify_mcp_after_thread_error(room, "thread/start") + raise self._safe_rpc_error(room, error) from None thread_id = _result_id(started, "thread") if not thread_id: raise RuntimeError("Codex thread/start returned no thread id") diff --git a/src/band/core/retention/room.py b/src/band/core/retention/room.py index 749572ea8..202bd221c 100644 --- a/src/band/core/retention/room.py +++ b/src/band/core/retention/room.py @@ -299,6 +299,8 @@ async def use( """ self.require_admission() holder = asyncio.current_task() + if holder is not None and holder.cancelling(): + raise asyncio.CancelledError self._holders[holder] += 1 self._schedule() try: @@ -383,6 +385,8 @@ async def _join( f"Room {self.room_id} transition still running after {timeout}s" ) if transition.cancelled(): + if initiator: + raise asyncio.CancelledError from None # Closed admission means disposal cancelled it; otherwise recovery # did, and the caller re-evaluates the room. self.require_admission() @@ -408,7 +412,8 @@ def _abandon_acquisition(self, transition: asyncio.Task[None]) -> None: self.room_id, self._generation, ) - transition.cancel(f"Room {self.room_id} acquisition was abandoned") + if not transition.cancelling(): + transition.cancel(f"Room {self.room_id} acquisition was abandoned") def _request(self, cause: ReleaseCause) -> None: self._cause = stronger_cause(self._cause, cause) @@ -585,7 +590,11 @@ def _start_idle_release(self) -> None: def enrolls(self, task: asyncio.Task[Any] | None) -> bool: """Whether ``task`` is accepted work this room's disposal would await.""" - return task in self._owned or task in self._holders + return ( + task in self._owned + or task in self._holders + or (task is self._transition and self._state == Residency.ACQUIRING) + ) def _spare_requester(self) -> bool: """Exempt enrolled work requesting cleanup from cancellation; it is @@ -656,7 +665,11 @@ async def _settle_accepted(self) -> None: self._work_cancelled = True for task in self._uncancelled_work(): task.cancel() - if self._state == Residency.ACQUIRING and self._transition is not None: + if ( + self._state == Residency.ACQUIRING + and self._transition is not None + and not self._transition.cancelling() + ): self._transition.cancel() await self._drained.wait() diff --git a/src/band/core/turn_lifecycle.py b/src/band/core/turn_lifecycle.py index 5c877bc5a..bcb66a52b 100644 --- a/src/band/core/turn_lifecycle.py +++ b/src/band/core/turn_lifecycle.py @@ -65,31 +65,58 @@ class TurnTask: def __init__(self, task: asyncio.Task[Any]) -> None: self.task = task - if task.done(): - self._retrieve_result(task) - else: - task.add_done_callback(self._retrieve_result) + self._cancel_baseline = task.cancelling() self.deadline: float | None = None self._cancel_delivered = False self._grace_expired = False self._timer: asyncio.Timeout | None = None self._cleanup_timer = False + self._expiry_timer: asyncio.TimerHandle | None = None + self._closed = False + if task.done(): + self._finish(task) + else: + task.add_done_callback(self._finish) @staticmethod def _retrieve_result(task: asyncio.Task[Any]) -> None: if not task.cancelled(): task.exception() + def _finish(self, task: asyncio.Task[None]) -> None: + self.close() + self._retrieve_result(task) + + def close(self) -> None: + """End this owner's lifetime without cancelling its task.""" + if self._closed: + return + self._closed = True + if self._expiry_timer is not None: + self._expiry_timer.cancel() + self._expiry_timer = None + if self._timer is not None and not self._timer.expired(): + self._timer.reschedule(None) + self.task.remove_done_callback(self._finish) + + def _schedule_expiry(self, deadline: float) -> None: + if self._expiry_timer is not None: + self._expiry_timer.cancel() + self._expiry_timer = asyncio.get_running_loop().call_at(deadline, self._expire) + def terminate( self, *, deadline: float | None = None, cooperative: bool = False ) -> None: """Start the cleanup grace once; later requests never renew it.""" + if self._closed: + return + self._adopt_cancellation() if self.deadline is not None: return bound = deadline if deadline is not None else terminal_deadline() self.deadline = self._release_timer(bound) # The owner enforces its own deadline: a detached worker may have no drainer. - asyncio.get_running_loop().call_at(self.deadline, self._expire) + self._schedule_expiry(self.deadline) if not cooperative: self._cancel() @@ -99,19 +126,24 @@ def retry(self, *, deadline: float) -> None: A task that swallowed its cancel would otherwise stay unfinished on every cleanup retry; within an unexpired grace this is ``terminate``. """ + if self._closed: + return if not self._grace_expired: self.terminate(deadline=deadline) return self.deadline = deadline self._grace_expired = False self._cancel_delivered = False - asyncio.get_running_loop().call_at(deadline, self._expire) + self._schedule_expiry(deadline) self._cancel() def _expire(self) -> None: + if self._closed: + return # The loop's own timer marks expiry: it can fire up to one clock # resolution before loop.time() reaches the deadline. self._grace_expired = True + self._adopt_cancellation() self._cancel() def _release_timer(self, bound: float) -> float: @@ -127,8 +159,13 @@ def _release_timer(self, bound: float) -> float: timer.reschedule(None) return bound + def _adopt_cancellation(self) -> None: + # Counts predating this owner belong to a previous cancellation lifetime. + if self.task.cancelling() > self._cancel_baseline: + self._cancel_delivered = True + def _cancel(self) -> None: - if not self._cancel_delivered and not self.task.done(): + if not self._closed and not self._cancel_delivered and not self.task.done(): self._cancel_delivered = True self.task.cancel() @@ -136,6 +173,8 @@ async def drain(self) -> bool: """Wait within the original terminal grace without cancelling a waiter.""" if self.task.done(): return True + if self._closed: + return False if self.deadline is None: raise RuntimeError("Turn draining requires termination") remaining = max(0.0, self.deadline - asyncio.get_running_loop().time()) @@ -147,7 +186,7 @@ async def timed( self, timeout: float | None, *, cleanup: bool = False ) -> AsyncIterator[asyncio.Timeout]: """Coordinate an existing provider timer with terminal cancellation.""" - if self.deadline is not None: + if self._closed or self.deadline is not None: raise asyncio.CancelledError if self._timer is not None or asyncio.current_task() is not self.task: raise RuntimeError("A turn timer must run in its owner, without nesting") diff --git a/src/band/integrations/codex/mcp.py b/src/band/integrations/codex/mcp.py new file mode 100644 index 000000000..3d539c605 --- /dev/null +++ b/src/band/integrations/codex/mcp.py @@ -0,0 +1,512 @@ +"""Immutable room-scoped MCP descriptions and secret-safe native Codex launches. + +These are live runtime objects, not adapter configuration. Pydantic serialization +excludes secret fields; stdlib dataclass serialization refuses their frozen maps. +""" + +from __future__ import annotations + +import json +import logging +import os +import re +import sys +import tomllib +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from heapq import heapify, heappop, heappush +from types import MappingProxyType +from typing import Annotated, Any, NoReturn, SupportsIndex, TypeAlias +from uuid import uuid4 + +from pydantic import Field + +logger = logging.getLogger(__name__) + +__all__ = [ + "CodexMcpHttpServer", + "CodexMcpLaunch", + "CodexMcpRedactor", + "CodexMcpResolver", + "CodexMcpServer", + "CodexMcpSnapshot", + "CodexMcpStartupError", + "CodexMcpStdioServer", + "freeze_mcp_servers", +] + +_SERVER_NAME = re.compile(r"[A-Za-z0-9_-]+") +# Codex manages this server separately from caller-configured transports. +# Native source: codex-rs/codex-mcp/src/mcp/mod.rs, CODEX_APPS_MCP_SERVER_NAME. +_RESERVED_SERVER_NAMES = frozenset({"codex_apps"}) +# Preserve native stdio filtering even after alias remapping. This is the union +# in protocol/src/shell_environment.rs for Codex 0.160.0 and 0.162.1. +_NON_INHERITABLE_STDIO_ENV = frozenset( + { + "CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN", + "NODE_REPL_AUTH_TOKEN", + "CODEX_GUARDIAN_DECISIONS_API_KEY", + "OPENAI_FEDERATION_RULE_ID", + "OPENAI_IDENTITY_TOKEN_FILE", + "OPENAI_WORKLOAD_IDENTITY_CONTEXT", + } +) +_REDACTED = "[REDACTED]" +_STARTUP_STAGES = frozenset( + { + "startup", + "retirement", + "resolution", + "snapshot", + "launch", + "connect", + "initialize", + "skills", + "model", + "readiness", + "thread/start", + "thread/resume", + } +) +_VALIDATION_MESSAGES = { + "unsupported-transport": "unsupported transport (only stdio and streamable HTTP are supported)", + "unsupported-platform": "secret-bearing stdio MCP servers require POSIX exec", + "invalid-name": "invalid server name", + "duplicate-name": "duplicate server name", + "reserved-name": "reserved server name", + "invalid-descriptor": "invalid server descriptor", +} +_TOML_ESCAPES = { + '"': '\\"', + "\\": "\\\\", + "\b": "\\b", + "\t": "\\t", + "\n": "\\n", + "\f": "\\f", + "\r": "\\r", +} +# Native env_vars forwards same-named values, never aliases (rmcp-client/src/utils.rs). +# POSIX exec preserves native ownership while removing consumed secret aliases. +_STDIO_ENV_EXEC = ( + "import json,os,sys; " + "aliases=json.loads(sys.argv[1]); env=os.environ.copy(); " + "values={name:env.pop(source) for name,source in aliases.items()}; " + "env.update(values); argv=sys.argv[2:]; os.execvpe(argv[0],argv,env)" +) + + +class CodexMcpStartupError(RuntimeError): + """A startup stage or owned validation category, never callback diagnostics.""" + + def __init__(self, stage: str) -> None: + known = type(stage) is str and ( + stage in _STARTUP_STAGES or stage in _VALIDATION_MESSAGES + ) + self.stage = stage if known else "startup" + category = _VALIDATION_MESSAGES.get(self.stage) + message = ( + f"Codex MCP startup failed: {category}." + if category is not None + else f"Codex MCP startup failed during {self.stage}." + ) + super().__init__(message) + + +def _mapping_entries( + mapping: Mapping[Any, Any], stage: str +) -> tuple[tuple[Any, Any], ...]: + if not isinstance(mapping, Mapping): + raise CodexMcpStartupError(stage) + try: + return tuple((name, value) for name, value in mapping.items()) + except Exception: + # Mapping iteration and access are caller code, not owned validation. + safe = CodexMcpStartupError(stage) + logger.debug("Codex MCP mapping refused", exc_info=(type(safe), safe, None)) + raise safe from None + + +def _valid_text(value: Any, *, nonempty: bool = False) -> bool: + return ( + type(value) is str + and (bool(value) or not nonempty) + and "\0" not in value + and not any(0xD800 <= ord(char) <= 0xDFFF for char in value) + ) + + +def _string_mapping( + mapping: Mapping[str, str], stage: str, *, environment: bool = False +) -> dict[str, str]: + copied: dict[str, str] = {} + for name, value in _mapping_entries(mapping, stage): + if ( + not _valid_text(name, nonempty=True) + or not _valid_text(value) + or (environment and "=" in name) + ): + raise CodexMcpStartupError(stage) + copied[name] = value + return copied + + +def _argv( + sequence: Sequence[str], stage: str, *, nonempty: bool = False +) -> tuple[str, ...]: + if not isinstance(sequence, Sequence) or isinstance(sequence, str | bytes): + raise CodexMcpStartupError(stage) + try: + copied = tuple(sequence) + except Exception: + safe = CodexMcpStartupError(stage) + logger.debug("Codex MCP arguments refused", exc_info=(type(safe), safe, None)) + raise safe from None + if (nonempty and not copied) or any(not _valid_text(item) for item in copied): + raise CodexMcpStartupError(stage) + return copied + + +@dataclass(frozen=True, slots=True) +class CodexMcpStdioServer: + """A native stdio server; env is a frozen copy of child-only secret values.""" + + command: str + args: tuple[str, ...] = () + env: Annotated[Mapping[str, str], Field(exclude=True)] = field( + default_factory=dict, repr=False, compare=False + ) + cwd: str | None = None + + def __post_init__(self) -> None: + if not _valid_text(self.command, nonempty=True) or ( + self.cwd is not None and not _valid_text(self.cwd) + ): + raise CodexMcpStartupError("invalid-descriptor") + object.__setattr__(self, "args", _argv(self.args, "invalid-descriptor")) + child_env = _string_mapping(self.env, "invalid-descriptor", environment=True) + if any(name.upper() in _NON_INHERITABLE_STDIO_ENV for name in child_env): + raise CodexMcpStartupError("invalid-descriptor") + object.__setattr__(self, "env", MappingProxyType(child_env)) + + +@dataclass(frozen=True, slots=True) +class CodexMcpHttpServer: + """A streamable HTTP server; bearer and header values stay out of diagnostics.""" + + url: str + bearer_token: Annotated[str | None, Field(exclude=True)] = field( + default=None, repr=False, compare=False + ) + headers: Annotated[Mapping[str, str], Field(exclude=True)] = field( + default_factory=dict, repr=False, compare=False + ) + + def __post_init__(self) -> None: + if not _valid_text(self.url, nonempty=True) or ( + self.bearer_token is not None and not _valid_text(self.bearer_token) + ): + raise CodexMcpStartupError("invalid-descriptor") + object.__setattr__( + self, + "headers", + MappingProxyType(_string_mapping(self.headers, "invalid-descriptor")), + ) + + +CodexMcpServer: TypeAlias = CodexMcpStdioServer | CodexMcpHttpServer +CodexMcpResolver: TypeAlias = Callable[[str], Awaitable[Mapping[str, CodexMcpServer]]] + + +def _escaped_secrets(secret: str) -> set[str]: + variants = {secret} + frontier = {secret} + # JSON payloads can themselves appear as escaped strings in RPC diagnostics. + for _ in range(2): + escaped: set[str] = set() + for value in frontier: + for rendered in ( + json.dumps(value, ensure_ascii=True), + json.dumps(value, ensure_ascii=False), + repr(value), + ascii(value), + ): + escaped.add(rendered) + escaped.add(rendered[1:-1]) + escaped.add(value.encode("unicode_escape").decode("ascii")) + escaped.add(json.dumps(value, ensure_ascii=False)[1:-1].replace("/", "\\/")) + frontier = escaped - variants + variants.update(escaped) + return variants + + +class CodexMcpRedactor: + """Exact-value redaction scoped to one snapshot and its owned child launches.""" + + __slots__ = ("_secrets", "_variants") + + def __reduce_ex__(self, protocol: SupportsIndex, /) -> NoReturn: + raise TypeError( + "Codex MCP redactors are runtime-only and cannot be serialized." + ) + + def __init__(self, secrets: Iterable[str] = ()) -> None: + self._secrets: set[str] = set() + self._variants: tuple[str, ...] = () + self._register(secrets) + + def _register(self, secrets: Iterable[str]) -> None: + previous_count = len(self._secrets) + for secret in secrets: + if type(secret) is not str: + raise CodexMcpStartupError("snapshot") + if secret: + self._secrets.add(secret) + if len(self._secrets) == previous_count: + return + variants = { + variant for secret in self._secrets for variant in _escaped_secrets(secret) + } + # Secret-bearing regexes would outlive clear() in Python's global regex cache. + self._variants = tuple(sorted(variants, key=len, reverse=True)) + + def text(self, text: str) -> str: + matches = [ + (start, -len(secret), secret) + for secret in self._variants + if (start := text.find(secret)) >= 0 + ] + if not matches: + return text + heapify(matches) + parts: list[str] = [] + cursor = 0 + while matches: + start, negative_length, secret = heappop(matches) + if start >= cursor: + parts.extend((text[cursor:start], _REDACTED)) + cursor = start - negative_length + next_start = text.find(secret, cursor) + if next_start >= 0: + heappush(matches, (next_start, negative_length, secret)) + parts.append(text[cursor:]) + return "".join(parts) + + def value(self, value: Any) -> Any: + """Copy nested diagnostic data, including keys, without formatting unknown objects.""" + if isinstance(value, str): + return self.text(value) + if isinstance(value, Mapping): + redacted: dict[Any, Any] = {} + for key, item in value.items(): + safe_key = self.value(key) + unique_key = safe_key + suffix = 2 + while unique_key in redacted: + unique_key = f"{safe_key}#{suffix}" + suffix += 1 + redacted[unique_key] = self.value(item) + return redacted + if isinstance(value, list): + return [self.value(item) for item in value] + if isinstance(value, tuple): + return tuple(self.value(item) for item in value) + if ( + value is None + or type(value) is bool + or type(value) is int + or type(value) is float + ): + return value + return _REDACTED + + def clear(self) -> None: + """Release tracked values only after the owner has torn down its children.""" + self._secrets.clear() + self._variants = () + + +@dataclass(frozen=True, slots=True) +class CodexMcpLaunch: + """A private command/env copy for one native child, with no secret serialization.""" + + command: tuple[str, ...] + env: Annotated[Mapping[str, str] | None, Field(exclude=True)] = field( + default=None, repr=False, compare=False + ) + + def __post_init__(self) -> None: + object.__setattr__( + self, "command", _argv(self.command, "launch", nonempty=True) + ) + if self.env is not None: + object.__setattr__( + self, + "env", + MappingProxyType(_string_mapping(self.env, "launch", environment=True)), + ) + + +def _toml_string(value: str) -> str: + return ( + '"' + + "".join( + _TOML_ESCAPES.get(char, f"\\u{ord(char):04x}") + if ord(char) < 0x20 or ord(char) == 0x7F or char in _TOML_ESCAPES + else char + for char in value + ) + + '"' + ) + + +def _toml_value(value: Any) -> str: + match value: + case str(): + return _toml_string(value) + case bool(): + return "true" if value else "false" + case Mapping(): + return ( + "{" + + ", ".join( + f"{_toml_string(key)} = {_toml_value(item)}" + for key, item in value.items() + ) + + "}" + ) + case tuple() | list(): + return "[" + ", ".join(_toml_value(item) for item in value) + "]" + case _: + raise CodexMcpStartupError("launch") + + +def _server_secret_values(servers: Mapping[str, CodexMcpServer]) -> Iterable[str]: + for server in servers.values(): + match server: + case CodexMcpStdioServer(): + yield from server.env.values() + case CodexMcpHttpServer(): + if server.bearer_token is not None: + yield server.bearer_token + yield from server.headers.values() + + +@dataclass(frozen=True, slots=True) +class CodexMcpSnapshot: + """Frozen servers, expected native names, and a pre-registered diagnostic boundary.""" + + servers: Mapping[str, CodexMcpServer] + expected_names: frozenset[str] = field(init=False) + redactor: CodexMcpRedactor = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + entries = _mapping_entries(self.servers, "snapshot") + frozen: dict[str, CodexMcpServer] = {} + for name, server in entries: + if type(name) is not str or _SERVER_NAME.fullmatch(name) is None: + raise CodexMcpStartupError("invalid-name") + if name in _RESERVED_SERVER_NAMES: + raise CodexMcpStartupError("reserved-name") + if name in frozen: + raise CodexMcpStartupError("duplicate-name") + if ( + type(server) is not CodexMcpStdioServer + and type(server) is not CodexMcpHttpServer + ): + raise CodexMcpStartupError("unsupported-transport") + match server: + case CodexMcpStdioServer(): + if server.env and os.name != "posix": + raise CodexMcpStartupError("unsupported-platform") + frozen[name] = CodexMcpStdioServer( + server.command, server.args, server.env, server.cwd + ) + case CodexMcpHttpServer(): + frozen[name] = CodexMcpHttpServer( + server.url, server.bearer_token, server.headers + ) + object.__setattr__(self, "servers", MappingProxyType(frozen)) + object.__setattr__(self, "expected_names", frozenset(frozen)) + object.__setattr__( + self, "redactor", CodexMcpRedactor(_server_secret_values(frozen)) + ) + + def launch( + self, command: Sequence[str], env: Mapping[str, str] | None = None + ) -> CodexMcpLaunch: + """Build private overrides with aliases reserved against the full child env.""" + argv = list(_argv(command, "launch", nonempty=True)) + child_env = ( + _string_mapping(env, "launch", environment=True) if env is not None else {} + ) + self.redactor._register(child_env.values()) + namespace = uuid4().hex.upper() + used_names = {name.upper() for name in os.environ} + used_names.update(name.upper() for name in child_env) + counter = 0 + + def secret_ref(value: str) -> str: + nonlocal counter + while True: + name = f"BAND_CODEX_MCP_{namespace}_{counter}" + counter += 1 + if name not in used_names: + break + used_names.add(name) + child_env[name] = value + return name + + for name, server in self.servers.items(): + settings: dict[str, Any] = { + "enabled": True, + "required": True, + "startup_readiness": "connection", + } + match server: + case CodexMcpStdioServer(): + settings.update( + command=server.command, args=server.args, env_vars=[] + ) + if server.cwd is not None: + settings["cwd"] = server.cwd + if server.env: + aliases = { + key: secret_ref(value) for key, value in server.env.items() + } + settings.update( + command=sys.executable, + args=( + "-I", + "-c", + _STDIO_ENV_EXEC, + json.dumps(aliases, ensure_ascii=False), + server.command, + *server.args, + ), + env_vars=list(aliases.values()), + ) + case CodexMcpHttpServer(): + settings["url"] = server.url + settings["env_http_headers"] = { + key: secret_ref(value) for key, value in server.headers.items() + } + if server.bearer_token is not None: + settings["bearer_token_env_var"] = secret_ref( + server.bearer_token + ) + override = f"mcp_servers.{name}={_toml_value(settings)}" + try: + parsed = tomllib.loads(override) + except (tomllib.TOMLDecodeError, ValueError): + raise CodexMcpStartupError("launch") from None + if parsed["mcp_servers"][name] != json.loads(json.dumps(settings)): + raise CodexMcpStartupError("launch") + argv.extend(("-c", override)) + return CodexMcpLaunch( + tuple(argv), child_env if env is not None or child_env else None + ) + + +def freeze_mcp_servers(servers: Mapping[str, CodexMcpServer]) -> CodexMcpSnapshot: + """Validate and snapshot one resolver result before constructing a native client.""" + return CodexMcpSnapshot(servers) diff --git a/src/band/integrations/codex/rpc_base.py b/src/band/integrations/codex/rpc_base.py index 2fd379ff3..3261a3839 100644 --- a/src/band/integrations/codex/rpc_base.py +++ b/src/band/integrations/codex/rpc_base.py @@ -10,6 +10,7 @@ from dataclasses import dataclass from typing import Any, Literal +from band.integrations.codex.mcp import CodexMcpRedactor from band.integrations.codex.types import CodexRequestMethod logger = logging.getLogger(__name__) @@ -82,9 +83,11 @@ def __init__( *, retry_policy: OverloadRetryPolicy | None = None, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + redactor: CodexMcpRedactor | None = None, ) -> None: self.retry_policy = retry_policy or OverloadRetryPolicy() self._sleep = sleep + self._redactor = redactor self._pending: dict[JsonRpcId, asyncio.Future[dict[str, Any]]] = {} self._events: asyncio.Queue[RpcEvent] = asyncio.Queue(maxsize=10_000) @@ -92,6 +95,9 @@ def __init__( self._connected = False self._transport_label = "payload" # overridden by subclasses for diagnostics + def _redact_value(self, value: Any) -> Any: + return self._redactor.value(value) if self._redactor is not None else value + # ------------------------------------------------------------------ # Transport methods # ------------------------------------------------------------------ @@ -263,11 +269,21 @@ async def _dispatch_rpc_message(self, text: str) -> None: try: message = json.loads(text) except json.JSONDecodeError: - logger.warning("Skipping non-JSON %s: %s", self._transport_label, text) + logger.warning( + "Skipping non-JSON %s (type=%s, length=%d)", + self._transport_label, + type(text).__name__, + len(text), + ) return if not isinstance(message, dict): - logger.warning("Skipping non-object JSON-RPC payload: %s", message) + logger.warning( + "Skipping non-object JSON-RPC %s (type=%s, length=%d)", + self._transport_label, + type(message).__name__, + len(text), + ) return msg_id = message.get("id") @@ -285,7 +301,7 @@ async def _dispatch_rpc_message(self, text: str) -> None: ) ) except asyncio.QueueFull: - logger.warning("Event queue full; dropping server request %s", method) + logger.warning("Event queue full; dropping server request") return if method: @@ -300,16 +316,24 @@ async def _dispatch_rpc_message(self, text: str) -> None: ) ) except asyncio.QueueFull: - logger.warning("Event queue full; dropping notification %s", method) + logger.warning("Event queue full; dropping notification") return if msg_id is not None: future = self._pending.get(msg_id) if future is not None and not future.done(): + if "error" in message: + # Freeze diagnostics before teardown releases the secret scope. + message["error"] = self._redact_value(message["error"]) future.set_result(message) return - logger.debug("Ignoring unknown JSON-RPC payload: %s", message) + logger.debug( + "Ignoring unknown JSON-RPC %s (type=%s, length=%d)", + self._transport_label, + type(message).__name__, + len(text), + ) def _next_request_id(self) -> int: self._request_id += 1 diff --git a/src/band/integrations/codex/stdio_client.py b/src/band/integrations/codex/stdio_client.py index c87961e5a..61149657c 100644 --- a/src/band/integrations/codex/stdio_client.py +++ b/src/band/integrations/codex/stdio_client.py @@ -9,11 +9,11 @@ import shutil import signal from collections.abc import Awaitable, Callable, Mapping, Sequence -from contextlib import suppress from pathlib import Path from typing import Any -from band.core.exceptions import BandConnectionError +from band.core.exceptions import BandConnectionError, CleanupIncompleteError +from band.integrations.codex.mcp import CodexMcpRedactor from band.integrations.codex.rpc_base import BaseJsonRpcClient, OverloadRetryPolicy from band.integrations.codex.winspawn import ( MISSING_EXECUTABLE, @@ -27,6 +27,10 @@ logger = logging.getLogger(__name__) +_PROCESS_EXIT_TIMEOUT_S = 1.0 +_PROCESS_GROUP_EXIT_TIMEOUT_S = 1.0 +_PROCESS_GROUP_POLL_S = 0.02 + # Long enough for the child watcher to record an exit, short enough that a # held stdio pipe cannot stall room cleanup. _EXIT_OBSERVE_S = 2.0 @@ -98,18 +102,53 @@ async def _wait_readers(readers: Sequence[asyncio.Task[None]]) -> None: logger.warning("Codex stdio readers still pending after tree kill") +async def _wait_for_process_group_exit(pgid: int) -> bool: + loop = asyncio.get_running_loop() + deadline = loop.time() + _PROCESS_GROUP_EXIT_TIMEOUT_S + last_error: OSError | None = None + while True: + try: + os.killpg(pgid, 0) + except ProcessLookupError: + return True + except OSError as error: + last_error = error + else: + last_error = None + remaining = deadline - loop.time() + if remaining <= 0: + if last_error is not None: + raise CleanupIncompleteError( + "Codex owned process-group cleanup could not be confirmed" + ) from last_error + return False + await asyncio.sleep(min(_PROCESS_GROUP_POLL_S, remaining)) + + async def kill_process_tree(process: asyncio.subprocess.Process) -> None: - """Force-stop a Unix process and everything in its session. + """Settle the original owned POSIX group without signaling a reused PID. - The codex launcher is a shim whose real CLI is its child. On Unix the - child leads a session, so ``killpg`` still reaches that group after the - leader has exited. + Descendants can keep the group alive after its leader exits. POSIX has no + generation-bound group handle, so reuse between a probe and signal remains + a residual race. """ - if os.name != "nt": - if not _session_pid_reused(process): - with suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) + if os.name == "nt" or _session_pid_reused(process): return + pgid = process.pid + for signum in (signal.SIGTERM, signal.SIGKILL): + try: + os.killpg(pgid, signum) + except ProcessLookupError: + return + except OSError as error: + raise CleanupIncompleteError( + "Codex owned process group could not be signalled" + ) from error + if await _wait_for_process_group_exit(pgid): + return + raise CleanupIncompleteError( + "Codex owned process group remained unsettled after TERM and KILL" + ) class CodexStdioClient(BaseJsonRpcClient): @@ -123,8 +162,9 @@ def __init__( env: Mapping[str, str] | None = None, retry_policy: OverloadRetryPolicy | None = None, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + redactor: CodexMcpRedactor | None = None, ) -> None: - super().__init__(retry_policy=retry_policy, sleep=sleep) + super().__init__(retry_policy=retry_policy, sleep=sleep, redactor=redactor) self._transport_label = "codex output line" self.command = tuple(command) if command else self._default_codex_command() self.cwd = cwd @@ -239,10 +279,11 @@ async def close(self) -> None: self._close_task.add_done_callback(self._log_close_failure) await asyncio.shield(self._close_task) - @staticmethod - def _log_close_failure(task: asyncio.Task[None]) -> None: + def _log_close_failure(self, task: asyncio.Task[None]) -> None: if not task.cancelled() and (error := task.exception()) is not None: - logger.error("Codex subprocess cleanup failed", exc_info=error) + logger.error( + "Codex subprocess cleanup failed (type=%s)", type(error).__name__ + ) async def _close_process(self) -> None: readers = [ @@ -258,15 +299,19 @@ async def _close_process(self) -> None: await self._windows_process.close(_EXIT_OBSERVE_S) self._windows_process = None elif (process := self._proc) is not None: - await kill_process_tree(process) await close_stdin(process) + await await_returncode(process, _PROCESS_EXIT_TIMEOUT_S) + await kill_process_tree(process) if not await await_returncode(process, _EXIT_OBSERVE_S): - logger.warning( - "Codex process %s still alive after tree kill", process.pid + raise CleanupIncompleteError( + "Codex process did not report exit after tree cleanup" ) self._fail_pending("Codex stdio client closed") finally: await _wait_readers(readers) + if self._redactor is not None: + self._redactor.clear() + self.env = None async def _send_json(self, payload: dict[str, Any]) -> None: if not self._proc or not self._proc.stdin: @@ -303,9 +348,8 @@ async def _read_stderr_loop(self) -> None: line = await self._proc.stderr.readline() if not line: return - text = line.decode("utf-8", errors="replace").rstrip() - if text: - logger.debug("codex stderr: %s", text) + if line.rstrip(): + logger.debug("codex stderr (type=bytes, length=%d)", len(line)) except asyncio.CancelledError: raise except Exception: diff --git a/tests/adapters/codex/kit.py b/tests/adapters/codex/kit.py index d6b19736f..ad803cbb8 100644 --- a/tests/adapters/codex/kit.py +++ b/tests/adapters/codex/kit.py @@ -176,7 +176,7 @@ def patch_codex_clients_by_room( ) -> None: """Give each room its own fake app-server, as the real adapter does.""" - def _build(_workspace: str) -> FakeCodexClient: + def _build(_room: object, _launch: object) -> FakeCodexClient: room_id = adapter._active_room.get() assert room_id is not None return clients[room_id] diff --git a/tests/adapters/codex/lifecycleprofile.py b/tests/adapters/codex/lifecycleprofile.py index 34ca58b5e..ca01ea35f 100644 --- a/tests/adapters/codex/lifecycleprofile.py +++ b/tests/adapters/codex/lifecycleprofile.py @@ -12,6 +12,7 @@ from band.adapters.codex import CodexAdapter, CodexAdapterConfig from band.core.protocols import TurnResultAlreadyReported from band.integrations.codex import CodexJsonRpcError, CodexStdioClient +from band.integrations.codex.mcp import CodexMcpRedactor from tests.adapters.codex.kit import residency_of, thread_of from tests.adapters.hold import Hold from tests.adapters.lifecycle.kit import ( @@ -75,8 +76,9 @@ def child( command: Sequence[str] | None, cwd: str | None, env: Mapping[str, str] | None, + redactor: CodexMcpRedactor | None, ) -> CodexPeer: - peer = CodexPeer(command=command, cwd=cwd, env=env) + peer = CodexPeer(command=command, cwd=cwd, env=env, redactor=redactor) host.spawned(peer, cwd) return peer diff --git a/tests/adapters/codex/mcpkit.py b/tests/adapters/codex/mcpkit.py new file mode 100644 index 000000000..7a8660449 --- /dev/null +++ b/tests/adapters/codex/mcpkit.py @@ -0,0 +1,130 @@ +"""Shared room-startup doubles for native Codex MCP tests.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable, Iterator, Mapping +from contextlib import suppress +from pathlib import Path +from typing import Any + +from band.adapters.codex import CodexAdapter, CodexAdapterConfig +from band.integrations.codex.mcp import ( + CodexMcpHttpServer, + CodexMcpLaunch, + CodexMcpServer, +) +from band.integrations.codex.types import CodexRequestMethod +from tests.adapters.codexturns import FakeCodexClient + + +def ready_row(name: str = "private") -> dict[str, Any]: + return { + "name": name, + "serverInfo": {"name": "fixture", "version": "1"}, + "serverCapabilities": {}, + "toolsError": None, + "runtimeStatus": None, + "tools": {}, + "authStatus": "bearerToken", + } + + +class StatusClient(FakeCodexClient): + def __init__( + self, + pages: list[dict[str, Any]] | None = None, + checkpoint: Callable[[str], Awaitable[None]] | None = None, + ) -> None: + super().__init__() + self.pages = list(pages if pages is not None else [{"data": [ready_row()]}]) + self.checkpoint = checkpoint + + async def pause(self, stage: str) -> None: + if self.checkpoint is not None: + await self.checkpoint(stage) + + async def connect(self) -> None: + await self.pause("connect") + await super().connect() + + async def initialize(self, **kwargs: Any) -> dict[str, Any]: + await self.pause("initialize") + return await super().initialize(**kwargs) + + async def request( + self, method: str, params: dict[str, Any] | None = None, **kwargs: Any + ) -> dict[str, Any]: + if method == "mcpServerStatus/list": + await self.pause("readiness") + self.requests.append((method, params or {})) + return self.pages.pop(0) + await self.pause( + "skills" if method == CodexRequestMethod.SKILLS_EXTRA_ROOTS_SET else "model" + ) + return await super().request(method, params, **kwargs) + + +def mcp_adapter( + tmp_path: Path, + clients: list[FakeCodexClient], + resolver: Callable[[str], Awaitable[Mapping[str, CodexMcpServer]]], + **config: Any, +) -> tuple[CodexAdapter, list[CodexMcpLaunch | None]]: + adapter = CodexAdapter( + CodexAdapterConfig( + workspace_for_room=lambda room: str(tmp_path / room), **config + ), + resolve_mcp_servers=resolver, + ) + launches: list[CodexMcpLaunch | None] = [] + + def build(_room: object, launch: CodexMcpLaunch | None) -> FakeCodexClient: + launches.append(launch) + return clients[len(launches) - 1] + + adapter._build_client = build + return adapter, launches + + +async def start(adapter: CodexAdapter, room_id: str = "a") -> None: + room = adapter._room_client(room_id) + adapter._active_room.set(room_id) + try: + async with room.retained.use(): + pass + except BaseException: + # A direct helper caller observes the owned recovery, not just its waiter. + if room.retained._transition is not None: + await asyncio.wait({room.retained._transition}) + with suppress(RuntimeError): + await room.retained._settle() + raise + + +async def private_servers(_workspace: str) -> Mapping[str, CodexMcpServer]: + return { + "private": CodexMcpHttpServer( + "http://localhost/mcp", bearer_token="secret-sentinel" + ) + } + + +class DuplicateServers(Mapping[str, Any]): + """A caller mapping can yield the same name twice from items().""" + + def __getitem__(self, key: str) -> Any: + if key != "private": + raise KeyError(key) + return CodexMcpHttpServer("http://localhost/mcp") + + def __iter__(self) -> Iterator[str]: + return iter(("private", "private")) + + def __len__(self) -> int: + return 2 + + +class SseServer: + transport = "sse" + url = "http://localhost/sse" diff --git a/tests/adapters/codex/test_mcp_acquisition.py b/tests/adapters/codex/test_mcp_acquisition.py new file mode 100644 index 000000000..3eb9cb35f --- /dev/null +++ b/tests/adapters/codex/test_mcp_acquisition.py @@ -0,0 +1,38 @@ +"""External MCP resolution follows the retained room's acquisition boundary.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from band.adapters.codex import CodexAdapter, CodexAdapterConfig +from band.integrations.codex.mcp import CodexMcpHttpServer, CodexMcpStartupError +from tests.adapters.codexturns import FakeCodexClient + + +class TestMcpAcquisition: + async def test_resolves_once_per_process_before_readiness( + self, tmp_path: Path + ) -> None: + resolved: list[str] = [] + + async def resolve(workspace: str): + resolved.append(workspace) + return {"private": CodexMcpHttpServer("http://localhost/mcp")} + + adapter = CodexAdapter( + CodexAdapterConfig(workspace_for_room=lambda _: str(tmp_path)), + resolve_mcp_servers=resolve, + ) + client = FakeCodexClient() + adapter._build_client = lambda _room, _launch: client + room = adapter._room_client("a") + try: + with pytest.raises(CodexMcpStartupError): + async with room.retained.use(): + pass + assert resolved == [str(tmp_path)] + assert "mcpServerStatus/list" in client.request_methods + finally: + await adapter.cleanup_all() diff --git a/tests/adapters/codex/test_mcp_cancellation.py b/tests/adapters/codex/test_mcp_cancellation.py new file mode 100644 index 000000000..a89b3f915 --- /dev/null +++ b/tests/adapters/codex/test_mcp_cancellation.py @@ -0,0 +1,204 @@ +"""Public invocation cancellation keeps settlement and discovery custody.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Mapping +from typing import Any + +import pytest + +from band.core.exceptions import CleanupIncompleteError +from band.integrations.codex import CodexJsonRpcError +from band.integrations.codex.mcp import CodexMcpServer +from band.integrations.codex.types import CodexRequestMethod, CodexSessionState +from band.testing import FakeAgentTools +from tests.adapters.codex.kit import make_platform_message +from tests.adapters.codex.mcpkit import ( + StatusClient, + mcp_adapter, + private_servers, +) + + +def message_task(adapter) -> asyncio.Task[None]: + return asyncio.create_task( + adapter.on_message( + msg=make_platform_message(room_id="a"), + tools=FakeAgentTools(), + history=CodexSessionState(), + participants_msg=None, + contacts_msg=None, + is_session_bootstrap=True, + room_id="a", + ) + ) + + +class TestMcpCancellation: + @pytest.mark.looptime + async def test_external_startup_cancel_outlasting_grace_retains_custody( + self, + tmp_path, + ) -> None: + entered, settling, release, settled = (asyncio.Event() for _ in range(4)) + + async def resolve(_workspace: str) -> Mapping[str, CodexMcpServer]: + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + settling.set() + await release.wait() + settled.set() + raise + + adapter, launches = mcp_adapter( + tmp_path, [StatusClient()], resolve, turn_settle_timeout_s=0.02 + ) + await adapter.on_started("Fixture", "Cancellation custody") + running = message_task(adapter) + try: + await entered.wait() + room = adapter._room_clients["a"] + invocation = room.startup_owner.task + assert invocation is not None + turn_release = adapter._turn_release["a"] + invocation.cancel() + await settling.wait() + outcome = await asyncio.gather( + adapter.on_cleanup("a"), return_exceptions=True + ) + assert isinstance(outcome[0], CleanupIncompleteError) + assert invocation.cancelling() == 1 + assert not invocation.done() + assert not settled.is_set() + assert not turn_release.done() + assert adapter._turn_release["a"] is turn_release + with pytest.raises(ValueError): + adapter._workspaces.claim("b", room.workspace) + assert launches == [] + release.set() + await asyncio.gather(running, return_exceptions=True) + await adapter.on_cleanup("a") + assert settled.is_set() + assert "a" not in adapter._workspaces.rooms + finally: + release.set() + await asyncio.gather(running, return_exceptions=True) + await adapter.cleanup_all() + adapter._workspaces.release("b") + + @pytest.mark.looptime + async def test_cleanup_disarms_readiness_deadline_during_settlement( + self, tmp_path, monkeypatch + ) -> None: + entered, settling, release, settled = (asyncio.Event() for _ in range(4)) + monkeypatch.setattr("band.adapters.codex._MCP_READINESS_TIMEOUT_S", 0.1) + + class SettlingClient(StatusClient): + async def request( + self, method: str, params: dict[str, Any] | None = None, **kwargs: Any + ) -> dict[str, Any]: + if method == "mcpServerStatus/list": + self.requests.append((method, params or {})) + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + settling.set() + await release.wait() + settled.set() + raise + return await super().request(method, params, **kwargs) + + client = SettlingClient() + adapter, _ = mcp_adapter( + tmp_path, [client], private_servers, turn_settle_timeout_s=0.02 + ) + await adapter.on_started("Fixture", "Readiness custody") + running = message_task(adapter) + cleanup = None + try: + await entered.wait() + room = adapter._room_clients["a"] + invocation = room.startup_owner.task + assert invocation is not None + turn_release = adapter._turn_release["a"] + cleanup = asyncio.create_task(adapter.on_cleanup("a")) + await settling.wait() + await asyncio.sleep( + 0.2 + ) # Cross the original readiness deadline while settling. + assert invocation.cancelling() == 1 + assert not invocation.done() + assert not settled.is_set() + assert not client.closed + assert room.client is client + assert not turn_release.done() + assert adapter._turn_release["a"] is turn_release + with pytest.raises(ValueError): + adapter._workspaces.claim("b", room.workspace) + outcome = await asyncio.gather(cleanup, return_exceptions=True) + assert isinstance(outcome[0], CleanupIncompleteError) + release.set() + await asyncio.gather(running, return_exceptions=True) + await adapter.on_cleanup("a") + assert settled.is_set() + assert client.closed + finally: + release.set() + await asyncio.gather(running, return_exceptions=True) + if cleanup is not None: + await asyncio.gather(cleanup, return_exceptions=True) + await adapter.cleanup_all() + adapter._workspaces.release("b") + + @pytest.mark.looptime + async def test_thread_error_after_cleanup_issues_no_discovery( + self, tmp_path + ) -> None: + entered, cancelling, release = (asyncio.Event() for _ in range(3)) + + class LateThreadError(StatusClient): + async def request( + self, method: str, params: dict[str, Any] | None = None, **kwargs: Any + ) -> dict[str, Any]: + if method == CodexRequestMethod.THREAD_START: + self.requests.append((method, params or {})) + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelling.set() + await release.wait() + raise CodexJsonRpcError(code=-32603, message="thread failed") + return await super().request(method, params, **kwargs) + + client = LateThreadError() + adapter, _ = mcp_adapter( + tmp_path, [client], private_servers, turn_settle_timeout_s=0.02 + ) + await adapter.on_started("Fixture", "Retirement discovery fence") + running = message_task(adapter) + cleanup = None + try: + await entered.wait() + initial_discovery = client.request_methods.count("mcpServerStatus/list") + assert initial_discovery == 1 + cleanup = asyncio.create_task(adapter.on_cleanup("a")) + await cancelling.wait() + release.set() + await asyncio.gather(running, cleanup, return_exceptions=True) + assert ( + client.request_methods.count("mcpServerStatus/list") + == initial_discovery + ) + await adapter.on_cleanup("a") + assert client.closed + finally: + release.set() + await asyncio.gather(running, return_exceptions=True) + if cleanup is not None: + await asyncio.gather(cleanup, return_exceptions=True) + await adapter.cleanup_all() diff --git a/tests/adapters/codex/test_mcp_launch.py b/tests/adapters/codex/test_mcp_launch.py new file mode 100644 index 000000000..51b5c7038 --- /dev/null +++ b/tests/adapters/codex/test_mcp_launch.py @@ -0,0 +1,659 @@ +"""Consumer-visible snapshot, launch and diagnostic safety for Codex MCP.""" + +from __future__ import annotations + +import json +import os +import pickle +import subprocess +import sys +import tomllib +import traceback +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import FrozenInstanceError, asdict +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from uuid import UUID + +import pytest +from pydantic import TypeAdapter + +from band.integrations.codex import CodexStdioClient +from band.integrations.codex import mcp as codex_mcp +from band.integrations.codex.mcp import ( + CodexMcpHttpServer, + CodexMcpLaunch, + CodexMcpRedactor, + CodexMcpServer, + CodexMcpStartupError, + CodexMcpStdioServer, + freeze_mcp_servers, +) + + +def launched_servers( + launch: CodexMcpLaunch, original_command: Sequence[str] +) -> dict[str, Any]: + tail = launch.command[len(original_command) :] + assert len(tail) % 2 == 0 + assert all(flag == "-c" for flag in tail[::2]) + parsed = tomllib.loads("\n".join(tail[1::2])) + return parsed.get("mcp_servers", {}) + + +def run_stdio_entry(entry: dict[str, Any], launch: CodexMcpLaunch) -> dict[str, Any]: + assert launch.env is not None + inherited_names = {"SYSTEMROOT", "WINDIR", "PATH", "HOME", "TEMP", "TMP", "TMPDIR"} + child_env = { + name: value + for name, value in os.environ.items() + if name.upper() in inherited_names + } + child_env.update({name: launch.env[name] for name in entry["env_vars"]}) + completed = subprocess.run( + [entry["command"], *entry["args"]], + cwd=entry["cwd"], + env=child_env, + capture_output=True, + text=True, + check=True, + ) + return json.loads(completed.stdout) + + +class TestCodexMcpLaunch: + def test_caller_mutation_cannot_change_snapshot_or_launch( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + codex_mcp, "os", SimpleNamespace(name="posix", environ=os.environ) + ) + args = ["--mode", "original"] + env = {"TOKEN": "stdio-original-grant"} + headers = {"X-Grant": "http-original-grant"} + stdio = CodexMcpStdioServer("relative-runner", args=args, env=env) + http = CodexMcpHttpServer("http://localhost/mcp", headers=headers) + servers: dict[str, CodexMcpServer] = {"stdio": stdio, "http": http} + snapshot = freeze_mcp_servers(servers) + command = ["codex", "app-server", "--listen", "stdio://"] + base_env = {"EXISTING": "existing-value"} + launch = snapshot.launch(command, base_env) + + args.append("mutated") + env["TOKEN"] = "stdio-replacement-grant" + headers["X-Grant"] = "http-replacement-grant" + servers.clear() + command.append("--changed") + base_env["EXISTING"] = "changed" + + assert snapshot.expected_names == frozenset({"stdio", "http"}) + assert snapshot.servers["stdio"].args == ("--mode", "original") + entries = launched_servers(launch, command[:-1]) + assert "mutated" not in entries["stdio"]["args"] + assert launch.env is not None + assert launch.env["EXISTING"] == "existing-value" + assert launch.env[entries["http"]["env_http_headers"]["X-Grant"]] == ( + "http-original-grant" + ) + assert "stdio-original-grant" in launch.env.values() + assert "stdio-replacement-grant" not in launch.env.values() + assert launch.command[:4] == ("codex", "app-server", "--listen", "stdio://") + assert snapshot.redactor.text("stdio-original-grant") == "[REDACTED]" + + with pytest.raises(TypeError): + snapshot.servers["new"] = http + with pytest.raises(TypeError): + stdio.env["TOKEN"] = "changed" + with pytest.raises(TypeError): + http.headers["X-Grant"] = "changed" + with pytest.raises(FrozenInstanceError): + stdio.command = "changed" + with pytest.raises(FrozenInstanceError): + http.url = "changed" + with pytest.raises(TypeError): + launch.env["EXISTING"] = "changed" + + def test_empty_snapshot_does_not_add_overrides_or_child_env(self) -> None: + command = ["codex", "app-server", "--listen", "stdio://"] + snapshot = freeze_mcp_servers({}) + launch = snapshot.launch(command, None) + assert launch.command == tuple(command) + assert launch.env is None + assert snapshot.expected_names == frozenset() + assert snapshot.redactor.text("unregistered value") == "unregistered value" + + base_env = {"BASE": "base-secret"} + launch_with_env = snapshot.launch(command, base_env) + base_env.clear() + assert launch_with_env.env == {"BASE": "base-secret"} + assert "base-secret" not in repr(launch_with_env) + + @pytest.mark.parametrize( + "unsupported", + [ + SimpleNamespace(transport="sse", url="secret-url"), + {"url": "secret-url", "transport": "http"}, + type("secret-callback-type", (), {})(), + type("HttpSubclass", (CodexMcpHttpServer,), {})("http://localhost/mcp"), + ], + ids=["sse", "untyped-mapping", "unknown-type", "descriptor-subclass"], + ) + def test_unsupported_descriptors_have_named_secret_free_errors( + self, unsupported: Any + ) -> None: + with pytest.raises(CodexMcpStartupError) as caught: + freeze_mcp_servers({"tools": unsupported}) + assert caught.value.stage == "unsupported-transport" + assert "unsupported transport" in str(caught.value).lower() + assert "secret-url" not in str(caught.value) + assert "secret-callback-type" not in repr(caught.value) + + def test_a_callback_type_cannot_spoof_supported_descriptor_identity(self) -> None: + class PretendType(type): + def __eq__(self, other: object) -> bool: + return True + + class PretendDescriptor(metaclass=PretendType): + pass + + with pytest.raises(CodexMcpStartupError) as caught: + freeze_mcp_servers({"tools": PretendDescriptor()}) + assert caught.value.stage == "unsupported-transport" + + @pytest.mark.parametrize( + "name", ["", "with.dot", "with space", "with/slash", "café", "secret-name\n", 7] + ) + def test_invalid_server_names_have_safe_named_errors(self, name: Any) -> None: + with pytest.raises(CodexMcpStartupError) as caught: + freeze_mcp_servers({name: CodexMcpHttpServer("http://localhost/mcp")}) + assert caught.value.stage == "invalid-name" + assert "invalid server name" in str(caught.value).lower() + assert "secret-name" not in repr(caught.value) + + def test_native_host_owned_name_is_reserved_without_invented_names(self) -> None: + server = CodexMcpHttpServer("http://localhost/mcp") + with pytest.raises(CodexMcpStartupError) as caught: + freeze_mcp_servers({"codex_apps": server}) + assert caught.value.stage == "reserved-name" + snapshot = freeze_mcp_servers( + {name: server for name in ("band", "codex", "Codex_apps", "tools-1_v2")} + ) + launch = snapshot.launch(("codex", "app-server"), None) + assert set(launched_servers(launch, ("codex", "app-server"))) == { + "band", + "codex", + "Codex_apps", + "tools-1_v2", + } + + @pytest.mark.parametrize("value", [None, [], "private-nonmapping-data"]) + def test_nonmapping_resolver_results_have_safe_snapshot_errors( + self, value: Any + ) -> None: + with pytest.raises(CodexMcpStartupError) as caught: + freeze_mcp_servers(value) + assert caught.value.stage == "snapshot" + assert "private-nonmapping-data" not in repr(caught.value) + + @pytest.mark.parametrize( + "make_server", + [ + lambda: CodexMcpStdioServer("runner", args="private-invalid-args"), + lambda: CodexMcpStdioServer("runner", env={"TOKEN": 7}), + lambda: CodexMcpStdioServer( + "runner", env={"BAD=NAME": "private-env-value"} + ), + lambda: CodexMcpHttpServer("http://localhost/mcp", bearer_token=object()), + lambda: CodexMcpHttpServer("http://localhost/mcp", headers={"X-Grant": 7}), + ], + ids=[ + "string-args", + "nonstring-env", + "invalid-env-key", + "nonstring-bearer", + "nonstring-header", + ], + ) + def test_invalid_descriptor_data_is_rejected_without_input_diagnostics( + self, make_server: Any + ) -> None: + with pytest.raises(CodexMcpStartupError) as caught: + make_server() + assert caught.value.stage == "invalid-descriptor" + assert "private-invalid-args" not in repr(caught.value) + assert "private-env-value" not in repr(caught.value) + + @pytest.mark.parametrize( + "name", + [ + "CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN", + "node_repl_auth_token", + "CODEX_GUARDIAN_DECISIONS_API_KEY", + "OPENAI_FEDERATION_RULE_ID", + "OPENAI_IDENTITY_TOKEN_FILE", + "OPENAI_WORKLOAD_IDENTITY_CONTEXT", + ], + ) + def test_stdio_alias_remapping_cannot_restore_native_noninheritable_env( + self, name: str + ) -> None: + with pytest.raises(CodexMcpStartupError) as caught: + CodexMcpStdioServer("runner", env={name: "private-native-launch-context"}) + assert caught.value.stage == "invalid-descriptor" + assert "private-native-launch-context" not in repr(caught.value) + + def test_duplicate_mapping_entries_are_not_silently_overwritten(self) -> None: + class RepeatedNames(Mapping[str, CodexMcpServer]): + def __iter__(self) -> Iterator[str]: + return iter(("repeated", "repeated")) + + def __len__(self) -> int: + return 2 + + def __getitem__(self, key: str) -> CodexMcpServer: + return CodexMcpHttpServer("http://localhost/mcp") + + with pytest.raises(CodexMcpStartupError) as caught: + freeze_mcp_servers(RepeatedNames()) + assert caught.value.stage == "duplicate-name" + assert "duplicate server name" in str(caught.value).lower() + + @pytest.mark.parametrize("failure", ["iteration", "lookup", "malformed-items"]) + def test_caller_mapping_failures_are_normalized_before_validation( + self, failure: str + ) -> None: + secret = "private-callback-error" + + class BrokenMapping(Mapping[str, CodexMcpServer]): + def __iter__(self) -> Iterator[str]: + if failure == "iteration": + raise RuntimeError(secret) + return iter(("tools",)) + + def __len__(self) -> int: + return 1 + + def __getitem__(self, key: str) -> CodexMcpServer: + raise ValueError(secret) + + def items(self) -> Any: + if failure == "malformed-items": + return [ + ("tools", CodexMcpHttpServer("http://localhost/mcp"), secret) + ] + return super().items() + + with pytest.raises(CodexMcpStartupError) as caught: + freeze_mcp_servers(BrokenMapping()) + assert caught.value.stage == "snapshot" + rendered = "".join(traceback.format_exception(caught.value)) + assert secret not in str(caught.value) + assert secret not in repr(caught.value) + assert secret not in rendered + + @pytest.mark.parametrize( + "command", + [ + ["codex", "app-server", "--listen", "stdio://", "-c", 'model="base"'], + [ + "python", + "-m", + "band_agent", + "codex-relay", + "--skills-root", + "skills", + "--", + "codex", + "app-server", + "--listen", + "stdio://", + "--extra-arg", + ], + ], + ids=["default-command-extra-config", "relay-boundary-extra-args"], + ) + def test_native_overrides_follow_existing_default_or_relay_argv( + self, command: list[str] + ) -> None: + original = command.copy() + snapshot = freeze_mcp_servers( + {"tools": CodexMcpHttpServer("http://localhost/mcp")} + ) + launch = snapshot.launch(command, None) + assert command == original + assert launch.command[: len(original)] == tuple(original) + assert launched_servers(launch, original)["tools"] == { + "url": "http://localhost/mcp", + "env_http_headers": {}, + "enabled": True, + "required": True, + "startup_readiness": "connection", + } + + def test_toml_preserves_nonsecret_strings_without_an_unneeded_trampoline( + self, + ) -> None: + args = ('quote" and backslash\\', "line\n\tbreak", "emoji-🦉", "delete-\x7f") + server = CodexMcpStdioServer("./relative-runner", args=args, cwd="relative dir") + command = ("codex", "app-server") + launch = freeze_mcp_servers({"tools": server}).launch(command, None) + entry = launched_servers(launch, command)["tools"] + assert entry["command"] == "./relative-runner" + assert entry["args"] == list(args) + assert entry["cwd"] == "relative dir" + assert entry["env_vars"] == [] + assert entry["required"] is True + assert entry["startup_readiness"] == "connection" + assert launch.env is None + + def test_secrets_only_exist_in_the_private_child_environment( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + codex_mcp, "os", SimpleNamespace(name="posix", environ=os.environ) + ) + bearer = 'bearer-secret-"\\-🦉' + header = "header-secret-with\nnewline" + stdio_value = "stdio-secret-value" + servers = { + "remote": CodexMcpHttpServer( + "http://localhost/mcp", bearer_token=bearer, headers={"X.Grant": header} + ), + "local": CodexMcpStdioServer("runner", env={"TOKEN": stdio_value}), + } + command = ["codex", "app-server"] + base_env = {"BASE": "base-secret-value"} + process_env_before = dict(os.environ) + snapshot = freeze_mcp_servers(servers) + launch = snapshot.launch(command, base_env) + entries = launched_servers(launch, command) + assert launch.env is not None + remote = entries["remote"] + assert launch.env[remote["bearer_token_env_var"]] == bearer + assert launch.env[remote["env_http_headers"]["X.Grant"]] == header + assert [launch.env[name] for name in entries["local"]["env_vars"]] == [ + stdio_value + ] + assert "TOKEN" not in launch.env + assert "bearer_token" not in remote + assert "http_headers" not in remote + assert "env" not in entries["local"] + assert dict(os.environ) == process_env_before + assert base_env == {"BASE": "base-secret-value"} + + diagnostic = repr((snapshot, snapshot.redactor, launch, *servers.values())) + for secret in (bearer, header, stdio_value, "base-secret-value"): + assert secret not in " ".join(launch.command) + assert secret not in diagnostic + assert snapshot.redactor.text(secret) == "[REDACTED]" + + def test_descriptor_equality_and_serialization_do_not_use_secret_fields( + self, + ) -> None: + left = CodexMcpHttpServer( + "http://localhost/mcp", + bearer_token="first-secret", + headers={"X-Grant": "one"}, + ) + right = CodexMcpHttpServer( + "http://localhost/mcp", + bearer_token="second-secret", + headers={"X-Grant": "two"}, + ) + stdio = CodexMcpStdioServer("runner", env={"TOKEN": "stdio-secret"}) + assert left == right + assert hash(left) == hash(right) + assert stdio == CodexMcpStdioServer("runner", env={"TOKEN": "other-secret"}) + assert TypeAdapter(CodexMcpHttpServer).dump_python(left) == { + "url": "http://localhost/mcp" + } + assert TypeAdapter(CodexMcpStdioServer).dump_python(stdio) == { + "command": "runner", + "args": (), + "cwd": None, + } + with pytest.raises(TypeError): + asdict(left) + with pytest.raises(TypeError): + asdict(stdio) + launch = freeze_mcp_servers({"remote": left}).launch( + ("codex", "app-server"), None + ) + assert ( + "first-secret" not in TypeAdapter(CodexMcpLaunch).dump_json(launch).decode() + ) + for runtime in (left, stdio, launch, CodexMcpRedactor(("first-secret",))): + with pytest.raises(TypeError) as caught: + pickle.dumps(runtime) + assert "first-secret" not in repr(caught.value) + + def test_windows_allows_env_free_stdio_and_secret_http( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + codex_mcp, "os", SimpleNamespace(name="nt", environ=os.environ) + ) + probe = ( + "import json,os,sys; " + "print(json.dumps({'argv':sys.argv[1:], 'env_names':list(os.environ)}))" + ) + command = ("codex", "app-server") + snapshot = freeze_mcp_servers( + { + "local": CodexMcpStdioServer( + sys.executable, + args=("-I", "-c", probe, "literal argument"), + cwd=str(tmp_path), + ), + "remote": CodexMcpHttpServer( + "http://localhost/mcp", + bearer_token="private-http-platform-grant", + headers={"X-Grant": "private-http-header-grant"}, + ), + } + ) + launch = snapshot.launch(command, None) + entries = launched_servers(launch, command) + assert set(entries) == {"local", "remote"} + observed = run_stdio_entry(entries["local"], launch) + assert observed["argv"] == ["literal argument"] + assert "TOKEN" not in observed["env_names"] + assert not any( + name.startswith("BAND_CODEX_MCP_") for name in observed["env_names"] + ) + + @pytest.mark.skipif( + os.name != "posix", + reason="secret-bearing stdio alias ownership requires POSIX exec", + ) + def test_stdio_env_aliases_deliver_distinct_same_named_grants_to_relative_targets( + self, + ) -> None: + target = str(Path(".") / Path(sys.executable).name) + if not target.startswith("."): + target = os.path.join(".", target) + cwd = str(Path(sys.executable).parent) + probe = ( + "import json,os,sys; " + "print(json.dumps({'token':os.environ['TOKEN'], 'argv':sys.argv[1:], " + "'cwd':os.getcwd(), 'env_names':list(os.environ)}))" + ) + command = ("codex", "app-server") + snapshot = freeze_mcp_servers( + { + "first": CodexMcpStdioServer( + target, + args=("-I", "-c", probe, "literal argument"), + env={"TOKEN": "first-server-grant"}, + cwd=cwd, + ), + "second": CodexMcpStdioServer( + target, + args=("-I", "-c", probe, "literal argument"), + env={"TOKEN": "second-server-grant"}, + cwd=cwd, + ), + } + ) + launch = snapshot.launch(command, None) + entries = launched_servers(launch, command) + assert set(entries["first"]["env_vars"]).isdisjoint( + entries["second"]["env_vars"] + ) + for name, expected in ( + ("first", "first-server-grant"), + ("second", "second-server-grant"), + ): + observed = run_stdio_entry(entries[name], launch) + assert observed["token"] == expected + assert observed["argv"] == ["literal argument"] + assert Path(observed["cwd"]).samefile(cwd) + assert set(observed["env_names"]).isdisjoint(entries[name]["env_vars"]) + + def test_fresh_launches_do_not_reuse_generated_secret_names(self) -> None: + snapshot = freeze_mcp_servers( + {"tools": CodexMcpHttpServer("http://localhost/mcp", bearer_token="grant")} + ) + command = ("codex", "app-server") + first = snapshot.launch(command, None) + second = snapshot.launch(command, None) + first_ref = launched_servers(first, command)["tools"]["bearer_token_env_var"] + second_ref = launched_servers(second, command)["tools"]["bearer_token_env_var"] + assert first_ref != second_ref + assert first.env[first_ref] == second.env[second_ref] == "grant" + + @pytest.mark.parametrize( + "supplied_env", [False, True], ids=["inherited", "overrides"] + ) + def test_secret_aliases_reserve_effective_stdio_child_environment( + self, monkeypatch: pytest.MonkeyPatch, supplied_env: bool + ) -> None: + namespace = UUID("00000000-0000-0000-0000-000000000005") + inherited_name = "BAND_CODEX_MCP_00000000000000000000000000000005_0" + override_name = "BAND_CODEX_MCP_00000000000000000000000000000005_1" + command = ("codex", "app-server") + snapshot = freeze_mcp_servers( + { + "tools": CodexMcpHttpServer( + "http://localhost/mcp", bearer_token="correct-server-grant" + ) + } + ) + with monkeypatch.context() as scoped: + scoped.setattr("band.integrations.codex.mcp.uuid4", lambda: namespace) + scoped.setenv(inherited_name, "inherited-value") + scoped.setenv(override_name, "inherited-override-value") + env = {override_name: "supplied-override-value"} if supplied_env else None + launch = snapshot.launch(command, env) + client = CodexStdioClient( + command=launch.command, env=launch.env, redactor=snapshot.redactor + ) + reference = launched_servers(launch, command)["tools"]["bearer_token_env_var"] + assert client.env is not None + assert client.env[inherited_name] == "inherited-value" + assert client.env[override_name] == ( + "supplied-override-value" if supplied_env else "inherited-override-value" + ) + assert reference not in {inherited_name, override_name} + assert client.env[reference] == "correct-server-grant" + + @pytest.mark.parametrize( + "stage", + [ + "resolution", + "snapshot", + "launch", + "readiness", + "thread/start", + "thread/resume", + ], + ) + def test_owned_startup_stages_have_fixed_safe_text(self, stage: str) -> None: + error = CodexMcpStartupError(stage) + assert error.stage == stage + assert stage in str(error) + assert str(error) == error.args[0] + + def test_untrusted_startup_stage_text_is_not_reflected(self) -> None: + error = CodexMcpStartupError("secret-callback-message") + assert error.stage == "startup" + assert "secret-callback-message" not in str(error) + assert "secret-callback-message" not in repr(error) + + @pytest.mark.parametrize( + "render", + [ + lambda value: value, + lambda value: json.dumps(value, ensure_ascii=True)[1:-1], + lambda value: json.dumps(value, ensure_ascii=False)[1:-1], + lambda value: repr(value)[1:-1], + lambda value: ascii(value)[1:-1], + lambda value: value.encode("unicode_escape").decode(), + lambda value: json.dumps(json.dumps(value))[1:-1], + ], + ids=[ + "raw", + "json-ascii", + "json-unicode", + "repr", + "ascii-repr", + "unicode-escape", + "nested-json", + ], + ) + def test_redaction_covers_raw_and_escaped_secret_forms(self, render: Any) -> None: + secret = 'private-"quote\\slash\nline-é-🦉' + snapshot = freeze_mcp_servers( + {"tools": CodexMcpHttpServer("http://localhost/mcp", bearer_token=secret)} + ) + message = f"before {render(secret)} after" + assert snapshot.redactor.text(message) == "before [REDACTED] after" + assert secret not in repr(snapshot.redactor) + + def test_redaction_keeps_distinct_fields_when_secret_keys_collide(self) -> None: + redactor = CodexMcpRedactor(("first-grant", "second-grant")) + assert redactor.value( + {"first-grant": "first error", "second-grant": "second error"} + ) == { + "[REDACTED]": "first error", + "[REDACTED]#2": "second error", + } + + def test_redaction_preserves_nested_nonsecret_data_and_does_not_mutate_input( + self, + ) -> None: + secret = "private-nested-grant" + redactor = CodexMcpRedactor((secret,)) + payload = { + "message": f"received {secret}", + secret: [secret, {"nested": (secret, 42, True, None)}], + "safe": "public explanation", + } + assert redactor.value(payload) == { + "message": "received [REDACTED]", + "[REDACTED]": ["[REDACTED]", {"nested": ("[REDACTED]", 42, True, None)}], + "safe": "public explanation", + } + assert payload["message"] == f"received {secret}" + assert payload[secret][0] == secret + assert redactor.value(SimpleNamespace(secret=secret)) == "[REDACTED]" + + def test_redaction_handles_overlapping_values_without_rewriting_its_marker( + self, + ) -> None: + redactor = CodexMcpRedactor(("grant", "grant-longer", "REDACTED")) + assert redactor.text("grant-longer grant REDACTED") == ( + "[REDACTED] [REDACTED] [REDACTED]" + ) + + def test_redaction_scope_clears_only_its_own_secret_references(self) -> None: + first = CodexMcpRedactor(("first-grant", "")) + second = CodexMcpRedactor(("second-grant",)) + assert first.text("first-grant second-grant") == "[REDACTED] second-grant" + assert second.text("first-grant second-grant") == "first-grant [REDACTED]" + first.clear() + first.clear() + assert first.text("first-grant") == "first-grant" + assert first.value({"detail": "first-grant"}) == {"detail": "first-grant"} + assert second.text("second-grant") == "[REDACTED]" diff --git a/tests/adapters/codex/test_mcp_owner_lifecycle.py b/tests/adapters/codex/test_mcp_owner_lifecycle.py new file mode 100644 index 000000000..be1ae6f64 --- /dev/null +++ b/tests/adapters/codex/test_mcp_owner_lifecycle.py @@ -0,0 +1,150 @@ +"""Owner installation and invocation lifetime never strand or cancel callers.""" + +from __future__ import annotations + +import asyncio +from contextlib import nullcontext +from typing import Any + +import pytest + +from band.core.exceptions import InputNotAccepted +from band.core.turn_lifecycle import TurnTask +from band.integrations.codex.types import CodexSessionState +from band.runtime.cycle import TurnScope +from band.testing import FakeAgentTools +from tests.adapters.codex.kit import make_platform_message +from tests.adapters.codex.mcpkit import ( + StatusClient, + mcp_adapter, + private_servers, + start, +) +from tests.conftest import eager_tasks, requires_eager_tasks + + +async def invoke(adapter, tools: FakeAgentTools) -> None: + await adapter.on_message( + msg=make_platform_message(room_id="a"), + tools=tools, + history=CodexSessionState(), + participants_msg=None, + contacts_msg=None, + is_session_bootstrap=True, + room_id="a", + ) + + +class TestMcpOwnerLifecycle: + @pytest.mark.parametrize( + "eager", [False, pytest.param(True, marks=requires_eager_tasks)] + ) + async def test_failed_registrar_unwinds_uninstalled_worker( + self, tmp_path, eager: bool + ) -> None: + original = RuntimeError("registrar refused") + workers: list[asyncio.Task[None]] = [] + + class RefusingTools(FakeAgentTools): + def register_turn_task(self, owner: TurnTask) -> None: + workers.append(owner.task) + raise original + + adapter, launches = mcp_adapter(tmp_path, [StatusClient()], private_servers) + await adapter.on_started("Fixture", "Registration custody") + try: + with eager_tasks() if eager else nullcontext(): + with pytest.raises(RuntimeError) as error: + await invoke(adapter, RefusingTools()) + assert error.value is original + await adapter.cleanup_all() + assert workers[0].done() + assert adapter._turn_tasks == {} + assert adapter._turn_release == {} + assert adapter._room_clients == {} + assert launches == [] + finally: + for worker in workers: + if not worker.done(): + worker.cancel() + await asyncio.gather(*workers, return_exceptions=True) + await adapter.cleanup_all() + + @pytest.mark.parametrize( + "eager", [False, pytest.param(True, marks=requires_eager_tasks)] + ) + async def test_terminating_scope_unwinds_cancelled_installation( + self, tmp_path, eager: bool + ) -> None: + scope = TurnScope(control_revision=0) + scope.terminate() + workers: list[asyncio.Task[None]] = [] + + class TerminatingTools(FakeAgentTools): + def register_turn_task(self, owner: TurnTask) -> None: + workers.append(owner.task) + scope.bind(owner) + + adapter, launches = mcp_adapter(tmp_path, [StatusClient()], private_servers) + await adapter.on_started("Fixture", "Terminating registration") + try: + result: BaseException | None = None + with eager_tasks() if eager else nullcontext(): + try: + await invoke(adapter, TerminatingTools()) + except (asyncio.CancelledError, asyncio.InvalidStateError) as error: + result = error + await adapter.cleanup_all() + assert isinstance(result, asyncio.CancelledError) + assert workers[0].done() + assert adapter._turn_tasks == {} + assert adapter._turn_release == {} + assert adapter._room_clients == {} + assert launches == [] + finally: + for worker in workers: + if not worker.done(): + worker.cancel() + await asyncio.gather(*workers, return_exceptions=True) + await adapter.cleanup_all() + + @pytest.mark.looptime + async def test_self_cleanup_temporary_owner_does_not_cancel_continuing_caller( + self, + tmp_path, + ) -> None: + continued = asyncio.Event() + + async def resolve(workspace: str) -> dict[str, Any]: + await adapter.on_cleanup("a") + return dict(await private_servers(workspace)) + + adapter, launches = mcp_adapter( + tmp_path, [StatusClient()], resolve, turn_settle_timeout_s=0.02 + ) + await adapter.on_started("Fixture", "Temporary owner lifetime") + + async def caller() -> None: + try: + await start(adapter) + except (asyncio.CancelledError, InputNotAccepted): + pass + await asyncio.sleep( + 0.1 + ) # Continue beyond the invocation's cleanup deadline. + continued.set() + + running = asyncio.create_task(caller()) + try: + await asyncio.gather(running, return_exceptions=True) + assert continued.is_set() + assert running.cancelling() == 0 + assert launches == [] + await adapter.cleanup_all() + assert adapter._room_clients == {} + assert "a" not in adapter._workspaces.rooms + finally: + if not running.done(): + running.cancel() + await asyncio.gather(running, return_exceptions=True) + await adapter.cleanup_all() diff --git a/tests/adapters/codex/test_mcp_recovery.py b/tests/adapters/codex/test_mcp_recovery.py new file mode 100644 index 000000000..b58aa6cb6 --- /dev/null +++ b/tests/adapters/codex/test_mcp_recovery.py @@ -0,0 +1,367 @@ +"""Recovery, settlement and deadline regressions at the actual Codex seams.""" + +from __future__ import annotations + +import asyncio +import tomllib +from collections.abc import Mapping +from typing import Any + +import pytest + +from band.core.retention import Residency +from band.integrations.codex import CodexJsonRpcError +from band.integrations.codex.mcp import ( + CodexMcpHttpServer, + CodexMcpLaunch, + CodexMcpServer, + CodexMcpStartupError, +) +from band.integrations.codex.types import CodexRequestMethod, CodexSessionState +from band.testing import FakeAgentTools +from tests.adapters.codex.kit import make_platform_message +from tests.adapters.codex.mcpkit import ( + StatusClient, + mcp_adapter, + private_servers, + ready_row, + start, +) +from tests.adapters.codexturns import tool_call_request, turn_completed + + +class TestMcpRecovery: + @pytest.mark.parametrize( + "method", [CodexRequestMethod.THREAD_START, CodexRequestMethod.THREAD_RESUME] + ) + async def test_thread_failure_next_turn_sends_rotated_grant( + self, tmp_path, method + ) -> None: + grant = {"value": "grant-A"} + receipts: list[tuple[str, int]] = [] + + async def endpoint( + reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + try: + request = await reader.readuntil(b"\r\n\r\n") + headers = dict( + line.split(b": ", 1) + for line in request.split(b"\r\n")[1:] + if b": " in line + ) + supplied = ( + headers.get(b"Authorization", b"").decode().removeprefix("Bearer ") + ) + status = 200 if supplied == grant["value"] else 401 + receipts.append((supplied, status)) + writer.write( + f"HTTP/1.1 {status} Fixture\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".encode() + ) + await writer.drain() + finally: + writer.close() + await writer.wait_closed() + + server = await asyncio.start_server(endpoint, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + resolved: list[str] = [] + clients: list[StatusClient] = [] + + async def resolve(_workspace: str) -> Mapping[str, CodexMcpServer]: + resolved.append(grant["value"]) + return { + "private": CodexMcpHttpServer( + f"http://127.0.0.1:{port}/mcp", bearer_token=grant["value"] + ) + } + + class EndpointClient(StatusClient): + def __init__(self, token: str) -> None: + super().__init__() + self.token = token + self._events.extend( + ( + tool_call_request( + 1, + "band_send_message", + {"content": "recovered", "mentions": ["user"]}, + ), + turn_completed(), + ) + ) + + async def authorized(self) -> bool: + reader, writer = await asyncio.open_connection("127.0.0.1", port) + try: + writer.write( + f"GET /mcp HTTP/1.1\r\nHost: localhost\r\nAuthorization: Bearer {self.token}\r\n\r\n".encode() + ) + await writer.drain() + status = await reader.readline() + await reader.read() + return status.split()[1] == b"200" + finally: + writer.close() + await writer.wait_closed() + + async def request( + self, sent: str, params: dict[str, Any] | None = None, **kwargs: Any + ) -> dict[str, Any]: + if sent == "mcpServerStatus/list": + self.requests.append((sent, params or {})) + row = ready_row() + if not await self.authorized(): + row.update( + serverInfo=None, + serverCapabilities=None, + runtimeStatus="authenticationRequired", + ) + return {"data": [row]} + if sent == method: + grant["value"] = "grant-B" + if not await self.authorized(): + self.requests.append((sent, params or {})) + raise CodexJsonRpcError( + code=-32603, message="Required MCP connection failed" + ) + return await super().request(sent, params, **kwargs) + + adapter, _ = mcp_adapter(tmp_path, [], resolve, model="fixture") + + def build(_room: object, launch: CodexMcpLaunch | None) -> EndpointClient: + assert launch is not None + overrides = [ + launch.command[index + 1] + for index, item in enumerate(launch.command) + if item == "-c" + ] + settings = tomllib.loads("\n".join(overrides))["mcp_servers"]["private"] + client = EndpointClient(launch.env[settings["bearer_token_env_var"]]) + clients.append(client) + return client + + adapter._build_client = build + tools = FakeAgentTools() + history = ( + CodexSessionState(thread_id="saved") + if method == CodexRequestMethod.THREAD_RESUME + else CodexSessionState() + ) + arguments = { + "msg": make_platform_message(room_id="a"), + "tools": tools, + "history": history, + "participants_msg": None, + "contacts_msg": None, + "is_session_bootstrap": True, + "room_id": "a", + } + try: + await adapter.on_started("Fixture", "Grant recovery fixture") + with pytest.raises(CodexMcpStartupError): + await adapter.on_message(**arguments) + second_failure = None + try: + await adapter.on_message(**arguments) + except CodexMcpStartupError as error: + second_failure = error + assert resolved == ["grant-A", "grant-B"] + assert len(clients) == 2 + assert clients[0].closed + assert second_failure is None + assert ("grant-A", 401) in receipts + assert receipts[-1] == ("grant-B", 200) + finally: + await adapter.cleanup_all() + server.close() + await server.wait_closed() + + @pytest.mark.parametrize("public_turn", [False, True]) + async def test_cleanup_joins_already_cancelling_resolver( + self, tmp_path, public_turn: bool + ) -> None: + entered, settling, release, settled = (asyncio.Event() for _ in range(4)) + + async def resolve(workspace: str) -> Mapping[str, CodexMcpServer]: + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + settling.set() + await release.wait() + settled.set() + raise + + adapter, launches = mcp_adapter( + tmp_path, [StatusClient()], resolve, turn_settle_timeout_s=2 + ) + await adapter.on_started("Fixture", "Cancellation settlement fixture") + if public_turn: + running = asyncio.create_task( + adapter.on_message( + msg=make_platform_message(room_id="a"), + tools=FakeAgentTools(), + history=CodexSessionState(), + participants_msg=None, + contacts_msg=None, + is_session_bootstrap=True, + room_id="a", + ) + ) + else: + running = asyncio.create_task(start(adapter)) + cleanup = None + try: + await asyncio.wait_for(entered.wait(), timeout=1) + room = adapter._room_clients["a"] + invocation = room.startup_owner.task + assert invocation is not None + running.cancel() + await asyncio.wait_for(settling.wait(), timeout=1) + cleanup = asyncio.create_task(adapter.on_cleanup("a")) + await asyncio.sleep(0) + await asyncio.sleep(0) + assert invocation.cancelling() == 1 + assert not cleanup.done() + assert not settled.is_set() + assert not (room.startup_owner is None) + with pytest.raises(ValueError): + adapter._workspaces.claim("b", room.workspace) + release.set() + await asyncio.gather(running, return_exceptions=True) + await asyncio.wait_for(cleanup, timeout=1) + assert settled.is_set() + assert launches == [] + assert "a" not in adapter._room_clients + assert "a" not in adapter._workspaces.rooms + finally: + release.set() + await asyncio.gather(running, return_exceptions=True) + if cleanup is not None: + await asyncio.gather(cleanup, return_exceptions=True) + await adapter.cleanup_all() + adapter._workspaces.release("b") + + async def test_resume_failure_without_mcp_word_checks_readiness_before_fallback( + self, + tmp_path, + ) -> None: + class LostRequiredServer(StatusClient): + failed = False + + async def request( + self, sent: str, params: dict[str, Any] | None = None, **kwargs: Any + ) -> dict[str, Any]: + if sent == CodexRequestMethod.THREAD_RESUME: + self.failed = True + self.requests.append((sent, params or {})) + raise CodexJsonRpcError( + code=-32603, + message="error resuming thread: connection failed", + data=None, + ) + if sent == "mcpServerStatus/list": + self.requests.append((sent, params or {})) + return { + "data": [ + {**ready_row(), "runtimeStatus": "failed"} + if self.failed + else ready_row() + ] + } + return await super().request(sent, params, **kwargs) + + client = LostRequiredServer() + adapter, _ = mcp_adapter(tmp_path, [client], private_servers) + try: + await start(adapter) + with pytest.raises(CodexMcpStartupError) as error: + await adapter._ensure_thread( + room=adapter._room_clients["a"], + history=CodexSessionState(thread_id="saved"), + tools=FakeAgentTools(), + is_session_bootstrap=True, + ) + assert error.value.stage == "thread/resume" + assert CodexRequestMethod.THREAD_START not in client.request_methods + assert "a" not in adapter._needs_history_injection + assert adapter._room_clients["a"].retained.snapshot.pending_cleanup + finally: + await adapter.cleanup_all() + + @pytest.mark.parametrize( + "method", [CodexRequestMethod.THREAD_START, CodexRequestMethod.THREAD_RESUME] + ) + async def test_unrelated_error_mentioning_mcp_is_not_startup_failure( + self, tmp_path, method + ) -> None: + class ModelFailure(StatusClient): + async def request( + self, sent: str, params: dict[str, Any] | None = None, **kwargs: Any + ) -> dict[str, Any]: + if sent == method: + self.requests.append((sent, params or {})) + raise CodexJsonRpcError( + code=-32603, + message="MCP was mentioned in an unrelated model validation error", + ) + if sent == "mcpServerStatus/list": + self.requests.append((sent, params or {})) + return {"data": [ready_row()]} + return await super().request(sent, params, **kwargs) + + client = ModelFailure() + adapter, _ = mcp_adapter(tmp_path, [client], private_servers) + try: + await start(adapter) + kwargs = { + "room": adapter._room_clients["a"], + "history": CodexSessionState(thread_id="saved") + if method == CodexRequestMethod.THREAD_RESUME + else CodexSessionState(), + "tools": FakeAgentTools(), + "is_session_bootstrap": True, + } + if method == CodexRequestMethod.THREAD_RESUME: + assert await adapter._ensure_thread(**kwargs) == "thr-1" + assert "a" in adapter._needs_history_injection + else: + with pytest.raises(CodexJsonRpcError): + await adapter._ensure_thread(**kwargs) + assert not adapter._room_clients["a"].retained.snapshot.pending_cleanup + finally: + await adapter.cleanup_all() + + @pytest.mark.looptime + async def test_consumed_readiness_timeout_cannot_publish_initialized( + self, tmp_path + ) -> None: + class ConsumingClient(StatusClient): + async def request( + self, sent: str, params: dict[str, Any] | None = None, **kwargs: Any + ) -> dict[str, Any]: + if sent == "mcpServerStatus/list": + self.requests.append((sent, params or {})) + try: + await asyncio.sleep(31) + except asyncio.CancelledError: + task = asyncio.current_task() + assert task is not None + task.uncancel() + return {"data": [ready_row()]} + return await super().request(sent, params, **kwargs) + + client = ConsumingClient() + adapter, _ = mcp_adapter(tmp_path, [client], private_servers) + try: + with pytest.raises(CodexMcpStartupError) as error: + await start(adapter) + assert error.value.stage == "readiness" + assert ( + adapter._room_clients["a"].retained.snapshot.state != Residency.RESIDENT + ) + assert client.closed + assert adapter._room_threads == {} + finally: + await adapter.cleanup_all() diff --git a/tests/adapters/codex/test_mcp_startup.py b/tests/adapters/codex/test_mcp_startup.py new file mode 100644 index 000000000..ec5828da2 --- /dev/null +++ b/tests/adapters/codex/test_mcp_startup.py @@ -0,0 +1,649 @@ +"""Runtime MCP startup ownership and positive readiness regressions.""" + +from __future__ import annotations + +import asyncio +import os +from collections.abc import Mapping +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + +from band.adapters.codex import CodexAdapter +from band.core.exceptions import InputNotAccepted +from band.core.retention import Residency +from band.integrations.codex import CodexJsonRpcError +from band.integrations.codex import mcp as codex_mcp +from band.integrations.codex.mcp import ( + CodexMcpHttpServer, + CodexMcpServer, + CodexMcpStartupError, + CodexMcpStdioServer, +) +from band.integrations.codex.types import CodexRequestMethod, CodexSessionState +from band.testing import FakeAgentTools, reported_failures +from tests.adapters.codex.kit import make_platform_message +from tests.adapters.codex.mcpkit import ( + DuplicateServers, + SseServer, + StatusClient, + mcp_adapter, + private_servers, + ready_row, + start, +) +from tests.adapters.codexturns import FakeCodexClient +from tests.conftest import eager_tasks, requires_eager_tasks + + +class TestMcpStartup: + async def test_reuse_skips_resolution_and_retirement_resolves_fresh( + self, + tmp_path: Path, + ) -> None: + calls = [] + + async def resolve(workspace: str) -> Mapping[str, CodexMcpServer]: + calls.append(workspace) + return await private_servers(workspace) + + first, second = StatusClient(), StatusClient() + adapter, launches = mcp_adapter(tmp_path, [first, second], resolve) + await start(adapter) + original = adapter._room_clients["a"].client + await start(adapter) + assert adapter._room_clients["a"].client is original + assert len(calls) == 1 + assert len(launches) == 1 + room = adapter._room_clients["a"] + room.retained.recover(generation=room.generation, reason="test recovery") + await start(adapter) + assert first.closed + assert adapter._room_clients["a"].client is second + assert calls == [str(tmp_path / "a"), str(tmp_path / "a")] + await adapter.cleanup_all() + + async def test_held_resolver_does_not_block_another_room( + self, tmp_path: Path + ) -> None: + entered, release = asyncio.Event(), asyncio.Event() + + async def resolve(workspace: str) -> Mapping[str, CodexMcpServer]: + if Path(workspace).name == "a": + entered.set() + await release.wait() + return await private_servers(workspace) + + adapter, launches = mcp_adapter( + tmp_path, [StatusClient(), StatusClient()], resolve + ) + waiting = asyncio.create_task(start(adapter, "a")) + try: + await asyncio.wait_for(entered.wait(), timeout=1) + await asyncio.wait_for(start(adapter, "b"), timeout=1) + assert ( + adapter._room_clients["b"].retained.snapshot.state == Residency.RESIDENT + ) + assert ( + adapter._room_clients["a"].retained.snapshot.state != Residency.RESIDENT + ) + assert adapter._room_clients["a"].client is None + assert not waiting.done() + assert len(launches) == 1 + release.set() + await asyncio.wait_for(waiting, timeout=1) + assert ( + adapter._room_clients["a"].retained.snapshot.state == Residency.RESIDENT + ) + assert len(launches) == 2 + finally: + release.set() + await waiting + await adapter.cleanup_all() + + @pytest.mark.parametrize( + "stage", ["resolution", "connect", "initialize", "skills", "model", "readiness"] + ) + @pytest.mark.parametrize("swallow", [False, True]) + async def test_cancelled_startup_cannot_publish( + self, tmp_path: Path, stage: str, swallow: bool + ) -> None: + entered, release = asyncio.Event(), asyncio.Event() + + async def checkpoint(current: str) -> None: + if current == stage: + entered.set() + try: + await release.wait() + except asyncio.CancelledError: + if not swallow: + raise + + async def resolve(workspace: str) -> Mapping[str, CodexMcpServer]: + await checkpoint("resolution") + return await private_servers(workspace) + + client = StatusClient(checkpoint=checkpoint) + adapter, launches = mcp_adapter( + tmp_path, [client], resolve, skill_roots=[str(tmp_path)] + ) + running = asyncio.create_task(start(adapter)) + try: + await entered.wait() + running.cancel() + with pytest.raises(asyncio.CancelledError): + await running + assert ( + adapter._room_clients["a"].retained.snapshot.state != Residency.RESIDENT + ) + assert adapter._room_clients["a"].client is None + assert adapter._room_threads == {} + if stage == "resolution": + assert launches == [] + else: + assert client.closed + finally: + release.set() + await asyncio.gather(running, return_exceptions=True) + await adapter.cleanup_all() + + async def test_cancelled_at_admission_never_resolves(self, tmp_path: Path) -> None: + calls = [] + + async def resolve(workspace: str) -> Mapping[str, CodexMcpServer]: + calls.append(workspace) + return {} + + adapter, launches = mcp_adapter(tmp_path, [StatusClient()], resolve) + + async def cancelled() -> None: + task = asyncio.current_task() + assert task is not None + task.cancel() + await start(adapter) + + with pytest.raises(asyncio.CancelledError): + await asyncio.create_task(cancelled()) + assert calls == [] + assert launches == [] + await adapter.cleanup_all() + + async def test_resolver_self_cleanup_unwinds_without_self_wait( + self, tmp_path: Path + ) -> None: + adapter: CodexAdapter + + async def resolve(workspace: str) -> Mapping[str, CodexMcpServer]: + await adapter.on_cleanup("a") + return await private_servers(workspace) + + adapter, launches = mcp_adapter(tmp_path, [StatusClient()], resolve) + + with pytest.raises(InputNotAccepted): + await asyncio.wait_for(start(adapter), timeout=1) + await asyncio.wait_for(adapter.on_cleanup("a"), timeout=1) + assert launches == [] + assert "a" not in adapter._room_clients + + async def test_failed_close_retains_secret_scope_claim_and_primary_error( + self, tmp_path: Path, caplog: pytest.LogCaptureFixture + ) -> None: + class FailedClient(StatusClient): + refuse_close = True + + async def initialize(self, **kwargs: Any) -> dict[str, Any]: + raise RuntimeError("primary secret-sentinel") + + async def close(self) -> None: + if self.refuse_close: + raise RuntimeError("cleanup secret-sentinel") + await super().close() + + first = FailedClient() + adapter, _ = mcp_adapter(tmp_path, [first, StatusClient()], private_servers) + with pytest.raises(CodexMcpStartupError) as error: + await start(adapter) + assert error.value.stage == "initialize" + assert "secret-sentinel" not in str(error.value) + assert "secret-sentinel" not in caplog.text + room = adapter._room_clients["a"] + assert room.client is first + assert room.mcp_snapshot is not None + assert room.retained.snapshot.pending_cleanup + assert adapter._workspaces.claim("a", room.workspace) == room.workspace + first.refuse_close = False + await start(adapter) + assert first.closed + assert room.retained.snapshot.state == Residency.RESIDENT + assert room.client is not first + await adapter.cleanup_all() + + async def test_status_pagination_requires_all_expected_names( + self, tmp_path: Path + ) -> None: + async def resolve(_workspace: str) -> Mapping[str, CodexMcpServer]: + return { + "one": CodexMcpHttpServer("http://localhost/one"), + "two": CodexMcpHttpServer("http://localhost/two"), + } + + client = StatusClient( + [ + {"data": [ready_row("one")], "nextCursor": "next"}, + {"data": [ready_row("two")], "nextCursor": None}, + ] + ) + adapter, _ = mcp_adapter(tmp_path, [client], resolve) + await start(adapter) + assert adapter._room_clients["a"].retained.snapshot.state == Residency.RESIDENT + assert [ + params + for method, params in client.requests + if method == "mcpServerStatus/list" + ] == [ + {"cursor": None, "limit": 100, "detail": "full"}, + {"cursor": "next", "limit": 100, "detail": "full"}, + ] + await adapter.cleanup_all() + + @pytest.mark.parametrize( + "page", + [ + {}, + {"data": None}, + {"data": {}}, + {"data": [None]}, + {"data": [{"name": 7}]}, + {"data": []}, + {"data": [ready_row(), ready_row()]}, + {"data": [ready_row()], "nextCursor": 1}, + {"data": [ready_row()], "nextCursor": ""}, + {"data": [{"name": "private", "authStatus": "bearerToken", "tools": {}}]}, + {"data": [{**ready_row(), "serverInfo": None}]}, + {"data": [{**ready_row(), "serverInfo": {}}]}, + {"data": [{**ready_row(), "serverCapabilities": None}]}, + {"data": [{**ready_row(), "toolsError": "secret-sentinel"}]}, + ], + ) + async def test_bad_catalog_never_marks_room_ready( + self, tmp_path: Path, page: dict[str, Any] + ) -> None: + client = StatusClient([page]) + adapter, _ = mcp_adapter(tmp_path, [client], private_servers) + with pytest.raises(CodexMcpStartupError) as error: + await start(adapter) + assert error.value.stage == "readiness" + assert "secret-sentinel" not in str(error.value) + assert adapter._room_clients["a"].retained.snapshot.state != Residency.RESIDENT + assert adapter._room_threads == {} + assert client.closed + await adapter.cleanup_all() + + @pytest.mark.parametrize( + "status", + [ + "notStarted", + "starting", + "failed", + "authenticationRequired", + "authRequired", + "disabled", + "cancelled", + "garbage", + False, + ], + ) + async def test_negative_runtime_status_is_not_readiness( + self, tmp_path: Path, status: Any + ) -> None: + client = StatusClient([{"data": [{**ready_row(), "runtimeStatus": status}]}]) + adapter, _ = mcp_adapter(tmp_path, [client], private_servers) + with pytest.raises(CodexMcpStartupError): + await start(adapter) + assert client.closed + assert adapter._room_clients["a"].retained.snapshot.state != Residency.RESIDENT + await adapter.cleanup_all() + + async def test_repeated_cursor_is_refused(self, tmp_path: Path) -> None: + client = StatusClient( + [ + {"data": [ready_row()], "nextCursor": "same"}, + {"data": [], "nextCursor": "same"}, + ] + ) + adapter, _ = mcp_adapter(tmp_path, [client], private_servers) + with pytest.raises(CodexMcpStartupError): + await start(adapter) + assert client.closed + assert ( + len( + [ + method + for method, _ in client.requests + if method == "mcpServerStatus/list" + ] + ) + == 2 + ) + await adapter.cleanup_all() + + @pytest.mark.parametrize( + "method", [CodexRequestMethod.THREAD_START, CodexRequestMethod.THREAD_RESUME] + ) + async def test_required_mcp_failure_never_starts_history_fallback( + self, tmp_path: Path, method: CodexRequestMethod + ) -> None: + class RequiredFailureClient(StatusClient): + failed = False + + async def request( + self, sent: str, params: dict[str, Any] | None = None, **kwargs: Any + ) -> dict[str, Any]: + if sent == method: + self.failed = True + self.requests.append((sent, params or {})) + raise CodexJsonRpcError( + code=-32002, + message="Required MCP startup failed; thread not found secret-sentinel", + data={"token": "secret-sentinel"}, + ) + if sent == "mcpServerStatus/list" and self.failed: + self.requests.append((sent, params or {})) + return {"data": [{**ready_row(), "runtimeStatus": "failed"}]} + return await super().request(sent, params, **kwargs) + + client = RequiredFailureClient() + adapter, _ = mcp_adapter(tmp_path, [client], private_servers) + await start(adapter) + history = ( + CodexSessionState(thread_id="saved") + if method == CodexRequestMethod.THREAD_RESUME + else CodexSessionState() + ) + with pytest.raises(CodexMcpStartupError) as error: + await adapter._ensure_thread( + room=adapter._room_clients["a"], + history=history, + tools=FakeAgentTools(), + is_session_bootstrap=True, + ) + assert error.value.stage == method + assert "secret-sentinel" not in str(error.value) + assert "a" not in adapter._needs_history_injection + assert adapter._room_threads == {} + if method == CodexRequestMethod.THREAD_RESUME: + assert CodexRequestMethod.THREAD_START not in client.request_methods + await adapter.cleanup_all() + + async def test_callback_exception_metadata_cannot_forge_startup_stage( + self, + tmp_path: Path, + ) -> None: + async def resolve(_workspace: str) -> Mapping[str, CodexMcpServer]: + raise CodexMcpStartupError("connect") + + adapter, launches = mcp_adapter(tmp_path, [StatusClient()], resolve) + with pytest.raises(CodexMcpStartupError) as error: + await start(adapter) + assert error.value.stage == "resolution" + assert launches == [] + await adapter.cleanup_all() + + async def test_empty_snapshot_has_no_mcp_readiness_requirement( + self, tmp_path: Path + ) -> None: + async def resolve(_workspace: str) -> Mapping[str, CodexMcpServer]: + return {} + + client = FakeCodexClient() + adapter, launches = mcp_adapter(tmp_path, [client], resolve) + await start(adapter) + assert adapter._room_clients["a"].retained.snapshot.state == Residency.RESIDENT + assert "mcpServerStatus/list" not in client.request_methods + assert launches[0] is not None + assert "-c" not in launches[0].command + await adapter.cleanup_all() + + async def test_snapshot_is_stable_while_connect_suspends( + self, tmp_path: Path + ) -> None: + entered, release = asyncio.Event(), asyncio.Event() + servers: dict[str, CodexMcpServer] = { + "private": CodexMcpHttpServer( + "http://localhost/original", bearer_token="original-sentinel" + ) + } + + async def checkpoint(stage: str) -> None: + if stage == "connect": + entered.set() + await release.wait() + + async def resolve(_workspace: str) -> Mapping[str, CodexMcpServer]: + return servers + + adapter, launches = mcp_adapter( + tmp_path, [StatusClient(checkpoint=checkpoint)], resolve + ) + running = asyncio.create_task(start(adapter)) + try: + await entered.wait() + servers.clear() + servers["replacement"] = CodexMcpHttpServer( + "http://localhost/replacement", bearer_token="replacement-sentinel" + ) + release.set() + await running + room = adapter._room_clients["a"] + assert room.mcp_snapshot is not None + assert room.mcp_snapshot.expected_names == frozenset({"private"}) + launch = launches[0] + assert launch is not None and launch.env is not None + assert "original-sentinel" in launch.env.values() + assert "replacement-sentinel" not in launch.env.values() + assert room.retained.snapshot.state == Residency.RESIDENT + finally: + release.set() + await asyncio.gather(running, return_exceptions=True) + await adapter.cleanup_all() + + @pytest.mark.looptime + async def test_readiness_pages_share_one_total_deadline( + self, tmp_path: Path + ) -> None: + class SlowPages(StatusClient): + async def request( + self, method: str, params: dict[str, Any] | None = None, **kwargs: Any + ) -> dict[str, Any]: + if method == "mcpServerStatus/list": + self.requests.append((method, params or {})) + await asyncio.sleep(20) + cursor = params.get("cursor") if params is not None else None + return { + "data": [ready_row()] if cursor is None else [], + "nextCursor": "next" if cursor is None else None, + } + return await super().request(method, params, **kwargs) + + client = SlowPages() + adapter, _ = mcp_adapter(tmp_path, [client], private_servers) + with pytest.raises(CodexMcpStartupError) as error: + await start(adapter) + assert error.value.stage == "readiness" + assert adapter._room_clients["a"].retained.snapshot.state != Residency.RESIDENT + assert client.closed + await adapter.cleanup_all() + + async def test_missing_history_still_resumes_via_fallback_with_mcp( + self, + tmp_path: Path, + ) -> None: + class MissingHistory(StatusClient): + async def request( + self, method: str, params: dict[str, Any] | None = None, **kwargs: Any + ) -> dict[str, Any]: + if method == CodexRequestMethod.THREAD_RESUME: + self.requests.append((method, params or {})) + raise CodexJsonRpcError(code=-32002, message="Thread not found") + if method == "mcpServerStatus/list": + self.requests.append((method, params or {})) + return {"data": [ready_row()]} + return await super().request(method, params, **kwargs) + + client = MissingHistory() + adapter, _ = mcp_adapter(tmp_path, [client], private_servers) + await start(adapter) + thread = await adapter._ensure_thread( + room=adapter._room_clients["a"], + history=CodexSessionState(thread_id="missing"), + tools=FakeAgentTools(), + is_session_bootstrap=True, + ) + assert thread == "thr-1" + assert "a" in adapter._needs_history_injection + params = client.params_of(CodexRequestMethod.THREAD_START)[0] + assert "dynamicTools" in params + assert params["cwd"] == str(tmp_path / "a") + await adapter.cleanup_all() + + @pytest.mark.parametrize( + ("servers", "stage"), + [ + ({"private": SseServer()}, "unsupported-transport"), + ( + {"invalid.name": CodexMcpHttpServer("http://localhost/mcp")}, + "invalid-name", + ), + (DuplicateServers(), "duplicate-name"), + ( + {"codex_apps": CodexMcpHttpServer("http://localhost/mcp")}, + "reserved-name", + ), + ], + ) + async def test_invalid_resolver_result_refuses_before_construction( + self, tmp_path: Path, servers: Mapping[str, Any], stage: str + ) -> None: + async def resolve(_workspace: str) -> Mapping[str, CodexMcpServer]: + return servers + + adapter, launches = mcp_adapter(tmp_path, [StatusClient()], resolve) + with pytest.raises(CodexMcpStartupError) as error: + await start(adapter) + assert error.value.stage == stage + assert launches == [] + assert adapter._room_clients["a"].client is None + await adapter.cleanup_all() + + async def test_windows_secret_stdio_refuses_before_client_construction( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + codex_mcp, "os", SimpleNamespace(name="nt", environ=os.environ) + ) + secret = "private-stdio-platform-grant" + + async def resolve(_workspace: str) -> Mapping[str, CodexMcpServer]: + return {"private": CodexMcpStdioServer("runner", env={"TOKEN": secret})} + + adapter, launches = mcp_adapter(tmp_path, [StatusClient()], resolve) + try: + with pytest.raises(CodexMcpStartupError) as error: + await start(adapter) + assert error.value.stage == "unsupported-platform" + assert secret not in str(error.value) + assert launches == [] + assert adapter._room_clients["a"].client is None + finally: + await adapter.cleanup_all() + + @requires_eager_tasks + async def test_cancelled_admission_refuses_before_eager_worker( + self, tmp_path: Path + ) -> None: + resolved = [] + + async def resolve(workspace: str) -> Mapping[str, CodexMcpServer]: + resolved.append(workspace) + return {} + + adapter, launches = mcp_adapter(tmp_path, [], resolve) + + async def cancelled_caller() -> None: + task = asyncio.current_task() + assert task is not None + task.cancel() + try: + await asyncio.sleep(0) + except asyncio.CancelledError: + pass + with pytest.raises(asyncio.CancelledError): + await adapter.on_message( + msg=make_platform_message(room_id="a"), + tools=FakeAgentTools(), + history=CodexSessionState(), + participants_msg=None, + contacts_msg=None, + is_session_bootstrap=True, + room_id="a", + ) + + with eager_tasks(): + caller = asyncio.create_task(cancelled_caller()) + await caller + assert resolved == [] + assert launches == [] + assert adapter._room_clients == {} + assert adapter._turn_tasks == {} + + async def test_callback_cancellation_cannot_leak_its_reason( + self, tmp_path: Path + ) -> None: + async def resolve(_workspace: str) -> Mapping[str, CodexMcpServer]: + raise asyncio.CancelledError("private-callback-sentinel") + + adapter, launches = mcp_adapter(tmp_path, [], resolve) + with pytest.raises(asyncio.CancelledError) as error: + await start(adapter) + assert error.value.args == () + assert launches == [] + await adapter.cleanup_all() + + async def test_rpc_failure_forwarded_to_band_is_redacted( + self, tmp_path: Path + ) -> None: + class FailingTurnClient(StatusClient): + async def request( + self, method: str, params: dict[str, Any] | None = None, **kwargs: Any + ) -> dict[str, Any]: + if method == CodexRequestMethod.TURN_START: + raise CodexJsonRpcError( + code=-32001, + message="Provider rejected secret-sentinel", + data={"nested": ["secret-sentinel"]}, + ) + return await super().request(method, params, **kwargs) + + adapter, _ = mcp_adapter(tmp_path, [FailingTurnClient()], private_servers) + tools = FakeAgentTools() + await adapter.on_started("Fixture", "RPC diagnostic fixture") + with pytest.raises(CodexJsonRpcError) as error: + await adapter.on_message( + msg=make_platform_message(room_id="a"), + tools=tools, + history=CodexSessionState(), + participants_msg=None, + contacts_msg=None, + is_session_bootstrap=True, + room_id="a", + ) + assert "secret-sentinel" not in str(error.value) + assert error.value.data == {"nested": ["[REDACTED]"]} + failures = reported_failures(tools) + assert len(failures) == 1 + assert failures[0]["provider"] == "codex" + assert "Provider rejected [REDACTED]" in failures[0]["message"] + assert "secret-sentinel" not in str(failures) + await adapter.cleanup_all() diff --git a/tests/adapters/codexturns.py b/tests/adapters/codexturns.py index 46daceefa..bce92e291 100644 --- a/tests/adapters/codexturns.py +++ b/tests/adapters/codexturns.py @@ -267,7 +267,7 @@ def patch_codex_client( to the next process started after a loss; the last one keeps serving.""" processes = deque([client, *replacements]) - def _build(_workspace: str) -> FakeCodexClient: + def _build(_room: object, _launch: object) -> FakeCodexClient: return processes.popleft() if len(processes) > 1 else processes[0] adapter._build_client = _build # type: ignore[method-assign] diff --git a/tests/adapters/opencode/helpers.py b/tests/adapters/opencode/helpers.py index f6cbdab4a..0fe1ecc21 100644 --- a/tests/adapters/opencode/helpers.py +++ b/tests/adapters/opencode/helpers.py @@ -443,12 +443,12 @@ async def time_passes(seconds: float) -> None: async def wait_for(predicate: Callable[[], bool], timeout_s: float = 1.0) -> None: + """Observe readiness before timing out when a delayed poll resumes.""" deadline = asyncio.get_running_loop().time() + timeout_s - while asyncio.get_running_loop().time() < deadline: - if predicate(): - return + while not predicate(): + if asyncio.get_running_loop().time() >= deadline: + pytest.fail("Timed out waiting for condition") await asyncio.sleep(0.01) - pytest.fail("Timed out waiting for condition") async def run_single_turn( diff --git a/tests/adapters/opencode/test_waiting.py b/tests/adapters/opencode/test_waiting.py new file mode 100644 index 000000000..24106528f --- /dev/null +++ b/tests/adapters/opencode/test_waiting.py @@ -0,0 +1,44 @@ +"""OpenCode synchronization observes readiness after a delayed poll wakes.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from tests.adapters.opencode.helpers import wait_for + + +class TestWaitFor: + @pytest.mark.parametrize("wake_at", [1.0, 1.5]) + @pytest.mark.parametrize("completed", [False, True]) + async def test_condition_is_observed_before_expired_deadline( + self, wake_at: float, completed: bool + ) -> None: + now = 0.0 + ready = False + loop = SimpleNamespace(time=lambda: now) + + async def resume_poll(_interval: float) -> None: + nonlocal now, ready + assert now == 0.0, "A late poll must not start another wait" + now = wake_at + ready = completed + + with ( + patch( + "tests.adapters.opencode.helpers.asyncio.get_running_loop", + return_value=loop, + ), + patch( + "tests.adapters.opencode.helpers.asyncio.sleep", side_effect=resume_poll + ), + ): + if completed: + await wait_for(lambda: ready, timeout_s=1.0) + else: + with pytest.raises(pytest.fail.Exception): + await wait_for(lambda: ready, timeout_s=1.0) + assert now == wake_at + assert ready is completed diff --git a/tests/core/test_turn_lifecycle.py b/tests/core/test_turn_lifecycle.py index 50d274235..30cfa0b4e 100644 --- a/tests/core/test_turn_lifecycle.py +++ b/tests/core/test_turn_lifecycle.py @@ -232,3 +232,99 @@ async def swallow_first_cancel() -> None: owner.retry(deadline=terminal_deadline()) assert await owner.drain() assert cancels == 2 + + @pytest.mark.looptime + async def test_new_owner_terminates_task_with_historical_cancellation(self) -> None: + started, handled, release = (asyncio.Event() for _ in range(3)) + cancellations = 0 + + async def work() -> None: + nonlocal cancellations + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancellations += 1 + handled.set() + try: + await release.wait() + except asyncio.CancelledError: + cancellations += 1 + raise + + task = asyncio.create_task(work()) + owner = None + try: + await started.wait() + task.cancel() + await handled.wait() + assert task.cancelling() == 1 + owner = TurnTask(task) + owner.terminate() + await owner.drain() + assert cancellations == 2 + assert task.cancelled() + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + + @pytest.mark.looptime + async def test_closed_owner_cannot_deliver_pending_or_new_expiry(self) -> None: + release = asyncio.Event() + task = asyncio.create_task(release.wait()) + owner = TurnTask(task) + try: + owner.terminate( + deadline=asyncio.get_running_loop().time() + 0.02, cooperative=True + ) + owner.close() + owner.close() + await asyncio.sleep(0.04) # Cross the scheduled cleanup expiry. + assert task.cancelling() == 0 + assert not task.done() + owner.terminate() + owner.retry(deadline=asyncio.get_running_loop().time() + 0.02) + await asyncio.sleep(0.04) + assert task.cancelling() == 0 + release.set() + await task + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + + @pytest.mark.looptime + async def test_external_cancel_during_cooperative_grace_is_adopted(self) -> None: + entered, settling, release = (asyncio.Event() for _ in range(3)) + cancellations = 0 + + async def work() -> None: + nonlocal cancellations + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancellations += 1 + settling.set() + try: + await release.wait() + except asyncio.CancelledError: + cancellations += 1 + raise + + task = asyncio.create_task(work()) + owner = TurnTask(task) + try: + await entered.wait() + owner.terminate( + deadline=asyncio.get_running_loop().time() + 0.02, cooperative=True + ) + task.cancel() + await settling.wait() + assert not await owner.drain() + assert cancellations == 1 + assert not task.done() + release.set() + await task + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) diff --git a/tests/integrations/codex/test_cleanup.py b/tests/integrations/codex/test_cleanup.py index b382d7644..5164e25d7 100644 --- a/tests/integrations/codex/test_cleanup.py +++ b/tests/integrations/codex/test_cleanup.py @@ -57,11 +57,8 @@ time.sleep({LEAKED_LIFETIME_S}) """ ) -# Written only if the launcher runs past stdin EOF. A tree kill that lands -# while it is blocked in that read never creates the file. # The worker does not inherit stdin: on Windows a child that inherits the # pipe its parent is blocked reading does not run until that read finishes. -STDIN_EOF_MARKER = "stdin-eof" EXITING_SHIM = textwrap.dedent( f""" import subprocess, sys @@ -70,7 +67,6 @@ stdin=subprocess.DEVNULL, ) sys.stdin.read() - open({STDIN_EOF_MARKER!r}, "w").close() """ ) # Starts the worker and exits. The worker keeps the workdir after this process is gone. @@ -188,6 +184,24 @@ async def test_overlapping_close_waits_for_the_same_child_to_exit( assert process.returncode is not None await first + async def test_forced_close_waits_for_an_exit_reported_late( + self, + child: tuple[CodexStdioClient, asyncio.subprocess.Process], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A forced kill always ends the leader, but a loaded event loop (as on + Windows CI) can report its exit seconds later. Close must keep waiting + for that report instead of failing the cleanup it already forced.""" + client, process = child + + async def kill_reported_late(target: asyncio.subprocess.Process) -> None: + asyncio.get_running_loop().call_later(1.5, target.kill) + + monkeypatch.setattr(stdio_client, "kill_process_tree", kill_reported_late) + await asyncio.wait_for(client.close(), timeout=CLOSE_BOUND_S) + + assert process.returncode is not None + async def test_reconnect_reaps_the_previous_child_before_spawning( self, child: tuple[CodexStdioClient, asyncio.subprocess.Process] ) -> None: @@ -237,7 +251,6 @@ async def test_close_reaps_a_worker_when_its_launcher_exits_on_stdin_eof( await asyncio.wait_for(wait_for_heartbeat(heartbeat), timeout=CLOSE_BOUND_S) await asyncio.wait_for(client.close(), timeout=CLOSE_BOUND_S) await assert_heartbeat_stopped(heartbeat) - assert not (workdir / STDIN_EOF_MARKER).exists() shutil.rmtree(workdir) finally: await asyncio.wait_for(client.close(), timeout=CLOSE_BOUND_S) @@ -387,32 +400,24 @@ def test_windows_command_does_not_resolve_a_binary_in_the_room( assert Path(resolved[resolved.index("--") + 1]).is_absolute() @pytest.mark.skipif(os.name == "nt", reason="session-group teardown is Unix") - @pytest.mark.parametrize( - ("pid_alive", "signals_group"), - [(False, True), (True, False)], - ) - async def test_close_signals_a_session_only_while_its_leader_pid_is_free( - self, monkeypatch: pytest.MonkeyPatch, pid_alive: bool, signals_group: bool + async def test_reused_leader_pid_does_not_stop_an_unrelated_process( + self, child: tuple[CodexStdioClient, asyncio.subprocess.Process] ) -> None: - process = await asyncio.create_subprocess_exec( + client, live = child + exited = await asyncio.create_subprocess_exec( sys.executable, "-c", "raise SystemExit(0)", start_new_session=True, ) - await process.wait() - signaled: list[int] = [] - - def record_killpg(pid: int, _sig: int) -> None: - signaled.append(pid) - - def probe(pid: int, sig: int) -> None: - if sig != 0: - raise AssertionError(sig) - if not pid_alive: - raise ProcessLookupError(pid) - - monkeypatch.setattr(stdio_client.os, "killpg", record_killpg) - monkeypatch.setattr(stdio_client.os, "kill", probe) - await kill_process_tree(process) - assert signaled == ([process.pid] if signals_group else []) + await exited.wait() + # Model PID reuse with a real live group, not a fake signal handler. + exited.pid = live.pid + await kill_process_tree(exited) + await asyncio.wait_for( + client.initialize( + client_name="alive", client_title="alive", client_version="1" + ), + timeout=CLOSE_BOUND_S, + ) + assert live.returncode is None diff --git a/tests/integrations/codex/test_mcp_redaction.py b/tests/integrations/codex/test_mcp_redaction.py new file mode 100644 index 000000000..6c90172b4 --- /dev/null +++ b/tests/integrations/codex/test_mcp_redaction.py @@ -0,0 +1,252 @@ +"""Secret-safe Codex transport diagnostics and forwarded RPC failures.""" + +from __future__ import annotations + +import asyncio +import json +import logging +from types import SimpleNamespace +from typing import Any + +import pytest + +from band.integrations.codex import CodexJsonRpcError, CodexStdioClient, stdio_client +from band.integrations.codex.mcp import CodexMcpHttpServer, freeze_mcp_servers +from band.integrations.codex.rpc_base import BaseJsonRpcClient + +SECRET = 'sentinel-"quoted"-\\escaped-credential' +_SYNTHETIC_PGID = 123_456_789 + + +@pytest.fixture +def exited_process_group(monkeypatch: pytest.MonkeyPatch) -> int: + def group_ended(pgid: int, _signum: int) -> None: + if pgid != _SYNTHETIC_PGID: + raise AssertionError("redaction fixture cannot settle another group") + raise ProcessLookupError("the synthetic owned group has exited") + + monkeypatch.setattr(stdio_client.os, "killpg", group_ended, raising=False) + return _SYNTHETIC_PGID + + +def redactor() -> Any: + return freeze_mcp_servers( + {"private": CodexMcpHttpServer("http://localhost/mcp", bearer_token=SECRET)} + ).redactor + + +class TestMcpRedaction: + async def test_stderr_and_malformed_rpc_log_metadata_without_payloads( + self, + caplog: pytest.LogCaptureFixture, + ) -> None: + caplog.set_level(logging.DEBUG, logger="band.integrations.codex.stdio_client") + caplog.set_level(logging.DEBUG, logger="band.integrations.codex.rpc_base") + client = CodexStdioClient(redactor=redactor()) + reader = asyncio.StreamReader() + reader.feed_data(f"connection failed {SECRET}\n".encode()) + reader.feed_eof() + client._proc = SimpleNamespace(stderr=reader) + await client._read_stderr_loop() + await client._dispatch_rpc_message("malformed " + SECRET) + await client._dispatch_rpc_message(json.dumps([SECRET])) + await client._dispatch_rpc_message(json.dumps({"unexpected": SECRET})) + rendered = caplog.text + assert SECRET not in rendered + assert json.dumps(SECRET)[1:-1] not in rendered + assert "connection failed" not in rendered + assert "malformed" not in rendered + assert "unexpected" not in rendered + assert "type=str" in rendered + assert "type=list" in rendered + assert "type=dict" in rendered + for payload in ( + "malformed " + SECRET, + json.dumps([SECRET]), + json.dumps({"unexpected": SECRET}), + ): + assert f"length={len(payload)}" in rendered + + @pytest.mark.parametrize("request_id", [None, 7]) + @pytest.mark.parametrize( + "method", + ["item/tool/call", "FirstServerGrant2026", {"token": SECRET}, 7], + ) + async def test_full_rpc_queue_does_not_log_method_or_parameter_contents( + self, + caplog: pytest.LogCaptureFixture, + request_id: int | None, + method: Any, + ) -> None: + caplog.set_level(logging.WARNING, logger="band.integrations.codex.rpc_base") + private = freeze_mcp_servers( + { + "private": CodexMcpHttpServer( + "http://localhost/mcp", bearer_token="FirstServerGrant2026" + ) + } + ) + client = CodexStdioClient(redactor=private.redactor) + client._events = asyncio.Queue(maxsize=1) + await client._dispatch_rpc_message(json.dumps({"method": "occupied"})) + await client._dispatch_rpc_message( + json.dumps( + {"id": request_id, "method": method, "params": {"token": SECRET}} + ) + ) + assert [record.levelno for record in caplog.records] == [logging.WARNING] + rendered = caplog.text + assert SECRET not in rendered + assert "FirstServerGrant2026" not in rendered + assert "token" not in rendered + if isinstance(method, str): + assert method not in rendered + + async def test_rpc_error_message_and_nested_data_are_redacted( + self, + ) -> None: + class Peer(BaseJsonRpcClient): + async def _send_json(self, payload: dict[str, Any]) -> None: + await self._dispatch_rpc_message( + json.dumps( + { + "id": payload["id"], + "error": { + "code": -32000, + "message": "server denied " + SECRET, + "data": { + SECRET: ["Bearer " + SECRET, {"cause": SECRET}] + }, + }, + } + ) + ) + + peer = Peer(redactor=redactor()) + with pytest.raises(CodexJsonRpcError) as error: + await peer.request("thread/start", {}) + assert error.value.code == -32000 + assert "server denied" in str(error.value) + assert SECRET not in str(error.value) + assert SECRET not in repr(error.value.data) + assert json.dumps(SECRET)[1:-1] not in json.dumps(error.value.data) + + async def test_success_results_and_tool_arguments_are_not_redacted( + self, + ) -> None: + class Peer(BaseJsonRpcClient): + async def _send_json(self, payload: dict[str, Any]) -> None: + assert payload["params"] == {"arguments": {"token": SECRET}} + await self._dispatch_rpc_message( + json.dumps({"id": payload["id"], "result": {"content": [SECRET]}}) + ) + + peer = Peer(redactor=redactor()) + result = await peer.request("tools/call", {"arguments": {"token": SECRET}}) + assert result == {"content": [SECRET]} + await peer._dispatch_rpc_message( + json.dumps( + { + "id": "server-request", + "method": "item/tool/call", + "params": {"arguments": {"token": SECRET}}, + } + ) + ) + event = await peer.recv_event(timeout_s=1) + assert event.params == {"arguments": {"token": SECRET}} + assert event.raw["params"] == {"arguments": {"token": SECRET}} + + async def test_rpc_error_delivered_before_close_is_safe_when_consumed_after( + self, + exited_process_group: int, + ) -> None: + scope = redactor() + client = CodexStdioClient(redactor=scope) + sent = asyncio.Event() + consume = asyncio.Event() + outgoing: list[dict[str, Any]] = [] + + def write_request(data: bytes) -> None: + outgoing.append(json.loads(data)) + sent.set() + + async def drain() -> None: + await consume.wait() + + async def wait() -> None: + return None + + client._proc = SimpleNamespace( + pid=exited_process_group, + stdin=SimpleNamespace(write=write_request, drain=drain, close=lambda: None), + wait=wait, + ) + request = asyncio.create_task(client.request("thread/start", {})) + try: + await asyncio.wait_for(sent.wait(), timeout=1) + await client._dispatch_rpc_message( + json.dumps( + { + "id": outgoing[0]["id"], + "error": { + "code": -32000, + "message": "server denied " + SECRET, + "data": {SECRET: ["Bearer " + SECRET, {"cause": SECRET}]}, + }, + } + ) + ) + assert not request.done() + await client.close() + assert scope.text(SECRET) == SECRET + assert not request.done() + consume.set() + with pytest.raises(CodexJsonRpcError) as error: + await asyncio.wait_for(request, timeout=1) + assert error.value.code == -32000 + assert error.value.message == "server denied [REDACTED]" + assert error.value.data == { + "[REDACTED]": ["Bearer [REDACTED]", {"cause": "[REDACTED]"}] + } + assert SECRET not in str(error.value) + assert json.dumps(SECRET)[1:-1] not in json.dumps(error.value.data) + finally: + consume.set() + await asyncio.gather(request, return_exceptions=True) + await client.close() + + async def test_secret_scope_is_retained_until_physical_teardown( + self, + exited_process_group: int, + ) -> None: + scope = redactor() + client = CodexStdioClient(env={"GENERATED": SECRET}, redactor=scope) + exited = asyncio.Event() + + class Process: + pid = exited_process_group + stdin = None + + async def wait(self) -> None: + await exited.wait() + + client._proc = Process() + closing = asyncio.create_task(client.close()) + try: + while client._close_task is None: + await asyncio.sleep(0) + assert scope.text(SECRET) != SECRET + assert client.env is not None and client.env["GENERATED"] == SECRET + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + assert scope.text(SECRET) != SECRET + exited.set() + await client.close() + assert client.env is None + assert scope.text(SECRET) == SECRET + finally: + exited.set() + await asyncio.gather(closing, return_exceptions=True) + await client.close() diff --git a/tests/integrations/codex/test_process_group_cleanup.py b/tests/integrations/codex/test_process_group_cleanup.py new file mode 100644 index 000000000..61aebe4c3 --- /dev/null +++ b/tests/integrations/codex/test_process_group_cleanup.py @@ -0,0 +1,510 @@ +"""Real local process groups must settle before Codex workspace custody ends. + +Each scenario runs in an isolated Python worker. On Linux that worker is a +subreaper: an orphan zombie still makes killpg(pgid, 0) succeed, and container +PID 1 is not necessarily a reaper. Only the recorded descendant is waitpid'd; +pytest's process-wide child watcher and subreaper state are never changed. +""" + +from __future__ import annotations + +import asyncio +import ctypes +import errno +import logging +import os +import signal +import sys +import textwrap +from collections.abc import AsyncIterator, Awaitable, Callable +from contextlib import asynccontextmanager, suppress +from pathlib import Path +from unittest.mock import patch + +import pytest +import pytest_asyncio + +from band.adapters.codex import CodexAdapter, CodexAdapterConfig +from band.core.exceptions import CleanupIncompleteError +from band.integrations.codex import CodexStdioClient +from tests.paths import REPO_ROOT + +logger = logging.getLogger(__name__) + +POLL_S = 0.02 +QUIET_S = 0.2 +READINESS_BOUND_S = 5 +CLOSE_BOUND_S = 8 +CLEANUP_BOUND_S = 4 +WORKER_BOUND_S = 20 +FAILSAFE_LIFETIME_S = 45 +SUBREAPER_UNAVAILABLE_EXIT = 77 +PR_SET_CHILD_SUBREAPER = 36 + +LEADER_READY = "leader-ready" +DESCENDANT_READY = "descendant-ready" +DESCENDANT_PID = "descendant-pid" +LEADER_EOF = "leader-eof" +HEARTBEAT = "heartbeat" +TERM_OBSERVED = "term-observed" +FAILSAFE_EXPIRED = "failsafe-expired" +ROOM_ID = "owned-room" +OTHER_ROOM_ID = "other-room" +NORMAL_CLOSE = "normal-close" +IGNORE_TERM = "ignore-term" +UNKNOWN_EMPTY_PROBE = "unknown-empty-probe" +TRANSIENT_EMPTY_PROBE = "transient-empty-probe" +WORKER_MODULE = "tests.integrations.codex.test_process_group_cleanup" + +pytestmark = pytest.mark.skipif( + os.name != "posix" or not hasattr(os, "fork") or not hasattr(os, "killpg"), + reason="requires POSIX sessions, fork, and process-group signals", +) + +LEADER = textwrap.dedent( + f""" + import os + import signal + import sys + import time + from pathlib import Path + + def publish(name, value): + temporary = Path(name + ".tmp") + temporary.write_text(value, encoding="utf-8") + temporary.replace(name) + + def expire(_signal, _frame): + publish({FAILSAFE_EXPIRED!r}, str(os.getpid())) + os._exit(99) + + signal.signal(signal.SIGALRM, expire) + signal.alarm({FAILSAFE_LIFETIME_S}) + descendant = os.fork() + if descendant == 0: + # Inherited pipes would keep process.wait() pending after normal EOF, + # accidentally exercising the leader-timeout path instead. + null = os.open(os.devnull, os.O_RDWR) + for descriptor in (0, 1, 2): + os.dup2(null, descriptor) + if null > 2: + os.close(null) + signal.alarm({FAILSAFE_LIFETIME_S}) + ignore_term = sys.argv[1] == {IGNORE_TERM!r} + + def on_term(_signal, _frame): + publish({TERM_OBSERVED!r}, "ignored" if ignore_term else "accepted") + if not ignore_term: + os._exit(0) + + signal.signal(signal.SIGTERM, on_term) + publish( + {DESCENDANT_READY!r}, + f"{{os.getpid()}} {{os.getpgrp()}} {{os.getsid(0)}}", + ) + while True: + with open({HEARTBEAT!r}, "a", encoding="utf-8") as heartbeat: + heartbeat.write(".") + time.sleep({POLL_S}) + + publish({DESCENDANT_PID!r}, str(descendant)) + publish( + {LEADER_READY!r}, + f"{{os.getpid()}} {{os.getpgrp()}} {{os.getsid(0)}}", + ) + for _line in sys.stdin: + pass + publish({LEADER_EOF!r}, "normal") + """ +) + +ProcessGroup = tuple[ + CodexStdioClient, + asyncio.subprocess.Process, + asyncio.Task[int | None], +] +ScenarioRunner = Callable[[str], Awaitable[None]] +GroupSignal = Callable[[int, int], None] + + +def leader_command(*, ignore_term: bool = False) -> tuple[str, ...]: + mode = IGNORE_TERM if ignore_term else NORMAL_CLOSE + return (sys.executable, "-u", "-c", LEADER, mode) + + +def group_is_empty(pgid: int, killpg: GroupSignal) -> bool: + try: + killpg(pgid, 0) + except ProcessLookupError: + return True + except PermissionError: + return False + return False + + +async def wait_until(observed: Callable[[], bool]) -> None: + async with asyncio.timeout(READINESS_BOUND_S): + while not observed(): + await asyncio.sleep(POLL_S) + + +def poll_descendant_exit(descendant_pid: int) -> int | None: + """Poll only the recorded adopted child without blocking the worker loop.""" + try: + reaped, status = os.waitpid(descendant_pid, os.WNOHANG) + except ChildProcessError: + return None + return status if reaped else None + + +async def reap_descendant( + descendant_pid: int, + pgid: int, + killpg: GroupSignal, +) -> int | None: + if sys.platform != "linux": + while not group_is_empty(pgid, killpg): + await asyncio.sleep(POLL_S) + return None + while True: + status = poll_descendant_exit(descendant_pid) + if status is not None: + return status + await asyncio.sleep(POLL_S) + + +async def clean_owned_group( + client: CodexStdioClient, + process: asyncio.subprocess.Process, + descendant: asyncio.Task[int | None] | None, + killpg: GroupSignal, +) -> None: + # A live recorded descendant pins this original group identity. Once both + # recorded processes are reaped, never signal a potentially recycled PGID. + descendant_settled = ( + descendant is not None + and descendant.done() + and not descendant.cancelled() + and descendant.exception() is None + ) + still_owned = process.returncode is None or not descendant_settled + if still_owned and not group_is_empty(process.pid, killpg): + with suppress(ProcessLookupError): + killpg(process.pid, signal.SIGKILL) + await process.wait() + if descendant is not None: + await descendant + assert group_is_empty(process.pid, killpg), "fixture's original group did not empty" + await client.close() + + +@asynccontextmanager +async def owned_group( + workspace: Path, + *, + ignore_term: bool = False, + client: CodexStdioClient | None = None, +) -> AsyncIterator[ProcessGroup]: + client = client or CodexStdioClient( + command=leader_command(ignore_term=ignore_term), cwd=str(workspace) + ) + native_killpg = os.killpg + descendant: asyncio.Task[int | None] | None = None + await client.connect() + process = client._proc + assert process is not None + try: + await wait_until( + lambda: ( + (workspace / LEADER_READY).exists() + and (workspace / DESCENDANT_READY).exists() + and (workspace / HEARTBEAT).exists() + ) + ) + leader_identity = tuple( + int(value) for value in (workspace / LEADER_READY).read_text().split() + ) + descendant_pid, descendant_group, descendant_session = ( + int(value) for value in (workspace / DESCENDANT_READY).read_text().split() + ) + recorded_pid = int((workspace / DESCENDANT_PID).read_text()) + descendant = asyncio.create_task( + reap_descendant(recorded_pid, process.pid, native_killpg) + ) + assert leader_identity == (process.pid, process.pid, process.pid) + assert descendant_pid == recorded_pid + assert descendant_pid != process.pid + assert (descendant_group, descendant_session) == (process.pid, process.pid) + initial_heartbeat = (workspace / HEARTBEAT).stat().st_size + await wait_until( + lambda: (workspace / HEARTBEAT).stat().st_size > initial_heartbeat + ) + assert process.returncode is None + assert not group_is_empty(process.pid, native_killpg) + yield client, process, descendant + finally: + primary_error = sys.exception() + try: + if descendant is None and (workspace / DESCENDANT_PID).exists(): + recorded_pid = int((workspace / DESCENDANT_PID).read_text()) + descendant = asyncio.create_task( + reap_descendant(recorded_pid, process.pid, native_killpg) + ) + async with asyncio.timeout(CLEANUP_BOUND_S): + await clean_owned_group(client, process, descendant, native_killpg) + except BaseException: + # Teardown must not turn the behavioral regression into a different + # cleanup exception. With no primary failure, cleanup still fails loudly. + if primary_error is None: + raise + logger.exception("Local process-group fixture cleanup failed") + + +async def assert_settled( + workspace: Path, + group: ProcessGroup, + killpg: GroupSignal, +) -> int | None: + _client, process, descendant = group + assert process.returncode == 0, "leader did not exit normally on stdin EOF" + assert (workspace / LEADER_EOF).read_text() == "normal" + assert group_is_empty(process.pid, killpg), ( + "normal leader exit left its originally owned group alive" + ) + status = await asyncio.wait_for(asyncio.shield(descendant), CLEANUP_BOUND_S) + assert (workspace / TERM_OBSERVED).exists(), "owned descendant never received TERM" + assert not (workspace / FAILSAFE_EXPIRED).exists(), ( + "failsafe, not close(), stopped work" + ) + stopped_at = (workspace / HEARTBEAT).stat().st_size + # Quiescence is itself the observation: span multiple scheduled heartbeats. + await asyncio.sleep(QUIET_S) + assert (workspace / HEARTBEAT).stat().st_size == stopped_at + return status + + +async def normal_close_scenario(workspace: Path) -> None: + async with owned_group(workspace) as group: + client, _process, _descendant = group + await asyncio.wait_for(client.close(), CLOSE_BOUND_S) + status = await assert_settled(workspace, group, os.killpg) + assert (workspace / TERM_OBSERVED).read_text() == "accepted" + if status is not None: + assert os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0 + + +async def ignoring_term_scenario(workspace: Path) -> None: + async with owned_group(workspace, ignore_term=True) as group: + client, _process, _descendant = group + await asyncio.wait_for(client.close(), CLOSE_BOUND_S) + status = await assert_settled(workspace, group, os.killpg) + assert (workspace / TERM_OBSERVED).read_text() == "ignored" + if status is not None: + assert os.WIFSIGNALED(status) and os.WTERMSIG(status) == signal.SIGKILL + + +async def transient_empty_probe_scenario(workspace: Path) -> None: + native_killpg = os.killpg + async with owned_group(workspace) as group: + client, process, _descendant = group + uncertain_probes = [] + + def transient_probe(pgid: int, signum: int) -> None: + try: + native_killpg(pgid, signum) + except ProcessLookupError: + if pgid == process.pid and signum == 0 and not uncertain_probes: + uncertain_probes.append(pgid) + raise PermissionError( + errno.EPERM, "transient empty-group observation" + ) + raise + + with patch.object(os, "killpg", transient_probe): + await asyncio.wait_for(client.close(), CLOSE_BOUND_S) + assert uncertain_probes, "the transient unknown observation was not exercised" + await assert_settled(workspace, group, native_killpg) + + +async def unknown_empty_probe_scenario(workspace: Path) -> None: + adapter = CodexAdapter( + config=CodexAdapterConfig( + codex_command=leader_command(), + workspace_for_room=lambda _room_id: str(workspace), + client_close_timeout_s=CLOSE_BOUND_S, + ) + ) + room = adapter._room_client(ROOM_ID) + adapter._active_room.set(ROOM_ID) + client = adapter._build_client(room, None) + room.client = client + + async def acquire(_generation: int, _resume: bool) -> None: + await client.connect() + + room.retained._acquire = acquire + async with room.retained.use(): + pass + native_killpg = os.killpg + uncertain_probes: list[int] = [] + async with owned_group(workspace, client=client) as group: + _client, process, _descendant = group + original_pgid = process.pid + + def uncertain_empty_probe(pgid: int, signum: int) -> None: + try: + native_killpg(pgid, signum) + except ProcessLookupError: + if pgid != original_pgid or signum != 0: + raise + # Only the observation of a genuinely empty owned group fails. + # TERM/KILL delivery and observations of live groups stay real. + uncertain_probes.append(pgid) + raise PermissionError( + errno.EPERM, "group-empty observation unavailable" + ) + + try: + with patch.object(os, "killpg", uncertain_empty_probe): + with pytest.raises(CleanupIncompleteError): + await asyncio.wait_for(adapter.on_cleanup(ROOM_ID), CLOSE_BOUND_S) + assert uncertain_probes, ( + "cleanup never tried to confirm its empty group" + ) + assert adapter._room_clients[ROOM_ID] is room + assert room.client is client + assert adapter._workspaces.workspace(ROOM_ID) == str(workspace) + assert room.retained.closing + with pytest.raises(ValueError, match="assigned .* to both"): + adapter._workspaces.claim(OTHER_ROOM_ID, room.workspace) + await assert_settled(workspace, group, native_killpg) + + await asyncio.wait_for(adapter.on_cleanup(ROOM_ID), CLOSE_BOUND_S) + assert room.client is None + assert ROOM_ID not in adapter._room_clients + assert ROOM_ID not in adapter._workspaces.rooms + assert ( + adapter._workspaces.claim(OTHER_ROOM_ID, room.workspace) + == room.workspace + ) + await assert_settled(workspace, group, native_killpg) + finally: + primary_error = sys.exception() + try: + await asyncio.wait_for(adapter.on_cleanup(ROOM_ID), CLOSE_BOUND_S) + except BaseException: + if primary_error is None: + raise + logger.exception("Local adapter cleanup failed") + finally: + adapter._workspaces.release(OTHER_ROOM_ID) + + +def enable_worker_subreaper() -> None: + if sys.platform != "linux": + return + libc = ctypes.CDLL(None, use_errno=True) + prctl = getattr(libc, "prctl", None) + if prctl is None: + logger.error("Linux test worker has no prctl child-subreaper support") + raise SystemExit(SUBREAPER_UNAVAILABLE_EXIT) + prctl.argtypes = [ + ctypes.c_int, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ctypes.c_ulong, + ] + prctl.restype = ctypes.c_int + if prctl(PR_SET_CHILD_SUBREAPER, 1, 0, 0, 0) != 0: + error = ctypes.get_errno() + logger.error( + "Linux test worker cannot enable child subreaping: %s", os.strerror(error) + ) + raise SystemExit(SUBREAPER_UNAVAILABLE_EXIT) + + +async def run_worker(scenario: str, workspace: Path) -> None: + task = asyncio.current_task() + assert task is not None + loop = asyncio.get_running_loop() + loop.add_signal_handler(signal.SIGTERM, task.cancel) + try: + scenarios = { + NORMAL_CLOSE: normal_close_scenario, + IGNORE_TERM: ignoring_term_scenario, + UNKNOWN_EMPTY_PROBE: unknown_empty_probe_scenario, + TRANSIENT_EMPTY_PROBE: transient_empty_probe_scenario, + } + try: + exercise = scenarios[scenario] + except KeyError as error: + raise ValueError(f"unknown process-group scenario: {scenario}") from error + await exercise(workspace) + finally: + loop.remove_signal_handler(signal.SIGTERM) + + +@pytest_asyncio.fixture(loop_scope="function") +async def scenario_runner(tmp_path: Path) -> AsyncIterator[ScenarioRunner]: + workers: list[asyncio.subprocess.Process] = [] + + async def run(scenario: str) -> None: + process = await asyncio.create_subprocess_exec( + sys.executable, + "-m", + WORKER_MODULE, + scenario, + str(tmp_path), + cwd=str(REPO_ROOT), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + start_new_session=True, + ) + workers.append(process) + stdout, stderr = await asyncio.wait_for(process.communicate(), WORKER_BOUND_S) + diagnostic = (stdout + stderr).decode("utf-8", errors="replace") + if process.returncode == SUBREAPER_UNAVAILABLE_EXIT: + pytest.skip(diagnostic) + assert process.returncode == 0, diagnostic + + try: + yield run + finally: + for process in workers: + if process.returncode is None: + with suppress(ProcessLookupError): + process.terminate() + try: + await asyncio.wait_for(process.wait(), CLEANUP_BOUND_S + 1) + except TimeoutError: + with suppress(ProcessLookupError): + process.kill() + await asyncio.wait_for(process.wait(), CLEANUP_BOUND_S) + + +class TestProcessGroupCleanup: + async def test_normal_close_settles_descendant_after_leader_stdin_eof( + self, scenario_runner: ScenarioRunner + ) -> None: + await scenario_runner(NORMAL_CLOSE) + + async def test_term_ignoring_descendant_is_killed_after_normal_leader_exit( + self, scenario_runner: ScenarioRunner + ) -> None: + await scenario_runner(IGNORE_TERM) + + async def test_unknown_group_empty_probe_retains_real_adapter_client_and_claim_until_retry( + self, scenario_runner: ScenarioRunner + ) -> None: + await scenario_runner(UNKNOWN_EMPTY_PROBE) + + async def test_transient_unknown_empty_probe_is_confirmed_within_close_bound( + self, scenario_runner: ScenarioRunner + ) -> None: + await scenario_runner(TRANSIENT_EMPTY_PROBE) + + +if __name__ == "__main__": + enable_worker_subreaper() + asyncio.run(run_worker(sys.argv[1], Path(sys.argv[2]))) From ac006bbff8757dc1a316067ba89099107bf823cc Mon Sep 17 00:00:00 2001 From: Natan Zalkin Date: Thu, 15 Oct 2026 16:44:28 +0300 Subject: [PATCH 02/13] fix(codex): retain startup cancellation ownership through settlement Port the reviewed cancellation correction from dd8f3af8..942759f4. Delegate acquisition cancellation explicitly to the matching startup owner, preserve first abandonment/cleanup diagnostics, and cover public caller/cleanup orderings after uncancel. Keep Claude-owned retention changes in the stacked Claude PR. --- CHANGELOG.md | 9 + docs/adapters/codex.md | 13 +- src/band/adapters/codex.py | 17 ++ src/band/core/retention/owner.py | 10 +- src/band/core/retention/room.py | 32 +++- src/band/core/turn_lifecycle.py | 16 +- tests/adapters/codex/test_mcp_cancellation.py | 158 ++++++++++++++++++ 7 files changed, 239 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index edac55b1f..150bdfc07 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,15 @@ The format is based on [Conventional Commits](https://www.conventionalcommits.or and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). This changelog is automatically generated by [Release Please](https://github.com/googleapis/release-please). +## [Unreleased] + +### Bug Fixes + +* **codex:** keep readiness cancellation from interrupting startup settlement or + closing its client early when the public message caller is cancelled or room + cleanup joins startup settlement, even when settlement consumes the task's + cancellation count. + ## [6.0.1](https://github.com/band-ai/band-sdk-python/compare/band-sdk-v6.0.0...band-sdk-v6.0.1) (2026-10-20) diff --git a/docs/adapters/codex.md b/docs/adapters/codex.md index 2b0c1f0cc..ff8040648 100644 --- a/docs/adapters/codex.md +++ b/docs/adapters/codex.md @@ -157,10 +157,15 @@ Pass the keyword-only `resolve_mcp_servers` callback to `CodexAdapter`, not names, repeated cursors, negative statuses and consumed timeout cancellation fail startup. - Cancellation and retirement fence each stage and every discovery request. - A resolver may request its own room cleanup without self-waiting. Cleanup - disarms readiness deadlines and joins an already-cancelling acquisition - without a second cancellation. Settlement beyond the reporting bound retains - acquisition, child, workspace and diagnostic scope until cleanup succeeds. + A resolver may request its own room cleanup without self-waiting. Cleanup and + acquisition abandonment disarm readiness deadlines through the startup + cancellation owner, which records and delivers the first cancellation with + its abandonment or cleanup reason. Later cleanup joins that same cancellation + authority even when settlement consumes the task's cancellation count; it + never sends a fallback cancellation into owned settlement. Explicit cleanup + retries retain their separate cancellation semantics. Settlement beyond the + reporting bound retains acquisition, child, workspace and diagnostic scope + until cleanup succeeds. Failed registration unwinds the gated worker before resolution can start. Invocation-only cancellation timers end with their owner, never later work. - Thread start/resume errors with adopted servers trigger bounded readiness diff --git a/src/band/adapters/codex.py b/src/band/adapters/codex.py index 5643b5982..41ecfea40 100644 --- a/src/band/adapters/codex.py +++ b/src/band/adapters/codex.py @@ -1044,10 +1044,25 @@ def _room_client(self, room_id: str) -> RoomCodexClient: acquire=partial(self._acquire, room), release=partial(self._release, room), dispose=partial(self._dispose, room), + before_cancel_acquire=partial(self._prepare_acquire_cancel, room), ) self._room_clients[room_id] = room return room + def _prepare_acquire_cancel( + self, room: RoomCodexClient, task: asyncio.Task[None], reason: str + ) -> bool: + """Keep a live startup's cancellation with its owner, even after delivery.""" + owner = room.startup_owner + if owner is None or owner.task is not task or task.done(): + return False + owner.terminate( + deadline=asyncio.get_running_loop().time() + + self.config.turn_settle_timeout_s, + reason=reason, + ) + return True + def _active_client_state(self) -> RoomCodexClient | None: room_id = self._active_room.get() return self._room_clients.get(room_id) if room_id is not None else None @@ -1966,6 +1981,7 @@ async def on_cleanup(self, room_id: str) -> None: deadline=asyncio.get_running_loop().time() + self.config.turn_settle_timeout_s, cooperative=True, + reason=f"Room {room_id} acquisition was cancelled for cleanup", ) try: await room.retained.dispose() @@ -1992,6 +2008,7 @@ async def cleanup_all(self) -> None: deadline=asyncio.get_running_loop().time() + self.config.turn_settle_timeout_s, cooperative=True, + reason=f"Room {room_id} acquisition was cancelled for cleanup", ) self._clear_pending_approvals_for_room(room_id) await self._retention.stop() diff --git a/src/band/core/retention/owner.py b/src/band/core/retention/owner.py index e5947ca53..920c8322a 100644 --- a/src/band/core/retention/owner.py +++ b/src/band/core/retention/owner.py @@ -50,11 +50,19 @@ def register( acquire: Callable[[int, bool], Awaitable[None]], release: Callable[[ReleaseCause], Awaitable[ReleaseOutcome]], dispose: Callable[[], Awaitable[None]], + before_cancel_acquire: Callable[[asyncio.Task[None], str], bool] | None = None, ) -> RetainedRoom: self.require_admission() if room_id in self._rooms: raise ValueError(f"Room {room_id} already has a residency owner") - room = RetainedRoom(self, room_id, acquire, release, dispose) + room = RetainedRoom( + self, + room_id, + acquire, + release, + dispose, + before_cancel_acquire=before_cancel_acquire, + ) self._rooms[room_id] = room self.reschedule() return room diff --git a/src/band/core/retention/room.py b/src/band/core/retention/room.py index 202bd221c..19a4a8cde 100644 --- a/src/band/core/retention/room.py +++ b/src/band/core/retention/room.py @@ -65,6 +65,10 @@ class RetainedRoom: release must be a qualified, non-cancelling boundary, and every retry must finish partially acquired or closed handles. ``dispose`` runs only after a confirmed ``FINAL`` release and erases continuation and workspace. + + ``before_cancel_acquire`` receives the task and cancellation reason and returns + whether the provider owns its cancellation, even after delivery. Only ``False`` + or no callback permits a fallback; generation and release ordering stay here. """ def __init__( @@ -74,12 +78,15 @@ def __init__( acquire: Callable[[int, bool], Awaitable[None]], release: Callable[[ReleaseCause], Awaitable[ReleaseOutcome]], dispose: Callable[[], Awaitable[None]], + *, + before_cancel_acquire: Callable[[asyncio.Task[None], str], bool] | None = None, ) -> None: self._owner = owner self.room_id = room_id self._acquire = acquire self._release = release self._dispose = dispose + self._before_cancel_acquire = before_cancel_acquire self._state = Residency.SUSPENDED self._generation = 0 self._leases = 0 @@ -412,8 +419,18 @@ def _abandon_acquisition(self, transition: asyncio.Task[None]) -> None: self.room_id, self._generation, ) - if not transition.cancelling(): - transition.cancel(f"Room {self.room_id} acquisition was abandoned") + self._cancel_acquire( + transition, reason=f"Room {self.room_id} acquisition was abandoned" + ) + + def _cancel_acquire(self, transition: asyncio.Task[None], *, reason: str) -> None: + """Delegate provider cancellation, or deliver an unowned fallback once.""" + if self._before_cancel_acquire is not None and self._before_cancel_acquire( + transition, reason + ): + return + if not transition.cancelling(): + transition.cancel(reason) def _request(self, cause: ReleaseCause) -> None: self._cause = stronger_cause(self._cause, cause) @@ -665,12 +682,11 @@ async def _settle_accepted(self) -> None: self._work_cancelled = True for task in self._uncancelled_work(): task.cancel() - if ( - self._state == Residency.ACQUIRING - and self._transition is not None - and not self._transition.cancelling() - ): - self._transition.cancel() + if self._state == Residency.ACQUIRING and self._transition is not None: + self._cancel_acquire( + self._transition, + reason=f"Room {self.room_id} acquisition was cancelled for cleanup", + ) await self._drained.wait() def _uncancelled_work(self) -> set[asyncio.Task[Any]]: diff --git a/src/band/core/turn_lifecycle.py b/src/band/core/turn_lifecycle.py index bcb66a52b..2b0747fb0 100644 --- a/src/band/core/turn_lifecycle.py +++ b/src/band/core/turn_lifecycle.py @@ -68,6 +68,7 @@ def __init__(self, task: asyncio.Task[Any]) -> None: self._cancel_baseline = task.cancelling() self.deadline: float | None = None self._cancel_delivered = False + self._cancel_reason: str | None = None self._grace_expired = False self._timer: asyncio.Timeout | None = None self._cleanup_timer = False @@ -105,11 +106,20 @@ def _schedule_expiry(self, deadline: float) -> None: self._expiry_timer = asyncio.get_running_loop().call_at(deadline, self._expire) def terminate( - self, *, deadline: float | None = None, cooperative: bool = False + self, + *, + deadline: float | None = None, + cooperative: bool = False, + reason: str | None = None, ) -> None: - """Start the cleanup grace once; later requests never renew it.""" + """Start the cleanup grace once; later requests never renew it. + + Keep the first supplied reason for delivery, including cooperative expiry. + """ if self._closed: return + if self._cancel_reason is None: + self._cancel_reason = reason self._adopt_cancellation() if self.deadline is not None: return @@ -167,7 +177,7 @@ def _adopt_cancellation(self) -> None: def _cancel(self) -> None: if not self._closed and not self._cancel_delivered and not self.task.done(): self._cancel_delivered = True - self.task.cancel() + self.task.cancel(self._cancel_reason) async def drain(self) -> bool: """Wait within the original terminal grace without cancelling a waiter.""" diff --git a/tests/adapters/codex/test_mcp_cancellation.py b/tests/adapters/codex/test_mcp_cancellation.py index a89b3f915..004e11f12 100644 --- a/tests/adapters/codex/test_mcp_cancellation.py +++ b/tests/adapters/codex/test_mcp_cancellation.py @@ -202,3 +202,161 @@ async def request( if cleanup is not None: await asyncio.gather(cleanup, return_exceptions=True) await adapter.cleanup_all() + + @pytest.mark.looptime + async def test_public_cancel_disarms_readiness_before_startup_settlement( + self, tmp_path, monkeypatch + ) -> None: + entered, settling, release = (asyncio.Event() for _ in range(3)) + cancellations: list[int] = [] + order: list[str] = [] + monkeypatch.setattr("band.adapters.codex._MCP_READINESS_TIMEOUT_S", 0.1) + + class SettlingClient(StatusClient): + async def request( + self, method: str, params: dict[str, Any] | None = None, **kwargs: Any + ) -> dict[str, Any]: + if method != "mcpServerStatus/list": + return await super().request(method, params, **kwargs) + self.requests.append((method, params or {})) + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + task = asyncio.current_task() + assert task is not None + cancellations.append(task.cancelling()) + order.append("settling") + settling.set() + try: + await release.wait() + except asyncio.CancelledError: + cancellations.append(task.cancelling()) + order.append("cancelled-again") + raise + order.append("settled") + raise + + async def close(self) -> None: + order.append("close") + await super().close() + + client = SettlingClient() + adapter, _ = mcp_adapter( + tmp_path, [client], private_servers, turn_settle_timeout_s=0.02 + ) + await adapter.on_started("Fixture", "Public cancellation custody") + running = message_task(adapter) + try: + await entered.wait() + room = adapter._room_clients["a"] + snapshot = room.mcp_snapshot + running.cancel() + await settling.wait() + await asyncio.sleep(0.2) # Cross the original readiness deadline. + assert cancellations == [1] + assert order == ["settling"] + assert not client.closed + assert room.client is client and room.mcp_snapshot is snapshot + release.set() + await asyncio.gather(running, return_exceptions=True) + await adapter.on_cleanup("a") + assert order == ["settling", "settled", "close"] + assert client.closed + assert "a" not in adapter._workspaces.rooms + finally: + release.set() + await asyncio.gather(running, return_exceptions=True) + await adapter.cleanup_all() + + @pytest.mark.looptime + @pytest.mark.parametrize("blocked_stage", ["resolution", "readiness"]) + @pytest.mark.parametrize( + "cancel_order", + ["caller", "cleanup", "caller-then-cleanup", "cleanup-then-caller"], + ) + async def test_public_cancel_consumed_count_keeps_settlement_owned( + self, tmp_path, monkeypatch, blocked_stage: str, cancel_order: str + ) -> None: + entered, settling, release = (asyncio.Event() for _ in range(3)) + cancellations: list[int] = [] + order: list[str] = [] + monkeypatch.setattr("band.adapters.codex._MCP_READINESS_TIMEOUT_S", 0.1) + + async def checkpoint(stage: str) -> None: + if stage != blocked_stage: + return + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + task = asyncio.current_task() + assert task is not None + cancellations.append(task.cancelling()) + task.uncancel() + order.append("settling") + settling.set() + try: + await release.wait() + except asyncio.CancelledError: + cancellations.append(task.cancelling()) + order.append("cancelled-again") + raise + order.append("settled") + raise + + async def resolve(workspace: str) -> Mapping[str, CodexMcpServer]: + await checkpoint("resolution") + return await private_servers(workspace) + + class ClosingClient(StatusClient): + async def close(self) -> None: + order.append("close") + await super().close() + + client = ClosingClient(checkpoint=checkpoint) + adapter, launches = mcp_adapter( + tmp_path, [client], resolve, turn_settle_timeout_s=0.02 + ) + await adapter.on_started("Fixture", "Consumed cancellation custody") + running = message_task(adapter) + cleanup = None + try: + await entered.wait() + room = adapter._room_clients["a"] + snapshot, startup_client = room.mcp_snapshot, room.client + if cancel_order.startswith("caller"): + running.cancel() + else: + cleanup = asyncio.create_task(adapter.on_cleanup("a")) + await settling.wait() + if cancel_order == "caller-then-cleanup": + cleanup = asyncio.create_task(adapter.on_cleanup("a")) + elif cancel_order == "cleanup-then-caller": + running.cancel() + await asyncio.sleep(0.2) # Cross both grace and readiness expiry. + assert cancellations == [1] + assert order == ["settling"] + assert room.client is startup_client and room.mcp_snapshot is snapshot + assert not client.closed + with pytest.raises(ValueError): + adapter._workspaces.claim("b", room.workspace) + release.set() + await asyncio.gather(running, return_exceptions=True) + if cleanup is not None: + await asyncio.gather(cleanup, return_exceptions=True) + await adapter.on_cleanup("a") + if blocked_stage == "readiness": + assert order == ["settling", "settled", "close"] + assert client.closed + else: + assert order == ["settling", "settled"] + assert launches == [] + assert "a" not in adapter._workspaces.rooms + finally: + release.set() + await asyncio.gather(running, return_exceptions=True) + if cleanup is not None: + await asyncio.gather(cleanup, return_exceptions=True) + await adapter.cleanup_all() + adapter._workspaces.release("b") From ec7edf6b3cc12dba29cb04ef46f751d15eabdd29 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Sun, 18 Oct 2026 14:25:18 +0300 Subject: [PATCH 03/13] fix(codex): retain cleanup ownership across process probe failures --- src/band/integrations/codex/stdio_client.py | 10 +++----- tests/integrations/codex/test_cleanup.py | 25 +++++++++++++++++++ .../integrations/codex/test_mcp_redaction.py | 17 +++++++------ 3 files changed, 39 insertions(+), 13 deletions(-) diff --git a/src/band/integrations/codex/stdio_client.py b/src/band/integrations/codex/stdio_client.py index 61149657c..64f3b2fc6 100644 --- a/src/band/integrations/codex/stdio_client.py +++ b/src/band/integrations/codex/stdio_client.py @@ -72,12 +72,10 @@ def _session_pid_reused(process: asyncio.subprocess.Process) -> bool: os.kill(process.pid, 0) except ProcessLookupError: return False - except OSError: - logger.warning( - "Codex process %s pid is not signalable; leaving its session", - process.pid, - ) - return True + except OSError as error: + raise CleanupIncompleteError( + "Codex owned process identity could not be confirmed" + ) from error logger.warning( "Codex process %s pid was reused; not signaling that session", process.pid, diff --git a/tests/integrations/codex/test_cleanup.py b/tests/integrations/codex/test_cleanup.py index 5164e25d7..b65d5f199 100644 --- a/tests/integrations/codex/test_cleanup.py +++ b/tests/integrations/codex/test_cleanup.py @@ -18,6 +18,7 @@ import pytest import pytest_asyncio +from band.core.exceptions import CleanupIncompleteError from band.integrations.codex import CodexStdioClient, stdio_client from band.integrations.codex.stdio_client import HELPER_PATH, kill_process_tree from band.integrations.codex.winspawn import ( @@ -421,3 +422,27 @@ async def test_reused_leader_pid_does_not_stop_an_unrelated_process( timeout=CLOSE_BOUND_S, ) assert live.returncode is None + + @pytest.mark.skipif(os.name == "nt", reason="session-group teardown is Unix") + async def test_unknown_leader_pid_probe_keeps_cleanup_retryable( + self, + child: tuple[CodexStdioClient, asyncio.subprocess.Process], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + client, process = child + process.kill() + await process.wait() + native_kill = os.kill + + def unknown_pid(pid: int, signum: int) -> None: + if pid == process.pid and signum == 0: + raise PermissionError("owned PID observation unavailable") + native_kill(pid, signum) + + with monkeypatch.context() as patch: + patch.setattr(stdio_client.os, "kill", unknown_pid) + with pytest.raises(CleanupIncompleteError): + await asyncio.wait_for(client.close(), timeout=CLOSE_BOUND_S) + + await asyncio.wait_for(client.close(), timeout=CLOSE_BOUND_S) + assert process.returncode is not None diff --git a/tests/integrations/codex/test_mcp_redaction.py b/tests/integrations/codex/test_mcp_redaction.py index 6c90172b4..1d78c48ae 100644 --- a/tests/integrations/codex/test_mcp_redaction.py +++ b/tests/integrations/codex/test_mcp_redaction.py @@ -174,13 +174,15 @@ def write_request(data: bytes) -> None: async def drain() -> None: await consume.wait() - async def wait() -> None: - return None - client._proc = SimpleNamespace( pid=exited_process_group, - stdin=SimpleNamespace(write=write_request, drain=drain, close=lambda: None), - wait=wait, + returncode=0, + stdin=SimpleNamespace( + write=write_request, + drain=drain, + close=lambda: None, + is_closing=lambda: False, + ), ) request = asyncio.create_task(client.request("thread/start", {})) try: @@ -228,8 +230,9 @@ class Process: pid = exited_process_group stdin = None - async def wait(self) -> None: - await exited.wait() + @property + def returncode(self) -> int | None: + return 0 if exited.is_set() else None client._proc = Process() closing = asyncio.create_task(client.close()) From ad3e843e57a490f6aa8ddeac1665fabd45654cdd Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Mon, 19 Oct 2026 04:10:42 +0300 Subject: [PATCH 04/13] test(codex): cover credential-bearing MCP endpoints --- tests/adapters/codex/test_mcp_launch.py | 129 +++++++++++++++++++++++ tests/adapters/codex/test_mcp_startup.py | 37 +++++++ 2 files changed, 166 insertions(+) diff --git a/tests/adapters/codex/test_mcp_launch.py b/tests/adapters/codex/test_mcp_launch.py index 51b5c7038..0554e3655 100644 --- a/tests/adapters/codex/test_mcp_launch.py +++ b/tests/adapters/codex/test_mcp_launch.py @@ -224,6 +224,83 @@ def test_invalid_descriptor_data_is_rejected_without_input_diagnostics( assert "private-invalid-args" not in repr(caught.value) assert "private-env-value" not in repr(caught.value) + @pytest.mark.parametrize( + "url", + [ + "https://url-only-grant@mcp.example.test/mcp", + "https://user:url-only-grant@mcp.example.test/mcp", + "https://user%3Aurl-only-grant@mcp.example.test/mcp", + "https://mcp.example.test/mcp?api_key=url-only-grant", + "https://mcp.example.test/mcp?opaque=url-only-grant", + "https://mcp.example.test/mcp?%61pi_key=url-only-grant", + "https://mcp.example.test/mcp#url-only-grant", + "https://@mcp.example.test/mcp", + "https://mcp.example.test/mcp?", + "https://mcp.example.test/mcp#", + "https://url-only-grant／mcp.example.test/mcp", + "https://[url-only-grant/mcp", + ], + ids=[ + "userinfo-token", + "userinfo-password", + "encoded-userinfo", + "api-key-query", + "unknown-query", + "encoded-query-key", + "fragment", + "empty-userinfo", + "empty-query", + "empty-fragment", + "parser-normalization-error", + "parser-bracket-error", + ], + ) + @pytest.mark.parametrize( + "entry_point", ["constructor", "pydantic-python", "pydantic-json"] + ) + def test_credential_urls_are_refused_before_public_descriptors_exist( + self, url: str, entry_point: str, caplog: pytest.LogCaptureFixture + ) -> None: + with pytest.raises(CodexMcpStartupError) as caught: + match entry_point: + case "constructor": + CodexMcpHttpServer(url) + case "pydantic-python": + TypeAdapter(CodexMcpHttpServer).validate_python({"url": url}) + case "pydantic-json": + TypeAdapter(CodexMcpHttpServer).validate_json( + json.dumps({"url": url}) + ) + assert caught.value.stage == "invalid-descriptor" + diagnostics = ( + str(caught.value) + + repr(caught.value) + + "".join(traceback.format_exception(caught.value)) + + caplog.text + ) + assert url not in diagnostics + assert "url-only-grant" not in diagnostics + + @pytest.mark.parametrize( + "url", + [ + "https://url-only-grant@mcp.example.test/mcp", + "https://mcp.example.test/mcp?api_key=url-only-grant", + "https://mcp.example.test/mcp#url-only-grant", + ], + ids=["userinfo", "query", "fragment"], + ) + def test_snapshot_revalidates_url_before_building_native_overrides( + self, url: str + ) -> None: + server = CodexMcpHttpServer("https://mcp.example.test/mcp") + object.__setattr__(server, "url", url) + with pytest.raises(CodexMcpStartupError) as caught: + freeze_mcp_servers({"remote": server}).launch(("codex", "app-server")) + assert caught.value.stage == "invalid-descriptor" + diagnostics = "".join(traceback.format_exception(caught.value)) + assert "url-only-grant" not in diagnostics + @pytest.mark.parametrize( "name", [ @@ -388,6 +465,58 @@ def test_secrets_only_exist_in_the_private_child_environment( assert secret not in diagnostic assert snapshot.redactor.text(secret) == "[REDACTED]" + @pytest.mark.parametrize( + "url", + [ + "http://localhost/mcp", + "http://127.0.0.1:8765/mcp", + "http://[::1]:8765/mcp", + "https://mcp.example.test", + "HTTPS://mcp.example.test/public%20endpoint", + "https://mcp.example.test/public@path", + ], + ) + def test_public_endpoints_keep_protected_auth_out_of_real_child_argv( + self, url: str + ) -> None: + bearer = "protected-bearer-grant" + header = "protected-header-grant" + server = CodexMcpHttpServer( + url, bearer_token=bearer, headers={"X-Grant": header} + ) + snapshot = freeze_mcp_servers({"remote": server}) + probe = ( + "import json,os,sys,tomllib; " + "entry=tomllib.loads('\\n'.join(sys.argv[2::2]))['mcp_servers']['remote']; " + "refs=[entry['bearer_token_env_var'],*entry['env_http_headers'].values()]; " + "print(json.dumps({'argv':sys.argv[1:]," + "'grants':[os.environ[name] for name in refs]}))" + ) + command = (sys.executable, "-I", "-c", probe) + launch = snapshot.launch(command) + assert launch.env is not None + completed = subprocess.run( + launch.command, + env={**os.environ, **launch.env}, + capture_output=True, + text=True, + check=True, + ) + observed = json.loads(completed.stdout) + assert observed["argv"] == list(launch.command[len(command) :]) + assert observed["grants"] == [bearer, header] + assert launched_servers(launch, command)["remote"]["url"] == url + assert TypeAdapter(CodexMcpHttpServer).dump_python(server) == {"url": url} + diagnostics = ( + repr((server, snapshot, launch)) + + TypeAdapter(CodexMcpHttpServer).dump_json(server).decode() + + TypeAdapter(CodexMcpLaunch).dump_json(launch).decode() + + snapshot.redactor.text(f"failed {bearer} {header}") + ) + for secret in (bearer, header): + assert secret not in " ".join(observed["argv"]) + assert secret not in diagnostics + def test_descriptor_equality_and_serialization_do_not_use_secret_fields( self, ) -> None: diff --git a/tests/adapters/codex/test_mcp_startup.py b/tests/adapters/codex/test_mcp_startup.py index ec5828da2..4b29b2fd5 100644 --- a/tests/adapters/codex/test_mcp_startup.py +++ b/tests/adapters/codex/test_mcp_startup.py @@ -4,6 +4,7 @@ import asyncio import os +import traceback from collections.abc import Mapping from pathlib import Path from types import SimpleNamespace @@ -537,6 +538,42 @@ async def resolve(_workspace: str) -> Mapping[str, CodexMcpServer]: assert adapter._room_clients["a"].client is None await adapter.cleanup_all() + @pytest.mark.parametrize( + "url", + [ + "https://url-only-grant@mcp.example.test/mcp", + "https://mcp.example.test/mcp?api_key=url-only-grant", + "https://mcp.example.test/mcp#url-only-grant", + ], + ids=["userinfo", "query", "fragment"], + ) + async def test_credential_url_refuses_before_client_construction( + self, + tmp_path: Path, + url: str, + caplog: pytest.LogCaptureFixture, + ) -> None: + async def resolve(_workspace: str) -> Mapping[str, CodexMcpServer]: + return {"private": CodexMcpHttpServer(url)} + + adapter, launches = mcp_adapter(tmp_path, [StatusClient()], resolve) + try: + with pytest.raises(CodexMcpStartupError) as error: + await start(adapter) + assert error.value.stage == "resolution" + assert launches == [] + assert adapter._room_clients["a"].client is None + diagnostics = ( + str(error.value) + + repr(error.value) + + "".join(traceback.format_exception(error.value)) + + caplog.text + ) + assert url not in diagnostics + assert "url-only-grant" not in diagnostics + finally: + await adapter.cleanup_all() + async def test_windows_secret_stdio_refuses_before_client_construction( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: From b98fe8f06b08087a972e3b3ff8ef2383954480d5 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Mon, 19 Oct 2026 04:11:57 +0300 Subject: [PATCH 05/13] fix(codex): reject URL authentication before MCP launch --- docs/adapters/codex.md | 12 ++++++++++-- src/band/integrations/codex/mcp.py | 26 +++++++++++++++++++++++--- 2 files changed, 33 insertions(+), 5 deletions(-) diff --git a/docs/adapters/codex.md b/docs/adapters/codex.md index ff8040648..f40d3a3b8 100644 --- a/docs/adapters/codex.md +++ b/docs/adapters/codex.md @@ -133,9 +133,17 @@ Pass the keyword-only `resolve_mcp_servers` callback to `CodexAdapter`, not protected environment names are refused. On Windows, nonempty stdio `env` fails before client construction with stage `unsupported-platform`: `secret-bearing stdio MCP servers require POSIX exec`. Env-free Windows stdio - and all HTTP servers remain supported; no Windows supervisor is included. + and HTTP servers remain supported; no Windows supervisor is included. - `CodexMcpHttpServer(url, bearer_token=None, headers={})` describes streamable - HTTP. Other descriptor types, including SSE, are refused. Names must match + HTTP. `url` must be an absolute `http://` or `https://` endpoint with a host. + Its origin and path are public: they appear in argv, repr and supported + Pydantic serialization, so never put credentials there. URL authentication + is unsupported: userinfo, every query string (regardless of parameter names) + and fragments are refused, including empty `?` or `#` suffixes. Put credentials + in `bearer_token` or `headers` instead. Invalid URLs raise + `CodexMcpStartupError` with stage `invalid-descriptor` without echoing the URL; + descriptor construction inside a resolver is reported as stage `resolution`. + Other descriptor types, including SSE, are refused. Names must match `[A-Za-z0-9_-]+`, be unique, and not be the reserved native name `codex_apps`. - Resolution runs inside the retained room's acquisition, once per new process after any predecessor is confirmed closed. Reusing a process does not resolve diff --git a/src/band/integrations/codex/mcp.py b/src/band/integrations/codex/mcp.py index 3d539c605..58be3f16a 100644 --- a/src/band/integrations/codex/mcp.py +++ b/src/band/integrations/codex/mcp.py @@ -17,6 +17,7 @@ from heapq import heapify, heappop, heappush from types import MappingProxyType from typing import Annotated, Any, NoReturn, SupportsIndex, TypeAlias +from urllib.parse import urlsplit from uuid import uuid4 from pydantic import Field @@ -135,6 +136,23 @@ def _valid_text(value: Any, *, nonempty: bool = False) -> bool: ) +def _public_http_url(value: str) -> bool: + # URL overrides are public argv, unlike protected bearer/header references. + # Refuse every query and fragment rather than guessing credential names. + if "?" in value or "#" in value: + return False + try: + endpoint = urlsplit(value) + return ( + endpoint.scheme in ("http", "https") + and bool(endpoint.hostname) + and endpoint.username is None + and endpoint.password is None + ) + except ValueError: + return False + + def _string_mapping( mapping: Mapping[str, str], stage: str, *, environment: bool = False ) -> dict[str, str]: @@ -191,7 +209,7 @@ def __post_init__(self) -> None: @dataclass(frozen=True, slots=True) class CodexMcpHttpServer: - """A streamable HTTP server; bearer and header values stay out of diagnostics.""" + """A public HTTP endpoint with protected bearer and header credentials.""" url: str bearer_token: Annotated[str | None, Field(exclude=True)] = field( @@ -202,8 +220,10 @@ class CodexMcpHttpServer: ) def __post_init__(self) -> None: - if not _valid_text(self.url, nonempty=True) or ( - self.bearer_token is not None and not _valid_text(self.bearer_token) + if ( + not _valid_text(self.url, nonempty=True) + or not _public_http_url(self.url) + or (self.bearer_token is not None and not _valid_text(self.bearer_token)) ): raise CodexMcpStartupError("invalid-descriptor") object.__setattr__( From 10178db98b811c4f3a696e39f373518b62c5cbb9 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Mon, 19 Oct 2026 04:31:22 +0300 Subject: [PATCH 06/13] docs(codex): separate POSIX cleanup from Windows ownership --- docs/adapters/codex.md | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/adapters/codex.md b/docs/adapters/codex.md index f40d3a3b8..67ee99402 100644 --- a/docs/adapters/codex.md +++ b/docs/adapters/codex.md @@ -41,13 +41,6 @@ For existing integrations, follow the [SDK 4.x to 5.x migration](../acp.md#upgra exited before terminating and draining the job; an empty job alone does not establish helper exit. Failed cleanup retains that same owner and job for retry. -- **POSIX cleanup settles the original owned process group.** Close allows the - leader one second to exit on stdin EOF, then sends TERM to its original group, - waits up to one second, and escalates to KILL with one more second to confirm - `ESRCH`. Unconfirmed cleanup retains the client and workspace for retry. - A reused leader PID is never signaled. POSIX groups lack generation-bound - handles: reuse between a probe and signal remains a residual race. - Descendants that leave the group are outside this guarantee. - **Leaving is final.** Leaving a room stops it admitting work at once. A leave that fails or runs over keeps the process, thread and workspace claim, and the room refuses work until a retried leave succeeds; a later rejoin @@ -87,6 +80,13 @@ For existing integrations, follow the [SDK 4.x to 5.x migration](../acp.md#upgra First-prompt context refreshes communication instructions, not saved tool schemas. - **`reasoning_effort` is not validated.** The valid values depend on the model and the Codex CLI version, and the backend rejects unknown ones. +- **POSIX cleanup settles the original owned process group.** Close allows the + leader one second to exit on stdin EOF, then sends TERM to its original group, + waits up to one second, and escalates to KILL with one more second to confirm + `ESRCH`. Unconfirmed cleanup retains the client and workspace for retry. + A reused leader PID is never signaled. POSIX groups lack generation-bound + handles: reuse between a probe and signal remains a residual race. + Descendants that leave the group are outside this guarantee. - **`skill_roots` need Codex CLI 0.136.0 or newer.** Codex has no config key for extra skill folders, so the adapter sends them to each room's app-server with `skills/extraRoots/set` before the room's thread starts. An older CLI From fab3fb205c6df955f870950a7edf22b2846f47a4 Mon Sep 17 00:00:00 2001 From: Natan Zalkin Date: Thu, 15 Oct 2026 05:34:49 +0300 Subject: [PATCH 07/13] feat(claude): resolve ambient servers under retained runtime custody --- docs/adapters/claude_sdk.md | 64 +- src/band/adapters/claude_sdk.py | 94 +- src/band/core/retention/owner.py | 15 +- src/band/core/retention/room.py | 8 +- .../integrations/claude_sdk/mcp_startup.py | 338 +++++++ .../claude_sdk/session_manager.py | 327 +++++-- src/band/integrations/claude_sdk/transport.py | 88 +- tests/adapters/claude_sdk/external_mcp.py | 81 ++ tests/adapters/claude_sdk/fakecli.py | 77 +- tests/adapters/claude_sdk/helpers.py | 36 +- .../lifecycle/test_external_cancellation.py | 141 +++ .../lifecycle/test_startup_owner.py | 489 ++++++++++ .../lifecycle/test_teardown_custody.py | 104 +++ tests/adapters/claude_sdk/mcpkit.py | 365 ++++++++ tests/adapters/claude_sdk/process.py | 7 + tests/adapters/claude_sdk/test_cli_options.py | 191 ++-- .../adapters/claude_sdk/test_mcp_resolver.py | 873 ++++++++++++++++++ .../claude_sdk/test_mcp_startup_ownership.py | 668 ++++++++++++++ tests/core/retention/test_shutdown.py | 20 + .../claude_sdk/test_confirmed_close.py | 108 +++ tests/mcpclient.py | 33 +- 21 files changed, 3947 insertions(+), 180 deletions(-) create mode 100644 src/band/integrations/claude_sdk/mcp_startup.py create mode 100644 tests/adapters/claude_sdk/external_mcp.py create mode 100644 tests/adapters/claude_sdk/lifecycle/test_external_cancellation.py create mode 100644 tests/adapters/claude_sdk/lifecycle/test_startup_owner.py create mode 100644 tests/adapters/claude_sdk/lifecycle/test_teardown_custody.py create mode 100644 tests/adapters/claude_sdk/mcpkit.py create mode 100644 tests/adapters/claude_sdk/test_mcp_resolver.py create mode 100644 tests/adapters/claude_sdk/test_mcp_startup_ownership.py create mode 100644 tests/integrations/claude_sdk/test_confirmed_close.py diff --git a/docs/adapters/claude_sdk.md b/docs/adapters/claude_sdk.md index 42a742a92..c144686b8 100644 --- a/docs/adapters/claude_sdk.md +++ b/docs/adapters/claude_sdk.md @@ -78,6 +78,65 @@ reuse the failed client or take another room's claimed directory. such as `settings`, `agents`, `bare` and `dangerously-skip-permissions` that would bypass Band's wiring, permission gating or host-config isolation. `env` reaches only the Claude CLI process; the host's `os.environ` is untouched. +- **`resolve_mcp_servers=` adds per-connection MCP servers.** This keyword-only + constructor argument is an async callback `(workspace) -> {name: McpServerConfig}` + that returns the complete set of external servers (native stdio, `http` or `sse` + configs) for a new Claude connection in that room's workspace. + - It runs once per new connection, under the room's retained acquisition; reusing a + room's connection never calls it, so a rotated grant reaches only a later + connection. An empty result means no external servers. A failure to + resolve, read or validate the servers, or to see them connect, fails the + turn with `ClaudeMcpStartupError` and never resumes without the room's + history; only the CLI itself failing to start (a session that can't be + resumed, for one) is retried as a fresh session, which resolves again. + - Band's server is added separately; the name `band`, in-process `sdk` + servers, names that aren't plain strings and a `type` that isn't a string + are refused by name. The servers reach the CLI only as the path of a + private file, never on its command line. On POSIX, its directory is `0700` + and its file is `0600`. On Windows, privacy relies on the user-profile and + temp-directory ACLs, not POSIX mode bits. The file stays until the CLI's exit + is confirmed and the file is removed. + - The session is used only after every server, Band's included, reports + `connected` through the CLI's `mcp_status`; anything else within 30 seconds + (each status request included) fails the turn. Errors and logs name the + stage and server, and the exception type only when it is exactly one of a + few builtin classes (such as `LookupError` or `OSError`); they never read + an exception's own name or module, and never carry callback values, + credentials or native error text, even for a `ClaudeMcpStartupError` the + callback raised, or when the check of a held client's servers is what + fails, with or without a resolver. Such a failure never starts the room's + session again without its history. + - External tools stay under `permission_mode` and `approvals`; they are never + added to the tool allowlist. +- **A starting connection is owned until it settles.** Accepted operations share + the room's retained acquisition, after any outdated CLI is closed. If the + initiating operation stops waiting, the acquisition is cancelled (even if the + callback swallows it), or the room is cleaned up, the startup is retired: it + publishes nothing, is cancelled at most once, and keeps its process, file and + workspace claim until it settles. A resolver may itself request cleanup: + adapter cleanup returns without waiting on that acquisition, while direct + manager cleanup reports `CleanupIncompleteError` to avoid waiting on itself. + The request fences admission immediately, and the retained room completes + disposal once the resolver and its accepted turn finish unwinding. + Admission is rechecked, without calling out, after every callback the startup + runs: a callback that closes admission or retires the room stops any later + reuse, client construction or publication. The reuse decision is settled to a + plain yes/no before that last check, so reading it again calls nothing. + The shared turn owner is bound before tool binding or asynchronous preparation, + so STOP and room retirement also end a turn held in resolution or readiness. + A still-settling turn keeps the room fenced; its provider timeout starts only + at the query, not during preparation. +- **Teardown keeps what it can't release.** `cleanup_all()` refuses new turns + and session startups as soon as it is called, and keeps refusing them until a + successful `on_started`. A CLI whose exit can't be confirmed within 5 seconds of + being killed, or a private file that can't be removed (including one only + partly written), is kept with its workspace claim and the Band server it may + dial, and the cleanup raises a safe error; calling it again (or `on_started`) + retries. A failed room cleanup keeps new input unaccepted until `on_cleanup()` + succeeds on an explicit retry. A file or client left by a failed startup is + retried before a later turn starts a replacement connection. + Concurrent `cleanup_all()` calls share one teardown that a cancelled caller + can't abandon. ```python import pytest @@ -156,8 +215,9 @@ so a slow CLI in one room never holds up another. discards that start. - **Shutdown shares one budget.** `cleanup_all` closes every room, then the shared Band MCP server, within one `retention.transition_timeout_s`. An - unfinished or failed step raises, stays owned, and a retry joins it; the - server stays up while any room's CLI may still dial it. + unfinished or failed step raises and stays owned; a retry joins it. The whole + shutdown continues after a cancelled caller or an expired reporting budget, + and the server stays up while any room's CLI may still dial it. ```python import pytest diff --git a/src/band/adapters/claude_sdk.py b/src/band/adapters/claude_sdk.py index 04e5f6ab2..3548b07c6 100644 --- a/src/band/adapters/claude_sdk.py +++ b/src/band/adapters/claude_sdk.py @@ -84,6 +84,7 @@ ) from band.core.adapterconfig import BaseAdapterConfig from band.core.defaultmodels import ANTHROPIC_MODEL +from band.core.exceptions import CleanupIncompleteError from band.core.protocols import ( FAILURE_CODE_TIMEOUT, GENERIC_PROVIDER_FAILURE_MESSAGE, @@ -94,6 +95,7 @@ RoomRetentionConfig, require_disabled_retention, ) +from band.core.retention.cleanup import finish_by, rejoin_owned, report_owned from band.core.simple_adapter import SimpleAdapter from band.core.turn import judge_detached_turn from band.core.turn_lifecycle import ( @@ -118,6 +120,11 @@ DEFAULT_DEDUP_TTL_SECONDS, DedupingAgentTools, ) +from band.integrations.claude_sdk.mcp_startup import ( + ClaudeMcpStartupError, + McpServerResolver, + is_safe_startup_error, +) from band.integrations.claude_sdk.prompts import generate_claude_sdk_agent_prompt from band.integrations.claude_sdk.session_manager import ( CLI_STARTUP_TIMEOUT_S, # noqa: F401 -- re-exported for adapter configs @@ -612,6 +619,7 @@ def __init__( history_converter: ClaudeSDKHistoryConverter | None = None, additional_tools: list[CustomToolDef] | None = None, workspace_for_room: WorkspaceResolver | None = None, + resolve_mcp_servers: McpServerResolver | None = None, **features: Unpack[FeatureKwargs], ) -> None: """ @@ -625,6 +633,12 @@ def __init__( history_converter: Optional custom history converter additional_tools: Optional list of custom tools as (PydanticModel, callable) tuples. These are converted to MCP tools internally. + resolve_mcp_servers: Async callback returning the complete external + MCP servers for each new connection, given its workspace. Never + called on reuse; Band's server is separate and reserved. Servers + reach the CLI through a private file, and every server must + report connected within 30 seconds before the first prompt. + External tools keep their normal permission mode and approvals. """ if not _CLAUDE_SDK_AVAILABLE: raise ImportError( @@ -645,6 +659,9 @@ def __init__( # Created in on_started. self._session_manager: ClaudeSessionManager | None = None + self._resolve_mcp_servers = resolve_mcp_servers + # One owned whole-adapter close continues after its caller leaves. + self._teardown: asyncio.Task[None] | None = None self._mcp = SharedBandMCPBackend(self._mcp_settings) @@ -709,10 +726,6 @@ def _approvals(self) -> ClaudeApprovalOptions: raise RuntimeError("chat-based approvals are disabled") return approvals - async def _stop_sessions(self) -> None: - if self._session_manager is not None: - await self._session_manager.stop() - @property def _approval_label(self) -> str: approvals = self.config.approvals @@ -726,6 +739,7 @@ async def on_started(self, agent_name: str, agent_description: str) -> None: await self.cleanup_all() await super().on_started(agent_name, agent_description) + self._teardown = None await self._mcp.reopen() mcp_backend = await self._mcp.ensure() @@ -795,6 +809,7 @@ async def on_started(self, agent_name: str, agent_description: str) -> None: can_use_tool_factory=can_use_tool_factory, mcp_servers_factory=self._room_mcp_servers, workspaces=self._workspaces, + resolve_mcp_servers=self._resolve_mcp_servers, retention=self.config.retention, settle_grace_s=self.config.turn_settle_timeout_s, on_disposed=self._forget_room, @@ -922,8 +937,15 @@ async def on_message( tools=tools, room_id=room_id, ) + try: + owner = bind_turn_task(tools, turn_task) + except BaseException: + # The eager worker is parked on this gate until registration succeeds. + owner_ready.cancel() + await asyncio.wait({turn_task}) + raise self._turn_release[room_id] = release_future - self._turn_tasks[room_id] = owner = bind_turn_task(tools, turn_task) + self._turn_tasks[room_id] = owner owner_ready.set_result(owner) try: await asyncio.wait( @@ -1004,6 +1026,8 @@ async def _run_turn( """Run one turn under the room's lease; always releases ``release_future``.""" owner = await owner_ready try: + if owner.deadline is not None: + raise asyncio.CancelledError async with ( self.backend_turns.running(session.room_id, session.room_id), self._turn_lease(session, tools, resume_from) as client, @@ -1040,15 +1064,26 @@ async def _turn_lease( try: # Before the lease: a replaced server retires the room's CLI, and # that recovery runs only while no operation leases the room. + if self._session_manager is not None: + self._session_manager.require_admission() await self._mcp.ensure() session.retire_outdated() prepare = partial(self._bind_turn, session, tools, resume_from) async with session.using(prepare) as client: started = True yield client - except Exception: + except Exception as error: if not started and not session.retained.closing: - logger.exception("Room %s: Session creation failed", session.room_id) + if isinstance(error, ClaudeMcpStartupError): + logger.error( + "Room %s: Session creation failed: %s", + session.room_id, + error if is_safe_startup_error(error) else "an untrusted error", + ) + else: + logger.exception( + "Room %s: Session creation failed", session.room_id + ) await tools.send_failure(_TURN_FAILURE) raise @@ -1642,7 +1677,11 @@ async def on_cleanup(self, room_id: str) -> None: if manager is None or manager.session(room_id) is None: self._forget_room(room_id) return - await manager.cleanup_session(room_id) + try: + await manager.cleanup_session(room_id) + except CleanupIncompleteError: + if not manager.owns_task(asyncio.current_task()): + raise def _forget_room(self, room_id: str) -> None: self._room_tools.pop(room_id, None) @@ -1655,18 +1694,43 @@ def _forget_room(self, room_id: str) -> None: logger.debug("Room %s: Cleaned up Claude SDK session", room_id) async def cleanup_all(self) -> None: - """Stop the session manager, then the shared Band MCP server, within - one ``retention.transition_timeout_s`` budget. + """Own the complete session-then-MCP shutdown under one reporting budget. - Stopping closes admission for every room at once; rooms close first - because their CLIs dial the server until they do. An unfinished or - failed step stays owned for a retry. + A cancelled caller leaves the whole close running. The retained room + contract reports self-requested cleanup instead of waiting on its own + acquisition; that request still fences admission immediately. """ + manager = self._session_manager + if manager is not None: + manager.close_admission() + if manager.owns_task(asyncio.current_task()): + try: + await manager.stop() + except CleanupIncompleteError: + pass # the retained owner joins this work once it unwinds + self._teardown = rejoin_owned( + self._teardown, self._tear_down, label="Claude adapter shutdown" + ) + return for room_id in list(self._pending_approvals): self._clear_pending_approvals_for_room(room_id) - await self._mcp.close_after( - self._stop_sessions(), retention=self.config.retention + self._teardown = rejoin_owned( + self._teardown, self._tear_down, label="Claude adapter shutdown" ) + deadline = ( + asyncio.get_running_loop().time() + + self.config.retention.transition_timeout_s + ) + await finish_by( + deadline, + report_owned(self._teardown, None), + what="Claude sessions and MCP shutdown", + ) + + async def _tear_down(self) -> None: + if self._session_manager is not None: + await self._session_manager.stop(report_timeout_s=None) + await self._mcp.close(final=True) await self.backend_turns.drain() # ------------------------------------------------------------------ diff --git a/src/band/core/retention/owner.py b/src/band/core/retention/owner.py index 920c8322a..b73b75dbd 100644 --- a/src/band/core/retention/owner.py +++ b/src/band/core/retention/owner.py @@ -5,6 +5,7 @@ import asyncio from collections.abc import Awaitable, Callable from functools import partial +from typing import Literal from band.core.exceptions import CleanupIncompleteError, RoomClosingError from band.core.retention.cleanup import ( @@ -101,11 +102,16 @@ def _release_due_rooms(self, due: float) -> None: room._start_idle_release() self.reschedule() - async def stop(self) -> None: + async def stop( + self, *, report_timeout_s: float | None | Literal["configured"] = "configured" + ) -> None: """Dispose every room under one reporting deadline; retain failures. Work enrolled in a room cannot await this shutdown; it is told cleanup is pending, is not cancelled by it, and is joined once it unwinds. + By default the caller uses ``transition_timeout_s``. An enclosing + owner may pass ``None`` to await settlement without a second reporting + deadline, while its own callers still receive bounded reports. """ requested_from_inside = [ room._spare_requester() for room in tuple(self._rooms.values()) @@ -119,7 +125,12 @@ async def stop(self) -> None: raise CleanupIncompleteError( "Room retention shutdown waits for the calling work to finish" ) - await report_owned(shutdown, self.config.transition_timeout_s) + timeout = ( + self.config.transition_timeout_s + if isinstance(report_timeout_s, str) + else report_timeout_s + ) + await report_owned(shutdown, timeout) def renewed(self) -> RoomRetention: """The owner to serve with after a restart. diff --git a/src/band/core/retention/room.py b/src/band/core/retention/room.py index 19a4a8cde..3afad04f8 100644 --- a/src/band/core/retention/room.py +++ b/src/band/core/retention/room.py @@ -619,7 +619,7 @@ def _spare_requester(self) -> bool: current = asyncio.current_task() if not self.enrolls(current): return False - if current is not None and current in self._owned: + if current is not None: self._unwinding.add(current) return True @@ -682,7 +682,11 @@ async def _settle_accepted(self) -> None: self._work_cancelled = True for task in self._uncancelled_work(): task.cancel() - if self._state == Residency.ACQUIRING and self._transition is not None: + if ( + self._state == Residency.ACQUIRING + and self._transition is not None + and self._transition not in self._unwinding + ): self._cancel_acquire( self._transition, reason=f"Room {self.room_id} acquisition was cancelled for cleanup", diff --git a/src/band/integrations/claude_sdk/mcp_startup.py b/src/band/integrations/claude_sdk/mcp_startup.py new file mode 100644 index 000000000..cc642af07 --- /dev/null +++ b/src/band/integrations/claude_sdk/mcp_startup.py @@ -0,0 +1,338 @@ +"""Connection-time MCP servers for Claude sessions: one frozen snapshot per new +connection, written to a private file, and positive per-server readiness. + +Nothing here renders callback values, native status payloads or native error +text: they can carry credentials. Errors name the stage, a known server name, +a fixed reason and the exception type only, and only errors created here are +trusted to say so -- one raised by a callback, even a ``ClaudeMcpStartupError``, +is replaced. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import shutil +import tempfile +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any, ForwardRef, TypeAlias + +if TYPE_CHECKING: + from claude_agent_sdk.types import McpServerConfig + + McpServerResolver: TypeAlias = Callable[ + [str], Awaitable[Mapping[str, McpServerConfig]] + ] +else: + # A real alias, so annotations such as ``McpServerResolver | None`` + # evaluate at runtime; the native config type is named by reference and + # resolved from the SDK only when introspected, keeping the SDK optional. + McpServerResolver = Callable[ + [str], + Awaitable[ + Mapping[str, ForwardRef("McpServerConfig", module="claude_agent_sdk.types")] + ], + ] +"""``resolve_mcp_servers(workspace)``: the complete external MCP servers for a +new Claude connection in ``workspace``, as native ``McpServerConfig`` values.""" + +MCP_READINESS_TIMEOUT_SECONDS = 30.0 +_STATUS_POLL_SECONDS = 0.25 +_REMOTE_TYPES = frozenset({"http", "sse"}) +_FAILED_STATUSES = frozenset({"failed", "needs-auth", "disabled"}) +_KNOWN_STATUSES = _FAILED_STATUSES | {"connected", "pending"} + + +class ClaudeMcpStartupError(RuntimeError): + """A Claude session's MCP servers could not be resolved, launched, confirmed + connected or cleaned up. Carries no callback, native status or error text. + + ``resume_eligible`` is set only when the native CLI failed to start, the one + case a resumed session may be retried as a fresh one. + """ + + def __init__( + self, + stage: str, + reason: str, + *, + server: str | None = None, + error_type: str | None = None, + resume_eligible: bool = False, + ) -> None: + self.stage = stage + self.reason = reason + self.server = server + self.error_type = error_type + self.resume_eligible = resume_eligible + parts = [f"Claude MCP startup failed at {stage}: {reason}"] + if server is not None: + parts.append(f"server {server!r}") + if error_type is not None: + parts.append(f"({error_type})") + super().__init__(" ".join(parts)) + + +class _OwnedStartupError(ClaudeMcpStartupError): + """Built here from owner-known metadata only.""" + + +def startup_error( + stage: str, + reason: str, + error: BaseException | None = None, + *, + server: str | None = None, + resume_eligible: bool = False, +) -> ClaudeMcpStartupError: + """A fresh safe error for ``error``, whatever its type; raise it ``from + None`` so the original text, which may hold credentials, never joins the + chain. The error's type is named only when it is a builtin or comes from a + library this adapter runs: a callback can name its own classes anything.""" + return _OwnedStartupError( + stage, + reason, + server=server, + error_type=_known_type_name(error), + resume_eligible=resume_eligible, + ) + + +# Literal labels for exactly these classes. An error is matched by the +# identity of its class alone -- never its name, module or other metadata, +# which a callback's own class can set (or compute) to anything. +_LABELLED_ERRORS: tuple[tuple[type[BaseException], str], ...] = ( + (LookupError, "LookupError"), + (KeyError, "KeyError"), + (IndexError, "IndexError"), + (TypeError, "TypeError"), + (ValueError, "ValueError"), + (RuntimeError, "RuntimeError"), + (OSError, "OSError"), + (PermissionError, "PermissionError"), + (FileNotFoundError, "FileNotFoundError"), + (ConnectionError, "ConnectionError"), + (TimeoutError, "TimeoutError"), +) + + +def _known_type_name(error: BaseException | None) -> str | None: + if error is None: + return None + cls = type(error) + for known, label in _LABELLED_ERRORS: + if cls is known: + return label + return None + + +def is_safe_startup_error(error: BaseException) -> bool: + """Whether ``error`` was built by the owner, not handed in by a callback.""" + return isinstance(error, _OwnedStartupError) + + +_NOT_A_MAPPING = object() + + +def freeze_external_servers( + resolved: object, protected: Mapping[str, Any] +) -> dict[str, dict[str, Any]]: + """An owned deep copy of the callback's servers, validated for the native CLI. + + Later changes to the mappings, lists or headers the callback returned can't + reach the launch. Reading the callback's objects runs their code, so every + exception raised while copying them -- whatever its type -- is replaced by + a safe error; only the plain copy is then validated, and in-process ``sdk`` + servers and names Band reserves are refused by name. + """ + return _validate(_copy_servers(resolved), protected) + + +def _copy_servers(resolved: object) -> dict[str, object]: + """Plain JSON values copied out of the callback's result. + + Every operation on the callback's objects -- the type check, iterating its + items, unpacking each item, reading and serializing each config -- runs + callback code, so any exception from one of them, whatever its type, is + replaced by a safe error. Only exact ``str`` names and plain copies leave; + no callback code runs after this returns. + """ + pairs: list[tuple[object, object]] | None = None + try: + if isinstance(resolved, Mapping): + pairs = [] + for item in resolved.items(): + name, config = item + pairs.append((name, config)) + except Exception as error: # noqa: BLE001 -- untrusted, whatever its type + raise startup_error( + "resolution", "the resolved servers could not be read", error + ) from None + if pairs is None: + raise startup_error( + "resolution", "resolve_mcp_servers must return a mapping of servers" + ) + copied: dict[str, object] = {} + for name, config in pairs: + # Exactly ``str`` (checked without calling into it): a subclass could + # render or compare itself however it likes. + if type(name) is not str or not name: + raise startup_error("resolution", "server names must be strings") + if name in copied: + raise startup_error("resolution", "duplicate server name", server=name) + try: + if isinstance(config, Mapping): + plain: object = json.loads(json.dumps(dict(config))) + else: + plain = _NOT_A_MAPPING + except (TypeError, ValueError) as error: + raise startup_error( + "resolution", "server config is not serializable", error, server=name + ) from None + except Exception as error: # noqa: BLE001 -- untrusted, whatever its type + raise startup_error( + "resolution", + "the resolved servers could not be read", + error, + server=name, + ) from None + copied[name] = plain + return copied + + +def _validate( + copied: dict[str, object], protected: Mapping[str, Any] +) -> dict[str, dict[str, Any]]: + frozen: dict[str, dict[str, Any]] = {} + for name, config in copied.items(): + if name in protected: + raise startup_error( + "resolution", "the name is reserved for Band", server=name + ) + if not isinstance(config, dict): + raise startup_error( + "resolution", "server config must be a mapping", server=name + ) + kind = config.get("type", "stdio") + if not isinstance(kind, str): + raise startup_error( + "resolution", "server type must be a string", server=name + ) + if kind == "sdk": + raise startup_error( + "resolution", "in-process sdk servers are not supported", server=name + ) + if kind not in _REMOTE_TYPES and kind != "stdio": + raise startup_error("resolution", "unsupported server type", server=name) + required = "command" if kind == "stdio" else "url" + if not isinstance(config.get(required), str): + raise startup_error( + "resolution", f"server config needs a string {required}", server=name + ) + frozen[name] = config + return frozen + + +@dataclass +class McpConfigFile: + """The private native ``--mcp-config`` file one connection launches with. + + ``allocate`` creates only its private directory, so the caller owns it + before anything can fail; ``write`` then fills it. Owned until its child's + exit is confirmed; then ``release`` removes it. A release that fails raises + and leaves it owned, for a later retry. + """ + + directory: Path + path: Path + + @classmethod + def allocate(cls) -> McpConfigFile: + directory = Path(tempfile.mkdtemp(prefix="band-claude-mcp-")) + return cls(directory=directory, path=directory / "mcp.json") + + def write(self, servers: Mapping[str, Any]) -> None: + os.chmod(self.directory, 0o700) + fd = os.open(self.path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump({"mcpServers": dict(servers)}, handle) + + def release(self) -> None: + try: + shutil.rmtree(self.directory) + except FileNotFoundError: + return + except OSError as error: + raise startup_error( + "cleanup", "the private MCP file could not be removed", error + ) from None + + +def _check_status( + response: object, expected: frozenset[str] +) -> tuple[bool, str | None]: + """``(ready, waiting_on)``; raises when the evidence is unusable or a + server failed. Only a unique ``connected`` entry per expected name counts.""" + servers = response.get("mcpServers") if isinstance(response, Mapping) else None + if not isinstance(servers, list): + raise startup_error("status", "malformed status response") + seen: dict[str, str] = {} + for entry in servers: + if not isinstance(entry, Mapping): + raise startup_error("status", "malformed status entry") + name, status = entry.get("name"), entry.get("status") + if not isinstance(name, str) or not isinstance(status, str): + raise startup_error("status", "malformed status entry") + if name not in expected: + continue + if name in seen: + raise startup_error("status", "duplicate status", server=name) + if status not in _KNOWN_STATUSES: + raise startup_error("status", "unknown status", server=name) + if status in _FAILED_STATUSES: + raise startup_error("status", f"server is {status}", server=name) + seen[name] = status + for name in sorted(expected): + if seen.get(name) != "connected": + return False, name + return True, None + + +async def await_mcp_ready( + get_status: Callable[[], Awaitable[object]], + expected: frozenset[str], + *, + timeout: float = MCP_READINESS_TIMEOUT_SECONDS, +) -> None: + """Wait until every expected server reports ``connected``, within one total + deadline that also bounds each status request.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + waiting_on: str | None = None + while True: + remaining = deadline - loop.time() + if remaining <= 0: + raise startup_error( + "status", "not connected before the deadline", server=waiting_on + ) + try: + async with asyncio.timeout(remaining) as budget: + response = await get_status() + except TimeoutError: + raise startup_error( + "status", "not connected before the deadline", server=waiting_on + ) from None + except Exception as error: # noqa: BLE001 -- any native status failure is a typed startup failure + raise startup_error("status", "status request failed", error) from None + if budget.expired(): + # A callback may catch the timer's cancellation and return anyway. + raise startup_error( + "status", "not connected before the deadline", server=waiting_on + ) + ready, waiting_on = _check_status(response, expected) + if ready: + return + await asyncio.sleep(min(_STATUS_POLL_SECONDS, max(0.0, deadline - loop.time()))) diff --git a/src/band/integrations/claude_sdk/session_manager.py b/src/band/integrations/claude_sdk/session_manager.py index 5f663a87c..d6eed97c3 100644 --- a/src/band/integrations/claude_sdk/session_manager.py +++ b/src/band/integrations/claude_sdk/session_manager.py @@ -18,7 +18,7 @@ from contextlib import asynccontextmanager, contextmanager from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Any, Literal from band.core.exceptions import RoomClosingError from band.core.retention import ( @@ -29,6 +29,15 @@ RoomRetentionConfig, ) from band.core.turn_lifecycle import TURN_SETTLE_GRACE_S +from band.integrations.claude_sdk.mcp_startup import ( + ClaudeMcpStartupError, + McpConfigFile, + McpServerResolver, + await_mcp_ready, + freeze_external_servers, + is_safe_startup_error, + startup_error, +) from band.workspaces import RoomWorkspaces try: @@ -123,6 +132,10 @@ def _start_failure(error: Exception) -> BaseException: return error +class _StartupReplaced(Exception): + """The Band endpoint changed before this connection was ready.""" + + class ClaudeRoomSession: """A room's Claude CLI session and everything it keeps until closure. @@ -144,6 +157,9 @@ def __init__( self._reader: asyncio.Task[None] | None = None self._turn: TurnStream | None = None self._forget_session = False + self._config_file: McpConfigFile | None = None + self._protected_servers: dict[str, McpServerConfig] | None = None + self._acquisition: asyncio.Task[None] | None = None self.retained: RetainedRoom = retention.register( room_id, acquire=self.acquire, release=self.release, dispose=self.dispose ) @@ -172,10 +188,27 @@ async def using( ``prepare`` binds the operation before a CLI is started for it. """ - async with self.retained.use(prepare): - if self.client is None: - raise RuntimeError(f"Room {self.room_id} is resident without a client") - yield self.client + try: + async with self.retained.use(prepare): + if self.client is None: + raise RuntimeError( + f"Room {self.room_id} is resident without a client" + ) + yield self.client + except asyncio.CancelledError: + acquisition = self._acquisition + if ( + self._manager._resolve_mcp_servers is not None + and acquisition is not None + and ( + acquisition.cancelling() + or not self.retained.has_authority(self.generation) + ) + ): + # The turn's runtime owner still answers for preparation until + # its cancelled callback and failed launch finish settling. + await asyncio.wait({acquisition}) + raise @contextmanager def listening(self) -> Iterator[TurnStream]: @@ -210,38 +243,67 @@ def _retire(self, generation: int, reason: str, *, forget_session: bool) -> bool return True def retire_outdated(self) -> None: - """A CLI dialing MCP servers the room no longer has resumes on the - current ones.""" - if self.client is None: - return - if self.client.options.mcp_servers != self._manager.mcp_servers(self.room_id): + """Compare only the protected servers: reuse never resolves external ones.""" + changed = False + try: + if self.client is not None or self._config_file is not None: + held: object = self._protected_servers + if ( + self._manager._resolve_mcp_servers is None + and self.client is not None + ): + held = self.client.options.mcp_servers + current = ( + self._manager.protected_servers(self.room_id) + if self._manager._resolve_mcp_servers is not None + else self._manager.mcp_servers(self.room_id) + ) + changed = bool(held != current) + except Exception as error: # noqa: BLE001 -- comparisons can run callback code + raise startup_error( + "reservation", "the room's protected servers could not be read", error + ) from None + self._manager.require_admission() + if changed: logger.info("Room %s: MCP servers changed; replacing its CLI", self.room_id) self.recover("Band MCP servers changed") async def acquire(self, generation: int, resume: bool) -> None: - """Start the room's CLI. ``resume`` restores the retained session - strictly; otherwise an unresumable earlier session gives way to a new - one, the bootstrap and crash policy, unless the room began closing.""" + """Resolve and ready each new connection under the retained acquisition.""" + self._acquisition = asyncio.current_task() try: - await self._connect(generation, resume=self.session_id) - except Exception as error: - if not self._may_start_fresh(resume): - raise - logger.warning( - "Room %s: Session resume failed (session_id=%s): %s. " - "Creating new session", - self.room_id, - self.session_id, - error, - ) - # The conversation is dropped only once a new one is running. - await self._connect(generation, resume=None) - self.session_id = None + while True: + try: + await self._connect(generation, resume=self.session_id) + except _StartupReplaced: + # Close the obsolete launch before resolving a fresh snapshot. + await self._close() + self._check_admitted(generation) + continue + except Exception as error: + if ( + isinstance(error, ClaudeMcpStartupError) + and ( + not is_safe_startup_error(error) + or not error.resume_eligible + ) + ) or not self._may_start_fresh(resume): + raise + logger.warning( + "Room %s: Session resume failed (session_id=%s): %s. " + "Creating new session", + self.room_id, + self.session_id, + error, + ) + await self._connect(generation, resume=None) + self.session_id = None + return + finally: + self._acquisition = None def _may_start_fresh(self, resume: bool) -> bool: - """Whether a failed resume may give way to a new session: never on a - strict resume or once closing, and not while the failed CLI, which - could not be closed, is still held for recovery.""" + """Only a failed native start may discard history, after its child closes.""" return ( not resume and self.session_id is not None @@ -249,6 +311,44 @@ def _may_start_fresh(self, resume: bool) -> bool: and self.client is None ) + def _check_admitted(self, generation: int) -> None: + """Callback-free fencing, including cancellation a callback swallowed.""" + try: + self._manager.require_admission() + except ClaudeSessionManagerStoppedError: + raise asyncio.CancelledError from None + current = asyncio.current_task() + if current is not None and current.cancelling(): + raise asyncio.CancelledError from None + if self.retained.closing: + if self._manager._resolve_mcp_servers is None: + self.retained.require_admission() + raise startup_error("retirement", "the room's startup was retired") + if ( + self._manager._resolve_mcp_servers is not None + and not self.retained.has_authority(generation) + ): + raise asyncio.CancelledError from None + + def _check_current(self, generation: int) -> None: + """Materialize comparisons before the final callback-free guard.""" + changed = False + if self._manager._resolve_mcp_servers is not None: + try: + changed = bool( + self._protected_servers + != self._manager.protected_servers(self.room_id) + ) + except Exception as error: # noqa: BLE001 -- a protected-server factory + raise startup_error( + "reservation", + "the room's protected servers could not be read", + error, + ) from None + self._check_admitted(generation) + if changed: + raise _StartupReplaced + async def release(self, cause: ReleaseCause) -> ReleaseOutcome: """Close the room's CLI for recovery or final cleanup.""" await self._close() @@ -264,19 +364,87 @@ async def dispose(self) -> None: self._manager.discard(self) async def _connect(self, generation: int, *, resume: str | None) -> None: - options = self._manager.options_for(self, generation, resume=resume) - transport = create_transport(options) - client = ClaudeSDKClient(options=options, transport=transport) - # Retain both before connect: SDK disconnect cannot close a subprocess - # when startup fails before its Query is constructed. - self.client, self.transport = client, transport + # A failed removal can leave only a file. Retry that obligation before + # resolving or allocating anything for the next attempt. + if self._config_file is not None: + await self._close() + self._check_admitted(generation) + stage = "reservation" + native_start = False + resolving = self._manager._resolve_mcp_servers is not None + ready_servers: frozenset[str] = frozenset() try: + if resolving: + workspace = self._manager.workspace(self.room_id) + protected = self._manager.protected_servers(self.room_id) + self._protected_servers = protected + self._check_current(generation) + stage = "resolution" + resolver = self._manager._resolve_mcp_servers + assert resolver is not None + try: + resolved = await resolver(workspace) + except Exception as error: # noqa: BLE001 -- callback errors are never trusted + raise startup_error( + "resolution", "resolve_mcp_servers failed", error + ) from None + external = freeze_external_servers(resolved, protected) + ready_servers = frozenset({*external, *protected}) + self._check_current(generation) + stage = "launch" + self._config_file = McpConfigFile.allocate() + self._config_file.write({**external, **protected}) + options = self._manager.options_for( + self, + generation, + resume=resume, + mcp_servers=str(self._config_file.path), + ) + else: + options = self._manager.options_for(self, generation, resume=resume) + self._check_current(generation) + self.transport = transport = create_transport(options) + self.client = client = ClaudeSDKClient(options=options, transport=transport) + # Retain both before connect: SDK disconnect cannot close a child + # when startup fails before its Query is constructed. + native_start = True await client.connect() + native_start = False + self._check_current(generation) + if resolving: + stage = "status" + try: + await await_mcp_ready(client.get_mcp_status, ready_servers) + except ClaudeMcpStartupError: + self._check_current(generation) + raise + self._check_current(generation) + except asyncio.CancelledError: + if resolving: + await self._close_failed_start() + raise asyncio.CancelledError from None + except _StartupReplaced: + raise except Exception as error: await self._close_failed_start() - if (failure := _start_failure(error)) is error: + self._check_admitted(generation) + failure = _start_failure(error) + if not resolving: + if failure is error: + raise + raise failure + if isinstance(error, ClaudeMcpStartupError) and is_safe_startup_error( + error + ): raise - raise failure + raise startup_error( + stage, + "the Claude CLI did not start" + if native_start + else "the Claude session could not be constructed", + failure, + resume_eligible=native_start, + ) from None self._reader = asyncio.create_task(self._read(client, generation)) logger.info( "Room %s: Claude session started (resume=%s)", self.room_id, options.resume @@ -291,20 +459,32 @@ async def _close_failed_start(self) -> None: logger.warning( "Room %s: closing a failed Claude start failed; recovery retries it", self.room_id, - exc_info=True, + exc_info=self._manager._resolve_mcp_servers is None, ) async def _close(self) -> None: - """Close the retained CLI; a failure keeps every handle for a retry.""" - if self.client is None or self.transport is None: - return + """Keep every unresolved child or private-file obligation for a retry.""" try: - await self.client.disconnect() - finally: - await self.transport.close() - if self._reader is not None: - await self._reader - self.client = self.transport = self._reader = None + if self.client is not None or self.transport is not None: + try: + if self.client is not None: + await self.client.disconnect() + finally: + if self.transport is not None: + await self.transport.close() + if self._reader is not None: + await self._reader + self.client = self.transport = self._reader = None + except Exception as error: + if self._manager._resolve_mcp_servers is None: + raise + raise startup_error( + "cleanup", "the Claude CLI could not be confirmed closed", error + ) from None + if self._config_file is not None: + self._config_file.release() + self._config_file = None + self._protected_servers = None async def _read(self, client: ClaudeSDKClient, generation: int) -> None: """The CLI's only reader, across turns, from connect to disconnect.""" @@ -373,6 +553,7 @@ def __init__( retention: RoomRetentionConfig | None = None, settle_grace_s: float = TURN_SETTLE_GRACE_S, on_disposed: Callable[[str], None] | None = None, + resolve_mcp_servers: McpServerResolver | None = None, ) -> None: """ Initialize session manager. @@ -392,22 +573,36 @@ def __init__( down on its own before cancelling it. on_disposed: Called with a room's id once its final release is confirmed, even if the leave that started it stopped waiting. + resolve_mcp_servers: Complete external servers for each new connection; + never called on reuse, and never allowed to replace Band's servers. """ self.base_options = base_options self._can_use_tool_factory = can_use_tool_factory self._mcp_servers_factory = mcp_servers_factory self._workspaces = workspaces self._on_disposed = on_disposed + self._resolve_mcp_servers = resolve_mcp_servers self._retention = RoomRetention( retention or default_claude_retention(), settle_grace_s=settle_grace_s, ) self._rooms: dict[str, ClaudeRoomSession] = {} - def room(self, room_id: str) -> ClaudeRoomSession: - """The room's session record, registered on first use.""" + def require_admission(self) -> None: if self._retention.stopping: raise ClaudeSessionManagerStoppedError() + + def close_admission(self) -> None: + """Fence new work now; the retained owner settles it through ``stop``.""" + self._retention.stopping = True + self._retention.reschedule() + + def owns_task(self, task: asyncio.Task[Any] | None) -> bool: + return any(session.retained.enrolls(task) for session in self._rooms.values()) + + def room(self, room_id: str) -> ClaudeRoomSession: + """The room's session record, registered on first use.""" + self.require_admission() if (session := self._rooms.get(room_id)) is None: session = self._rooms[room_id] = ClaudeRoomSession( self, room_id, self._retention @@ -423,14 +618,36 @@ def mcp_servers(self, room_id: str) -> dict[str, McpServerConfig] | str | Path: return self.base_options.mcp_servers return self._mcp_servers_factory(room_id) + def protected_servers(self, room_id: str) -> dict[str, McpServerConfig]: + """Band's protected servers, separate from the resolver's complete result.""" + return ( + {} + if self._mcp_servers_factory is None + else self._mcp_servers_factory(room_id) + ) + + def workspace(self, room_id: str) -> str: + if self._workspaces is not None: + return self._workspaces.claim(room_id) + return str(self.base_options.cwd or Path.cwd()) + def options_for( - self, session: ClaudeRoomSession, generation: int, *, resume: str | None + self, + session: ClaudeRoomSession, + generation: int, + *, + resume: str | None, + mcp_servers: str | None = None, ) -> ClaudeAgentOptions: """``base_options`` with one room session's overrides applied. ``dataclasses.replace()`` shallow-copies every base field. """ - overrides: dict[str, Any] = {"mcp_servers": self.mcp_servers(session.room_id)} + overrides: dict[str, Any] = { + "mcp_servers": self.mcp_servers(session.room_id) + if mcp_servers is None + else mcp_servers + } if self._workspaces is not None: overrides["cwd"] = self._workspaces.claim(session.room_id) if resume: @@ -457,10 +674,12 @@ async def cleanup_session(self, room_id: str) -> None: if (session := self._rooms.get(room_id)) is not None: await session.retained.dispose() - async def stop(self) -> None: + async def stop( + self, *, report_timeout_s: float | None | Literal["configured"] = "configured" + ) -> None: """Dispose every room; a failed or unfinished room stays owned for a retry, and the manager never serves again.""" - await self._retention.stop() + await self._retention.stop(report_timeout_s=report_timeout_s) def has_session(self, room_id: str) -> bool: """Whether the room has a CLI process held.""" diff --git a/src/band/integrations/claude_sdk/transport.py b/src/band/integrations/claude_sdk/transport.py index a157c0478..fe810306c 100644 --- a/src/band/integrations/claude_sdk/transport.py +++ b/src/band/integrations/claude_sdk/transport.py @@ -1,13 +1,97 @@ -"""Claude subprocess construction, including SDK permission normalization.""" +"""Claude subprocess construction, including SDK permission normalization, and +custody of the CLI child until its exit is confirmed.""" from __future__ import annotations +import asyncio +import contextlib +from collections.abc import AsyncIterator +from typing import Any + from claude_agent_sdk import ClaudeAgentOptions from claude_agent_sdk._internal.transport import Transport from claude_agent_sdk._internal.transport.subprocess_cli import SubprocessCLITransport from claude_agent_sdk.types import _configure_can_use_tool +CHILD_EXIT_TIMEOUT_SECONDS = 5.0 +"""How long a child the native close left running may take to die once killed.""" + + +class ClaudeProcessCloseError(RuntimeError): + """The Claude CLI child could not be confirmed exited; its transport keeps + the child for a later ``close()`` to retry.""" + + +class ConfirmedCloseTransport(Transport): + """The SDK's subprocess transport, holding the CLI child until it is reaped. + + The native ``close()`` can return after its waits time out with the child + still running, and it forgets the child either way, so neither its return + nor a cleared handle proves exit. This wrapper keeps the child it saw + spawned, kills and reaps one that a normal close left running, and keeps + the child (raising ``ClaudeProcessCloseError``) when that cannot be + confirmed. A close that raised or was cancelled also keeps the child, so a + later ``close()`` retries. + """ + + def __init__(self, inner: Transport) -> None: + self._inner = inner + self._child: Any = None + + @property + def close_pending(self) -> bool: + """Whether a child is held whose exit is not yet confirmed.""" + return self._child is not None + + def _observe_child(self) -> None: + if (process := getattr(self._inner, "_process", None)) is not None: + self._child = process + + async def connect(self) -> None: + try: + await self._inner.connect() + finally: + self._observe_child() + + async def write(self, data: str) -> None: + await self._inner.write(data) + + def read_messages(self) -> AsyncIterator[dict[str, Any]]: + return self._inner.read_messages() + + def is_ready(self) -> bool: + return self._inner.is_ready() + + async def end_input(self) -> None: + await self._inner.end_input() + + async def close(self) -> None: + self._observe_child() + await self._inner.close() + await self._confirm_exit() + + async def _confirm_exit(self) -> None: + """Kill and reap a child the native close left running, within + ``CHILD_EXIT_TIMEOUT_SECONDS``; otherwise keep it and raise.""" + child = self._child + if child is None: + return + if child.returncode is None: + try: + with contextlib.suppress(ProcessLookupError): + child.kill() + async with asyncio.timeout(CHILD_EXIT_TIMEOUT_SECONDS): + await child.wait() + except Exception as error: # noqa: BLE001 -- unconfirmed exit, whatever the cause + raise ClaudeProcessCloseError( + "the Claude CLI process could not be confirmed exited" + f" ({type(error).__name__})" + ) from None + self._child = None + def create_transport(options: ClaudeAgentOptions) -> Transport: """Keep permission routing identical to the SDK's default transport.""" - return SubprocessCLITransport(prompt="", options=_configure_can_use_tool(options)) + return ConfirmedCloseTransport( + SubprocessCLITransport(prompt="", options=_configure_can_use_tool(options)) + ) diff --git a/tests/adapters/claude_sdk/external_mcp.py b/tests/adapters/claude_sdk/external_mcp.py new file mode 100644 index 000000000..c9a3f318a --- /dev/null +++ b/tests/adapters/claude_sdk/external_mcp.py @@ -0,0 +1,81 @@ +"""A real local MCP server outside Band that only accepts its current bearer +grant, recording every request it sees (for connection-time resolver tests).""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass, field +from typing import Any + +import uvicorn +from mcp.server.fastmcp import FastMCP + + +@dataclass +class ExternalMcp: + url: str + grants: set[str] + # (grant presented, HTTP status) per request. + receipts: list[tuple[str, int]] = field(default_factory=list) + + def server(self, grant: str, **extra: Any) -> dict[str, Any]: + """The native config a resolver returns for this server and grant.""" + return { + "type": "http", + "url": self.url, + "headers": {"Authorization": f"Bearer {grant}"}, + **extra, + } + + +@asynccontextmanager +async def external_mcp(*grants: str) -> AsyncIterator[ExternalMcp]: + mcp = FastMCP("external", stateless_http=True) + + @mcp.tool() + def echo(text: str) -> str: + return text + + inner = mcp.streamable_http_app() + state = ExternalMcp(url="", grants=set(grants)) + + async def app(scope: Any, receive: Any, send: Any) -> None: + if scope["type"] != "http": + await inner(scope, receive, send) + return + auth = dict(scope["headers"]).get(b"authorization", b"").decode() + grant = auth.removeprefix("Bearer ") + if grant not in state.grants: + state.receipts.append((grant, 401)) + await send({"type": "http.response.start", "status": 401, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + return + status: dict[str, int] = {} + + async def recording(message: Any) -> None: + if message["type"] == "http.response.start": + status["code"] = message["status"] + await send(message) + + await inner(scope, receive, recording) + state.receipts.append((grant, status.get("code", 0))) + + server = uvicorn.Server( + uvicorn.Config( + app, host="127.0.0.1", port=0, log_level="warning", lifespan="on" + ) + ) + serving = asyncio.create_task(server.serve()) + try: + while not server.started: + if serving.done(): + await serving + await asyncio.sleep(0.01) + port = server.servers[0].sockets[0].getsockname()[1] + state.url = f"http://127.0.0.1:{port}/mcp" + yield state + finally: + server.should_exit = True + await serving diff --git a/tests/adapters/claude_sdk/fakecli.py b/tests/adapters/claude_sdk/fakecli.py index 7e851c9d6..3c2cd63cc 100644 --- a/tests/adapters/claude_sdk/fakecli.py +++ b/tests/adapters/claude_sdk/fakecli.py @@ -17,6 +17,7 @@ from collections.abc import AsyncIterator, Callable, Iterator, Sequence from contextlib import AbstractAsyncContextManager, contextmanager from dataclasses import dataclass +from pathlib import Path from typing import Any from unittest.mock import patch @@ -158,6 +159,11 @@ def __init__(self) -> None: # Parks the next close; the session that reaches it takes it. self.closing: Hold | None = None self.refuse_close = False + # Scripted status replies repeat the last answer, as the native CLI does. + self.mcp_statuses: deque[list[dict[str, Any]]] = deque() + self.status_requests = 0 + self.status_hold: Hold | None = None + self.status_unsupported = False def script(self, *turns: Turn) -> None: self._turns.extend(turns) @@ -251,6 +257,10 @@ async def close(self) -> None: self.alive = False if self._turn is not None: self._turn.cancel() + for task in tuple(self._background): + task.cancel() + if self._background: + await asyncio.wait(self._background) self._outbox.put_nowait(None) def die(self) -> None: @@ -283,6 +293,11 @@ def _answer_sdk(self, message: dict[str, Any]) -> None: request = message["request"] if request["subtype"] == "initialize": self._hooks = request.get("hooks") or {} + if request["subtype"] == "mcp_status": + task = asyncio.create_task(self._answer_status(message["request_id"])) + self._background.add(task) + task.add_done_callback(self._background.discard) + return self._emit( type="control_response", response={ @@ -532,12 +547,66 @@ async def _execute(self, call: ToolCall) -> tuple[Any, bool]: content = [block.model_dump(exclude_none=True) for block in result.content] return content, result.isError + @property + def configured_servers(self) -> dict[str, Any]: + """The configured servers, read from the private file as the CLI reads it.""" + servers = self.options.mcp_servers + if isinstance(servers, dict): + return servers + return json.loads(Path(str(servers)).read_text())["mcpServers"] + def _mcp(self, server: str) -> AbstractAsyncContextManager[ClientSession]: """A session to the MCP server the options name, dialed by its URL.""" - servers = self.options.mcp_servers - assert isinstance(servers, dict), servers - config = servers[server] - return mcp_session(config["url"], BandMCPTransport(config["type"])) + config = self.configured_servers[server] + return mcp_session( + config["url"], BandMCPTransport(config["type"]), config.get("headers") + ) + + async def _mcp_status(self) -> dict[str, Any]: + """Report real handshakes, or scripted status with native credential fields.""" + self.claude.status_requests += 1 + if self.claude.mcp_statuses: + statuses = self.claude.mcp_statuses + entries = statuses.popleft() if len(statuses) > 1 else statuses[0] + return {"mcpServers": entries} + entries = [] + for name, config in self.configured_servers.items(): + entry: dict[str, Any] = {"name": name, "config": config} + if config.get("type") in ("http", "sse"): + try: + async with asyncio.timeout(5), self._mcp(name): + pass + entry["status"] = "connected" + except Exception as error: # noqa: BLE001 -- the CLI reports any handshake failure + entry["status"] = "failed" + entry["error"] = f"{type(error).__name__}: {config}" + else: + entry["status"] = "pending" + entries.append(entry) + return {"mcpServers": entries} + + async def _answer_status(self, request_id: str) -> None: + if (hold := self.claude.status_hold) is not None: + hold.reached.set() + await hold.released.wait() + if self.claude.status_unsupported: + self._emit( + type="control_response", + response={ + "subtype": "error", + "request_id": request_id, + "error": "Unsupported control request subtype: mcp_status", + }, + ) + return + self._emit( + type="control_response", + response={ + "subtype": "success", + "request_id": request_id, + "response": await self._mcp_status(), + }, + ) def _tool_result(self, tool_use_id: str, content: Any, *, is_error: bool) -> None: self._emit( diff --git a/tests/adapters/claude_sdk/helpers.py b/tests/adapters/claude_sdk/helpers.py index 27bd0e697..9ad47fdd9 100644 --- a/tests/adapters/claude_sdk/helpers.py +++ b/tests/adapters/claude_sdk/helpers.py @@ -6,11 +6,16 @@ import copy import itertools import json -from collections.abc import Awaitable, Callable +import shutil +from collections.abc import Awaitable, Callable, Iterator +from contextlib import contextmanager from datetime import UTC, datetime +from pathlib import Path from typing import Any from urllib.parse import urlsplit +import pytest +from claude_agent_sdk import ClaudeSDKClient from claude_agent_sdk.types import Message from band.adapters.claude_sdk import ( @@ -61,6 +66,35 @@ NO_SETTLE_GRACE = ClaudeSDKAdapterConfig(turn_settle_timeout_s=0) +@contextmanager +def refuse_mcp_removal(path: Path) -> Iterator[None]: + """Refuse deletion before any private file is removed.""" + remove = shutil.rmtree + + def refuse(directory: str | Path, *args: Any, **kwargs: Any) -> None: + if Path(directory) == path.parent: + raise PermissionError("private MCP removal refused") + remove(directory, *args, **kwargs) + + with pytest.MonkeyPatch.context() as patch: + patch.setattr(shutil, "rmtree", refuse) + yield + + +async def lease_client( + manager: ClaudeSessionManager, + room_id: str, + *, + resume_session_id: str | None = None, +) -> ClaudeSDKClient: + """Acquire a resident client through the current retained-room contract.""" + session = manager.room(room_id) + session.retire_outdated() + session.adopt(resume_session_id) + async with session.using() as client: + return client + + async def exchange(session: ClaudeRoomSession, prompt: str) -> list[Message]: """One leased query through the room's single reader, up to its result.""" async with session.using() as client: diff --git a/tests/adapters/claude_sdk/lifecycle/test_external_cancellation.py b/tests/adapters/claude_sdk/lifecycle/test_external_cancellation.py new file mode 100644 index 000000000..df5c8b3bb --- /dev/null +++ b/tests/adapters/claude_sdk/lifecycle/test_external_cancellation.py @@ -0,0 +1,141 @@ +"""External worker cancellation keeps Claude and runtime custody through settlement.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable, Mapping +from datetime import UTC, datetime +from typing import Any + +import pytest + +from band.adapters.claude_sdk import ClaudeSDKAdapterConfig +from band.core.exceptions import CleanupIncompleteError +from band.core.turn_lifecycle import CYCLE_CANCEL_GRACE_SECONDS +from band.core.types import AgentInput, HistoryProvider, PlatformMessage +from band.platform.link import BandLink +from band.runtime.execution import ExecutionContext +from band.runtime.tools import AgentTools +from tests.adapters.claude_sdk.fakecli import Hold +from tests.adapters.claude_sdk.helpers import REPORTING_BUDGET, ClaudeRoom +from tests.conftest import make_message_event +from tests.runtime.helpers import LifecyclePlatform, rest_client_over + +OpenRoom = Callable[..., Awaitable[ClaudeRoom]] + + +class TestExternalCancellation: + @pytest.mark.parametrize("cleanup", ["room", "adapter"]) + async def test_external_worker_cancel_keeps_handler_and_owner_until_resolver_settles( + self, claude_room: OpenRoom, monkeypatch: pytest.MonkeyPatch, cleanup: str + ) -> None: + resolving, settling = Hold(), Hold() + native_constructions: list[tuple[Any, ...]] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + resolving.reached.set() + try: + await resolving.released.wait() + except asyncio.CancelledError: + settling.reached.set() + await settling.released.wait() + raise + return {} + + def refuse_native_client(*args: Any, **kwargs: Any) -> None: + native_constructions.append(args) + raise AssertionError( + "retired resolution must not construct a native client" + ) + + monkeypatch.setattr( + "band.integrations.claude_sdk.session_manager.ClaudeSDKClient", + refuse_native_client, + ) + room = await claude_room( + ClaudeSDKAdapterConfig(retention=REPORTING_BUDGET, turn_settle_timeout_s=0), + resolve_mcp_servers=resolver, + ) + manager = room.adapter._session_manager + assert manager is not None + peer = LifecyclePlatform() + async with rest_client_over(peer.answer) as rest: + link = BandLink(agent_id="agent-123", api_key="test") + link.rest = rest + + async def handle(ctx: ExecutionContext, event: Any) -> None: + await room.adapter.on_event( + AgentInput( + msg=PlatformMessage( + id=event.payload.id, + room_id=room.room_id, + content="start", + sender_id="user-456", + sender_type="User", + sender_name="User", + message_type="text", + metadata={}, + created_at=datetime.now(UTC), + ), + tools=AgentTools.from_context(ctx), + history=HistoryProvider(raw=[]), + participants_msg=None, + contacts_msg=None, + is_session_bootstrap=True, + room_id=room.room_id, + ) + ) + + ctx = ExecutionContext(room.room_id, link, handle) + scope = await ctx._begin_scope() + assert scope is not None + running = asyncio.create_task( + ctx._run_cycle(scope, make_message_event(room_id=room.room_id), None) + ) + cleaning: asyncio.Task[None] | None = None + try: + await asyncio.wait_for(resolving.reached.wait(), timeout=5.0) + owner = room.adapter._turn_tasks[room.room_id] + handler = scope.handler + assert handler is not None and owner in scope.workers + owner.task.cancel() + await asyncio.wait_for(settling.reached.wait(), timeout=5.0) + cleaning = asyncio.create_task( + room.leave() if cleanup == "room" else room.adapter.cleanup_all() + ) + await asyncio.wait({cleaning}, timeout=2 * CYCLE_CANCEL_GRACE_SECONDS) + assert owner.task.cancelling() == 1 + assert room.adapter._turn_tasks.get(room.room_id) is owner + assert owner in scope.workers and not owner.task.done() + assert ( + ctx._active_cycle_task is handler.task and not handler.task.done() + ) + assert not scope.settled + assert manager._rooms + assert not settling.released.is_set() + assert native_constructions == [] + assert room.claude.sessions == [] + assert cleaning.done(), "cleanup must report bounded incompleteness" + assert isinstance(cleaning.exception(), CleanupIncompleteError) + finally: + resolving.released.set() + settling.released.set() + await asyncio.gather(running, return_exceptions=True) + if cleaning is not None: + await asyncio.gather(cleaning, return_exceptions=True) + await room.leave() + await room.adapter.cleanup_all() + + assert owner.task.cancelling() == 1 + assert owner.task.done() and handler.task.done() + assert scope.settled and ctx._active_cycle_task is None + assert ctx._detached_cycle_tasks == {} + assert not manager._rooms + assert room.adapter._turn_tasks == {} + assert room.adapter._turn_release == {} + assert room.adapter._room_tools == {} + assert room.adapter._mcp_room_tools == {} + assert room.adapter._mcp._closing is None + assert room.adapter._mcp.current is None + assert native_constructions == [] + assert room.claude.sessions == [] diff --git a/tests/adapters/claude_sdk/lifecycle/test_startup_owner.py b/tests/adapters/claude_sdk/lifecycle/test_startup_owner.py new file mode 100644 index 000000000..a65e559cf --- /dev/null +++ b/tests/adapters/claude_sdk/lifecycle/test_startup_owner.py @@ -0,0 +1,489 @@ +"""The room's shared turn owner covers asynchronous Claude preparation too.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable, Mapping +from contextlib import nullcontext +from pathlib import Path +from typing import Any + +import pytest + +from band.adapters.claude_sdk import ClaudeSDKAdapterConfig +from band.client.streaming import ControlMode +from band.core.exceptions import CleanupIncompleteError, InputNotAccepted +from band.core.turn_lifecycle import TurnTask +from band.integrations.claude_sdk.mcp_startup import ClaudeMcpStartupError +from tests.adapters.claude_sdk.fakecli import Hold +from tests.adapters.claude_sdk.helpers import REPORTING_BUDGET, ClaudeRoom +from tests.conftest import eager_tasks, requires_eager_tasks + +OpenRoom = Callable[..., Awaitable[ClaudeRoom]] + + +class TestStartupOwner: + @pytest.mark.looptime + async def test_the_provider_deadline_starts_after_preparation( + self, claude_room: OpenRoom + ) -> None: + hold = Hold() + + async def resolver(workspace: str) -> Mapping[str, Any]: + hold.reached.set() + await hold.released.wait() + return {} + + room = await claude_room( + ClaudeSDKAdapterConfig(turn_timeout_s=2.0), + resolve_mcp_servers=resolver, + ) + room.claude.script([room.model_reply("Ready.")]) + sending = room.send_in_background("hello") + try: + await hold.reached.wait() + await asyncio.wait({sending}, timeout=4.0) + assert not sending.done() + assert room.claude.prompts == [] + finally: + hold.released.set() + await asyncio.gather(sending, return_exceptions=True) + await sending + assert room.chat == ["Ready."] + assert room.failures == [] + + async def test_a_resolving_turn_refuses_overlap_without_rebinding_tools( + self, claude_room: OpenRoom + ) -> None: + hold = Hold() + calls: list[str] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + calls.append(workspace) + hold.reached.set() + await hold.released.wait() + return {} + + room = await claude_room( + ClaudeSDKAdapterConfig(retention=REPORTING_BUDGET, turn_settle_timeout_s=0), + resolve_mcp_servers=resolver, + ) + accepted, busy = room.fresh_tools(), room.fresh_tools() + room.claude.script([room.model_reply("First.")]) + sending = asyncio.create_task(room.send("first", tools=accepted)) + try: + await hold.reached.wait() + with pytest.raises(InputNotAccepted): + await room.send("overlap", tools=busy) + assert not busy.turn.complete + assert room.claude.prompts == [] + finally: + hold.released.set() + await sending + + assert accepted.turn.replied + assert not busy.turn.replied + assert len(calls) == 1 + assert room.chat == ["First."] + room.claude.script([room.model_reply("Second.")]) + await room.send("again", tools=busy) + assert busy.turn.replied + assert len(calls) == 1 + assert len(room.claude.sessions) == 1 + + @pytest.mark.parametrize("phase", ["resolution", "readiness"]) + async def test_stop_retains_startup_custody_until_settlement( + self, claude_room: OpenRoom, phase: str + ) -> None: + hold, closing = Hold(), Hold() + cancelled = asyncio.Event() + cancellations: list[str] = [] + calls: list[str] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + calls.append(workspace) + if phase == "resolution" and len(calls) == 1: + hold.reached.set() + try: + await hold.released.wait() + except asyncio.CancelledError: + cancellations.append(workspace) + cancelled.set() + await hold.released.wait() + return {} + + room = await claude_room( + ClaudeSDKAdapterConfig(retention=REPORTING_BUDGET, turn_settle_timeout_s=0), + resolve_mcp_servers=resolver, + ) + if phase == "readiness": + room.claude.status_hold = hold + room.claude.closing = closing + sending = room.send_in_background("start") + stopping: asyncio.Task[None] | None = None + config_file: Path | None = None + try: + await hold.reached.wait() + if phase == "readiness": + config_file = Path(str(room.claude.sessions[0].options.mcp_servers)) + assert config_file.exists() + stopping = asyncio.create_task( + room.adapter.on_interrupt(room.room_id, ControlMode.STOP) + ) + if phase == "resolution": + await cancelled.wait() + else: + await closing.reached.wait() + await asyncio.gather( + room.adapter.on_interrupt(room.room_id, ControlMode.STOP), + room.adapter.on_interrupt(room.room_id, ControlMode.STOP), + ) + with pytest.raises(CleanupIncompleteError): + await room.leave() + with pytest.raises(InputNotAccepted): + await room.send("overlap") + manager = room.adapter._session_manager + assert manager is not None and bool(manager._rooms) + assert room.claude.prompts == [] + if phase == "resolution": + assert len(cancellations) == 1 + assert room.claude.sessions == [] + else: + assert room.claude.sessions[0].alive + assert config_file is not None and config_file.exists() + finally: + hold.released.set() + closing.released.set() + room.claude.status_hold = None + room.claude.closing = None + if stopping is not None: + await stopping + await asyncio.gather(sending, return_exceptions=True) + + with pytest.raises(asyncio.CancelledError): + await sending + await room.settled() + await room.leave() + assert manager is not None and not bool(manager._rooms) + if config_file is not None: + assert not config_file.exists() + assert not room.claude.sessions[0].alive + assert room.failures == [] + room.claude.script([room.model_reply("Fresh.")]) + await room.send("again") + assert len(calls) == 2 + assert room.chat == ["Fresh."] + + async def test_a_held_room_retirement_does_not_block_another_room( + self, claude_room: OpenRoom + ) -> None: + hold = Hold() + cancelled = asyncio.Event() + calls: list[str] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + calls.append(workspace) + if workspace.endswith("room-1") and not hold.released.is_set(): + hold.reached.set() + try: + await hold.released.wait() + except asyncio.CancelledError: + cancelled.set() + await hold.released.wait() + return {} + + room = await claude_room( + ClaudeSDKAdapterConfig(retention=REPORTING_BUDGET, turn_settle_timeout_s=0), + resolve_mcp_servers=resolver, + ) + sending = room.send_in_background("start") + leaving: asyncio.Task[None] | None = None + try: + await hold.reached.wait() + leaving = asyncio.create_task(room.leave()) + await cancelled.wait() + other = room.beside("room-2") + room.claude.script([other.model_reply("Independent.")]) + await other.send("hello") + with pytest.raises(InputNotAccepted): + await room.send("overlap") + assert room.chat == [] + assert other.chat == ["Independent."] + finally: + hold.released.set() + await asyncio.gather(sending, return_exceptions=True) + if leaving is not None: + try: + await leaving + except CleanupIncompleteError: + pass + + with pytest.raises(asyncio.CancelledError): + await sending + await room.settled() + await room.leave() + room.claude.script([room.model_reply("Rejoined.")]) + await room.send("again") + assert len(calls) == 3 + assert room.chat == ["Rejoined."] + assert room.failures == [] + + async def test_registration_failure_launches_nothing_and_leaves_no_worker( + self, claude_room: OpenRoom, monkeypatch: pytest.MonkeyPatch + ) -> None: + calls: list[str] = [] + workers: list[asyncio.Task[None]] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + calls.append(workspace) + return {} + + def refuse_registration(owner: TurnTask) -> None: + workers.append(owner.task) + raise RuntimeError("turn registration refused") + + room = await claude_room( + ClaudeSDKAdapterConfig(retention=REPORTING_BUDGET, turn_settle_timeout_s=0), + resolve_mcp_servers=resolver, + ) + refused = room.fresh_tools() + monkeypatch.setattr( + refused, "register_turn_task", refuse_registration, raising=False + ) + with pytest.raises(RuntimeError, match="turn registration refused"): + await room.send("refused", tools=refused) + + assert workers and all(worker.done() for worker in workers) + assert not room.adapter._turn_running(room.room_id) + assert room.claude.sessions == [] + assert calls == [] + assert room.chat == [] + assert room.failures == [] + room.claude.script([room.model_reply("Accepted.")]) + await room.send("again") + assert len(calls) == 1 + assert room.chat == ["Accepted."] + + @requires_eager_tasks + @pytest.mark.parametrize("cooperative", [False, True]) + async def test_a_scope_stopped_at_registration_never_starts_preparation( + self, + claude_room: OpenRoom, + monkeypatch: pytest.MonkeyPatch, + cooperative: bool, + ) -> None: + calls: list[str] = [] + workers: list[asyncio.Task[None]] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + calls.append(workspace) + return {} + + def stopped_registration(owner: TurnTask) -> None: + workers.append(owner.task) + owner.terminate(cooperative=cooperative) + + room = await claude_room( + ClaudeSDKAdapterConfig(retention=REPORTING_BUDGET, turn_settle_timeout_s=0), + resolve_mcp_servers=resolver, + ) + stopped = room.fresh_tools() + monkeypatch.setattr( + stopped, "register_turn_task", stopped_registration, raising=False + ) + with eager_tasks(), pytest.raises(asyncio.CancelledError): + await room.send("stopped", tools=stopped) + + assert workers and all(worker.done() for worker in workers) + assert not room.adapter._turn_running(room.room_id) + assert room.claude.sessions == [] + assert calls == [] + assert room.failures == [] + room.claude.script([room.model_reply("Accepted.")]) + await room.send("again") + assert room.chat == ["Accepted."] + + @requires_eager_tasks + async def test_resolver_preparation_runs_under_the_eager_bound_owner( + self, claude_room: OpenRoom + ) -> None: + calls: list[str] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + calls.append(workspace) + return {} + + room = await claude_room( + ClaudeSDKAdapterConfig(retention=REPORTING_BUDGET, turn_settle_timeout_s=0), + resolve_mcp_servers=resolver, + ) + room.claude.script([room.model_reply("Ready.")]) + with eager_tasks(): + await room.send("hello") + + assert len(calls) == 1 + assert room.chat == ["Ready."] + assert room.failures == [] + + @pytest.mark.parametrize( + "eager", + [ + pytest.param(False, id="default"), + pytest.param(True, marks=requires_eager_tasks, id="eager"), + ], + ) + async def test_a_resolver_can_clean_up_its_own_adapter_room( + self, claude_room: OpenRoom, tmp_path: Path, eager: bool + ) -> None: + room_ref: list[ClaudeRoom] = [] + returned = asyncio.Event() + calls: list[str] = [] + workspace = tmp_path / "shared" + marker = workspace / "keep.txt" + + async def resolver(path: str) -> Mapping[str, Any]: + calls.append(path) + if len(calls) == 1: + room = room_ref[0] + manager = room.adapter._session_manager + assert manager is not None and bool(manager._rooms) + marker.write_text("membership files stay", encoding="utf-8") + await room.adapter.on_cleanup(room.room_id) + returned.set() + return {} + + room = await claude_room( + ClaudeSDKAdapterConfig(retention=REPORTING_BUDGET, turn_settle_timeout_s=0), + workspace_for_room=lambda room_id: str(workspace), + resolve_mcp_servers=resolver, + ) + room_ref.append(room) + with eager_tasks() if eager else nullcontext(): + async with asyncio.timeout(5): + [outcome] = await asyncio.gather( + room.send("retire this startup"), return_exceptions=True + ) + + assert returned.is_set(), "self-cleanup must return to its resolver" + assert isinstance(outcome, (ClaudeMcpStartupError, asyncio.CancelledError)) + if isinstance(outcome, ClaudeMcpStartupError): + assert outcome.stage == "retirement" + assert room.claude.sessions == [] + assert room.claude.prompts == [] + assert room.chat == [] + assert room.failures == [] + await room.leave() + manager = room.adapter._session_manager + assert manager is not None and not bool(manager._rooms) + assert marker.read_text(encoding="utf-8") == "membership files stay" + + other = room.beside("room-2") + room.claude.script([other.model_reply("Reclaimed.")]) + await other.send("use the released workspace") + assert calls == [str(workspace), str(workspace)] + assert room.chat == [] + assert other.chat == ["Reclaimed."] + + async def test_a_resolver_closes_a_foreign_room_without_deferring_its_turn( + self, claude_room: OpenRoom + ) -> None: + room_ref: list[ClaudeRoom] = [] + returned = asyncio.Event() + + async def resolver(workspace: str) -> Mapping[str, Any]: + if workspace.endswith("room-1"): + await room_ref[0].adapter.on_cleanup("room-2") + returned.set() + return {} + + room = await claude_room( + ClaudeSDKAdapterConfig(retention=REPORTING_BUDGET, turn_settle_timeout_s=0), + resolve_mcp_servers=resolver, + ) + room_ref.append(room) + other = room.beside("room-2") + work = Hold() + room.claude.script( + [work], + [room.model_reply("Current.")], + [other.model_reply("Fresh.")], + ) + sending = other.send_in_background("hold this turn") + try: + await work.reached.wait() + await room.send("clean up the other room") + finally: + work.released.set() + await asyncio.gather(sending, return_exceptions=True) + + with pytest.raises(asyncio.CancelledError): + await sending + assert returned.is_set() + assert room.chat == ["Current."] + assert other.chat == [] + assert room.failures == [] + assert other.failures == [] + assert not room.claude.sessions[0].alive + await other.send("rejoin") + assert other.chat == ["Fresh."] + assert room.claude.resumed == [None, None, None] + + @pytest.mark.parametrize("stop_after_cleanup", [False, True]) + async def test_self_cleanup_keeps_ownership_while_its_resolver_settles( + self, claude_room: OpenRoom, tmp_path: Path, stop_after_cleanup: bool + ) -> None: + room_ref: list[ClaudeRoom] = [] + hold = Hold() + calls: list[str] = [] + workspace = tmp_path / "shared" + + async def resolver(path: str) -> Mapping[str, Any]: + calls.append(path) + if len(calls) == 1: + room = room_ref[0] + await room.adapter.on_cleanup(room.room_id) + hold.reached.set() + await hold.released.wait() + return {} + + room = await claude_room( + ClaudeSDKAdapterConfig(retention=REPORTING_BUDGET, turn_settle_timeout_s=0), + workspace_for_room=lambda room_id: str(workspace), + resolve_mcp_servers=resolver, + ) + room_ref.append(room) + manager = room.adapter._session_manager + assert manager is not None + sending = room.send_in_background("retire this startup") + try: + await hold.reached.wait() + assert bool(manager._rooms) + assert room.adapter._turn_running(room.room_id) + with pytest.raises(InputNotAccepted): + await room.send("overlap") + if stop_after_cleanup: + await room.adapter.on_interrupt(room.room_id, ControlMode.STOP) + assert room.adapter._turn_running(room.room_id) + assert not sending.done() + with pytest.raises(CleanupIncompleteError): + await room.leave() + assert room.claude.sessions == [] + assert room.claude.prompts == [] + finally: + hold.released.set() + [outcome] = await asyncio.gather(sending, return_exceptions=True) + + assert isinstance(outcome, (ClaudeMcpStartupError, asyncio.CancelledError)) + if isinstance(outcome, ClaudeMcpStartupError): + assert outcome.stage == "retirement" + assert room.failures == [] + await room.leave() + await room.adapter.backend_turns.settled(room.room_id) + assert not bool(manager._rooms) + assert not room.adapter._turn_running(room.room_id) + other = room.beside("room-2") + room.claude.script([other.model_reply("Reclaimed.")]) + await other.send("use the released workspace") + assert calls == [str(workspace), str(workspace)] + assert room.chat == [] + assert other.chat == ["Reclaimed."] diff --git a/tests/adapters/claude_sdk/lifecycle/test_teardown_custody.py b/tests/adapters/claude_sdk/lifecycle/test_teardown_custody.py new file mode 100644 index 000000000..7c5498d0a --- /dev/null +++ b/tests/adapters/claude_sdk/lifecycle/test_teardown_custody.py @@ -0,0 +1,104 @@ +"""The adapter's teardown is one durable sequence that keeps the Band MCP +server until every Claude process it may serve is confirmed closed.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from typing import Any + +import pytest + +from tests.adapters.claude_sdk.fakecli import FakeClaude, Hold +from tests.adapters.claude_sdk.helpers import ClaudeRoom + +OpenRoom = Callable[..., Awaitable[ClaudeRoom]] + + +def hold_stop(room: ClaudeRoom) -> tuple[Hold, list[str]]: + backend = room.adapter._mcp.current + assert backend is not None + stop = backend.stop + hold, stops = Hold(), [] + + async def held() -> None: + hold.reached.set() + await hold.released.wait() + await stop() + stops.append("stopped") + + object.__setattr__(backend, "stop", held) + return hold, stops + + +class TestTeardownCustody: + async def test_a_cancelled_teardown_caller_does_not_abandon_the_band_server( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + room = await claude_room() + claude.script([room.model_reply("Hi.")]) + await room.send("hi") + hold, stops = hold_stop(room) + + caller = asyncio.create_task(room.adapter.cleanup_all()) + try: + await hold.reached.wait() + caller.cancel() + await asyncio.gather(caller, return_exceptions=True) + finally: + hold.released.set() + await room.adapter.cleanup_all() + + assert stops == ["stopped"] + assert room.adapter._mcp._closing is None + assert not claude.sessions[0].alive + + async def test_a_session_that_will_not_close_keeps_the_band_server_until_retry( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + room = await claude_room() + claude.script([room.model_reply("Hi.")]) + await room.send("hi") + backend = room.adapter._mcp.current + claude.refuse_close = True + + with pytest.raises(ExceptionGroup) as caught: + await room.adapter.cleanup_all() + assert caught.value.subgroup(RuntimeError) is not None + # The CLI may still dial it: the Band server keeps running. + assert room.adapter._mcp.current is backend + assert backend is not None and backend.is_running + + claude.refuse_close = False + await room.adapter.cleanup_all() + assert room.adapter._mcp.current is None + assert not backend.is_running + + async def test_a_failed_band_stop_is_retried_before_a_restart( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + room = await claude_room() + backend = room.adapter._mcp.current + assert backend is not None + stop = backend.stop + attempts: list[Any] = [] + + async def failing_once() -> None: + attempts.append(None) + if len(attempts) == 1: + raise RuntimeError("stop failed") + await stop() + + object.__setattr__(backend, "stop", failing_once) + with pytest.raises(RuntimeError, match="stop failed"): + await room.adapter.cleanup_all() + assert room.adapter._mcp._closing is backend + + await room.adapter.on_started("Test Agent", "restarted") + + assert len(attempts) == 2 + assert room.adapter._mcp._closing is None + assert not backend.is_running + claude.script([room.model_reply("Back.")]) + await room.send("hi") + assert room.chat == ["Back."] diff --git a/tests/adapters/claude_sdk/mcpkit.py b/tests/adapters/claude_sdk/mcpkit.py new file mode 100644 index 000000000..e04cf59ad --- /dev/null +++ b/tests/adapters/claude_sdk/mcpkit.py @@ -0,0 +1,365 @@ +"""Connection snapshots, hostile callback values and retained startup doubles.""" + +from __future__ import annotations + +import asyncio +import tempfile +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import pytest +from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient + +from band.core.exceptions import CleanupIncompleteError +from band.integrations.claude_sdk.mcp_startup import startup_error +from band.integrations.claude_sdk.session_manager import ClaudeSessionManager +from band.runtime.tools import BAND_MCP_SERVER_NAME +from band.workspaces import RoomWorkspaces +from tests.adapters.claude_sdk.fakecli import FakeClaude, FakeCLISession +from tests.adapters.claude_sdk.helpers import REPORTING_BUDGET, lease_client + +GRANT_A = "grant-a-5ecret" +GRANT_B = "grant-b-5ecret" + + +class Resolver: + """Returns ``servers`` and records each workspace it was asked for.""" + + def __init__(self, servers: Mapping[str, Any]) -> None: + self.servers = servers + self.calls: list[str] = [] + + async def __call__(self, workspace: str) -> Mapping[str, Any]: + self.calls.append(workspace) + return self.servers + + +def launch_file(session: FakeCLISession) -> Path: + servers = session.options.mcp_servers + assert isinstance(servers, str), "a resolver's servers go through a file" + return Path(servers) + + +def connected(*names: str) -> list[dict[str, Any]]: + return [{"name": name, "status": "connected"} for name in names] + + +def assert_secret_free(secret: str, error: BaseException, logs: str) -> None: + rendered = [str(error), repr(error), logs] + chained: BaseException | None = error + while chained is not None: + rendered.append(str(chained)) + chained = chained.__cause__ or ( + None if chained.__suppress_context__ else chained.__context__ + ) + assert not any(secret in text for text in rendered) + + +class ExplodingServers(Mapping[str, Any]): + """A mapping that fails while being read, with a credential in its error.""" + + def __getitem__(self, key: str) -> Any: + raise LookupError(f"vault said {GRANT_A}") + + def __iter__(self) -> Any: + raise LookupError(f"vault said {GRANT_A}") + + def __len__(self) -> int: + return 1 + + +def private_files() -> set[Path]: + return set(Path(tempfile.gettempdir()).glob("band-claude-mcp-*/mcp.json")) + + +class MarkedServers(Mapping[str, Any]): + """A lazy mapping raising the module's own safe-looking error type, with a + credential in its reason and resume eligibility set.""" + + def __getitem__(self, key: str) -> Any: + raise startup_error("launch", f"token {GRANT_A}", resume_eligible=True) + + def __iter__(self) -> Any: + raise startup_error("launch", f"token {GRANT_A}", resume_eligible=True) + + def __len__(self) -> int: + return 1 + + +def item_failure(kind: str) -> BaseException: + if kind == "ordinary": + return LookupError(f"vault said {GRANT_A}") + if kind == "marked": + return startup_error("launch", f"token {GRANT_A}", resume_eligible=True) + # A callback-defined exception whose very class name carries the secret. + return type(f"Leak_{GRANT_A.replace('-', '_')}", (Exception,), {})() + + +class BadItem(tuple): + """A ``(name, config)`` item that fails when unpacked.""" + + failure: BaseException + + def __iter__(self) -> Any: + raise self.failure + + +class HostileItems(Mapping[str, Any]): + """A mapping whose items fail only once a caller unpacks them.""" + + def __init__(self, failure: BaseException) -> None: + self.failure = failure + + def __getitem__(self, key: str) -> Any: + return {"type": "stdio", "command": "tool"} + + def __iter__(self) -> Any: + return iter(["tool"]) + + def __len__(self) -> int: + return 1 + + def items(self) -> Any: + item = BadItem(("tool", {"type": "stdio", "command": "tool"})) + item.failure = self.failure + return [item] + + +class _ExplodingMeta(type): + """A metaclass whose class metadata raises a secret-bearing error.""" + + def __getattribute__(cls, name: str) -> Any: + if name in {"__module__", "__name__", "__qualname__"}: + raise LookupError(f"meta said {GRANT_A}") + return super().__getattribute__(name) + + +def metadata_attack(kind: str) -> BaseException: + secret_name = f"Leak_{GRANT_A.replace('-', '_')}" + if kind == "spoofed-module": + cls = type(secret_name, (Exception,), {"__module__": "builtins"}) + elif kind == "non-string-module": + cls = type(secret_name, (Exception,), {"__module__": 3}) + else: + cls = _ExplodingMeta(secret_name, (Exception,), {}) + return cls() + + +BAND = {BAND_MCP_SERVER_NAME: {"type": "http", "url": "http://127.0.0.1:9/band"}} +SECRET = "close-5ecret" + + +class HeldResolver: + """Resolves only when released; records calls and cancellations. With + ``swallow`` it finishes its own work despite a cancellation and returns.""" + + def __init__(self, *, swallow: bool = False) -> None: + self.entered = asyncio.Event() + self.release = asyncio.Event() + self.swallow = swallow + self.calls: list[str] = [] + self.cancellations = 0 + self.settled = asyncio.Event() + + async def __call__(self, workspace: str) -> Mapping[str, Any]: + self.calls.append(workspace) + self.entered.set() + try: + while True: + try: + await self.release.wait() + break + except asyncio.CancelledError: + self.cancellations += 1 + if not self.swallow: + raise + return {} + finally: + self.settled.set() + + +def manager_for( + tmp_path: Path, + resolver: Any, + claude: FakeClaude, + band: dict[str, Any] | None = None, +) -> ClaudeSessionManager: + """A manager whose protected Band server is ``band`` (changed in place to + model a replaced Band server).""" + protected = BAND if band is None else band + claude.mcp_statuses.append([{"name": BAND_MCP_SERVER_NAME, "status": "connected"}]) + return ClaudeSessionManager( + ClaudeAgentOptions(), + mcp_servers_factory=lambda room_id: dict(protected), + workspaces=RoomWorkspaces(lambda room_id: str(tmp_path / room_id)), + resolve_mcp_servers=resolver, + retention=REPORTING_BUDGET, + settle_grace_s=0, + ) + + +class SettlingResolver: + """Already being cancelled (by itself or its own room cleanup) when it + enters a settlement it must finish; logs any further cancellation.""" + + def __init__(self, *, retire_own_room: bool = False) -> None: + self.retire_own_room = retire_own_room + self.manager: ClaudeSessionManager | None = None + self.settling = asyncio.Event() + self.release = asyncio.Event() + self.log: list[str] = [] + + async def __call__(self, workspace: str) -> Mapping[str, Any]: + if self.retire_own_room: + assert self.manager is not None + with pytest.raises(CleanupIncompleteError): + await self.manager.cleanup_session("room-a") + else: + task = asyncio.current_task() + assert task is not None + task.cancel() + try: + await asyncio.sleep(0) + except asyncio.CancelledError: + pass # caught, never uncancelled + self.settling.set() + try: + await self.release.wait() + except asyncio.CancelledError: + self.log.append("settlement interrupted") + raise + self.log.append("settled") + return {} + + +class CountingResolver: + def __init__(self) -> None: + self.calls: list[str] = [] + + async def __call__(self, workspace: str) -> Mapping[str, Any]: + self.calls.append(workspace) + return {} + + +async def outdated_session( + tmp_path: Path, claude: FakeClaude +) -> tuple[ClaudeSessionManager, CountingResolver, dict[str, Any]]: + """A manager whose room-a session (A) dials a Band server since replaced.""" + resolver = CountingResolver() + band = dict(BAND) + manager = manager_for(tmp_path, resolver, claude, band) + await lease_client(manager, "room-a") + band[BAND_MCP_SERVER_NAME] = {"type": "http", "url": "http://127.0.0.1:9/band-2"} + return manager, resolver, band + + +class FailingBand: + """A protected-server factory that can be made to fail, with a secret.""" + + def __init__(self) -> None: + self.failing = False + + def __call__(self, room_id: str) -> dict[str, Any]: + if self.failing: + raise LookupError(f"band backend said {SECRET}") + return dict(BAND) + + +class ReuseProbe: + """The servers a held client dials, as a callback would hand them over: the + comparison answers with this callback's own object, whose truth is therefore + consulted wherever the decision is read (an ``and`` chain coerces each + earlier operand and leaves the final one, and the settled result is read + again at the branch). Closing admission from the second read models a + callback acting when the decision is read after the final admission guard; + closing on the first read is the early control. The client itself stays the + manager's real one, so its teardown stays owned.""" + + def __init__(self, manager: ClaudeSessionManager, mode: str) -> None: + self.manager = manager + self.mode = mode + self.reads = 0 + + def _decide(self) -> bool: + self.reads += 1 + closing_now = ( + self.reads == 1 if self.mode == "materialized" else self.reads >= 2 + ) + if closing_now: + self.manager.close_admission() + return True + + def __eq__(self, other: object) -> Any: + return self + + def __bool__(self) -> bool: + return self._decide() + + __hash__ = None # type: ignore[assignment] + + +class HeldServers: + """A stand-in for the held client's options: only its servers are read.""" + + def __init__(self, servers: Any) -> None: + self.mcp_servers = servers + + +class ClientCustody: + """What teardown must leave behind on an owned client: one disconnect, its + Query closed on both sides, and every reference dropped. A manager that + reports nothing held while the client it owned stays connected is the + failure this pins, so the checks read the client, not the manager.""" + + def __init__(self, client: ClaudeSDKClient) -> None: + self.client = client + self.query = client._query + self.disconnects = 0 + self.calls: list[str] = [] + assert self.query is not None, "the client was never connected" + query, real_disconnect = self.query, client.disconnect + real_close, real_close_receive = query.close, query.close_receive_stream + + async def disconnect() -> None: + self.disconnects += 1 + await real_disconnect() + + async def close() -> None: + self.calls.append("close") + await real_close() + + def close_receive() -> None: + self.calls.append("receive") + real_close_receive() + + client.disconnect = disconnect # type: ignore[method-assign] + query.close = close # type: ignore[method-assign] + query.close_receive_stream = close_receive # type: ignore[method-assign] + + def assert_disconnected(self) -> None: + assert self.disconnects == 1, f"disconnect ran {self.disconnects} times" + assert self.calls == ["close", "receive"], self.calls + assert self.query._closed, "the client's Query stayed open" + assert self.client._query is None + assert self.client._transport is None + assert self.client._materialized is None + + +class RaisingTruth: + """A callback's servers whose comparison answers with the callback's own + object and whose truth test raises a secret-bearing error of a builtin + type, as a callback-computed step could.""" + + def __init__(self, text: str) -> None: + self.text = text + self.reads = 0 + + def __eq__(self, other: object) -> Any: + return self + + def __bool__(self) -> bool: + self.reads += 1 + raise LookupError(self.text) + + __hash__ = None # type: ignore[assignment] diff --git a/tests/adapters/claude_sdk/process.py b/tests/adapters/claude_sdk/process.py index 001a995b0..0d240250a 100644 --- a/tests/adapters/claude_sdk/process.py +++ b/tests/adapters/claude_sdk/process.py @@ -101,3 +101,10 @@ def launch(*, prompt: str, options: ClaudeAgentOptions) -> Child: monkeypatch.setattr(transport, "SubprocessCLITransport", launch) return children + + +class ForgetfulProcess(WorkspaceProcess): + """A native close that returns while its child still runs and forgets it.""" + + async def close(self) -> None: + self._process = None diff --git a/tests/adapters/claude_sdk/test_cli_options.py b/tests/adapters/claude_sdk/test_cli_options.py index 258f26cb0..9b3c7305f 100644 --- a/tests/adapters/claude_sdk/test_cli_options.py +++ b/tests/adapters/claude_sdk/test_cli_options.py @@ -22,6 +22,7 @@ ClaudeSDKAdapterConfig, ) from band.core.types import Emit +from band.integrations.claude_sdk import transport as transport_module from band.integrations.claude_sdk.transport import create_transport from tests.adapters.claude_sdk.helpers import ( SEND_MESSAGE_MCP_NAME, @@ -31,21 +32,6 @@ from tests.paths import host_absolute_path -def test_owned_transport_preserves_sdk_permission_routing() -> None: - async def allow( - tool_name: str, tool_input: dict[str, Any], context: ToolPermissionContext - ) -> PermissionResultAllow: - return PermissionResultAllow() - - options = ClaudeAgentOptions(cli_path=sys.executable, can_use_tool=allow) - transport = create_transport(options) - assert isinstance(transport, SubprocessCLITransport) - command = transport._build_command() - assert command[command.index("--permission-prompt-tool") + 1] == "stdio" - assert options.permission_prompt_tool_name is None - assert options.can_use_tool is allow - - async def _cli_options( room: ClaudeRoom, session_id: str | None = None ) -> ClaudeAgentOptions: @@ -70,87 +56,118 @@ def _emitted_flags(options: ClaudeAgentOptions) -> set[str]: } -async def test_cli_options_reach_the_cli(claude_room: OpenRoom, tmp_path: Path) -> None: - cli = ClaudeCLIOptions( - cli_path=host_absolute_path("opt", "claude", "bin", "claude"), - plugin_dirs=(str(tmp_path / "plugin"),), - add_dirs=(str(tmp_path / "shared"),), - env={"ANTHROPIC_API_KEY": "per-agent-key"}, - extra_args={"debug-to-stderr": None}, - ) - room = await claude_room(ClaudeSDKAdapterConfig(cli=cli)) - - options = await _cli_options(room) - - assert options.plugins == [{"type": "local", "path": str(tmp_path / "plugin")}] - assert options.cli_path == cli.cli_path - assert options.env == {"ANTHROPIC_API_KEY": "per-agent-key"} - assert options.add_dirs == [str(tmp_path / "shared")] - assert options.extra_args == {"debug-to-stderr": None} +class TestCliOptions: + def test_owned_transport_preserves_sdk_permission_routing( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """The CLI the adapter spawns is told to route permission prompts to the + SDK, without the caller's options being rewritten.""" + + async def allow( + tool_name: str, tool_input: dict[str, Any], context: ToolPermissionContext + ) -> PermissionResultAllow: + return PermissionResultAllow() + + spawned: list[SubprocessCLITransport] = [] + + def spawn( + *, prompt: str, options: ClaudeAgentOptions + ) -> SubprocessCLITransport: + native = SubprocessCLITransport(prompt=prompt, options=options) + spawned.append(native) + return native + + monkeypatch.setattr(transport_module, "SubprocessCLITransport", spawn) + options = ClaudeAgentOptions(cli_path=sys.executable, can_use_tool=allow) + create_transport(options) + + [native] = spawned + command = native._build_command() + assert command[command.index("--permission-prompt-tool") + 1] == "stdio" + assert options.permission_prompt_tool_name is None + assert options.can_use_tool is allow + + async def test_cli_options_reach_the_cli( + self, claude_room: OpenRoom, tmp_path: Path + ) -> None: + cli = ClaudeCLIOptions( + cli_path=host_absolute_path("opt", "claude", "bin", "claude"), + plugin_dirs=(str(tmp_path / "plugin"),), + add_dirs=(str(tmp_path / "shared"),), + env={"ANTHROPIC_API_KEY": "per-agent-key"}, + extra_args={"debug-to-stderr": None}, + ) + room = await claude_room(ClaudeSDKAdapterConfig(cli=cli)) + options = await _cli_options(room) -async def test_plugins_leave_band_wiring_in_place( - claude_room: OpenRoom, tmp_path: Path -) -> None: - room = await claude_room( - ClaudeSDKAdapterConfig(cli=ClaudeCLIOptions(plugin_dirs=(str(tmp_path),))) - ) + assert options.plugins == [{"type": "local", "path": str(tmp_path / "plugin")}] + assert options.cli_path == cli.cli_path + assert options.env == {"ANTHROPIC_API_KEY": "per-agent-key"} + assert options.add_dirs == [str(tmp_path / "shared")] + assert options.extra_args == {"debug-to-stderr": None} - options = await _cli_options(room) - - assert list(options.mcp_servers) == ["band"] - assert SEND_MESSAGE_MCP_NAME in options.allowed_tools - assert options.setting_sources == [] - - -async def test_sdk_owned_flags_match_what_the_adapter_makes_the_sdk_emit( - claude_room: OpenRoom, tmp_path: Path -) -> None: - room = await claude_room( - ClaudeSDKAdapterConfig( - model="sonnet", - fallback_model="haiku", - effort="high", - max_thinking_tokens=1024, - cwd=tmp_path, - approvals=ClaudeApprovalOptions(mode="manual"), - cli=ClaudeCLIOptions( - plugin_dirs=(str(tmp_path),), add_dirs=(str(tmp_path),) - ), + async def test_plugins_leave_band_wiring_in_place( + self, claude_room: OpenRoom, tmp_path: Path + ) -> None: + room = await claude_room( + ClaudeSDKAdapterConfig(cli=ClaudeCLIOptions(plugin_dirs=(str(tmp_path),))) ) - ) - options = await _cli_options(room, session_id="sess-1") + options = await _cli_options(room) + + assert list(options.mcp_servers) == ["band"] + assert SEND_MESSAGE_MCP_NAME in options.allowed_tools + assert options.setting_sources == [] + + async def test_sdk_owned_flags_match_what_the_adapter_makes_the_sdk_emit( + self, claude_room: OpenRoom, tmp_path: Path + ) -> None: + room = await claude_room( + ClaudeSDKAdapterConfig( + model="sonnet", + fallback_model="haiku", + effort="high", + max_thinking_tokens=1024, + cwd=tmp_path, + approvals=ClaudeApprovalOptions(mode="manual"), + cli=ClaudeCLIOptions( + plugin_dirs=(str(tmp_path),), add_dirs=(str(tmp_path),) + ), + ) + ) - assert _emitted_flags(options) == SDK_OWNED_CLI_FLAGS + options = await _cli_options(room, session_id="sess-1") + assert _emitted_flags(options) == SDK_OWNED_CLI_FLAGS -@pytest.mark.parametrize( - ("emit", "display"), - [(Emit.THOUGHTS, "summarized"), (Emit.TOOL_CALLS, None)], - ids=["thoughts-posted", "thoughts-not-posted"], -) -async def test_thinking_text_is_requested_only_when_thoughts_are_posted( - claude_room: OpenRoom, emit: Emit, display: str | None -) -> None: - room = await claude_room( - ClaudeSDKAdapterConfig(max_thinking_tokens=1024), emit=emit + @pytest.mark.parametrize( + ("emit", "display"), + [(Emit.THOUGHTS, "summarized"), (Emit.TOOL_CALLS, None)], + ids=["thoughts-posted", "thoughts-not-posted"], ) + async def test_thinking_text_is_requested_only_when_thoughts_are_posted( + self, claude_room: OpenRoom, emit: Emit, display: str | None + ) -> None: + room = await claude_room( + ClaudeSDKAdapterConfig(max_thinking_tokens=1024), emit=emit + ) - options = await _cli_options(room) - - assert options.thinking is not None - assert options.thinking.get("display") == display - - -@pytest.mark.parametrize( - "flag", ["mcp-config", "effort", "settings", "dangerously-skip-permissions"] -) -def test_extra_args_cannot_replace_adapter_owned_flags(flag: str) -> None: - with pytest.raises(ValueError, match="adapter-owned CLI flags"): - ClaudeCLIOptions(extra_args={flag: "x"}) + options = await _cli_options(room) + assert options.thinking is not None + assert options.thinking.get("display") == display -def test_extra_args_keys_are_bare_flag_names() -> None: - with pytest.raises(ValueError, match="bare flag names"): - ClaudeCLIOptions(extra_args={"--debug-to-stderr": None}) + @pytest.mark.parametrize( + "flag", ["mcp-config", "effort", "settings", "dangerously-skip-permissions"] + ) + def test_extra_args_cannot_replace_adapter_owned_flags(self, flag: str) -> None: + with pytest.raises(ValueError, match="adapter-owned CLI flags"): + ClaudeCLIOptions(extra_args={flag: "x"}) + + def test_extra_args_keys_are_bare_flag_names( + self, + ) -> None: + with pytest.raises(ValueError, match="bare flag names"): + ClaudeCLIOptions(extra_args={"--debug-to-stderr": None}) diff --git a/tests/adapters/claude_sdk/test_mcp_resolver.py b/tests/adapters/claude_sdk/test_mcp_resolver.py new file mode 100644 index 000000000..d1a71ff2b --- /dev/null +++ b/tests/adapters/claude_sdk/test_mcp_resolver.py @@ -0,0 +1,873 @@ +"""``resolve_mcp_servers``: one frozen snapshot per new Claude connection, sent +through a private file, used only once every server reports connected.""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import logging +import os +import stat +import subprocess +import sys +import tempfile +import typing +from collections.abc import Awaitable, Callable, Mapping +from pathlib import Path +from typing import Any + +import pytest +from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient +from claude_agent_sdk._internal.transport.subprocess_cli import SubprocessCLITransport +from claude_agent_sdk.types import McpServerConfig + +from band.adapters.claude_sdk import ClaudeSDKAdapter +from band.core.exceptions import InputNotAccepted +from band.core.retention import Residency +from band.integrations.claude_sdk.mcp_startup import ( + ClaudeMcpStartupError, + McpConfigFile, + await_mcp_ready, +) +from band.integrations.claude_sdk.session_manager import ( + ClaudeSessionManager, + ClaudeSessionManagerStoppedError, +) +from band.runtime.tools import BAND_MCP_SERVER_NAME +from tests.adapters.claude_sdk.external_mcp import external_mcp +from tests.adapters.claude_sdk.fakecli import FakeClaude, Hold +from tests.adapters.claude_sdk.helpers import ( + ClaudeRoom, + refuse_mcp_removal, + with_approvals, +) +from tests.baseline.decisions import ModelDecision + +OpenRoom = Callable[..., Awaitable[ClaudeRoom]] + +from tests.adapters.claude_sdk.mcpkit import ( + GRANT_A, + GRANT_B, + ExplodingServers, + HostileItems, + MarkedServers, + Resolver, + assert_secret_free, + connected, + item_failure, + launch_file, + metadata_attack, + private_files, +) + + +@pytest.fixture(autouse=True) +def private_tempdir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + """Private MCP files go to test-owned storage, so their inventory never + sees another process's files in the host's temp directory.""" + owned = tmp_path / "tmp" + owned.mkdir() + monkeypatch.setattr(tempfile, "tempdir", str(owned)) + return owned + + +class TestMcpResolver: + async def test_a_new_connection_dials_its_snapshot_and_reuse_never_resolves( + self, claude_room: OpenRoom, claude: FakeClaude, tmp_path: Path + ) -> None: + async with external_mcp(GRANT_A) as ext: + resolver = Resolver({"ext": ext.server(GRANT_A)}) + room = await claude_room( + workspace_for_room=lambda room_id: str(tmp_path / room_id), + resolve_mcp_servers=resolver, + ) + claude.script([room.model_reply("One.")], [room.model_reply("Two.")]) + await room.send("first") + await room.send("second") + + assert resolver.calls == [str(tmp_path / "room-1")] + [session] = claude.sessions + assert set(session.configured_servers) == {"ext", BAND_MCP_SERVER_NAME} + # A real MCP handshake with grant A happened before the first prompt. + assert (GRANT_A, 200) in ext.receipts + assert room.chat == ["One.", "Two."] + + async def test_an_empty_result_launches_only_band( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + room = await claude_room(resolve_mcp_servers=Resolver({})) + claude.script([room.model_reply("Hi.")]) + await room.send("hi") + + assert set(claude.sessions[0].configured_servers) == {BAND_MCP_SERVER_NAME} + + async def test_the_launch_ignores_changes_to_what_the_callback_returned( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + async with external_mcp(GRANT_A) as ext: + returned = {"ext": ext.server(GRANT_A)} + + async def resolver(workspace: str) -> Mapping[str, Any]: + # The callback keeps changing its objects after handing them over. + asyncio.get_running_loop().call_soon( + returned["ext"]["headers"].__setitem__, "Authorization", "Bearer x" + ) + return returned + + room = await claude_room(resolve_mcp_servers=resolver) + claude.script([room.model_reply("Hi.")]) + await room.send("hi") + + sent = json.loads(launch_file(claude.sessions[0]).read_text()) + assert sent["mcpServers"]["ext"]["headers"] == { + "Authorization": f"Bearer {GRANT_A}" + } + + async def test_a_rotated_grant_reaches_only_a_new_connection( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + async with external_mcp(GRANT_A) as ext: + resolver = Resolver({"ext": ext.server(GRANT_A)}) + room = await claude_room(resolve_mcp_servers=resolver) + claude.script( + [room.model_reply("A.")], + [room.model_reply("Still A.")], + [room.model_reply("B.")], + ) + await room.send("one") + ext.grants = {GRANT_B} + resolver.servers = {"ext": ext.server(GRANT_B)} + await room.send("two") # reused: no resolution, no restart + assert len(resolver.calls) == 1 + assert len(claude.sessions) == 1 + + await room.leave() + await room.send("three") + + assert len(resolver.calls) == 2 + assert len(claude.sessions) == 2 + assert not claude.sessions[0].alive + assert (GRANT_B, 200) in ext.receipts + assert room.chat == ["A.", "Still A.", "B."] + + async def test_a_replaced_band_server_gets_a_new_connection( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + resolver = Resolver({}) + room = await claude_room(resolve_mcp_servers=resolver) + claude.script([room.model_reply("Before.")], [room.model_reply("After.")]) + await room.send("one") + first = claude.sessions[0].configured_servers[BAND_MCP_SERVER_NAME]["url"] + await room.crash_band_server() + await room.send("two") + + assert len(claude.sessions) == 2 + assert not claude.sessions[0].alive + second = claude.sessions[1].configured_servers[BAND_MCP_SERVER_NAME]["url"] + assert first != second + assert len(resolver.calls) == 2 + assert room.chat == ["Before.", "After."] + + async def test_the_first_prompt_waits_for_every_server_to_connect( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + pending = [ + {"name": "ext", "status": "pending"}, + {"name": BAND_MCP_SERVER_NAME, "status": "connected"}, + ] + claude.mcp_statuses.extend( + [pending, pending, connected("ext", BAND_MCP_SERVER_NAME)] + ) + resolver = Resolver({"ext": {"type": "stdio", "command": "ext-server"}}) + room = await claude_room(resolve_mcp_servers=resolver) + claude.script([room.model_reply("Ready.")]) + await room.send("hi") + + assert claude.status_requests == 3 + assert room.chat == ["Ready."] + + @pytest.mark.parametrize( + "statuses", + [ + pytest.param( + [{"name": "ext", "status": "failed", "error": GRANT_A}], id="failed" + ), + pytest.param([{"name": "ext", "status": "needs-auth"}], id="needs-auth"), + pytest.param([{"name": "ext", "status": "disabled"}], id="disabled"), + pytest.param( + [*connected("ext", BAND_MCP_SERVER_NAME), *connected("ext")], + id="duplicate", + ), + pytest.param([{"name": "ext"}], id="malformed"), + pytest.param([{"name": "ext", "status": "weird"}], id="unknown"), + ], + ) + async def test_a_server_that_cannot_connect_fails_the_turn_safely( + self, + claude_room: OpenRoom, + claude: FakeClaude, + caplog: pytest.LogCaptureFixture, + statuses: list[dict[str, Any]], + ) -> None: + claude.mcp_statuses.append( + [{"config": {"headers": {"Authorization": GRANT_A}}, **s} for s in statuses] + ) + resolver = Resolver( + {"ext": {"type": "stdio", "command": "ext", "env": {"T": GRANT_A}}} + ) + room = await claude_room(resolve_mcp_servers=resolver) + + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ClaudeMcpStartupError) as caught, + ): + await room.send("hi") + + assert claude.prompts == [] + [session] = claude.sessions + assert not session.alive + assert not launch_file(session).exists() + assert len(room.failures) == 1 + assert_secret_free(GRANT_A, caught.value, caplog.text) + + @pytest.mark.looptime + async def test_servers_still_pending_at_the_deadline_fail_the_turn( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + claude.mcp_statuses.append( + [{"name": BAND_MCP_SERVER_NAME, "status": "connected"}] + ) + resolver = Resolver({"ext": {"type": "stdio", "command": "ext"}}) + room = await claude_room(resolve_mcp_servers=resolver) + loop = asyncio.get_running_loop() + started = loop.time() + + with pytest.raises(ClaudeMcpStartupError, match="deadline"): + await room.send("hi") + + assert 30 <= loop.time() - started < 31 + assert claude.prompts == [] + assert not claude.sessions[0].alive + + async def test_a_failing_resolver_fails_once_without_falling_back( + self, + claude_room: OpenRoom, + claude: FakeClaude, + caplog: pytest.LogCaptureFixture, + ) -> None: + async def resolver(workspace: str) -> Mapping[str, Any]: + raise LookupError(f"grant store said {GRANT_A}") + + room = await claude_room(resolve_mcp_servers=resolver) + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ClaudeMcpStartupError) as caught, + ): + await room.send("hi", session_id="sess-earlier") + + # No retry without the history, and no CLI was ever started. + assert claude.sessions == [] + assert caught.value.stage == "resolution" + assert caught.value.error_type == "LookupError" + assert_secret_free(GRANT_A, caught.value, caplog.text) + + @pytest.mark.parametrize( + "servers", + [ + pytest.param( + {BAND_MCP_SERVER_NAME: {"type": "http", "url": "http://x"}}, + id="band-name", + ), + pytest.param({"inproc": {"type": "sdk", "name": "inproc"}}, id="sdk"), + pytest.param({"bad": {"type": "http"}}, id="no-url"), + ], + ) + async def test_unacceptable_servers_are_refused_by_name_before_launch( + self, claude_room: OpenRoom, claude: FakeClaude, servers: dict[str, Any] + ) -> None: + room = await claude_room(resolve_mcp_servers=Resolver(servers)) + with pytest.raises(ClaudeMcpStartupError) as caught: + await room.send("hi") + + assert claude.sessions == [] + assert caught.value.server == next(iter(servers)) + + async def test_an_external_tool_still_needs_host_permission( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + async with external_mcp(GRANT_A) as ext: + room = await claude_room( + resolve_mcp_servers=Resolver({"ext": ext.server(GRANT_A)}) + ) + claude.script( + [ + ModelDecision.call("mcp__ext__echo", text="hello"), + room.model_reply("Done."), + ] + ) + await room.send("hi") + + assert room.tool_errors == { + "mcp__ext__echo": True, + "band_send_message": False, + } + assert room.chat == ["Done."] + + async def test_the_native_launch_carries_a_private_file_path_not_credentials( + self, claude_room: OpenRoom, claude: FakeClaude, tmp_path: Path + ) -> None: + resolver = Resolver( + { + "ext": { + "type": "http", + "url": "http://127.0.0.1:9/mcp", + "headers": {"Authorization": f"Bearer {GRANT_A}"}, + } + } + ) + claude.mcp_statuses.append(connected("ext", BAND_MCP_SERVER_NAME)) + room = await claude_room(resolve_mcp_servers=resolver) + claude.script([room.model_reply("Hi.")]) + await room.send("hi") + + [session] = claude.sessions + path = launch_file(session) + native = SubprocessCLITransport( + prompt="", + options=ClaudeAgentOptions( + cli_path=sys.executable, mcp_servers=session.options.mcp_servers + ), + ) + argv = native._build_command() + assert argv[argv.index("--mcp-config") + 1] == str(path) + assert not any(GRANT_A in arg for arg in argv) + + await room.leave() + assert not path.exists() + assert not path.parent.exists() + + @pytest.mark.skipif( + os.name != "posix", + reason="POSIX modes do not prove Windows user-profile/temp-directory ACL privacy", + ) + async def test_the_private_launch_file_has_posix_permissions( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + room = await claude_room(resolve_mcp_servers=Resolver({})) + claude.script([room.model_reply("Hi.")]) + await room.send("hi") + path = launch_file(claude.sessions[0]) + try: + assert stat.S_IMODE(os.stat(path.parent).st_mode) == 0o700 + assert stat.S_IMODE(os.stat(path).st_mode) == 0o600 + finally: + await room.leave() + + async def test_a_missing_history_still_falls_back_to_a_fresh_session( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + claude.unresumable = {"sess-gone"} + resolver = Resolver({}) + room = await claude_room(resolve_mcp_servers=resolver) + claude.script([room.model_reply("Fresh.")]) + await room.send("hi", session_id="sess-gone") + + assert claude.resumed == ["sess-gone", None] + # Each genuine connection resolved its own snapshot. + assert len(resolver.calls) == 2 + assert room.chat == ["Fresh."] + + @pytest.mark.parametrize("failure", ["unreadable-result", "typed-error"]) + async def test_callback_failures_are_sanitized_and_never_retried_fresh( + self, + claude_room: OpenRoom, + claude: FakeClaude, + caplog: pytest.LogCaptureFixture, + failure: str, + ) -> None: + calls: list[str] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + calls.append(workspace) + if failure == "typed-error": + raise ClaudeMcpStartupError( + "resolution", f"token {GRANT_A}", server=GRANT_A + ) + return ExplodingServers() + + room = await claude_room(resolve_mcp_servers=resolver) + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ClaudeMcpStartupError) as caught, + ): + await room.send("hi", session_id="sess-earlier") + + assert len(calls) == 1 + assert claude.sessions == [] + assert_secret_free(GRANT_A, caught.value, caplog.text) + + async def test_no_prompt_before_a_held_status_reports_connected( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + hold = claude.status_hold = Hold() + room = await claude_room(resolve_mcp_servers=Resolver({})) + claude.script([room.model_reply("Ready.")]) + sending = room.send_in_background("hi") + try: + await hold.reached.wait() + assert claude.prompts == [] + session = room.session_manager.session(room.room_id) + assert ( + session is not None + and session.retained.snapshot.state == Residency.ACQUIRING + ) + finally: + hold.released.set() + await sending + assert room.chat == ["Ready."] + + @pytest.mark.looptime + async def test_a_status_request_that_never_answers_fails_at_the_readiness_deadline( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + hold = claude.status_hold = Hold() + room = await claude_room(resolve_mcp_servers=Resolver({})) + loop = asyncio.get_running_loop() + started = loop.time() + try: + with pytest.raises(ClaudeMcpStartupError, match="deadline"): + await room.send("hi") + # Bounded by the 30 s total, not the SDK's own 60 s per request. + assert 30 <= loop.time() - started < 31 + assert claude.prompts == [] + finally: + hold.released.set() + + async def test_a_cli_without_mcp_status_fails_safely( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + claude.status_unsupported = True + room = await claude_room(resolve_mcp_servers=Resolver({})) + with pytest.raises(ClaudeMcpStartupError, match="status request failed"): + await room.send("hi") + assert claude.prompts == [] + assert not claude.sessions[0].alive + + @pytest.mark.parametrize("held", ["resolution", "readiness"]) + async def test_a_band_server_replaced_during_startup_gets_a_new_connection( + self, claude_room: OpenRoom, claude: FakeClaude, held: str + ) -> None: + entered, release = asyncio.Event(), asyncio.Event() + calls: list[str] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + calls.append(workspace) + if held == "resolution" and len(calls) == 1: + entered.set() + await release.wait() + return {} + + if held == "readiness": + hold = claude.status_hold = Hold() + room = await claude_room(resolve_mcp_servers=resolver) + claude.script([room.model_reply("Live.")]) + sending = room.send_in_background("hi") + try: + if held == "resolution": + await entered.wait() + else: + await hold.reached.wait() + old_url = room.adapter._room_mcp_servers(room.room_id)[ + BAND_MCP_SERVER_NAME + ]["url"] + await room.crash_band_server() + await room.adapter._mcp.ensure() + finally: + release.set() + if held == "readiness": + claude.status_hold = None + hold.released.set() + await sending + + live = claude.sessions[-1] + assert live.alive + assert live.configured_servers[BAND_MCP_SERVER_NAME]["url"] != old_url + assert all(not s.alive for s in claude.sessions[:-1]) + assert len(calls) == 2 + assert room.chat == ["Live."] + + async def test_a_resolver_requesting_teardown_behind_a_settling_turn_publishes_nothing( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + adapter_ref: list[Any] = [] + requested = asyncio.Event() + + async def resolver(workspace: str) -> Mapping[str, Any]: + if workspace.endswith("room-1"): + return {} + await adapter_ref[0].cleanup_all() + requested.set() + return {} + + room = await claude_room(resolve_mcp_servers=resolver) + adapter_ref.append(room.adapter) + other = room.beside("room-2") + turn_hold = Hold() + claude.script([turn_hold]) + first = room.send_in_background("hold on") + try: + await turn_hold.reached.wait() + sessions_before = len(claude.sessions) + + with pytest.raises( + ( + asyncio.CancelledError, + ClaudeMcpStartupError, + ClaudeSessionManagerStoppedError, + ) + ): + await other.send("hi") + assert requested.is_set() + finally: + # Released and drained whatever the assertions above found. + turn_hold.released.set() + await asyncio.gather(first, return_exceptions=True) + await room.adapter.cleanup_all() + + # The room-2 startup asked for the teardown and was refused on its own: + # nothing launched after the request. + assert len(claude.sessions) == sessions_before + assert not bool(room.adapter._session_manager._rooms) + + async def test_a_failed_teardown_keeps_turns_out_until_a_restart( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + room = await claude_room(resolve_mcp_servers=Resolver({})) + claude.script([room.model_reply("Hi.")]) + await room.send("hi") + backend = room.adapter._mcp.current + bound = room.adapter._mcp_room_tools.get(room.room_id) + claude.refuse_close = True + try: + with pytest.raises(BaseException, match="cleanup"): + await room.adapter.cleanup_all() + # The retained Band server dies; a late turn must not start another. + await room.crash_band_server() + with pytest.raises(ClaudeSessionManagerStoppedError): + await room.send("again") + assert room.adapter._mcp.current is backend + assert room.adapter._mcp_room_tools.get(room.room_id) is bound + finally: + claude.refuse_close = False + await room.adapter.on_started("Test Agent", "restarted") + claude.script([room.model_reply("Back.")]) + await room.send("later") + assert room.chat == ["Hi.", "Back."] + + async def test_a_file_left_behind_fences_new_connections_until_removed( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + before = private_files() + resolver = Resolver({}) + room = await claude_room(resolve_mcp_servers=resolver) + claude.script([room.model_reply("A.")], [room.model_reply("B.")]) + await room.send("first") + first = launch_file(claude.sessions[0]) + with refuse_mcp_removal(first): + with pytest.raises(ClaudeMcpStartupError, match="cleanup"): + await room.leave() + with pytest.raises(InputNotAccepted, match="closing"): + await room.send("second") + with pytest.raises(ClaudeMcpStartupError, match="cleanup"): + await room.leave() + assert len(claude.sessions) == 1 + assert len(resolver.calls) == 1 + assert first.exists() + assert bool(room.adapter._session_manager._rooms) + + await room.leave() + await room.send("third") + assert not first.exists() + assert room.chat == ["A.", "B."] + await room.leave() + assert private_files() == before + assert not bool(room.adapter._session_manager._rooms) + + async def test_a_resume_fallback_does_not_strand_the_failed_attempts_file( + self, claude_room: OpenRoom, claude: FakeClaude + ) -> None: + before = private_files() + claude.unresumable = {"sess-gone"} + resolver = Resolver({}) + room = await claude_room(resolve_mcp_servers=resolver) + hold = claude.connecting = Hold() + sending = asyncio.create_task(room.send("hi", session_id="sess-gone")) + try: + await hold.reached.wait() + attempt = launch_file(claude.sessions[0]) + with refuse_mcp_removal(attempt): + claude.connecting = None + hold.released.set() + # The resumed attempt fails; the fresh fallback finds its file + # still held and stops there. + with pytest.raises(ClaudeMcpStartupError, match="cleanup"): + await sending + assert claude.resumed == ["sess-gone"] + assert attempt.exists() + assert bool(room.adapter._session_manager._rooms) + finally: + claude.connecting = None + hold.released.set() + await asyncio.gather(sending, return_exceptions=True) + + claude.script([room.model_reply("Fresh.")]) + await room.send("again") + assert not attempt.exists() + assert room.chat == ["Fresh."] + await room.leave() + assert private_files() == before + + async def test_a_partly_written_file_stays_owned_until_removed( + self, + claude_room: OpenRoom, + claude: FakeClaude, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + ) -> None: + before = private_files() + written: list[Path] = [] + real_write = McpConfigFile.write + removal_failures = contextlib.ExitStack() + + def interrupted(self: McpConfigFile, servers: Mapping[str, Any]) -> None: + self.path.write_text(f'{{"partial": "{GRANT_A}"') + written.append(self.path) + removal_failures.enter_context(refuse_mcp_removal(self.path)) + raise OSError(f"disk full writing {GRANT_A}") + + monkeypatch.setattr(McpConfigFile, "write", interrupted) + room = await claude_room(resolve_mcp_servers=Resolver({})) + try: + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ClaudeMcpStartupError) as caught, + ): + await room.send("hi") + assert caught.value.stage == "launch" + assert claude.sessions == [] + assert_secret_free(GRANT_A, caught.value, caplog.text) + [partial] = written + assert partial.exists() + assert bool(room.adapter._session_manager._rooms) + finally: + removal_failures.close() + monkeypatch.setattr(McpConfigFile, "write", real_write) + + await room.leave() + assert not written[0].exists() + assert private_files() == before + assert not bool(room.adapter._session_manager._rooms) + + async def test_a_failing_permission_factory_is_sanitized_and_never_retried_fresh( + self, + claude_room: OpenRoom, + claude: FakeClaude, + caplog: pytest.LogCaptureFixture, + ) -> None: + before = private_files() + resolver = Resolver({}) + room = await claude_room(with_approvals(), resolve_mcp_servers=resolver) + + def leaking_factory(room_id: str) -> Any: + raise LookupError(f"policy store said {GRANT_A}") + + room.adapter._session_manager._can_use_tool_factory = leaking_factory + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ClaudeMcpStartupError) as caught, + ): + await room.send("hi", session_id="sess-earlier") + + assert caught.value.stage == "launch" + assert not caught.value.resume_eligible + assert len(resolver.calls) == 1 + assert claude.sessions == [] + assert_secret_free(GRANT_A, caught.value, caplog.text) + assert private_files() == before + + async def test_an_owned_looking_error_from_the_result_is_still_untrusted( + self, + claude_room: OpenRoom, + claude: FakeClaude, + caplog: pytest.LogCaptureFixture, + ) -> None: + calls: list[str] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + calls.append(workspace) + return MarkedServers() + + room = await claude_room(resolve_mcp_servers=resolver) + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ClaudeMcpStartupError) as caught, + ): + await room.send("hi", session_id="sess-earlier") + + assert caught.value.stage == "resolution" + assert len(calls) == 1 + assert claude.sessions == [] + assert_secret_free(GRANT_A, caught.value, caplog.text) + + def test_the_public_constructors_have_evaluable_native_resolver_types( + self, + ) -> None: + for constructor in (ClaudeSDKAdapter.__init__, ClaudeSessionManager.__init__): + hint = typing.get_type_hints(constructor)["resolve_mcp_servers"] + resolver, none = typing.get_args(hint) + assert none is type(None) + [returned] = typing.get_args(typing.get_args(resolver)[1]) + assert typing.get_args(returned) == (str, McpServerConfig) + + def test_the_resolver_type_needs_no_claude_sdk_at_import( + self, + ) -> None: + probe = ( + "import sys\n" + "class Block:\n" + " def find_spec(self, name, path=None, target=None):\n" + " if name.split('.')[0] == 'claude_agent_sdk':\n" + " raise ImportError(name)\n" + "sys.meta_path.insert(0, Block())\n" + "from band.integrations.claude_sdk.mcp_startup import McpServerResolver\n" + "McpServerResolver | None\n" + ) + subprocess.run([sys.executable, "-c", probe], check=True) + + @pytest.mark.parametrize("kind", ["ordinary", "marked", "named-class"]) + async def test_an_item_that_fails_when_unpacked_is_untrusted( + self, + claude_room: OpenRoom, + claude: FakeClaude, + caplog: pytest.LogCaptureFixture, + kind: str, + ) -> None: + before = private_files() + calls: list[str] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + calls.append(workspace) + return HostileItems(item_failure(kind)) + + room = await claude_room(resolve_mcp_servers=resolver) + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ClaudeMcpStartupError) as caught, + ): + await room.send("hi", session_id="sess-earlier") + + assert caught.value.stage == "resolution" + assert not caught.value.resume_eligible + assert len(calls) == 1 + assert claude.sessions == [] + assert_secret_free(GRANT_A, caught.value, caplog.text) + assert_secret_free(GRANT_A.replace("-", "_"), caught.value, caplog.text) + await room.leave() + assert private_files() == before + assert not bool(room.adapter._session_manager._rooms) + + @pytest.mark.parametrize( + "kind", [[], {}, 3, None], ids=["list", "dict", "int", "null"] + ) + async def test_a_non_string_server_type_is_refused_by_name( + self, claude_room: OpenRoom, claude: FakeClaude, kind: Any + ) -> None: + calls: list[str] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + calls.append(workspace) + return {"tool": {"type": kind, "command": "tool", "url": "http://x"}} + + room = await claude_room(resolve_mcp_servers=resolver) + with pytest.raises( + ClaudeMcpStartupError, match="type must be a string" + ) as caught: + await room.send("hi", session_id="sess-earlier") + + assert caught.value.server == "tool" + assert not caught.value.resume_eligible + assert len(calls) == 1 + assert claude.sessions == [] + + @pytest.mark.parametrize( + "kind", ["spoofed-module", "non-string-module", "raising-metaclass"] + ) + async def test_an_error_classs_own_metadata_is_never_read( + self, + claude_room: OpenRoom, + claude: FakeClaude, + caplog: pytest.LogCaptureFixture, + kind: str, + ) -> None: + calls: list[str] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + calls.append(workspace) + raise metadata_attack(kind) + + room = await claude_room(resolve_mcp_servers=resolver) + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ClaudeMcpStartupError) as caught, + ): + await room.send("hi", session_id="sess-earlier") + + assert caught.value.stage == "resolution" + assert caught.value.error_type is None + assert not caught.value.resume_eligible + assert len(calls) == 1 + assert claude.sessions == [] + assert_secret_free(GRANT_A, caught.value, caplog.text) + assert_secret_free(GRANT_A.replace("-", "_"), caught.value, caplog.text) + await room.leave() + assert not bool(room.adapter._session_manager._rooms) + + @pytest.mark.looptime + async def test_a_status_callback_swallowing_the_deadline_cannot_publish_readiness( + self, + ) -> None: + async def late_status() -> object: + try: + await asyncio.sleep(31) + except asyncio.CancelledError: + return {"mcpServers": []} + return {"mcpServers": []} + + with pytest.raises(ClaudeMcpStartupError, match="deadline"): + await await_mcp_ready(late_status, frozenset()) + + async def test_a_failed_client_constructor_closes_the_owned_transport( + self, + claude_room: OpenRoom, + claude: FakeClaude, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + ) -> None: + def refuse(self: ClaudeSDKClient, *args: Any, **kwargs: Any) -> None: + raise LookupError(f"client construction failed with {GRANT_A}") + + monkeypatch.setattr(ClaudeSDKClient, "__init__", refuse) + room = await claude_room(resolve_mcp_servers=Resolver({})) + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ClaudeMcpStartupError) as caught, + ): + await room.send("hi") + assert caught.value.stage == "launch" + assert not caught.value.resume_eligible + assert_secret_free(GRANT_A, caught.value, caplog.text) + assert len(claude.sessions) == 1 + assert not claude.sessions[0].alive + assert not launch_file(claude.sessions[0]).exists() + await room.leave() diff --git a/tests/adapters/claude_sdk/test_mcp_startup_ownership.py b/tests/adapters/claude_sdk/test_mcp_startup_ownership.py new file mode 100644 index 000000000..f9247d965 --- /dev/null +++ b/tests/adapters/claude_sdk/test_mcp_startup_ownership.py @@ -0,0 +1,668 @@ +"""Who owns a Claude connection while its MCP servers resolve: the room's +startup, until it publishes or settles. The CLI is the scripted fake behind +the real SDK client.""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import tempfile +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import pytest +from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient + +from band.core.exceptions import CleanupIncompleteError, RoomClosingError +from band.integrations.claude_sdk.mcp_startup import ClaudeMcpStartupError +from band.integrations.claude_sdk.session_manager import ( + ClaudeSessionManager, + ClaudeSessionManagerStoppedError, +) +from band.runtime.tools import BAND_MCP_SERVER_NAME +from band.workspaces import RoomWorkspaces +from tests.adapters.claude_sdk.fakecli import FakeClaude, Hold +from tests.adapters.claude_sdk.helpers import ( + lease_client, + refuse_mcp_removal, +) +from tests.adapters.claude_sdk.mcpkit import ( + BAND, + SECRET, + ClientCustody, + CountingResolver, + FailingBand, + HeldResolver, + HeldServers, + RaisingTruth, + ReuseProbe, + SettlingResolver, + manager_for, + outdated_session, +) + + +class TestMcpStartupOwnership: + async def test_concurrent_callers_share_one_resolution( + self, tmp_path: Path, claude: FakeClaude + ) -> None: + resolver = HeldResolver() + manager = manager_for(tmp_path, resolver, claude) + try: + first = asyncio.create_task(lease_client(manager, "room-a")) + second = asyncio.create_task(lease_client(manager, "room-a")) + await resolver.entered.wait() + resolver.release.set() + a, b = await asyncio.gather(first, second) + assert a is b + assert len(resolver.calls) == 1 + assert len(claude.sessions) == 1 + finally: + resolver.release.set() + await manager.stop() + + async def test_a_held_resolution_never_blocks_another_room( + self, tmp_path: Path, claude: FakeClaude + ) -> None: + held = HeldResolver() + + async def resolver(workspace: str) -> Mapping[str, Any]: + if workspace.endswith("room-a"): + return await held(workspace) + return {} + + manager = manager_for(tmp_path, resolver, claude) + try: + blocked = asyncio.create_task(lease_client(manager, "room-a")) + await held.entered.wait() + other = await lease_client(manager, "room-b") + await manager.cleanup_session("room-b") + assert not blocked.done() + assert other is not None + finally: + held.release.set() + await asyncio.gather(blocked, return_exceptions=True) + await manager.stop() + + async def test_the_last_waiter_leaving_retires_the_startup_without_a_late_client( + self, tmp_path: Path, claude: FakeClaude + ) -> None: + resolver = HeldResolver(swallow=True) + manager = manager_for(tmp_path, resolver, claude) + try: + staying = asyncio.create_task(lease_client(manager, "room-a")) + leaving = asyncio.create_task(lease_client(manager, "room-a")) + await resolver.entered.wait() + leaving.cancel() + await asyncio.gather(leaving, return_exceptions=True) + # One waiter still wants it: nothing was cancelled. + assert resolver.cancellations == 0 + staying.cancel() + while resolver.cancellations == 0: + await asyncio.sleep(0) + assert not staying.done() + # The callback swallows that one cancellation and keeps going; the room + # keeps its claim until it settles, and nothing launches afterwards. + assert bool(manager._rooms) + resolver.release.set() + await resolver.settled.wait() + await asyncio.gather(staying, return_exceptions=True) + await manager.cleanup_session("room-a") + assert resolver.cancellations == 1 + assert claude.sessions == [] + assert not bool(manager._rooms) + finally: + resolver.release.set() + await manager.stop() + + async def test_a_request_cancelled_before_starting_claims_nothing( + self, tmp_path: Path, claude: FakeClaude + ) -> None: + resolver = HeldResolver() + manager = manager_for(tmp_path, resolver, claude) + try: + request = asyncio.create_task(lease_client(manager, "room-a")) + request.cancel() + await asyncio.gather(request, return_exceptions=True) + # No room or resolver is entered by the cancelled operation. + await manager.cleanup_session("room-other") + assert resolver.calls == [] + assert not bool(manager._rooms) + finally: + await manager.stop() + + @pytest.mark.parametrize("swallow", [False, True]) + async def test_cleanup_retires_a_resolving_startup_and_holds_the_room_until_it_settles( + self, tmp_path: Path, claude: FakeClaude, swallow: bool + ) -> None: + resolver = HeldResolver(swallow=swallow) + manager = manager_for(tmp_path, resolver, claude) + try: + waiting = asyncio.create_task(lease_client(manager, "room-a")) + await resolver.entered.wait() + cleanup = asyncio.create_task(manager.cleanup_session("room-a")) + arrival = asyncio.create_task(lease_client(manager, "room-a")) + await asyncio.wait((cleanup,), timeout=0) + if swallow: + assert not cleanup.done() + resolver.release.set() + await cleanup + with pytest.raises((ClaudeMcpStartupError, asyncio.CancelledError)): + await waiting + # Closing admission refuses arrivals; the host can redeliver after + # the retained disposal has released the previous room. + with pytest.raises(RoomClosingError): + await arrival + assert resolver.cancellations == 1 + resolver.release.set() + client = await lease_client(manager, "room-a") + assert len(resolver.calls) == 2 + assert len(claude.sessions) == 1 + assert claude.sessions[0].alive and client is not None + finally: + resolver.release.set() + await manager.stop() + + async def test_a_resolver_cleaning_up_its_own_room_does_not_wait_on_itself( + self, tmp_path: Path, claude: FakeClaude + ) -> None: + manager_ref: list[ClaudeSessionManager] = [] + returned = asyncio.Event() + + async def resolver(workspace: str) -> Mapping[str, Any]: + with pytest.raises(CleanupIncompleteError): + await manager_ref[0].cleanup_session("room-a") + returned.set() + return {} + + manager = manager_for(tmp_path, resolver, claude) + manager_ref.append(manager) + try: + with pytest.raises((ClaudeMcpStartupError, asyncio.CancelledError)): + async with asyncio.timeout(5): + await lease_client(manager, "room-a") + assert returned.is_set() + assert claude.sessions == [] + async with asyncio.timeout(5): + while bool(manager._rooms): + await asyncio.sleep(0) + finally: + await manager.stop() + + async def test_a_resolver_stopping_the_manager_does_not_wait_on_itself( + self, tmp_path: Path, claude: FakeClaude + ) -> None: + manager_ref: list[ClaudeSessionManager] = [] + + async def resolver(workspace: str) -> Mapping[str, Any]: + with pytest.raises(CleanupIncompleteError): + await manager_ref[0].stop() + return {} + + manager = manager_for(tmp_path, resolver, claude) + manager_ref.append(manager) + with pytest.raises((ClaudeSessionManagerStoppedError, asyncio.CancelledError)): + async with asyncio.timeout(5): + await lease_client(manager, "room-a") + async with asyncio.timeout(5): + await manager.stop() + assert claude.sessions == [] + assert not bool(manager._rooms) + + async def test_stop_wins_over_a_startup_finishing_late( + self, tmp_path: Path, claude: FakeClaude + ) -> None: + resolver = HeldResolver(swallow=True) + manager = manager_for(tmp_path, resolver, claude) + waiting = asyncio.create_task(lease_client(manager, "room-a")) + await resolver.entered.wait() + stopping = asyncio.create_task(manager.stop()) + await asyncio.wait((stopping,), timeout=0) + resolver.release.set() + await stopping + with pytest.raises((ClaudeSessionManagerStoppedError, asyncio.CancelledError)): + await waiting + assert claude.sessions == [] + assert not bool(manager._rooms) + + async def test_a_failed_close_keeps_the_file_and_claim_until_a_retry( + self, tmp_path: Path, claude: FakeClaude + ) -> None: + async def resolver(workspace: str) -> Mapping[str, Any]: + return {} + + manager = manager_for(tmp_path, resolver, claude) + try: + await lease_client(manager, "room-a") + [session] = claude.sessions + path = Path(str(session.options.mcp_servers)) + claude.refuse_close = True + with pytest.raises(ClaudeMcpStartupError, match="cleanup"): + await manager.cleanup_session("room-a") + assert path.exists() + assert bool(manager._rooms) + + claude.refuse_close = False + await manager.cleanup_session("room-a") + assert not path.exists() + assert not bool(manager._rooms) + finally: + claude.refuse_close = False + await manager.stop() + + @pytest.mark.parametrize( + "retire_own_room", [False, True], ids=["self-cancel", "self-cleanup"] + ) + async def test_a_settling_startup_is_never_cancelled_again( + self, tmp_path: Path, claude: FakeClaude, retire_own_room: bool + ) -> None: + resolver = SettlingResolver(retire_own_room=retire_own_room) + manager = manager_for(tmp_path, resolver, claude) + resolver.manager = manager + try: + waiting = asyncio.create_task(lease_client(manager, "room-a")) + await resolver.settling.wait() + # An outside cleanup takes the startup over while it settles. + cleanup = asyncio.create_task(manager.cleanup_session("room-a")) + await asyncio.sleep(0) + resolver.release.set() + await cleanup + with pytest.raises((ClaudeMcpStartupError, asyncio.CancelledError)): + await waiting + assert resolver.log == ["settled"] + assert claude.sessions == [] + assert not bool(manager._rooms) + finally: + resolver.release.set() + await manager.stop() + + async def test_a_replacement_whose_caller_left_during_the_old_close_never_publishes( + self, tmp_path: Path, claude: FakeClaude + ) -> None: + manager, resolver, _ = await outdated_session(tmp_path, claude) + hold = claude.closing = Hold() + try: + replacing = asyncio.create_task(lease_client(manager, "room-a")) + await hold.reached.wait() # the old client's close is in progress + replacing.cancel() + await asyncio.gather(replacing, return_exceptions=True) + hold.released.set() + await manager.cleanup_session("room-a") + + # The replacement was retired: never resolved, launched or published. + assert len(resolver.calls) == 1 + assert len(claude.sessions) == 1 + assert not claude.sessions[0].alive + assert not bool(manager._rooms) + finally: + hold.released.set() + claude.closing = None + await manager.stop() + + async def test_stop_during_a_replacements_old_close_leaves_nothing_behind( + self, tmp_path: Path, claude: FakeClaude + ) -> None: + manager, resolver, _ = await outdated_session(tmp_path, claude) + hold = claude.closing = Hold() + try: + replacing = asyncio.create_task(lease_client(manager, "room-a")) + await hold.reached.wait() + stopping = asyncio.create_task(manager.stop()) + hold.released.set() + await stopping + with pytest.raises(RoomClosingError): + await replacing + await manager.stop() + + assert len(resolver.calls) == 1 + assert len(claude.sessions) == 1 + assert not bool(manager._rooms) + finally: + hold.released.set() + claude.closing = None + + async def test_a_resolver_that_swallows_its_own_cancellation_publishes_nothing( + self, tmp_path: Path, claude: FakeClaude + ) -> None: + async def resolver(workspace: str) -> Mapping[str, Any]: + task = asyncio.current_task() + assert task is not None + task.cancel() + try: + await asyncio.sleep(0) + except asyncio.CancelledError: + pass # caught, never uncancelled, and a result returned anyway + return {} + + manager = manager_for(tmp_path, resolver, claude) + try: + with pytest.raises(asyncio.CancelledError): + await lease_client(manager, "room-a") + assert claude.sessions == [] + assert not manager.has_session("room-a") + finally: + await manager.stop() + + async def test_a_finished_keep_claim_retirement_does_not_stall_a_releasing_cleanup( + self, tmp_path: Path, claude: FakeClaude + ) -> None: + manager = manager_for(tmp_path, CountingResolver(), claude) + await lease_client(manager, "room-a") + try: + session = manager.room("room-a") + assert session.recover("retire before final disposal") + + await manager.cleanup_session("room-a") + + assert not bool(manager._rooms) + finally: + await manager.stop() + + async def test_a_file_that_cannot_be_removed_stays_owned_until_a_retry( + self, tmp_path: Path, claude: FakeClaude + ) -> None: + manager = manager_for(tmp_path, CountingResolver(), claude) + await lease_client(manager, "room-a") + path = Path(str(claude.sessions[0].options.mcp_servers)) + try: + with refuse_mcp_removal(path): + with pytest.raises(ClaudeMcpStartupError, match="cleanup"): + await manager.cleanup_session("room-a") + assert path.exists() + assert bool(manager._rooms) + + await manager.cleanup_session("room-a") + assert not path.exists() + assert not bool(manager._rooms) + finally: + await manager.stop() + + async def test_a_close_failure_reports_no_native_text( + self, tmp_path: Path, claude: FakeClaude, caplog: pytest.LogCaptureFixture + ) -> None: + manager = manager_for(tmp_path, CountingResolver(), claude) + await lease_client(manager, "room-a") + [session] = claude.sessions + real_close = session.close + + async def leaking_close() -> None: + raise OSError(f"close failed with {SECRET}") + + session.close = leaking_close # type: ignore[method-assign] + try: + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ClaudeMcpStartupError) as caught, + ): + await manager.cleanup_session("room-a") + rendered = [str(caught.value), repr(caught.value), caplog.text] + assert caught.value.__cause__ is None and caught.value.__suppress_context__ + assert not any(SECRET in text for text in rendered) + assert bool(manager._rooms) + finally: + session.close = real_close # type: ignore[method-assign] + await manager.stop() + + @pytest.mark.parametrize("retained", ["file", "child"]) + async def test_a_failed_reservation_keeps_what_the_room_still_holds( + self, + tmp_path: Path, + claude: FakeClaude, + caplog: pytest.LogCaptureFixture, + retained: str, + ) -> None: + resolver = CountingResolver() + band = FailingBand() + claude.mcp_statuses.append( + [{"name": BAND_MCP_SERVER_NAME, "status": "connected"}] + ) + manager = ClaudeSessionManager( + ClaudeAgentOptions(), + mcp_servers_factory=band, + workspaces=RoomWorkspaces(lambda room_id: str(tmp_path / room_id)), + resolve_mcp_servers=resolver, + ) + await lease_client(manager, "room-a") + path = Path(str(claude.sessions[0].options.mcp_servers)) + refused_removal = contextlib.ExitStack() + try: + # The room's cleanup fails, leaving its file (or its whole child) held. + if retained == "file": + refused_removal.enter_context(refuse_mcp_removal(path)) + else: + claude.refuse_close = True + with pytest.raises(ClaudeMcpStartupError, match="cleanup"): + session = manager.room("room-a") + session.recover("retire before replacement") + await lease_client(manager, "room-a") + if retained == "file": + refused_removal.close() + else: + claude.refuse_close = False + + band.failing = True + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(ClaudeMcpStartupError) as caught, + ): + await lease_client(manager, "room-a") + assert caught.value.stage == "reservation" + assert SECRET not in str(caught.value) and SECRET not in caplog.text + # Nothing new was resolved, allocated or spawned, and the old custody + # is still owned. + assert len(resolver.calls) == 1 + assert len(claude.sessions) == 1 + assert path.exists() + assert bool(manager._rooms) + + band.failing = False + await manager.cleanup_session("room-a") + assert not path.exists() + assert not bool(manager._rooms) + await lease_client(manager, "room-a") + assert len(claude.sessions) == 2 + finally: + refused_removal.close() + claude.refuse_close = False + band.failing = False + await manager.stop() + assert not bool(manager._rooms) + + @pytest.mark.parametrize( + "phase", ["reservation", "reuse", "pre-launch", "construction", "post-status"] + ) + async def test_a_callback_closing_admission_stops_everything_after_it( + self, + tmp_path: Path, + claude: FakeClaude, + monkeypatch: pytest.MonkeyPatch, + phase: str, + ) -> None: + """A callback the manager runs closes its admission mid-startup: nothing + is constructed, reused or published after that.""" + owned_tmp = tmp_path / "tmp" + owned_tmp.mkdir() + monkeypatch.setattr(tempfile, "tempdir", str(owned_tmp)) + armed = {"on": phase == "reservation"} + manager_ref: list[ClaudeSessionManager] = [] + resolutions: list[str] = [] + + def band(room_id: str) -> dict[str, Any]: + if armed["on"] or (phase == "post-status" and claude.status_requests > 0): + manager_ref[0].close_admission() + return dict(BAND) + + async def resolver(workspace: str) -> Mapping[str, Any]: + resolutions.append(workspace) + if phase == "pre-launch": + armed["on"] = True + return {} + + async def allow(*args: Any) -> Any: + return None + + def permissions(session: Any, generation: int) -> Any: + if phase == "construction" and resolutions: + manager_ref[0].close_admission() + return allow + + claude.mcp_statuses.append( + [{"name": BAND_MCP_SERVER_NAME, "status": "connected"}] + ) + manager = ClaudeSessionManager( + ClaudeAgentOptions(), + can_use_tool_factory=permissions, + mcp_servers_factory=band, + workspaces=RoomWorkspaces(lambda room_id: str(tmp_path / room_id)), + resolve_mcp_servers=resolver, + ) + manager_ref.append(manager) + try: + if phase == "reuse": + await lease_client(manager, "room-a") + armed["on"] = True + sessions_before = len(claude.sessions) + with pytest.raises( + (ClaudeSessionManagerStoppedError, asyncio.CancelledError) + ): + async with asyncio.timeout(10): + await lease_client(manager, "room-a") + + if phase == "post-status": + # Built and confirmed, then refused instead of published. + assert len(claude.sessions) == 1 + assert not claude.sessions[0].alive + else: + assert len(claude.sessions) == sessions_before + if phase == "reservation": + assert resolutions == [] + assert not manager.has_session("room-a") or phase == "reuse" + finally: + await manager.stop() + assert not bool(manager._rooms) + assert all(not s.alive for s in claude.sessions) + assert list(owned_tmp.glob("band-claude-mcp-*")) == [] + + @pytest.mark.parametrize("mode", ["materialized", "deferred"]) + async def test_the_reuse_decision_is_a_plain_bool_by_the_final_guard( + self, + tmp_path: Path, + claude: FakeClaude, + mode: str, + ) -> None: + """The reuse decision is reduced to a plain bool once, inside the fallible + prerequisite boundary, before the callback-free admission guard. An object + answering the currentness comparison is therefore never asked its truth + again: a callback that would close admission at the branch test cannot run + there at all (the deferred case), and closing on the first read is refused + outright (the control).""" + + def band(room_id: str) -> Any: + return dict(BAND) + + manager = ClaudeSessionManager( + ClaudeAgentOptions(), + mcp_servers_factory=band, + workspaces=RoomWorkspaces(lambda room_id: str(tmp_path / room_id)), + ) + claude.mcp_statuses.append( + [{"name": BAND_MCP_SERVER_NAME, "status": "connected"}] + ) + client = await lease_client(manager, "room-a") + custody = ClientCustody(client) + real_options = client.options + probe = ReuseProbe(manager, mode) + sessions_before = len(claude.sessions) + try: + client.options = HeldServers(probe) + if mode == "materialized": # closes while the decision is materialized + with pytest.raises(ClaudeSessionManagerStoppedError): + await lease_client(manager, "room-a") + else: # a second read, which the settled decision must never take + returned = await lease_client(manager, "room-a") + assert probe.reads == 1, f"the decision was read {probe.reads} times" + assert not manager._retention.stopping + assert returned is client + finally: + client.options = real_options + await manager.stop() + assert all(not s.alive for s in claude.sessions) + assert len(claude.sessions) == sessions_before + assert not bool(manager._rooms) + custody.assert_disconnected() + + @pytest.mark.parametrize("with_resolver", [True, False]) + async def test_a_failing_reuse_decision_stays_a_safe_non_resume_failure( + self, + tmp_path: Path, + claude: FakeClaude, + caplog: pytest.LogCaptureFixture, + with_resolver: bool, + ) -> None: + """A truth test that raises while the reuse decision is worked out is a + reservation failure the owner built, with or without a resolver: its text + carries no callback value, and the adapter never retries the room without + its history.""" + secret = "res-ul3t-9a7c" + + def band(room_id: str) -> Any: + return dict(BAND) + + async def resolver(workspace: str) -> Mapping[str, Any]: + return {} + + manager = ClaudeSessionManager( + ClaudeAgentOptions(), + mcp_servers_factory=band, + workspaces=RoomWorkspaces(lambda room_id: str(tmp_path / room_id)), + resolve_mcp_servers=resolver if with_resolver else None, + ) + calls: list[str | None] = [] + + async def counted( + room_id: str, *, resume_session_id: str | None = None + ) -> ClaudeSDKClient: + calls.append(resume_session_id) + return await lease_client( + manager, room_id, resume_session_id=resume_session_id + ) + + claude.mcp_statuses.append( + [{"name": BAND_MCP_SERVER_NAME, "status": "connected"}] + ) + client = await lease_client(manager, "room-a") + custody = ClientCustody(client) + probe = RaisingTruth(secret) + real_options = client.options + session = manager.room("room-a") + real_identity = session._protected_servers + try: + if with_resolver: # the retained protected servers are the comparison + session._protected_servers = probe + else: # what the held client dials is the comparison + client.options = HeldServers(probe) + # The decision's truth test runs inside the fallible boundary; the + # branch test after the callback-free guard reads the settled bool. + with ( + caplog.at_level(logging.WARNING), + pytest.raises(ClaudeMcpStartupError) as caught, + ): + await counted("room-a", resume_session_id="hist-1") + finally: + client.options = real_options + if real_identity is not None: + session._protected_servers = real_identity + await manager.stop() + assert probe.reads == 1, f"the decision was read {probe.reads} times" + error = caught.value + assert error.stage == "reservation" and not error.resume_eligible + assert secret not in str(error) + assert secret not in caplog.text + assert calls == ["hist-1"], ( + "the room's session was started again without history" + ) + assert all(not s.alive for s in claude.sessions) + assert not bool(manager._rooms) + custody.assert_disconnected() diff --git a/tests/core/retention/test_shutdown.py b/tests/core/retention/test_shutdown.py index 2defd8d11..c084b0141 100644 --- a/tests/core/retention/test_shutdown.py +++ b/tests/core/retention/test_shutdown.py @@ -122,3 +122,23 @@ async def test_stops_that_stopped_waiting_leave_one_shutdown_failure( ), ) ] + + async def test_owned_shutdown_can_wait_beyond_the_callers_reporting_budget( + self, impatient_host: ResidenceHost + ) -> None: + resident = impatient_host.room() + await resident.turn() + resident.leaving = Hold() + stopping = asyncio.create_task(impatient_host.owner.stop(report_timeout_s=None)) + try: + await asyncio.wait( + {stopping}, timeout=2 * impatient_host.owner.config.transition_timeout_s + ) + if stopping.done(): + await stopping + pytest.fail("owned shutdown returned before disposal settled") + assert resident.identity is not None + finally: + resident.leaving.released.set() + await stopping + assert resident.identity is None diff --git a/tests/integrations/claude_sdk/test_confirmed_close.py b/tests/integrations/claude_sdk/test_confirmed_close.py new file mode 100644 index 000000000..c5d0d128c --- /dev/null +++ b/tests/integrations/claude_sdk/test_confirmed_close.py @@ -0,0 +1,108 @@ +"""The native transport cannot forget a child before its exit is confirmed.""" + +from __future__ import annotations + +import asyncio +import contextlib +import sys +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from pathlib import Path + +import pytest +from claude_agent_sdk import ClaudeAgentOptions + +from band.integrations.claude_sdk import transport +from tests.adapters.claude_sdk.process import ForgetfulProcess, WorkspaceProcess + + +@asynccontextmanager +async def started( + child_type: type[WorkspaceProcess], tmp_path: Path +) -> AsyncIterator[tuple[transport.ConfirmedCloseTransport, WorkspaceProcess]]: + native = child_type( + prompt="", + options=ClaudeAgentOptions( + cli_path=sys.executable, cwd=str(tmp_path), env={"PROBE_MARKER": "close"} + ), + ) + owned = transport.ConfirmedCloseTransport(native) + await owned.connect() + child = native._child + kill, wait = type(child).kill, type(child).wait + try: + yield owned, native + finally: + # Reap with unpatched operations even when the assertion under test fails. + if child.returncode is None: + with contextlib.suppress(ProcessLookupError): + kill(child) + async with asyncio.timeout(10): + await wait(child) + + +class TestConfirmedClose: + async def test_a_close_that_falsely_succeeds_still_reaps_the_child( + self, tmp_path: Path + ) -> None: + async with started(ForgetfulProcess, tmp_path) as (owned, native): + await owned.close() + assert native.exited + assert not owned.close_pending + + async def test_an_unconfirmed_exit_keeps_the_child_for_a_retry( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + async with started(ForgetfulProcess, tmp_path) as (owned, native): + child = native._child + real_kill = child.kill + monkeypatch.setattr( + child, "kill", lambda: (_ for _ in ()).throw(PermissionError()) + ) + with pytest.raises(transport.ClaudeProcessCloseError): + await owned.close() + assert owned.close_pending + assert not native.exited + + monkeypatch.setattr(child, "kill", real_kill) + await owned.close() + assert native.exited + assert not owned.close_pending + + async def test_an_exit_that_never_confirms_fails_within_the_bound( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(transport, "CHILD_EXIT_TIMEOUT_SECONDS", 0.2) + async with started(ForgetfulProcess, tmp_path) as (owned, native): + child = native._child + real_wait = child.wait + + async def never() -> int: + await asyncio.Event().wait() + return 0 + + monkeypatch.setattr(child, "wait", never) + async with asyncio.timeout(5): + with pytest.raises(transport.ClaudeProcessCloseError): + await owned.close() + assert owned.close_pending + + monkeypatch.setattr(child, "wait", real_wait) + await owned.close() + assert native.exited + assert not owned.close_pending + + async def test_a_close_that_raised_keeps_the_child_for_a_retry( + self, tmp_path: Path + ) -> None: + async with started(WorkspaceProcess, tmp_path) as (owned, native): + native.refuse_close = True + with pytest.raises(RuntimeError, match="cleanup failed"): + await owned.close() + assert owned.close_pending + assert not native.exited + + native.refuse_close = False + await owned.close() + assert native.exited + assert not owned.close_pending diff --git a/tests/mcpclient.py b/tests/mcpclient.py index 1765c590f..d52356c24 100644 --- a/tests/mcpclient.py +++ b/tests/mcpclient.py @@ -4,12 +4,13 @@ import asyncio from collections.abc import AsyncIterator, Sequence -from contextlib import AbstractAsyncContextManager, asynccontextmanager +from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager from typing import Any from mcp import ClientSession from mcp.client.sse import sse_client from mcp.client.streamable_http import streamable_http_client +from mcp.shared._httpx_utils import create_mcp_http_client from mcp.types import Tool from band.integrations.mcp import ( @@ -94,19 +95,29 @@ async def crash_backend(owner: SharedBandMCPBackend) -> BandMCPBackend: @asynccontextmanager async def mcp_session( - url: str, transport: BandMCPTransport = BandMCPTransport.HTTP + url: str, + transport: BandMCPTransport = BandMCPTransport.HTTP, + headers: dict[str, str] | None = None, ) -> AsyncIterator[ClientSession]: """An initialized MCP client session to ``url``, closed on exit.""" streams: AbstractAsyncContextManager[tuple[Any, ...]] - match transport: - case BandMCPTransport.HTTP: - streams = streamable_http_client(url) - case BandMCPTransport.SSE: - streams = sse_client(url) - async with ( - streams as (read_stream, write_stream, *_), - ClientSession(read_stream, write_stream) as session, - ): + async with AsyncExitStack() as stack: + match transport: + case BandMCPTransport.HTTP: + http_client = ( + None + if headers is None + else await stack.enter_async_context( + create_mcp_http_client(headers=headers) + ) + ) + streams = streamable_http_client(url, http_client=http_client) + case BandMCPTransport.SSE: + streams = sse_client(url, headers=headers) + read_stream, write_stream, *_ = await stack.enter_async_context(streams) + session = await stack.enter_async_context( + ClientSession(read_stream, write_stream) + ) await session.initialize() yield session From 04ee857ae903d66a2c68348b0d9edb3f4d3fd476 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Mon, 19 Oct 2026 04:19:14 +0300 Subject: [PATCH 08/13] test(claude): cover explicit physical exit custody handoff --- tests/adapters/claude_sdk/process.py | 15 +++++ .../claude_sdk/test_confirmed_close.py | 56 ++++++++++++++++++- 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/tests/adapters/claude_sdk/process.py b/tests/adapters/claude_sdk/process.py index 0d240250a..dcf71299c 100644 --- a/tests/adapters/claude_sdk/process.py +++ b/tests/adapters/claude_sdk/process.py @@ -10,6 +10,7 @@ from claude_agent_sdk import ClaudeAgentOptions from claude_agent_sdk._internal.transport.subprocess_cli import SubprocessCLITransport +from band.core.exceptions import CleanupIncompleteError from band.integrations.claude_sdk import transport from band.runtime.tools import BAND_MCP_SERVER_NAME from tests.adapters.roompeer import CLEANUP_FAILURE @@ -108,3 +109,17 @@ class ForgetfulProcess(WorkspaceProcess): async def close(self) -> None: self._process = None + + +class DelegatingProcess(ForgetfulProcess): + """A native close that requires a physical-exit custodian before returning.""" + + _exit_confirmation_deferred = False + + def _defer_exit_confirmation(self) -> None: + self._exit_confirmation_deferred = True + + async def close(self) -> None: + if not self._exit_confirmation_deferred: + raise CleanupIncompleteError("Claude CLI process exit is unconfirmed") + await super().close() diff --git a/tests/integrations/claude_sdk/test_confirmed_close.py b/tests/integrations/claude_sdk/test_confirmed_close.py index c5d0d128c..05be979ce 100644 --- a/tests/integrations/claude_sdk/test_confirmed_close.py +++ b/tests/integrations/claude_sdk/test_confirmed_close.py @@ -13,7 +13,11 @@ from claude_agent_sdk import ClaudeAgentOptions from band.integrations.claude_sdk import transport -from tests.adapters.claude_sdk.process import ForgetfulProcess, WorkspaceProcess +from tests.adapters.claude_sdk.process import ( + DelegatingProcess, + ForgetfulProcess, + WorkspaceProcess, +) @asynccontextmanager @@ -106,3 +110,53 @@ async def test_a_close_that_raised_keeps_the_child_for_a_retry( await owned.close() assert native.exited assert not owned.close_pending + + async def test_a_native_exit_check_defers_to_the_physical_custodian( + self, tmp_path: Path + ) -> None: + async with started(DelegatingProcess, tmp_path) as (owned, native): + await owned.close() + assert native.exited + assert not owned.close_pending + + async def test_delegated_exit_custody_survives_a_failed_kill( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + async with started(DelegatingProcess, tmp_path) as (owned, native): + child = native._child + with monkeypatch.context() as faults: + faults.setattr( + child, "kill", lambda: (_ for _ in ()).throw(PermissionError()) + ) + with pytest.raises(transport.ClaudeProcessCloseError): + await owned.close() + assert owned.close_pending + assert not native.exited + + await owned.close() + assert native.exited + assert not owned.close_pending + + async def test_a_cancelled_native_close_keeps_physical_custody_for_retry( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + async with started(WorkspaceProcess, tmp_path) as (owned, native): + entered = asyncio.Event() + + async def held_close() -> None: + entered.set() + await asyncio.Event().wait() + + with monkeypatch.context() as faults: + faults.setattr(native, "close", held_close) + closing = asyncio.create_task(owned.close()) + await entered.wait() + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + assert owned.close_pending + assert not native.exited + + await owned.close() + assert native.exited + assert not owned.close_pending From 93d2c1828b3b9538613b83c6f2d209b733e0d52e Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Mon, 19 Oct 2026 04:20:00 +0300 Subject: [PATCH 09/13] fix(claude): hand native exit checks to the physical custodian --- docs/adapters/claude_sdk.md | 2 ++ src/band/integrations/claude_sdk/transport.py | 24 +++++++++++++++---- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/docs/adapters/claude_sdk.md b/docs/adapters/claude_sdk.md index c144686b8..a956ea33a 100644 --- a/docs/adapters/claude_sdk.md +++ b/docs/adapters/claude_sdk.md @@ -137,6 +137,8 @@ reuse the failed client or take another room's claimed directory. retried before a later turn starts a replacement connection. Concurrent `cleanup_all()` calls share one teardown that a cancelled caller can't abandon. + The physical-exit owner accepts the native child before a transport can defer + its own exit check; a deferred check never replaces the owner's kill and reap. ```python import pytest diff --git a/src/band/integrations/claude_sdk/transport.py b/src/band/integrations/claude_sdk/transport.py index fe810306c..c24e8c4a5 100644 --- a/src/band/integrations/claude_sdk/transport.py +++ b/src/band/integrations/claude_sdk/transport.py @@ -5,7 +5,8 @@ import asyncio import contextlib -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Callable +from functools import wraps from typing import Any from claude_agent_sdk import ClaudeAgentOptions @@ -67,6 +68,12 @@ async def end_input(self) -> None: async def close(self) -> None: self._observe_child() + # Accept the child before asking an inner transport to defer its own + # exit assertion; this owner can then kill and reap a forgotten child. + if self._child is not None and ( + defer := getattr(self._inner, "_defer_exit_confirmation", None) + ): + defer() await self._inner.close() await self._confirm_exit() @@ -90,8 +97,17 @@ async def _confirm_exit(self) -> None: self._child = None +def _confirmed_close( + factory: Callable[[ClaudeAgentOptions], Transport], +) -> Callable[[ClaudeAgentOptions], Transport]: + @wraps(factory) + def create(options: ClaudeAgentOptions) -> Transport: + return ConfirmedCloseTransport(factory(options)) + + return create + + +@_confirmed_close def create_transport(options: ClaudeAgentOptions) -> Transport: """Keep permission routing identical to the SDK's default transport.""" - return ConfirmedCloseTransport( - SubprocessCLITransport(prompt="", options=_configure_can_use_tool(options)) - ) + return SubprocessCLITransport(prompt="", options=_configure_can_use_tool(options)) From bcb1eb373ed2b16a75c143cf80ea6b8f3ffab188 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Mon, 19 Oct 2026 04:23:07 +0300 Subject: [PATCH 10/13] test(claude): retain physical close after native wait exhaustion --- .../claude_sdk/test_confirmed_close.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/integrations/claude_sdk/test_confirmed_close.py b/tests/integrations/claude_sdk/test_confirmed_close.py index 05be979ce..79a47e226 100644 --- a/tests/integrations/claude_sdk/test_confirmed_close.py +++ b/tests/integrations/claude_sdk/test_confirmed_close.py @@ -11,6 +11,7 @@ import pytest from claude_agent_sdk import ClaudeAgentOptions +from claude_agent_sdk._internal.transport.subprocess_cli import SubprocessCLITransport from band.integrations.claude_sdk import transport from tests.adapters.claude_sdk.process import ( @@ -160,3 +161,28 @@ async def held_close() -> None: await owned.close() assert native.exited assert not owned.close_pending + + async def test_the_physical_owner_reaps_after_native_waits_forget_the_child( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + close_native = SubprocessCLITransport.close + + async def unconfirmed_wait() -> int: + raise TimeoutError + + async def forgotten_exit(native: SubprocessCLITransport) -> None: + child = native._process + assert child is not None + with monkeypatch.context() as faults: + faults.setattr(child, "wait", unconfirmed_wait) + faults.setattr(child, "terminate", lambda: None) + faults.setattr(child, "kill", lambda: None) + await close_native(native) + + async with started(WorkspaceProcess, tmp_path) as (owned, native): + with monkeypatch.context() as faults: + faults.setattr(SubprocessCLITransport, "close", forgotten_exit) + await owned.close() + assert native.exited + assert not owned.close_pending + await owned.close() From 1cbf38d72618074a289390873b0e0a0f385b52ac Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Mon, 19 Oct 2026 04:29:07 +0300 Subject: [PATCH 11/13] refactor(claude): separate native construction from exit custody --- src/band/integrations/claude_sdk/transport.py | 24 +++++++------------ 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/src/band/integrations/claude_sdk/transport.py b/src/band/integrations/claude_sdk/transport.py index c24e8c4a5..72256f9f9 100644 --- a/src/band/integrations/claude_sdk/transport.py +++ b/src/band/integrations/claude_sdk/transport.py @@ -5,8 +5,7 @@ import asyncio import contextlib -from collections.abc import AsyncIterator, Callable -from functools import wraps +from collections.abc import AsyncIterator from typing import Any from claude_agent_sdk import ClaudeAgentOptions @@ -14,6 +13,12 @@ from claude_agent_sdk._internal.transport.subprocess_cli import SubprocessCLITransport from claude_agent_sdk.types import _configure_can_use_tool + +def _create_native_transport(options: ClaudeAgentOptions) -> Transport: + """Keep permission routing identical to the SDK's default transport.""" + return SubprocessCLITransport(prompt="", options=_configure_can_use_tool(options)) + + CHILD_EXIT_TIMEOUT_SECONDS = 5.0 """How long a child the native close left running may take to die once killed.""" @@ -97,17 +102,6 @@ async def _confirm_exit(self) -> None: self._child = None -def _confirmed_close( - factory: Callable[[ClaudeAgentOptions], Transport], -) -> Callable[[ClaudeAgentOptions], Transport]: - @wraps(factory) - def create(options: ClaudeAgentOptions) -> Transport: - return ConfirmedCloseTransport(factory(options)) - - return create - - -@_confirmed_close def create_transport(options: ClaudeAgentOptions) -> Transport: - """Keep permission routing identical to the SDK's default transport.""" - return SubprocessCLITransport(prompt="", options=_configure_can_use_tool(options)) + """Keep native permission routing under physical child custody.""" + return ConfirmedCloseTransport(_create_native_transport(options)) From 2a2361057d4806fea14c224c4b8e9e05965dfb56 Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Mon, 19 Oct 2026 04:35:45 +0300 Subject: [PATCH 12/13] refactor(claude): own room subprocess closure at session construction --- src/band/integrations/claude_sdk/custody.py | 93 ++++++++++++++++++ .../claude_sdk/session_manager.py | 5 +- src/band/integrations/claude_sdk/transport.py | 98 +------------------ .../claude_sdk/test_confirmed_close.py | 14 +-- 4 files changed, 106 insertions(+), 104 deletions(-) create mode 100644 src/band/integrations/claude_sdk/custody.py diff --git a/src/band/integrations/claude_sdk/custody.py b/src/band/integrations/claude_sdk/custody.py new file mode 100644 index 000000000..613bb2653 --- /dev/null +++ b/src/band/integrations/claude_sdk/custody.py @@ -0,0 +1,93 @@ +"""Physical custody of a room's Claude CLI until its exit is confirmed.""" + +from __future__ import annotations + +import asyncio +import contextlib +from collections.abc import AsyncIterator +from typing import Any + +from claude_agent_sdk._internal.transport import Transport + + +CHILD_EXIT_TIMEOUT_SECONDS = 5.0 +"""How long a child the native close left running may take to die once killed.""" + + +class ClaudeProcessCloseError(RuntimeError): + """The Claude CLI child could not be confirmed exited; its transport keeps + the child for a later ``close()`` to retry.""" + + +class ConfirmedCloseTransport(Transport): + """The SDK's subprocess transport, holding the CLI child until it is reaped. + + The native ``close()`` can return after its waits time out with the child + still running, and it forgets the child either way, so neither its return + nor a cleared handle proves exit. This wrapper keeps the child it saw + spawned, kills and reaps one that a normal close left running, and keeps + the child (raising ``ClaudeProcessCloseError``) when that cannot be + confirmed. A close that raised or was cancelled also keeps the child, so a + later ``close()`` retries. + """ + + def __init__(self, inner: Transport) -> None: + self._inner = inner + self._child: Any = None + + @property + def close_pending(self) -> bool: + """Whether a child is held whose exit is not yet confirmed.""" + return self._child is not None + + def _observe_child(self) -> None: + if (process := getattr(self._inner, "_process", None)) is not None: + self._child = process + + async def connect(self) -> None: + try: + await self._inner.connect() + finally: + self._observe_child() + + async def write(self, data: str) -> None: + await self._inner.write(data) + + def read_messages(self) -> AsyncIterator[dict[str, Any]]: + return self._inner.read_messages() + + def is_ready(self) -> bool: + return self._inner.is_ready() + + async def end_input(self) -> None: + await self._inner.end_input() + + async def close(self) -> None: + self._observe_child() + # Accept the child before asking an inner transport to defer its own + # exit assertion; this owner can then kill and reap a forgotten child. + if self._child is not None and ( + defer := getattr(self._inner, "_defer_exit_confirmation", None) + ): + defer() + await self._inner.close() + await self._confirm_exit() + + async def _confirm_exit(self) -> None: + """Kill and reap a child the native close left running, within + ``CHILD_EXIT_TIMEOUT_SECONDS``; otherwise keep it and raise.""" + child = self._child + if child is None: + return + if child.returncode is None: + try: + with contextlib.suppress(ProcessLookupError): + child.kill() + async with asyncio.timeout(CHILD_EXIT_TIMEOUT_SECONDS): + await child.wait() + except Exception as error: # noqa: BLE001 -- unconfirmed exit, whatever the cause + raise ClaudeProcessCloseError( + "the Claude CLI process could not be confirmed exited" + f" ({type(error).__name__})" + ) from None + self._child = None diff --git a/src/band/integrations/claude_sdk/session_manager.py b/src/band/integrations/claude_sdk/session_manager.py index d6eed97c3..eb35f4f91 100644 --- a/src/band/integrations/claude_sdk/session_manager.py +++ b/src/band/integrations/claude_sdk/session_manager.py @@ -54,6 +54,7 @@ Message, ) + from band.integrations.claude_sdk.custody import ConfirmedCloseTransport from band.integrations.claude_sdk.transport import create_transport _CLAUDE_SDK_AVAILABLE = True @@ -403,7 +404,9 @@ async def _connect(self, generation: int, *, resume: str | None) -> None: else: options = self._manager.options_for(self, generation, resume=resume) self._check_current(generation) - self.transport = transport = create_transport(options) + self.transport = transport = ConfirmedCloseTransport( + create_transport(options) + ) self.client = client = ClaudeSDKClient(options=options, transport=transport) # Retain both before connect: SDK disconnect cannot close a child # when startup fails before its Query is constructed. diff --git a/src/band/integrations/claude_sdk/transport.py b/src/band/integrations/claude_sdk/transport.py index 72256f9f9..a157c0478 100644 --- a/src/band/integrations/claude_sdk/transport.py +++ b/src/band/integrations/claude_sdk/transport.py @@ -1,107 +1,13 @@ -"""Claude subprocess construction, including SDK permission normalization, and -custody of the CLI child until its exit is confirmed.""" +"""Claude subprocess construction, including SDK permission normalization.""" from __future__ import annotations -import asyncio -import contextlib -from collections.abc import AsyncIterator -from typing import Any - from claude_agent_sdk import ClaudeAgentOptions from claude_agent_sdk._internal.transport import Transport from claude_agent_sdk._internal.transport.subprocess_cli import SubprocessCLITransport from claude_agent_sdk.types import _configure_can_use_tool -def _create_native_transport(options: ClaudeAgentOptions) -> Transport: +def create_transport(options: ClaudeAgentOptions) -> Transport: """Keep permission routing identical to the SDK's default transport.""" return SubprocessCLITransport(prompt="", options=_configure_can_use_tool(options)) - - -CHILD_EXIT_TIMEOUT_SECONDS = 5.0 -"""How long a child the native close left running may take to die once killed.""" - - -class ClaudeProcessCloseError(RuntimeError): - """The Claude CLI child could not be confirmed exited; its transport keeps - the child for a later ``close()`` to retry.""" - - -class ConfirmedCloseTransport(Transport): - """The SDK's subprocess transport, holding the CLI child until it is reaped. - - The native ``close()`` can return after its waits time out with the child - still running, and it forgets the child either way, so neither its return - nor a cleared handle proves exit. This wrapper keeps the child it saw - spawned, kills and reaps one that a normal close left running, and keeps - the child (raising ``ClaudeProcessCloseError``) when that cannot be - confirmed. A close that raised or was cancelled also keeps the child, so a - later ``close()`` retries. - """ - - def __init__(self, inner: Transport) -> None: - self._inner = inner - self._child: Any = None - - @property - def close_pending(self) -> bool: - """Whether a child is held whose exit is not yet confirmed.""" - return self._child is not None - - def _observe_child(self) -> None: - if (process := getattr(self._inner, "_process", None)) is not None: - self._child = process - - async def connect(self) -> None: - try: - await self._inner.connect() - finally: - self._observe_child() - - async def write(self, data: str) -> None: - await self._inner.write(data) - - def read_messages(self) -> AsyncIterator[dict[str, Any]]: - return self._inner.read_messages() - - def is_ready(self) -> bool: - return self._inner.is_ready() - - async def end_input(self) -> None: - await self._inner.end_input() - - async def close(self) -> None: - self._observe_child() - # Accept the child before asking an inner transport to defer its own - # exit assertion; this owner can then kill and reap a forgotten child. - if self._child is not None and ( - defer := getattr(self._inner, "_defer_exit_confirmation", None) - ): - defer() - await self._inner.close() - await self._confirm_exit() - - async def _confirm_exit(self) -> None: - """Kill and reap a child the native close left running, within - ``CHILD_EXIT_TIMEOUT_SECONDS``; otherwise keep it and raise.""" - child = self._child - if child is None: - return - if child.returncode is None: - try: - with contextlib.suppress(ProcessLookupError): - child.kill() - async with asyncio.timeout(CHILD_EXIT_TIMEOUT_SECONDS): - await child.wait() - except Exception as error: # noqa: BLE001 -- unconfirmed exit, whatever the cause - raise ClaudeProcessCloseError( - "the Claude CLI process could not be confirmed exited" - f" ({type(error).__name__})" - ) from None - self._child = None - - -def create_transport(options: ClaudeAgentOptions) -> Transport: - """Keep native permission routing under physical child custody.""" - return ConfirmedCloseTransport(_create_native_transport(options)) diff --git a/tests/integrations/claude_sdk/test_confirmed_close.py b/tests/integrations/claude_sdk/test_confirmed_close.py index 79a47e226..6ea078a35 100644 --- a/tests/integrations/claude_sdk/test_confirmed_close.py +++ b/tests/integrations/claude_sdk/test_confirmed_close.py @@ -13,7 +13,7 @@ from claude_agent_sdk import ClaudeAgentOptions from claude_agent_sdk._internal.transport.subprocess_cli import SubprocessCLITransport -from band.integrations.claude_sdk import transport +from band.integrations.claude_sdk import custody from tests.adapters.claude_sdk.process import ( DelegatingProcess, ForgetfulProcess, @@ -24,14 +24,14 @@ @asynccontextmanager async def started( child_type: type[WorkspaceProcess], tmp_path: Path -) -> AsyncIterator[tuple[transport.ConfirmedCloseTransport, WorkspaceProcess]]: +) -> AsyncIterator[tuple[custody.ConfirmedCloseTransport, WorkspaceProcess]]: native = child_type( prompt="", options=ClaudeAgentOptions( cli_path=sys.executable, cwd=str(tmp_path), env={"PROBE_MARKER": "close"} ), ) - owned = transport.ConfirmedCloseTransport(native) + owned = custody.ConfirmedCloseTransport(native) await owned.connect() child = native._child kill, wait = type(child).kill, type(child).wait @@ -64,7 +64,7 @@ async def test_an_unconfirmed_exit_keeps_the_child_for_a_retry( monkeypatch.setattr( child, "kill", lambda: (_ for _ in ()).throw(PermissionError()) ) - with pytest.raises(transport.ClaudeProcessCloseError): + with pytest.raises(custody.ClaudeProcessCloseError): await owned.close() assert owned.close_pending assert not native.exited @@ -77,7 +77,7 @@ async def test_an_unconfirmed_exit_keeps_the_child_for_a_retry( async def test_an_exit_that_never_confirms_fails_within_the_bound( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr(transport, "CHILD_EXIT_TIMEOUT_SECONDS", 0.2) + monkeypatch.setattr(custody, "CHILD_EXIT_TIMEOUT_SECONDS", 0.2) async with started(ForgetfulProcess, tmp_path) as (owned, native): child = native._child real_wait = child.wait @@ -88,7 +88,7 @@ async def never() -> int: monkeypatch.setattr(child, "wait", never) async with asyncio.timeout(5): - with pytest.raises(transport.ClaudeProcessCloseError): + with pytest.raises(custody.ClaudeProcessCloseError): await owned.close() assert owned.close_pending @@ -129,7 +129,7 @@ async def test_delegated_exit_custody_survives_a_failed_kill( faults.setattr( child, "kill", lambda: (_ for _ in ()).throw(PermissionError()) ) - with pytest.raises(transport.ClaudeProcessCloseError): + with pytest.raises(custody.ClaudeProcessCloseError): await owned.close() assert owned.close_pending assert not native.exited From 17c575a7edc4b8c83e464844da373ef07ed7523d Mon Sep 17 00:00:00 2001 From: bandzalkin Date: Mon, 19 Oct 2026 04:42:20 +0300 Subject: [PATCH 13/13] test(claude): capture native launches across transport wrappers --- src/band/integrations/claude_sdk/custody.py | 1 - tests/adapters/claude_sdk/test_cli_options.py | 13 ++++++------- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/src/band/integrations/claude_sdk/custody.py b/src/band/integrations/claude_sdk/custody.py index 613bb2653..27a9ff554 100644 --- a/src/band/integrations/claude_sdk/custody.py +++ b/src/band/integrations/claude_sdk/custody.py @@ -9,7 +9,6 @@ from claude_agent_sdk._internal.transport import Transport - CHILD_EXIT_TIMEOUT_SECONDS = 5.0 """How long a child the native close left running may take to die once killed.""" diff --git a/tests/adapters/claude_sdk/test_cli_options.py b/tests/adapters/claude_sdk/test_cli_options.py index 9b3c7305f..bf41cbd2e 100644 --- a/tests/adapters/claude_sdk/test_cli_options.py +++ b/tests/adapters/claude_sdk/test_cli_options.py @@ -22,7 +22,6 @@ ClaudeSDKAdapterConfig, ) from band.core.types import Emit -from band.integrations.claude_sdk import transport as transport_module from band.integrations.claude_sdk.transport import create_transport from tests.adapters.claude_sdk.helpers import ( SEND_MESSAGE_MCP_NAME, @@ -70,15 +69,15 @@ async def allow( return PermissionResultAllow() spawned: list[SubprocessCLITransport] = [] + initialize = SubprocessCLITransport.__init__ - def spawn( - *, prompt: str, options: ClaudeAgentOptions - ) -> SubprocessCLITransport: - native = SubprocessCLITransport(prompt=prompt, options=options) + def capture( + native: SubprocessCLITransport, *, prompt: str, options: ClaudeAgentOptions + ) -> None: + initialize(native, prompt=prompt, options=options) spawned.append(native) - return native - monkeypatch.setattr(transport_module, "SubprocessCLITransport", spawn) + monkeypatch.setattr(SubprocessCLITransport, "__init__", capture) options = ClaudeAgentOptions(cli_path=sys.executable, can_use_tool=allow) create_transport(options)