diff --git a/CHANGELOG.md b/CHANGELOG.md index b3fac25b..390da074 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -161,6 +161,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- **Exhaust upstream `tools/list` pagination before drift and provenance + comparisons** (#631). HTTP and stdio share bounded acquisition: later-page + failures, malformed or ambiguous listings, and cursor cycles are unchecked, + never a comparison against a partial catalog. Existing drift policy and + unchecked-call behavior are unchanged. Observed by solloek369-arch on #566 + and confirmed in #631 by Imran Siddique. + Drift and provenance now share one completed discovery acquisition per + server/authority per session, including unchecked outcomes, avoiding a second + full pagination walk on cold calls with provenance configured. Concurrent + readers wait for completion; cancelled reads are not cached. + Session rebinding resets acquisition and comparison caches together. The + duplicate-fetch cost was identified by qubeena07 during review of #633. + - TLS pinning test fixtures set `minimum_version = TLSv1_2`; the server was built with `PROTOCOL_TLS_SERVER` and no floor, leaving TLSv1 and TLSv1.1 reachable in the test that asserts the gateway's transport rules. diff --git a/LIMITATIONS.md b/LIMITATIONS.md index 78b6411e..2d4d2594 100644 --- a/LIMITATIONS.md +++ b/LIMITATIONS.md @@ -25,6 +25,10 @@ cMCP compares what each upstream server advertises against the approved catalog Separately, the approved description rather than the live one is what the gateway serves to the agent on `tools/list`, so a mutated description does not reach the model through cMCP even in the windows above. That is a structural property of proxying an approved catalog, not a detection result, and it does not extend to the tool's behaviour once called. +HTTP and stdio discovery exhaust `tools/list` pagination before comparing either drift or provenance. A later-page failure, malformed discovery shape, duplicate tool name, repeated/cyclic cursor, or continuation beyond 1,000 pages makes the entire acquisition unchecked; no partial list is compared. Cursors are passed back unchanged, including an empty string. This bound limits page count, not total elapsed time or response bytes, and pagination does not establish an atomic snapshot of a changing server. This is acquisition validation, not full MCP schema validation or a new approval/hash policy. The unchecked-call behavior above is unchanged. + +Drift and provenance share the completed first-contact acquisition for the same server and publisher authority within a session, rather than independently walking all pages. An unchecked acquisition is also cached until the next session. Cancellation leaves no cached acquisition; existing stdio child-close behavior is unchanged, not an automatic child restart. Close/reset drains admitted calls before cleanup invalidates both comparisons and the shared acquisition, even if resource cleanup subsequently fails. This avoids duplicate discovery work; it does not add continuous monitoring or make the listing an atomic snapshot. + **Phase 2 completeness: server-side attestation** Phase 1 attests the gateway boundary. It does not attest what happens on the other side of that boundary. The `tool_transcript.hash` field in the TRACE Claim records a hash of the audit chain tip, but the tool transcript binding that ties a specific tool execution to a specific response is Phase 2 work. Phase 1 partially addresses P1.4 (transitive trust into upstream dependencies) and P4.1 (typosquatted packages added to catalog) -- both are fully closed by Phase 2. Any compliance claim that relies on server-side proof must wait for Phase 2. diff --git a/src/cmcp_runtime/mcp/discovery.py b/src/cmcp_runtime/mcp/discovery.py new file mode 100644 index 00000000..c8c2028b --- /dev/null +++ b/src/cmcp_runtime/mcp/discovery.py @@ -0,0 +1,68 @@ +"""Acquire a complete tools/list before drift or provenance comparison. + +This is acquisition validation, not approval or a catalog hash construction. +No partial result escapes: every page must be attributable and well-shaped, +and pagination must terminate within the local page budget. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import Any + +# A page bound also terminates servers that issue endlessly distinct cursors. +# It is not a wall-clock deadline or a guarantee of an atomic server snapshot. +MAX_DISCOVERY_PAGES = 1000 + +PageFetcher = Callable[[str, dict[str, Any]], Awaitable[dict[str, Any]]] + + +class DiscoveryError(ValueError): + """A bounded local reason; never embed upstream payloads or cursor values.""" + + +async def collect_tools(fetch_page: PageFetcher) -> list[dict[str, Any]]: + """Return only an exhausted, unambiguous listing; otherwise raise. + + Transport failures and cancellation propagate to the owning transport. + Cursors are opaque strings, including the empty string. Only absence of + nextCursor terminates a listing; an empty page alone does not. + """ + tools: list[dict[str, Any]] = [] + names: set[str] = set() + cursors: set[str] = set() + params: dict[str, Any] = {} + for page in range(MAX_DISCOVERY_PAGES): + request_id = f"provenance-tools-list-{page}" + body = await fetch_page(request_id, params) + if ( + not isinstance(body, dict) + or body.get("jsonrpc") != "2.0" + or body.get("id") != request_id + or "error" in body + ): + raise DiscoveryError("invalid response envelope") + result = body.get("result") + if not isinstance(result, dict) or not isinstance(result.get("tools"), list): + raise DiscoveryError("invalid tools page") + for tool in result["tools"]: + if ( + not isinstance(tool, dict) + or not isinstance(tool.get("name"), str) + or not tool["name"] + ): + raise DiscoveryError("invalid tool name") + if tool["name"] in names: + raise DiscoveryError("duplicate tool name") + names.add(tool["name"]) + tools.append(tool) + if "nextCursor" not in result: + return tools + cursor = result["nextCursor"] + if not isinstance(cursor, str): + raise DiscoveryError("invalid continuation cursor") + if cursor in cursors: + raise DiscoveryError("repeated continuation cursor") + cursors.add(cursor) + params = {"cursor": cursor} + raise DiscoveryError("discovery page limit exceeded") diff --git a/src/cmcp_runtime/mcp/proxy.py b/src/cmcp_runtime/mcp/proxy.py index 1b2c251f..11a4b324 100644 --- a/src/cmcp_runtime/mcp/proxy.py +++ b/src/cmcp_runtime/mcp/proxy.py @@ -42,6 +42,7 @@ ) from cmcp_runtime.execution import valid_execution_id from cmcp_runtime.mcp import tls_pinning +from cmcp_runtime.mcp.discovery import DiscoveryError, collect_tools from cmcp_runtime.mcp.stdio import StdioServer from cmcp_runtime.mcp.streamable_http import ( build_request, @@ -291,18 +292,9 @@ def __init__( # in memory would carry it from one agent's session into the next, and # the audit chain cannot see that happen (docs/spec/stdio-transport.md). self._stdio_servers: dict[tuple[str, ...], StdioServer] = {} - # Provenance outcome per server, decided once per session on first use. - # Cached because the answer cannot change within a session without the - # server being replaced underneath us, and re-listing tools on every call - # would make the check expensive enough to be turned off. - self._provenance: dict[tuple[str, ...], ProvenanceResult] = {} + self._reset_upstream_checks() # Servers already warned about unenforceable pinning (warn once each). self._tls_pin_warned: set[str] = set() - # #521: servers whose advertised tool definitions have been compared against - # the catalog. Cached per server for the same reason provenance is: one - # tools/list round trip per server per session is affordable, one per call - # 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`. @@ -328,6 +320,17 @@ def __init__( self._failed_terminal_call: str | None = None self._shutting_down = False + def _reset_upstream_checks(self) -> None: + # Drift and provenance share one completed paginated acquisition per + # server/authority per session, including an unchecked (None) outcome. + # These are first-contact observations, not continuous monitoring. + # Replace, rather than clear: an in-flight acquisition retains its old + # cache identity and must retry before returning into a new session. + self._advertised: dict[tuple[str, ...], list[dict[str, Any]] | None] = {} + self._discovery_locks: dict[tuple[str, ...], asyncio.Lock] = {} + self._provenance: dict[tuple[str, ...], ProvenanceResult] = {} + self._drift_checked: set[tuple[str, ...]] = set() + def _ensure_running(self) -> None: if self._shutting_down: raise UpstreamUnavailable("gateway is shutting down") @@ -718,8 +721,7 @@ async def aclose(self) -> None: """ stdio_items = tuple(self._stdio_servers.items()) http_items = tuple(self._http_clients.items()) - self._provenance.clear() - self._drift_checked.clear() + self._reset_upstream_checks() if not stdio_items and not http_items: self._cleanup_incomplete = False @@ -766,30 +768,54 @@ async def aclose(self) -> None: 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. + """One first-contact acquisition shared by drift and provenance checks. Returns ``None`` when the server will not say, which the caller records as ``unchecked`` rather than as a pass. Never falls back to the catalog's own approved definitions: comparing a record against our approval instead of against the server is the substitution that turns the check into theatre. """ + key = _server_provenance_key(entry) + while True: + cache = self._advertised + lock = self._discovery_locks.setdefault(key, asyncio.Lock()) + async with lock: + if cache is not self._advertised: + continue + if key in cache: + return cache[key] + advertised = await self._discover_tools(entry) + if cache is not self._advertised: + continue + # Only completed acquisition outcomes are cached. In particular, + # cancellation propagates without storing partial data or None. + cache[key] = advertised + return advertised + + async def _discover_tools(self, entry: CatalogEntry) -> list[dict[str, Any]] | None: + """Acquire the entire listing, or None on an ordinary acquisition failure.""" if entry.server.is_stdio: return await (await self._stdio_for(entry)).list_tools() - try: + + async def fetch_page(request_id: str, params: dict[str, Any]) -> dict[str, Any]: + payload, headers = build_request(request_id, "tools/list", params) client = self._client_for_upstream(entry) - payload, headers = build_request("provenance-tools-list", "tools/list", {}) resp = await client.post( entry.server.url, json=payload, headers=headers, ) resp.raise_for_status() - result = parse_response(resp, "provenance-tools-list").get("result") - except Exception as exc: # noqa: BLE001 - any failure means "could not check" - logger.warning("could not list tools for provenance check: %s", exc) - return None - tools = result.get("tools") if isinstance(result, dict) else None - return tools if isinstance(tools, list) else None + return parse_response(resp, request_id) + + try: + return await collect_tools(fetch_page) + except DiscoveryError as exc: + logger.warning("tools discovery incomplete: %s", exc) + except Exception: # noqa: BLE001 - any acquisition failure means "could not check" + # Upstream exceptions may contain response bodies or opaque cursors. + logger.warning("tools discovery incomplete: upstream request failed") + return None async def _check_upstream_drift(self, entry: CatalogEntry) -> bool: """Compare what a server advertises against what we approved (P4.2). @@ -810,14 +836,17 @@ async def _check_upstream_drift(self, entry: CatalogEntry) -> bool: key = _server_provenance_key(entry) if key in self._drift_checked: return self._session.catalog_drift - self._drift_checked.add(key) - advertised = await self._advertised_tools(entry) + # Another caller may have completed the comparison while this one was + # waiting for discovery. An in-flight check is never marked completed. + if key in self._drift_checked: + return self._session.catalog_drift if advertised is None: logger.info( "upstream drift: server=%s outcome=unchecked (server would not list tools)", key, ) + self._drift_checked.add(key) return self._session.catalog_drift by_name = { @@ -840,6 +869,7 @@ async def _check_upstream_drift(self, entry: CatalogEntry) -> bool: if not drifted: logger.info("upstream drift: server=%s outcome=match", key) + self._drift_checked.add(key) return self._session.catalog_drift fail_closed = self._config.catalog.drift_policy is DriftPolicy.FAIL_CLOSED @@ -872,6 +902,7 @@ async def _check_upstream_drift(self, entry: CatalogEntry) -> bool: if fail_closed: self._session.catalog_drift = True + self._drift_checked.add(key) return self._session.catalog_drift async def _check_provenance(self, entry: CatalogEntry) -> ProvenanceResult: diff --git a/src/cmcp_runtime/mcp/stdio.py b/src/cmcp_runtime/mcp/stdio.py index a7b66655..379d5822 100644 --- a/src/cmcp_runtime/mcp/stdio.py +++ b/src/cmcp_runtime/mcp/stdio.py @@ -47,6 +47,7 @@ class rather than folded into hardware attestation. from typing import Any from cmcp_runtime.errors import ConfigError, UpstreamToolError, UpstreamUnavailable +from cmcp_runtime.mcp.discovery import DiscoveryError, collect_tools logger = logging.getLogger(__name__) @@ -244,14 +245,21 @@ async def list_tools(self) -> list[dict[str, Any]] | None: is one whose provenance could not be checked, and the two are recorded differently. """ + async def fetch_page(request_id: str, params: dict[str, Any]) -> dict[str, Any]: + return await self._request(request_id, "tools/list", params) + try: - body = await self._request("provenance-tools-list", "tools/list", {}) - except (UpstreamUnavailable, UpstreamToolError) as exc: - logger.warning("could not list tools for provenance check: %s", exc) - return None - result = body.get("result") - tools = result.get("tools") if isinstance(result, dict) else None - return tools if isinstance(tools, list) else None + return await collect_tools(fetch_page) + except asyncio.CancelledError: + # An interrupted read can leave a late reply in the pipe. Do not + # let a later tool call consume it as that call's response. + await self.close() + raise + except DiscoveryError as exc: + logger.warning("tools discovery incomplete: %s", exc) + except (UpstreamUnavailable, UpstreamToolError): + logger.warning("tools discovery incomplete: upstream request failed") + return None async def call(self, call_id: str, tool_name: str, arguments: dict[str, Any]) -> str: """One JSON-RPC ``tools/call`` over the child's stdin/stdout.""" diff --git a/tests/unit/test_discovery_session_lifecycle.py b/tests/unit/test_discovery_session_lifecycle.py new file mode 100644 index 00000000..43b0ca64 --- /dev/null +++ b/tests/unit/test_discovery_session_lifecycle.py @@ -0,0 +1,296 @@ +"""Public calls compose paginated discovery with session draining and cleanup. + +HTTP uses real owned httpx clients with MockTransport; stdio uses a measured +Python child and a loopback gate. Acquisition, lifecycle, audit finalization, +and signed TRACE provenance verification are not replaced. The existing proxy +helper supplies only unrelated policy/scanner seams. +""" + +from __future__ import annotations + +import asyncio +import json +import sys +import textwrap +from contextlib import asynccontextmanager, suppress +from types import SimpleNamespace + +import httpx +import pytest + +from cmcp_runtime.audit.chain import AuditChain +from cmcp_runtime.mcp.stdio import StdioSpawn, measure_executable, resolve_executable +from cmcp_runtime.provenance import ProvenanceOutcome +from cmcp_runtime.session.state import SessionState +from tests.unit.test_mcp_proxy import _make_proxy +from tests.unit.test_shared_discovery_transports import ( + _CURSOR, + _RESPONSE, + _SERVER_INFO, + _signed_catalog, +) +from tests.unit.test_upstream_catalog_drift import _advertise + +_CHILD = """ + import json + import socket + import sys + + with open(sys.argv[1], encoding="utf-8") as handle: + pages = json.load(handle) + for line in sys.stdin: + request = json.loads(line) + with open(sys.argv[2], "a", encoding="utf-8") as handle: + handle.write(json.dumps(request) + "\\n") + if request["method"] == "tools/list": + second = "cursor" in request["params"] + if second: + with socket.create_connection(("127.0.0.1", int(sys.argv[3]))) as gate: + if gate.recv(1) != b"x": + raise RuntimeError("test gate closed") + reply = pages["second" if second else "first"] + else: + assert request["method"] == "tools/call" + reply = {"result": {"content": [{"type": "text", "text": "customer found"}]}} + sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": request["id"], **reply}) + "\\n") + sys.stdout.flush() +""" + + +@asynccontextmanager +async def _gateway(tmp_path, monkeypatch, transport): + catalog, entry = _signed_catalog(tmp_path) + pages = { + "first": {"result": {"tools": [_SERVER_INFO], "nextCursor": _CURSOR}}, + "second": {"result": {"tools": _advertise()}}, + } + entered, release = asyncio.Event(), asyncio.Event() + gate_tasks = set() + + async def gate_connection(reader, writer): + task = asyncio.current_task() + gate_tasks.add(task) + try: + entered.set() + await release.wait() + writer.write(b"x") + await writer.drain() + except (ConnectionResetError, BrokenPipeError): + pass # A cancelled discovery deliberately terminates the child. + finally: + writer.close() + # Windows reports the deliberately terminated child's reset here too. + with suppress(ConnectionResetError, BrokenPipeError): + await writer.wait_closed() + gate_tasks.discard(task) + + gate = None + request_log = tmp_path / "requests.jsonl" + if transport == "stdio": + gate = await asyncio.start_server(gate_connection, "127.0.0.1", 0) + port = gate.sockets[0].getsockname()[1] + script = tmp_path / "upstream.py" + script.write_text(textwrap.dedent(_CHILD), encoding="utf-8") + page_path = tmp_path / "pages.json" + page_path.write_text(json.dumps(pages), encoding="utf-8") + entry.server.transport = "stdio" + entry.server.spawn = StdioSpawn( + command=sys.executable, + args=(str(script), str(page_path), str(request_log), str(port)), + measure_target=str(script), + binary_digest=measure_executable(resolve_executable(str(script))), + ) + + proxy, session, chain = _make_proxy(catalog=catalog) + del proxy._forward_to_upstream + proxy._mcp_gateway.intercept_tool_response.side_effect = lambda **kwargs: SimpleNamespace( + allowed=True, content=kwargs["response_content"], threats=[], action="allowed" + ) + proxy._config.attestation.required_provenance_kind = "publisher-asserted" + requests, clients = [], [] + + async def respond(request): + payload = json.loads(request.content) + requests.append(payload) + assert request.headers["Mcp-Method"] == payload["method"] + if payload["method"] == "tools/list": + second = "cursor" in payload["params"] + if second: + entered.set() + await release.wait() + reply = pages["second" if second else "first"] + else: + assert payload["method"] == "tools/call" + reply = {"result": {"content": [{"type": "text", "text": _RESPONSE}]}} + body = {"jsonrpc": "2.0", "id": payload["id"], **reply} + if transport == "sse": + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=f"data: {json.dumps(body)}\n\n", + ) + return httpx.Response(200, json=body) + + if transport != "stdio": + entry.server.url = "http://upstream.example/mcp" + real_client = httpx.AsyncClient + + def local_client(*args, **kwargs): + kwargs["transport"] = httpx.MockTransport(respond) + client = real_client(*args, **kwargs) + clients.append(client) + return client + + # Keep production resource ownership/keying/cleanup; replace only I/O. + monkeypatch.setattr("cmcp_runtime.mcp.proxy.httpx.AsyncClient", local_client) + + def observed_requests(): + if transport == "stdio": + return [json.loads(line) for line in request_log.read_text().splitlines()] + return list(requests) + + def owned_resource(): + if transport == "stdio": + return next(iter(proxy._stdio_servers.values()))._proc + return clients[-1] + + try: + yield SimpleNamespace( + proxy=proxy, + session=session, + chain=chain, + entered=entered, + release=release, + requests=observed_requests, + resource=owned_resource, + ) + finally: + release.set() + await proxy.shutdown(drain_timeout=0) + if gate is not None: + gate.close() + await gate.wait_closed() + if gate_tasks: + await asyncio.gather(*tuple(gate_tasks)) + + +@pytest.mark.parametrize("transport", ["json", "sse", "stdio"]) +@pytest.mark.parametrize("cancel", [False, True], ids=["normal-drain", "deadline-cancellation"]) +async def test_paginated_calls_drain_before_rebind_and_successor_discovers_fresh( + tmp_path, monkeypatch, transport, cancel +): + async with _gateway(tmp_path, monkeypatch, transport) as gateway: + proxy = gateway.proxy + calls = [asyncio.create_task(proxy.call_tool("old-owner", "lookup_customer", {}))] + await asyncio.wait_for(gateway.entered.wait(), timeout=5) + old_resource = gateway.resource() + old_cache = proxy._advertised + old_locks = proxy._discovery_locks + assert not old_cache + assert not proxy._drift_checked + assert not proxy._provenance + + if cancel: + admitted = asyncio.Event() + enter_call = proxy._enter_call + + async def observed_enter(): + await enter_call() + admitted.set() + + monkeypatch.setattr(proxy, "_enter_call", observed_enter) + calls.append(asyncio.create_task(proxy.call_tool("old-waiter", "lookup_customer", {}))) + await asyncio.wait_for(admitted.wait(), timeout=5) + assert proxy._active_calls == 2 + assert not calls[-1].done() + + draining = asyncio.Event() + drain_calls = proxy._drain_calls + + async def observed_drain(timeout): + draining.set() + await drain_calls(timeout) + + monkeypatch.setattr(proxy, "_drain_calls", observed_drain) + next_session = SessionState(session_id="successor") + next_chain = AuditChain(next_session.session_id) + next_genesis = list(next_chain.entries) + + async def rotate(): + async with proxy.session_rotation(drain_timeout=0 if cancel else 5) as acquired: + assert acquired + assert proxy._active_calls == 0 + if cancel: + assert not old_cache + assert not proxy._drift_checked + assert not proxy._provenance + faults = [e for e in gateway.chain.entries if e.entry_type == "fault"] + assert {e.call_id for e in faults} == {"old-owner", "old-waiter"} + assert all(e.detail["exception_type"] == "CancelledError" for e in faults) + else: + assert ( + next(iter(proxy._provenance.values())).outcome is ProvenanceOutcome.VERIFIED + ) + await proxy.rebind_session(next_session, next_chain) + + rotation = asyncio.create_task(rotate()) + try: + await asyncio.wait_for(draining.wait(), timeout=5) + assert proxy._session is gateway.session + assert not rotation.done() + if not cancel: + assert not calls[0].done() + if transport == "stdio": + assert old_resource.returncode is None + else: + assert not old_resource.is_closed + gateway.release.set() + results = await asyncio.wait_for( + asyncio.gather(*calls, return_exceptions=True), timeout=10 + ) + await asyncio.wait_for(rotation, timeout=10) + finally: + gateway.release.set() + for task in [*calls, rotation]: + if not task.done(): + task.cancel() + await asyncio.gather(*calls, rotation, return_exceptions=True) + + if cancel: + assert all(isinstance(result, asyncio.CancelledError) for result in results) + else: + assert results[0].allowed + assert results[0].response == _RESPONSE + assert proxy._active_calls == 0 + assert proxy._session is next_session + assert proxy._advertised is not old_cache + assert proxy._discovery_locks is not old_locks + assert not proxy._advertised and not proxy._discovery_locks + assert not proxy._drift_checked and not proxy._provenance + assert next_chain.entries == next_genesis + if transport == "stdio": + assert old_resource.returncode is not None + else: + assert old_resource.is_closed + + old_entries = list(gateway.chain.entries) + for call_id in ("next-cold", "next-warm"): + result = await proxy.call_tool(call_id, "lookup_customer", {}) + assert result.allowed + assert result.response == _RESPONSE + assert gateway.resource() is not old_resource + assert gateway.chain.entries == old_entries + assert {e.call_id for e in next_chain.entries if e.call_id} == {"next-cold", "next-warm"} + assert next(iter(proxy._provenance.values())).outcome is ProvenanceOutcome.VERIFIED + requests = gateway.requests() + listings = [request for request in requests if request["method"] == "tools/list"] + assert ["cursor" in request["params"] for request in listings] == [False, True, False, True] + assert [request["params"].get("cursor") for request in listings] == [ + None, + _CURSOR, + None, + _CURSOR, + ] + assert sum(request["method"] == "tools/call" for request in requests) == ( + 2 if cancel else 3 + ) diff --git a/tests/unit/test_shared_discovery_cache.py b/tests/unit/test_shared_discovery_cache.py new file mode 100644 index 00000000..01ec8b64 --- /dev/null +++ b/tests/unit/test_shared_discovery_cache.py @@ -0,0 +1,338 @@ +"""Session-scoped discovery sharing; transport/signature coverage lives separately. + +These unit tests replace only acquisition, not the cache under test. A completed +empty catalog, an unchecked result, and a cancelled acquisition are distinct. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import replace +from unittest.mock import AsyncMock, patch + +import pytest + +from cmcp_runtime.audit.chain import AuditChain +from cmcp_runtime.config import DriftPolicy +from cmcp_runtime.provenance import ProvenanceOutcome, ProvenanceResult +from cmcp_runtime.session.state import SessionState +from tests.unit.test_upstream_catalog_drift import _advertise, _catalog, _proxy + + +@pytest.mark.parametrize("outcome", [[], None], ids=["empty-catalog", "unchecked"]) +async def test_completed_empty_and_unchecked_results_are_cached(outcome): + catalog = _catalog() + proxy, _, _ = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + proxy._discover_tools = AsyncMock(return_value=outcome) + entry = catalog.entries["lookup_customer"] + + assert await proxy._advertised_tools(entry) == outcome + assert await proxy._advertised_tools(entry) == outcome + proxy._discover_tools.assert_awaited_once_with(entry) + + +async def test_concurrent_drift_checks_wait_for_one_completed_comparison(): + """A first-contact marker is not permission to skip an in-flight check.""" + catalog = _catalog() + proxy, session, chain = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + entry = catalog.entries["lookup_customer"] + started, release, second_started = asyncio.Event(), asyncio.Event(), asyncio.Event() + + async def discover(_entry): + started.set() + await release.wait() + return _advertise("Changed upstream definition") + + async def second_check(): + second_started.set() + return await proxy._check_upstream_drift(entry) + + proxy._discover_tools = AsyncMock(side_effect=discover) + first = asyncio.create_task(proxy._check_upstream_drift(entry)) + await asyncio.wait_for(started.wait(), timeout=1) + second = asyncio.create_task(second_check()) + try: + await asyncio.wait_for(second_started.wait(), timeout=1) + assert not first.done() + assert not second.done() + assert not proxy._drift_checked + finally: + release.set() + results = await asyncio.wait_for(asyncio.gather(first, second), timeout=1) + + assert results == [True, True] + assert session.catalog_drift is True + assert len([item for item in chain.entries if item.entry_type == "catalog_drift"]) == 1 + proxy._discover_tools.assert_awaited_once_with(entry) + + +async def test_cancelled_acquisition_is_not_cached_and_waiter_retries(): + catalog = _catalog() + proxy, session, _ = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + entry = catalog.entries["lookup_customer"] + started, second_started = asyncio.Event(), asyncio.Event() + attempts = 0 + + async def discover(_entry): + nonlocal attempts + attempts += 1 + if attempts == 1: + started.set() + await asyncio.Event().wait() + return _advertise() + + async def second_check(): + second_started.set() + return await proxy._check_upstream_drift(entry) + + proxy._discover_tools = AsyncMock(side_effect=discover) + first = asyncio.create_task(proxy._check_upstream_drift(entry)) + await asyncio.wait_for(started.wait(), timeout=1) + second = asyncio.create_task(second_check()) + await asyncio.wait_for(second_started.wait(), timeout=1) + assert not second.done() + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + assert await asyncio.wait_for(second, timeout=1) is False + assert await proxy._advertised_tools(entry) == _advertise() + assert session.catalog_drift is False + assert proxy._discover_tools.await_count == 2 + + +async def test_cancelled_waiter_does_not_cancel_the_shared_acquisition(): + catalog = _catalog() + proxy, _, _ = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + entry = catalog.entries["lookup_customer"] + started, release, waiter_started = asyncio.Event(), asyncio.Event(), asyncio.Event() + + async def discover(_entry): + started.set() + await release.wait() + return _advertise() + + async def wait_for_catalog(): + waiter_started.set() + return await proxy._advertised_tools(entry) + + proxy._discover_tools = AsyncMock(side_effect=discover) + owner = asyncio.create_task(proxy._advertised_tools(entry)) + await asyncio.wait_for(started.wait(), timeout=1) + waiter = asyncio.create_task(wait_for_catalog()) + try: + await asyncio.wait_for(waiter_started.wait(), timeout=1) + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + assert not owner.done() + finally: + release.set() + assert await asyncio.wait_for(owner, timeout=1) == _advertise() + + assert await proxy._advertised_tools(entry) == _advertise() + proxy._discover_tools.assert_awaited_once_with(entry) + + +@pytest.mark.parametrize( + "server_change", + [ + {"url": "https://other.example/mcp"}, + {"tls_fingerprint": "sha256:" + "d" * 64}, + {"provenance_record_path": "another-record.json"}, + {"publisher_jwk": {"kty": "EC", "kid": "another-authority"}}, + ], + ids=["endpoint", "tls-pin", "record", "publisher-authority"], +) +async def test_discovery_cache_does_not_cross_server_or_provenance_identity(server_change): + catalog = _catalog() + proxy, _, _ = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + entry = catalog.entries["lookup_customer"] + other = replace(entry, server=replace(entry.server, **server_change)) + proxy._discover_tools = AsyncMock(side_effect=[_advertise(), []]) + + assert await proxy._advertised_tools(entry) == _advertise() + assert await proxy._advertised_tools(other) == [] + assert await proxy._advertised_tools(entry) == _advertise() + assert await proxy._advertised_tools(other) == [] + assert proxy._discover_tools.await_count == 2 + + +async def test_tools_and_display_labels_sharing_server_identity_share_discovery(): + catalog = _catalog() + proxy, _, _ = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + entry = catalog.entries["lookup_customer"] + other = replace( + entry, tool_name="another_tool", server=replace(entry.server, display_name="alias") + ) + proxy._discover_tools = AsyncMock(return_value=_advertise()) + + assert await proxy._advertised_tools(entry) == _advertise() + assert await proxy._advertised_tools(other) == _advertise() + proxy._discover_tools.assert_awaited_once_with(entry) + + +async def test_slow_discovery_does_not_block_a_different_server(): + catalog = _catalog() + proxy, _, _ = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + entry = catalog.entries["lookup_customer"] + other = replace(entry, server=replace(entry.server, url="https://other.example/mcp")) + started, release = asyncio.Event(), asyncio.Event() + + async def discover(candidate): + if candidate.server.url == entry.server.url: + started.set() + await release.wait() + return _advertise() + return [] + + proxy._discover_tools = AsyncMock(side_effect=discover) + slow = asyncio.create_task(proxy._advertised_tools(entry)) + await asyncio.wait_for(started.wait(), timeout=1) + try: + assert await asyncio.wait_for(proxy._advertised_tools(other), timeout=1) == [] + assert not slow.done() + finally: + release.set() + assert await asyncio.wait_for(slow, timeout=1) == _advertise() + assert proxy._discover_tools.await_count == 2 + + +async def test_rebind_rechecks_discovery_drift_and_provenance(): + """Acquisition and both downstream verdicts expire at the session boundary.""" + catalog = _catalog() + proxy, old_session, old_chain = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + entry = catalog.entries["lookup_customer"] + entry.server.provenance_record_path = "unit-test-record.json" + proxy._discover_tools = AsyncMock(side_effect=[_advertise(), _advertise("Changed")]) + initial = ProvenanceResult(ProvenanceOutcome.VERIFIED, kind="publisher-asserted") + changed = ProvenanceResult(ProvenanceOutcome.CATALOG_MISMATCH, kind="publisher-asserted") + with patch( + "cmcp_runtime.mcp.proxy.check_server_provenance", side_effect=[initial, changed] + ) as check: + assert await proxy._check_upstream_drift(entry) is False + assert await proxy._check_provenance(entry) == initial + assert proxy._discover_tools.await_count == 1 + + session = SessionState(session_id="next-session") + chain = AuditChain(session_id=session.session_id) + async with proxy.session_rotation(): + await proxy.rebind_session(session, chain) + assert await proxy._check_upstream_drift(entry) is True + assert await proxy._check_provenance(entry) == changed + assert check.call_count == 2 + + assert proxy._discover_tools.await_count == 2 + assert old_session.catalog_drift is False + assert not any(item.entry_type == "catalog_drift" for item in old_chain.entries) + assert session.catalog_drift is True + assert len([item for item in chain.entries if item.entry_type == "catalog_drift"]) == 1 + + +async def test_cache_generation_discards_inflight_result_and_retries_waiters(): + """Private cache-unit coverage, not permission to rebind during public calls. + + The production lifecycle drains admitted calls before resetting these caches. + This directly exercises the acquisition's defensive cache-identity check. + """ + catalog = _catalog() + proxy, session, chain = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + entry = catalog.entries["lookup_customer"] + started, release, second_started = asyncio.Event(), asyncio.Event(), asyncio.Event() + attempts = 0 + + async def discover(_entry): + nonlocal attempts + attempts += 1 + if attempts == 1: + started.set() + await release.wait() + return _advertise("Old-session definition must not enter the new session") + return _advertise() + + async def old_waiter(): + second_started.set() + return await proxy._check_upstream_drift(entry) + + proxy._discover_tools = AsyncMock(side_effect=discover) + old_fetch = asyncio.create_task(proxy._check_upstream_drift(entry)) + await asyncio.wait_for(started.wait(), timeout=1) + old_waiting = asyncio.create_task(old_waiter()) + await asyncio.wait_for(second_started.wait(), timeout=1) + proxy._reset_upstream_checks() + try: + # The new cache generation must not reuse the old acquisition lock. + assert await asyncio.wait_for(proxy._check_upstream_drift(entry), timeout=1) is False + finally: + release.set() + results = await asyncio.wait_for(asyncio.gather(old_fetch, old_waiting), timeout=1) + + assert results == [False, False] + assert await proxy._advertised_tools(entry) == _advertise() + assert proxy._discover_tools.await_count == 2 + assert session.catalog_drift is False + assert not any(item.entry_type == "catalog_drift" for item in chain.entries) + + +@pytest.mark.parametrize("child_fails", [False, True], ids=["no-resources", "failed-child"]) +async def test_rebind_invalidates_all_discovery_caches_before_cleanup(child_fails): + """Empty cleanup and retryable failure both discard first-contact observations.""" + catalog = _catalog() + proxy, old_session, old_chain = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + entry = catalog.entries["lookup_customer"] + proxy._discover_tools = AsyncMock(return_value=_advertise()) + assert await proxy._check_upstream_drift(entry) is False + await proxy._check_provenance(entry) + old_caches = ( + proxy._advertised, + proxy._discovery_locks, + proxy._provenance, + proxy._drift_checked, + ) + assert all(old_caches) + + def assert_invalidated(): + for old, current in zip( + old_caches, + ( + proxy._advertised, + proxy._discovery_locks, + proxy._provenance, + proxy._drift_checked, + ), + strict=True, + ): + assert current is not old + assert not current + + session = SessionState(session_id="next-session") + chain = AuditChain(session_id=session.session_id) + if child_fails: + + async def fail_close(): + # Invalidate before the first resource-cleanup await, not just on success. + assert_invalidated() + raise OSError("injected child close failure") + + child = AsyncMock() + child.close.side_effect = fail_close + proxy._stdio_servers[("test-child",)] = child + with pytest.raises(OSError, match="injected child close failure"): + async with proxy.session_rotation(): + await proxy.rebind_session(session, chain) + assert proxy._session is old_session + assert proxy._audit is old_chain + assert proxy._stdio_servers[("test-child",)] is child + assert proxy._cleanup_incomplete + assert proxy._session_rotation_in_progress + assert not proxy._session_rebound + child.close.side_effect = None + + async with proxy.session_rotation(): + await proxy.rebind_session(session, chain) + assert_invalidated() + assert proxy._session is session + assert proxy._audit is chain + assert not proxy._stdio_servers + assert not proxy._cleanup_incomplete + assert not proxy._session_rotation_in_progress + assert proxy._session_rebound diff --git a/tests/unit/test_shared_discovery_transports.py b/tests/unit/test_shared_discovery_transports.py new file mode 100644 index 00000000..d55d335a --- /dev/null +++ b/tests/unit/test_shared_discovery_transports.py @@ -0,0 +1,179 @@ +"""Drift and provenance share one complete first-contact tools/list acquisition. + +The public call pipeline and TRACE provenance verification are real here. HTTP +uses a local MockTransport and stdio uses a measured Python subprocess; only the +unrelated policy/scanner seams are mocked by the existing proxy test helper. +""" + +from __future__ import annotations + +import json +import sys +import textwrap +from types import SimpleNamespace + +import httpx +import pytest +from agentrust_trace.provenance import build_record, sign_record +from agentrust_trace.sign import generate_key, key_to_jwk + +from cmcp_runtime.config import DriftPolicy +from cmcp_runtime.mcp.stdio import StdioSpawn, measure_executable, resolve_executable +from cmcp_runtime.provenance import ProvenanceOutcome +from tests.unit.test_mcp_proxy import _make_proxy +from tests.unit.test_upstream_catalog_drift import _advertise, _catalog + +_SERVER_INFO = {"name": "server_info", "description": "server information", "inputSchema": {}} +_CURSOR = " next/%2F+雪== " +_RESPONSE = "customer found" + +_STDIO_SERVER = """ + import json + import sys + + with open(sys.argv[1], encoding="utf-8") as handle: + pages = json.load(handle) + for line in sys.stdin: + request = json.loads(line) + with open(sys.argv[2], "a", encoding="utf-8") as handle: + handle.write(json.dumps(request) + "\\n") + if request["method"] == "tools/list": + reply = pages["second" if "cursor" in request["params"] else "first"] + elif request["method"] == "tools/call": + reply = {"result": {"content": [{"type": "text", "text": "customer found"}]}} + else: + raise AssertionError("unexpected method") + body = {"jsonrpc": "2.0", "id": request["id"], **reply} + sys.stdout.write(json.dumps(body) + "\\n") + sys.stdout.flush() +""" + + +def _signed_catalog(tmp_path): + key = generate_key() + record = build_record( + kind="publisher-asserted", + publisher="did:web:crm.example", + tools=[_SERVER_INFO, *_advertise()], + artifact={"package": "pkg:npm/crm@1.0.0", "digest": "sha256:" + "a" * 64}, + ) + record_path = tmp_path / "provenance.json" + record_path.write_text(json.dumps(sign_record(record, key)), encoding="utf-8") + catalog = _catalog() + entry = catalog.entries["lookup_customer"] + entry.server.provenance_record_path = str(record_path) + entry.server.publisher_jwk = key_to_jwk(key) + return catalog, entry + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", ["json", "sse", "stdio"]) +@pytest.mark.parametrize( + "case,outcome,required,allowed", + [ + ("matching", ProvenanceOutcome.VERIFIED, "publisher-asserted", True), + ("changed", ProvenanceOutcome.CATALOG_MISMATCH, None, True), + ("incomplete", ProvenanceOutcome.UNCHECKED, None, True), + ("incomplete", ProvenanceOutcome.UNCHECKED, "publisher-asserted", False), + ], +) +async def test_cold_tool_call_shares_paginated_discovery_between_drift_and_provenance( + tmp_path, monkeypatch, transport, case, outcome, required, allowed +): + catalog, entry = _signed_catalog(tmp_path) + pages = { + "first": {"result": {"tools": [_SERVER_INFO], "nextCursor": _CURSOR}}, + "second": {"result": {"tools": _advertise()}}, + } + if case == "changed": + pages["second"] = {"result": {"tools": _advertise("changed description")}} + elif case == "incomplete": + pages["second"] = {"error": {"code": -32603, "message": "listing failed"}} + + request_log = tmp_path / "requests.jsonl" + if transport == "stdio": + script = tmp_path / "upstream.py" + script.write_text(textwrap.dedent(_STDIO_SERVER), encoding="utf-8") + page_path = tmp_path / "pages.json" + page_path.write_text(json.dumps(pages), encoding="utf-8") + entry.server.transport = "stdio" + entry.server.spawn = StdioSpawn( + command=sys.executable, + args=(str(script), str(page_path), str(request_log)), + measure_target=str(script), + binary_digest=measure_executable(resolve_executable(str(script))), + ) + + proxy, session, chain = _make_proxy(catalog=catalog) + # The helper normally stubs forwarding. Restore the class method so this + # test exercises the exact drift -> forwarding -> provenance cold path. + del proxy._forward_to_upstream + proxy._mcp_gateway.intercept_tool_response.side_effect = lambda **kwargs: SimpleNamespace( + allowed=True, content=kwargs["response_content"], threats=[], action="allowed" + ) + proxy._config.attestation.required_provenance_kind = required + proxy._config.catalog.drift_policy = ( + DriftPolicy.WARN_ONLY if case == "changed" else DriftPolicy.FAIL_CLOSED + ) + requests = [] + + def respond(request): + payload = json.loads(request.content) + requests.append(payload) + assert request.headers["Mcp-Method"] == payload["method"] + if payload["method"] == "tools/list": + reply = pages["second" if "cursor" in payload["params"] else "first"] + else: + assert payload["method"] == "tools/call" + reply = {"result": {"content": [{"type": "text", "text": _RESPONSE}]}} + body = {"jsonrpc": "2.0", "id": payload["id"], **reply} + if transport == "sse": + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=f"data: {json.dumps(body)}\n\n", + ) + return httpx.Response(200, json=body) + + def observed_requests(): + if transport == "stdio": + return [json.loads(line) for line in request_log.read_text().splitlines()] + return requests + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + if transport != "stdio": + monkeypatch.setattr(proxy, "_client_for_upstream", lambda entry: client) + try: + result = await proxy.call_tool("cold", "lookup_customer", {"id": "customer-1"}) + + assert result.allowed is allowed + assert result.response == (_RESPONSE if allowed else None) + assert (await proxy._check_provenance(entry)).outcome is outcome + assert session.catalog_drift is False + assert session.upstream_drift_tools == ( + ["lookup_customer"] if case == "changed" else [] + ) + assert any(item.entry_type == "catalog_drift" for item in chain.entries) is ( + case == "changed" + ) + cold_requests = observed_requests() + discovery_requests = [item for item in cold_requests if item["method"] == "tools/list"] + assert len(discovery_requests) == 2 + assert "cursor" not in discovery_requests[0]["params"] + assert discovery_requests[1]["params"]["cursor"] == _CURSOR + tool_requests = [item for item in cold_requests if item["method"] == "tools/call"] + assert len(tool_requests) == int(allowed) + if allowed: + assert tool_requests[0]["params"]["name"] == "lookup_customer" + assert tool_requests[0]["params"]["arguments"] == {"id": "customer-1"} + else: + assert result.deny_reason == "upstream_error:UPSTREAM_UNAVAILABLE" + + # The shared acquisition is session-scoped, not a one-call shortcut. + repeated = await proxy.call_tool("warm", "lookup_customer", {"id": "customer-2"}) + assert repeated.allowed is allowed + all_requests = observed_requests() + assert sum(item["method"] == "tools/list" for item in all_requests) == 2 + assert sum(item["method"] == "tools/call" for item in all_requests) == 2 * int(allowed) + finally: + await proxy.aclose() diff --git a/tests/unit/test_stdio_upstream.py b/tests/unit/test_stdio_upstream.py index 4f5b28d5..7173a480 100644 --- a/tests/unit/test_stdio_upstream.py +++ b/tests/unit/test_stdio_upstream.py @@ -9,6 +9,7 @@ from __future__ import annotations +import asyncio import json import os import stat @@ -385,3 +386,164 @@ async def test_list_tools_returns_none_rather_than_raising(tmp_path) -> None: assert await server.list_tools() is None finally: await server.close() + + +def _paginated_list_server(tmp_path, pages): + return _script( + tmp_path, + f""" + import json, sys + pages = {pages!r} + requests = [] + for line in sys.stdin: + req = json.loads(line) + if req["method"] == "tools/list": + requests.append(req["params"]) + reply = pages[min(len(requests) - 1, len(pages) - 1)] + else: + reply = {{"result": {{"content": [ + {{"type": "text", "text": json.dumps(requests)}}, + ]}}}} + body = {{"jsonrpc": "2.0", "id": req["id"], **reply}} + sys.stdout.write(json.dumps(body) + "\\n") + sys.stdout.flush() + """, + ) + + +async def test_list_tools_exhausts_pages_with_opaque_and_empty_cursors(tmp_path) -> None: + tools = [ + {"name": "search", "description": "search", "inputSchema": {}}, + {"name": "fetch", "description": "fetch", "inputSchema": {}}, + ] + opaque_cursor = " page/%2F+雪== " + script = _paginated_list_server( + tmp_path, + [ + {"result": {"tools": tools[:1], "nextCursor": ""}}, + {"result": {"tools": [], "nextCursor": opaque_cursor}}, + {"result": {"tools": tools[1:]}}, + ], + ) + server = StdioServer(_spawn_for(script, None), allow_unmeasured=True) + await server.start() + try: + assert await asyncio.wait_for(server.list_tools(), timeout=2) == tools + requests = json.loads(await server.call("after-list", "requests", {})) + assert requests == [{}, {"cursor": ""}, {"cursor": opaque_cursor}] + finally: + await server.close() + + +@pytest.mark.parametrize( + "later_pages", + [ + [{"error": {"code": -32603, "message": "listing failed"}}], + [{"result": {"tools": {}}}], + [{"result": {"tools": [{"name": 7, "inputSchema": {}}]}}], + [{"result": {"tools": [], "nextCursor": None}}], + [{"result": {"tools": [], "nextCursor": 0}}], + [{"result": {"tools": [{"name": "search", "inputSchema": {}}]}}], + [{"result": {"tools": [], "nextCursor": "next"}}], + [ + {"result": {"tools": [], "nextCursor": "another"}}, + {"result": {"tools": [], "nextCursor": "next"}}, + ], + ], + ids=[ + "rpc-error", "malformed-tools", "malformed-name", "null-cursor", "numeric-cursor", + "duplicate-name", "repeated-cursor", "cursor-cycle", + ], +) +async def test_list_tools_discards_partial_pages_without_desynchronizing( + tmp_path, later_pages +) -> None: + script = _paginated_list_server( + tmp_path, + [ + {"result": {"tools": [{"name": "search", "inputSchema": {}}], + "nextCursor": "next"}}, + *later_pages, + ], + ) + server = StdioServer(_spawn_for(script, None), allow_unmeasured=True) + await server.start() + try: + assert await asyncio.wait_for(server.list_tools(), timeout=2) is None + requests = json.loads(await server.call("after-list", "requests", {})) + expected = [{}, {"cursor": "next"}] + if len(later_pages) == 2: + expected.append({"cursor": "another"}) + assert requests == expected + finally: + await server.close() + + +async def test_mismatched_later_list_response_discards_pages_and_closes_child( + tmp_path, caplog +) -> None: + script = _paginated_list_server( + tmp_path, + [ + {"result": {"tools": [{"name": "search", "inputSchema": {}}], + "nextCursor": "next"}}, + {"id": "wrong-page-id-do-not-log", "result": { + "tools": [{"name": "fetch", "inputSchema": {}}], + }}, + ], + ) + server = StdioServer(_spawn_for(script, None), allow_unmeasured=True) + await server.start() + proc = server._proc + assert proc is not None + try: + assert await asyncio.wait_for(server.list_tools(), timeout=2) is None + assert proc.returncode is not None + assert server._proc is None + assert "tools discovery incomplete: upstream request failed" in caplog.text + assert "wrong-page-id-do-not-log" not in caplog.text + with pytest.raises(UpstreamUnavailable, match="not running"): + await server.call("after-invalid-list", "search", {}) + finally: + await server.close() + + +async def test_cancelled_later_list_page_closes_child_before_another_call(tmp_path) -> None: + script = _script( + tmp_path, + """ + import json, sys + for line in sys.stdin: + req = json.loads(line) + if req["params"].get("cursor") == "second": + sys.stderr.buffer.write(b"later-page-received\\n") + sys.stderr.flush() + sys.stdin.readline() + result = {"tools": [{"name": "late", "inputSchema": {}}]} + else: + result = {"tools": [], "nextCursor": "second"} + sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": req["id"], + "result": result}) + "\\n") + sys.stdout.flush() + """, + ) + server = StdioServer(_spawn_for(script, None), allow_unmeasured=True) + await server.start() + proc = server._proc + assert proc is not None and proc.stderr is not None + task = asyncio.create_task(server.list_tools()) + try: + assert await asyncio.wait_for(proc.stderr.readline(), timeout=2) == ( + b"later-page-received\n" + ) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=2) + assert proc.returncode is not None + assert server._proc is None + with pytest.raises(UpstreamUnavailable, match="not running"): + await server.call("after-cancel", "search", {}) + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await server.close() diff --git a/tests/unit/test_tool_discovery.py b/tests/unit/test_tool_discovery.py new file mode 100644 index 00000000..8c2b9a99 --- /dev/null +++ b/tests/unit/test_tool_discovery.py @@ -0,0 +1,112 @@ +"""Issue #631: acquisition must finish before a catalog becomes comparable.""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock + +import pytest + +from cmcp_runtime.mcp import discovery +from cmcp_runtime.mcp.discovery import DiscoveryError, collect_tools + + +def _pages(*results): + remaining = iter(results) + + async def fetch(request_id, params): + return {"jsonrpc": "2.0", "id": request_id, "result": next(remaining)} + + return AsyncMock(side_effect=fetch) + + +@pytest.mark.parametrize("cursor", ["", "opaque /+==\n雪", "0"]) +async def test_empty_page_with_opaque_cursor_is_not_terminal(cursor): + tool = {"name": "lookup", "input_schema": {}, "extension": {"kept": True}} + fetch = _pages({"tools": [], "nextCursor": cursor}, {"tools": [tool]}) + assert await collect_tools(fetch) == [tool] + assert fetch.call_args_list[0].args[1] == {} + assert fetch.call_args_list[1].args[1] == {"cursor": cursor} + assert fetch.call_args_list[0].args[0] != fetch.call_args_list[1].args[0] + + +async def test_terminal_empty_page_completes_empty_catalog(): + assert await collect_tools(_pages({"tools": []})) == [] + + +@pytest.mark.parametrize( + "bad_page", + [ + None, + [], + {}, + {"tools": None}, + {"tools": {}}, + {"tools": [None]}, + {"tools": [{"name": 1}]}, + {"tools": [{}]}, + {"tools": [{"name": ""}]}, + {"tools": [], "nextCursor": None}, + {"tools": [], "nextCursor": 0}, + {"tools": [], "nextCursor": []}, + ], +) +async def test_malformed_later_page_never_returns_partial_catalog(bad_page): + fetch = _pages({"tools": [{"name": "lookup"}], "nextCursor": "next"}, bad_page) + with pytest.raises(DiscoveryError): + await collect_tools(fetch) + assert fetch.await_count == 2 + + +@pytest.mark.parametrize( + "body", + [ + None, + [], + {}, + {"jsonrpc": "1.0", "id": "provenance-tools-list-0", "result": {"tools": []}}, + {"jsonrpc": "2.0", "id": "wrong", "result": {"tools": []}}, + {"jsonrpc": "2.0", "id": "provenance-tools-list-0", "error": {"message": "secret"}}, + {"jsonrpc": "2.0", "id": "provenance-tools-list-0", "error": None, "result": {"tools": []}}, + ], +) +async def test_malformed_or_error_envelope_is_not_a_listing(body): + with pytest.raises(DiscoveryError, match="^invalid response envelope$"): + await collect_tools(AsyncMock(return_value=body)) + + +@pytest.mark.parametrize("cursors", [["secret", "secret"], ["secret", "other", "secret"]]) +async def test_repeated_and_cyclic_cursors_terminate_without_exposing_cursor(cursors): + fetch = _pages(*({"tools": [], "nextCursor": cursor} for cursor in cursors)) + with pytest.raises(DiscoveryError, match="^repeated continuation cursor$"): + await collect_tools(fetch) + assert fetch.await_count == len(cursors) + + +@pytest.mark.parametrize("same_page", [False, True]) +async def test_duplicate_names_are_ambiguous_even_with_matching_definitions(same_page): + tool = {"name": "lookup"} + fetch = ( + _pages({"tools": [tool, tool]}) + if same_page + else _pages({"tools": [tool], "nextCursor": "next"}, {"tools": [tool]}) + ) + with pytest.raises(DiscoveryError, match="^duplicate tool name$"): + await collect_tools(fetch) + + +async def test_exact_page_budget_can_complete_but_cannot_return_partial(monkeypatch): + monkeypatch.setattr(discovery, "MAX_DISCOVERY_PAGES", 2) + first = {"tools": [{"name": "lookup"}], "nextCursor": "one"} + complete = _pages(first, {"tools": []}) + assert await collect_tools(complete) == first["tools"] + continuing = _pages(first, {"tools": [], "nextCursor": "two"}) + with pytest.raises(DiscoveryError, match="^discovery page limit exceeded$"): + await collect_tools(continuing) + assert continuing.await_count == 2 + + +@pytest.mark.parametrize("failure", [RuntimeError("upstream failure"), asyncio.CancelledError()]) +async def test_transport_failure_and_cancellation_propagate(failure): + with pytest.raises(type(failure)): + await collect_tools(AsyncMock(side_effect=failure)) diff --git a/tests/unit/test_upstream_catalog_drift.py b/tests/unit/test_upstream_catalog_drift.py index 443d2436..195673f4 100644 --- a/tests/unit/test_upstream_catalog_drift.py +++ b/tests/unit/test_upstream_catalog_drift.py @@ -8,9 +8,13 @@ from __future__ import annotations +import asyncio +import json +import logging import uuid from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from cmcp_runtime.audit.chain import AuditChain @@ -32,6 +36,7 @@ TEEProvider, ) from cmcp_runtime.mcp.proxy import CMCPProxy +from cmcp_runtime.provenance import ProvenanceOutcome from cmcp_runtime.session.state import SessionState APPROVED_DESCRIPTION = "Look up a customer record by id." @@ -227,3 +232,280 @@ async def test_drift_is_caught_without_the_optional_scanner(): assert await proxy._check_upstream_drift(catalog.entries["lookup_customer"]) is True assert session.catalog_drift is True + + +# --- discovery must finish before it can produce a verdict (#631) ---------- + + +def _discovery_response(request: httpx.Request, reply: dict, response_format: str): + payload = json.loads(request.content) + assert payload["method"] == "tools/list" + assert request.headers["Accept"] == "application/json, text/event-stream" + assert request.headers["Mcp-Method"] == "tools/list" + assert payload["params"]["_meta"][ + "io.modelcontextprotocol/protocolVersion" + ] == request.headers["MCP-Protocol-Version"] + body = {"jsonrpc": "2.0", "id": payload["id"], **reply} + if response_format == "sse": + notification = {"jsonrpc": "2.0", "method": "notifications/progress"} + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=f"data: {json.dumps(notification)}\n\ndata: {json.dumps(body)}\n\n", + ) + return httpx.Response(200, json=body) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("response_format", ["json", "sse"]) +async def test_http_discovery_preserves_opaque_cursors_and_empty_pages( + monkeypatch, response_format +): + catalog = _catalog() + proxy, _, _ = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + opaque_cursor = " page/%2F+雪== " + pages = [ + {"tools": [], "nextCursor": ""}, + {"tools": [], "nextCursor": opaque_cursor}, + {"tools": _advertise()}, + ] + requests = [] + + def respond(request): + payload = json.loads(request.content) + requests.append(payload) + return _discovery_response( + request, {"result": pages[len(requests) - 1]}, response_format + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + monkeypatch.setattr(proxy, "_client_for_upstream", lambda entry: client) + assert await proxy._advertised_tools(catalog.entries["lookup_customer"]) == _advertise() + + assert len(requests) == 3 + assert "cursor" not in requests[0]["params"] + assert requests[1]["params"]["cursor"] == "" + assert requests[2]["params"]["cursor"] == opaque_cursor + + +@pytest.mark.asyncio +@pytest.mark.parametrize("response_format", ["json", "sse"]) +@pytest.mark.parametrize( + "description,drift_policy,denied", + [ + (APPROVED_DESCRIPTION, DriftPolicy.FAIL_CLOSED, False), + ("changed description", DriftPolicy.FAIL_CLOSED, True), + ("changed description", DriftPolicy.WARN_ONLY, False), + ], +) +async def test_http_drift_compares_the_tool_on_the_second_page( + monkeypatch, response_format, description, drift_policy, denied +): + catalog = _catalog() + proxy, session, chain = _proxy(catalog, drift_policy=drift_policy) + requests = [] + + def respond(request): + payload = json.loads(request.content) + requests.append(payload) + result = ( + {"tools": [], "nextCursor": "second"} + if "cursor" not in payload["params"] + else {"tools": _advertise(description)} + ) + return _discovery_response(request, {"result": result}, response_format) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + monkeypatch.setattr(proxy, "_client_for_upstream", lambda entry: client) + entry = catalog.entries["lookup_customer"] + assert await proxy._check_upstream_drift(entry) is denied + assert await proxy._check_upstream_drift(entry) is denied + + assert len(requests) == 2 + assert requests[1]["params"]["cursor"] == "second" + assert session.catalog_drift is denied + drift_entries = [e for e in chain.entries if e.entry_type == "catalog_drift"] + if description == APPROVED_DESCRIPTION: + assert session.upstream_drift_tools == [] + assert drift_entries == [] + else: + assert session.upstream_drift_tools == ["lookup_customer"] + assert len(drift_entries) == 1 + assert drift_entries[0].detail["kind"] == "definition_changed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("response_format", ["json", "sse"]) +@pytest.mark.parametrize( + "failure,first_tools", + [ + ("rpc_error", []), + ("http_error", _advertise()), + ("malformed_result", _advertise("changed description")), + ("malformed_tools", []), + ("malformed_name", _advertise()), + ("null_cursor", _advertise("changed description")), + ("numeric_cursor", []), + ("duplicate_name", _advertise()), + ("repeated_cursor", _advertise()), + ("cursor_cycle", _advertise("changed description")), + ], +) +async def test_http_incomplete_discovery_is_unchecked_not_drift( + monkeypatch, caplog, response_format, failure, first_tools +): + catalog = _catalog() + proxy, session, chain = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + requests = [] + + def respond(request): + payload = json.loads(request.content) + requests.append(payload) + if len(requests) == 1: + reply = {"result": {"tools": first_tools, "nextCursor": "next"}} + elif failure == "rpc_error": + reply = {"error": {"code": -32603, "message": "listing failed"}} + elif failure == "http_error": + return httpx.Response(503) + elif failure == "malformed_result": + reply = {"result": []} + elif failure == "malformed_tools": + reply = {"result": {"tools": {}}} + elif failure == "malformed_name": + reply = {"result": {"tools": [{"name": 7, "inputSchema": {}}]}} + elif failure == "null_cursor": + reply = {"result": {"tools": [], "nextCursor": None}} + elif failure == "numeric_cursor": + reply = {"result": {"tools": [], "nextCursor": 0}} + elif failure == "duplicate_name": + reply = {"result": {"tools": _advertise("changed description")}} + elif failure == "cursor_cycle" and len(requests) == 2: + reply = {"result": {"tools": [], "nextCursor": "another"}} + else: + reply = {"result": {"tools": [], "nextCursor": "next"}} + return _discovery_response(request, reply, response_format) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + monkeypatch.setattr(proxy, "_client_for_upstream", lambda entry: client) + with caplog.at_level(logging.INFO, logger="cmcp_runtime.mcp.proxy"): + entry = catalog.entries["lookup_customer"] + assert await proxy._check_upstream_drift(entry) is False + assert await proxy._check_upstream_drift(entry) is False + + assert len(requests) == (3 if failure == "cursor_cycle" else 2) + assert session.catalog_drift is False + assert session.upstream_drift_tools == [] + assert not any(e.entry_type == "catalog_drift" for e in chain.entries) + assert "outcome=unchecked" in caplog.text + assert "outcome=match" not in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("response_format", ["json", "sse"]) +@pytest.mark.parametrize( + "discovery,outcome", + [ + ("matching", ProvenanceOutcome.VERIFIED), + ("changed", ProvenanceOutcome.CATALOG_MISMATCH), + ("malformed", ProvenanceOutcome.UNCHECKED), + ("rpc_error", ProvenanceOutcome.UNCHECKED), + ("cursor_cycle", ProvenanceOutcome.UNCHECKED), + ], +) +async def test_http_paginated_discovery_preserves_signed_provenance_outcomes( + tmp_path, monkeypatch, response_format, discovery, outcome +): + from agentrust_trace.provenance import build_record, sign_record + from agentrust_trace.sign import generate_key, key_to_jwk + + first_tool = {"name": "server_info", "description": "server information", "inputSchema": {}} + key = generate_key() + record = build_record( + kind="publisher-asserted", + publisher="did:web:crm.example", + tools=[first_tool, *_advertise()], + artifact={"package": "pkg:npm/crm@1.0.0", "digest": "sha256:" + "a" * 64}, + ) + record_path = tmp_path / "provenance.json" + record_path.write_text(json.dumps(sign_record(record, key)), encoding="utf-8") + catalog = _catalog() + entry = catalog.entries["lookup_customer"] + entry.server.provenance_record_path = str(record_path) + entry.server.publisher_jwk = key_to_jwk(key) + proxy, _, _ = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + requests = [] + + def respond(request): + requests.append(json.loads(request.content)) + if len(requests) == 1: + reply = {"result": {"tools": [first_tool], "nextCursor": "next"}} + elif discovery == "matching": + reply = {"result": {"tools": _advertise()}} + elif discovery == "changed": + reply = {"result": {"tools": _advertise("changed description")}} + elif discovery == "malformed": + reply = {"result": {"tools": None}} + elif discovery == "rpc_error": + reply = {"error": {"code": -32603, "message": "listing failed"}} + else: + reply = {"result": {"tools": [], "nextCursor": "next"}} + return _discovery_response(request, reply, response_format) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + monkeypatch.setattr(proxy, "_client_for_upstream", lambda entry: client) + result = await proxy._check_provenance(entry) + assert result.outcome is outcome + assert result.kind == "publisher-asserted" + assert result.publisher == "did:web:crm.example" + assert await proxy._check_provenance(entry) is result + + assert len(requests) == 2 + assert requests[1]["params"]["cursor"] == "next" + + +@pytest.mark.asyncio +async def test_http_cancelled_later_page_does_not_cache_a_drift_check(monkeypatch): + catalog = _catalog() + proxy, session, chain = _proxy(catalog, drift_policy=DriftPolicy.FAIL_CLOSED) + entry = catalog.entries["lookup_customer"] + later_page_entered = asyncio.Event() + pending_response = asyncio.Event() + requests = [] + + async def respond(request): + payload = json.loads(request.content) + requests.append(payload) + if len(requests) == 2: + later_page_entered.set() + await pending_response.wait() + result = ( + {"tools": [], "nextCursor": "second"} + if "cursor" not in payload["params"] + else {"tools": _advertise("changed description")} + ) + return _discovery_response(request, {"result": result}, "json") + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + monkeypatch.setattr(proxy, "_client_for_upstream", lambda entry: client) + task = asyncio.create_task(proxy._check_upstream_drift(entry)) + try: + await asyncio.wait_for(later_page_entered.wait(), timeout=2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert proxy._drift_checked == set() + assert session.catalog_drift is False + assert session.upstream_drift_tools == [] + assert not any(e.entry_type == "catalog_drift" for e in chain.entries) + + assert await proxy._check_upstream_drift(entry) is True + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert len(requests) == 4 + assert "cursor" not in requests[2]["params"] + assert requests[3]["params"]["cursor"] == "second" + drift_entries = [e for e in chain.entries if e.entry_type == "catalog_drift"] + assert len(drift_entries) == 1 + assert drift_entries[0].detail["kind"] == "definition_changed"