[v2] Dispatcher/ServerRunner receive-path swap — replaces BaseSession (#2710)
This commit is contained in:
Max
GitHub
parent
b478bff56d
commit
5d826490b6
54 changed files with 6594 additions and 1701 deletions
1
.github/workflows/main.yml
vendored
1
.github/workflows/main.yml
vendored
|
|
@ -5,7 +5,6 @@ on:
|
|||
branches: ["main", "v1.x"]
|
||||
tags: ["v*.*.*"]
|
||||
pull_request:
|
||||
branches: ["main", "v1.x"]
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
|
|||
|
|
@ -472,7 +472,7 @@ async def my_tool(x: int, ctx: Context) -> str:
|
|||
|
||||
The internal layers (`ToolManager.call_tool`, `Tool.run`, `Prompt.render`, `ResourceTemplate.create_resource`, etc.) now require `context` as a positional argument.
|
||||
|
||||
### Registering lowlevel handlers on `MCPServer` (workaround)
|
||||
### Registering lowlevel handlers from `MCPServer`
|
||||
|
||||
`MCPServer` does not expose public APIs for `subscribe_resource`, `unsubscribe_resource`, or `set_logging_level` handlers. In v1, the workaround was to reach into the private lowlevel server and use its decorator methods:
|
||||
|
||||
|
|
@ -486,7 +486,7 @@ async def handle_set_logging_level(level: str) -> None:
|
|||
mcp._mcp_server.subscribe_resource()(handle_subscribe) # pyright: ignore[reportPrivateUsage]
|
||||
```
|
||||
|
||||
In v2, the lowlevel `Server` no longer has decorator methods (handlers are constructor-only), so the equivalent workaround is `_add_request_handler`:
|
||||
In v2, the lowlevel `Server` supports arbitrary request handlers directly via `add_request_handler` (the decorator methods are gone; handlers are otherwise constructor-only). From `MCPServer`, access it via `_lowlevel_server`:
|
||||
|
||||
**After (v2):**
|
||||
|
||||
|
|
@ -505,11 +505,11 @@ async def handle_subscribe(ctx: ServerRequestContext, params: SubscribeRequestPa
|
|||
return EmptyResult()
|
||||
|
||||
|
||||
mcp._lowlevel_server._add_request_handler("logging/setLevel", handle_set_logging_level) # pyright: ignore[reportPrivateUsage]
|
||||
mcp._lowlevel_server._add_request_handler("resources/subscribe", handle_subscribe) # pyright: ignore[reportPrivateUsage]
|
||||
mcp._lowlevel_server.add_request_handler("logging/setLevel", SetLevelRequestParams, handle_set_logging_level) # pyright: ignore[reportPrivateUsage]
|
||||
mcp._lowlevel_server.add_request_handler("resources/subscribe", SubscribeRequestParams, handle_subscribe) # pyright: ignore[reportPrivateUsage]
|
||||
```
|
||||
|
||||
This is a private API and may change. A public way to register these handlers on `MCPServer` is planned; until then, use this workaround or use the lowlevel `Server` directly.
|
||||
`_lowlevel_server` is private and may change. A public way to register these handlers on `MCPServer` is planned; until then, use this workaround or use the lowlevel `Server` directly.
|
||||
|
||||
### `MCPServer`'s `Context` logging: `message` renamed to `data`, `extra` removed
|
||||
|
||||
|
|
@ -664,6 +664,8 @@ ctx: ClientRequestContext
|
|||
server_ctx: ServerRequestContext[LifespanContextT, RequestT]
|
||||
```
|
||||
|
||||
`ServerRequestContext` is now a standalone dataclass — it no longer subclasses `RequestContext[ServerSession]`. It carries the same fields (`session`, `request_id`, `meta`, `lifespan_context`, `request`, `close_sse_stream`, `close_standalone_sse_stream`), so handler code is unaffected, but `isinstance(ctx, RequestContext)` checks and `RequestContext[ServerSession]` annotations need updating to `ServerRequestContext`.
|
||||
|
||||
The high-level `Context` class (injected into `@mcp.tool()` etc.) similarly dropped its `ServerSessionT` parameter: `Context[ServerSessionT, LifespanContextT, RequestT]` → `Context[LifespanContextT, RequestT]`. Both remaining parameters have defaults, so bare `Context` is usually sufficient:
|
||||
|
||||
**Before (v1):**
|
||||
|
|
@ -857,6 +859,54 @@ server = Server("my-server", on_list_tools=handle_list_tools)
|
|||
|
||||
If you need to check whether a handler is registered, track this yourself — there is currently no public introspection API.
|
||||
|
||||
### Lowlevel `Server`: `add_request_handler` is now public and takes `params_type`
|
||||
|
||||
The private `_add_request_handler(method, handler)` escape hatch is now the public `add_request_handler(method, params_type, handler)`, alongside a matching `add_notification_handler`. Each takes a `params_type` model that incoming params are validated against before the handler runs. A message with no `params` member validates `{}` against the model, so handlers never receive `None`: all-optional models arrive with their defaults, and models with required fields reject the message as `INVALID_PARAMS` before the handler runs (matching the Go SDK).
|
||||
|
||||
```python
|
||||
# Before (v1 / earlier v2 prereleases)
|
||||
server._add_request_handler("custom/method", my_handler)
|
||||
|
||||
# After (v2)
|
||||
server.add_request_handler("custom/method", MyParams, my_handler)
|
||||
server.add_notification_handler("notifications/custom", MyNotifyParams, my_notify_handler)
|
||||
```
|
||||
|
||||
### Lowlevel `Server`: private `_handle_*` dispatch methods removed
|
||||
|
||||
`Server._handle_message`, `_handle_request`, and `_handle_notification` have been removed. The receive loop and per-message dispatch now live in `JSONRPCDispatcher` and `ServerRunner`, which `Server.run()` drives internally.
|
||||
|
||||
These were private, but some users subclassed `Server` and overrode them to intercept requests. Use middleware instead:
|
||||
|
||||
```python
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
from mcp.server import Server, ServerRequestContext
|
||||
from mcp.server.context import CallNext, HandlerResult
|
||||
|
||||
|
||||
async def logging_middleware(
|
||||
ctx: ServerRequestContext[Any, Any], method: str, params: Mapping[str, Any] | None, call_next: CallNext
|
||||
) -> HandlerResult:
|
||||
print(f"handling {method}")
|
||||
result = await call_next()
|
||||
print(f"done {method}")
|
||||
return result
|
||||
|
||||
|
||||
server = Server("my-server", on_call_tool=...)
|
||||
server.middleware.append(logging_middleware)
|
||||
```
|
||||
|
||||
Middleware runs before params validation, so `params` is the raw inbound mapping (or `None`), and it also wraps unknown methods.
|
||||
|
||||
### Lowlevel `Server.run(raise_exceptions=True)`: transport errors no longer re-raised
|
||||
|
||||
`raise_exceptions=True` now only governs handler exceptions: an exception raised by an `on_*` handler propagates out of `run()`. The JSON-RPC error response is still written to the client first, regardless of the flag.
|
||||
|
||||
Previously it also re-raised exceptions yielded by the transport onto the read stream (e.g. JSON parse errors). Those are now debug-logged and dropped regardless of `raise_exceptions`. If you relied on `run()` exiting on a transport-level parse error, that no longer happens.
|
||||
|
||||
### Lowlevel `Server`: decorator-based handlers replaced with constructor `on_*` params
|
||||
|
||||
The lowlevel `Server` class no longer uses decorator methods for handler registration. Instead, handlers are passed as `on_*` keyword arguments to the constructor.
|
||||
|
|
@ -1083,6 +1133,39 @@ from mcp.server import ServerRequestContext
|
|||
# but None in notification handlers
|
||||
```
|
||||
|
||||
### `ServerSession` is now a thin proxy (no longer a `BaseSession`)
|
||||
|
||||
`ServerSession` no longer subclasses `BaseSession`. It is now a small connection-scoped proxy that exposes `send_request`, `send_notification`, the typed convenience helpers (`create_message`, `elicit_form`, `send_log_message`, `send_tool_list_changed`, ...), `client_params`, and `check_client_capability`. The receive loop, `initialize` handling, and per-request task isolation that previously lived in `ServerSession` have moved to `JSONRPCDispatcher` and `ServerRunner`.
|
||||
|
||||
`ServerSession` is normally constructed for you by `Server.run()` and reached via `ctx.session` in handlers, so most servers are unaffected. If you were constructing or subclassing it directly:
|
||||
|
||||
**Constructor change:**
|
||||
|
||||
```python
|
||||
# Before (v1)
|
||||
session = ServerSession(read_stream, write_stream, init_options, stateless=False)
|
||||
|
||||
# After (v2)
|
||||
session = ServerSession(dispatcher, connection, stateless=False)
|
||||
# where `dispatcher` is a JSONRPCDispatcher and `connection` is a Connection
|
||||
```
|
||||
|
||||
In practice, replace direct `ServerSession` use with `Server.run(read_stream, write_stream, init_options)` and let the framework wire it up.
|
||||
|
||||
**Removed from `mcp.server.session`:**
|
||||
|
||||
- `InitializationState` enum and `ServerSession._initialization_state` — initialization tracking is now on `Connection` (`connection.initialized` is an `anyio.Event`, `connection.client_params` holds the init params).
|
||||
- `ServerRequestResponder` type alias.
|
||||
- `ServerSession.incoming_messages` stream — there is no longer a public stream of inbound messages to iterate. Register handlers via the `on_*` constructor params (or `add_request_handler`) and use `Server.middleware` to observe every inbound request and notification (`initialize`, unknown methods, validation failures, and `notifications/initialized` included).
|
||||
- `ServerSession.__aenter__` / `__aexit__` — `ServerSession` is no longer an async context manager.
|
||||
- The private `_receive_loop`, `_received_request`, `_received_notification`, and `_handle_incoming` overrides — there is nothing to override on `ServerSession` anymore. To intercept inbound messages, use `Server.middleware` (see the `_handle_*` removal section above).
|
||||
|
||||
### `BaseSession` / `RequestResponder`: server-side cancellation tracking removed
|
||||
|
||||
`BaseSession._in_flight` and the `RequestResponder` members that supported it (`cancel()`, the `cancelled` and `in_flight` properties, the `on_complete` constructor argument, and the internal `CancelScope`) have been removed. These existed to let `ServerSession` cancel a handler when a `CancelledNotification` arrived; `ServerSession` no longer drives a receive loop, so they were dead code. Inbound-cancellation handling for the server now lives in `JSONRPCDispatcher`.
|
||||
|
||||
`BaseSession` is still used by `ClientSession`, which never relied on these members. `RequestResponder.respond()` is unchanged.
|
||||
|
||||
### Experimental Tasks support removed
|
||||
|
||||
Tasks (SEP-1686) have been removed from the MCP specification and are no longer part of this SDK. The `mcp.client.experimental`, `mcp.server.experimental`, `mcp.shared.experimental`, and `mcp.server.lowlevel.experimental` modules have been removed, along with all `Task*` types, the `tasks` capability fields, `Tool.execution`, and the `experimental` properties on `ClientSession`, `ServerSession`, `Server`, and `ServerRequestContext`.
|
||||
|
|
|
|||
|
|
@ -417,9 +417,15 @@ async def handle_unsubscribe(ctx: ServerRequestContext, params: UnsubscribeReque
|
|||
return EmptyResult()
|
||||
|
||||
|
||||
mcp._lowlevel_server._add_request_handler("logging/setLevel", handle_set_logging_level) # pyright: ignore[reportPrivateUsage]
|
||||
mcp._lowlevel_server._add_request_handler("resources/subscribe", handle_subscribe) # pyright: ignore[reportPrivateUsage]
|
||||
mcp._lowlevel_server._add_request_handler("resources/unsubscribe", handle_unsubscribe) # pyright: ignore[reportPrivateUsage]
|
||||
mcp._lowlevel_server.add_request_handler( # pyright: ignore[reportPrivateUsage]
|
||||
"logging/setLevel", SetLevelRequestParams, handle_set_logging_level
|
||||
)
|
||||
mcp._lowlevel_server.add_request_handler( # pyright: ignore[reportPrivateUsage]
|
||||
"resources/subscribe", SubscribeRequestParams, handle_subscribe
|
||||
)
|
||||
mcp._lowlevel_server.add_request_handler( # pyright: ignore[reportPrivateUsage]
|
||||
"resources/unsubscribe", UnsubscribeRequestParams, handle_unsubscribe
|
||||
)
|
||||
|
||||
|
||||
@mcp.completion()
|
||||
|
|
|
|||
|
|
@ -25,7 +25,11 @@ classifiers = [
|
|||
"Programming Language :: Python :: 3.14",
|
||||
]
|
||||
dependencies = [
|
||||
"anyio>=4.9",
|
||||
# anyio < 4.10 triggers a compile-time SyntaxWarning on Python 3.14 (PEP 765,
|
||||
# "'return' in a 'finally' block"); for stdio servers it lands on the child's
|
||||
# stderr (agronholm/anyio#816, fixed in 4.10).
|
||||
"anyio>=4.10; python_version >= '3.14'",
|
||||
"anyio>=4.9; python_version < '3.14'",
|
||||
"httpx>=0.27.1,<1.0.0",
|
||||
"httpx-sse>=0.4",
|
||||
"pydantic>=2.12.0",
|
||||
|
|
@ -90,6 +94,7 @@ dev = [
|
|||
"pillow>=12.0",
|
||||
"strict-no-cover",
|
||||
"logfire>=3.0.0",
|
||||
"opentelemetry-sdk>=1.39.1",
|
||||
]
|
||||
docs = [
|
||||
"mkdocs>=1.6.1",
|
||||
|
|
|
|||
|
|
@ -14,6 +14,9 @@ from mcp.server import Server
|
|||
from mcp.server.mcpserver import MCPServer
|
||||
from mcp.shared.memory import create_client_server_memory_streams
|
||||
|
||||
SERVER_SHUTDOWN_GRACE = 2.0
|
||||
"""Seconds to wait for the in-process server to exit on EOF before cancelling."""
|
||||
|
||||
|
||||
class InMemoryTransport:
|
||||
"""In-memory transport for testing MCP servers without network overhead.
|
||||
|
|
@ -48,21 +51,49 @@ class InMemoryTransport:
|
|||
client_read, client_write = client_streams
|
||||
server_read, server_write = server_streams
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
# Start server in background
|
||||
tg.start_soon(
|
||||
lambda: actual_server.run(
|
||||
server_done = anyio.Event()
|
||||
|
||||
async def _run_server() -> None:
|
||||
try:
|
||||
await actual_server.run(
|
||||
server_read,
|
||||
server_write,
|
||||
actual_server.create_initialization_options(),
|
||||
raise_exceptions=self._raise_exceptions,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
server_done.set()
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(_run_server)
|
||||
|
||||
try:
|
||||
yield client_read, client_write
|
||||
finally:
|
||||
tg.cancel_scope.cancel()
|
||||
# EOF the server (and our own read side) instead of
|
||||
# cancelling outright. The dispatcher's run() cancels its
|
||||
# own in-flight handlers on read-stream EOF, so for a
|
||||
# well-behaved server the task exits naturally and the
|
||||
# task-group join below is immediate. Cancelling here
|
||||
# unconditionally would `coro.throw()` into this task,
|
||||
# which on CPython 3.11 (gh-106749) drops `'call'` trace
|
||||
# events for the outer await chain and desyncs coverage's
|
||||
# CTracer past the test frame.
|
||||
await client_write.aclose()
|
||||
await server_write.aclose()
|
||||
# Backstop: the dispatcher exits on EOF, but the server's
|
||||
# own teardown (lifespan __aexit__, connection.exit_stack
|
||||
# callbacks) runs after that and is user code. If it never
|
||||
# completes the join would hang forever, so bound the wait
|
||||
# and fall back to cancelling. The healthy path returns
|
||||
# from wait() without the timeout firing, so the cancel is
|
||||
# never reached and gh-106749 stays avoided. If the cancel
|
||||
# does fire, the checkpoint at the end of
|
||||
# `create_client_server_memory_streams` resyncs the tracer.
|
||||
with anyio.move_on_after(SERVER_SHUTDOWN_GRACE):
|
||||
await server_done.wait()
|
||||
if not server_done.is_set():
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
async def __aenter__(self) -> TransportStreams:
|
||||
"""Connect to the server and return streams for communication."""
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from anyio.abc import TaskStatus
|
|||
from httpx_sse import SSEError, aconnect_sse
|
||||
|
||||
from mcp import types
|
||||
from mcp.shared._compat import resync_tracer
|
||||
from mcp.shared._context_streams import create_context_streams
|
||||
from mcp.shared._httpx_utils import McpHttpClientFactory, create_mcp_http_client
|
||||
from mcp.shared.message import SessionMessage
|
||||
|
|
@ -157,3 +158,4 @@ async def sse_client(
|
|||
|
||||
yield read_stream, write_stream
|
||||
tg.cancel_scope.cancel()
|
||||
await resync_tracer()
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from httpx_sse import EventSource, ServerSentEvent, aconnect_sse
|
|||
from pydantic import ValidationError
|
||||
|
||||
from mcp.client._transport import TransportStreams
|
||||
from mcp.shared._compat import resync_tracer
|
||||
from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams
|
||||
from mcp.shared._httpx_utils import create_mcp_http_client
|
||||
from mcp.shared.message import ClientMessageMetadata, SessionMessage
|
||||
|
|
@ -586,3 +587,4 @@ async def streamable_http_client(
|
|||
if transport.session_id and terminate_on_close:
|
||||
await transport.terminate_session(client)
|
||||
tg.cancel_scope.cancel()
|
||||
await resync_tracer()
|
||||
|
|
|
|||
|
|
@ -1,14 +1,11 @@
|
|||
import importlib.metadata
|
||||
import logging
|
||||
import sys
|
||||
import warnings
|
||||
|
||||
import anyio
|
||||
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.server.lowlevel.server import Server
|
||||
from mcp.server.stdio import stdio_server
|
||||
from mcp.types import ServerCapabilities
|
||||
|
||||
if not sys.warnoptions:
|
||||
warnings.simplefilter("ignore")
|
||||
|
|
@ -17,32 +14,10 @@ logging.basicConfig(level=logging.INFO)
|
|||
logger = logging.getLogger("server")
|
||||
|
||||
|
||||
async def receive_loop(session: ServerSession):
|
||||
logger.info("Starting receive loop")
|
||||
async for message in session.incoming_messages:
|
||||
if isinstance(message, Exception):
|
||||
logger.error("Error: %s", message)
|
||||
continue
|
||||
|
||||
logger.info("Received message from client: %s", message)
|
||||
|
||||
|
||||
async def main():
|
||||
version = importlib.metadata.version("mcp")
|
||||
async def main() -> None:
|
||||
server: Server[dict[str, object]] = Server("mcp")
|
||||
async with stdio_server() as (read_stream, write_stream):
|
||||
async with (
|
||||
ServerSession(
|
||||
read_stream,
|
||||
write_stream,
|
||||
InitializationOptions(
|
||||
server_name="mcp",
|
||||
server_version=version,
|
||||
capabilities=ServerCapabilities(),
|
||||
),
|
||||
) as session,
|
||||
write_stream,
|
||||
):
|
||||
await receive_loop(session)
|
||||
await server.run(read_stream, write_stream, server.create_initialization_options())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
249
src/mcp/server/connection.py
Normal file
249
src/mcp/server/connection.py
Normal file
|
|
@ -0,0 +1,249 @@
|
|||
"""`Connection` - per-client connection state and the standalone outbound channel.
|
||||
|
||||
Always present on `Context` (never `None`), even in stateless deployments.
|
||||
Holds peer info populated at `initialize` time, per-connection scratch
|
||||
`state` and an `exit_stack` for teardown, and an `Outbound` for the
|
||||
standalone stream (the SSE GET stream in streamable HTTP, or the single duplex
|
||||
stream in stdio).
|
||||
|
||||
`notify` is best-effort: it never raises. If there's no standalone channel
|
||||
(stateless HTTP) or the stream has been dropped, the notification is
|
||||
debug-logged and silently discarded - server-initiated notifications are
|
||||
inherently advisory. `send_raw_request` *does* raise `NoBackChannelError` when
|
||||
there's no channel; `ping` is the only spec-sanctioned standalone request.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from contextlib import AsyncExitStack
|
||||
from typing import Any, TypeVar, overload
|
||||
|
||||
import anyio
|
||||
from pydantic import BaseModel
|
||||
|
||||
from mcp.shared.dispatcher import CallOptions, Outbound
|
||||
from mcp.shared.exceptions import NoBackChannelError
|
||||
from mcp.shared.peer import Meta, dump_params
|
||||
from mcp.types import (
|
||||
ClientCapabilities,
|
||||
CreateMessageRequest,
|
||||
CreateMessageResult,
|
||||
ElicitRequest,
|
||||
ElicitResult,
|
||||
EmptyResult,
|
||||
InitializeRequestParams,
|
||||
ListRootsRequest,
|
||||
ListRootsResult,
|
||||
LoggingLevel,
|
||||
PingRequest,
|
||||
Request,
|
||||
)
|
||||
|
||||
__all__ = ["Connection"]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
ResultT = TypeVar("ResultT", bound=BaseModel)
|
||||
|
||||
# Result types for the spec's server-to-client request set, used by
|
||||
# `Connection.send_request` to infer the result type. If the spec's request
|
||||
# set grows substantially, consider declaring the result mapping on the
|
||||
# request types themselves (a `__mcp_result__` ClassVar read via a structural
|
||||
# protocol) so this table and the overload ladder don't need maintaining.
|
||||
_RESULT_FOR: dict[type[Request[Any, Any]], type[BaseModel]] = {
|
||||
CreateMessageRequest: CreateMessageResult,
|
||||
ElicitRequest: ElicitResult,
|
||||
ListRootsRequest: ListRootsResult,
|
||||
PingRequest: EmptyResult,
|
||||
}
|
||||
|
||||
|
||||
def _notification_params(payload: dict[str, Any] | None, meta: Meta | None) -> dict[str, Any] | None:
|
||||
if not meta:
|
||||
return payload
|
||||
out = dict(payload or {})
|
||||
out["_meta"] = meta
|
||||
return out
|
||||
|
||||
|
||||
class Connection:
|
||||
"""Per-client connection state and standalone-stream `Outbound`.
|
||||
|
||||
Constructed by `ServerRunner` once per connection. The peer-info fields
|
||||
are `None` until `initialize` completes; `initialized` is set later, when
|
||||
the client's `notifications/initialized` follow-up arrives. In stateless
|
||||
deployments the runner sets `initialized` immediately and peer-info
|
||||
remains `None` (no handshake reaches a stateless connection).
|
||||
"""
|
||||
|
||||
has_standalone_channel: bool
|
||||
session_id: str | None
|
||||
|
||||
client_params: InitializeRequestParams | None
|
||||
"""The full `initialize` request params; `None` before initialization."""
|
||||
|
||||
protocol_version: str | None
|
||||
|
||||
initialized: anyio.Event
|
||||
"""Set when `notifications/initialized` arrives (matches TS `oninitialized`);
|
||||
the point from which the spec permits server-initiated requests beyond
|
||||
ping/logging. Pre-set on stateless connections."""
|
||||
|
||||
state: dict[str, Any]
|
||||
"""Per-connection scratch state; persists across requests on this connection."""
|
||||
|
||||
exit_stack: AsyncExitStack
|
||||
"""Per-connection teardown, unwound LIFO (shielded) when the connection
|
||||
closes. Push cleanup from handlers or middleware; exceptions are logged
|
||||
and swallowed."""
|
||||
|
||||
def __init__(self, outbound: Outbound, *, has_standalone_channel: bool, session_id: str | None = None) -> None:
|
||||
self._outbound = outbound
|
||||
self.has_standalone_channel = has_standalone_channel
|
||||
self.session_id = session_id
|
||||
|
||||
self.client_params = None
|
||||
self.protocol_version = None
|
||||
self.initialized = anyio.Event()
|
||||
|
||||
self.state = {}
|
||||
|
||||
self.exit_stack = AsyncExitStack()
|
||||
|
||||
@property
|
||||
def initialize_accepted(self) -> bool:
|
||||
"""True once the inbound request gate is open: `initialize` recorded the
|
||||
peer info, or the handshake completed outright (stateless birth, or a
|
||||
bare `notifications/initialized`). Derived, never stored."""
|
||||
return self.client_params is not None or self.initialized.is_set()
|
||||
|
||||
async def send_raw_request(
|
||||
self,
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
opts: CallOptions | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Send a raw request on the standalone stream.
|
||||
|
||||
Low-level `Outbound` channel. Prefer the typed `send_request` or the
|
||||
convenience methods below; use this directly only for off-spec
|
||||
messages. `opts` carries per-call `timeout` / `on_progress` /
|
||||
resumption hints; see `CallOptions`.
|
||||
|
||||
Raises:
|
||||
MCPError: The peer responded with an error.
|
||||
NoBackChannelError: `has_standalone_channel` is `False`.
|
||||
"""
|
||||
if not self.has_standalone_channel:
|
||||
raise NoBackChannelError(method)
|
||||
return await self._outbound.send_raw_request(method, params, opts)
|
||||
|
||||
@overload
|
||||
async def send_request(
|
||||
self, req: CreateMessageRequest, *, opts: CallOptions | None = None
|
||||
) -> CreateMessageResult: ...
|
||||
@overload
|
||||
async def send_request(self, req: ElicitRequest, *, opts: CallOptions | None = None) -> ElicitResult: ...
|
||||
@overload
|
||||
async def send_request(self, req: ListRootsRequest, *, opts: CallOptions | None = None) -> ListRootsResult: ...
|
||||
@overload
|
||||
async def send_request(self, req: PingRequest, *, opts: CallOptions | None = None) -> EmptyResult: ...
|
||||
@overload
|
||||
async def send_request(
|
||||
self, req: Request[Any, Any], *, result_type: type[ResultT], opts: CallOptions | None = None
|
||||
) -> ResultT: ...
|
||||
async def send_request(
|
||||
self,
|
||||
req: Request[Any, Any],
|
||||
*,
|
||||
result_type: type[BaseModel] | None = None,
|
||||
opts: CallOptions | None = None,
|
||||
) -> BaseModel:
|
||||
"""Send a typed server-to-client request and return its typed result.
|
||||
|
||||
For spec request types the result type is inferred. For custom requests
|
||||
pass `result_type=` explicitly.
|
||||
|
||||
Raises:
|
||||
MCPError: The peer responded with an error.
|
||||
NoBackChannelError: No back-channel for server-initiated requests.
|
||||
pydantic.ValidationError: The peer's result does not match the expected result type.
|
||||
KeyError: `result_type` omitted for a non-spec request type.
|
||||
"""
|
||||
raw = await self.send_raw_request(req.method, dump_params(req.params), opts)
|
||||
cls = result_type if result_type is not None else _RESULT_FOR[type(req)]
|
||||
return cls.model_validate(raw, by_name=False)
|
||||
|
||||
async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
||||
"""Send a best-effort notification on the standalone stream.
|
||||
|
||||
Never raises. If there's no standalone channel or the stream is broken,
|
||||
the notification is dropped and debug-logged.
|
||||
"""
|
||||
if not self.has_standalone_channel:
|
||||
logger.debug("dropped %s: no standalone channel", method)
|
||||
return
|
||||
try:
|
||||
await self._outbound.notify(method, params)
|
||||
except (anyio.BrokenResourceError, anyio.ClosedResourceError):
|
||||
logger.debug("dropped %s: standalone stream closed", method)
|
||||
|
||||
async def ping(self, *, meta: Meta | None = None, opts: CallOptions | None = None) -> None:
|
||||
"""Send a `ping` request on the standalone stream.
|
||||
|
||||
Raises:
|
||||
MCPError: The peer responded with an error.
|
||||
NoBackChannelError: `has_standalone_channel` is `False`.
|
||||
"""
|
||||
await self.send_raw_request("ping", dump_params(None, meta), opts)
|
||||
|
||||
async def log(self, level: LoggingLevel, data: Any, logger: str | None = None, *, meta: Meta | None = None) -> None:
|
||||
"""Send a `notifications/message` log entry on the standalone stream. Best-effort."""
|
||||
params: dict[str, Any] = {"level": level, "data": data}
|
||||
if logger is not None:
|
||||
params["logger"] = logger
|
||||
await self.notify("notifications/message", _notification_params(params, meta))
|
||||
|
||||
async def send_tool_list_changed(self, *, meta: Meta | None = None) -> None:
|
||||
await self.notify("notifications/tools/list_changed", _notification_params(None, meta))
|
||||
|
||||
async def send_prompt_list_changed(self, *, meta: Meta | None = None) -> None:
|
||||
await self.notify("notifications/prompts/list_changed", _notification_params(None, meta))
|
||||
|
||||
async def send_resource_list_changed(self, *, meta: Meta | None = None) -> None:
|
||||
await self.notify("notifications/resources/list_changed", _notification_params(None, meta))
|
||||
|
||||
async def send_resource_updated(self, uri: str, *, meta: Meta | None = None) -> None:
|
||||
await self.notify("notifications/resources/updated", _notification_params({"uri": uri}, meta))
|
||||
|
||||
def check_capability(self, capability: ClientCapabilities) -> bool:
|
||||
"""Return whether the connected client declared the given capability.
|
||||
|
||||
Returns `False` if `initialize` hasn't completed yet.
|
||||
"""
|
||||
# TODO: redesign - mirrors v1 ServerSession.check_client_capability
|
||||
# verbatim for parity.
|
||||
if self.client_params is None:
|
||||
return False
|
||||
have = self.client_params.capabilities
|
||||
if capability.roots is not None:
|
||||
if have.roots is None:
|
||||
return False
|
||||
if capability.roots.list_changed and not have.roots.list_changed:
|
||||
return False
|
||||
if capability.sampling is not None:
|
||||
if have.sampling is None:
|
||||
return False
|
||||
if capability.sampling.context is not None and have.sampling.context is None:
|
||||
return False
|
||||
if capability.sampling.tools is not None and have.sampling.tools is None:
|
||||
return False
|
||||
if capability.elicitation is not None and have.elicitation is None:
|
||||
return False
|
||||
if capability.experimental is not None:
|
||||
if have.experimental is None:
|
||||
return False
|
||||
for k, v in capability.experimental.items():
|
||||
if k not in have.experimental or have.experimental[k] != v:
|
||||
return False
|
||||
return True
|
||||
|
|
@ -1,21 +1,165 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Generic
|
||||
from typing import Any, Generic, Protocol
|
||||
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import TypeVar
|
||||
|
||||
from mcp.server.connection import Connection
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.shared._context import RequestContext
|
||||
from mcp.shared.context import BaseContext
|
||||
from mcp.shared.dispatcher import DispatchContext
|
||||
from mcp.shared.message import CloseSSEStreamCallback
|
||||
from mcp.shared.peer import Meta
|
||||
from mcp.shared.transport_context import TransportContext
|
||||
from mcp.types import LoggingLevel, RequestId, RequestParamsMeta
|
||||
|
||||
# Invariant: parameterizes a mutable dataclass field; dict default matches the default lifespan.
|
||||
LifespanContextT = TypeVar("LifespanContextT", default=dict[str, Any])
|
||||
RequestT = TypeVar("RequestT", default=Any)
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ServerRequestContext(RequestContext[ServerSession], Generic[LifespanContextT, RequestT]):
|
||||
class ServerRequestContext(Generic[LifespanContextT, RequestT]):
|
||||
"""Per-request context handed to lowlevel request and notification handlers.
|
||||
|
||||
Built by `ServerRunner._make_context` for each inbound message. Carries the
|
||||
connection-scoped `ServerSession` (server-to-client requests and
|
||||
notifications), per-request metadata, and any per-message data the
|
||||
transport attached (the HTTP request, SSE stream-close callbacks).
|
||||
"""
|
||||
|
||||
session: ServerSession
|
||||
lifespan_context: LifespanContextT
|
||||
request_id: RequestId | None = None
|
||||
meta: RequestParamsMeta | None = None
|
||||
request: RequestT | None = None
|
||||
close_sse_stream: CloseSSEStreamCallback | None = None
|
||||
close_standalone_sse_stream: CloseSSEStreamCallback | None = None
|
||||
|
||||
|
||||
# Covariant: `lifespan` is exposed read-only, so a `Context[AppState]` passes as `Context[object]`.
|
||||
LifespanT_co = TypeVar("LifespanT_co", default=Any, covariant=True)
|
||||
|
||||
|
||||
class Context(BaseContext[TransportContext], Generic[LifespanT_co]):
|
||||
"""Server-side per-request context.
|
||||
|
||||
Extends `BaseContext` (transport metadata, the raw back-channel, progress
|
||||
reporting) with `lifespan`, `connection`, and request-scoped `log`.
|
||||
|
||||
Not currently constructed by `ServerRunner`, which hands handlers a
|
||||
`ServerRequestContext` instead.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dctx: DispatchContext[TransportContext],
|
||||
*,
|
||||
lifespan: LifespanT_co,
|
||||
connection: Connection,
|
||||
meta: RequestParamsMeta | None = None,
|
||||
) -> None:
|
||||
super().__init__(dctx, meta=meta)
|
||||
self._lifespan = lifespan
|
||||
self._connection = connection
|
||||
|
||||
@property
|
||||
def lifespan(self) -> LifespanT_co:
|
||||
"""The server-wide lifespan output (what `Server(..., lifespan=...)` yielded)."""
|
||||
return self._lifespan
|
||||
|
||||
@property
|
||||
def connection(self) -> Connection:
|
||||
"""The per-client `Connection` for this request's connection."""
|
||||
return self._connection
|
||||
|
||||
@property
|
||||
def session_id(self) -> str | None:
|
||||
"""The transport's session id for this connection, when one exists.
|
||||
|
||||
Convenience for `ctx.connection.session_id`. `None` on stdio and
|
||||
stateless HTTP.
|
||||
"""
|
||||
return self._connection.session_id
|
||||
|
||||
@property
|
||||
def headers(self) -> Mapping[str, str] | None:
|
||||
"""Request headers carried by this message, when the transport has them.
|
||||
|
||||
Convenience for `ctx.transport.headers`. `None` on stdio.
|
||||
"""
|
||||
return self.transport.headers
|
||||
|
||||
async def log(self, level: LoggingLevel, data: Any, logger: str | None = None, *, meta: Meta | None = None) -> None:
|
||||
"""Send a request-scoped `notifications/message` log entry.
|
||||
|
||||
Uses this request's back-channel (so the entry rides the request's SSE
|
||||
stream in streamable HTTP), not the standalone stream - use
|
||||
`ctx.connection.log(...)` for that.
|
||||
"""
|
||||
params: dict[str, Any] = {"level": level, "data": data}
|
||||
if logger is not None:
|
||||
params["logger"] = logger
|
||||
if meta:
|
||||
params["_meta"] = meta
|
||||
await self.notify("notifications/message", params)
|
||||
|
||||
|
||||
HandlerResult = BaseModel | dict[str, Any] | None
|
||||
"""What a request handler (or middleware) may return. `ServerRunner` serializes
|
||||
all three to a result dict."""
|
||||
|
||||
CallNext = Callable[[], Awaitable[HandlerResult]]
|
||||
|
||||
_MwLifespanT = TypeVar("_MwLifespanT")
|
||||
|
||||
|
||||
class ServerMiddleware(Protocol[_MwLifespanT]):
|
||||
"""Context-tier middleware: `(ctx, method, params, call_next) -> result`.
|
||||
|
||||
Runs at the top of `ServerRunner._on_request` / `_on_notify` after `ctx`
|
||||
is built but before any validation, lookup, or handshake. Wraps every
|
||||
inbound request and notification: `initialize`, the pre-init gate,
|
||||
`METHOD_NOT_FOUND`, params validation, the handler call, and
|
||||
`notifications/initialized` all run inside `call_next()`.
|
||||
`notifications/cancelled` is observed too; the dispatcher applies the
|
||||
cancellation itself, then forwards the notification. A request-side
|
||||
failure reaches the middleware as a raised `MCPError` (or
|
||||
`ValidationError` for malformed params) so observation/logging middleware
|
||||
can record it. Listed outermost-first on `Server.middleware`.
|
||||
|
||||
`ctx.request_id is None` distinguishes a notification from a request. For
|
||||
notifications `call_next()` returns `None` (a dropped or unhandled
|
||||
notification also returns `None`) and the middleware's own return value is
|
||||
discarded.
|
||||
|
||||
`params` is the raw inbound mapping (no model validation has happened
|
||||
yet). For typed inspection, validate against the model the middleware
|
||||
expects.
|
||||
|
||||
Warning: `initialize` is handled inline - the dispatcher does not read
|
||||
further inbound messages until the middleware chain returns. Awaiting a
|
||||
server-to-client request (`ctx.session.send_request`, `send_ping`, ...)
|
||||
while handling `initialize` therefore deadlocks the connection: the
|
||||
response can never be dequeued. Send-and-forget notifications are safe.
|
||||
|
||||
`Server[L].middleware` holds `ServerMiddleware[L]`, so an app-specific
|
||||
middleware sees `ctx.lifespan_context: L`. While the context is the
|
||||
mutable `ServerRequestContext` dataclass it is invariant in `L`, so a
|
||||
reusable middleware should be typed `ServerMiddleware[Any]` to register on
|
||||
any `Server[L]`.
|
||||
"""
|
||||
|
||||
# TODO(maxisbey): once `_make_context` returns the (covariant) `Context[L]`
|
||||
# again, restore `_MwLifespanT` to `contravariant=True` and retype `ctx`
|
||||
# below to `Context[_MwLifespanT]` so reusable middleware can be
|
||||
# `ServerMiddleware[object]` instead of `ServerMiddleware[Any]`.
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
ctx: ServerRequestContext[_MwLifespanT, Any],
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
call_next: CallNext,
|
||||
) -> HandlerResult: ...
|
||||
|
|
|
|||
|
|
@ -36,16 +36,14 @@ handler callables by method string.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextvars
|
||||
import logging
|
||||
import warnings
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
|
||||
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from importlib.metadata import version as importlib_version
|
||||
from typing import Any, Generic, cast
|
||||
from typing import Any, Generic
|
||||
|
||||
import anyio
|
||||
from opentelemetry.trace import SpanKind, StatusCode
|
||||
from pydantic import BaseModel
|
||||
from starlette.applications import Starlette
|
||||
from starlette.middleware import Middleware
|
||||
from starlette.middleware.authentication import AuthenticationMiddleware
|
||||
|
|
@ -58,22 +56,45 @@ from mcp.server.auth.middleware.bearer_auth import BearerAuthBackend, RequireAut
|
|||
from mcp.server.auth.provider import OAuthAuthorizationServerProvider, TokenVerifier
|
||||
from mcp.server.auth.routes import build_resource_metadata_url, create_auth_routes, create_protected_resource_routes
|
||||
from mcp.server.auth.settings import AuthSettings
|
||||
from mcp.server.context import ServerRequestContext
|
||||
from mcp.server.context import HandlerResult, ServerMiddleware, ServerRequestContext
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.server.runner import ServerRunner, otel_middleware
|
||||
from mcp.server.streamable_http import EventStore
|
||||
from mcp.server.streamable_http_manager import StreamableHTTPASGIApp, StreamableHTTPSessionManager
|
||||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
from mcp.shared._otel import extract_trace_context, otel_span
|
||||
from mcp.shared._stream_protocols import ReadStream, WriteStream
|
||||
from mcp.shared.exceptions import MCPError
|
||||
from mcp.shared.message import ServerMessageMetadata, SessionMessage
|
||||
from mcp.shared.session import RequestResponder
|
||||
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.transport_context import TransportContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
LifespanResultT = TypeVar("LifespanResultT", default=Any)
|
||||
|
||||
_ParamsT = TypeVar("_ParamsT", bound=BaseModel, default=BaseModel)
|
||||
|
||||
RequestHandler = Callable[[ServerRequestContext[LifespanResultT], _ParamsT], Awaitable[HandlerResult]]
|
||||
"""A registered request handler: `(ctx, params) -> result`."""
|
||||
|
||||
NotificationHandler = Callable[[ServerRequestContext[LifespanResultT], _ParamsT], Awaitable[None]]
|
||||
"""A registered notification handler: `(ctx, params) -> None`."""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class HandlerEntry(Generic[LifespanResultT]):
|
||||
"""A registered handler and the params model to validate incoming params against.
|
||||
|
||||
Stored in `Server._request_handlers` / `_notification_handlers` and consumed
|
||||
by `ServerRunner` to validate, build `Context`, and invoke. The handler's
|
||||
second-argument type is erased to `Any` in storage (each entry has a
|
||||
different concrete params type and `Callable` parameters are contravariant);
|
||||
the precise type is recoverable via `params_type`. The correlation is
|
||||
enforced at registration time by `Server.add_request_handler`.
|
||||
"""
|
||||
|
||||
params_type: type[BaseModel]
|
||||
handler: RequestHandler[LifespanResultT, Any]
|
||||
|
||||
|
||||
class NotificationOptions:
|
||||
def __init__(self, prompts_changed: bool = False, resources_changed: bool = False, tools_changed: bool = False):
|
||||
|
|
@ -83,7 +104,7 @@ class NotificationOptions:
|
|||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_: Server[LifespanResultT]) -> AsyncIterator[dict[str, Any]]:
|
||||
async def lifespan(_: Server[Any]) -> AsyncIterator[dict[str, Any]]:
|
||||
"""Default lifespan context manager that does nothing.
|
||||
|
||||
Returns:
|
||||
|
|
@ -191,53 +212,91 @@ class Server(Generic[LifespanResultT]):
|
|||
self.website_url = website_url
|
||||
self.icons = icons
|
||||
self.lifespan = lifespan
|
||||
self._request_handlers: dict[str, Callable[[ServerRequestContext[LifespanResultT], Any], Awaitable[Any]]] = {}
|
||||
self._notification_handlers: dict[
|
||||
str, Callable[[ServerRequestContext[LifespanResultT], Any], Awaitable[None]]
|
||||
] = {}
|
||||
self._request_handlers: dict[str, HandlerEntry[LifespanResultT]] = {}
|
||||
self._notification_handlers: dict[str, HandlerEntry[LifespanResultT]] = {}
|
||||
self._session_manager: StreamableHTTPSessionManager | None = None
|
||||
# Context-tier middleware: wraps every inbound request (including
|
||||
# `initialize`, lookup, validation, handler) with
|
||||
# `(ctx, method, params, call_next)`. Applied in `ServerRunner._on_request`.
|
||||
# TODO(maxisbey): provisional - signature and semantics change with the
|
||||
# Context/middleware rework (covariant `Context[L]`, outbound seam) before
|
||||
# v2 final.
|
||||
self.middleware: list[ServerMiddleware[LifespanResultT]] = []
|
||||
logger.debug("Initializing server %r", name)
|
||||
|
||||
# Populate internal handler dicts from on_* kwargs
|
||||
self._request_handlers.update(
|
||||
{
|
||||
method: handler
|
||||
for method, handler in {
|
||||
"ping": on_ping,
|
||||
"prompts/list": on_list_prompts,
|
||||
"prompts/get": on_get_prompt,
|
||||
"resources/list": on_list_resources,
|
||||
"resources/templates/list": on_list_resource_templates,
|
||||
"resources/read": on_read_resource,
|
||||
"resources/subscribe": on_subscribe_resource,
|
||||
"resources/unsubscribe": on_unsubscribe_resource,
|
||||
"tools/list": on_list_tools,
|
||||
"tools/call": on_call_tool,
|
||||
"logging/setLevel": on_set_logging_level,
|
||||
"completion/complete": on_completion,
|
||||
}.items()
|
||||
if handler is not None
|
||||
}
|
||||
)
|
||||
_spec_requests: list[tuple[str, type[BaseModel], RequestHandler[LifespanResultT, Any] | None]] = [
|
||||
("ping", types.RequestParams, on_ping),
|
||||
("prompts/list", types.PaginatedRequestParams, on_list_prompts),
|
||||
("prompts/get", types.GetPromptRequestParams, on_get_prompt),
|
||||
("resources/list", types.PaginatedRequestParams, on_list_resources),
|
||||
("resources/templates/list", types.PaginatedRequestParams, on_list_resource_templates),
|
||||
("resources/read", types.ReadResourceRequestParams, on_read_resource),
|
||||
("resources/subscribe", types.SubscribeRequestParams, on_subscribe_resource),
|
||||
("resources/unsubscribe", types.UnsubscribeRequestParams, on_unsubscribe_resource),
|
||||
("tools/list", types.PaginatedRequestParams, on_list_tools),
|
||||
("tools/call", types.CallToolRequestParams, on_call_tool),
|
||||
("logging/setLevel", types.SetLevelRequestParams, on_set_logging_level),
|
||||
("completion/complete", types.CompleteRequestParams, on_completion),
|
||||
]
|
||||
self._request_handlers.update({m: HandlerEntry(pt, h) for m, pt, h in _spec_requests if h is not None})
|
||||
|
||||
_spec_notifications: list[tuple[str, type[BaseModel], NotificationHandler[LifespanResultT, Any] | None]] = [
|
||||
("notifications/roots/list_changed", types.NotificationParams, on_roots_list_changed),
|
||||
("notifications/progress", types.ProgressNotificationParams, on_progress),
|
||||
]
|
||||
self._notification_handlers.update(
|
||||
{
|
||||
method: handler
|
||||
for method, handler in {
|
||||
"notifications/roots/list_changed": on_roots_list_changed,
|
||||
"notifications/progress": on_progress,
|
||||
}.items()
|
||||
if handler is not None
|
||||
}
|
||||
{m: HandlerEntry(pt, h) for m, pt, h in _spec_notifications if h is not None}
|
||||
)
|
||||
|
||||
def _add_request_handler(
|
||||
def add_request_handler(
|
||||
self,
|
||||
method: str,
|
||||
handler: Callable[[ServerRequestContext[LifespanResultT], Any], Awaitable[Any]],
|
||||
params_type: type[_ParamsT],
|
||||
handler: RequestHandler[LifespanResultT, _ParamsT],
|
||||
) -> None:
|
||||
"""Add a request handler, silently replacing any existing handler for the same method."""
|
||||
self._request_handlers[method] = handler
|
||||
"""Register a request handler for `method`.
|
||||
|
||||
`params_type` is the model incoming params are validated against
|
||||
before the handler is invoked. It should subclass `RequestParams` so
|
||||
`_meta` parses uniformly. A message with no `params` member validates
|
||||
`{}` against `params_type`: models with required fields reject it as
|
||||
INVALID_PARAMS, all-optional models reach the handler with their
|
||||
defaults - the handler never receives `None`. Replaces any existing
|
||||
handler for the same method, except `initialize`, which is reserved:
|
||||
the runner owns the handshake, so registering it raises `ValueError`.
|
||||
Use `Server.middleware` to observe or wrap initialization.
|
||||
"""
|
||||
if method == "initialize":
|
||||
raise ValueError(
|
||||
"'initialize' is handled by the server runner and cannot be overridden; "
|
||||
"use Server.middleware to observe or wrap initialization"
|
||||
)
|
||||
self._request_handlers[method] = HandlerEntry(params_type, handler)
|
||||
|
||||
def add_notification_handler(
|
||||
self,
|
||||
method: str,
|
||||
params_type: type[_ParamsT],
|
||||
handler: NotificationHandler[LifespanResultT, _ParamsT],
|
||||
) -> None:
|
||||
"""Register a notification handler for `method`.
|
||||
|
||||
`params_type` should subclass `NotificationParams` so `_meta`
|
||||
parses uniformly. Absent params follow the same contract as requests:
|
||||
`{}` is validated, so the handler receives the model with its defaults,
|
||||
never `None`. Replaces any existing handler. A handler for
|
||||
`notifications/initialized` runs after the runner has marked the
|
||||
connection initialized.
|
||||
"""
|
||||
self._notification_handlers[method] = HandlerEntry(params_type, handler)
|
||||
|
||||
def get_request_handler(self, method: str) -> HandlerEntry[LifespanResultT] | None:
|
||||
"""Return the registered entry for a request method, or `None`."""
|
||||
return self._request_handlers.get(method)
|
||||
|
||||
def get_notification_handler(self, method: str) -> HandlerEntry[LifespanResultT] | None:
|
||||
"""Return the registered entry for a notification method, or `None`."""
|
||||
return self._notification_handlers.get(method)
|
||||
|
||||
# TODO: Rethink capabilities API. Currently capabilities are derived from registered
|
||||
# handlers but require NotificationOptions to be passed externally for list_changed
|
||||
|
|
@ -347,167 +406,31 @@ class Server(Generic[LifespanResultT]):
|
|||
# the initialization lifecycle, but can do so with any available node
|
||||
# rather than requiring initialization for each connection.
|
||||
stateless: bool = False,
|
||||
):
|
||||
async with AsyncExitStack() as stack:
|
||||
lifespan_context = await stack.enter_async_context(self.lifespan(self))
|
||||
session = await stack.enter_async_context(
|
||||
ServerSession(
|
||||
read_stream,
|
||||
write_stream,
|
||||
initialization_options,
|
||||
stateless=stateless,
|
||||
)
|
||||
)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
try:
|
||||
async for message in session.incoming_messages:
|
||||
logger.debug("Received message: %s", message)
|
||||
|
||||
if isinstance(message, RequestResponder) and message.context is not None:
|
||||
context = message.context
|
||||
else:
|
||||
context = contextvars.copy_context()
|
||||
|
||||
context.run(
|
||||
tg.start_soon,
|
||||
self._handle_message,
|
||||
message,
|
||||
session,
|
||||
lifespan_context,
|
||||
raise_exceptions,
|
||||
)
|
||||
finally:
|
||||
# Transport closed: cancel in-flight handlers. Without this the
|
||||
# TG join waits for them, and when they eventually try to
|
||||
# respond they hit a closed write stream (the session's
|
||||
# _receive_loop closed it when the read stream ended).
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
async def _handle_message(
|
||||
self,
|
||||
message: RequestResponder[types.ClientRequest, types.ServerResult] | types.ClientNotification | Exception,
|
||||
session: ServerSession,
|
||||
lifespan_context: LifespanResultT,
|
||||
raise_exceptions: bool = False,
|
||||
):
|
||||
with warnings.catch_warnings(record=True) as w:
|
||||
match message:
|
||||
case RequestResponder() as responder:
|
||||
with responder:
|
||||
await self._handle_request(
|
||||
message, responder.request, session, lifespan_context, raise_exceptions
|
||||
)
|
||||
case Exception():
|
||||
logger.error(f"Received exception from stream: {message}")
|
||||
if raise_exceptions:
|
||||
raise message
|
||||
case _:
|
||||
await self._handle_notification(message, session, lifespan_context)
|
||||
|
||||
for warning in w: # pragma: lax no cover
|
||||
logger.info("Warning: %s: %s", warning.category.__name__, warning.message)
|
||||
|
||||
async def _handle_request(
|
||||
self,
|
||||
message: RequestResponder[types.ClientRequest, types.ServerResult],
|
||||
req: types.ClientRequest,
|
||||
session: ServerSession,
|
||||
lifespan_context: LifespanResultT,
|
||||
raise_exceptions: bool,
|
||||
):
|
||||
logger.info("Processing request of type %s", type(req).__name__)
|
||||
|
||||
target = getattr(req.params, "name", None) if req.params else None
|
||||
span_name = f"MCP handle {req.method} {target}" if target else f"MCP handle {req.method}"
|
||||
|
||||
# Extract W3C trace context from _meta (SEP-414).
|
||||
meta = cast(dict[str, Any] | None, getattr(req.params, "meta", None)) if req.params else None
|
||||
parent_context = extract_trace_context(meta) if meta is not None else None
|
||||
|
||||
with otel_span(
|
||||
span_name,
|
||||
kind=SpanKind.SERVER,
|
||||
attributes={"mcp.method.name": req.method, "jsonrpc.request.id": message.request_id},
|
||||
context=parent_context,
|
||||
) as span:
|
||||
if handler := self._request_handlers.get(req.method):
|
||||
logger.debug("Dispatching request of type %s", type(req).__name__)
|
||||
|
||||
try:
|
||||
# Extract request context and close_sse_stream from message metadata
|
||||
request_data = None
|
||||
close_sse_stream_cb = None
|
||||
close_standalone_sse_stream_cb = None
|
||||
if message.message_metadata is not None and isinstance(
|
||||
message.message_metadata, ServerMessageMetadata
|
||||
):
|
||||
request_data = message.message_metadata.request_context
|
||||
close_sse_stream_cb = message.message_metadata.close_sse_stream
|
||||
close_standalone_sse_stream_cb = message.message_metadata.close_standalone_sse_stream
|
||||
|
||||
ctx = ServerRequestContext(
|
||||
request_id=message.request_id,
|
||||
meta=message.request_meta,
|
||||
session=session,
|
||||
lifespan_context=lifespan_context,
|
||||
request=request_data,
|
||||
close_sse_stream=close_sse_stream_cb,
|
||||
close_standalone_sse_stream=close_standalone_sse_stream_cb,
|
||||
)
|
||||
response = await handler(ctx, req.params)
|
||||
except MCPError as err:
|
||||
response = err.error
|
||||
except anyio.get_cancelled_exc_class():
|
||||
if message.cancelled:
|
||||
# Client sent CancelledNotification; responder.cancel() already
|
||||
# sent an error response, so skip the duplicate.
|
||||
logger.info("Request %s cancelled - duplicate response suppressed", message.request_id)
|
||||
return
|
||||
# Transport-close cancellation from the TG in run(); re-raise so the
|
||||
# TG swallows its own cancellation.
|
||||
raise
|
||||
except Exception as err:
|
||||
if raise_exceptions: # pragma: no cover
|
||||
raise err
|
||||
response = types.ErrorData(code=0, message=str(err))
|
||||
else:
|
||||
response = types.ErrorData(code=types.METHOD_NOT_FOUND, message="Method not found")
|
||||
|
||||
if isinstance(response, types.ErrorData) and span is not None:
|
||||
span.set_status(StatusCode.ERROR, response.message)
|
||||
|
||||
try:
|
||||
await message.respond(response)
|
||||
except (anyio.BrokenResourceError, anyio.ClosedResourceError):
|
||||
# Transport closed between handler unblocking and respond. Happens
|
||||
# when _receive_loop's finally wakes a handler blocked on
|
||||
# send_request: the handler runs to respond() before run()'s TG
|
||||
# cancel fires, but after the write stream closed. Closed if our
|
||||
# end closed (_receive_loop's async-with exit); Broken if the peer
|
||||
# end closed first (streamable_http terminate()).
|
||||
logger.debug("Response for %s dropped - transport closed", message.request_id)
|
||||
return
|
||||
|
||||
logger.debug("Response sent")
|
||||
|
||||
async def _handle_notification(
|
||||
self,
|
||||
notify: types.ClientNotification,
|
||||
session: ServerSession,
|
||||
lifespan_context: LifespanResultT,
|
||||
) -> None:
|
||||
if handler := self._notification_handlers.get(notify.method):
|
||||
logger.debug("Dispatching notification of type %s", type(notify).__name__)
|
||||
|
||||
try:
|
||||
ctx = ServerRequestContext(
|
||||
session=session,
|
||||
lifespan_context=lifespan_context,
|
||||
)
|
||||
await handler(ctx, notify.params)
|
||||
except Exception: # pragma: no cover
|
||||
logger.exception("Uncaught exception in notification handler")
|
||||
async with self.lifespan(self) as lifespan_context:
|
||||
dispatcher: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(
|
||||
read_stream,
|
||||
write_stream,
|
||||
raise_handler_exceptions=raise_exceptions,
|
||||
# Handle `initialize` inline so a client that pipelines it with
|
||||
# the next request (spec says SHOULD NOT, not MUST NOT) sees
|
||||
# the initialized state instead of failing the init-gate.
|
||||
inline_methods=frozenset({"initialize"}),
|
||||
)
|
||||
runner = ServerRunner(
|
||||
server=self,
|
||||
dispatcher=dispatcher,
|
||||
lifespan_state=lifespan_context,
|
||||
init_options=initialization_options,
|
||||
# Stateless HTTP has no standalone GET stream, so server-initiated
|
||||
# requests on `runner.connection` must fail fast with
|
||||
# `NoBackChannelError` rather than write to a channel that will
|
||||
# never deliver a response.
|
||||
has_standalone_channel=not stateless,
|
||||
stateless=stateless,
|
||||
dispatch_middleware=[otel_middleware],
|
||||
)
|
||||
await runner.run()
|
||||
|
||||
def streamable_http_app(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -607,11 +607,7 @@ class MCPServer(Generic[LifespanResultT]):
|
|||
completion=result if result is not None else Completion(values=[], total=None, has_more=None),
|
||||
)
|
||||
|
||||
# TODO(maxisbey): remove private access — completion needs post-construction
|
||||
# handler registration, find a better pattern for this
|
||||
self._lowlevel_server._add_request_handler( # pyright: ignore[reportPrivateUsage]
|
||||
"completion/complete", handler
|
||||
)
|
||||
self._lowlevel_server.add_request_handler("completion/complete", CompleteRequestParams, handler)
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
|
|
|||
377
src/mcp/server/runner.py
Normal file
377
src/mcp/server/runner.py
Normal file
|
|
@ -0,0 +1,377 @@
|
|||
"""`ServerRunner` - per-connection orchestrator over a `Dispatcher`.
|
||||
|
||||
`ServerRunner` is the bridge between the dispatcher layer (`on_request` /
|
||||
`on_notify`, untyped dicts) and the user's handler layer (typed `Context`,
|
||||
typed params). One instance per client connection. It:
|
||||
|
||||
* handles the `initialize` handshake and populates `Connection`
|
||||
* gates requests until initialized (`ping` exempt)
|
||||
* looks up the handler in the server's registry, validates params, builds
|
||||
`Context`, runs the middleware chain, returns the result dict
|
||||
* drives `dispatcher.run()` and the per-connection lifespan
|
||||
|
||||
`ServerRunner` holds a `Server` directly - `Server` is the registry.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from functools import partial, reduce
|
||||
from typing import TYPE_CHECKING, Any, Generic, cast, get_args
|
||||
|
||||
import anyio.abc
|
||||
from opentelemetry.trace import SpanKind, StatusCode
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from typing_extensions import TypeVar
|
||||
|
||||
from mcp.server.connection import Connection
|
||||
from mcp.server.context import CallNext, HandlerResult, ServerMiddleware, ServerRequestContext
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.shared._otel import extract_trace_context, otel_span
|
||||
from mcp.shared.dispatcher import DispatchContext, DispatchMiddleware, OnRequest
|
||||
from mcp.shared.exceptions import MCPError
|
||||
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
|
||||
from mcp.shared.message import ServerMessageMetadata
|
||||
from mcp.shared.transport_context import TransportContext
|
||||
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
|
||||
from mcp.types import (
|
||||
INVALID_PARAMS,
|
||||
LATEST_PROTOCOL_VERSION,
|
||||
METHOD_NOT_FOUND,
|
||||
ClientRequest,
|
||||
ErrorData,
|
||||
Implementation,
|
||||
InitializeRequestParams,
|
||||
InitializeResult,
|
||||
NotificationParams,
|
||||
RequestParams,
|
||||
RequestParamsMeta,
|
||||
client_request_adapter,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mcp.server.lowlevel.server import Server
|
||||
|
||||
__all__ = ["CallNext", "ServerMiddleware", "ServerRunner", "otel_middleware"]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
LifespanT = TypeVar("LifespanT", default=Any)
|
||||
|
||||
|
||||
_INIT_EXEMPT: frozenset[str] = frozenset({"ping"})
|
||||
|
||||
_EXIT_STACK_CLOSE_TIMEOUT: float = 5
|
||||
"""Bound for the shielded exit-stack unwind in `run()`; a hung cleanup
|
||||
callback must not wedge shutdown."""
|
||||
|
||||
|
||||
def _extract_meta(params: Mapping[str, Any] | None) -> RequestParamsMeta | None:
|
||||
"""Lift `_meta` from raw params; `None` when absent or malformed, so
|
||||
context construction is independent of params validity."""
|
||||
if not params or "_meta" not in params:
|
||||
return None
|
||||
try:
|
||||
return RequestParams.model_validate(params, by_name=False).meta
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
_SPEC_CLIENT_METHODS: frozenset[str] = frozenset(
|
||||
cast(type[BaseModel], arm).model_fields["method"].default for arm in get_args(ClientRequest)
|
||||
)
|
||||
"""Method names in the spec `ClientRequest` union, derived from the
|
||||
discriminator literal on each arm. Used to gate upfront validation so custom
|
||||
methods registered via `add_request_handler` are not rejected."""
|
||||
|
||||
|
||||
def otel_middleware(next_on_request: OnRequest) -> OnRequest:
|
||||
"""Dispatch-tier middleware that wraps each request in an OpenTelemetry span.
|
||||
|
||||
Mirrors the span shape of the existing `Server._handle_request`: span name
|
||||
`"MCP handle <method> [<target>]"`, `mcp.method.name` attribute, W3C
|
||||
trace context extracted from `params._meta` (SEP-414), and an ERROR
|
||||
status if the handler raises.
|
||||
"""
|
||||
|
||||
async def wrapped(
|
||||
dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
|
||||
) -> dict[str, Any]:
|
||||
target: str | None
|
||||
match params:
|
||||
case {"name": str() as target}:
|
||||
pass
|
||||
case _:
|
||||
target = None
|
||||
parent: Any | None
|
||||
match params:
|
||||
case {"_meta": {**meta}}:
|
||||
parent = extract_trace_context(meta)
|
||||
case _:
|
||||
parent = None
|
||||
span_name = f"MCP handle {method}{f' {target}' if target else ''}"
|
||||
# `otel_middleware` wraps `on_request` only, so `request_id` is always set.
|
||||
attributes = {"mcp.method.name": method, "jsonrpc.request.id": str(dctx.request_id)}
|
||||
with otel_span(
|
||||
span_name,
|
||||
kind=SpanKind.SERVER,
|
||||
attributes=attributes,
|
||||
context=parent,
|
||||
record_exception=False,
|
||||
set_status_on_exception=False,
|
||||
) as span:
|
||||
try:
|
||||
return await next_on_request(dctx, method, params)
|
||||
except MCPError as e:
|
||||
span.set_status(StatusCode.ERROR, e.error.message)
|
||||
raise
|
||||
except ValidationError:
|
||||
# Mirror the sanitized wire response; pydantic messages carry client input.
|
||||
span.set_status(StatusCode.ERROR, "Invalid request parameters")
|
||||
raise
|
||||
except Exception as e:
|
||||
span.record_exception(e)
|
||||
span.set_status(StatusCode.ERROR, str(e))
|
||||
raise
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
def _dump_result(result: Any) -> dict[str, Any]:
|
||||
if result is None:
|
||||
return {}
|
||||
if isinstance(result, ErrorData):
|
||||
# ErrorData is a JSON-RPC error, not a success result. Handler returns
|
||||
# already raise in `_inner`; this catches middleware returning one.
|
||||
raise MCPError.from_error_data(result)
|
||||
if isinstance(result, BaseModel):
|
||||
return result.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
if isinstance(result, dict):
|
||||
return cast(dict[str, Any], result)
|
||||
raise TypeError(f"handler returned {type(result).__name__}; expected BaseModel, dict, or None")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServerRunner(Generic[LifespanT]):
|
||||
"""Per-connection orchestrator. One instance per client connection."""
|
||||
|
||||
server: Server[LifespanT]
|
||||
dispatcher: JSONRPCDispatcher[Any]
|
||||
lifespan_state: LifespanT
|
||||
has_standalone_channel: bool
|
||||
init_options: InitializationOptions | None = None
|
||||
"""`InitializeResult` payload. Defaults to `server.create_initialization_options()`."""
|
||||
session_id: str | None = None
|
||||
stateless: bool = False
|
||||
dispatch_middleware: list[DispatchMiddleware] = field(default_factory=list[DispatchMiddleware])
|
||||
|
||||
connection: Connection = field(init=False)
|
||||
session: ServerSession = field(init=False)
|
||||
"""Connection-scoped: the same instance reaches every request as `ctx.session`."""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.init_options is None:
|
||||
self.init_options = self.server.create_initialization_options()
|
||||
self.connection = Connection(
|
||||
self.dispatcher, has_standalone_channel=self.has_standalone_channel, session_id=self.session_id
|
||||
)
|
||||
if self.stateless:
|
||||
# No handshake ever arrives on a stateless connection; born ready.
|
||||
self.connection.initialized.set()
|
||||
self.session = ServerSession(self.dispatcher, self.connection, stateless=self.stateless)
|
||||
|
||||
async def run(self, *, task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED) -> None:
|
||||
"""Drive the dispatcher until the underlying channel closes.
|
||||
|
||||
Composes `dispatch_middleware` over `_on_request` and hands the result
|
||||
to `dispatcher.run()`. `task_status.started()` is forwarded so callers
|
||||
can `await tg.start(runner.run)` and resume once the dispatcher is
|
||||
ready to accept requests. Once the dispatcher exits,
|
||||
`connection.exit_stack` is unwound (shielded from outer cancellation,
|
||||
bounded by `_EXIT_STACK_CLOSE_TIMEOUT`) so any per-connection cleanup
|
||||
registered by handlers or middleware gets a chance to run without a
|
||||
misbehaving callback hanging shutdown indefinitely.
|
||||
"""
|
||||
try:
|
||||
await self.dispatcher.run(self._compose_on_request(), self._on_notify, task_status=task_status)
|
||||
finally:
|
||||
with anyio.move_on_after(_EXIT_STACK_CLOSE_TIMEOUT, shield=True) as scope:
|
||||
try:
|
||||
await self.connection.exit_stack.aclose()
|
||||
except Exception:
|
||||
# Raising here would mask dispatcher.run()'s exception and
|
||||
# crash stdio servers on normal disconnect.
|
||||
logger.exception("connection exit_stack cleanup raised")
|
||||
if scope.cancelled_caught:
|
||||
logger.warning(
|
||||
"connection exit_stack cleanup exceeded %s seconds; abandoning remaining callbacks",
|
||||
_EXIT_STACK_CLOSE_TIMEOUT,
|
||||
)
|
||||
|
||||
def _compose_on_request(self) -> OnRequest:
|
||||
"""Wrap `_on_request` in `dispatch_middleware`, outermost-first.
|
||||
|
||||
Dispatch-tier middleware sees raw `(dctx, method, params) -> dict`
|
||||
and wraps everything - initialize, METHOD_NOT_FOUND, validation
|
||||
failures included. `run()` calls this once and hands the result to
|
||||
`dispatcher.run()`.
|
||||
"""
|
||||
return reduce(lambda h, mw: mw(h), reversed(self.dispatch_middleware), self._on_request)
|
||||
|
||||
async def _on_request(
|
||||
self,
|
||||
dctx: DispatchContext[TransportContext],
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
ctx = self._make_context(dctx, _extract_meta(params))
|
||||
|
||||
async def _inner() -> HandlerResult:
|
||||
# TODO(maxisbey): pinned compat: spec methods are validated against
|
||||
# the ClientRequest union before lookup, so malformed params are
|
||||
# INVALID_PARAMS even with no handler registered.
|
||||
if method in _SPEC_CLIENT_METHODS:
|
||||
payload: dict[str, Any] = {"method": method}
|
||||
if params is not None:
|
||||
payload["params"] = dict(params)
|
||||
client_request_adapter.validate_python(payload, by_name=False)
|
||||
# TODO(maxisbey): the 2026-07-28 spec drops the handshake; this branch and
|
||||
# the gate become a per-version legacy path then. Initialize runs inline
|
||||
# (read loop parked), so awaiting the peer anywhere on this path deadlocks.
|
||||
if method == "initialize":
|
||||
return self._handle_initialize(params)
|
||||
if not self.connection.initialize_accepted and method not in _INIT_EXEMPT:
|
||||
# Pinned compat: the same error shape the union validation produced.
|
||||
raise MCPError(code=INVALID_PARAMS, message="Invalid request parameters", data="")
|
||||
entry = self.server.get_request_handler(method)
|
||||
if entry is None:
|
||||
raise MCPError(code=METHOD_NOT_FOUND, message="Method not found")
|
||||
# Absent params validate as {} (required fields still reject), so
|
||||
# the handler receives the model with its defaults, never None.
|
||||
typed_params = entry.params_type.model_validate({} if params is None else params, by_name=False)
|
||||
result = await entry.handler(ctx, typed_params)
|
||||
if isinstance(result, ErrorData):
|
||||
# Raise inside the chain so middleware observes the failure.
|
||||
raise MCPError.from_error_data(result)
|
||||
return result
|
||||
|
||||
call = self._compose_server_middleware(ctx, method, params, _inner)
|
||||
result = _dump_result(await call())
|
||||
if method == "initialize":
|
||||
# Commit only on chain success, so a middleware veto leaves no state.
|
||||
# Race-free: the read loop is parked until this call returns.
|
||||
self.connection.client_params, self.connection.protocol_version = self._negotiate_initialize(params)
|
||||
return result
|
||||
|
||||
async def _on_notify(
|
||||
self,
|
||||
dctx: DispatchContext[TransportContext],
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
) -> None:
|
||||
ctx = self._make_context(dctx, _extract_meta(params))
|
||||
|
||||
async def _inner() -> None:
|
||||
if method == "notifications/initialized":
|
||||
# Validate before committing so a malformed notification leaves
|
||||
# state untouched; then fall through so a registered handler
|
||||
# observes an initialized connection.
|
||||
if params is not None:
|
||||
try:
|
||||
NotificationParams.model_validate(params, by_name=False)
|
||||
except ValidationError:
|
||||
logger.warning("dropped %r: malformed params", method)
|
||||
return
|
||||
self.connection.initialized.set()
|
||||
elif not self.connection.initialize_accepted:
|
||||
logger.debug("dropped %s: received before initialization", method)
|
||||
return
|
||||
entry = self.server.get_notification_handler(method)
|
||||
if entry is None:
|
||||
logger.debug("no handler for notification %s", method)
|
||||
return
|
||||
# Same absent-params contract as requests.
|
||||
try:
|
||||
typed_params = entry.params_type.model_validate({} if params is None else params, by_name=False)
|
||||
except ValidationError:
|
||||
logger.warning("dropped %r: malformed params", method)
|
||||
return
|
||||
await entry.handler(ctx, typed_params)
|
||||
|
||||
call = self._compose_server_middleware(ctx, method, params, _inner)
|
||||
try:
|
||||
await call()
|
||||
except Exception:
|
||||
# A crashing handler must not cancel the dispatcher's task group;
|
||||
# middleware saw the raise out of call_next() first.
|
||||
logger.exception("notification handler for %r raised", method)
|
||||
|
||||
def _compose_server_middleware(
|
||||
self,
|
||||
ctx: ServerRequestContext[LifespanT, Any],
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
inner: CallNext,
|
||||
) -> CallNext:
|
||||
"""Wrap `inner` in `Server.middleware`, outermost-first.
|
||||
|
||||
Shared by `_on_request` and `_on_notify` so the same middleware chain
|
||||
observes every inbound message.
|
||||
"""
|
||||
call = inner
|
||||
for mw in reversed(self.server.middleware):
|
||||
call = partial(mw, ctx, method, params, call)
|
||||
return call
|
||||
|
||||
def _make_context(
|
||||
self, dctx: DispatchContext[TransportContext], meta: RequestParamsMeta | None
|
||||
) -> ServerRequestContext[LifespanT, Any]:
|
||||
# TODO(maxisbey): remove for Context rework. Reads the SHTTP per-request
|
||||
# data off the raw `dctx.message_metadata` carrier; replace with the
|
||||
# per-transport context once that lands.
|
||||
md = dctx.message_metadata
|
||||
if isinstance(md, ServerMessageMetadata):
|
||||
request = md.request_context
|
||||
close_sse_stream = md.close_sse_stream
|
||||
close_standalone_sse_stream = md.close_standalone_sse_stream
|
||||
else:
|
||||
request = close_sse_stream = close_standalone_sse_stream = None
|
||||
return ServerRequestContext(
|
||||
session=self.session,
|
||||
lifespan_context=self.lifespan_state,
|
||||
request_id=dctx.request_id,
|
||||
meta=meta,
|
||||
request=request,
|
||||
close_sse_stream=close_sse_stream,
|
||||
close_standalone_sse_stream=close_standalone_sse_stream,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _negotiate_initialize(params: Mapping[str, Any] | None) -> tuple[InitializeRequestParams, str]:
|
||||
"""Validate `initialize` params and pick the protocol version."""
|
||||
init = InitializeRequestParams.model_validate(params or {}, by_name=False)
|
||||
requested = init.protocol_version
|
||||
negotiated = requested if requested in SUPPORTED_PROTOCOL_VERSIONS else LATEST_PROTOCOL_VERSION
|
||||
return init, negotiated
|
||||
|
||||
def _handle_initialize(self, params: Mapping[str, Any] | None) -> InitializeResult:
|
||||
"""Build the `initialize` result; state commits later in `_on_request`."""
|
||||
_, negotiated = self._negotiate_initialize(params)
|
||||
assert self.init_options is not None
|
||||
opts = self.init_options
|
||||
return InitializeResult(
|
||||
protocol_version=negotiated,
|
||||
capabilities=opts.capabilities,
|
||||
server_info=Implementation(
|
||||
name=opts.server_name,
|
||||
title=opts.title,
|
||||
description=opts.description,
|
||||
version=opts.server_version,
|
||||
website_url=opts.website_url,
|
||||
icons=opts.icons,
|
||||
),
|
||||
instructions=opts.instructions,
|
||||
)
|
||||
|
|
@ -1,188 +1,104 @@
|
|||
"""ServerSession Module
|
||||
"""`ServerSession`: server-to-client requests and notifications.
|
||||
|
||||
This module provides the ServerSession class, which manages communication between the
|
||||
server and client in the MCP (Model Context Protocol) framework. It is most commonly
|
||||
used in MCP servers to interact with the client.
|
||||
A thin proxy over `JSONRPCDispatcher` and `Connection`. One instance per
|
||||
client connection (built by `ServerRunner`). Handlers reach it as
|
||||
`ctx.session` and use the typed helpers (`create_message`, `elicit_form`,
|
||||
`send_log_message`, ...) to call back to the client.
|
||||
|
||||
Common usage pattern:
|
||||
```
|
||||
async def handle_call_tool(ctx: RequestContext, params: CallToolRequestParams) -> CallToolResult:
|
||||
# Check client capabilities before proceeding
|
||||
if ctx.session.check_client_capability(
|
||||
types.ClientCapabilities(experimental={"advanced_tools": dict()})
|
||||
):
|
||||
result = await perform_advanced_tool_operation(params.arguments)
|
||||
else:
|
||||
result = await perform_basic_tool_operation(params.arguments)
|
||||
return result
|
||||
|
||||
async def handle_list_prompts(ctx: RequestContext, params) -> ListPromptsResult:
|
||||
if ctx.session.client_params:
|
||||
return ListPromptsResult(prompts=generate_custom_prompts(ctx.session.client_params))
|
||||
return ListPromptsResult(prompts=default_prompts)
|
||||
|
||||
server = Server(name, on_call_tool=handle_call_tool, on_list_prompts=handle_list_prompts)
|
||||
```
|
||||
|
||||
The ServerSession class is typically used internally by the Server class and should not
|
||||
be instantiated directly by users of the MCP framework.
|
||||
The receive-loop, initialize handling, and per-request task isolation that
|
||||
used to live here are now owned by `JSONRPCDispatcher` and `ServerRunner`.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any, TypeVar, overload
|
||||
|
||||
import anyio
|
||||
import anyio.lowlevel
|
||||
from anyio.streams.memory import MemoryObjectReceiveStream
|
||||
from pydantic import AnyUrl, TypeAdapter
|
||||
from pydantic import AnyUrl, BaseModel
|
||||
|
||||
from mcp import types
|
||||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.connection import Connection
|
||||
from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages
|
||||
from mcp.shared._stream_protocols import ReadStream, WriteStream
|
||||
from mcp.shared.exceptions import StatelessModeNotSupported
|
||||
from mcp.shared.message import ServerMessageMetadata, SessionMessage
|
||||
from mcp.shared.session import (
|
||||
BaseSession,
|
||||
RequestResponder,
|
||||
)
|
||||
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
|
||||
from mcp.shared.dispatcher import CallOptions, ProgressFnT
|
||||
from mcp.shared.exceptions import NoBackChannelError, StatelessModeNotSupported
|
||||
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
|
||||
from mcp.shared.message import ServerMessageMetadata
|
||||
|
||||
__all__ = ["ServerSession"]
|
||||
|
||||
ResultT = TypeVar("ResultT", bound=BaseModel)
|
||||
|
||||
|
||||
class InitializationState(Enum):
|
||||
NotInitialized = 1
|
||||
Initializing = 2
|
||||
Initialized = 3
|
||||
class ServerSession:
|
||||
"""Connection-scoped proxy for server-to-client requests and notifications.
|
||||
|
||||
|
||||
ServerSessionT = TypeVar("ServerSessionT", bound="ServerSession")
|
||||
|
||||
ServerRequestResponder = (
|
||||
RequestResponder[types.ClientRequest, types.ServerResult] | types.ClientNotification | Exception
|
||||
)
|
||||
|
||||
|
||||
class ServerSession(
|
||||
BaseSession[
|
||||
types.ServerRequest,
|
||||
types.ServerNotification,
|
||||
types.ServerResult,
|
||||
types.ClientRequest,
|
||||
types.ClientNotification,
|
||||
]
|
||||
):
|
||||
_initialized: InitializationState = InitializationState.NotInitialized
|
||||
_client_params: types.InitializeRequestParams | None = None
|
||||
`send_request` / `send_notification` model-dump their argument and forward
|
||||
to the dispatcher; the typed helpers below are unchanged from the previous
|
||||
implementation and only call those two methods.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
read_stream: ReadStream[SessionMessage | Exception],
|
||||
write_stream: WriteStream[SessionMessage],
|
||||
init_options: InitializationOptions,
|
||||
dispatcher: JSONRPCDispatcher[Any],
|
||||
connection: Connection,
|
||||
*,
|
||||
stateless: bool = False,
|
||||
) -> None:
|
||||
super().__init__(read_stream, write_stream)
|
||||
self._dispatcher = dispatcher
|
||||
self._connection = connection
|
||||
self._stateless = stateless
|
||||
self._initialization_state = (
|
||||
InitializationState.Initialized if stateless else InitializationState.NotInitialized
|
||||
)
|
||||
|
||||
self._init_options = init_options
|
||||
self._incoming_message_stream_writer, self._incoming_message_stream_reader = anyio.create_memory_object_stream[
|
||||
ServerRequestResponder
|
||||
](0)
|
||||
self._exit_stack.push_async_callback(lambda: self._incoming_message_stream_reader.aclose())
|
||||
|
||||
@property
|
||||
def _receive_request_adapter(self) -> TypeAdapter[types.ClientRequest]:
|
||||
return types.client_request_adapter
|
||||
|
||||
@property
|
||||
def _receive_notification_adapter(self) -> TypeAdapter[types.ClientNotification]:
|
||||
return types.client_notification_adapter
|
||||
|
||||
@property
|
||||
def client_params(self) -> types.InitializeRequestParams | None:
|
||||
return self._client_params
|
||||
"""The client's `initialize` request params; `None` before initialization."""
|
||||
return self._connection.client_params
|
||||
|
||||
async def send_request(
|
||||
self,
|
||||
request: types.ServerRequest,
|
||||
result_type: type[ResultT],
|
||||
request_read_timeout_seconds: float | None = None,
|
||||
metadata: ServerMessageMetadata | None = None,
|
||||
progress_callback: ProgressFnT | None = None,
|
||||
) -> ResultT:
|
||||
"""Send a typed server-to-client request and validate the result.
|
||||
|
||||
`metadata.related_request_id` (when supplied) routes the outgoing
|
||||
message onto the originating request's response stream over
|
||||
streamable HTTP; it is the only metadata field honored here.
|
||||
|
||||
Raises:
|
||||
MCPError: The peer responded with an error.
|
||||
NoBackChannelError: If there is no related request to ride on and
|
||||
the connection has no standalone channel (stateless HTTP), so
|
||||
a response could never arrive.
|
||||
pydantic.ValidationError: The peer's result does not match `result_type`.
|
||||
"""
|
||||
data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
opts: CallOptions = {}
|
||||
if request_read_timeout_seconds:
|
||||
opts["timeout"] = request_read_timeout_seconds
|
||||
if progress_callback is not None:
|
||||
opts["on_progress"] = progress_callback
|
||||
related = metadata.related_request_id if metadata is not None else None
|
||||
if related is None and not self._connection.has_standalone_channel:
|
||||
# Fail fast instead of parking forever on a response that cannot
|
||||
# arrive; matches `Connection.send_raw_request`.
|
||||
raise NoBackChannelError(data["method"])
|
||||
result = await self._dispatcher.send_raw_request(
|
||||
data["method"], data.get("params"), opts or None, _related_request_id=related
|
||||
)
|
||||
return result_type.model_validate(result, by_name=False)
|
||||
|
||||
async def send_notification(
|
||||
self,
|
||||
notification: types.ServerNotification,
|
||||
related_request_id: types.RequestId | None = None,
|
||||
) -> None:
|
||||
"""Send a typed server-to-client notification."""
|
||||
data = notification.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
await self._dispatcher.notify(data["method"], data.get("params"), _related_request_id=related_request_id)
|
||||
|
||||
def check_client_capability(self, capability: types.ClientCapabilities) -> bool:
|
||||
"""Check if the client supports a specific capability."""
|
||||
if self._client_params is None: # pragma: lax no cover
|
||||
return False
|
||||
|
||||
client_caps = self._client_params.capabilities
|
||||
|
||||
if capability.roots is not None: # pragma: lax no cover
|
||||
if client_caps.roots is None:
|
||||
return False
|
||||
if capability.roots.list_changed and not client_caps.roots.list_changed:
|
||||
return False
|
||||
|
||||
if capability.sampling is not None: # pragma: lax no cover
|
||||
if client_caps.sampling is None:
|
||||
return False
|
||||
if capability.sampling.context is not None and client_caps.sampling.context is None:
|
||||
return False
|
||||
if capability.sampling.tools is not None and client_caps.sampling.tools is None:
|
||||
return False
|
||||
|
||||
if capability.elicitation is not None and client_caps.elicitation is None: # pragma: lax no cover
|
||||
return False
|
||||
|
||||
if capability.experimental is not None: # pragma: lax no cover
|
||||
if client_caps.experimental is None:
|
||||
return False
|
||||
for exp_key, exp_value in capability.experimental.items():
|
||||
if exp_key not in client_caps.experimental or client_caps.experimental[exp_key] != exp_value:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
async def _receive_loop(self) -> None:
|
||||
async with self._incoming_message_stream_writer:
|
||||
await super()._receive_loop()
|
||||
|
||||
async def _received_request(self, responder: RequestResponder[types.ClientRequest, types.ServerResult]):
|
||||
match responder.request:
|
||||
case types.InitializeRequest(params=params):
|
||||
requested_version = params.protocol_version
|
||||
self._initialization_state = InitializationState.Initializing
|
||||
self._client_params = params
|
||||
with responder:
|
||||
await responder.respond(
|
||||
types.InitializeResult(
|
||||
protocol_version=requested_version
|
||||
if requested_version in SUPPORTED_PROTOCOL_VERSIONS
|
||||
else types.LATEST_PROTOCOL_VERSION,
|
||||
capabilities=self._init_options.capabilities,
|
||||
server_info=types.Implementation(
|
||||
name=self._init_options.server_name,
|
||||
title=self._init_options.title,
|
||||
description=self._init_options.description,
|
||||
version=self._init_options.server_version,
|
||||
website_url=self._init_options.website_url,
|
||||
icons=self._init_options.icons,
|
||||
),
|
||||
instructions=self._init_options.instructions,
|
||||
)
|
||||
)
|
||||
self._initialization_state = InitializationState.Initialized
|
||||
case types.PingRequest():
|
||||
# Ping requests are allowed at any time
|
||||
pass
|
||||
case _:
|
||||
if self._initialization_state != InitializationState.Initialized:
|
||||
raise RuntimeError("Received request before initialization was complete")
|
||||
|
||||
async def _received_notification(self, notification: types.ClientNotification) -> None:
|
||||
# Need this to avoid ASYNC910
|
||||
await anyio.lowlevel.checkpoint()
|
||||
match notification:
|
||||
case types.InitializedNotification():
|
||||
self._initialization_state = InitializationState.Initialized
|
||||
case _:
|
||||
if self._initialization_state != InitializationState.Initialized: # pragma: no cover
|
||||
raise RuntimeError("Received notification before initialization was complete")
|
||||
return self._connection.check_capability(capability)
|
||||
|
||||
async def send_log_message(
|
||||
self,
|
||||
|
|
@ -293,7 +209,7 @@ class ServerSession(
|
|||
"""
|
||||
if self._stateless:
|
||||
raise StatelessModeNotSupported(method="sampling")
|
||||
client_caps = self._client_params.capabilities if self._client_params else None
|
||||
client_caps = self.client_params.capabilities if self.client_params else None
|
||||
validate_sampling_tools(client_caps, tools, tool_choice)
|
||||
validate_tool_use_result_messages(messages)
|
||||
|
||||
|
|
@ -313,7 +229,6 @@ class ServerSession(
|
|||
)
|
||||
metadata_obj = ServerMessageMetadata(related_request_id=related_request_id)
|
||||
|
||||
# Use different result types based on whether tools are provided
|
||||
if tools is not None:
|
||||
return await self.send_request(
|
||||
request=request,
|
||||
|
|
@ -488,10 +403,3 @@ class ServerSession(
|
|||
),
|
||||
related_request_id,
|
||||
)
|
||||
|
||||
async def _handle_incoming(self, req: ServerRequestResponder) -> None:
|
||||
await self._incoming_message_stream_writer.send(req)
|
||||
|
||||
@property
|
||||
def incoming_messages(self) -> MemoryObjectReceiveStream[ServerRequestResponder]:
|
||||
return self._incoming_message_stream_reader
|
||||
|
|
|
|||
|
|
@ -633,7 +633,10 @@ class StreamableHTTPServerTransport:
|
|||
finally:
|
||||
await sse_stream_reader.aclose()
|
||||
|
||||
except Exception as err:
|
||||
except Exception as err: # pragma: lax no cover
|
||||
# Reached only when something raises during POST handling outside
|
||||
# the per-SSE-stream guard above; whether tests reach this depends
|
||||
# on client teardown timing.
|
||||
logger.exception("Error handling POST request")
|
||||
response = self._create_error_response(
|
||||
f"Error handling POST request: {err}",
|
||||
|
|
@ -641,7 +644,7 @@ class StreamableHTTPServerTransport:
|
|||
INTERNAL_ERROR,
|
||||
)
|
||||
await response(scope, receive, send)
|
||||
if writer: # pragma: no cover
|
||||
if writer:
|
||||
await writer.send(Exception(err))
|
||||
return # pragma: no cover
|
||||
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from mcp.server.streamable_http import (
|
|||
StreamableHTTPServerTransport,
|
||||
)
|
||||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
from mcp.shared._compat import resync_tracer
|
||||
from mcp.types import INVALID_REQUEST, ErrorData, JSONRPCError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -139,6 +140,7 @@ class StreamableHTTPSessionManager:
|
|||
# Clear any remaining server instances
|
||||
self._server_instances.clear()
|
||||
self._session_owners.clear()
|
||||
await resync_tracer()
|
||||
|
||||
async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
"""Process ASGI request with proper session handling and transport setup.
|
||||
|
|
|
|||
19
src/mcp/shared/_compat.py
Normal file
19
src/mcp/shared/_compat.py
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
"""Workarounds for CPython interpreter bugs the SDK papers over."""
|
||||
|
||||
import anyio.lowlevel
|
||||
|
||||
__all__ = ["resync_tracer"]
|
||||
|
||||
|
||||
async def resync_tracer() -> None:
|
||||
"""Resync coverage tracing after a cancelled task-group join.
|
||||
|
||||
A cancel delivered at a join resumes the awaiting coroutine chain via
|
||||
`coro.throw()`; on CPython 3.11 (python/cpython#106749) that drops the
|
||||
`'call'` trace events for the outer frames and desyncs coverage's CTracer
|
||||
until the chain next suspends and resumes normally. Yielding once here
|
||||
resumes via `.send()`, which re-stamps the missing events. Shielded so a
|
||||
pending outer cancel is not re-delivered at this point; behaviorally a
|
||||
no-op. Delete this module when Python 3.11 support ends (EOL 2027-10).
|
||||
"""
|
||||
await anyio.lowlevel.cancel_shielded_checkpoint()
|
||||
|
|
@ -20,9 +20,18 @@ def otel_span(
|
|||
kind: SpanKind,
|
||||
attributes: dict[str, Any] | None = None,
|
||||
context: Context | None = None,
|
||||
record_exception: bool = True,
|
||||
set_status_on_exception: bool = True,
|
||||
) -> Iterator[Any]:
|
||||
"""Create an OTel span."""
|
||||
with _tracer.start_as_current_span(name, kind=kind, attributes=attributes, context=context) as span:
|
||||
with _tracer.start_as_current_span(
|
||||
name,
|
||||
kind=kind,
|
||||
attributes=attributes,
|
||||
context=context,
|
||||
record_exception=record_exception,
|
||||
set_status_on_exception=set_status_on_exception,
|
||||
) as span:
|
||||
yield span
|
||||
|
||||
|
||||
|
|
@ -31,6 +40,13 @@ def inject_trace_context(meta: dict[str, Any]) -> None:
|
|||
inject(meta)
|
||||
|
||||
|
||||
def extract_trace_context(meta: dict[str, Any]) -> Context:
|
||||
"""Extract W3C trace context from a `_meta` dict."""
|
||||
return extract(meta)
|
||||
def extract_trace_context(meta: dict[str, Any]) -> Context | None:
|
||||
"""Extract W3C trace context from a `_meta` dict.
|
||||
|
||||
Returns `None` when the carrier is malformed; telemetry parsing must
|
||||
never fail the request it annotates.
|
||||
"""
|
||||
try:
|
||||
return extract(meta)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
|
|
|||
85
src/mcp/shared/context.py
Normal file
85
src/mcp/shared/context.py
Normal file
|
|
@ -0,0 +1,85 @@
|
|||
"""`BaseContext` - the user-facing per-request context.
|
||||
|
||||
Composition over a `DispatchContext`: forwards the transport metadata, the
|
||||
back-channel (`send_raw_request`/`notify`), progress reporting, and the cancel
|
||||
event. Adds `meta` (the inbound request's `_meta` field).
|
||||
|
||||
Satisfies `Outbound`, so `ClientPeer` can wrap it. Shared between client and
|
||||
server: the server's `Context` extends this with `lifespan`/`connection`;
|
||||
`ClientContext` is just an alias.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Generic
|
||||
|
||||
import anyio
|
||||
from typing_extensions import TypeVar
|
||||
|
||||
from mcp.shared.dispatcher import CallOptions, DispatchContext
|
||||
from mcp.shared.transport_context import TransportContext
|
||||
from mcp.types import RequestParamsMeta
|
||||
|
||||
__all__ = ["BaseContext"]
|
||||
|
||||
TransportT = TypeVar("TransportT", bound=TransportContext, default=TransportContext, covariant=True)
|
||||
|
||||
|
||||
class BaseContext(Generic[TransportT]):
|
||||
"""Per-request context wrapping a `DispatchContext`.
|
||||
|
||||
`ServerRunner` constructs one per inbound request and passes it to the
|
||||
user's handler.
|
||||
"""
|
||||
|
||||
def __init__(self, dctx: DispatchContext[TransportT], meta: RequestParamsMeta | None = None) -> None:
|
||||
self._dctx = dctx
|
||||
self._meta = meta
|
||||
|
||||
@property
|
||||
def transport(self) -> TransportT:
|
||||
"""Transport-specific metadata for this inbound request."""
|
||||
return self._dctx.transport
|
||||
|
||||
@property
|
||||
def cancel_requested(self) -> anyio.Event:
|
||||
"""Set when the peer sends `notifications/cancelled` for this request."""
|
||||
return self._dctx.cancel_requested
|
||||
|
||||
@property
|
||||
def can_send_request(self) -> bool:
|
||||
"""Whether the back-channel can currently deliver server-initiated requests.
|
||||
|
||||
`False` when the transport has no back-channel, or when the underlying
|
||||
dispatch context has been closed because the inbound request finished.
|
||||
"""
|
||||
return self._dctx.can_send_request
|
||||
|
||||
@property
|
||||
def meta(self) -> RequestParamsMeta | None:
|
||||
"""The inbound request's `_meta` field, if present."""
|
||||
return self._meta
|
||||
|
||||
async def send_raw_request(
|
||||
self,
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
opts: CallOptions | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Send a request to the peer on the back-channel.
|
||||
|
||||
Raises:
|
||||
MCPError: The peer responded with an error.
|
||||
NoBackChannelError: `can_send_request` is `False`.
|
||||
"""
|
||||
return await self._dctx.send_raw_request(method, params, opts)
|
||||
|
||||
async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
||||
"""Send a notification to the peer on the back-channel."""
|
||||
await self._dctx.notify(method, params)
|
||||
|
||||
async def report_progress(self, progress: float, total: float | None = None, message: str | None = None) -> None:
|
||||
"""Report progress for this request, if the peer supplied a progress token.
|
||||
|
||||
A no-op when no token was supplied.
|
||||
"""
|
||||
await self._dctx.progress(progress, total, message)
|
||||
197
src/mcp/shared/direct_dispatcher.py
Normal file
197
src/mcp/shared/direct_dispatcher.py
Normal file
|
|
@ -0,0 +1,197 @@
|
|||
"""In-memory `Dispatcher` that wires two peers together with no transport.
|
||||
|
||||
`DirectDispatcher` is the simplest possible `Dispatcher` implementation: a
|
||||
request on one side directly invokes the other side's `on_request`. There is no
|
||||
serialization, no JSON-RPC framing, and no streams. It exists to:
|
||||
|
||||
* prove the `Dispatcher` Protocol is implementable without JSON-RPC
|
||||
* provide a fast substrate for testing the layers above the dispatcher
|
||||
(`ServerRunner`, `Context`, `Connection`) without wire-level moving parts
|
||||
* embed a server in-process when the JSON-RPC overhead is unnecessary
|
||||
|
||||
Unlike `JSONRPCDispatcher`, exceptions raised in a handler propagate directly
|
||||
to the caller - there is no exception-to-`ErrorData` boundary here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import anyio
|
||||
import anyio.abc
|
||||
from pydantic import ValidationError
|
||||
|
||||
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
|
||||
|
||||
__all__ = ["DirectDispatcher", "create_direct_dispatcher_pair"]
|
||||
|
||||
DIRECT_TRANSPORT_KIND = "direct"
|
||||
|
||||
|
||||
_Request = Callable[[str, Mapping[str, Any] | None, CallOptions | None], Awaitable[dict[str, Any]]]
|
||||
_Notify = Callable[[str, Mapping[str, Any] | None], Awaitable[None]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _DirectDispatchContext:
|
||||
"""`DispatchContext` for an inbound request on a `DirectDispatcher`.
|
||||
|
||||
The back-channel callables target the *originating* side, so a handler's
|
||||
`send_raw_request` reaches the peer that made the inbound request.
|
||||
"""
|
||||
|
||||
transport: TransportContext
|
||||
_back_request: _Request
|
||||
_back_notify: _Notify
|
||||
request_id: RequestId | None = None
|
||||
"""Always `None`: direct dispatch has no wire-level request id."""
|
||||
message_metadata: MessageMetadata = None # TODO(maxisbey): remove for Context rework
|
||||
"""Always `None`: in-memory dispatch attaches no transport metadata."""
|
||||
_on_progress: ProgressFnT | None = None
|
||||
cancel_requested: anyio.Event = field(default_factory=anyio.Event)
|
||||
|
||||
@property
|
||||
def can_send_request(self) -> bool:
|
||||
return self.transport.can_send_request
|
||||
|
||||
async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
||||
await self._back_notify(method, params)
|
||||
|
||||
async def send_raw_request(
|
||||
self,
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
opts: CallOptions | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if not self.can_send_request:
|
||||
raise NoBackChannelError(method)
|
||||
return await self._back_request(method, params, opts)
|
||||
|
||||
async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None:
|
||||
if self._on_progress is not None:
|
||||
await self._on_progress(progress, total, message)
|
||||
|
||||
|
||||
class DirectDispatcher:
|
||||
"""A `Dispatcher` that calls a peer's handlers directly, in-process.
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
def __init__(self, transport_ctx: TransportContext):
|
||||
self._transport_ctx = transport_ctx
|
||||
self._peer: DirectDispatcher | None = None
|
||||
self._on_request: OnRequest | None = None
|
||||
self._on_notify: OnNotify | None = None
|
||||
self._ready = anyio.Event()
|
||||
self._closed = anyio.Event()
|
||||
|
||||
def connect_to(self, peer: DirectDispatcher) -> None:
|
||||
self._peer = peer
|
||||
|
||||
async def send_raw_request(
|
||||
self,
|
||||
method: str,
|
||||
params: Mapping[str, Any] | None,
|
||||
opts: CallOptions | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if self._peer is None:
|
||||
raise RuntimeError("DirectDispatcher has no peer; use create_direct_dispatcher_pair()")
|
||||
return await self._peer._dispatch_request(method, params, opts)
|
||||
|
||||
async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
||||
if self._peer is None:
|
||||
raise RuntimeError("DirectDispatcher has no peer; use create_direct_dispatcher_pair()")
|
||||
await self._peer._dispatch_notify(method, params)
|
||||
|
||||
async def run(
|
||||
self,
|
||||
on_request: OnRequest,
|
||||
on_notify: OnNotify,
|
||||
*,
|
||||
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()
|
||||
|
||||
def close(self) -> None:
|
||||
self._closed.set()
|
||||
|
||||
def _make_context(self, on_progress: ProgressFnT | 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),
|
||||
_on_progress=on_progress,
|
||||
)
|
||||
|
||||
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")):
|
||||
try:
|
||||
return await self._on_request(dctx, method, params)
|
||||
except MCPError:
|
||||
raise
|
||||
except ValidationError as e:
|
||||
# Same shape JSONRPCDispatcher writes, so runner-over-direct
|
||||
# tests see what runner-over-JSONRPC would.
|
||||
raise MCPError(code=INVALID_PARAMS, message="Invalid request parameters", data="") from e
|
||||
except Exception as e:
|
||||
raise MCPError(code=INTERNAL_ERROR, message=str(e)) from e
|
||||
except TimeoutError:
|
||||
raise MCPError(
|
||||
code=REQUEST_TIMEOUT,
|
||||
message=f"Timed out after {opts.get('timeout')}s waiting for {method!r}",
|
||||
) from None
|
||||
|
||||
async def _dispatch_notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
||||
await self._ready.wait()
|
||||
assert self._on_notify is not None
|
||||
dctx = self._make_context()
|
||||
await self._on_notify(dctx, method, params)
|
||||
|
||||
|
||||
def create_direct_dispatcher_pair(
|
||||
*,
|
||||
can_send_request: bool = True,
|
||||
headers: Mapping[str, str] | None = None,
|
||||
) -> tuple[DirectDispatcher, DirectDispatcher]:
|
||||
"""Create two `DirectDispatcher` instances wired to each other.
|
||||
|
||||
Args:
|
||||
can_send_request: Sets `TransportContext.can_send_request` on both
|
||||
sides. Pass `False` to simulate a transport with no back-channel.
|
||||
headers: Sets `TransportContext.headers` on both sides.
|
||||
|
||||
Returns:
|
||||
A `(client, server)` pair. The wiring is symmetric, so the roles
|
||||
are conventional only.
|
||||
"""
|
||||
ctx = TransportContext(kind=DIRECT_TRANSPORT_KIND, can_send_request=can_send_request, headers=headers)
|
||||
client = DirectDispatcher(ctx)
|
||||
server = DirectDispatcher(ctx)
|
||||
client.connect_to(server)
|
||||
server.connect_to(client)
|
||||
return client, server
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Reference in a new issue