_connect.py

"""Transport-parametrized connection factories for the interaction suite.

The `connect` fixture (see conftest.py) hands tests one of these factories so the same test body
runs over each transport without naming any of them: the factory is a drop-in replacement for
constructing `Client(server, ...)` and yields the connected client. The HTTP factories drive the
server's real Starlette app through the in-process streaming bridge, so the full transport layer
(session ids, SSE encoding, session management) runs with no sockets, threads, or subprocesses.
"""

from collections.abc import AsyncIterator, Awaitable, Callable, Iterable
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from typing import Any, Protocol

import httpx
from httpx_sse import ServerSentEvent, aconnect_sse
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 mcp.types import (
    LATEST_PROTOCOL_VERSION,
    ClientCapabilities,
    Implementation,
    InitializeRequestParams,
    JSONRPCMessage,
    JSONRPCRequest,
    JSONRPCResponse,
    SamplingCapability,
    jsonrpc_message_adapter,
)
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 protection validates Host/Origin headers against a real network attack that cannot
# exist for an in-process ASGI app, so the in-process factories disable it; tests that exercise the
# protection itself pass explicit settings (or transport_security=None to get the localhost
# auto-enable behaviour).
NO_DNS_REBINDING_PROTECTION = TransportSecuritySettings(enable_dns_rebinding_protection=False)


class Connect(Protocol):
    """Connect a Client to a server over the transport selected by the `connect` fixture.

    Accepts the same keyword arguments as `Client` and yields the connected client.
    """

    def __call__(
        self,
        server: Server | MCPServer,
        *,
        read_timeout_seconds: float | None = None,
        sampling_callback: SamplingFnT | None = None,
        sampling_capabilities: SamplingCapability | 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,
    ) -> AbstractAsyncContextManager[Client]: ...


@asynccontextmanager
async def connect_in_memory(
    server: Server | MCPServer,
    *,
    read_timeout_seconds: float | None = None,
    sampling_callback: SamplingFnT | None = None,
    sampling_capabilities: SamplingCapability | 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,
) -> AsyncIterator[Client]:
    """Yield a Client connected to the server over the in-memory transport."""
    async with Client(
        server,
        read_timeout_seconds=read_timeout_seconds,
        sampling_callback=sampling_callback,
        sampling_capabilities=sampling_capabilities,
        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,
    sampling_capabilities: SamplingCapability | 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,
) -> AsyncIterator[Client]:
    """Yield a Client connected to the server's streamable HTTP app, entirely in process.

    With the defaults this is the matrix leg (stateful sessions, SSE responses); the
    transport-specific tests pass `stateless_http` or `json_response` to select the other
    server modes, and the resumability tests pass an `event_store` (with `retry_interval=0` so
    the client's reconnection wait is a no-op).
    """
    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),
            read_timeout_seconds=read_timeout_seconds,
            sampling_callback=sampling_callback,
            sampling_capabilities=sampling_capabilities,
            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 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,
    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.

    Yields the httpx client (rooted at the in-process origin) and the live session manager. Tests
    use this in two ways: for raw-httpx assertions (status codes, headers, SSE bytes) the test
    speaks HTTP through the yielded client directly; for client-driven assertions the test wraps
    that client in `client_via_http(http)`, which lets several `Client`s share the one mounted
    session manager. `on_request` records every outgoing HTTP request before it leaves the
    yielded client.

    DNS-rebinding protection is disabled by default; pass explicit settings (or `None` for the
    localhost auto-enable behaviour) to test the protection itself.
    """
    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 = {"request": [on_request]} if on_request is not None else None
    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.

    Use with `mounted_app(...)` so several `Client`s share the one session manager, or so a
    client-driven assertion can sit alongside raw-httpx assertions in the same test. The
    underlying `httpx.AsyncClient` is left open when the `Client` exits.
    """
    transport = streamable_http_client(f"{BASE_URL}/mcp", http_client=http_client)
    async with Client(
        transport,
        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 HTTP response (for header/status assertions) and the parsed JSON-RPC messages
    that arrived on the response's SSE stream. Only meaningful for requests the server answers
    with `text/event-stream`; for error responses or 202 notification acknowledgements, use
    `httpx.AsyncClient.post` directly and assert on the response.
    """
    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 request headers for raw-httpx streamable-HTTP tests.

    Every well-formed request carries these (Accept covering both response representations,
    Content-Type for POST bodies, MCP-Protocol-Version at the latest revision, and the session
    ID once one exists), so a test that wants to assert a specific rejection only varies the one
    header under test.
    """
    headers = {
        "accept": "application/json, text/event-stream",
        "content-type": "application/json",
        "mcp-protocol-version": LATEST_PROTOCOL_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_PROTOCOL_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:
    """Perform the initialize handshake over a raw `httpx.AsyncClient` and return the session ID.

    Validates the SSE response and sends the `notifications/initialized` follow-up, so the server
    is fully ready for subsequent feature requests when this returns.
    """
    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/.

    `MCPServer.sse_app()` exists but does not expose the underlying `SseServerTransport`, which
    the SSE-specific tests need; building the app explicitly here gives both server flavours the
    same routing while keeping that handle.
    """
    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,
    sampling_capabilities: SamplingCapability | 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,
) -> 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:
        # The SSE server transport's connect_sse runs the entire MCP session inside the GET
        # request and only releases its streams after that request observes a disconnect, so the
        # bridge must let the application drain rather than cancelling at close.
        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,
        read_timeout_seconds=read_timeout_seconds,
        sampling_callback=sampling_callback,
        sampling_capabilities=sampling_capabilities,
        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