From e38e7ffb63dfd2daeeb415c9da3a6550d8ef9ed0 Mon Sep 17 00:00:00 2001 From: Natan Zalkin Date: Thu, 15 Oct 2026 05:34:49 +0300 Subject: [PATCH 1/2] feat(codex): resolve ambient servers inside owned room startup --- docs/adapters/codex.md | 67 ++ 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 | 80 ++- 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 | 18 + .../integrations/codex/test_mcp_redaction.py | 252 +++++++ .../codex/test_process_group_cleanup.py | 510 ++++++++++++++ 23 files changed, 4182 insertions(+), 80 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 1377ce05d..a6ae71a34 100644 --- a/docs/adapters/codex.md +++ b/docs/adapters/codex.md @@ -35,6 +35,15 @@ For existing integrations, follow the [SDK 4.x to 5.x migration](../acp.md#upgra close errors and close timeouts retain the claim for retry. A timed-out or cancelled caller does not abandon subprocess cleanup; retry room cleanup to finish releasing ownership. Shutdown attempts every room and reports failures. +- **POSIX cleanup settles the original owned process group.** After normal leader + exit or an exit timeout, close sends TERM, waits up to one second, then sends + KILL and waits up to one more second for `ESRCH`. Unconfirmed cleanup retains + the client and workspace for retry. The group ID comes from the original + subprocess handle, never a lookup of a possibly recycled PID. POSIX groups + lack generation-bound handles: reuse between the group emptying and a final + probe or signal remains a residual race. Descendants that leave the group are + outside this guarantee. Windows retains best-effort `taskkill /F /T` on leader + timeout, not confirmation of orphaned or escaped descendants. - **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 @@ -105,3 +114,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 b9115e676..2444bb6c5 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, @@ -68,7 +69,6 @@ ApprovalInterruptMixin, TurnSettleTimeout, TurnTask, - bind_turn_task, start_releasable_turn, ) from band.core.types import ( @@ -87,6 +87,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, @@ -617,6 +624,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. @@ -633,6 +669,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 @@ -909,6 +947,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], @@ -935,6 +974,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. @@ -1180,16 +1220,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 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, @@ -1203,15 +1250,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. @@ -1220,6 +1287,7 @@ async def on_message( async def _run_turn( self, *, + owner_ready: asyncio.Future[TurnTask], room: RoomCodexClient, msg: PlatformMessage, tools: AgentToolsProtocol, @@ -1230,6 +1298,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( @@ -1252,12 +1321,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. @@ -1479,6 +1549,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, @@ -1866,7 +1939,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 @@ -1878,38 +1964,167 @@ 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() + @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. @@ -2115,7 +2330,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 @@ -2153,11 +2375,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: @@ -2179,6 +2406,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, *, @@ -2241,11 +2514,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) @@ -2293,7 +2567,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 fef3c34b8..e8b65e7e4 100644 --- a/src/band/core/retention/room.py +++ b/src/band/core/retention/room.py @@ -245,6 +245,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: @@ -329,6 +331,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() @@ -354,7 +358,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) @@ -531,7 +536,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 @@ -602,7 +611,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 58657f596..7c8afab97 100644 --- a/src/band/core/turn_lifecycle.py +++ b/src/band/core/turn_lifecycle.py @@ -54,31 +54,58 @@ class TurnTask: def __init__(self, task: asyncio.Task[None]) -> 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[None]) -> 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() @@ -88,19 +115,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: @@ -116,8 +148,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() @@ -125,6 +162,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()) @@ -136,7 +175,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 6fb3602c4..d755613e8 100644 --- a/src/band/integrations/codex/stdio_client.py +++ b/src/band/integrations/codex/stdio_client.py @@ -12,26 +12,64 @@ from contextlib import suppress 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 logger = logging.getLogger(__name__) _PROCESS_EXIT_TIMEOUT_S = 1.0 +_PROCESS_GROUP_EXIT_TIMEOUT_S = 1.0 +_PROCESS_GROUP_POLL_S = 0.02 + + +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 ``process`` and everything it spawned. + """Settle the original owned POSIX group; Windows tree kill is best effort. - The ``codex`` launcher is a shim (``codex.cmd`` on Windows, a node script - elsewhere) whose real CLI is its child; killing only the shim leaves that - child running, and on Windows its open workdir cannot be removed. + ``start_new_session`` makes the original process handle's PID its group ID. + Never rediscover a group from a current PID, which may already be recycled. + POSIX has no generation-bound group handle: reuse after the last member exits + and before the final probe remains a residual race. """ if os.name != "nt": - # The child leads its own session, so its pid is the group id. - with suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) - 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" + ) taskkill = await asyncio.create_subprocess_exec( "taskkill", "/F", @@ -57,8 +95,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 @@ -131,10 +170,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 = [ @@ -151,9 +191,18 @@ async def _close_process(self) -> None: await asyncio.wait_for(process.wait(), timeout=_PROCESS_EXIT_TIMEOUT_S) except TimeoutError: await kill_process_tree(process) + # The kill is forced, so the leader does exit; a loaded loop can + # report it late. Callers are bounded by the room owner's budget, + # which keeps this close owned for a retry to join. await process.wait() + else: + if os.name != "nt": + await kill_process_tree(process) self._fail_pending("Codex stdio client closed") + 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: @@ -190,9 +239,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 cda37d249..b89ab4b6f 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 8974badda..baa370adc 100644 --- a/tests/adapters/codexturns.py +++ b/tests/adapters/codexturns.py @@ -257,7 +257,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 ab9d0669c..d3b7b396f 100644 --- a/tests/integrations/codex/test_cleanup.py +++ b/tests/integrations/codex/test_cleanup.py @@ -136,6 +136,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: 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 1b61d7ca1a440e9be0e86d51fb5cb3c6b2e59ba8 Mon Sep 17 00:00:00 2001 From: Natan Zalkin Date: Thu, 15 Oct 2026 16:44:28 +0300 Subject: [PATCH 2/2] 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 ed6027494..0ef9da7bf 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. + ## [5.1.0](https://github.com/band-ai/band-sdk-python/compare/band-sdk-v5.0.0...band-sdk-v5.1.0) (2026-10-13) diff --git a/docs/adapters/codex.md b/docs/adapters/codex.md index a6ae71a34..f8d47f354 100644 --- a/docs/adapters/codex.md +++ b/docs/adapters/codex.md @@ -153,10 +153,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 2444bb6c5..307f70598 100644 --- a/src/band/adapters/codex.py +++ b/src/band/adapters/codex.py @@ -1027,10 +1027,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 @@ -1944,6 +1959,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() @@ -1970,6 +1986,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 e8b65e7e4..27c40b49f 100644 --- a/src/band/core/retention/room.py +++ b/src/band/core/retention/room.py @@ -63,6 +63,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__( @@ -72,12 +76,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 @@ -358,8 +365,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) @@ -611,12 +628,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 7c8afab97..ee105eb83 100644 --- a/src/band/core/turn_lifecycle.py +++ b/src/band/core/turn_lifecycle.py @@ -57,6 +57,7 @@ def __init__(self, task: asyncio.Task[None]) -> 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 @@ -94,11 +95,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 @@ -156,7 +166,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")