454 lines
18 KiB
Python
454 lines
18 KiB
Python
"""MCP unified conformance test client.
|
|
|
|
This client is designed to work with the @modelcontextprotocol/conformance npm package.
|
|
It handles all conformance test scenarios via environment variables and CLI arguments.
|
|
|
|
Contract:
|
|
- MCP_CONFORMANCE_SCENARIO env var -> scenario name
|
|
- MCP_CONFORMANCE_CONTEXT env var -> optional JSON (for client-credentials scenarios)
|
|
- MCP_CONFORMANCE_PROTOCOL_VERSION env var -> spec version the harness mock
|
|
server is speaking (e.g. "2025-11-25", "2026-07-28"). Always set; when
|
|
--spec-version is omitted the harness picks per-scenario (LATEST_SPEC_VERSION
|
|
for active scenarios, DRAFT_PROTOCOL_VERSION for draft-only ones).
|
|
- Server URL as last CLI argument (sys.argv[1])
|
|
- Must exit 0 within 30 seconds
|
|
|
|
Scenarios:
|
|
initialize - Connect, initialize, list tools, close
|
|
tools_call - Connect, call add_numbers(a=5, b=3), close
|
|
sse-retry - Connect, call test_reconnection, close
|
|
json-schema-ref-no-deref - Connect, list tools (no $ref deref)
|
|
request-metadata - Connect with all callbacks; client stamps _meta
|
|
http-standard-headers - Connect, call a tool (Mcp-* headers checked)
|
|
elicitation-sep1034-client-defaults - Elicitation with default accept callback
|
|
auth/client-credentials-jwt - Client credentials with private_key_jwt
|
|
auth/client-credentials-basic - Client credentials with client_secret_basic
|
|
auth/* - Authorization code flow (default for auth scenarios)
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
from collections.abc import Callable, Coroutine
|
|
from typing import Any, cast
|
|
from urllib.parse import parse_qs, urlparse
|
|
|
|
import httpx
|
|
from pydantic import AnyUrl
|
|
|
|
from mcp import types
|
|
from mcp.client.auth import OAuthClientProvider, TokenStorage
|
|
from mcp.client.auth.extensions.client_credentials import (
|
|
ClientCredentialsOAuthProvider,
|
|
PrivateKeyJWTOAuthProvider,
|
|
SignedJWTParameters,
|
|
)
|
|
from mcp.client.client import Client
|
|
from mcp.client.context import ClientRequestContext
|
|
from mcp.client.streamable_http import streamable_http_client
|
|
from mcp.shared.auth import AuthorizationCodeResult, OAuthClientInformationFull, OAuthClientMetadata, OAuthToken
|
|
from mcp.shared.version import MODERN_PROTOCOL_VERSIONS
|
|
|
|
# Set up logging to stderr (stdout is for conformance test output)
|
|
logging.basicConfig(
|
|
level=logging.DEBUG,
|
|
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
|
stream=sys.stderr,
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
#: Spec version the harness is running this scenario at (e.g. "2025-11-25",
|
|
#: "2026-07-28"). The harness always sets this (when --spec-version is omitted
|
|
#: it picks per-scenario: LATEST_SPEC_VERSION for active scenarios,
|
|
#: DRAFT_PROTOCOL_VERSION for draft-only ones), so None means we were invoked
|
|
#: outside the harness.
|
|
PROTOCOL_VERSION: str | None = os.environ.get("MCP_CONFORMANCE_PROTOCOL_VERSION")
|
|
|
|
|
|
def client_mode() -> str:
|
|
"""Pick the Client(mode=) for the harness leg.
|
|
|
|
On a modern leg (2026-07-28+) -> 'auto' so Client.discover() runs and the
|
|
_meta envelope + MCP-Protocol-Version header are stamped on every request.
|
|
On a handshake-era leg -> 'legacy' so the initialize handshake runs exactly
|
|
as before (no server/discover probe is sent against a mock that would 400 it).
|
|
Outside the harness -> 'auto' (probe + fallback).
|
|
"""
|
|
if PROTOCOL_VERSION is None or PROTOCOL_VERSION in MODERN_PROTOCOL_VERSIONS:
|
|
return "auto"
|
|
return "legacy"
|
|
|
|
|
|
# Type for async scenario handler functions
|
|
ScenarioHandler = Callable[[str], Coroutine[Any, None, None]]
|
|
|
|
# Registry of scenario handlers
|
|
HANDLERS: dict[str, ScenarioHandler] = {}
|
|
|
|
|
|
def register(name: str) -> Callable[[ScenarioHandler], ScenarioHandler]:
|
|
"""Register a scenario handler."""
|
|
|
|
def decorator(fn: ScenarioHandler) -> ScenarioHandler:
|
|
HANDLERS[name] = fn
|
|
return fn
|
|
|
|
return decorator
|
|
|
|
|
|
def get_conformance_context() -> dict[str, Any]:
|
|
"""Load conformance test context from MCP_CONFORMANCE_CONTEXT environment variable."""
|
|
context_json = os.environ.get("MCP_CONFORMANCE_CONTEXT")
|
|
if not context_json:
|
|
raise RuntimeError(
|
|
"MCP_CONFORMANCE_CONTEXT environment variable not set. "
|
|
"Expected JSON with client_id, client_secret, and/or private_key_pem."
|
|
)
|
|
try:
|
|
return json.loads(context_json)
|
|
except json.JSONDecodeError as e:
|
|
raise RuntimeError(f"Failed to parse MCP_CONFORMANCE_CONTEXT as JSON: {e}") from e
|
|
|
|
|
|
class InMemoryTokenStorage(TokenStorage):
|
|
"""Simple in-memory token storage for conformance testing."""
|
|
|
|
def __init__(self) -> None:
|
|
self._tokens: OAuthToken | None = None
|
|
self._client_info: OAuthClientInformationFull | None = None
|
|
|
|
async def get_tokens(self) -> OAuthToken | None:
|
|
return self._tokens
|
|
|
|
async def set_tokens(self, tokens: OAuthToken) -> None:
|
|
self._tokens = tokens
|
|
|
|
async def get_client_info(self) -> OAuthClientInformationFull | None:
|
|
return self._client_info
|
|
|
|
async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
|
|
self._client_info = client_info
|
|
|
|
|
|
class ConformanceOAuthCallbackHandler:
|
|
"""OAuth callback handler that automatically fetches the authorization URL
|
|
and extracts the auth code, without requiring user interaction.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._auth_code: str | None = None
|
|
self._state: str | None = None
|
|
self._iss: str | None = None
|
|
|
|
async def handle_redirect(self, authorization_url: str) -> None:
|
|
"""Fetch the authorization URL and extract the auth code from the redirect."""
|
|
logger.debug(f"Fetching authorization URL: {authorization_url}")
|
|
|
|
async with httpx.AsyncClient() as client:
|
|
response = await client.get(
|
|
authorization_url,
|
|
follow_redirects=False,
|
|
)
|
|
|
|
if response.status_code in (301, 302, 303, 307, 308):
|
|
location = cast(str, response.headers.get("location"))
|
|
if location:
|
|
redirect_url = urlparse(location)
|
|
query_params: dict[str, list[str]] = parse_qs(redirect_url.query)
|
|
|
|
if "code" in query_params:
|
|
self._auth_code = query_params["code"][0]
|
|
state_values = query_params.get("state")
|
|
self._state = state_values[0] if state_values else None
|
|
iss_values = query_params.get("iss")
|
|
self._iss = iss_values[0] if iss_values else None
|
|
logger.debug(f"Got auth code from redirect: {self._auth_code[:10]}...")
|
|
return
|
|
else:
|
|
raise RuntimeError(f"No auth code in redirect URL: {location}")
|
|
else:
|
|
raise RuntimeError(f"No redirect location received from {authorization_url}")
|
|
else:
|
|
raise RuntimeError(f"Expected redirect response, got {response.status_code} from {authorization_url}")
|
|
|
|
async def handle_callback(self) -> AuthorizationCodeResult:
|
|
"""Return the captured auth code, state, and iss."""
|
|
if self._auth_code is None:
|
|
raise RuntimeError("No authorization code available - was handle_redirect called?")
|
|
result = AuthorizationCodeResult(code=self._auth_code, state=self._state, iss=self._iss)
|
|
self._auth_code = None
|
|
self._state = None
|
|
self._iss = None
|
|
return result
|
|
|
|
|
|
# --- Stub callbacks (declare capabilities in _meta without doing real work) ---
|
|
|
|
|
|
async def stub_sampling_callback(
|
|
context: ClientRequestContext,
|
|
params: types.CreateMessageRequestParams,
|
|
) -> types.CreateMessageResult | types.ErrorData:
|
|
return types.CreateMessageResult(
|
|
role="assistant",
|
|
content=types.TextContent(type="text", text=""),
|
|
model="conformance-stub",
|
|
)
|
|
|
|
|
|
async def stub_list_roots_callback(context: ClientRequestContext) -> types.ListRootsResult | types.ErrorData:
|
|
return types.ListRootsResult(roots=[])
|
|
|
|
|
|
async def default_elicitation_callback(
|
|
context: ClientRequestContext,
|
|
params: types.ElicitRequestParams,
|
|
) -> types.ElicitResult | types.ErrorData:
|
|
"""Accept elicitation and apply defaults from the schema (SEP-1034)."""
|
|
content: dict[str, str | int | float | bool | list[str] | None] = {}
|
|
|
|
# For form mode, extract defaults from the requested_schema
|
|
if isinstance(params, types.ElicitRequestFormParams):
|
|
schema = params.requested_schema
|
|
logger.debug(f"Elicitation schema: {schema}")
|
|
properties = schema.get("properties", {})
|
|
for prop_name, prop_schema in properties.items():
|
|
if "default" in prop_schema:
|
|
content[prop_name] = prop_schema["default"]
|
|
logger.debug(f"Applied defaults: {content}")
|
|
|
|
return types.ElicitResult(action="accept", content=content)
|
|
|
|
|
|
# --- Scenario Handlers ---
|
|
|
|
|
|
@register("initialize")
|
|
async def run_initialize(server_url: str) -> None:
|
|
"""Connect, initialize, list tools, close."""
|
|
async with Client(server_url, mode=client_mode()) as client:
|
|
logger.debug("Initialized successfully")
|
|
await client.list_tools()
|
|
logger.debug("Listed tools successfully")
|
|
|
|
|
|
@register("json-schema-ref-no-deref")
|
|
async def run_json_schema_ref_no_deref(server_url: str) -> None:
|
|
"""Initialize and list tools; the scenario fails only if the client fetches a network $ref.
|
|
|
|
The client never walks inputSchema or resolves $refs, so listing is enough (SEP-2106).
|
|
Pinned to mode='legacy': the harness reports PROTOCOL_VERSION=2026-07-28 for this
|
|
scenario but its mock server only speaks the handshake-era lifecycle and 400s a
|
|
modern-stamped tools/list. The check is lifecycle-agnostic so this is harmless.
|
|
"""
|
|
async with Client(server_url, mode="legacy") as client:
|
|
await client.list_tools()
|
|
|
|
|
|
@register("tools_call")
|
|
async def run_tools_call(server_url: str) -> None:
|
|
"""Connect, list tools, call add_numbers(a=5, b=3), close."""
|
|
async with Client(server_url, mode=client_mode()) as client:
|
|
await client.list_tools()
|
|
result = await client.call_tool("add_numbers", {"a": 5, "b": 3})
|
|
logger.debug(f"add_numbers result: {result}")
|
|
|
|
|
|
@register("sse-retry")
|
|
async def run_sse_retry(server_url: str) -> None:
|
|
"""Connect, list tools, call test_reconnection, close."""
|
|
async with Client(server_url, mode=client_mode()) as client:
|
|
await client.list_tools()
|
|
result = await client.call_tool("test_reconnection", {})
|
|
logger.debug(f"test_reconnection result: {result}")
|
|
|
|
|
|
@register("request-metadata")
|
|
async def run_request_metadata(server_url: str) -> None:
|
|
"""Connect on the modern path with every client capability declared.
|
|
|
|
The scenario inspects every request's `_meta` envelope (SEP-2575) for
|
|
protocolVersion / clientInfo / clientCapabilities, and the matching
|
|
MCP-Protocol-Version header. mode='auto' makes the SDK send
|
|
server/discover (covering the unsupported-version retry check), then adopt
|
|
and stamp the envelope on the follow-up requests.
|
|
"""
|
|
async with Client(
|
|
server_url,
|
|
mode=client_mode(),
|
|
sampling_callback=stub_sampling_callback,
|
|
list_roots_callback=stub_list_roots_callback,
|
|
elicitation_callback=default_elicitation_callback,
|
|
) as client:
|
|
await client.list_tools()
|
|
result = await client.call_tool("add_numbers", {"a": 5, "b": 3})
|
|
logger.debug(f"add_numbers result: {result}")
|
|
|
|
|
|
@register("http-standard-headers")
|
|
async def run_http_standard_headers(server_url: str) -> None:
|
|
"""Connect on the modern path so Mcp-Method / Mcp-Name / MCP-Protocol-Version are sent (SEP-2243)."""
|
|
async with Client(server_url, mode=client_mode()) as client:
|
|
await client.list_tools()
|
|
result = await client.call_tool("add_numbers", {"a": 5, "b": 3})
|
|
logger.debug(f"add_numbers result: {result}")
|
|
|
|
|
|
@register("elicitation-sep1034-client-defaults")
|
|
async def run_elicitation_defaults(server_url: str) -> None:
|
|
"""Connect with elicitation callback that applies schema defaults."""
|
|
async with Client(server_url, mode=client_mode(), elicitation_callback=default_elicitation_callback) as client:
|
|
await client.list_tools()
|
|
result = await client.call_tool("test_client_elicitation_defaults", {})
|
|
logger.debug(f"test_client_elicitation_defaults result: {result}")
|
|
|
|
|
|
@register("auth/client-credentials-jwt")
|
|
async def run_client_credentials_jwt(server_url: str) -> None:
|
|
"""Client credentials flow with private_key_jwt authentication."""
|
|
context = get_conformance_context()
|
|
client_id = context.get("client_id")
|
|
private_key_pem = context.get("private_key_pem")
|
|
signing_algorithm = context.get("signing_algorithm", "ES256")
|
|
|
|
if not client_id:
|
|
raise RuntimeError("MCP_CONFORMANCE_CONTEXT missing 'client_id'")
|
|
if not private_key_pem:
|
|
raise RuntimeError("MCP_CONFORMANCE_CONTEXT missing 'private_key_pem'")
|
|
|
|
jwt_params = SignedJWTParameters(
|
|
issuer=client_id,
|
|
subject=client_id,
|
|
signing_algorithm=signing_algorithm,
|
|
signing_key=private_key_pem,
|
|
)
|
|
|
|
oauth_auth = PrivateKeyJWTOAuthProvider(
|
|
server_url=server_url,
|
|
storage=InMemoryTokenStorage(),
|
|
client_id=client_id,
|
|
assertion_provider=jwt_params.create_assertion_provider(),
|
|
)
|
|
|
|
await _run_auth_session(server_url, oauth_auth)
|
|
|
|
|
|
@register("auth/client-credentials-basic")
|
|
async def run_client_credentials_basic(server_url: str) -> None:
|
|
"""Client credentials flow with client_secret_basic authentication."""
|
|
context = get_conformance_context()
|
|
client_id = context.get("client_id")
|
|
client_secret = context.get("client_secret")
|
|
|
|
if not client_id:
|
|
raise RuntimeError("MCP_CONFORMANCE_CONTEXT missing 'client_id'")
|
|
if not client_secret:
|
|
raise RuntimeError("MCP_CONFORMANCE_CONTEXT missing 'client_secret'")
|
|
|
|
oauth_auth = ClientCredentialsOAuthProvider(
|
|
server_url=server_url,
|
|
storage=InMemoryTokenStorage(),
|
|
client_id=client_id,
|
|
client_secret=client_secret,
|
|
token_endpoint_auth_method="client_secret_basic",
|
|
)
|
|
|
|
await _run_auth_session(server_url, oauth_auth)
|
|
|
|
|
|
async def run_auth_code_client(server_url: str) -> None:
|
|
"""Authorization code flow (default for auth/* scenarios)."""
|
|
callback_handler = ConformanceOAuthCallbackHandler()
|
|
storage = InMemoryTokenStorage()
|
|
|
|
# Check for pre-registered client credentials from context
|
|
context_json = os.environ.get("MCP_CONFORMANCE_CONTEXT")
|
|
if context_json:
|
|
try:
|
|
context = json.loads(context_json)
|
|
client_id = context.get("client_id")
|
|
client_secret = context.get("client_secret")
|
|
if client_id:
|
|
await storage.set_client_info(
|
|
OAuthClientInformationFull(
|
|
client_id=client_id,
|
|
client_secret=client_secret,
|
|
redirect_uris=[AnyUrl("http://localhost:3000/callback")],
|
|
token_endpoint_auth_method="client_secret_basic" if client_secret else "none",
|
|
)
|
|
)
|
|
logger.debug(f"Pre-loaded client credentials: client_id={client_id}")
|
|
except json.JSONDecodeError:
|
|
logger.exception("Failed to parse MCP_CONFORMANCE_CONTEXT")
|
|
|
|
oauth_auth = OAuthClientProvider(
|
|
server_url=server_url,
|
|
client_metadata=OAuthClientMetadata(
|
|
client_name="conformance-client",
|
|
redirect_uris=[AnyUrl("http://localhost:3000/callback")],
|
|
grant_types=["authorization_code", "refresh_token"],
|
|
response_types=["code"],
|
|
),
|
|
storage=storage,
|
|
redirect_handler=callback_handler.handle_redirect,
|
|
callback_handler=callback_handler.handle_callback,
|
|
client_metadata_url="https://conformance-test.local/client-metadata.json",
|
|
)
|
|
|
|
await _run_auth_session(server_url, oauth_auth)
|
|
|
|
|
|
async def _run_auth_session(server_url: str, oauth_auth: OAuthClientProvider) -> None:
|
|
"""Common session logic for all OAuth flows."""
|
|
http_client = httpx.AsyncClient(auth=oauth_auth, timeout=30.0)
|
|
transport = streamable_http_client(url=server_url, http_client=http_client)
|
|
async with Client(transport, mode=client_mode(), elicitation_callback=default_elicitation_callback) as client:
|
|
logger.debug("Initialized successfully")
|
|
|
|
tools_result = await client.list_tools()
|
|
logger.debug(f"Listed tools: {[t.name for t in tools_result.tools]}")
|
|
|
|
# Call the first available tool (different tests have different tools)
|
|
if tools_result.tools:
|
|
tool_name = tools_result.tools[0].name
|
|
try:
|
|
result = await client.call_tool(tool_name, {})
|
|
logger.debug(f"Called {tool_name}, result: {result}")
|
|
except Exception as e:
|
|
logger.debug(f"Tool call result/error: {e}")
|
|
|
|
logger.debug("Connection closed successfully")
|
|
|
|
|
|
def main() -> None:
|
|
"""Main entry point for the conformance client."""
|
|
if len(sys.argv) < 2:
|
|
print(f"Usage: {sys.argv[0]} <server-url>", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
server_url = sys.argv[1]
|
|
scenario = os.environ.get("MCP_CONFORMANCE_SCENARIO")
|
|
logger.debug(f"Conformance protocol version: {PROTOCOL_VERSION!r} -> mode={client_mode()!r}")
|
|
|
|
if scenario:
|
|
logger.debug(f"Running explicit scenario '{scenario}' against {server_url}")
|
|
handler = HANDLERS.get(scenario)
|
|
if handler:
|
|
asyncio.run(handler(server_url))
|
|
elif scenario.startswith("auth/"):
|
|
asyncio.run(run_auth_code_client(server_url))
|
|
else:
|
|
# Unhandled scenarios:
|
|
# - sep-2322-client-request-state (SEP-2322 / S6: MRTR client loop)
|
|
# - http-custom-headers, http-invalid-tool-headers (SEP-2243 / S8: Mcp-Param-* headers)
|
|
print(f"Unknown scenario: {scenario}", file=sys.stderr)
|
|
sys.exit(1)
|
|
else:
|
|
logger.debug(f"Running default auth flow against {server_url}")
|
|
asyncio.run(run_auth_code_client(server_url))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|