diff --git a/CHANGELOG.md b/CHANGELOG.md index f47ec0ac..b3fac25b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,37 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 unset those routes still fall back to the bearer token, so an existing single-token deployment keeps working until it sets the new variable. +- **Session close retained stdio children and per-session upstream caches.** + (`#625`) `rebind_session()` rotated the audit chain and call logs but never + called `aclose()`, so the stdio child, pooled HTTP clients, and the + provenance/drift-checked caches all survived a close and were reused by the + next session - the exact cross-session contamination `docs/spec/stdio-transport.md` + names as the reason children are scoped to one session rather than pooled. + Close now serializes session transitions, drains admitted calls before + signing, cleans up resources before rebinding, and preserves claim and + resource ownership across retryable failures. Partial claim failures keep + admission sealed and are reported for operator investigation. A failed + cancellation drain also keeps admission sealed until a close retry can + drain the remaining work. Graceful shutdown rejects new work and resource + acquisition, drains active calls, and coordinates spawning with cleanup; + an incomplete drain is reported as failure. Concurrent first-use stdio + spawning is serialized to avoid creating an untracked second child. + Cancellation during session hydration is finalized before signing; a failed + terminal audit write prevents signing or rotating an incomplete claim. + `POST /sessions/{id}/reset` retires a session id and opens a successor, so it + leaked the same resources for the same reason; it now drains admitted calls + and releases them before recording the boundary, which also leaves a failed + reset retryable with the session untouched. A reset naming an already-rotated + session is rejected before draining, so it cannot cancel the successor's + in-flight calls. A child that fails to close is retained for retry; a pooled + HTTP client that fails to close is dropped and logged instead, because an + `AsyncClient` marks itself closed and HTTPcore empties its pool before the + streams are released, leaving nothing a retry could reach. A call arriving + during a transition still waits for the successor, but the wait is bounded: + a close that cannot resolve, such as one whose successor creation keeps + failing, now answers callers with the reason instead of blocking them + indefinitely. + ### Added - **The accumulated session-sensitivity value can now live in a shared, diff --git a/LIMITATIONS.md b/LIMITATIONS.md index 700f4ed0..78b6411e 100644 --- a/LIMITATIONS.md +++ b/LIMITATIONS.md @@ -67,6 +67,15 @@ The TEE prevents plaintext from leaving the enclave to any destination not cover **Tool name collision via malicious catalog entries** The catalog binds each tool name to a specific upstream server identity, which prevents routing ambiguity for approved servers. It does not prevent a typosquatted or look-alike package from being added to the catalog in the first place. Catalog approval is human-gated. The gateway trusts the catalog; it cannot detect that a catalog entry was added via a compromised reviewer or a social engineering attack. +**Session cleanup is bounded by cooperation and by time** +Session-scoped resources, meaning the stdio child, the pooled HTTP clients, and the provenance and drift caches, are released on every path that ends a session: `POST /sessions/{id}/close`, `POST /sessions/{id}/reset`, and graceful shutdown. None of those paths released anything before this was implemented, so the first three limits below are what remains rather than what was added. The last two are deliberate trades the behaviour introduces. + +- **A pooled HTTP client that fails to close leaks its connections.** `AsyncClient` marks itself closed, and HTTPcore empties its pool, before the underlying streams are released, so nothing a retry could reach survives a failed close. The client is dropped so the successor cannot reuse it, and the failure is logged. A child process that fails to close is retained instead, and a close retry can still reap it. +- **Graceful shutdown can outlast a deployment's termination grace period.** It waits for any in-flight close, then drains again on the same budget, so with the defaults cleanup can begin as late as seventy seconds in. A shorter grace period ends in SIGKILL and none of this runs. Size the grace period above twice `CMCP_SESSION_CLOSE_DRAIN_SECONDS`, or lower that deadline. +- **Cancellation is cooperative, so a failed drain does not prove a call stopped.** Close requests cancellation at the deadline and allows a further five seconds to unwind. A call that does not honour it leaves the drain incomplete, which seals admission rather than signing a claim that omits an outcome. +- **A failed terminal audit write leaves the session unavailable, with no repair.** A deliberate trade: it blocks signing, rotation, reset, and further admission for that session, because the alternative is a signed claim missing a call the gateway made. Restoring the writer does not reconstruct the missing outcome, and none is provided. Recovery is a new session. +- **A close that trips the kill switch leaves the gateway with no live session.** Also deliberate. The claim for the closed session is signed and retrievable, but no successor can be created until an operator unblocks that agent identity. This is the kill switch working as specified, at the cost of availability. + ## What Level 0 (CMCP_DEV_MODE) does not provide `CMCP_DEV_MODE=1` uses a software-only TEE provider. It is suitable for development, testing, and demo scenarios. It does not satisfy production governance requirements because: diff --git a/docs/spec/stdio-transport.md b/docs/spec/stdio-transport.md index 51bc6507..ece5570e 100644 --- a/docs/spec/stdio-transport.md +++ b/docs/spec/stdio-transport.md @@ -111,7 +111,46 @@ the launch measurement and must be reported as a distinct evidence class, not fo 1. **Lifecycle.** Implemented as children scoped to a session, reused by execution identity within it, and closed with that session. The original alternative was a pool across sessions. A pool is faster and leaks state between sessions, which is exactly the kind of cross-session - contamination the audit chain cannot see. + contamination the audit chain cannot see. "Closed with that session" covers every way a + session ends: an explicit `POST /sessions/{id}/close`, the gateway process exiting + with a session still live during graceful shutdown, and `POST /sessions/{id}/reset`, + which also retires a session id and opens a successor. Reset drains admitted calls + and releases the same session-scoped resources before it records the boundary, so + the successor never inherits a child, a pooled client, or a provenance entry from + the session it replaced. If a child fails to close, the current session ID and + audit boundary remain unchanged, admission stays sealed, and a retry closes the + retained child before recording the reset. + A call that arrives during a transition waits for it and is admitted to the + successor, but that wait is bounded: a transition that has already failed is + lifted only by a close retry or operator action, so a call waiting past the + bound is answered with the reason rather than held on an open socket. + Close blocks new calls, + waits up to `CMCP_SESSION_CLOSE_DRAIN_SECONDS` (default 30 seconds), then + requests cancellation and allows a further five seconds for calls to unwind. + If calls remain, close fails with `SessionDrainIncomplete` and admission stays + sealed; a retry must drain them before signing and rebinding. Partial claim + failure also seals admission and requires operator investigation. Successful + task completion alone does not prove audit completeness: a failed terminal + audit write prevents signing, rotation, reset, and further call admission. + Hydration failures and cancellations are included in terminal finalization. + Shutdown can still release resources without signing an incomplete claim. + Successful cleanup precedes rebinding; a child that fails to close is retained + for retry. Pooled HTTP clients are closed on a best-effort basis instead: an + `AsyncClient` marks itself closed and HTTPcore empties its pool before the + underlying streams are released, so a failed close leaves connections no retry + reaches through any public API. Such a client is dropped and the failure logged + rather than sealing the session, because it is reuse by the successor, not the + socket, that this lifecycle rule exists to prevent. + Graceful shutdown permanently rejects new work and resource acquisition, + drains active calls, and serializes spawning with cleanup. An incomplete + drain, or a child that could not be closed, is reported as shutdown failure + rather than success; a pooled client that could not be closed is logged and + does not fail the shutdown, for the reason given above. Hard + process termination cannot run this cleanup, and a shutdown that waits out an + in-flight close can need twice the drain budget before cleanup begins, so a + deployment's termination grace period has to exceed it or the cleanup is cut + short by the kill. These drain deadlines do not + bound arbitrary signing or transport cleanup time. 2. **stderr.** The implementation logs diagnostics through the gateway logger and records a byte count in evidence. MCP servers write diagnostics there. Capturing it into the audit chain risks payload leakage into an artifact meant to be shareable; discarding it loses the only signal when a child misbehaves. diff --git a/src/cmcp_runtime/errors.py b/src/cmcp_runtime/errors.py index f76509f6..75e7a819 100644 --- a/src/cmcp_runtime/errors.py +++ b/src/cmcp_runtime/errors.py @@ -130,6 +130,24 @@ class TeeFault(CMCPError): http_status = 500 +class SessionCloseIncomplete(CMCPError): + """Terminal audit or close bookkeeping failed without a safe recovery. + + Repeating accounting/signing is unsafe; operator investigation is required.""" + + code = "SESSION_CLOSE_INCOMPLETE" + http_status = 500 + + +class SessionDrainIncomplete(CMCPError): + """Calls remain active after the drain deadline and cancellation grace. + + Admission stays sealed; a transition retry must finish draining first.""" + + code = "SESSION_DRAIN_INCOMPLETE" + http_status = 503 + + class UpstreamUnavailable(CMCPError): code = "UPSTREAM_UNAVAILABLE" http_status = 502 diff --git a/src/cmcp_runtime/mcp/proxy.py b/src/cmcp_runtime/mcp/proxy.py index fd90eba0..1b2c251f 100644 --- a/src/cmcp_runtime/mcp/proxy.py +++ b/src/cmcp_runtime/mcp/proxy.py @@ -15,10 +15,12 @@ import hashlib import json import logging +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import UTC, datetime from enum import StrEnum -from typing import Any +from typing import Any, NoReturn import httpx @@ -31,7 +33,13 @@ ) from cmcp_runtime.catalog.scanner import CatalogScanner from cmcp_runtime.config import Config, DriftPolicy -from cmcp_runtime.errors import PolicyDeny, UpstreamToolError, UpstreamUnavailable +from cmcp_runtime.errors import ( + PolicyDeny, + SessionCloseIncomplete, + SessionDrainIncomplete, + UpstreamToolError, + UpstreamUnavailable, +) from cmcp_runtime.execution import valid_execution_id from cmcp_runtime.mcp import tls_pinning from cmcp_runtime.mcp.stdio import StdioServer @@ -49,6 +57,7 @@ logger = logging.getLogger(__name__) + _EXTERNAL_EVIDENCE_FIELDS: frozenset[str] = frozenset( { "issuer", @@ -205,6 +214,14 @@ def _extract_external_execution_evidence(response_text: str) -> dict[str, str] | return {field: receipt[field] for field in sorted(_EXTERNAL_EVIDENCE_FIELDS)} +# Default matches the existing upstream HTTP timeout. +SESSION_CLOSE_DRAIN_SECONDS = 30.0 + +# Cancellation is cooperative; failure after this grace keeps admission sealed. +SESSION_CANCELLATION_GRACE_SECONDS = 5.0 + + + class CMCPProxy: """ Enforces every tool call through the cMCP runtime gateway: @@ -287,18 +304,296 @@ def __init__( # is not, and a check expensive enough to hurt is a check that gets disabled. self._drift_checked: set[tuple[str, ...]] = set() self._catalog_scanner = catalog_scanner + # #625: serialises first-use stdio spawns so two calls racing on the + # same server's first use cannot both spawn a child - see `_stdio_for`. + self._stdio_spawn_lock = asyncio.Lock() + + # Calls drain before signing; failed transitions may keep admission sealed. + self._lifecycle_condition = asyncio.Condition() + # Ownership ends on failure; admission may remain sealed for retry. + self._transition_lock = asyncio.Lock() + self._active_calls = 0 + self._active_call_tasks: set[asyncio.Task[Any]] = set() + self._session_rotation_in_progress = False + self._session_rebound = False + self._close_committed = False + self._drain_incomplete = False + self._cleanup_incomplete = False + # The drain deadline the current transition runs under. A waiting call + # bounds itself by this, so raising the configured deadline does not + # start rejecting calls an ordinary close would have admitted. + self._transition_drain_s = SESSION_CLOSE_DRAIN_SECONDS + # No safe reconstruction is available for a terminal that failed to + # persist. Keep this separate from drain state: shutdown can still reap. + self._failed_terminal_call: str | None = None + self._shutting_down = False + + def _ensure_running(self) -> None: + if self._shutting_down: + raise UpstreamUnavailable("gateway is shutting down") + + def _ensure_terminal_audit_complete(self) -> None: + if self._failed_terminal_call is not None: + raise SessionCloseIncomplete( + "A call's terminal audit write failed; signing or rotating this " + "session would omit an outcome. Operator investigation is required.", + detail=self._failed_terminal_call, + ) + + def _admission_wait_s(self) -> float: + """How long to wait out a transition: its own deadline, plus the grace.""" + return self._transition_drain_s + SESSION_CANCELLATION_GRACE_SECONDS - def rebind_session(self, session: SessionState, audit_chain: AuditChain) -> None: + def _raise_stuck_transition(self) -> NoReturn: + """Explain the transition a waiting call gave up on. + + Waiting out a rotation is ordinary back-pressure; the successor admits + the call. Waiting on one that has already failed is not: only a close + retry or operator action lifts it, so the caller is told what to fix. + """ + if self._drain_incomplete: + raise SessionDrainIncomplete( + "calls from the closing session did not drain; a close retry " + "must finish draining before this gateway admits work again" + ) + if self._close_committed: + raise SessionCloseIncomplete( + "the session was closed but no successor was adopted; retry " + "the close once the cause of the failure is cleared" + ) + raise SessionCloseIncomplete( + f"a session transition did not complete within " + f"{self._admission_wait_s():g}s; the gateway is not admitting calls" + ) + + async def _enter_call(self) -> None: + async with self._lifecycle_condition: + # Only a transition in flight is worth waiting on, and waiting is + # the whole cost here: the deadline is set up per call otherwise. + if self._session_rotation_in_progress: + try: + await asyncio.wait_for( + self._lifecycle_condition.wait_for( + lambda: self._shutting_down + or self._failed_terminal_call is not None + or not self._session_rotation_in_progress + ), + timeout=self._admission_wait_s(), + ) + except TimeoutError: + self._raise_stuck_transition() + self._ensure_running() + self._ensure_terminal_audit_complete() + self._active_calls += 1 + task = asyncio.current_task() + if task is not None: + self._active_call_tasks.add(task) + + async def _leave_call(self) -> None: + async with self._lifecycle_condition: + self._active_calls -= 1 + self._active_call_tasks.discard(asyncio.current_task()) + if self._active_calls == 0 or self._failed_terminal_call is not None: + self._lifecycle_condition.notify_all() + + async def _begin_session_rotation( + self, expected_session_id: str | None, *, drain_timeout: float + ) -> bool: + """Acquire admission exclusion; retry a failed drain without reopening it. + + The caller owns _transition_lock. A committed close already drained; + an incomplete drain must still wait for the outstanding calls. + """ + async with self._lifecycle_condition: + self._ensure_running() + if expected_session_id is not None and self._session.session_id != expected_session_id: + return False + if self._close_committed: + self._session_rotation_in_progress = True + return True + await self._seal_and_drain(drain_timeout) + return True + + async def _seal_and_drain(self, drain_timeout: float) -> None: + """Block admission, then wait out the calls already admitted. + + The caller owns _lifecycle_condition. A failure that needs operator + recovery keeps admission sealed; any other failure reopens it so one + failed transition does not take the gateway down with it. + """ + self._session_rotation_in_progress = True + self._transition_drain_s = drain_timeout + try: + await self._drain_calls(drain_timeout) + self._ensure_terminal_audit_complete() + except BaseException: + if ( + not self._drain_incomplete + and not self._cleanup_incomplete + and self._failed_terminal_call is None + ): + self._session_rotation_in_progress = False + self._lifecycle_condition.notify_all() + raise + + async def _drain_calls(self, drain_timeout: float) -> None: + """Drain while holding the condition; wait releases it for call finalizers. + + Once cancellation is requested, admission stays sealed on any failure. + A later transition can retry draining; only zero active calls clears + the incomplete state. Cancellation cannot forcibly stop a coroutine. + """ + try: + await asyncio.wait_for( + self._lifecycle_condition.wait_for(lambda: self._active_calls == 0), + timeout=drain_timeout, + ) + except TimeoutError: + self._drain_incomplete = True + for task in tuple(self._active_call_tasks): + task.cancel("session lifecycle drain deadline exceeded") + try: + await asyncio.wait_for( + self._lifecycle_condition.wait_for(lambda: self._active_calls == 0), + timeout=SESSION_CANCELLATION_GRACE_SECONDS, + ) + except TimeoutError: + raise SessionDrainIncomplete( + f"{self._active_calls} call(s) did not honor cancellation within " + f"{SESSION_CANCELLATION_GRACE_SECONDS}s of the drain deadline", + detail=str(self._active_calls), + ) from None + self._drain_incomplete = False + + async def shutdown(self, *, drain_timeout: float = SESSION_CLOSE_DRAIN_SECONDS) -> None: + """Permanently reject admission, drain calls, then close owned resources. + + Failed drain or cleanup is reported, not a successful shutdown. State + and resource ownership remain available for a shutdown retry. The spawn + lock also excludes a first-use start from the cleanup snapshot. + """ + async with self._lifecycle_condition: + self._shutting_down = True + self._lifecycle_condition.notify_all() + async with self._transition_lock: + async with self._lifecycle_condition: + self._session_rotation_in_progress = True + await self._drain_calls(drain_timeout) + async with self._stdio_spawn_lock: + await self.aclose() + + @asynccontextmanager + async def session_rotation( + self, + *, + expected_session_id: str | None = None, + drain_timeout: float = SESSION_CLOSE_DRAIN_SECONDS, + ) -> AsyncIterator[bool]: + """Serialize close attempts while blocking admission across retries. + + Active calls drain before yielding. Once the caller marks irreversible + close work, failure keeps admission sealed, but releases transition + ownership so another close request can recover. A stale session ID + yields False. Only successful rebind reopens a closing session. + """ + async with self._transition_lock: + acquired = await self._begin_session_rotation( + expected_session_id, drain_timeout=drain_timeout + ) + if not acquired: + yield False + return + self._session_rebound = False + try: + yield True + finally: + async with self._lifecycle_condition: + if ( + not self._shutting_down + and not self._drain_incomplete + and not self._cleanup_incomplete + and self._failed_terminal_call is None + and (self._session_rebound or not self._close_committed) + ): + self._session_rotation_in_progress = False + self._lifecycle_condition.notify_all() + + def mark_close_committed(self) -> None: + """Seal admission after irreversible close work, including partial failure. + + The caller must hold session_rotation. This does not assert that a + signature exists; it prevents further mutation once close has started. """ - Point the proxy at a fresh session after the previous one was closed. + if not self._session_rotation_in_progress: + raise RuntimeError("mark_close_committed requires an active session_rotation") + self._close_committed = True - Call logs are recreated for the new session id; catalog, policy - evaluator, and gateway are unchanged. + @asynccontextmanager + async def exclude_session_transition( + self, + *, + expected_session_id: str | None = None, + drain_timeout: float = SESSION_CLOSE_DRAIN_SECONDS, + ) -> AsyncIterator[bool]: + """Serialize reset against close, and drain before the caller's boundary. + + Reset ends a session and opens a successor, so it releases the same + session-scoped resources close does (#625). Draining first is what + makes releasing them safe: an admitted call may still hold the child, + and closing it underneath that call would break it. + + Reject a reset of a session awaiting close recovery. Waiting while + holding transition ownership would prevent that recovery from running. """ + async with self._transition_lock: + if ( + expected_session_id is not None + and self._session.session_id != expected_session_id + ): + yield False + return + async with self._lifecycle_condition: + self._ensure_running() + self._ensure_terminal_audit_complete() + if self._drain_incomplete: + if self._active_calls: + raise SessionDrainIncomplete("session drain requires recovery") + # A prior reset timed out, but every cancelled call has now + # reached a terminal state. Retry the same reset safely. + self._drain_incomplete = False + if self._close_committed: + raise SessionCloseIncomplete("session close requires recovery") + await self._seal_and_drain(drain_timeout) + try: + yield True + finally: + async with self._lifecycle_condition: + # Mirrors session_rotation: a failure needing operator + # recovery keeps admission sealed rather than serving the + # successor from a session that did not finish unwinding. + if ( + not self._shutting_down + and not self._drain_incomplete + and not self._cleanup_incomplete + and self._failed_terminal_call is None + ): + self._session_rotation_in_progress = False + self._lifecycle_condition.notify_all() + + async def rebind_session(self, session: SessionState, audit_chain: AuditChain) -> None: + """Close resources before adopting the successor inside session_rotation.""" + if not self._session_rotation_in_progress: + raise RuntimeError("rebind_session requires an active session_rotation") + if self._session_rebound: + raise RuntimeError("session has already been rebound") + self._ensure_terminal_audit_complete() + await self.aclose() self._session = session self._audit = audit_chain self._call_log = CallLog(session_id=session.session_id) self._session_call_log = SessionCallLog(session_id=session.session_id) + self._session_rebound = True + self._close_committed = False def _warn_pin_unenforced(self, server_url: str, reason: str) -> None: """Log TLS_PIN_UNENFORCED once per server URL (#281, dev/demo paths).""" @@ -338,6 +633,7 @@ def _client_for_upstream(self, entry: CatalogEntry) -> httpx.AsyncClient: pinned are not the same peer, and must not share state meant to be scoped to one. """ + self._ensure_running() server_url = entry.server.url fingerprint = entry.server.tls_fingerprint scheme = httpx.URL(server_url).scheme.lower() @@ -379,10 +675,19 @@ def _client_for_upstream(self, entry: CatalogEntry) -> httpx.AsyncClient: return client async def _stdio_for(self, entry: CatalogEntry) -> StdioServer: - """The child for this server, spawned on first use in this session.""" + """Return the session-owned child, serializing first use with shutdown.""" + self._ensure_running() key = _server_execution_key(entry) server = self._stdio_servers.get(key) - if server is None: + if server is not None: + return server + async with self._stdio_spawn_lock: + self._ensure_running() + # Re-check: a coroutine that waited for the lock may find another + # already finished spawning this key while it waited. + server = self._stdio_servers.get(key) + if server is not None: + return server if entry.server.spawn is None: raise UpstreamUnavailable( f"catalog entry {entry.tool_name!r} declares stdio transport with no " @@ -394,13 +699,71 @@ async def _stdio_for(self, entry: CatalogEntry) -> StdioServer: ) await server.start() self._stdio_servers[key] = server - return server + return server async def aclose(self) -> None: - """Terminate spawned children. A session that ends leaves nothing running.""" - for server in self._stdio_servers.values(): - await server.close() - self._stdio_servers.clear() + """Close owned resources: children retryably, HTTP clients best-effort. + + Callers must exclude admission and resource creation before cleanup. + + A child that fails to close is retained and its failure propagates: the + process is still live, still owns the handle, and a retry can still reap + it. An HTTP client cannot offer the same guarantee. `AsyncClient` marks + itself closed and HTTPcore empties its pool before the underlying + streams are released, so a failed close leaves connections that no retry + reaches through any public API. Retaining such a client would advertise a + recovery that does not exist, so it is dropped and the failure logged. + Dropping it is what the successor needs anyway: it is the reuse, not the + socket, that #625 is about. + """ + stdio_items = tuple(self._stdio_servers.items()) + http_items = tuple(self._http_clients.items()) + self._provenance.clear() + self._drift_checked.clear() + + if not stdio_items and not http_items: + self._cleanup_incomplete = False + return + try: + results = await asyncio.gather( + *(server.close() for _, server in stdio_items), + *(client.aclose() for _, client in http_items), + return_exceptions=True, + ) + except BaseException: + self._cleanup_incomplete = True + raise + # Every close is attempted even if an earlier one raises (gather with + # return_exceptions=True), so one stuck child cannot leak the rest. + stdio_results = results[: len(stdio_items)] + http_results = results[len(stdio_items) :] + failures: list[BaseException] = [] + for (stdio_key, _), result in zip(stdio_items, stdio_results, strict=True): + if isinstance(result, BaseException): + failures.append(result) + else: + del self._stdio_servers[stdio_key] + for (http_key, _), result in zip(http_items, http_results, strict=True): + del self._http_clients[http_key] + if isinstance(result, BaseException): + logger.error( + "session-owned HTTP client for %s failed to close and was " + "dropped; its connections may be leaked: %r", + http_key, + result, + ) + + # Only the first failure propagates; it is logged with how many others + # also failed so a partial-cleanup failure is not read as a single one. + if failures: + self._cleanup_incomplete = True + if len(failures) > 1: + logger.error( + "%d session-owned children failed to close; raising the first", + len(failures), + ) + raise failures[0] + self._cleanup_incomplete = False async def _advertised_tools(self, entry: CatalogEntry) -> list[dict[str, Any]] | None: """What the server offers *this gateway*, for the provenance comparison. @@ -911,6 +1274,10 @@ def _finalize_unexpected_call_failure( external_execution_evidence=(finalization.external_execution_evidence), ) except (Exception, asyncio.CancelledError) as persistence_exc: + # A finished task is not proof of a recorded outcome. Retain the + # first failure before admission accounting can let close proceed. + if self._failed_terminal_call is None: + self._failed_terminal_call = call_id exc.add_note( "terminal audit persistence failed with " f"{type(persistence_exc).__name__} during " @@ -929,38 +1296,33 @@ async def call_tool( ) -> CallResult: """Run one call and guarantee one terminal on failure or cancellation.""" finalization = _CallFinalizationState() - # Adopt the session's shared value before anything evaluates this call. - # An instance joining a session another instance opened, or one that has - # restarted, would otherwise evaluate the first call against its own - # empty copy and permit what the session's accumulated value forbids. - # No-op when no shared store is configured. - await self._session.hydrate() - # The session generation this call was issued under, read after hydration - # so a reset performed on another instance is already visible. A reset - # arriving mid-call closes that session, and this response must not raise - # the successor. - finalization.reset_count = self._session.reset_count + await self._enter_call() try: - return await self._call_tool_impl( - call_id, - tool_name, - arguments, - workflow_id, - declared_data_class, - execution_id=execution_id, - _finalization=finalization, - ) - except BaseException as exc: - if not isinstance(exc, (Exception, asyncio.CancelledError)): + try: + await self._session.hydrate() + finalization.reset_count = self._session.reset_count + return await self._call_tool_impl( + call_id, + tool_name, + arguments, + workflow_id, + declared_data_class, + execution_id=execution_id, + _finalization=finalization, + ) + except BaseException as exc: + if not isinstance(exc, (Exception, asyncio.CancelledError)): + raise + self._finalize_unexpected_call_failure( + finalization, + exc, + call_id=call_id, + tool_name=tool_name, + workflow_id=workflow_id, + ) raise - self._finalize_unexpected_call_failure( - finalization, - exc, - call_id=call_id, - tool_name=tool_name, - workflow_id=workflow_id, - ) - raise + finally: + await self._leave_call() async def _call_tool_impl( self, diff --git a/src/cmcp_runtime/mcp/server.py b/src/cmcp_runtime/mcp/server.py index aa574070..fba9a170 100644 --- a/src/cmcp_runtime/mcp/server.py +++ b/src/cmcp_runtime/mcp/server.py @@ -18,6 +18,8 @@ import time import uuid from collections import defaultdict +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from typing import TYPE_CHECKING, Any from starlette.applications import Starlette @@ -28,7 +30,7 @@ from starlette.routing import Route from cmcp_runtime.catalog.loader import ApprovedDefinition, CatalogEntry, ServerIdentity -from cmcp_runtime.mcp.proxy import CMCPProxy +from cmcp_runtime.mcp.proxy import SESSION_CLOSE_DRAIN_SECONDS, CMCPProxy if TYPE_CHECKING: from cmcp_runtime.audit.chain import AuditChain @@ -341,6 +343,14 @@ def __init__( # Chains of closed sessions, kept so /audit/export still serves them # after the live session rotates. self._closed_chains: dict[str, AuditChain] = {} + # Preserve the successor across failed resource cleanup. The manager + # independently retains claim/partial-close state if creation fails. + # At most one close can be pending: admission stays sealed until it + # resolves, and a close naming any other session is rejected before it + # reaches commit. One slot, so a close nobody retries cannot accumulate. + self._pending_close: ( + tuple[str, dict[str, Any], SessionState, AuditChain] | None + ) = None self._kernel = StatelessKernel() # NET-002: rate-limit unauthenticated /health before auth middleware runs. # Starlette applies middleware outermost-first (first in list = first to run). @@ -368,8 +378,14 @@ def __init__( self._cleanup_interval_s: int = int( os.environ.get("CMCP_SESSION_CLEANUP_INTERVAL_SECONDS", "60") ) + self._session_close_drain_s: float = float( + os.environ.get( + "CMCP_SESSION_CLOSE_DRAIN_SECONDS", str(SESSION_CLOSE_DRAIN_SECONDS) + ) + ) self.app = Starlette( + lifespan=self._lifespan, routes=[ Route("/mcp", self._handle_mcp, methods=["POST"]), Route("/health", self._health, methods=["GET"]), @@ -397,6 +413,14 @@ def __init__( exception_handlers={Exception: _unhandled_error_handler}, ) + @asynccontextmanager + async def _lifespan(self, app: Starlette) -> AsyncIterator[None]: + """Drain admitted calls and close session resources on graceful shutdown.""" + try: + yield + finally: + await self._proxy.shutdown(drain_timeout=self._session_close_drain_s) + async def _parse_mcp_envelope(self, request: Request) -> dict[str, Any] | Response: """Read, size-check, and parse the request body. @@ -831,21 +855,53 @@ async def _session_close(self, request: Request) -> Response: status_code=404, ) - claim = self._session_manager.close_session( - session_id, - self._session, - self._audit_chain, - call_log=getattr(self._proxy, "_call_log", None), - session_call_log=getattr(self._proxy, "_session_call_log", None), - ) - self._closed_chains[session_id] = self._audit_chain - - # Rotate onto a fresh session so the gateway keeps serving. - new_session, new_chain = self._session_manager.create_session() - self._session = new_session - self._audit_chain = new_chain - self._audit = new_chain - self._proxy.rebind_session(new_session, new_chain) + async with self._proxy.session_rotation( + expected_session_id=session_id, drain_timeout=self._session_close_drain_s + ) as acquired: + if not acquired: + return JSONResponse( + { + "error": "session_not_found", + "message": ( + f"No open session with id '{session_id}'. It may already be " + "closed, or you passed the _cmcp.session_id label instead of " + "the internal session id. Look up the internal id via " + "GET /audit/export?session_id=