Client-side 2026-07-28 support: .discover()/.adopt() + Client(mode=); request-metadata green (#2950)
This commit is contained in:
parent
ae13ede143
commit
f226d00d0a
1 changed files with 2533 additions and 993 deletions
|
|
@ -2,27 +2,34 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from contextlib import AsyncExitStack
|
||||
from dataclasses import KW_ONLY, dataclass, field
|
||||
from typing import Any
|
||||
from typing import Any, Literal, TypeVar
|
||||
|
||||
import anyio
|
||||
from typing_extensions import deprecated
|
||||
|
||||
from mcp import types
|
||||
from mcp.client._memory import InMemoryTransport
|
||||
from mcp.client._probe import negotiate_auto
|
||||
from mcp.client._transport import Transport
|
||||
from mcp.client.session import ClientSession, ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
|
||||
from mcp.client.streamable_http import streamable_http_client
|
||||
from mcp.server import Server
|
||||
from mcp.server.mcpserver import MCPServer
|
||||
from mcp.shared.dispatcher import ProgressFnT
|
||||
from mcp.server.runner import modern_on_request
|
||||
from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair
|
||||
from mcp.shared.dispatcher import Dispatcher, ProgressFnT
|
||||
from mcp.shared.exceptions import MCPDeprecationWarning
|
||||
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
|
||||
from mcp.shared.version import HANDSHAKE_PROTOCOL_VERSIONS, MODERN_PROTOCOL_VERSIONS
|
||||
from mcp.types import (
|
||||
CallToolResult,
|
||||
CompleteResult,
|
||||
EmptyResult,
|
||||
GetPromptResult,
|
||||
Implementation,
|
||||
InitializeResult,
|
||||
ListPromptsResult,
|
||||
ListResourcesResult,
|
||||
ListResourceTemplatesResult,
|
||||
|
|
@ -33,8 +40,87 @@ from mcp.types import (
|
|||
ReadResourceResult,
|
||||
RequestParamsMeta,
|
||||
ResourceTemplateReference,
|
||||
ServerCapabilities,
|
||||
)
|
||||
|
||||
ConnectMode = Literal["legacy", "auto"] | str
|
||||
"""``mode=`` value: ``"legacy"`` (initialize handshake), ``"auto"`` (discover, fall back to
|
||||
initialize), or a modern protocol-version string (adopt directly). The ``str`` arm is for
|
||||
forward-compat; ``Client.__post_init__`` rejects anything outside that set at construction."""
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
_Connector = Callable[[AsyncExitStack, ConnectMode, bool], Awaitable["Dispatcher[Any]"]]
|
||||
"""Resolved at ``__post_init__`` from the shape of ``server`` alone: enter whatever resources
|
||||
are needed onto the exit stack and hand back the ``Dispatcher`` ``ClientSession`` will drive.
|
||||
``mode`` and ``raise_exceptions`` are passed at call time so they're read at the same moment
|
||||
``__aenter__`` reads them for the handshake step."""
|
||||
|
||||
|
||||
def _connect_transport(transport: Transport) -> _Connector:
|
||||
"""Connector for the stream-backed paths (URL, user-supplied ``Transport``)."""
|
||||
|
||||
async def connect(exit_stack: AsyncExitStack, _mode: ConnectMode, _raise_exceptions: bool) -> Dispatcher[Any]:
|
||||
read_stream, write_stream = await exit_stack.enter_async_context(transport)
|
||||
return JSONRPCDispatcher(read_stream, write_stream)
|
||||
|
||||
return connect
|
||||
|
||||
|
||||
def _connect_inproc(server: Server[Any]) -> _Connector:
|
||||
"""Connector for an in-process ``Server``: legacy mode drives the stream loop via
|
||||
``InMemoryTransport``; any other mode drives the modern per-request path through a
|
||||
``DirectDispatcher`` peer pair (no streams, no JSON-RPC framing, no initialize handshake)."""
|
||||
|
||||
async def connect(exit_stack: AsyncExitStack, mode: ConnectMode, raise_exceptions: bool) -> Dispatcher[Any]:
|
||||
if mode == "legacy":
|
||||
transport = InMemoryTransport(server, raise_exceptions=raise_exceptions)
|
||||
read_stream, write_stream = await exit_stack.enter_async_context(transport)
|
||||
return JSONRPCDispatcher(read_stream, write_stream)
|
||||
lifespan_state = await exit_stack.enter_async_context(server.lifespan(server))
|
||||
client_disp, server_disp = create_direct_dispatcher_pair(raise_handler_exceptions=raise_exceptions)
|
||||
tg = await exit_stack.enter_async_context(anyio.create_task_group())
|
||||
exit_stack.callback(server_disp.close)
|
||||
on_request = modern_on_request(server, lifespan_state)
|
||||
await tg.start(server_disp.run, on_request, _no_inbound_client_notifications)
|
||||
return client_disp
|
||||
|
||||
return connect
|
||||
|
||||
|
||||
def _connected(value: _T | None) -> _T:
|
||||
"""Narrow a post-handshake session attribute from ``T | None`` to ``T``.
|
||||
|
||||
``Client.__aenter__`` only assigns ``_session`` after the handshake succeeds, so inside
|
||||
``async with Client(...)`` these attributes are always populated; the ``.session`` gate
|
||||
raises before this is reached otherwise. The guard exists for pyright, not runtime.
|
||||
"""
|
||||
if value is None: # pragma: no cover
|
||||
raise RuntimeError("Client must be used within an async context manager")
|
||||
return value
|
||||
|
||||
|
||||
def _synthesize_discover(protocol_version: str) -> types.DiscoverResult:
|
||||
return types.DiscoverResult(
|
||||
supported_versions=[protocol_version],
|
||||
capabilities=types.ServerCapabilities(),
|
||||
server_info=types.Implementation(name="", version=""),
|
||||
result_type="complete",
|
||||
ttl_ms=0,
|
||||
cache_scope="public",
|
||||
)
|
||||
|
||||
|
||||
async def _no_inbound_client_notifications(_dctx: Any, _method: str, _params: Mapping[str, Any] | None) -> None:
|
||||
"""Server-side inbound ``OnNotify`` for the modern in-process path — receives nothing.
|
||||
|
||||
At 2026-07-28 the spec defines no client→server notifications: ``initialized`` and
|
||||
``roots/list_changed`` are removed, and cancellation is structural (anyio scope cancel
|
||||
through the direct await, not a notify). Server→client notifications (progress, log
|
||||
messages) flow the other way via the per-request ``DispatchContext`` into the client's
|
||||
callbacks, and are not seen here.
|
||||
"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class Client:
|
||||
|
|
@ -65,7 +151,7 @@ class Client:
|
|||
server: Server[Any] | MCPServer | Transport | str
|
||||
"""The MCP server to connect to.
|
||||
|
||||
If the server is a `Server` or `MCPServer` instance, it will be wrapped in an `InMemoryTransport`.
|
||||
If the server is a `Server` or `MCPServer` instance, it will be connected in-process.
|
||||
If the server is a URL string, it will be used as the URL for a `streamable_http_client` transport.
|
||||
If the server is a `Transport` instance, it will be used directly.
|
||||
"""
|
||||
|
|
@ -95,58 +181,83 @@ class Client:
|
|||
client_info: Implementation | None = None
|
||||
"""Client implementation info to send to server."""
|
||||
|
||||
protocol_version: str | None = None
|
||||
"""Pin the protocol version instead of negotiating it.
|
||||
mode: ConnectMode = "auto"
|
||||
"""How to negotiate the protocol version.
|
||||
|
||||
Pinning to ``2026-07-28`` or later selects the stateless transport era: no initialize
|
||||
handshake is sent on the wire (the session synthesizes its `InitializeResult` locally),
|
||||
and for HTTP the ``MCP-Protocol-Version`` header is set from the first request. A modern
|
||||
pin currently requires a URL or `Transport`; the in-memory `Server`/`MCPServer` path
|
||||
does not yet have a modern entry point.
|
||||
Leave as ``None`` to negotiate the version via the initialize handshake.
|
||||
"""
|
||||
'auto' (the default) probes `server/discover` and falls back to the initialize handshake on legacy servers;
|
||||
for an in-process `Server`/`MCPServer` it dispatches directly without JSON-RPC framing. 'legacy' forces the
|
||||
initialize handshake (byte-identical pre-2026 behavior). A modern protocol-version string (e.g. '2026-07-28')
|
||||
adopts that version directly without a probe — supply `prior_discover` to reuse a known DiscoverResult, or
|
||||
omit it to synthesize a minimal one."""
|
||||
|
||||
prior_discover: types.DiscoverResult | None = None
|
||||
"""A previously-obtained DiscoverResult to install via .adopt() when mode is a version pin.
|
||||
Ignored when mode='legacy'."""
|
||||
|
||||
elicitation_callback: ElicitationFnT | None = None
|
||||
"""Callback for handling elicitation requests."""
|
||||
|
||||
_entered: bool = field(init=False, default=False)
|
||||
_session: ClientSession | None = field(init=False, default=None)
|
||||
_exit_stack: AsyncExitStack | None = field(init=False, default=None)
|
||||
_transport: Transport = field(init=False)
|
||||
_connect: _Connector = field(init=False, repr=False, compare=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if isinstance(self.server, Server | MCPServer):
|
||||
self._transport = InMemoryTransport(self.server, raise_exceptions=self.raise_exceptions)
|
||||
elif isinstance(self.server, str):
|
||||
self._transport = streamable_http_client(self.server, protocol_version=self.protocol_version)
|
||||
if self.mode not in ("legacy", "auto") and self.mode not in MODERN_PROTOCOL_VERSIONS:
|
||||
hint = (
|
||||
f" ({self.mode!r} is a handshake-era version — use mode='legacy')"
|
||||
if self.mode in HANDSHAKE_PROTOCOL_VERSIONS
|
||||
else ""
|
||||
)
|
||||
raise ValueError(
|
||||
f"mode must be 'legacy', 'auto', or one of {list(MODERN_PROTOCOL_VERSIONS)}; got {self.mode!r}{hint}"
|
||||
)
|
||||
|
||||
srv = self.server
|
||||
if isinstance(srv, MCPServer):
|
||||
srv = srv._lowlevel_server # pyright: ignore[reportPrivateUsage]
|
||||
if isinstance(srv, Server):
|
||||
self._connect = _connect_inproc(srv)
|
||||
elif isinstance(srv, str):
|
||||
self._connect = _connect_transport(streamable_http_client(srv))
|
||||
else:
|
||||
self._transport = self.server
|
||||
self._connect = _connect_transport(srv)
|
||||
|
||||
async def _build_session(self, exit_stack: AsyncExitStack) -> ClientSession:
|
||||
"""Enter the resolved connector and return an un-entered ClientSession."""
|
||||
dispatcher = await self._connect(exit_stack, self.mode, self.raise_exceptions)
|
||||
return ClientSession(
|
||||
dispatcher=dispatcher,
|
||||
read_timeout_seconds=self.read_timeout_seconds,
|
||||
sampling_callback=self.sampling_callback,
|
||||
list_roots_callback=self.list_roots_callback,
|
||||
logging_callback=self.logging_callback,
|
||||
message_handler=self.message_handler,
|
||||
client_info=self.client_info,
|
||||
elicitation_callback=self.elicitation_callback,
|
||||
)
|
||||
|
||||
async def __aenter__(self) -> Client:
|
||||
"""Enter the async context manager."""
|
||||
if self._session is not None:
|
||||
if self._entered:
|
||||
raise RuntimeError("Client is already entered; cannot reenter")
|
||||
self._entered = True
|
||||
|
||||
async with AsyncExitStack() as exit_stack:
|
||||
read_stream, write_stream = await exit_stack.enter_async_context(self._transport)
|
||||
session = await self._build_session(exit_stack)
|
||||
session = await exit_stack.enter_async_context(session)
|
||||
|
||||
self._session = await exit_stack.enter_async_context(
|
||||
ClientSession(
|
||||
read_stream=read_stream,
|
||||
write_stream=write_stream,
|
||||
read_timeout_seconds=self.read_timeout_seconds,
|
||||
sampling_callback=self.sampling_callback,
|
||||
list_roots_callback=self.list_roots_callback,
|
||||
logging_callback=self.logging_callback,
|
||||
message_handler=self.message_handler,
|
||||
client_info=self.client_info,
|
||||
elicitation_callback=self.elicitation_callback,
|
||||
protocol_version=self.protocol_version,
|
||||
)
|
||||
)
|
||||
if self.mode == "legacy":
|
||||
await session.initialize()
|
||||
elif self.mode == "auto":
|
||||
await negotiate_auto(session)
|
||||
else:
|
||||
session.adopt(self.prior_discover or _synthesize_discover(self.mode))
|
||||
|
||||
await self._session.initialize()
|
||||
|
||||
# Transfer ownership to self for __aexit__ to handle
|
||||
# Only publish the session after the handshake succeeds, so `_session is not None`
|
||||
# implies the protocol_version/server_info/server_capabilities are populated. If the
|
||||
# handshake raised above, the local exit_stack unwinds the transport for us.
|
||||
self._session = session
|
||||
self._exit_stack = exit_stack.pop_all()
|
||||
return self
|
||||
|
||||
|
|
@ -169,22 +280,42 @@ class Client:
|
|||
raise RuntimeError("Client must be used within an async context manager")
|
||||
return self._session
|
||||
|
||||
# TODO(maxisbey): the by-construction shape is for __aenter__ to return a connected-view
|
||||
# type whose protocol_version/server_info/server_capabilities are non-Optional fields,
|
||||
# eliminating these guards (and the one in .session). Same family as resolving the
|
||||
# transport/connector at __post_init__ so the Optional internal fields disappear.
|
||||
@property
|
||||
def initialize_result(self) -> InitializeResult:
|
||||
"""The server's InitializeResult.
|
||||
def protocol_version(self) -> str:
|
||||
"""Negotiated protocol version (set by initialize/discover/adopt during ``__aenter__``)."""
|
||||
return _connected(self.session.protocol_version)
|
||||
|
||||
Contains server_info, capabilities, instructions, and the negotiated protocol_version.
|
||||
Raises RuntimeError if accessed outside the context manager.
|
||||
"""
|
||||
result = self.session.initialize_result
|
||||
if result is None: # pragma: no cover
|
||||
raise RuntimeError("Client must be used within an async context manager")
|
||||
return result
|
||||
@property
|
||||
def server_info(self) -> Implementation:
|
||||
"""Server name/version (set by initialize/discover/adopt during ``__aenter__``)."""
|
||||
return _connected(self.session.server_info)
|
||||
|
||||
@property
|
||||
def server_capabilities(self) -> ServerCapabilities:
|
||||
"""Server capabilities (set by initialize/discover/adopt during ``__aenter__``)."""
|
||||
return _connected(self.session.server_capabilities)
|
||||
|
||||
@property
|
||||
def instructions(self) -> str | None:
|
||||
"""Server-provided instructions text, if any."""
|
||||
return self.session.instructions
|
||||
|
||||
@deprecated(
|
||||
"ping is removed as of 2026-07-28; the method only works under mode='legacy'.",
|
||||
category=MCPDeprecationWarning,
|
||||
)
|
||||
async def send_ping(self, *, meta: RequestParamsMeta | None = None) -> EmptyResult:
|
||||
"""Send a ping request to the server."""
|
||||
return await self.session.send_ping(meta=meta)
|
||||
|
||||
@deprecated(
|
||||
"Client-to-server progress is deprecated as of 2026-07-28; progress is server-to-client only.",
|
||||
category=MCPDeprecationWarning,
|
||||
)
|
||||
async def send_progress_notification(
|
||||
self,
|
||||
progress_token: str | int,
|
||||
|
|
@ -193,7 +324,7 @@ class Client:
|
|||
message: str | None = None,
|
||||
) -> None:
|
||||
"""Send a progress notification to the server."""
|
||||
await self.session.send_progress_notification(
|
||||
await self.session.send_progress_notification( # pyright: ignore[reportDeprecated]
|
||||
progress_token=progress_token,
|
||||
progress=progress,
|
||||
total=total,
|
||||
|
|
|
|||
Loading…
Reference in a new issue