[v2] Dispatcher/ServerRunner receive-path swap — replaces BaseSession#2710
Draft
maxisbey wants to merge 27 commits into
Draft
[v2] Dispatcher/ServerRunner receive-path swap — replaces BaseSession#2710maxisbey wants to merge 27 commits into
maxisbey wants to merge 27 commits into
Conversation
Introduces the Dispatcher abstraction that decouples MCP request/response handling from JSON-RPC framing. A Dispatcher exposes call/notify for outbound messages and run(on_call, on_notify) for inbound dispatch, with no knowledge of MCP types or wire encoding. - shared/dispatcher.py: Dispatcher, DispatchContext, RequestSender Protocols; CallOptions, OnCall/OnNotify, ProgressFnT, DispatchMiddleware - shared/transport_context.py: TransportContext base dataclass - shared/direct_dispatcher.py: in-memory Dispatcher impl that wires two peers with no transport; serves as a fast test substrate and second-impl proof - shared/exceptions.py: NoBackChannelError(MCPError) for transports without a server-to-client request channel - types: REQUEST_CANCELLED SDK error code The JSON-RPC implementation and ServerRunner that consume this Protocol land in follow-up PRs.
- tests: replace unreachable 'return {}' with 'raise NotImplementedError'
(already in coverage exclude_also) and collapse send_request+return into
one statement
- dispatcher: RequestSender docstring no longer claims Dispatcher satisfies it
(Dispatcher exposes call(), not send_request())
…er with Outbound The design doc's `send_request = call` alias only makes the concrete class satisfy RequestSender, not the abstract Dispatcher Protocol — so any consumer typed against `Dispatcher[TT]` (Connection, ServerRunner) couldn't pass it to something expecting a RequestSender without a cast or hand-written bridge. RequestSender was also half a contract: every implementor (Dispatcher, DispatchContext, Connection, Context) has `notify` too, and PeerMixin needs both for its typed sugar (elicit/sample are requests, log is a notification). Outbound(Protocol) declares both methods; Dispatcher and DispatchContext extend it. PeerMixin will wrap an Outbound. One verb everywhere, no aliases, no extra Protocols. - Dispatcher.call -> send_request - OnCall -> OnRequest, on_call -> on_request - RequestSender -> Outbound (now also declares notify) - Dispatcher(Outbound, Protocol[TT]), DispatchContext(Outbound, Protocol[TT])
The dispatcher-layer raw channel is now `send_raw_request(method, params) -> dict`. This frees the `send_request` name for the typed surface (`send_request(req: Request) -> Result`) that Connection/Context/Client add in later PRs. Mechanical rename across Outbound, Dispatcher, DispatchContext, DirectDispatcher, _DirectDispatchContext, and all tests. `can_send_request` (the transport capability flag) is unchanged — it names the capability, not the method.
Chunk (a) of JSONRPCDispatcher: constructor, _Pending/_InFlight/_JSONRPCDispatchContext, send_request/notify and helpers. run() is stubbed. The Dispatcher contract tests are now parametrized over a pair_factory fixture (direct + jsonrpc). The 9 jsonrpc cases are strict-xfail until run()/ _handle_request land in the next commits; once those pass, strict xfail flips to XPASS and forces removal of the marker. Factories return (client, server, close) so running_pair can shut down any implementation uniformly.
run() drives the receive loop in a per-request task group; task_status.started() fires once send_request is usable. _dispatch routes each inbound message synchronously (no awaits — send_nowait/_spawn only) to avoid head-of-line blocking. _spawn propagates the sender's contextvars via Context.run(tg.start_soon, ...) so auth/OTel set by ASGI middleware survive. _fan_out_closed wakes pending send_request waiters with CONNECTION_CLOSED on shutdown (called both post-EOF and in finally; idempotent). Wire-param extraction (progressToken, cancelled.requestId, progress fields) uses structural match patterns — runtime narrowing, no casts, no mcp.types model coupling; malformed input fails to match and the correlation is skipped. _handle_request is happy-path only here (run on_request, write response); the exception-to-wire boundary lands in the next commit. Dispatcher.run() Protocol gained a task_status kwarg (it's a contract-level guarantee). DirectDispatcher.run() updated to match. running_pair now uses tg.start so the test body runs only once the dispatcher is ready. 20 contract tests pass; the 2 needing the exception boundary are strict-xfail.
_handle_request is now the single exception-to-wire boundary: - MCPError -> JSONRPCError(e.error) - pydantic ValidationError -> INVALID_PARAMS - Exception -> INTERNAL_ERROR(str(e)), logged, optionally re-raised - outer-cancel (run() TG shutdown) -> shielded REQUEST_CANCELLED write, re-raise - peer-cancel (notifications/cancelled) -> scope swallows, no response written dctx.close() runs in an inner finally so the back-channel shuts the moment the handler exits. _write_result/_write_error swallow Broken/ClosedResourceError so a dropped connection during the response write doesn't crash the dispatcher. All 22 contract tests now pass against both DirectDispatcher and JSONRPCDispatcher; chunk-c xfail markers removed.
Covers behaviors with no DirectDispatcher analog: out-of-order response correlation, INTERNAL_ERROR over the wire, peer-cancel in interrupt and signal modes, CONNECTION_CLOSED on stream EOF mid-await, late-response drop, raise_handler_exceptions propagation, ServerMessageMetadata tagging on ctx.send_request, null-id JSONRPCError drop, ValidationError->INVALID_PARAMS, contextvar propagation via _spawn, and the defensive Broken/Closed/WouldBlock catches. Two small src tweaks for coverage: - _cancel_outbound: combine the two except arms into one tuple - _dispatch: pragma no-branch on the final case (match is exhaustive over JSONRPCMessage; the no-match arc is unreachable) 43 tests, 100% coverage on all PR2 modules, 0.15s wall-clock.
The pull_request branch filter meant the test/lint/coverage matrix only ran on PRs targeting main or v1.x. Stacked PRs (targeting feature branches) only got the conformance checks, which are continue-on-error and don't exercise unit tests. Removing the filter so the full matrix runs on every PR.
3.14: nested async-with arc misreporting on three create_task_group lines (the documented AGENTS.md case) — pragma: no branch. 3.11: lines after async-CM exit with pytest.raises mis-traced in one test — moved the asserts inside the context manager.
Follows the Outbound Protocol rename in the previous commit. Mechanical rename across JSONRPCDispatcher, _JSONRPCDispatchContext, and tests.
PeerMixin defines the typed server-to-client request methods (sample with overloads, elicit_form, elicit_url, list_roots, ping) once. Each method constrains `self: Outbound` so any class with send_request/notify can mix it in — pyright checks the host structurally at the call site. The mixin does no capability gating; that's the host's send_request's job. Peer is a trivial standalone wrapper for when you have a bare Outbound (e.g. a dispatcher) and want the typed sugar without writing your own host class. 6 tests over DirectDispatcher, 0.03s.
Composition over a DispatchContext: forwards transport/cancel_requested/ send_request/notify/progress and adds meta. Satisfies Outbound so PeerMixin works on it (proven by Peer(bctx).ping() round-tripping). The server Context (next commit) extends this with lifespan/connection; ClientContext will be an alias once ClientSession is reworked.
…ntext PeerMixin methods and Peer/BaseContext now call/expose send_raw_request. The typed send_request lands on Connection/Context in the next commit.
TypedServerRequestMixin (server/_typed_request.py) provides shape-2 typed send_request: per-spec overloads (CreateMessage/Elicit/ListRoots/Ping) infer the result type; custom requests pass result_type explicitly. Mixed into both Connection and the server Context. Connection (server/connection.py) wraps an Outbound for the standalone stream. notify is best-effort (never raises); send_raw_request gated on has_standalone_channel; check_capability mirrors v1 for now (FOLLOWUP). Holds peer info populated at initialize time and the per-connection lifespan state. Context (server/context.py, alongside v1's ServerRequestContext) composes BaseContext + PeerMixin + TypedServerRequestMixin and adds lifespan/connection. Request-scoped log() rides the request's back-channel; ctx.connection.log() uses the standalone stream. dump_params(model, meta) merges user-supplied meta into _meta; threaded through every PeerMixin and Connection convenience method. 31 tests, 0.06s.
- Connection.check_capability per-field branches (parametrized) - Context.log with logger and meta supplied - Peer.notify forwards to wrapped Outbound
coverage.py on Python 3.11 doesn't record statements after an 'async with running_pair(...)' exit when there's a nested 'with anyio.fail_after()' inside. Same workaround as 0a8f0f4 in PR2 — move the asserts inside the async-with block.
Remove references to PR numbers, internal scratch notes, and design-spike shorthand that won't make sense to a fresh reader of the codebase.
LifespanT and TransportT are only exposed via read-only properties (lifespan, transport), so covariance is sound. This lets a Context[AppState, HttpTC] be passed where a Context[object, TransportContext] is expected — needed for ServerRunner's middleware chain to compose without casts, and for reusable middleware to be typed Context[object, TransportContext] instead of relying on Any-slack.
ServerRunner is the per-connection orchestrator over a Dispatcher. This commit lands the skeleton: ServerRegistry Protocol, _on_request (lookup → validate → build Context → call handler → dump), _handle_initialize (populates Connection, opens the init-gate), and a basic _on_notify. Additive methods on lowlevel Server (get_request_handler / get_notification_handler / middleware / connection_lifespan) so it satisfies ServerRegistry without touching the existing run() path. _PARAMS_FOR_METHOD is scaffolding (marked TODO) until the registry stores params types directly. 5 tests over DirectDispatcher + a real lowlevel Server.
ContextMiddleware is a Protocol[L] (contravariant) so Server[L].middleware: list[ContextMiddleware[L]] is properly typed. App-specific middleware sees ctx.lifespan: L; reusable middleware typed ContextMiddleware[object] registers on any Server via contravariance. Context's covariance (previous PR3 commit) makes Context[L, ST] <: Context[L, TransportContext] so the chain composes without casts. dispatch_middleware (DispatchMiddleware list on ServerRunner) wraps the raw _on_request and sees everything including initialize/METHOD_NOT_FOUND. server.middleware (ContextMiddleware) runs inside _on_request after validation/ctx-build and wraps registered handlers only. _on_notify routes notifications/initialized (sets the flag), drops before-init and unknown methods, otherwise builds Context and calls the registered handler. 11 tests over DirectDispatcher + a real lowlevel Server.
run() composes dispatch_middleware over _on_request and forwards task_status to dispatcher.run() so callers can 'await tg.start(runner.run)'. otel_middleware is a DispatchMiddleware that wraps each request in a span, mirroring the existing Server._handle_request span shape: name 'MCP handle <method> [<target>]', mcp.method.name attribute, W3C trace context extracted from params._meta (SEP-414), and ERROR status if the handler raises. connection_lifespan plumbing (the enter-late dance) is deferred to a separate commit since Server.connection_lifespan is None today.
…d_runner harness
- Add opentelemetry-sdk as a dev dep and a tests/server/conftest.py 'spans'
fixture (TracerProvider + InMemorySpanExporter) so otel_middleware's span
contract is observable.
- Replace the otel pass-through test with four span-asserting tests (name +
target, _meta traceparent → parent, MCPError → ERROR status without
traceback, unexpected exception → ERROR status + exception event). These
surfaced that start_as_current_span's default set_status_on_exception /
record_exception was overwriting the middleware's explicit set_status and
attaching tracebacks to protocol-level MCPErrors — now disabled and handled
explicitly.
- Add handler-return contract tests (None → {}, unsupported → INTERNAL_ERROR).
- Introduce connected_runner async-contextmanager test harness and retrofit all
tests through runner.run(); drop two tests made redundant by that. Harness
closes dispatchers gracefully and re-raises body exceptions outside the task
group so failures aren't ExceptionGroup-wrapped (and to avoid a coverage.py
trace-loss false-negative on cancel-during-aexit).
- Remove the unused Server.connection_lifespan placeholder; it lands with its
consumer.
The previous tests/server/conftest.py called trace.set_tracer_provider() directly, which is set-once per process and raced against logfire's capfire fixture (tests/shared/test_otel.py) under xdist — whichever ran first in a worker won, the other's tests broke. Converge on capfire as the single span-capture owner since logfire.configure() already handles repeat calls by swapping span processors instead of re-setting the provider: - tests/conftest.py: set LOGFIRE_DISTRIBUTED_TRACING=true so propagation tests don't trip logfire's 'found propagated trace context' RuntimeWarning. - tests/server/conftest.py: SpanCapture adapter over capfire.exporter — filters to the mcp-python-sdk instrumentation scope and excludes logfire's pending_span markers, so tests assert on raw ReadableSpan without importing logfire types. - tests/shared/test_otel.py: drop the now-unneeded filterwarnings decorator.
…er[L] directly Server is generic in LifespanResultT only — no TransportContextT. Spike (scratch/spike-tt-on-server) found a third generic breaks bare-Server plumbing helpers via invariance and only buys one None-check; it remains additive later via PEP 696 default if demand materialises. TT stays on the transport layer (Dispatcher/DispatchContext/BaseContext in mcp.shared); the server layer (Server/Context/ServerRunner/ServerMiddleware) consumes base TransportContext. - HandlerEntry[L] frozen dataclass (params_type, handler) replaces bare callables in the registry; params type erased to Any in storage, correlated at add_request_handler[P] - Public add_request_handler/add_notification_handler; capabilities() zero-arg (notification_options/experimental_capabilities now ctor kwargs) - ServerRunner drops the ServerRegistry Protocol scaffold and reads Server[L] directly; _make_context no longer narrows dctx - ServerMiddleware[L] (one contravariant param) - Context[L] (BaseContext[TransportContext] fixed)
…tContext.headers Per-connection state without a connection_lifespan CM or a second Server generic. Stateless is the default deployment, where a per-connection lifespan would wrap a single request; the enter-late mechanics it would need (race init vs dispatcher-done, ready-gate) were more machinery than the use case warrants. - Connection.session_id: str | None — set by the mount via ServerRunner(session_id=...); per-connection, not per-message - Connection.state: dict[str, Any] — scratch that persists across requests; handlers/middleware read and write freely - Connection.exit_stack: AsyncExitStack — handlers/middleware push CMs or callbacks for per-connection teardown; ServerRunner.run() unwinds it (shielded) in a finally after dispatcher.run() returns - TransportContext.headers: Mapping[str, str] | None on the base — populated by HTTP transports, None on stdio - Context.session_id / Context.headers convenience properties - create_direct_dispatcher_pair(headers=...) and connected_runner(session_id=..., headers=...) for tests
…r correlation Matches BaseSession._normalize_request_id and the TypeScript SDK: a peer that echoes the request ID as a JSON string still resolves the waiter. Applied at both lookup sites (_resolve_pending and the progress-token match). Parity prep for the PR6 e2e suite.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Status: WIP — opened as a base for the swap work; commits will accumulate here.
V2 server-side receive path:
Transport → JSONRPCDispatcher → ServerRunner → Serverregistry, replacing theBaseSession/ServerSessionmessage loop. Consolidates the previously-stacked PR #2562 (and the four below it) into a single PR targetingmain.Goal
The
tests/interaction/suite (#2691) is the bar: this PR is done when that suite passes on the new runtime path withBaseSession/ServerSessionremoved and the handler-facing context object shimmed to the existingServerRequestContextsurface.What's already here (from the consolidated stack)
Dispatcher/DispatchContext/OutboundProtocols (mcp/shared/dispatcher.py)JSONRPCDispatcherover the existingSessionMessagestream contract — request-ID correlation, per-request task isolation, cancel/progress interception, exception→wire boundaryServerRunner[L]per-connection orchestrator;Connection(state, exit_stack, session_id); newContext[L]Server[L]registry:HandlerEntrydataclass,add_request_handler, zero-argcapabilities()DirectDispatcherin-memory pair;ServerMiddleware/DispatchMiddlewareTransportContext.headers;ctx.session_id/ctx.headerspropertiesWhat's coming
Server.run()rewritten to driveJSONRPCDispatcher+ServerRunnerStreamableHTTPSessionManager/SseServerTransportroute throughServerRunnerServerRequestContextcompat shim so the 200+ handler annotations intests/interaction/keep typechecking and the runtime object satisfies the surface_metaparity onsend_raw_requestBaseSession/ServerSession/ the old_handle_*pathSupersedes
#2562 and the four stacked below it. Those stay open as drafts for reference until this is reviewable.
AI Disclaimer