forked from modelcontextprotocol/python-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathcontext.py
More file actions
165 lines (130 loc) · 6.55 KB
/
context.py
File metadata and controls
165 lines (130 loc) · 6.55 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
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 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(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: ...