Skip to content

Commit 535d621

Browse files
committed
feat: JSONRPCDispatcher.inline_methods; otel_middleware sets jsonrpc.request.id
inline_methods: request methods in this set are awaited inline in the read loop instead of spawned, so their side effects are visible to the next dequeued message. Server.run() passes {"initialize"} so a client that pipelines initialize + the next request without awaiting InitializeResult (spec says SHOULD NOT, not MUST NOT) sees the initialized state instead of failing the init-gate. Matches the go-sdk's explicit carve-out. _dispatch / _dispatch_request are now async (only await for inline methods). otel_middleware: restore the jsonrpc.request.id span attribute that the previous Server._handle_request set.
1 parent 07d94ee commit 535d621

5 files changed

Lines changed: 66 additions & 9 deletions

File tree

‎src/mcp/server/lowlevel/server.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -396,6 +396,10 @@ async def run(
396396
read_stream,
397397
write_stream,
398398
raise_handler_exceptions=raise_exceptions,
399+
# Handle `initialize` inline so a client that pipelines it with
400+
# the next request (spec says SHOULD NOT, not MUST NOT) sees
401+
# the initialized state instead of failing the init-gate.
402+
inline_methods=frozenset({"initialize"}),
399403
)
400404
runner = ServerRunner(
401405
server=self,

‎src/mcp/server/runner.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,10 +94,13 @@ async def wrapped(
9494
case _:
9595
parent = None
9696
span_name = f"MCP handle {method}{f' {target}' if target else ''}"
97+
attributes: dict[str, str | int] = {"mcp.method.name": method}
98+
if dctx.request_id is not None:
99+
attributes["jsonrpc.request.id"] = dctx.request_id
97100
with otel_span(
98101
span_name,
99102
kind=SpanKind.SERVER,
100-
attributes={"mcp.method.name": method},
103+
attributes=attributes,
101104
context=parent,
102105
record_exception=False,
103106
set_status_on_exception=False,

‎src/mcp/shared/jsonrpc_dispatcher.py‎

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -213,6 +213,7 @@ def __init__(
213213
*,
214214
peer_cancel_mode: PeerCancelMode = "interrupt",
215215
raise_handler_exceptions: bool = False,
216+
inline_methods: frozenset[str] = frozenset(),
216217
) -> None: ...
217218
@overload
218219
def __init__(
@@ -223,6 +224,7 @@ def __init__(
223224
transport_builder: Callable[[MessageMetadata], TransportT],
224225
peer_cancel_mode: PeerCancelMode = "interrupt",
225226
raise_handler_exceptions: bool = False,
227+
inline_methods: frozenset[str] = frozenset(),
226228
) -> None: ...
227229
def __init__(
228230
self,
@@ -232,6 +234,7 @@ def __init__(
232234
transport_builder: Callable[[MessageMetadata], TransportT] | None = None,
233235
peer_cancel_mode: PeerCancelMode = "interrupt",
234236
raise_handler_exceptions: bool = False,
237+
inline_methods: frozenset[str] = frozenset(),
235238
) -> None:
236239
self._read_stream = read_stream
237240
self._write_stream = write_stream
@@ -244,6 +247,13 @@ def __init__(
244247
)
245248
self._peer_cancel_mode: PeerCancelMode = peer_cancel_mode
246249
self._raise_handler_exceptions = raise_handler_exceptions
250+
self._inline_methods = inline_methods
251+
"""Request methods handled inline in the read loop (awaited before the
252+
next message is dequeued) instead of spawned concurrently. Use for
253+
methods whose side effects must be observable to the next message,
254+
e.g. `initialize`, so a pipelined follow-up sees the initialized state.
255+
Only suitable for handlers that complete quickly, since inline handling
256+
blocks dequeuing."""
247257

248258
self._next_id = 0
249259
self._pending: dict[RequestId, _Pending] = {}
@@ -384,7 +394,7 @@ async def run(
384394
sender_ctx: contextvars.Context | None = getattr(
385395
self._read_stream, "last_context", None
386396
)
387-
self._dispatch(item, on_request, on_notify, sender_ctx)
397+
await self._dispatch(item, on_request, on_notify, sender_ctx)
388398
except anyio.ClosedResourceError:
389399
# The transport closed our receive end and we looped
390400
# back to `__anext__` on the now-closed stream
@@ -409,17 +419,19 @@ async def run(
409419
self._tg = None
410420
self._fan_out_closed()
411421

412-
def _dispatch(
422+
async def _dispatch(
413423
self,
414424
item: SessionMessage | Exception,
415425
on_request: OnRequest,
416426
on_notify: OnNotify,
417427
sender_ctx: contextvars.Context | None,
418428
) -> None:
419-
"""Route one inbound item. Synchronous: never awaits.
429+
"""Route one inbound item.
420430
421-
Everything here is `send_nowait` or `_spawn`. An `await` would let one
422-
slow message head-of-line block the entire read loop.
431+
Everything here is `send_nowait` or `_spawn`; the only `await` is for
432+
`inline_methods` requests, which deliberately block dequeuing until
433+
handled. Any other `await` would let one slow message head-of-line
434+
block the entire read loop.
423435
"""
424436
if isinstance(item, Exception):
425437
logger.debug("transport yielded exception: %r", item)
@@ -428,7 +440,7 @@ def _dispatch(
428440
msg = item.message
429441
match msg:
430442
case JSONRPCRequest():
431-
self._dispatch_request(msg, metadata, on_request, sender_ctx)
443+
await self._dispatch_request(msg, metadata, on_request, sender_ctx)
432444
case JSONRPCNotification():
433445
self._dispatch_notification(msg, metadata, on_notify, sender_ctx)
434446
case JSONRPCResponse():
@@ -439,7 +451,7 @@ def _dispatch(
439451
# on this final case is unreachable.
440452
self._resolve_pending(msg.id, msg.error)
441453

442-
def _dispatch_request(
454+
async def _dispatch_request(
443455
self,
444456
req: JSONRPCRequest,
445457
metadata: MessageMetadata,
@@ -462,7 +474,10 @@ def _dispatch_request(
462474
)
463475
scope = anyio.CancelScope()
464476
self._in_flight[req.id] = _InFlight(scope=scope, dctx=dctx)
465-
self._spawn(self._handle_request, req, dctx, scope, on_request, sender_ctx=sender_ctx)
477+
if req.method in self._inline_methods:
478+
await self._handle_request(req, dctx, scope, on_request)
479+
else:
480+
self._spawn(self._handle_request, req, dctx, scope, on_request, sender_ctx=sender_ctx)
466481

467482
def _dispatch_notification(
468483
self,

‎tests/server/test_runner.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -411,6 +411,7 @@ async def call_tool(ctx: Ctx, params: CallToolRequestParams) -> dict[str, Any]:
411411
assert span.name == "MCP handle tools/call mytool"
412412
assert span.attributes is not None
413413
assert span.attributes["mcp.method.name"] == "tools/call"
414+
assert isinstance(span.attributes["jsonrpc.request.id"], int)
414415
assert span.status.status_code == StatusCode.UNSET
415416

416417

‎tests/shared/test_jsonrpc_dispatcher.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from typing import Any
1212

1313
import anyio
14+
import anyio.lowlevel
1415
import pytest
1516

1617
from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream
@@ -461,6 +462,39 @@ async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] |
461462
assert "progress callback raised" in caplog.text
462463

463464

465+
@pytest.mark.anyio
466+
async def test_inline_methods_are_handled_before_next_message_is_dequeued():
467+
"""A method in `inline_methods` runs to completion before subsequent
468+
messages are dispatched, so its side effects are visible to them."""
469+
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
470+
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
471+
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(
472+
c2s_recv, s2c_send, inline_methods=frozenset({"first"})
473+
)
474+
state = {"initialized": False}
475+
seen: list[bool] = []
476+
477+
async def on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
478+
if method == "first":
479+
await anyio.lowlevel.checkpoint()
480+
state["initialized"] = True
481+
else:
482+
seen.append(state["initialized"])
483+
return {}
484+
485+
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
486+
raise NotImplementedError
487+
488+
# Buffer both requests before run() reads anything.
489+
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="first", params=None)))
490+
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=2, method="second", params=None)))
491+
c2s_send.close()
492+
with anyio.fail_after(5):
493+
await server.run(on_request, on_notify)
494+
assert seen == [True]
495+
s2c_recv.close()
496+
497+
464498
@pytest.mark.anyio
465499
async def test_send_raw_request_always_carries_meta_on_the_wire():
466500
"""Outbound requests always include `params._meta` (otel injection per SEP-414).

0 commit comments

Comments
 (0)