Max Isbey
1b74b06753
Cut comment and docstring volume roughly in half across src, tests, examples, and docs_src: removed comments that restate the adjacent code, leftover development narration, section banners, and self-evident Args/Returns blocks, and compressed the remaining docstrings to a Google-style summary line plus only the detail that earns its place. Kept (and tightened) the load-bearing content: Raises sections, deprecation and version-availability notes, spec/RFC/issue references, why-comments for non-obvious decisions, and all coverage pragmas. The generated mcp_types.v* wire modules are untouched.
358 lines
15 KiB
Python
358 lines
15 KiB
Python
"""Transport-parametrized connection factories for the interaction suite.
|
|
|
|
The `connect` fixture (conftest.py) hands tests one of these factories as a drop-in for
|
|
`Client(server, ...)`, so one test body runs over every transport. The HTTP factories drive the
|
|
server's real Starlette app through the in-process streaming bridge -- no sockets, threads, or
|
|
subprocesses.
|
|
"""
|
|
|
|
from collections.abc import AsyncIterator, Awaitable, Callable, Iterable
|
|
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
|
from functools import partial
|
|
from typing import Any, Protocol
|
|
|
|
import httpx
|
|
from httpx_sse import ServerSentEvent, aconnect_sse
|
|
from mcp_types import (
|
|
ClientCapabilities,
|
|
Implementation,
|
|
InitializeRequestParams,
|
|
JSONRPCMessage,
|
|
JSONRPCRequest,
|
|
JSONRPCResponse,
|
|
jsonrpc_message_adapter,
|
|
)
|
|
from mcp_types.version import LATEST_HANDSHAKE_VERSION, MODERN_PROTOCOL_VERSIONS
|
|
from starlette.applications import Starlette
|
|
from starlette.requests import Request
|
|
from starlette.responses import Response
|
|
from starlette.routing import Mount, Route
|
|
|
|
from mcp.client.client import Client
|
|
from mcp.client.session import ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
|
|
from mcp.client.sse import sse_client
|
|
from mcp.client.streamable_http import streamable_http_client
|
|
from mcp.server import Server
|
|
from mcp.server.auth.provider import OAuthAuthorizationServerProvider, TokenVerifier
|
|
from mcp.server.auth.settings import AuthSettings
|
|
from mcp.server.mcpserver import MCPServer
|
|
from mcp.server.sse import SseServerTransport
|
|
from mcp.server.streamable_http import EventStore
|
|
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
|
|
from mcp.server.transport_security import TransportSecuritySettings
|
|
from tests.interaction.transports._bridge import StreamingASGITransport
|
|
|
|
# The in-process app is mounted at this origin purely so URLs are well-formed; nothing listens here.
|
|
BASE_URL = "http://127.0.0.1:8000"
|
|
|
|
# DNS rebinding cannot reach an in-process ASGI app, so the factories disable the Host/Origin checks;
|
|
# tests of the protection itself pass explicit settings (or None for the localhost auto-enable).
|
|
NO_DNS_REBINDING_PROTECTION = TransportSecuritySettings(enable_dns_rebinding_protection=False)
|
|
|
|
|
|
class Connect(Protocol):
|
|
"""Connect a Client over the fixture-selected transport; accepts `Client` kwargs, yields the client."""
|
|
|
|
def __call__(
|
|
self,
|
|
server: Server | MCPServer,
|
|
*,
|
|
read_timeout_seconds: float | None = None,
|
|
sampling_callback: SamplingFnT | None = None,
|
|
list_roots_callback: ListRootsFnT | None = None,
|
|
logging_callback: LoggingFnT | None = None,
|
|
message_handler: MessageHandlerFnT | None = None,
|
|
client_info: Implementation | None = None,
|
|
elicitation_callback: ElicitationFnT | None = None,
|
|
spec_version: str = LATEST_HANDSHAKE_VERSION,
|
|
) -> AbstractAsyncContextManager[Client]: ...
|
|
|
|
|
|
@asynccontextmanager
|
|
async def connect_in_memory(
|
|
server: Server | MCPServer,
|
|
*,
|
|
read_timeout_seconds: float | None = None,
|
|
sampling_callback: SamplingFnT | None = None,
|
|
list_roots_callback: ListRootsFnT | None = None,
|
|
logging_callback: LoggingFnT | None = None,
|
|
message_handler: MessageHandlerFnT | None = None,
|
|
client_info: Implementation | None = None,
|
|
elicitation_callback: ElicitationFnT | None = None,
|
|
spec_version: str = LATEST_HANDSHAKE_VERSION,
|
|
) -> AsyncIterator[Client]:
|
|
"""Yield a Client connected to the server over the in-memory transport.
|
|
|
|
Modern (2026-07-28+) `spec_version`s open with `mode=<version>`: the DirectDispatcher
|
|
peer-pair (per-request `serve_one`, no initialize handshake) instead of the legacy streams.
|
|
"""
|
|
async with Client(
|
|
server,
|
|
mode=spec_version if spec_version in MODERN_PROTOCOL_VERSIONS else "legacy",
|
|
read_timeout_seconds=read_timeout_seconds,
|
|
sampling_callback=sampling_callback,
|
|
list_roots_callback=list_roots_callback,
|
|
logging_callback=logging_callback,
|
|
message_handler=message_handler,
|
|
client_info=client_info,
|
|
elicitation_callback=elicitation_callback,
|
|
) as client:
|
|
yield client
|
|
|
|
|
|
@asynccontextmanager
|
|
async def connect_over_streamable_http(
|
|
server: Server | MCPServer,
|
|
*,
|
|
stateless_http: bool = False,
|
|
json_response: bool = False,
|
|
event_store: EventStore | None = None,
|
|
retry_interval: int | None = None,
|
|
read_timeout_seconds: float | None = None,
|
|
sampling_callback: SamplingFnT | None = None,
|
|
list_roots_callback: ListRootsFnT | None = None,
|
|
logging_callback: LoggingFnT | None = None,
|
|
message_handler: MessageHandlerFnT | None = None,
|
|
client_info: Implementation | None = None,
|
|
elicitation_callback: ElicitationFnT | None = None,
|
|
spec_version: str = LATEST_HANDSHAKE_VERSION,
|
|
) -> AsyncIterator[Client]:
|
|
"""Yield a Client connected to the server's streamable HTTP app, entirely in process.
|
|
|
|
Resumability tests pass an `event_store` with `retry_interval=0` so the client's reconnection
|
|
wait is a no-op. Modern (2026-07-28+) `spec_version`s open with `mode=<version>`, adopting a
|
|
synthesized DiscoverResult instead of running the legacy initialize handshake.
|
|
"""
|
|
app = server.streamable_http_app(
|
|
stateless_http=stateless_http,
|
|
json_response=json_response,
|
|
event_store=event_store,
|
|
retry_interval=retry_interval,
|
|
transport_security=NO_DNS_REBINDING_PROTECTION,
|
|
)
|
|
async with (
|
|
server.session_manager.run(),
|
|
httpx.AsyncClient(transport=StreamingASGITransport(app), base_url=BASE_URL) as http_client,
|
|
Client(
|
|
streamable_http_client(f"{BASE_URL}/mcp", http_client=http_client),
|
|
mode=spec_version if spec_version in MODERN_PROTOCOL_VERSIONS else "legacy",
|
|
read_timeout_seconds=read_timeout_seconds,
|
|
sampling_callback=sampling_callback,
|
|
list_roots_callback=list_roots_callback,
|
|
logging_callback=logging_callback,
|
|
message_handler=message_handler,
|
|
client_info=client_info,
|
|
elicitation_callback=elicitation_callback,
|
|
) as client,
|
|
):
|
|
yield client
|
|
|
|
|
|
connect_over_streamable_http_stateless: Connect = partial(connect_over_streamable_http, stateless_http=True)
|
|
"""The streamable-http matrix arm with the server stateless: fresh transport per request, no
|
|
session id, no standalone GET stream. The one shared Server instance backs every request."""
|
|
|
|
|
|
@asynccontextmanager
|
|
async def mounted_app(
|
|
server: Server | MCPServer,
|
|
*,
|
|
stateless_http: bool = False,
|
|
json_response: bool = False,
|
|
event_store: EventStore | None = None,
|
|
retry_interval: int | None = None,
|
|
transport_security: TransportSecuritySettings | None = NO_DNS_REBINDING_PROTECTION,
|
|
on_request: Callable[[httpx.Request], Awaitable[None]] | None = None,
|
|
on_response: Callable[[httpx.Response], Awaitable[None]] | None = None,
|
|
headers: dict[str, str] | None = None,
|
|
auth: AuthSettings | None = None,
|
|
token_verifier: TokenVerifier | None = None,
|
|
auth_server_provider: OAuthAuthorizationServerProvider[Any, Any, Any] | None = None,
|
|
) -> AsyncIterator[tuple[httpx.AsyncClient, StreamableHTTPSessionManager]]:
|
|
"""Mount the server's streamable HTTP app on the in-process bridge and yield an httpx client.
|
|
|
|
Tests speak raw HTTP through the yielded client (status codes, headers, SSE bytes) or wrap it
|
|
in `client_via_http(http)` so several `Client`s share the one mounted session manager.
|
|
`on_response` fires as each response's headers arrive -- SSE bodies are not yet read then.
|
|
"""
|
|
lowlevel = server._lowlevel_server if isinstance(server, MCPServer) else server
|
|
app = lowlevel.streamable_http_app(
|
|
stateless_http=stateless_http,
|
|
json_response=json_response,
|
|
event_store=event_store,
|
|
retry_interval=retry_interval,
|
|
transport_security=transport_security,
|
|
auth=auth,
|
|
token_verifier=token_verifier,
|
|
auth_server_provider=auth_server_provider,
|
|
)
|
|
event_hooks: dict[str, list[Callable[..., Awaitable[None]]]] = {}
|
|
if on_request is not None:
|
|
event_hooks["request"] = [on_request]
|
|
if on_response is not None:
|
|
event_hooks["response"] = [on_response]
|
|
async with (
|
|
server.session_manager.run(),
|
|
httpx.AsyncClient(
|
|
transport=StreamingASGITransport(app), base_url=BASE_URL, event_hooks=event_hooks, headers=headers
|
|
) as http_client,
|
|
):
|
|
yield http_client, server.session_manager
|
|
|
|
|
|
@asynccontextmanager
|
|
async def client_via_http(
|
|
http_client: httpx.AsyncClient,
|
|
*,
|
|
logging_callback: LoggingFnT | None = None,
|
|
message_handler: MessageHandlerFnT | None = None,
|
|
elicitation_callback: ElicitationFnT | None = None,
|
|
) -> AsyncIterator[Client]:
|
|
"""Connect a `Client` over an already-mounted streamable HTTP app (see `mounted_app`).
|
|
|
|
The underlying `httpx.AsyncClient` is left open when the `Client` exits, so several `Client`s
|
|
can sit alongside raw-httpx assertions in the same test.
|
|
"""
|
|
transport = streamable_http_client(f"{BASE_URL}/mcp", http_client=http_client)
|
|
async with Client(
|
|
transport,
|
|
# Callers assert the legacy HTTP wire shape (session-id header, standalone GET stream,
|
|
# closing DELETE); the modern flow is sessionless and would silently change the subject.
|
|
mode="legacy",
|
|
logging_callback=logging_callback,
|
|
message_handler=message_handler,
|
|
elicitation_callback=elicitation_callback,
|
|
) as client:
|
|
yield client
|
|
|
|
|
|
def parse_sse_messages(events: Iterable[ServerSentEvent]) -> list[JSONRPCMessage]:
|
|
"""Decode SSE events into JSON-RPC messages, skipping priming events that carry no data."""
|
|
return [jsonrpc_message_adapter.validate_json(event.data) for event in events if event.data]
|
|
|
|
|
|
async def post_jsonrpc(
|
|
http: httpx.AsyncClient, body: dict[str, object], *, session_id: str | None = None
|
|
) -> tuple[httpx.Response, list[JSONRPCMessage]]:
|
|
"""POST a JSON-RPC body and read its SSE response stream to completion.
|
|
|
|
Returns the response plus the parsed messages from its SSE stream. Only meaningful when the
|
|
server answers with `text/event-stream`; for errors or 202 acks use `httpx.AsyncClient.post`.
|
|
"""
|
|
async with aconnect_sse(http, "POST", "/mcp", json=body, headers=base_headers(session_id=session_id)) as source:
|
|
events = [event async for event in source.aiter_sse()]
|
|
return source.response, parse_sse_messages(events)
|
|
|
|
|
|
def base_headers(*, session_id: str | None = None) -> dict[str, str]:
|
|
"""Standard headers for raw-httpx streamable-HTTP requests; rejection tests vary only the header under test."""
|
|
headers = {
|
|
"accept": "application/json, text/event-stream",
|
|
"content-type": "application/json",
|
|
"mcp-protocol-version": LATEST_HANDSHAKE_VERSION,
|
|
}
|
|
if session_id is not None:
|
|
headers["mcp-session-id"] = session_id
|
|
return headers
|
|
|
|
|
|
def initialize_body(request_id: int = 1) -> dict[str, object]:
|
|
"""A wire-level initialize JSON-RPC request body, exactly as an SDK client would send it."""
|
|
params = InitializeRequestParams(
|
|
protocol_version=LATEST_HANDSHAKE_VERSION,
|
|
capabilities=ClientCapabilities(),
|
|
client_info=Implementation(name="raw", version="0.0.0"),
|
|
)
|
|
return JSONRPCRequest(
|
|
jsonrpc="2.0", id=request_id, method="initialize", params=params.model_dump(by_alias=True, exclude_none=True)
|
|
).model_dump(by_alias=True, exclude_none=True)
|
|
|
|
|
|
async def initialize_via_http(http: httpx.AsyncClient) -> str:
|
|
"""Initialize over raw httpx, send `notifications/initialized`, and return the session ID."""
|
|
async with aconnect_sse(http, "POST", "/mcp", json=initialize_body(), headers=base_headers()) as source:
|
|
assert source.response.status_code == 200
|
|
# An event-store-backed server opens the stream with a priming event (empty data); skip it.
|
|
events = [event async for event in source.aiter_sse() if event.data]
|
|
assert len(events) == 1
|
|
assert JSONRPCResponse.model_validate_json(events[0].data).id == 1
|
|
session_id = source.response.headers["mcp-session-id"]
|
|
initialized = await http.post(
|
|
"/mcp",
|
|
json={"jsonrpc": "2.0", "method": "notifications/initialized"},
|
|
headers=base_headers(session_id=session_id),
|
|
)
|
|
assert initialized.status_code == 202
|
|
return session_id
|
|
|
|
|
|
def build_sse_app(server: Server | MCPServer) -> tuple[Starlette, SseServerTransport]:
|
|
"""Mount a server on a Starlette app exposing the legacy SSE transport at /sse and /messages/.
|
|
|
|
Built by hand because `MCPServer.sse_app()` hides the `SseServerTransport` handle SSE tests need.
|
|
"""
|
|
sse = SseServerTransport(
|
|
"/messages/", security_settings=TransportSecuritySettings(enable_dns_rebinding_protection=False)
|
|
)
|
|
lowlevel = server._lowlevel_server if isinstance(server, MCPServer) else server
|
|
|
|
async def handle_sse(request: Request) -> Response:
|
|
async with sse.connect_sse(request.scope, request.receive, request._send) as (read, write):
|
|
await lowlevel.run(read, write, lowlevel.create_initialization_options())
|
|
return Response()
|
|
|
|
app = Starlette(
|
|
routes=[
|
|
Route("/sse", endpoint=handle_sse, methods=["GET"]),
|
|
Mount("/messages/", app=sse.handle_post_message),
|
|
],
|
|
)
|
|
return app, sse
|
|
|
|
|
|
@asynccontextmanager
|
|
async def connect_over_sse(
|
|
server: Server | MCPServer,
|
|
*,
|
|
read_timeout_seconds: float | None = None,
|
|
sampling_callback: SamplingFnT | None = None,
|
|
list_roots_callback: ListRootsFnT | None = None,
|
|
logging_callback: LoggingFnT | None = None,
|
|
message_handler: MessageHandlerFnT | None = None,
|
|
client_info: Implementation | None = None,
|
|
elicitation_callback: ElicitationFnT | None = None,
|
|
spec_version: str = LATEST_HANDSHAKE_VERSION,
|
|
) -> AsyncIterator[Client]:
|
|
"""Yield a Client connected to the server's legacy SSE transport, entirely in process."""
|
|
app, _ = build_sse_app(server)
|
|
|
|
def httpx_client_factory(
|
|
headers: dict[str, str] | None = None,
|
|
timeout: httpx.Timeout | None = None,
|
|
auth: httpx.Auth | None = None,
|
|
) -> httpx.AsyncClient:
|
|
# connect_sse runs the entire MCP session inside the GET request and releases its streams
|
|
# only after that request observes a disconnect, so the bridge must drain, not cancel.
|
|
return httpx.AsyncClient(
|
|
transport=StreamingASGITransport(app, cancel_on_close=False),
|
|
base_url=BASE_URL,
|
|
headers=headers,
|
|
timeout=timeout,
|
|
auth=auth,
|
|
)
|
|
|
|
transport = sse_client(f"{BASE_URL}/sse", httpx_client_factory=httpx_client_factory)
|
|
async with Client(
|
|
transport,
|
|
# SSE is a legacy-only transport; the modern path has no SSE story.
|
|
mode="legacy",
|
|
read_timeout_seconds=read_timeout_seconds,
|
|
sampling_callback=sampling_callback,
|
|
list_roots_callback=list_roots_callback,
|
|
logging_callback=logging_callback,
|
|
message_handler=message_handler,
|
|
client_info=client_info,
|
|
elicitation_callback=elicitation_callback,
|
|
) as client:
|
|
yield client
|