modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[v2] ClientSession runs on JSONRPCDispatcher; BaseSession removed (#2838)

Browse source
This commit is contained in:
Max 2026-06-15 14:46:34 +01:00 • committed by GitHub
commit 1012d60004
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 2886 additions and 1734 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/mcp/client/client.py
View file

@ -12,7 +12,7 @@ from mcp.client.session import ClientSession, ElicitationFnT, ListRootsFnT, Logg
from mcp.client.streamable_http import streamable_http_client
from mcp.server import Server
from mcp.server.mcpserver import MCPServer
from mcp.shared.session import ProgressFnT
from mcp.shared.dispatcher import ProgressFnT
from mcp.types import (
CallToolResult,
CompleteResult,

15
src/mcp/client/context.py
View file

@ -1,16 +1,5 @@
"""Request context for MCP client handlers."""
from mcp.client.session import ClientSession
from mcp.shared._context import RequestContext
from mcp.client.session import ClientRequestContext
ClientRequestContext = RequestContext[ClientSession]
"""Context for handling incoming requests in a client session.
This context is passed to client-side callbacks (sampling, elicitation, list_roots) when the server sends requests
to the client.
Attributes:
request_id: The unique identifier for this request.
meta: Optional metadata associated with the request.
session: The client session handling this request.
"""
__all__ = ["ClientRequestContext"]

277
src/mcp/client/session.py
View file

@ -1,28 +1,49 @@
from __future__ import annotations
import logging
from collections.abc import Mapping
from dataclasses import dataclass
from types import TracebackType
from typing import Any, Protocol, cast, get_args
import anyio
import anyio.abc
import anyio.lowlevel
from pydantic import BaseModel, TypeAdapter
from pydantic import BaseModel, TypeAdapter, ValidationError
from typing_extensions import Self, TypeVar
from mcp import types
from mcp.client._transport import ReadStream, WriteStream
from mcp.shared._context import RequestContext
from mcp.shared.message import SessionMessage
from mcp.shared.session import BaseSession, ProgressFnT, RequestResponder
from mcp.shared._compat import resync_tracer
from mcp.shared.dispatcher import CallOptions, DispatchContext, Dispatcher, ProgressFnT
from mcp.shared.exceptions import MCPError
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
from mcp.shared.message import ClientMessageMetadata, SessionMessage
from mcp.shared.session import RequestResponder
from mcp.shared.transport_context import TransportContext
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
from mcp.types._types import RequestParamsMeta
from mcp.types import RequestId, RequestParamsMeta
DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0")
logger = logging.getLogger("client")
ReceiveResultT = TypeVar("ReceiveResultT", bound=BaseModel)
@dataclass(kw_only=True)
class ClientRequestContext:
"""Context for a server-initiated request, passed to the sampling/elicitation/list-roots callbacks."""
session: ClientSession
request_id: RequestId
meta: RequestParamsMeta | None = None
class SamplingFnT(Protocol):
async def __call__(
self,
context: RequestContext[ClientSession],
context: ClientRequestContext,
params: types.CreateMessageRequestParams,
) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData: ... # pragma: no branch
@ -30,14 +51,14 @@ class SamplingFnT(Protocol):
class ElicitationFnT(Protocol):
async def __call__(
self,
context: RequestContext[ClientSession],
context: ClientRequestContext,
params: types.ElicitRequestParams,
) -> types.ElicitResult | types.ErrorData: ... # pragma: no branch
class ListRootsFnT(Protocol):
async def __call__(
self, context: RequestContext[ClientSession]
self, context: ClientRequestContext
) -> types.ListRootsResult | types.ErrorData: ... # pragma: no branch
@ -59,7 +80,7 @@ async def _default_message_handler(
async def _default_sampling_callback(
context: RequestContext[ClientSession],
context: ClientRequestContext,
params: types.CreateMessageRequestParams,
) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData:
return types.ErrorData(
@ -69,7 +90,7 @@ async def _default_sampling_callback(
async def _default_elicitation_callback(
context: RequestContext[ClientSession],
context: ClientRequestContext,
params: types.ElicitRequestParams,
) -> types.ElicitResult | types.ErrorData:
return types.ErrorData(
@ -79,7 +100,7 @@ async def _default_elicitation_callback(
async def _default_list_roots_callback(
context: RequestContext[ClientSession],
context: ClientRequestContext,
) -> types.ListRootsResult | types.ErrorData:
return types.ErrorData(
code=types.INVALID_REQUEST,
@ -104,19 +125,21 @@ spec methods this SDK deliberately doesn't model, like `tasks/*` — are
answered with METHOD_NOT_FOUND instead of failing union validation."""
class ClientSession(
BaseSession[
types.ClientRequest,
types.ClientNotification,
types.ClientResult,
types.ServerRequest,
types.ServerNotification,
]
):
class ClientSession:
"""Client half of an MCP connection, running on a `Dispatcher`.
Construct it over a transport's stream pair (or pass a pre-built
`dispatcher=`), enter as an async context manager, then call
`initialize()`. The dispatcher owns the receive loop and request
correlation; this class owns the typed MCP layer and the constructor
callbacks. Transport `Exception` items reach `message_handler` only when
the session builds its own dispatcher from a stream pair.
"""
def __init__(
self,
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[SessionMessage],
read_stream: ReadStream[SessionMessage | Exception] | None = None,
write_stream: WriteStream[SessionMessage] | None = None,
read_timeout_seconds: float | None = None,
sampling_callback: SamplingFnT | None = None,
elicitation_callback: ElicitationFnT | None = None,
@ -126,8 +149,9 @@ class ClientSession(
client_info: types.Implementation | None = None,
*,
sampling_capabilities: types.SamplingCapability | None = None,
dispatcher: Dispatcher[Any] | None = None,
) -> None:
super().__init__(read_stream, write_stream, read_timeout_seconds=read_timeout_seconds)
self._session_read_timeout_seconds = read_timeout_seconds
self._client_info = client_info or DEFAULT_CLIENT_INFO
self._sampling_callback = sampling_callback or _default_sampling_callback
self._sampling_capabilities = sampling_capabilities
@ -137,18 +161,99 @@ class ClientSession(
self._message_handler = message_handler or _default_message_handler
self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
self._initialize_result: types.InitializeResult | None = None
self._task_group: anyio.abc.TaskGroup | None = None
if dispatcher is not None:
if read_stream is not None or write_stream is not None:
raise ValueError("pass read_stream/write_stream or dispatcher, not both")
self._dispatcher: Dispatcher[Any] = dispatcher
else:
if read_stream is None or write_stream is None:
raise ValueError("read_stream and write_stream are required when no dispatcher is given")
# Built eagerly so notifications can be sent before entering the context manager.
self._dispatcher = JSONRPCDispatcher(
read_stream, write_stream, on_stream_exception=self._on_stream_exception
)
@property
def _receive_request_adapter(self) -> TypeAdapter[types.ServerRequest]:
return types.server_request_adapter
async def __aenter__(self) -> Self:
self._task_group = anyio.create_task_group()
await self._task_group.__aenter__()
try:
await self._task_group.start(self._dispatcher.run, self._on_request, self._on_notify)
except BaseException:
# Unwind the entered task group before propagating: a cancellation
# landing here (e.g. `move_on_after` around connect) would abandon
# it and anyio would later raise "exited non-innermost cancel scope".
task_group = self._task_group
self._task_group = None
task_group.cancel_scope.cancel()
# Shield the group's own scope (a new one would break LIFO exit)
# so a pending outer cancellation cannot re-fire inside __aexit__.
task_group.cancel_scope.shield = True
await task_group.__aexit__(None, None, None)
raise
return self
@property
def _receive_request_methods(self) -> frozenset[str]:
return _SERVER_REQUEST_METHODS
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> bool | None:
# Exit must not block: cancel the dispatcher and in-flight callbacks.
assert self._task_group is not None
self._task_group.cancel_scope.cancel()
result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
await resync_tracer()
return result
@property
def _receive_notification_adapter(self) -> TypeAdapter[types.ServerNotification]:
return types.server_notification_adapter
async def send_request(
self,
request: types.ClientRequest,
result_type: type[ReceiveResultT],
request_read_timeout_seconds: float | None = None,
metadata: ClientMessageMetadata | None = None,
progress_callback: ProgressFnT | None = None,
) -> ReceiveResultT:
"""Send a request and wait for its typed result.
Args:
metadata: Streamable HTTP resumption hints.
Raises:
MCPError: Error response, read timeout, or connection closed.
RuntimeError: Called before entering the context manager.
"""
data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
method: str = data["method"]
opts: CallOptions = {}
timeout = (
request_read_timeout_seconds
if request_read_timeout_seconds is not None
else self._session_read_timeout_seconds
)
if timeout is not None:
opts["timeout"] = timeout
if progress_callback is not None:
opts["on_progress"] = progress_callback
if metadata is not None:
if metadata.resumption_token is not None:
opts["resumption_token"] = metadata.resumption_token
if metadata.on_resumption_token_update is not None:
opts["on_resumption_token"] = metadata.on_resumption_token_update
if method == "initialize":
# The spec forbids cancelling initialize.
opts["cancel_on_abandon"] = False
raw = await self._dispatcher.send_raw_request(method, data.get("params"), opts)
return result_type.model_validate(raw, by_name=False)
async def send_notification(self, notification: types.ClientNotification) -> None:
"""Send a one-way notification. Usable before entering the context manager.
Fire-and-forget: after the connection has closed, the notification is
dropped with a debug log instead of raising.
"""
data = notification.model_dump(by_alias=True, mode="json", exclude_none=True)
await self._dispatcher.notify(data["method"], data.get("params"))
async def initialize(self) -> types.InitializeResult:
sampling = (
@ -397,49 +502,75 @@ class ClientSession(
"""Send a roots/list_changed notification."""
await self.send_notification(types.RootsListChangedNotification())
async def _received_request(self, responder: RequestResponder[types.ServerRequest, types.ClientResult]) -> None:
ctx = RequestContext[ClientSession](request_id=responder.request_id, meta=responder.request_meta, session=self)
async def _on_request(
self, dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
) -> dict[str, Any]:
"""Answer a server-initiated request via the registered callbacks."""
if method not in _SERVER_REQUEST_METHODS:
raise MCPError(code=types.METHOD_NOT_FOUND, message="Method not found", data=method)
payload: dict[str, Any] = {"method": method}
if params is not None:
payload["params"] = dict(params)
request = types.server_request_adapter.validate_python(payload, by_name=False)
match responder.request:
case types.CreateMessageRequest(params=params):
with responder:
response = await self._sampling_callback(ctx, params)
client_response = ClientResponse.validate_python(response)
await responder.respond(client_response)
case types.ElicitRequest(params=params):
with responder:
response = await self._elicitation_callback(ctx, params)
client_response = ClientResponse.validate_python(response)
await responder.respond(client_response)
case types.ListRootsRequest():
with responder:
response: types.ClientResult | types.ErrorData
if isinstance(request, types.PingRequest):
# Answered without a context: ping has no callback that would need one.
response = types.EmptyResult()
else:
assert dctx.request_id is not None # the callback-driving dispatchers always assign ids
ctx = ClientRequestContext(
session=self, request_id=dctx.request_id, meta=request.params.meta if request.params else None
)
match request:
case types.CreateMessageRequest(params=sampling_params):
response = await self._sampling_callback(ctx, sampling_params)
case types.ElicitRequest(params=elicit_params):
response = await self._elicitation_callback(ctx, elicit_params)
case types.ListRootsRequest(): # pragma: no branch
response = await self._list_roots_callback(ctx)
client_response = ClientResponse.validate_python(response)
await responder.respond(client_response)
client_response = ClientResponse.validate_python(response)
if isinstance(client_response, types.ErrorData):
raise MCPError.from_error_data(client_response)
return client_response.model_dump(by_alias=True, mode="json", exclude_none=True)
case types.PingRequest(): # pragma: no branch
with responder:
await responder.respond(types.EmptyResult())
async def _handle_incoming(
self,
req: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
async def _on_notify(
self, dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
) -> None:
"""Handle incoming messages by forwarding to the message handler."""
await self._message_handler(req)
"""Route a server notification: validate, run the typed callback, tee to message_handler."""
payload: dict[str, Any] = {"method": method}
if params is not None:
payload["params"] = dict(params)
try:
notification = types.server_notification_adapter.validate_python(payload, by_name=False)
except ValidationError:
logger.warning("Failed to validate notification: %s", payload, exc_info=True)
return
if isinstance(notification, types.CancelledNotification):
# The dispatcher already applied the cancellation; not surfaced to message_handler.
return
try:
if isinstance(notification, types.LoggingMessageNotification):
await self._logging_callback(notification.params)
await self._message_handler(notification)
except Exception:
# Contain here, not in the dispatcher: DirectDispatcher awaits this
# handler inline in the peer's notify() call, so a raising callback
# would otherwise fail the peer's send. A raising logging_callback
# skips the message_handler tee for that notification (v1 parity).
logger.exception("notification callback for %r raised", method)
async def _received_notification(self, notification: types.ServerNotification) -> None:
"""Handle notifications from the server."""
# Process specific notification types
match notification:
case types.LoggingMessageNotification(params=params):
await self._logging_callback(params)
case types.ElicitCompleteNotification(params=params):
# Handle elicitation completion notification
# Clients MAY use this to retry requests or update UI
# The notification contains the elicitationId of the completed elicitation
pass
case _:
pass
async def _on_stream_exception(self, exc: Exception) -> None:
"""Deliver a transport-level fault to message_handler via a spawned task.
Running the handler inline would park the dispatcher's read loop and
deadlock handlers that await session I/O.
"""
assert self._task_group is not None
self._task_group.start_soon(self._deliver_stream_exception, exc)
async def _deliver_stream_exception(self, exc: Exception) -> None:
try:
await self._message_handler(exc)
except Exception:
logger.exception("message_handler raised on transport exception")

2
src/mcp/server/session.py
View file

@ -84,7 +84,7 @@ class ServerSession:
"""
data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
opts: CallOptions = {}
if request_read_timeout_seconds:
if request_read_timeout_seconds is not None:
opts["timeout"] = request_read_timeout_seconds
if progress_callback is not None:
opts["on_progress"] = progress_callback

3
src/mcp/server/streamable_http.py
View file

@ -717,6 +717,9 @@ class StreamableHTTPServerTransport:
# Send the message via SSE
event_data = self._create_event_data(event_message)
await sse_stream_writer.send(event_data)
except anyio.ClosedResourceError:
# Session teardown can close the stream while the writer is between dequeues.
pass
except Exception:
logger.exception("Error in standalone SSE writer") # pragma: no cover
finally:

24
src/mcp/shared/_context.py
View file

@ -1,24 +0,0 @@
"""Request context for MCP handlers."""
from dataclasses import dataclass
from typing import Any, Generic
from typing_extensions import TypeVar
from mcp.shared.session import BaseSession
from mcp.types import RequestId, RequestParamsMeta
SessionT = TypeVar("SessionT", bound=BaseSession[Any, Any, Any, Any, Any])
@dataclass(kw_only=True)
class RequestContext(Generic[SessionT]):
"""Common context for handling incoming requests.
For request handlers, request_id is always populated.
For notification handlers, request_id is None.
"""
session: SessionT
request_id: RequestId | None = None
meta: RequestParamsMeta | None = None

109
src/mcp/shared/direct_dispatcher.py
View file

@ -15,6 +15,7 @@ to the caller - there is no exception-to-`ErrorData` boundary here.
from __future__ import annotations
import logging
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass, field
from typing import Any
@ -27,7 +28,9 @@ from mcp.shared.dispatcher import CallOptions, OnNotify, OnRequest, ProgressFnT
from mcp.shared.exceptions import MCPError, NoBackChannelError
from mcp.shared.message import MessageMetadata
from mcp.shared.transport_context import TransportContext
from mcp.types import INTERNAL_ERROR, INVALID_PARAMS, REQUEST_TIMEOUT, RequestId
from mcp.types import CONNECTION_CLOSED, INTERNAL_ERROR, INVALID_PARAMS, REQUEST_TIMEOUT, RequestId
logger = logging.getLogger(__name__)
__all__ = ["DirectDispatcher", "create_direct_dispatcher_pair"]
@ -50,7 +53,7 @@ class _DirectDispatchContext:
_back_request: _Request
_back_notify: _Notify
request_id: RequestId | None = None
"""Always `None`: direct dispatch has no wire-level request id."""
"""A dispatcher-synthesized id for requests; `None` for notifications."""
message_metadata: MessageMetadata = None # TODO(maxisbey): remove for Context rework
"""Always `None`: in-memory dispatch attaches no transport metadata."""
_on_progress: ProgressFnT | None = None
@ -84,6 +87,13 @@ class DirectDispatcher:
Two instances are wired together with `create_direct_dispatcher_pair`; each
holds a reference to the other. `send_raw_request` on one awaits the peer's
`on_request`. `run` parks until `close` is called.
Lifecycle mirrors `JSONRPCDispatcher`: `send_raw_request` requires `run()`
to have started, and once a side has closed - via `close()` or `run()`
ending - `send_raw_request` raises `MCPError` (`CONNECTION_CLOSED`) and
inbound requests fail the peer's call the same way instead of invoking the
handler. Notifications are fire-and-forget in both directions: after close
they are silently dropped.
"""
def __init__(self, transport_ctx: TransportContext):
@ -91,8 +101,11 @@ class DirectDispatcher:
self._peer: DirectDispatcher | None = None
self._on_request: OnRequest | None = None
self._on_notify: OnNotify | None = None
self._next_id = 0
self._ready = anyio.Event()
self._closed = anyio.Event()
self._close_event = anyio.Event()
self._running = False
self._closed = False
def connect_to(self, peer: DirectDispatcher) -> None:
self._peer = peer
@ -103,13 +116,35 @@ class DirectDispatcher:
params: Mapping[str, Any] | None,
opts: CallOptions | None = None,
) -> dict[str, Any]:
"""Send a request by invoking the peer's `on_request` directly.
Raises:
MCPError: The peer's handler raised; `REQUEST_TIMEOUT` if
`opts["timeout"]` elapsed; `CONNECTION_CLOSED` if either
side has closed.
RuntimeError: Called before `run()`.
"""
if self._peer is None:
raise RuntimeError("DirectDispatcher has no peer; use create_direct_dispatcher_pair()")
# Post-close sends get the same CONNECTION_CLOSED contract as JSONRPCDispatcher.
if self._closed:
raise MCPError(code=CONNECTION_CLOSED, message="Connection closed")
if not self._running:
raise RuntimeError("DirectDispatcher.send_raw_request called before run()")
return await self._peer._dispatch_request(method, params, opts)
async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
"""Send a notification by invoking the peer's `on_notify` directly.
Fire-and-forget: usable before `run()` (delivery waits for the peer to
start), and after close it is silently dropped, matching
`JSONRPCDispatcher.notify`.
"""
if self._peer is None:
raise RuntimeError("DirectDispatcher has no peer; use create_direct_dispatcher_pair()")
if self._closed:
logger.debug("dropped notification %r on closed DirectDispatcher", method)
return
await self._peer._dispatch_notify(method, params)
async def run(
@ -119,37 +154,77 @@ class DirectDispatcher:
*,
task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
) -> None:
self._on_request = on_request
self._on_notify = on_notify
self._ready.set()
task_status.started()
await self._closed.wait()
"""Mark this side ready and park until `close()` is called.
Single-shot, like `JSONRPCDispatcher.run`: once it returns the
dispatcher stays closed and cannot be restarted.
"""
try:
self._on_request = on_request
self._on_notify = on_notify
self._running = True
self._ready.set()
task_status.started()
await self._close_event.wait()
finally:
self._running = False
self._closed = True
# run() may end via cancellation without close() ever being
# called; setting the event wakes `_wait_ready` waiters so they
# observe the closed state instead of parking forever.
self._close_event.set()
def close(self) -> None:
self._closed.set()
self._closed = True
self._close_event.set()
def _make_context(self, on_progress: ProgressFnT | None = None) -> _DirectDispatchContext:
def _make_context(
self, on_progress: ProgressFnT | None = None, request_id: RequestId | None = None
) -> _DirectDispatchContext:
assert self._peer is not None
peer = self._peer
return _DirectDispatchContext(
transport=self._transport_ctx,
_back_request=lambda m, p, o: peer._dispatch_request(m, p, o),
_back_notify=lambda m, p: peer._dispatch_notify(m, p),
request_id=request_id,
_on_progress=on_progress,
)
async def _wait_ready(self) -> None:
"""Park until `run()` has started, waking early if this side closes.
Raises:
MCPError: `CONNECTION_CLOSED` if this side has closed.
"""
if not self._ready.is_set() and not self._close_event.is_set():
async with anyio.create_task_group() as tg:
async def wake_on(event: anyio.Event) -> None:
await event.wait()
tg.cancel_scope.cancel()
tg.start_soon(wake_on, self._ready)
tg.start_soon(wake_on, self._close_event)
if self._closed:
raise MCPError(code=CONNECTION_CLOSED, message="Connection closed")
async def _dispatch_request(
self,
method: str,
params: Mapping[str, Any] | None,
opts: CallOptions | None,
) -> dict[str, Any]:
await self._ready.wait()
assert self._on_request is not None
opts = opts or {}
dctx = self._make_context(on_progress=opts.get("on_progress"))
try:
with anyio.fail_after(opts.get("timeout")):
# Inside the timeout scope, so a configured timeout also bounds
# waiting on a peer whose run() has not started yet.
await self._wait_ready()
assert self._on_request is not None
# Synthesize an id: the DispatchContext contract reserves None for notifications.
self._next_id += 1
dctx = self._make_context(on_progress=opts.get("on_progress"), request_id=self._next_id)
try:
return await self._on_request(dctx, method, params)
except MCPError:
@ -167,7 +242,13 @@ class DirectDispatcher:
) from None
async def _dispatch_notify(self, method: str, params: Mapping[str, Any] | None) -> None:
await self._ready.wait()
try:
await self._wait_ready()
except MCPError:
# Notifications are fire-and-forget: a notify to a closed peer is
# dropped, not raised back into the sender's call.
logger.debug("dropped notification %r to closed DirectDispatcher", method)
return
assert self._on_notify is not None
dctx = self._make_context()
await self._on_notify(dctx, method, params)

12
src/mcp/shared/dispatcher.py
View file

@ -55,6 +55,13 @@ class CallOptions(TypedDict, total=False):
timeout: float
"""Seconds to wait for a result before raising and sending `notifications/cancelled`."""
cancel_on_abandon: bool
"""Whether abandoning this request (timeout or caller cancellation) sends `notifications/cancelled`.
Defaults to `True`. Set `False` for requests the protocol forbids cancelling, such as `initialize`.
Also suppressed when resumption hints reach the transport, or when the request was never written.
"""
on_progress: ProgressFnT
"""Receive `notifications/progress` updates for this request."""
@ -97,9 +104,6 @@ class Outbound(Protocol):
) -> dict[str, Any]:
"""Send a request and await its raw result dict.
`opts` carries per-call `timeout` / `on_progress` / resumption hints;
see `CallOptions`.
Raises:
MCPError: If the peer responded with an error, or the handler
raised. Implementations normalize all handler exceptions to
@ -187,6 +191,8 @@ class Dispatcher(Outbound, Protocol[TransportT_co]):
Implementations own correlation of outbound requests to inbound results, the
receive loop, per-request concurrency, and cancellation/progress wiring.
The lifecycle surface is provisional; `run()` may change before v2 stable.
"""
async def run(

558
src/mcp/shared/jsonrpc_dispatcher.py
View file

@ -1,21 +1,8 @@
"""JSON-RPC `Dispatcher` implementation.
"""JSON-RPC `Dispatcher` over the `SessionMessage` stream contract all transports speak.
Consumes the existing `SessionMessage`-based stream contract that all current
transports (stdio, SSE, streamable HTTP) speak. Owns request-id correlation,
the receive loop, per-request task isolation, cancellation/progress wiring, and
the single exception-to-wire boundary.
The MCP type layer (`ServerRunner`, `Context`, `Client`) sits above this and
sees only `(ctx, method, params) -> dict`. Transports sit below and see only
`SessionMessage` reads/writes.
The dispatcher is *mostly* MCP-agnostic - methods/params are opaque strings and
dicts - but it intercepts `notifications/cancelled` and
`notifications/progress` because request correlation, cancellation and
progress are exactly the wiring this layer exists to provide. Those few wire
shapes are extracted with structural `match` patterns (no casts, no
`mcp.types` model coupling); a malformed payload simply fails to match and
the correlation is skipped.
Owns request-id correlation, the receive loop, per-request task isolation,
cancellation/progress wiring, and the single exception-to-wire boundary;
methods and params are otherwise opaque strings and dicts.
"""
from __future__ import annotations
@ -24,17 +11,21 @@ import contextvars
import logging
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass, field
from typing import Any, Generic, Literal, TypeVar, cast, overload
from functools import partial
from typing import Any, Generic, Literal, cast
import anyio
import anyio.abc
import anyio.lowlevel
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from opentelemetry.trace import SpanKind
from pydantic import ValidationError
from typing_extensions import TypeVar
from mcp.shared._compat import resync_tracer
from mcp.shared._otel import inject_trace_context, otel_span
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.dispatcher import CallOptions, Dispatcher, OnNotify, OnRequest, ProgressFnT
from mcp.shared.dispatcher import CallOptions, DispatchContext, Dispatcher, OnNotify, OnRequest, ProgressFnT
from mcp.shared.exceptions import MCPError, NoBackChannelError
from mcp.shared.message import (
ClientMessageMetadata,
@ -47,7 +38,6 @@ from mcp.types import (
CONNECTION_CLOSED,
INTERNAL_ERROR,
INVALID_PARAMS,
REQUEST_CANCELLED,
REQUEST_TIMEOUT,
ErrorData,
JSONRPCError,
@ -63,23 +53,22 @@ __all__ = ["JSONRPCDispatcher"]
logger = logging.getLogger(__name__)
TransportT = TypeVar("TransportT", bound=TransportContext)
_ABANDON_WRITE_TIMEOUT: float = 5
"""Bound for courtesy-cancel writes on the abandon paths; the caller-cancel
arm shields its write, so a wedged transport would otherwise hang it uncancellably."""
_SHUTDOWN_WRITE_TIMEOUT: float = 1
"""Tighter bound for the shutdown-arm error write so a wedged transport can't hold session close."""
TransportT = TypeVar("TransportT", bound=TransportContext, default=TransportContext)
PeerCancelMode = Literal["interrupt", "signal"]
"""How inbound `notifications/cancelled` is applied to a running handler.
`"interrupt"` (default) cancels the handler's scope. `"signal"` only sets
`ctx.cancel_requested` and lets the handler observe it cooperatively.
"""
"""How `notifications/cancelled` is applied: `"interrupt"` (default) cancels
the handler's scope; `"signal"` only sets `ctx.cancel_requested`."""
def _coerce_id(request_id: RequestId) -> RequestId:
"""Coerce a string request ID to int when it's a valid int literal.
`_allocate_id` only ever produces `int` keys for `_pending`, but a peer
may echo the ID back as a JSON string. The TypeScript SDK and `BaseSession`
both perform this coercion at lookup time so the response still correlates.
"""
"""Coerce a stringified int request ID back to int so a peer-echoed ID still correlates (matches the TS SDK)."""
if isinstance(request_id, str):
try:
return int(request_id)
@ -113,12 +102,7 @@ class _JSONRPCDispatchContext(Generic[TransportT]):
_dispatcher: JSONRPCDispatcher[TransportT]
_request_id: RequestId | None
message_metadata: MessageMetadata = None # TODO(maxisbey): remove for Context rework
"""The transport-attached `SessionMessage.metadata` for this inbound message.
Carries `ServerMessageMetadata` (HTTP request, SSE stream-close callbacks)
that the server lifts onto its request context. `None` for transports
that attach nothing.
"""
"""Transport-attached `SessionMessage.metadata` that the server lifts onto its request context."""
_progress_token: ProgressToken | None = None
_closed: bool = False
cancel_requested: anyio.Event = field(default_factory=anyio.Event)
@ -166,13 +150,7 @@ def _default_transport_builder(_meta: MessageMetadata) -> TransportContext:
def _shielded_progress(fn: ProgressFnT) -> ProgressFnT:
"""Wrap a user progress callback so it can't crash the dispatcher.
The callback runs as a bare task in the dispatcher's task group; an
uncaught exception would cancel every sibling (the read loop and all
in-flight requests). Swallow and log instead, matching the previous
receive-loop's behavior.
"""
"""Wrap a user progress callback so an exception can't cancel the dispatcher's task group."""
async def _wrapped(progress: float, total: float | None, message: str | None) -> None:
try:
@ -183,61 +161,57 @@ def _shielded_progress(fn: ProgressFnT) -> ProgressFnT:
return _wrapped
def _outbound_metadata(related_request_id: RequestId | None, opts: CallOptions | None) -> MessageMetadata:
"""Choose the `SessionMessage.metadata` for an outgoing request/notification.
def _contained_notify(fn: OnNotify) -> OnNotify:
"""Wrap a notification handler so it can't crash the dispatcher (same boundary as `_shielded_progress`)."""
`ServerMessageMetadata` tags a server-to-client message with the inbound
request it belongs to (so streamable-HTTP can route it onto that request's
SSE stream). `ClientMessageMetadata` carries resumption hints to the
client transport. `None` is the common case.
async def _wrapped(dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None) -> None:
try:
await fn(dctx, method, params)
except Exception:
logger.exception("notification handler for %r raised", method)
`SessionMessage.metadata` carries exactly one of these, so when
`related_request_id` is set it takes precedence and any resumption hints
in `opts` are dropped (with a debug log): requests made from a dispatch
context are routed onto the inbound request's stream, not resumed.
return _wrapped
@dataclass(slots=True, frozen=True)
class _OutboundPlan:
"""Outbound metadata plus whether abandoning the request sends a courtesy `notifications/cancelled`."""
metadata: MessageMetadata
cancel_on_abandon: bool
def _plan_outbound(related_request_id: RequestId | None, opts: CallOptions | None) -> _OutboundPlan:
"""Choose the outbound `SessionMessage.metadata` and the abandon-cancellation policy.
`related_request_id` wins over resumption hints (they are dropped). Only
hints that actually reach the transport suppress the courtesy cancel - a
request that is neither resumable nor cancelled would leak the peer's work.
"""
opts = opts or {}
cancel_on_abandon = opts.get("cancel_on_abandon", True)
token = opts.get("resumption_token")
on_token = opts.get("on_resumption_token")
if related_request_id is not None:
if opts and (opts.get("resumption_token") is not None or opts.get("on_resumption_token") is not None):
if token is not None or on_token is not None:
logger.debug(
"dropping resumption hints: related_request_id %r takes precedence on metadata", related_request_id
)
return ServerMessageMetadata(related_request_id=related_request_id)
if opts:
token = opts.get("resumption_token")
on_token = opts.get("on_resumption_token")
if token is not None or on_token is not None:
return ClientMessageMetadata(resumption_token=token, on_resumption_token_update=on_token)
return None
return _OutboundPlan(ServerMessageMetadata(related_request_id=related_request_id), cancel_on_abandon)
if token is not None or on_token is not None:
return _OutboundPlan(
ClientMessageMetadata(resumption_token=token, on_resumption_token_update=on_token),
cancel_on_abandon=False,
)
return _OutboundPlan(None, cancel_on_abandon)
class JSONRPCDispatcher(Dispatcher[TransportT]):
"""`Dispatcher` over the existing `SessionMessage` stream contract.
"""`Dispatcher` over the `SessionMessage` stream contract.
Inherits the `Dispatcher` Protocol explicitly so pyright checks
conformance at the class definition rather than at first use.
Explicit Protocol base so pyright checks conformance at the class definition.
"""
@overload
def __init__(
self: JSONRPCDispatcher[TransportContext],
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[SessionMessage],
*,
peer_cancel_mode: PeerCancelMode = "interrupt",
raise_handler_exceptions: bool = False,
inline_methods: frozenset[str] = frozenset(),
) -> None: ...
@overload
def __init__(
self,
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[SessionMessage],
*,
transport_builder: Callable[[MessageMetadata], TransportT],
peer_cancel_mode: PeerCancelMode = "interrupt",
raise_handler_exceptions: bool = False,
inline_methods: frozenset[str] = frozenset(),
) -> None: ...
def __init__(
self,
read_stream: ReadStream[SessionMessage | Exception],
@ -247,33 +221,41 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
peer_cancel_mode: PeerCancelMode = "interrupt",
raise_handler_exceptions: bool = False,
inline_methods: frozenset[str] = frozenset(),
on_stream_exception: Callable[[Exception], Awaitable[None]] | None = None,
) -> None:
"""Wire a dispatcher over a transport's `SessionMessage` stream pair.
Args:
transport_builder: Builds each message's `TransportContext` from
its `SessionMessage.metadata`.
raise_handler_exceptions: Re-raise handler exceptions out of
`run()` after the error response is written.
inline_methods: Methods awaited in the read loop before the next
message is dequeued (e.g. `initialize`); an inline handler
that awaits the peer deadlocks the parked loop.
on_stream_exception: Observer for `Exception` items on the read
stream; without it they are debug-logged and dropped. Awaited
inline in the read loop, so a slow observer stalls dispatch.
"""
self._read_stream = read_stream
self._write_stream = write_stream
# The overloads guarantee that when `transport_builder` is omitted,
# `TransportT` is `TransportContext`, so the default is type-correct;
# pyright can't see across overloads, hence the cast.
# With transport_builder omitted, TransportT defaults to
# TransportContext; pyright can't connect the two, hence the cast.
self._transport_builder = cast(
"Callable[[MessageMetadata], TransportT]",
transport_builder or _default_transport_builder,
)
self._peer_cancel_mode: PeerCancelMode = peer_cancel_mode
self._raise_handler_exceptions = raise_handler_exceptions
# Request methods handled inline in the read loop (awaited before the
# next message is dequeued) instead of spawned concurrently. Use for
# methods whose side effects must be observable to the next message,
# e.g. `initialize`, so a pipelined follow-up sees the initialized state.
# Only suitable for handlers that complete quickly, since inline handling
# blocks dequeuing; a handler that awaits the peer (`send_raw_request`)
# while inline will deadlock because the parked read loop cannot dequeue
# the response.
self._inline_methods = inline_methods
self._on_stream_exception = on_stream_exception
self._next_id = 0
self._pending: dict[RequestId, _Pending] = {}
self._in_flight: dict[RequestId, _InFlight[TransportT]] = {}
self._tg: anyio.abc.TaskGroup | None = None
self._running = False
self._closed = False
async def send_raw_request(
self,
@ -285,82 +267,105 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
) -> dict[str, Any]:
"""Send a JSON-RPC request and await its response.
`_related_request_id` is set only by `_JSONRPCDispatchContext` when a
handler makes a server-to-client request mid-flight; it routes the
outgoing message onto the correct per-request SSE stream (SHTTP) via
`ServerMessageMetadata`. Top-level callers leave it `None`.
`_related_request_id` is set only by `_JSONRPCDispatchContext` so that
mid-handler requests route onto the inbound request's SSE stream.
Raises:
MCPError: The peer responded with a JSON-RPC error; or
`REQUEST_TIMEOUT` if `opts["timeout"]` elapsed; or
`CONNECTION_CLOSED` if the dispatcher shut down while
awaiting the response.
RuntimeError: Called before `run()` has started or after it has
finished.
MCPError: Peer error response; `REQUEST_TIMEOUT` if
`opts["timeout"]` elapsed; `CONNECTION_CLOSED` if the
transport closed or the dispatcher shut down.
RuntimeError: Called before `run()`.
"""
# Post-close sends get the same CONNECTION_CLOSED contract as in-flight waiters.
if self._closed:
raise MCPError(code=CONNECTION_CLOSED, message="Connection closed")
if not self._running:
raise RuntimeError("JSONRPCDispatcher.send_raw_request called before run() / after close")
raise RuntimeError("JSONRPCDispatcher.send_raw_request called before run()")
opts = opts or {}
request_id = self._allocate_id()
out_params = dict(params) if params is not None else {}
out_meta = dict(out_params.get("_meta") or {})
on_progress = opts.get("on_progress")
if on_progress is not None:
# The caller wants progress updates. The spec mechanism is: include
# `_meta.progressToken` on the request; the peer echoes that token on
# any `notifications/progress` it sends. We use the request id as the
# token so the receive loop can find this `_Pending.on_progress` by
# `_pending[token]` without a second lookup table.
# The request id doubles as the progress token, so `_pending[token]` finds `on_progress` directly.
out_meta["progressToken"] = request_id
out_params["_meta"] = out_meta
# buffer=1: at most one outcome is ever delivered. A `WouldBlock` from
# `_resolve_pending`/`_fan_out_closed` means the waiter already has an
# outcome and dropping the late/redundant signal is correct. buffer=0
# is unsafe - there's a window between registering `_pending[id]` and
# parking in `receive()` where a close signal would be lost.
# buffer=1: a close signal can arrive before the waiter parks in receive();
# a WouldBlock later just means the waiter already has its one outcome.
send, receive = anyio.create_memory_object_stream[dict[str, Any] | ErrorData](1)
pending = _Pending(send=send, receive=receive, on_progress=on_progress)
self._pending[request_id] = pending
metadata = _outbound_metadata(_related_request_id, opts)
plan = _plan_outbound(_related_request_id, opts)
# Spec MUST: only previously-issued requests may be cancelled. A write
# interrupted by cancellation may still have delivered (a memory-stream
# send can hand its item to the receiver and still raise), so a started
# write counts as issued: the peer ignores a cancel for an id it never
# saw, while skipping it would leak a delivered request's handler.
request_write_started = False
timeout_armed = False
target = out_params.get("name")
span_name = f"MCP send {method}{f' {target}' if isinstance(target, str) else ''}"
# TODO(maxisbey): the otel span + inject below mirror
# BaseSession.send_request for parity. They belong in an outbound
# middleware (symmetric with otel_middleware on the inbound side) once
# that seam exists; the dispatcher should not own otel.
# TODO(maxisbey): move the otel span + inject into an outbound
# middleware once that seam exists; the dispatcher should not own otel.
try:
with otel_span(
span_name,
kind=SpanKind.CLIENT,
attributes={"mcp.method.name": method, "jsonrpc.request.id": str(request_id)},
):
# Inject W3C trace context into _meta (SEP-414). With a no-op
# tracer this writes nothing, but `_meta` itself is still
# present on the wire (and the interaction suite pins that).
# SEP-414: inject W3C trace context; `_meta` stays on the wire even with a no-op tracer.
inject_trace_context(out_meta)
msg = JSONRPCRequest(jsonrpc="2.0", id=request_id, method=method, params=out_params)
await self._write(msg, metadata)
# Surface a pre-existing cancellation while the request provably
# never started; past this point a cancelled write counts as issued.
await anyio.lowlevel.checkpoint_if_cancelled()
request_write_started = True
try:
await self._write(msg, plan.metadata)
except (anyio.BrokenResourceError, anyio.ClosedResourceError):
# Transport tore down before run() noticed EOF; surface the documented contract.
raise MCPError(code=CONNECTION_CLOSED, message="Connection closed") from None
with anyio.fail_after(opts.get("timeout")):
timeout_armed = True
outcome = await receive.receive()
except TimeoutError:
# Spec-recommended courtesy: tell the peer we've given up so it can
# stop work and free resources. v1's BaseSession.send_request does
# NOT do this; it's new behaviour.
await self._cancel_outbound(request_id, f"timed out after {opts.get('timeout')}s", _related_request_id)
if not timeout_armed:
# `fail_after` arms only after the write, so this TimeoutError is the
# transport's own bounded send() failing - a transport error, not
# `opts["timeout"]` elapsing. Propagate it raw (v1 kept the write
# outside the timeout-catching try and did the same).
raise
# Courtesy cancel (spec-recommended, new vs v1) so the peer stops work;
# unshielded so an outer caller cancellation can still interrupt the write.
if plan.cancel_on_abandon:
await self._final_write(
partial(
self._cancel_outbound,
request_id,
f"timed out after {opts.get('timeout')}s",
_related_request_id,
),
shield=False,
timeout=_ABANDON_WRITE_TIMEOUT,
describe=f"courtesy cancel for timed-out request {request_id!r}",
)
raise MCPError(code=REQUEST_TIMEOUT, message=f"Request {method!r} timed out") from None
except anyio.get_cancelled_exc_class():
# Our caller's scope was cancelled. We're already inside a cancelled
# scope, so any bare `await` here re-raises immediately - shield to
# let the courtesy cancel notification go out before we propagate.
with anyio.CancelScope(shield=True):
await self._cancel_outbound(request_id, "caller cancelled", _related_request_id)
# Caller cancelled: bare awaits re-raise here, so the shielded helper
# lets the courtesy cancel go out before we propagate.
if plan.cancel_on_abandon and request_write_started:
await self._final_write(
partial(self._cancel_outbound, request_id, "caller cancelled", _related_request_id),
shield=True,
timeout=_ABANDON_WRITE_TIMEOUT,
describe=f"courtesy cancel for caller-cancelled request {request_id!r}",
)
raise
finally:
# Always remove the waiter, even on cancel/timeout, so a late
# response from the peer (race) hits a closed stream and is dropped
# in `_dispatch` rather than leaking.
# Remove the waiter on every path so a late response is dropped, not leaked.
self._pending.pop(request_id, None)
send.close()
receive.close()
@ -376,15 +381,26 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
*,
_related_request_id: RequestId | None = None,
) -> None:
# Leave `params` unset (not explicitly None) when there are none:
# transports serialize with `exclude_unset=True`, and an explicit None
# would survive as `"params": null`, which JSON-RPC 2.0 forbids and
# strict peers (e.g. the TypeScript SDK's zod schemas) reject.
"""Send a fire-and-forget notification.
Fire-and-forget all the way: a post-close send or a write onto a
torn-down transport drops the notification with a debug log instead
of raising (same policy as the response writes and `ctx.notify`).
"""
if self._closed:
logger.debug("dropped %s: dispatcher closed", method)
return
# Leave `params` unset when None: with `exclude_unset=True` an explicit
# None would serialize as `"params": null`, which JSON-RPC 2.0 forbids.
if params is not None:
msg = JSONRPCNotification(jsonrpc="2.0", method=method, params=dict(params))
else:
msg = JSONRPCNotification(jsonrpc="2.0", method=method)
await self._write(msg, _outbound_metadata(_related_request_id, None))
try:
await self._write(msg, _plan_outbound(_related_request_id, None).metadata)
except (anyio.BrokenResourceError, anyio.ClosedResourceError):
# Transport tore down before run() noticed EOF.
logger.debug("dropped %s: write stream closed", method)
async def run(
self,
@ -395,50 +411,44 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
) -> None:
"""Drive the receive loop until the read stream closes.
Each inbound request is handled in its own task in an internal task
group; `task_status.started()` fires once that group is open, so
`await tg.start(dispatcher.run, ...)` resumes when `send_raw_request`
is usable.
`task_status.started()` fires once `send_raw_request` is usable.
Single-shot: once the loop ends the dispatcher stays closed and cannot be restarted.
"""
try:
async with anyio.create_task_group() as tg:
self._tg = tg
self._running = True
task_status.started()
try:
async with self._read_stream, self._write_stream:
try:
async for item in self._read_stream:
# Duck-typed: `_context_streams.ContextReceiveStream`
# exposes `.last_context` (the sender's contextvars
# snapshot per message). Plain memory streams don't.
sender_ctx: contextvars.Context | None = getattr(
self._read_stream, "last_context", None
)
await self._dispatch(item, on_request, on_notify, sender_ctx)
except anyio.ClosedResourceError:
# The transport closed our receive end and we looped
# back to `__anext__` on the now-closed stream
# (stateless SHTTP teardown). Same as EOF.
logger.debug("read stream closed by transport; treating as EOF")
# Read stream EOF: wake any blocked `send_raw_request` waiters
# (callers outside this task group) with CONNECTION_CLOSED.
self._running = False
self._fan_out_closed()
finally:
# Transport closed: cancel in-flight handlers. Without this
# the task-group join waits for them, and a handler that
# outlives its caller (its request timed out client-side, or
# the client disconnected mid-call) would keep `run()` from
# returning forever. Same behaviour as `Server.run()` before
# the dispatcher rework.
tg.cancel_scope.cancel()
# LIFO exits: the write stream closes only after the task-group join, so teardown writes still land.
async with self._write_stream:
async with anyio.create_task_group() as tg:
self._tg = tg
self._running = True
task_status.started()
try:
async with self._read_stream:
try:
async for item in self._read_stream:
# Duck-typed: only `ContextReceiveStream` carries the
# sender's per-message contextvars snapshot.
sender_ctx: contextvars.Context | None = getattr(
self._read_stream, "last_context", None
)
await self._dispatch(item, on_request, on_notify, sender_ctx)
except anyio.ClosedResourceError:
# Receive end closed under us (stateless SHTTP teardown); same as EOF.
logger.debug("read stream closed by transport; treating as EOF")
# EOF: wake blocked `send_raw_request` waiters with CONNECTION_CLOSED.
self._running = False
self._closed = True
self._fan_out_closed()
finally:
# Cancel in-flight handlers; otherwise the task-group join
# waits on handlers whose callers are already gone.
tg.cancel_scope.cancel()
finally:
# Covers the cancel/crash paths where the inline fan-out above is
# never reached. Idempotent.
# Covers cancel/crash paths that skip the inline fan-out; idempotent.
self._running = False
self._closed = True
self._tg = None
self._fan_out_closed()
await resync_tracer()
async def _dispatch(
self,
@ -449,13 +459,17 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
) -> None:
"""Route one inbound item.
Everything here is `send_nowait` or `_spawn`; the only `await` is for
`inline_methods` requests, which deliberately block dequeuing until
handled. Any other `await` would let one slow message head-of-line
block the entire read loop.
Only `inline_methods` requests and the `on_stream_exception` observer
are awaited; any other `await` would head-of-line block the read loop.
"""
if isinstance(item, Exception):
logger.debug("transport yielded exception: %r", item)
if self._on_stream_exception is None:
logger.debug("transport yielded exception: %r", item)
return
try:
await self._on_stream_exception(item)
except Exception:
logger.exception("on_stream_exception observer raised")
return
metadata = item.metadata
msg = item.message
@ -467,9 +481,7 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
case JSONRPCResponse():
self._resolve_pending(msg.id, msg.result)
case JSONRPCError(): # pragma: no branch
# `id` may be None per JSON-RPC (parse error before id known).
# The match is exhaustive over JSONRPCMessage; the no-match arc
# on this final case is unreachable.
# Exhaustive over JSONRPCMessage, so the no-match arc is unreachable.
self._resolve_pending(msg.id, msg.error)
async def _dispatch_request(
@ -481,8 +493,7 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
) -> None:
progress_token: ProgressToken | None
match req.params:
# The bool guard matters: `int()` patterns match bool (a subclass),
# and `True == 1` would alias dict lookups to request id 1.
# bool subclasses int: without the guard True would alias request id 1.
case {"_meta": {"progressToken": str() | int() as progress_token}} if not isinstance(progress_token, bool):
pass
case _:
@ -490,9 +501,7 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
try:
transport_ctx = self._transport_builder(metadata)
except Exception:
# Containment boundary for the user-supplied builder: a raising
# builder must cost only this message, not the whole connection
# (the exception would otherwise escape into run()'s read loop).
# A raising builder must cost only this message, not the connection.
logger.exception("transport_builder raised; rejecting request %r", req.id)
self._spawn(
self._write_error,
@ -509,20 +518,13 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
_progress_token=progress_token,
)
scope = anyio.CancelScope()
# TODO(maxisbey): the spec puts request-id uniqueness on the sender;
# neither v1 nor the TS SDK guards a duplicate id here, so for now we
# blind-overwrite (parity). Revisit rejecting with INVALID_REQUEST.
# Coerced key so `notifications/cancelled` correlates regardless of
# whether the peer stringifies the id between request and cancel
# (`_dispatch_notification` coerces at lookup; responses still echo
# `req.id` verbatim).
# TODO(maxisbey): duplicate ids blind-overwrite (v1/TS parity); revisit
# rejecting with INVALID_REQUEST. Key coerced so a stringified
# `notifications/cancelled` id still correlates.
self._in_flight[_coerce_id(req.id)] = _InFlight(scope=scope, dctx=dctx)
if req.method in self._inline_methods:
# Spawn (so `sender_ctx` applies, matching the concurrent path) but
# park the read loop until the handler returns; that's the inline
# ordering guarantee. Because the read loop is parked, a handler
# that awaits the peer here (e.g. `dctx.send_raw_request`) will
# deadlock: the response can never be dequeued.
# Spawn so `sender_ctx` applies, but park the read loop until the
# handler returns - that's the inline ordering guarantee.
done = anyio.Event()
async def _run_inline() -> None:
@ -546,17 +548,12 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
"""Route one inbound notification.
`notifications/cancelled` and `notifications/progress` are intercepted
here because they correlate against JSON-RPC request IDs - the
`_in_flight` / `_pending` tables this layer owns - so no higher layer
can act on them. Both are still teed to `on_notify` afterwards, so
middleware and registered notification handlers observe every inbound
notification. See the module docstring for the design rationale.
here (they correlate against the `_in_flight`/`_pending` tables this
layer owns) and still teed to `on_notify` afterwards.
"""
if msg.method == "notifications/cancelled":
match msg.params:
# The bool guards here and below matter: `int()` patterns match
# bool (a subclass), and `True == 1` would alias the dict lookup
# to the entry keyed by request id 1.
# bool subclasses int: the guards keep True from aliasing request id 1.
case {"requestId": str() | int() as rid} if (
not isinstance(rid, bool) and (in_flight := self._in_flight.get(_coerce_id(rid))) is not None
):
@ -565,9 +562,6 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
in_flight.scope.cancel()
case _:
pass
# fall through: cancelled is also teed to on_notify so middleware
# and registered handlers can observe it (matches DirectDispatcher,
# which forwards every notification).
elif msg.method == "notifications/progress":
match msg.params:
case {"progressToken": str() | int() as token, "progress": int() | float() as progress} if (
@ -587,18 +581,16 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
)
case _:
pass
# fall through: progress is also teed to on_notify
try:
transport_ctx = self._transport_builder(metadata)
except Exception:
# Same containment boundary as `_dispatch_request`: a raising
# builder drops this notification instead of killing the read loop.
# Same containment as `_dispatch_request`: drop the notification, keep the loop.
logger.exception("transport_builder raised; dropping notification %r", msg.method)
return
dctx = _JSONRPCDispatchContext(
transport=transport_ctx, _dispatcher=self, _request_id=None, message_metadata=metadata
)
self._spawn(on_notify, dctx, msg.method, msg.params, sender_ctx=sender_ctx)
self._spawn(_contained_notify(on_notify), dctx, msg.method, msg.params, sender_ctx=sender_ctx)
def _resolve_pending(self, request_id: RequestId | None, outcome: dict[str, Any] | ErrorData) -> None:
pending = self._pending.get(_coerce_id(request_id)) if request_id is not None else None
@ -618,10 +610,8 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
) -> None:
"""Schedule `fn(*args)` in the run() task group, propagating the sender's contextvars.
ASGI middleware (auth, OTel) sets contextvars on the request task that
wrote into the read stream. `Context.run(tg.start_soon, ...)` makes
the spawned handler inherit *that* context instead of the receive
loop's, so `auth_context_var` and OTel spans survive.
ASGI middleware (auth, OTel) sets contextvars on the task that wrote the
message; `Context.run` makes the spawned handler inherit that context.
"""
assert self._tg is not None
if sender_ctx is not None:
@ -632,10 +622,9 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
def _fan_out_closed(self) -> None:
"""Wake every pending `send_raw_request` waiter with `CONNECTION_CLOSED`.
Synchronous (uses `send_nowait`) because it's called from `finally`
which may be inside a cancelled scope. Idempotent.
Synchronous: callers may be inside a cancelled scope. Idempotent.
"""
closed = ErrorData(code=CONNECTION_CLOSED, message="connection closed")
closed = ErrorData(code=CONNECTION_CLOSED, message="Connection closed")
for pending in self._pending.values():
try:
pending.send.send_nowait(closed)
@ -652,61 +641,65 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
) -> None:
"""Run `on_request` for one inbound request and write its response.
This is the single exception-to-wire boundary: handler exceptions are
caught here and serialized to `JSONRPCError`. Nothing above this in
the stack constructs wire errors.
The single exception-to-wire boundary: handler exceptions become `JSONRPCError` here.
"""
answer_write_started = False
try:
with scope:
try:
result = await on_request(dctx, req.method, req.params)
finally:
# Handler done: close the back-channel (detached work that
# later calls `dctx.send_raw_request()` should see
# `NoBackChannelError`) and drop from `_in_flight` so a
# late `notifications/cancelled` is a no-op rather than
# racing the result write below. No checkpoint between
# handler return and the pop, so the cancel can't
# interleave there.
# Close the back-channel and drop from `_in_flight`; no checkpoint
# since handler return, so a peer cancel can't interleave.
# Identity guard: don't evict a duplicate id's newer entry.
dctx.close()
self._in_flight.pop(_coerce_id(req.id), None)
key = _coerce_id(req.id)
if (entry := self._in_flight.get(key)) is not None and entry.dctx is dctx:
del self._in_flight[key]
# A write interrupted by cancellation may still have delivered
# (a memory-stream send can hand its item to the receiver and
# still raise), so a started answer write counts as sent below:
# peers drop late responses, while a second answer for one id
# would break JSON-RPC.
answer_write_started = True
await self._write_result(req.id, result)
if scope.cancel_called:
# Peer-cancel: `_dispatch_notification` cancelled this scope
# while the handler was running. anyio swallows a scope's *own*
# cancel at __exit__, so execution lands here rather than the
# `except cancelled` arm below.
# TODO(maxisbey): spec says SHOULD NOT respond after cancel.
# The existing server always has, so match that for now.
if scope.cancelled_caught:
# anyio absorbs the scope's own cancel at __exit__, and
# `cancelled_caught` (unlike `cancel_called`) guarantees the
# result write above did not happen - no double response.
# TODO(maxisbey): spec says SHOULD NOT respond after cancel;
# the existing server always has, so match that for now.
answer_write_started = True
await self._write_error(req.id, ErrorData(code=0, message="Request cancelled"))
except anyio.get_cancelled_exc_class():
# Outer-cancel: run()'s task group is shutting down. Any bare
# `await` here re-raises immediately, so shield the courtesy write.
with anyio.CancelScope(shield=True):
await self._write_error(req.id, ErrorData(code=REQUEST_CANCELLED, message="Request cancelled"))
# Shutdown: answer the request so the peer isn't left waiting - unless
# an answer write already started (it may have reached the transport;
# prefer possibly-zero answers over possibly-two). The shielded helper
# is needed because bare awaits re-raise here.
if not answer_write_started:
await self._final_write(
partial(self._write_error, req.id, ErrorData(code=CONNECTION_CLOSED, message="Connection closed")),
shield=True,
timeout=_SHUTDOWN_WRITE_TIMEOUT,
describe=f"shutdown error response for request {req.id!r}",
)
raise
except MCPError as e:
await self._write_error(req.id, e.error)
except ValidationError:
# TODO(maxisbey): data="" is pinned compat with the existing
# server (which never leaked pydantic error text onto the wire).
# Consider putting the validation detail in `data` once the
# interaction suite's divergence entry is resolved.
# TODO(maxisbey): data="" pins existing-server compat (no pydantic
# text on the wire); revisit per the suite's divergence entry.
await self._write_error(
req.id, ErrorData(code=INVALID_PARAMS, message="Invalid request parameters", data="")
)
except Exception as e:
logger.exception("handler for %r raised", req.method)
# TODO(maxisbey): code=0 is pinned compat with the existing
# server's `_handle_request`. JSON-RPC says INTERNAL_ERROR
# (-32603); revisit once the suite's divergence entry is resolved.
# TODO(maxisbey): code=0 pins existing-server compat; JSON-RPC says
# INTERNAL_ERROR. Revisit per the suite's divergence entry.
await self._write_error(req.id, ErrorData(code=0, message=str(e)))
if self._raise_handler_exceptions:
raise
# No outer `_in_flight` pop here: the inner `finally` above already
# removes the entry on every path out of the handler, and a second
# pop after the awaited response writes could evict a newer request
# that reused the id during that window.
# No `_in_flight` pop here: the inner finally covers every path, and a late pop could evict a reused id.
def _allocate_id(self) -> int:
self._next_id += 1
@ -727,16 +720,31 @@ class JSONRPCDispatcher(Dispatcher[TransportT]):
except (anyio.BrokenResourceError, anyio.ClosedResourceError):
logger.debug("dropped error for %r: write stream closed", request_id)
async def _final_write(
self,
write: Callable[[], Awaitable[None]],
*,
shield: bool,
timeout: float,
describe: str,
) -> None:
"""Attempt one last write under the shared abandon/teardown policy.
`shield=True` is for arms already inside a cancelled scope (a bare
`await` would re-raise); the bound keeps a wedged transport write
from becoming an uncancellable hang.
"""
with anyio.move_on_after(timeout, shield=shield) as scope:
await write()
if scope.cancelled_caught:
logger.warning("%s gave up: transport write blocked", describe)
async def _cancel_outbound(self, request_id: RequestId, reason: str, related_request_id: RequestId | None) -> None:
# Thread `related_request_id` so streamable-HTTP routes the cancel onto
# the same per-request SSE stream as the request it cancels; without it
# the notification falls through to the standalone GET stream and is
# dropped when no GET stream is open.
try:
await self.notify(
"notifications/cancelled",
{"requestId": request_id, "reason": reason},
_related_request_id=related_request_id,
)
except (anyio.BrokenResourceError, anyio.ClosedResourceError):
pass
# Thread `related_request_id` so streamable HTTP routes the cancel onto
# the request's own SSE stream instead of a possibly-absent GET stream.
# `notify` swallows connection-state errors itself, so no guard here.
await self.notify(
"notifications/cancelled",
{"requestId": request_id, "reason": reason},
_related_request_id=related_request_id,
)

490
src/mcp/shared/session.py
View file

@ -1,487 +1,21 @@
from __future__ import annotations
"""Compatibility names that outlived the removed v1 session layer (`BaseSession`)."""
import contextvars
import logging
from contextlib import AsyncExitStack
from types import TracebackType
from typing import Any, Generic, Protocol, TypeVar
from typing import Generic, TypeVar
import anyio
from anyio.streams.memory import MemoryObjectSendStream
from opentelemetry.trace import SpanKind
from pydantic import BaseModel, TypeAdapter
from typing_extensions import Self
from mcp.shared._compat import resync_tracer
from mcp.shared._otel import inject_trace_context, otel_span
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.exceptions import MCPError
from mcp.shared.message import MessageMetadata, ServerMessageMetadata, SessionMessage
from mcp.types import (
CONNECTION_CLOSED,
INVALID_PARAMS,
METHOD_NOT_FOUND,
REQUEST_TIMEOUT,
CancelledNotification,
ClientNotification,
ClientRequest,
ClientResult,
ErrorData,
JSONRPCError,
JSONRPCNotification,
JSONRPCRequest,
JSONRPCResponse,
ProgressNotification,
ProgressToken,
RequestParamsMeta,
ServerNotification,
ServerRequest,
ServerResult,
)
SendRequestT = TypeVar("SendRequestT", ClientRequest, ServerRequest)
SendResultT = TypeVar("SendResultT", ClientResult, ServerResult)
SendNotificationT = TypeVar("SendNotificationT", ClientNotification, ServerNotification)
ReceiveRequestT = TypeVar("ReceiveRequestT", ClientRequest, ServerRequest)
ReceiveResultT = TypeVar("ReceiveResultT", bound=BaseModel)
ReceiveNotificationT = TypeVar("ReceiveNotificationT", ClientNotification, ServerNotification)
from mcp.shared.dispatcher import ProgressFnT as ProgressFnT
from mcp.shared.message import MessageMetadata
from mcp.types import RequestParamsMeta
RequestId = str | int
class ProgressFnT(Protocol):
"""Protocol for progress notification callbacks."""
async def __call__(
self, progress: float, total: float | None, message: str | None
) -> None: ... # pragma: no branch
ReceiveRequestT = TypeVar("ReceiveRequestT")
SendResultT = TypeVar("SendResultT")
class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
"""Handles responding to MCP requests and manages request lifecycle.
"""Typing stub for the v1 responder; the SDK never instantiates it."""
This class MUST be used as a context manager to ensure proper cleanup and
cancellation handling:
Example:
```python
with request_responder as resp:
await resp.respond(result)
```
The context manager ensures:
1. Proper cancellation scope setup and cleanup
2. Request completion tracking
3. Cleanup of in-flight requests
"""
def __init__(
self,
request_id: RequestId,
request_meta: RequestParamsMeta | None,
request: ReceiveRequestT,
session: BaseSession[SendRequestT, SendNotificationT, SendResultT, ReceiveRequestT, ReceiveNotificationT],
message_metadata: MessageMetadata = None,
context: contextvars.Context | None = None,
) -> None:
self.request_id = request_id
self.request_meta = request_meta
self.request = request
self.message_metadata = message_metadata
self.context = context
self._session = session
self._completed = False
self._entered = False # Track if we're in a context manager
def __enter__(self) -> RequestResponder[ReceiveRequestT, SendResultT]:
self._entered = True
return self
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> None:
self._entered = False
async def respond(self, response: SendResultT | ErrorData) -> None:
"""Send a response for this request.
Must be called within a context manager block.
Raises:
RuntimeError: If not used within a context manager
AssertionError: If request was already responded to
"""
if not self._entered: # pragma: no cover
raise RuntimeError("RequestResponder must be used as a context manager")
assert not self._completed, "Request already responded to"
self._completed = True
await self._session._send_response( # type: ignore[reportPrivateUsage]
request_id=self.request_id, response=response
)
class BaseSession(
Generic[
SendRequestT,
SendNotificationT,
SendResultT,
ReceiveRequestT,
ReceiveNotificationT,
],
):
"""Implements an MCP "session" on top of read/write streams, including features
like request/response linking, notifications, and progress.
This class is an async context manager that automatically starts processing
messages when entered.
"""
_response_streams: dict[RequestId, MemoryObjectSendStream[JSONRPCResponse | JSONRPCError]]
_request_id: int
_progress_callbacks: dict[RequestId, ProgressFnT]
def __init__(
self,
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[SessionMessage],
# If none, reading will never time out
read_timeout_seconds: float | None = None,
) -> None:
self._read_stream = read_stream
self._write_stream = write_stream
self._response_streams = {}
self._request_id = 0
self._session_read_timeout_seconds = read_timeout_seconds
self._progress_callbacks = {}
self._exit_stack = AsyncExitStack()
async def __aenter__(self) -> Self:
self._task_group = anyio.create_task_group()
await self._task_group.__aenter__()
self._task_group.start_soon(self._receive_loop)
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> bool | None:
await self._exit_stack.aclose()
# Using BaseSession as a context manager should not block on exit (this
# would be very surprising behavior), so make sure to cancel the tasks
# in the task group.
self._task_group.cancel_scope.cancel()
result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
await resync_tracer()
return result
async def send_request(
self,
request: SendRequestT,
result_type: type[ReceiveResultT],
request_read_timeout_seconds: float | None = None,
metadata: MessageMetadata = None,
progress_callback: ProgressFnT | None = None,
) -> ReceiveResultT:
"""Sends a request and waits for a response.
Raises an MCPError if the response contains an error. If a request read timeout is provided, it will take
precedence over the session read timeout.
Do not use this method to emit notifications! Use send_notification() instead.
"""
request_id = self._request_id
self._request_id = request_id + 1
response_stream, response_stream_reader = anyio.create_memory_object_stream[JSONRPCResponse | JSONRPCError](1)
self._response_streams[request_id] = response_stream
# Set up progress token if progress callback is provided
request_data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
if progress_callback is not None:
# Use request_id as progress token
if "params" not in request_data: # pragma: lax no cover
request_data["params"] = {}
if "_meta" not in request_data["params"]: # pragma: lax no cover
request_data["params"]["_meta"] = {}
request_data["params"]["_meta"]["progressToken"] = request_id
# Store the callback for this request
self._progress_callbacks[request_id] = progress_callback
try:
target = request_data.get("params", {}).get("name")
span_name = f"MCP send {request.method} {target}" if target else f"MCP send {request.method}"
with otel_span(
span_name,
kind=SpanKind.CLIENT,
attributes={"mcp.method.name": request.method, "jsonrpc.request.id": str(request_id)},
):
# Inject W3C trace context into _meta (SEP-414).
meta: dict[str, Any] = request_data.setdefault("params", {}).setdefault("_meta", {})
inject_trace_context(meta)
jsonrpc_request = JSONRPCRequest(jsonrpc="2.0", id=request_id, **request_data)
await self._write_stream.send(SessionMessage(message=jsonrpc_request, metadata=metadata))
# request read timeout takes precedence over session read timeout
timeout = request_read_timeout_seconds or self._session_read_timeout_seconds
try:
with anyio.fail_after(timeout):
response_or_error = await response_stream_reader.receive()
except TimeoutError:
class_name = request.__class__.__name__
message = f"Timed out while waiting for response to {class_name}. Waited {timeout} seconds."
raise MCPError(code=REQUEST_TIMEOUT, message=message)
if isinstance(response_or_error, JSONRPCError):
raise MCPError.from_jsonrpc_error(response_or_error)
else:
return result_type.model_validate(response_or_error.result, by_name=False)
finally:
self._response_streams.pop(request_id, None)
self._progress_callbacks.pop(request_id, None)
await response_stream.aclose()
await response_stream_reader.aclose()
async def send_notification(
self,
notification: SendNotificationT,
related_request_id: RequestId | None = None,
) -> None:
"""Emits a notification, which is a one-way message that does not expect a response."""
# Some transport implementations may need to set the related_request_id
# to attribute to the notifications to the request that triggered them.
jsonrpc_notification = JSONRPCNotification(
jsonrpc="2.0",
**notification.model_dump(by_alias=True, mode="json", exclude_none=True),
)
session_message = SessionMessage(
message=jsonrpc_notification,
metadata=ServerMessageMetadata(related_request_id=related_request_id) if related_request_id else None,
)
await self._write_stream.send(session_message)
async def _send_response(self, request_id: RequestId, response: SendResultT | ErrorData) -> None:
if isinstance(response, ErrorData):
jsonrpc_error = JSONRPCError(jsonrpc="2.0", id=request_id, error=response)
session_message = SessionMessage(message=jsonrpc_error)
await self._write_stream.send(session_message)
else:
jsonrpc_response = JSONRPCResponse(
jsonrpc="2.0",
id=request_id,
result=response.model_dump(by_alias=True, mode="json", exclude_none=True),
)
session_message = SessionMessage(message=jsonrpc_response)
await self._write_stream.send(session_message)
@property
def _receive_request_adapter(self) -> TypeAdapter[ReceiveRequestT]:
"""Each subclass must provide its own request adapter."""
raise NotImplementedError
@property
def _receive_request_methods(self) -> frozenset[str]:
"""Method names in the receive-request union; anything else is
answered with METHOD_NOT_FOUND before validation is attempted."""
raise NotImplementedError
@property
def _receive_notification_adapter(self) -> TypeAdapter[ReceiveNotificationT]:
raise NotImplementedError
async def _receive_loop(self) -> None:
async with self._read_stream, self._write_stream:
try:
async def _handle_session_message(message: SessionMessage) -> None:
sender_context: contextvars.Context | None = getattr(self._read_stream, "last_context", None)
if isinstance(message.message, JSONRPCRequest):
if message.message.method not in self._receive_request_methods:
# Unknown methods are METHOD_NOT_FOUND (-32601) per
# JSON-RPC 2.0, not validation failures (-32602).
error_response = JSONRPCError(
jsonrpc="2.0",
id=message.message.id,
error=ErrorData(
code=METHOD_NOT_FOUND, message="Method not found", data=message.message.method
),
)
await self._write_stream.send(SessionMessage(message=error_response))
return
try:
validated_request = self._receive_request_adapter.validate_python(
message.message.model_dump(by_alias=True, mode="json", exclude_none=True),
by_name=False,
)
responder = RequestResponder(
request_id=message.message.id,
request_meta=validated_request.params.meta if validated_request.params else None,
request=validated_request,
session=self,
message_metadata=message.metadata,
context=sender_context,
)
await self._received_request(responder)
except Exception:
# For request validation errors, send a proper JSON-RPC error
# response instead of crashing the server
logging.warning("Failed to validate request", exc_info=True)
logging.debug(f"Message that failed validation: {message.message}")
error_response = JSONRPCError(
jsonrpc="2.0",
id=message.message.id,
error=ErrorData(code=INVALID_PARAMS, message="Invalid request parameters", data=""),
)
session_message = SessionMessage(message=error_response)
await self._write_stream.send(session_message)
elif isinstance(message.message, JSONRPCNotification):
try:
notification = self._receive_notification_adapter.validate_python(
message.message.model_dump(by_alias=True, mode="json", exclude_none=True),
by_name=False,
)
if isinstance(notification, CancelledNotification):
# ClientSession runs server-initiated requests
# inline in this loop, so by the time a peer
# cancellation is read there is nothing left to
# cancel. Consume it here so message_handler
# keeps the contract it had before the
# dispatcher swap removed _in_flight.
return
# Handle progress notifications callback
if isinstance(notification, ProgressNotification):
progress_token = notification.params.progress_token
# If there is a progress callback for this token,
# call it with the progress information
if progress_token in self._progress_callbacks:
callback = self._progress_callbacks[progress_token]
try:
await callback(
notification.params.progress,
notification.params.total,
notification.params.message,
)
except Exception:
logging.exception("Progress callback raised an exception")
await self._received_notification(notification)
await self._handle_incoming(notification)
except Exception:
# For other validation errors, log and continue
logging.warning(
"Failed to validate notification: %s",
message.message,
exc_info=True,
)
else: # Response or error
await self._handle_response(message)
async for message in self._read_stream:
if isinstance(message, Exception):
await self._handle_incoming(message)
continue
await _handle_session_message(message)
except anyio.ClosedResourceError:
# This is expected when the client disconnects abruptly.
# Without this handler, the exception would propagate up and
# crash the server's task group.
logging.debug("Read stream closed by client")
except Exception as e:
# Other exceptions are not expected and should be logged. We purposefully
# catch all exceptions here to avoid crashing the server.
logging.exception(f"Unhandled exception in receive loop: {e}") # pragma: no cover
finally:
# after the read stream is closed, we need to send errors
# to any pending requests
# Snapshot: stream.send() wakes the waiter, whose finally pops
# from _response_streams before the next __next__() call.
for id, stream in list(self._response_streams.items()):
error = ErrorData(code=CONNECTION_CLOSED, message="Connection closed")
try:
await stream.send(JSONRPCError(jsonrpc="2.0", id=id, error=error))
await stream.aclose()
except Exception: # pragma: lax no cover
# Stream might already be closed
pass
self._response_streams.clear()
def _normalize_request_id(self, response_id: RequestId) -> RequestId:
"""Normalize a response ID to match how request IDs are stored.
Since the client always sends integer IDs, we normalize string IDs
to integers when possible. This matches the TypeScript SDK approach:
https://github.com/modelcontextprotocol/typescript-sdk/blob/a606fb17909ea454e83aab14c73f14ea45c04448/src/shared/protocol.ts#L861
Args:
response_id: The response ID from the incoming message.
Returns:
The normalized ID (int if possible, otherwise original value).
"""
if isinstance(response_id, str):
try:
return int(response_id)
except ValueError:
logging.warning(f"Response ID {response_id!r} cannot be normalized to match pending requests")
return response_id
async def _handle_response(self, message: SessionMessage) -> None:
"""Handle an incoming response or error message."""
# This check is always true at runtime: the caller (_receive_loop) only invokes
# this method in the else branch after checking for JSONRPCRequest and
# JSONRPCNotification. However, the type checker can't infer this from the
# method signature, so we need this guard for type narrowing.
if not isinstance(message.message, JSONRPCResponse | JSONRPCError):
return # pragma: no cover
if message.message.id is None:
# Narrows to JSONRPCError since JSONRPCResponse.id is always RequestId
error = message.message.error
logging.warning(f"Received error with null ID: {error.message}")
await self._handle_incoming(MCPError(error.code, error.message, error.data))
return
# Normalize response ID to handle type mismatches (e.g., "0" vs 0)
response_id = self._normalize_request_id(message.message.id)
stream = self._response_streams.pop(response_id, None)
if stream:
await stream.send(message.message)
else:
await self._handle_incoming(RuntimeError(f"Received response with an unknown request ID: {message}"))
async def _received_request(self, responder: RequestResponder[ReceiveRequestT, SendResultT]) -> None:
"""Can be overridden by subclasses to handle a request without needing to
listen on the message stream.
If the request is responded to within this method, it will not be
forwarded on to the message stream.
"""
async def _received_notification(self, notification: ReceiveNotificationT) -> None:
"""Can be overridden by subclasses to handle a notification without needing
to listen on the message stream.
"""
async def send_progress_notification(
self,
progress_token: ProgressToken,
progress: float,
total: float | None = None,
message: str | None = None,
) -> None:
"""Sends a progress notification for a request that is currently being processed."""
async def _handle_incoming(
self, req: RequestResponder[ReceiveRequestT, SendResultT] | ReceiveNotificationT | Exception
) -> None:
"""A generic handler for incoming messages. Overridden by subclasses."""
request_id: RequestId
request_meta: RequestParamsMeta | None
request: ReceiveRequestT
message_metadata: MessageMetadata

2
src/mcp/types/__init__.py
View file

@ -152,7 +152,6 @@ from mcp.types.jsonrpc import (
INVALID_REQUEST,
METHOD_NOT_FOUND,
PARSE_ERROR,
REQUEST_CANCELLED,
REQUEST_TIMEOUT,
URL_ELICITATION_REQUIRED,
ErrorData,
@ -320,7 +319,6 @@ __all__ = [
"INVALID_REQUEST",
"METHOD_NOT_FOUND",
"PARSE_ERROR",
"REQUEST_CANCELLED",
"REQUEST_TIMEOUT",
"URL_ELICITATION_REQUIRED",
"ErrorData",

1
src/mcp/types/jsonrpc.py
View file

@ -43,7 +43,6 @@ URL_ELICITATION_REQUIRED = -32042
# SDK error codes
CONNECTION_CLOSED = -32000
REQUEST_TIMEOUT = -32001
REQUEST_CANCELLED = -32002
# Standard JSON-RPC error codes
PARSE_ERROR = -32700