diff --git a/TRANSPORT_API_PLAN.md b/TRANSPORT_API_PLAN.md new file mode 100644 index 0000000000..580210a3e7 --- /dev/null +++ b/TRANSPORT_API_PLAN.md @@ -0,0 +1,298 @@ +# Extensible transport API plan + +Status: implementation in progress. New APIs still need final compatibility and native-binding review before release. + +## Review layout + +The review stack separates the SDK transport APIs, native gRPC adapter, MQTT adapter, and AMQP adapter into four pull requests, in that order. This document tracks the whole initiative; the SDK pull request contains no network adapter implementation or optional adapter dependencies. Each adapter pull request targets the preceding branch so its diff contains only that layer. + +## Progress + +| Work | Current evidence | Remaining | +| --- | --- | --- | +| Shared message contract | `mcp.shared.transport` exports stream, message, metadata, and transport types; existing client imports remain available | External adapter validation | +| Handler metadata | `TransportContext` reaches both handler APIs; concurrent principal-bound state, forged `_meta`, cross-peer replay, and missing-identity rejection are tested through the existing policy hook | Broker-specific identity-denial and cancellation scenarios | +| Multi-client lifecycle | Core tests cover shared lifespan, peer isolation, native readiness, and preservation of listener/startup failures across cleanup timeouts. `DirectDispatcher` now tracks operations on both peers and joins nested calls and handler cleanup; independent static review found no actionable issues | Final combined API and lifecycle review | +| Native dispatcher entry | Real gRPC binding uses protobuf envelopes with JSON payloads and no JSON-RPC frames. Maintained tests cover malformed traffic, capacity, cancellation, shutdown, error fidelity, subscriptions, and TLS. Both server APIs and multi-round trips also pass live examples on Python 3.10 and 3.14 | Final contract review, cross-platform validation, and the documented single-event-loop restriction | +| MQTT and AMQP adapters | Separate `examples/transports` package; both demos pass real two-peer calls for both server APIs and three client modes on Python 3.10 and 3.14; RabbitMQ cross-peer consumption is denied and borrowed channels survive cancellation | Full delivery/failure checks, MQTT provider limitations, record/replay, and adapter coverage | +| Release gates | Core: 6,027 tests pass with 100% branch coverage and `strict-no-cover`. Adapter package: 60 tests pass on Python 3.10 and 3.14. Every gRPC implementation and test file has 100% branch coverage on both interpreters. Ruff, both Pyright configurations, English docs, and all six HTTP conformance baseline legs pass | Whole adapter coverage remains 90%: AMQP is 31% and MQTT is 37%. Broker record/replay, remaining CI entries, and final API review remain | + +`ServerRuntime` replaces the unshipped `ServerHost` name following naming review. It is not an MCP host or network listener. `DispatcherTransport` replaces the planned client factory: the explicit wrapper avoids duplicating the `Client` constructor's options or guessing what an arbitrary context manager yields. Existing `Transport` objects still yield stream pairs. + +## Objective + +You can implement MQTT, AMQP, or gRPC adapters outside the SDK and use them with `Client`, low-level `Server`, and `MCPServer`. Adapters use only supported public APIs. They reuse MCP negotiation, validation, middleware, callbacks, and result handling instead of implementing those features again. + +There are two integration levels: + +- **Message transport:** carries MCP JSON-RPC messages over another communication channel. MQTT, AMQP, and a gRPC bidirectional stream can use this level. +- **Native binding:** maps MCP operations to another RPC system, such as gRPC methods and protobuf messages. It uses the dispatcher boundary rather than pretending to be a JSON-RPC stream. + +Supporting a native binding in the SDK does not make that binding an official MCP transport. Its wire format and interoperability claims need a separate specification review. + +## Scope and constraints + +- Preserve the existing `Transport` context manager and its stream-pair return value. +- Preserve `Client(...)`, `ClientSession(...)`, `Server.run(...)`, and built-in `MCPServer.run(...)` behavior. +- Add public entry points rather than removing or deprecating existing ones. +- Keep MQTT, AMQP, gRPC, and protobuf dependencies in adapter packages. +- Reuse `Dispatcher`, `JSONRPCDispatcher`, and the server runner functions. +- Do not introduce an interface with one method per MCP operation. +- Do not add a plugin registry, URL-scheme discovery, a universal broker configuration, or automatic request replay. +- Update relevant documentation with each public change. Do not add entries to the closed v1-to-v2 migration guide. + +## Existing foundations + +| Component | Existing extension point | Work needed | +| --- | --- | --- | +| `src/mcp/client/_transport.py` | Async context manager yielding `SessionMessage` streams | Document the full contract and expose supporting types through supported public imports | +| `src/mcp/client/client.py` | Accepts stream transports | Add an explicit lifecycle-managed dispatcher integration | +| `src/mcp/client/session.py` | Accepts `dispatcher=` | Preserve this entry point and make custom implementations supportable | +| `src/mcp/shared/dispatcher.py` | Request/notification boundary independent of wire encoding | Stabilize lifecycle, failures, ordering, and cancellation requirements | +| `src/mcp/server/runner.py` | Connection, stream, and single-request drivers | Expose hosting without requiring adapters to reconstruct the protocol pipeline | +| `src/mcp/server/lowlevel/server.py` | Runs one stream connection with its own lifespan | Support a host owning one lifespan across multiple peers | +| `src/mcp/server/mcpserver/server.py` | Built-in transport hosting | Make the same custom hosting surface available without private access | +| `src/mcp/shared/transport_context.py` | Transport-specific metadata | Connect it to actual user handler contexts | + +## Recommended ownership model + +| Layer | Owns | +| --- | --- | +| Adapter | Wire framing, physical connections, broker subscriptions, delivery settlement, routing, authenticated transport identity | +| Dispatcher | Request correlation, inbound scheduling, response delivery, notification ordering, progress and cancellation translation | +| MCP layer | Version negotiation, capability rules, typed validation, middleware, callbacks, result shaping | +| Server host | Application lifespan and supervision of active connections and requests | +| Application | Adapter configuration, authorization policy, and any operation-specific idempotency guarantees | + +A broker connection is not an MCP client connection. Each logical peer needs isolated request correlation and, for handshake-era protocols, isolated negotiated state. + +## Phase 1: Approve and pin the contracts + +### Deliverables + +- [ ] Inventory public imports, constructor forms, stream ownership, and observable failure behavior. Identify gaps in existing tests before adding more tests. +- [ ] Review real adapter call sites. Start with the [Google gRPC Python adapter](https://github.com/GoogleCloudPlatform/mcp-grpc-transport-py) and [Amazon MQ AMQP adapter](https://github.com/amazon-mq/mcp-amqp-transport). The latter is TypeScript and informs routing requirements, not Python API compatibility. +- [ ] Approve exact names and types for a shared transport import surface, an explicit client dispatcher factory, and a server hosting context. Names remain open until this review. +- [ ] Write complete client and server usage examples for both integration levels as design artifacts. Mark proposed calls as proposed until implemented. +- [ ] Define supported protocol versions and required features for each reference adapter. Start AMQP validation with AMQP 0.9.1 and MQTT validation with MQTT 5; do not imply support for other versions without testing them. +- [ ] Inspect the pinned conformance suite and map relevant existing scenarios to the work. SDK extension-point tests are separate from wire conformance tests. + +### Conformance evidence + +The pinned `@modelcontextprotocol/conformance@0.2.0-alpha.11` lists 69 frozen scenarios for 2026-07-28. Relevant shared-pipeline checks include `tools-call-with-progress`, `caching`, `request-metadata`, and the `input-required-result-*` scenarios. These existing features are reused, not reimplemented by the new extension points. + +Fresh local baseline runs pass for all six legs. Server 2026-07-28 has 151 passing checks and server 2025-11-25 has 84. Client 2026-07-28 has 387 and client 2025-11-25 has 224. The default server leg has 204 passing checks and 25 expected failures; the default client leg has 464 passing checks and nine expected failures. The existing baselines were not changed, and no solo retry was needed. These runs use the HTTP harness and do not certify custom broker or protobuf wire bindings. + +### Contract decisions + +Specify readiness, single-entry/re-entry rules, borrowed versus owned resources, EOF, cancellation, shutdown order, and failure propagation. Preserve current behavior on existing entry points. + +Distinguish malformed message observations, peer MCP errors, request timeouts, and terminal transport failures. A fatal receive or send failure must settle pending calls; yielding an exception item must not be mistaken for closing the channel. + +Define which notification ordering the dispatcher guarantees, including the existing receive-order intercept used by subscriptions. Do not require globally ordered request completion. + +Define how unsupported back-channels are reported. Transport capability never overrides a protocol-version prohibition. + +For any newly implemented 2026-07-28 feature, require a matching conformance-suite test. If none exists, stop that feature and report the missing test so an issue can be raised upstream. Do not silently substitute a local test for the repository's conformance requirement. + +### Exit condition + +A maintainer approves the contract and compatibility matrix before implementation. No proposed native wire binding is presented as standardized without evidence. + +## Phase 2: Expose shared types and transport metadata + +Depends on phase 1. + +### Shared transport deliverables + +- [x] Make `Transport`, stream types, message types, and required adapter metadata available through documented public imports. Keep existing imports working. +- [x] Carry `TransportContext` from adapters through dispatch to the actual `ServerRequestContext` and high-level `MCPServer` context. Add fields or properties without replacing handler argument types. +- [x] Expose the transport context builder on supported stream-hosting paths instead of requiring adapters to construct the internal dispatcher recipe. +- [ ] Preserve existing HTTP request access, headers, SSE callbacks, and unanswered-request settlement behavior. +- [x] Define a supported way to bind verified transport identity to a request. Audit request-state principal binding and context propagation so custom authentication does not accidentally become anonymous. + +### Shared transport acceptance checks + +- A handler can observe typed adapter metadata without a fake Starlette request or a private import. +- Concurrent requests from different principals retain the correct identity and metadata, including during cancellation. +- Broker-supplied reply destinations are authorized against the caller rather than trusted as arbitrary routing instructions. +- Existing HTTP, stdio, request-state, and handler-context tests remain unchanged and pass. + +## Phase 3: Add public multi-client server hosting + +Depends on phase 2. + +### Server runtime deliverables + +- [x] Add a hosting context available from both `Server` and `MCPServer`. It owns one application lifespan and exposes serving operations bound to that lifespan state. +- [x] Support one logical stream connection through the existing dual-era runner. Each connection retains its own protocol and correlation state. +- [ ] Define supervision: one peer disconnecting or sending malformed traffic does not cancel unrelated peers; a fatal listener failure is reported to the host owner. +- [x] Stop admission before shutdown, cancel and join active work, close connection resources, and finally exit application lifespan. Document resource-cleanup deadlines separately from cooperative handler joins. +- [ ] Keep connection/session admission limits and queue bounds configurable at the layer that owns them. Do not create an unbounded task per broker message. +- [ ] Preserve the existing single-connection `Server.run()` behavior. Avoid migrating all built-in hosting paths in the same change. + +### Server runtime acceptance checks + +- Two clients can both issue request ID `1` without cross-delivery or shared negotiation state. +- Application lifespan enters once and exits once while multiple connections come and go. +- One client can disconnect while another completes a call. +- Startup failure, idle connections, in-flight cancellation, and shutdown release resources without hanging. +- A custom adapter can host an `MCPServer` without accessing `_lowlevel_server`. + +## Phase 4: Validate message transports outside the SDK + +Depends on phase 3. Develop MQTT and AMQP adapters independently once the shared contract is settled. + +### Message transport deliverables + +- [x] Build or adapt an external-package-shaped MQTT 5 client and server adapter using only public SDK imports. +- [x] Build or adapt an AMQP 0.9.1 client and server adapter using only public SDK imports. +- [ ] Run the same client/server behavior checks through each adapter. Neither adapter may bypass the runner by calling tool methods directly. +- [ ] Document the wire binding, configuration, supported features, limits, failure behavior, and backend requirements for each adapter. + +### MQTT binding decisions + +Specify request and reply topics, subscription readiness, peer/session identity, and reply-topic authorization. Define QoS and duplicate handling. Do not retain command messages; define rejection of unexpected retained deliveries. Specify message expiry, disconnect detection, and reconnect behavior. + +### AMQP binding decisions + +Specify exchanges, queues, reply addresses, consumer prefetch, and publisher confirmation behavior. Define acknowledgment timing relative to request execution and response publication. Specify redelivery, poison-message handling, and expiry. Keep handshake-era traffic on the appropriate logical peer/worker; do not load-balance it blindly across independent sessions. + +### Delivery guarantees + +Broker delivery guarantees are not exactly-once tool execution. Document the crash window between a side effect, response publication, and message settlement. Do not retry arbitrary operations automatically. If deduplication is offered, define identity scope, retention, and behavior after process restart. + +Cancellation must reach the worker running the request. Reconnect must either restore explicitly supported state or fail the old calls and create a fresh logical connection. It must not silently replay calls. + +### Message transport acceptance checks + +Test a real broker for both adapters: multiple clients, out-of-order responses, duplicate delivery, disconnect/reconnect, backpressure, stale deliveries, and authenticated routing. Record supported external interactions and review recordings for secrets. Do not claim real broker behavior from handwritten mocks. + +Passing this phase establishes the JSON-RPC transport milestone. It does not complete native gRPC support. + +## Phase 5: Support native dispatcher integrations + +Depends on phases 1-3. This work can proceed alongside phase 4, but the final contract must incorporate findings from both paths. + +### Native dispatcher deliverables + +- [ ] Stabilize the custom `Dispatcher` lifecycle after reviewing the existing JSON-RPC and direct implementations and a native gRPC prototype. +- [x] Add an explicit `DispatcherTransport` wrapper accepting an async context manager yielding a dispatcher. Reuse existing client negotiation, caching, extensions, callback, and cleanup paths. Do not infer the integration type from an ambiguous context-manager return value. +- [x] Let the server runtime serve a dispatcher through the existing runner pipeline. Native requests retain inbound envelope/version validation at the untrusted entry boundary. +- [ ] Define native mappings for deadlines, transport cancellation, MCP errors, progress, notifications, request IDs, and subscriptions. Reuse existing call options; do not make HTTP-only options mandatory for native implementations. +- [x] Support required notification-intercept behavior and explicit request IDs used by subscriptions. Live regression tests cover acknowledgment/event routing, collisions, minted IDs, and ID reuse. +- [x] Verify arbitrary MCP method names and extension payloads survive the boundary. The gRPC binding carries arbitrary method names and preserves JSON integer precision; recorded calls also cover error codes outside int32. +- [ ] Preserve stream exception observations on the existing client paths. Define equivalent diagnostics for native transports without requiring `isinstance(JSONRPCDispatcher)` in third-party code. + +### Native dispatcher acceptance checks + +- The same high-level `Client` operations work through stream and native dispatcher integrations. +- Negotiation, multi-round-trip results, middleware, and result validation run through shared MCP code. +- Native cancellation and disconnects settle calls without requiring a fabricated JSON-RPC connection. +- MCP errors retain code, message, and data according to the approved mapping; gRPC status failures remain distinguishable where needed. +- The adapter has no duplicate client-session API and does not subclass `ClientSession` to reimplement every MCP method. + +## Phase 6: Validate native gRPC and publish the contract + +Depends on phases 4 and 5. + +### Native binding validation deliverables + +- [x] Implement a native gRPC client and server reference adapter with a documented protobuf-envelope binding. This is a reference binding, not interoperability with another project's schema. +- [x] Test notification/progress delivery, deadlines, cancellation, metadata, extension payloads, and error fidelity over a real gRPC connection. Replay checks are supplemented by tests executing the current server. +- [x] Document asyncio-only requirements for `grpc.aio`, including the single-loop restriction. Core extension points remain AnyIO-compatible; native adapters do not claim Trio support. +- [x] Validate both low-level `Server` and high-level `MCPServer` hosting, including concurrent peers and independent request delivery. +- [ ] Document stable import paths and complete runnable examples in the relevant existing pages: `docs/client/transports.md`, `docs/advanced/low-level-server.md`, `docs/run/index.md`, `docs/run/asgi.md`, and `docs/handlers/context.md`. Update lifespan, authorization, and client caching pages where their contracts are affected. +- [ ] Obtain fresh API-compatibility and adversarial lifecycle/security reviews. Reviewers should specifically challenge identity isolation, replay, callback deadlocks, and incomplete cleanup. + +### Native binding validation exit condition + +MQTT, AMQP, and native gRPC adapters work on both sides without private imports, duplicated MCP semantics, or new runtime dependencies in the core SDK. Compatibility and validation gates below pass. + +## Validation gates for every implementation slice + +Prefer existing public-API tests and add only missing behavior checks. Core lifecycle tests use in-memory execution, events, and bounded waits. Real transport semantics use real services. Keep test files aligned with the source tree and follow `.claude/skills/test-quality/SKILL.md`. + +Cover the following combinations where applicable: + +| Dimension | Cases | +| --- | --- | +| Server API | `Server`, `MCPServer` | +| Client integration | Existing stream transport, new dispatcher factory | +| Protocol | Legacy handshake, automatic discovery, pinned modern version | +| Messaging | Concurrent requests, peer errors, notifications, progress, subscriptions, extension methods | +| Lifecycle | Startup failure, timeout, peer cancellation, EOF, send failure, shutdown | +| Isolation | Repeated request IDs across peers, independent negotiated state, distinct identities | +| Delivery | Backpressure, duplicate and late messages, worker failure, reconnect | + +Transport bindings must document unsupported combinations rather than silently pass partial behavior as full support. + +Run repository checks with the frozen lockfile: + +```bash +uv run --frozen ruff check . +uv run --frozen ruff format --check . +uv run --frozen pyright +./scripts/test +``` + +Require 100% branch coverage and `strict-no-cover` for SDK changes. Run the existing client/server conformance jobs for both protocol eras and the default suite. These protect current wire behavior; they do not by themselves certify an MQTT, AMQP, or native protobuf binding. Validate cross-version and platform behavior in the repository CI matrix. + +## Review findings addressed + +A focused independent reviewer found two native lifecycle defects: cancellation was signalled only after handler cleanup, and a five-second join could abandon handlers before application lifespan closed. The server now signals cancellation before unwinding a handler task group. Both client and server join active work without abandoning it at a deadline. Live checks hold shielded cleanup beyond five seconds and prove that resources remain alive until it finishes. Shutdown can wait indefinitely for code that ignores cancellation; a process supervisor owns any hard termination deadline. + +Follow-up review confirmed those corrections and found a borrowed-channel closure race. It was reproduced through `Client.list_tools()` immediately after `channel.close()`, then corrected by consulting native channel state before constructing an RPC. The regression passes without making a network request. + +The `reviewer` Agent Hub profile is available. Earlier attempts used nonexistent profile names; subsequent focused reviews completed. Final review of the complete API and adapter work is still required. + +## Latest validation slice + +The four standalone gRPC shutdown programs are maintained AnyIO regressions under `examples/transports/tests/`. The tests execute real loopback servers and hold shielded cleanup beyond the former five-second join deadline. Additional cases exercise malformed frames, saturation, application and validation errors, callback isolation, request-ID collisions, and subscription routing. Cassette tests compare their recorded native results against the current in-process MCP handler, so replay does not leave their server handlers untested. + +`DirectDispatcher` previously returned from `run()` while a caller task was still unwinding a request or notification handler. Operations now register a cancellation scope and completion event on both peers. Closing either peer cancels them; `run()` joins completion before returning. Tests cover both closing peers, requests and notifications, nested back-channel calls, ordinary handler errors, and in-process client lifespan ordering. The dispatcher/client subset also passes all 108 tests on Python 3.10. + +The native adapter now exposes gRPC's verified `peer_identity_key` and immutable `peer_identities`. A fresh adversarial review found no direct spoofing path but requested stronger boundary assertions and authority documentation. Five live cases now cover plaintext, server-only TLS, mutual TLS, absent client certificates, and certificates signed by an untrusted key with the same issuer name. Middleware proves rejected peers never reach MCP dispatch. Accepted requests prove neither MCP client-info claims nor forged invocation metadata can supply or replace certificate identity. Empty identity does not prove plaintext, and names are not globally unique across independent trusted authorities. A fresh read-only closeout review found no remaining correctness, security, or test gaps in this TLS slice; it did not approve the combined API. + +TLS validation exposed a gRPC completion-queue limitation, reproduced without MCP in `examples/transports/reproduce_grpc_loop_shutdown.py`. With `grpcio==1.84.0` on macOS/Python 3.14.6, cancelled connectivity watches can complete after `channel.close()` and target a previously closed event loop. The adapter suite keeps one AnyIO runner for its session, while still closing per-test resources. The README records this support restriction, not an upstream fix; repeated loop lifetimes and final native-queue drainage are not certified. + +### Reproduce the validation + +```bash +./scripts/test +UV_PROJECT_ENVIRONMENT=examples/transports/.venv uv sync --frozen --package mcp-transport-examples --group dev +UV_PROJECT_ENVIRONMENT=examples/transports/.venv uv run --frozen --package mcp-transport-examples --group dev pytest -c examples/transports/pyproject.toml examples/transports/tests --record-mode=none +uv run --frozen pyright --project examples/transports +DOCS_LANGUAGES=en-only bash scripts/docs/build.sh +``` + +Run the broker commands in `examples/transports/README.md` against the pinned Compose fixtures. The shared-check workflow now runs the complete adapter suite and live broker programs on Python 3.10 and 3.14, and retains `transport-results-` JUnit artifacts. Results are attached to [the pull request's checks](https://github.com/modelcontextprotocol/python-sdk/pull/3517/checks), not machine-local log paths. The conformance workflow records all six baseline legs separately. + +Core coverage remains 100%. Whole adapter coverage is still incomplete because MQTT/AMQP failure paths are not cassette-backed; generated protobuf implementation is excluded as compiler output, not handwritten adapter code. Do not interpret a passing adapter pytest job as completion of that separate coverage gate. + +### Review corrections + +Native regressions now cover swallowed direct-handler cancellation, notification callback isolation, post-close notification drops, late request-scoped notifications, sanitized raw-dispatcher errors, strict progress fields, deep JSON and exponent overflow. HTTP framing supplies handler-visible context and headers without adding credentials to message representations; driver stream cleanup is shielded and bounded. The published principal-binding example is exercised directly. + +RabbitMQ no longer grants client writes to the default exchange. Each receiving queue has a dedicated direct exchange, and live checks reject both default-exchange injection and writes to another principal's exchange. MQTT examples configure Last Wills before CONNECT; a live broker-forced disconnect settles a pending MCP call without relying on its request timeout. The existing aiomqtt negative-publish-reason limitation remains a merge blocker: its public API does not expose those acknowledgement codes. + +Optional runtime design feedback remains separate from these corrections: configurable cleanup grace, exception-group behavior, and admission cancellation during runtime shutdown need a contract decision rather than an unreviewed change in semantics. + +## Next implementation work + +1. Settle broker record/replay. `cassetter` has no MQTT/AMQP interceptor. Its gRPC wrapper also omits parts of streaming cancellation and matches only RPC methods; current tests supplement matching with serialized-request assertions and keep each cassette to one RPC. Do not substitute handwritten broker mocks to clear coverage. +2. Exercise broker redelivery, expiry, connection loss, malformed frames, cancellation, saturation, and TLS identity denial. Resolve aiomqtt's queue-overflow drops and discarded negative publish reason codes, or choose a provider with the required failure signals. Bound executing broker work, not only queued messages and connections. +3. Complete remaining CI matrix coverage and external adapter compatibility checks. Existing HTTP conformance baselines and local native coverage do not certify broker delivery semantics or another protobuf binding. +4. Obtain final independent compatibility and lifecycle/security reviews of the combined API and adapters. The scoped DirectDispatcher and TLS reviews are not approval of the entire change. Keep the contract provisional until those gates and maintainer approval are complete. + +The local Mosquitto and RabbitMQ fixtures are pinned by image digest. The temporary `mcp-sdk-transport-check` containers and network were removed after verification. Their public test credentials and ACLs are in `examples/transports/brokers/`; they are not production configuration. + +## Delivery order + +1. Approve the contract, compatibility matrix, and usage examples. +2. Expose shared types and context propagation in small reviewable changes. +3. Add server hosting with lifecycle and isolation tests. +4. Validate MQTT and AMQP adapters while implementing dispatcher integration. +5. Complete native gRPC validation and review the combined public contract. + +Each implementation change includes its tests and affected documentation. Do not defer coverage or lifecycle verification to the last phase. Do not publish the contract as stable until both message-based and native integrations have exercised it. diff --git a/docs/advanced/low-level-server.md b/docs/advanced/low-level-server.md index 5d49846b5f..d8181f1a9a 100644 --- a/docs/advanced/low-level-server.md +++ b/docs/advanced/low-level-server.md @@ -65,6 +65,12 @@ The same text the `@mcp.tool()` version produced. Two honest differences: In a test you skip uvicorn and the port: `Client(server)` takes a low-level `Server` in-process exactly like it takes an `MCPServer`, and **[Testing](../get-started/testing.md)** is that pattern. +## Custom transports + +`Server.serve()` shares one application lifespan across multiple custom transport connections, just like `MCPServer.serve()`. Use the complete adapter example under [Running your server](../run/index.md#custom-transports). + +For a single connection, `Server.run(read_stream, write_stream, initialization_options, *, transport_builder=...)` remains available. The optional builder converts inbound message metadata into the `TransportContext` exposed as `ctx.transport`. Without it, stream dispatch uses the context supplied by the framing transport, falling back to generic JSON-RPC metadata. Built-in HTTP transports supply their kind and the current request's headers. Both paths retain the existing protocol-version handling; custom transport capabilities cannot enable features that the negotiated version forbids. + ## Nothing is checked for you `MCPServer` rejects a bad argument before your function ever runs, validating the call against the schema it generated (**[Tools](../servers/tools.md)**). diff --git a/docs/client/transports.md b/docs/client/transports.md index 6d9d30f90d..97050799a0 100644 --- a/docs/client/transports.md +++ b/docs/client/transports.md @@ -118,6 +118,8 @@ No subprocess, no port, no bytes on a wire. The client and the server are two ob The same form doubles as an embedding API: an application that constructs the server itself can call its tools without a network hop. +Closing the client cancels active in-process requests and waits for their handler cleanup before leaving application lifespan. A caller interrupted by connection closure receives `MCPError` with code `CONNECTION_CLOSED`. Handlers and callbacks must cooperate with cancellation; shielded cleanup keeps the application's resources alive until it finishes. + ## SSE `sse_client(url)`, from `mcp.client.sse`, is the HTTP transport that Streamable HTTP superseded. Wrap it the same way, `Client(sse_client("http://localhost:8000/sse"))`, to talk to a server that still speaks it, and don't build anything new on it. @@ -128,6 +130,40 @@ To `Client`, all of the above are the same thing. A **transport** is any async context manager that yields a `(read, write)` pair of message streams: formally, the `Transport` protocol in `mcp.client`. `Client` resolves its argument by type: a `str` becomes `streamable_http_client(url)`, a `StdioServerParameters` becomes `stdio_client(params)`, a server object connects in-process, and anything else is entered as a transport directly. That last rule is why `stdio_client(...)`, `streamable_http_client(...)` and `sse_client(...)` all drop into the same slot, and why you can write your own. +### Implement a message transport + +```python title="custom_transport.py" +--8<-- "docs_src/client_transports/tutorial005.py" +``` + +This example implements an in-memory adapter with two independent clients. A network adapter uses the same `TransportStreams` contract and replaces the memory channels with message readers and writers. You import the contract and its supporting types from `mcp.shared.transport`; the existing `mcp.client.Transport` import still works. + +Each stream pair represents **one logical peer**, not an entire broker. The adapter owns framing, routing, and its network resources. The SDK owns negotiation, request correlation, and MCP validation. + +Entering a transport opens its channel. Exiting stops its background tasks and closes resources it owns. The SDK also closes streams during connection shutdown, so their `aclose()` methods must be safe to call more than once. A network client supplied by the application remains owned by the application. + +An inbound item is a decoded `SessionMessage` or an exception describing a recoverable message error. An exception item alone does not disconnect the peer. End the read stream on connection loss so pending calls fail instead of waiting indefinitely. Make writes cancellable and apply backpressure rather than buffering without a bound. + +!!! warning "Delivery is not execution" + MQTT or AMQP delivery guarantees do not make a tool execute exactly once. A redelivered request can repeat a side effect. Define expiry, duplicate handling, and reconnect behavior in the adapter; do not silently replay unfinished calls. + +The server side of this example uses `server.serve()`. Its lifecycle and connection limits are covered under [Custom transports](../run/index.md#custom-transports). The repository's `examples/transports/README.md` contains live MQTT 5 and AMQP 0.9.1 examples, their binding rules, and the validation still needed before production use. + +### Integrate a native dispatcher + +```python title="dispatcher_transport.py" +--8<-- "docs_src/client_transports/tutorial006.py" +``` + +`DispatcherTransport` explicitly wraps an async context manager yielding a `Dispatcher`. `Client` enters that context, starts the dispatcher, and uses its ordinary MCP negotiation, callbacks, caching, and validation. It stops the dispatcher before exiting the connection context. You configure the client through the same constructor; there is no separate native client-session API. + +The example uses the SDK's `DirectDispatcher`. The native gRPC reference adapter is developed in a separate follow-up to this SDK API change. Native network bindings implement this dispatcher boundary instead of creating `SessionMessage` streams. The connection context acquires the transport resources; it must yield an unstarted dispatcher because the SDK owns `run()`. + +On the server, `runtime.connect(DispatcherTransport(...))` serves the modern per-request-envelope protocol. It rejects the legacy initialize handshake. Use `mode="auto"` or a supported modern version on the client. Message transports still support both eras. Native dispatchers supply their own contexts, so this server path rejects `session_id=` and `transport_builder=`. + +!!! warning "Native bindings remain experimental" + The custom `Dispatcher` lifecycle is still provisional pending validation against native network adapters. This wrapper is not an official gRPC wire binding. Define and test framing, cancellation, error mapping, notifications, and extension payloads in your adapter before claiming interoperability. + ## Recap * `Client("http://.../mcp")` (a URL) connects over Streamable HTTP, the production transport. diff --git a/docs/handlers/context.md b/docs/handlers/context.md index f43521aa05..a09b78f94a 100644 --- a/docs/handlers/context.md +++ b/docs/handlers/context.md @@ -63,6 +63,7 @@ The injected object is small. Besides `request_id`: * `await ctx.report_progress(progress, total, message)`: stream progress back to the caller during a long call. The whole story is in **[Progress](progress.md)**. * `await ctx.elicit(message, schema)` and `await ctx.elicit_url(...)`: pause the tool and ask the user a question. That's **[Elicitation](elicitation.md)**. * `ctx.session`: the server's side of the conversation with this client. Notifications you send to the client live here; the last section uses it. +* `ctx.transport`: transport metadata supplied by the dispatcher. Custom adapters can attach a `TransportContext` subclass; see [Custom transports](../run/index.md#custom-transports). The SDK populates it for dispatched requests; manually constructed request contexts may leave it `None`. It does not change the existing `ctx.headers` behavior. * `ctx.headers`: the request headers the transport carried, or `None` on stdio. Read a custom header with `(ctx.headers or {}).get("x-...")`. Headers are client-supplied input - fine for a locale or a feature flag, never an identity. * `ctx.request_context`: the raw per-request record. The field you'll reach for is `lifespan_context`, the object your startup code yielded (see **[Lifespan](lifespan.md)**). diff --git a/docs/run/authorization.md b/docs/run/authorization.md index fefd0ed34a..59299825c8 100644 --- a/docs/run/authorization.md +++ b/docs/run/authorization.md @@ -42,7 +42,7 @@ The SDK has no opinion about what a valid token looks like. You tell it, by impl ## What you get over HTTP -Authorization lives in HTTP headers, so it exists only on the HTTP transports. Run it on the one you deploy: `mcp.run(transport="streamable-http")` puts it on `http://127.0.0.1:8000/mcp`, and **[Running your server](index.md)** has the rest. The app now has two routes: +The SDK's built-in OAuth integration uses HTTP headers, so it applies only to HTTP transports. Run it on the one you deploy: `mcp.run(transport="streamable-http")` puts it on `http://127.0.0.1:8000/mcp`, and **[Running your server](index.md)** has the rest. The app now has two routes: ```text /mcp @@ -104,6 +104,21 @@ Call `whoami` with `Authorization: Bearer alice-token` and the model reads: alice (scopes: notes:read) ``` +## Custom transport identities + +```python title="server.py" +--8<-- "docs_src/authorization/tutorial003.py" +``` + +Have your adapter attach `VerifiedPeer` only after authenticating the caller. `runtime.connect(transport_builder=...)` passes that metadata to handlers as `ctx.transport`. Use a stable, namespaced principal that distinguishes the issuing authority and user, not a display name or a client-supplied `_meta` field. + +The existing `RequestStateSecurity.bind_principal` hook binds sealed request state to this identity. Another principal cannot replay it. Raising when verified metadata is absent prevents state from silently becoming anonymous. This hook protects multi-round-trip state; it does not authenticate connections or authorize ordinary tool calls. Those checks still belong at the adapter boundary and in your application policy. + +The generated key suits a single process. Share keys across workers when retries can reach another instance, as described in [Protecting request state](../handlers/multi-round-trip.md#protecting-requeststate). + +!!! warning "Broker credentials are not publisher identity" + A service's broker credentials authenticate the service, not every publisher. Bind peers through broker-enforced topic or queue permissions, or verify an end-user credential yourself. Validate reply destinations before sending data. The SDK's `get_access_token()` remains an HTTP OAuth helper; custom transport metadata does not populate it automatically. + ## The half the SDK doesn't do The SDK gives you the resource-server half: verify, advertise, refuse. It does not give you a login page, a consent screen, or a token. diff --git a/docs/run/index.md b/docs/run/index.md index fb23b4bb54..e162ff4d01 100644 --- a/docs/run/index.md +++ b/docs/run/index.md @@ -88,6 +88,31 @@ Each transport has its own keyword arguments, all on `run()`: `run()` is the short road. The moment you need more (your server mounted inside an existing app, two servers in one process, CORS for browser clients), you build the ASGI app yourself and hand it to any ASGI host. That is **[Add to an existing app](asgi.md)**. +## Custom transports + +```python title="custom_transport.py" +--8<-- "docs_src/client_transports/tutorial005.py" +``` + +`server.serve()` returns a context manager yielding a `ServerRuntime`. It starts application lifespan once and shares that state across the connections you supply. Both `MCPServer` and the low-level `Server` expose this API. It does not open a network listener or connect to a broker. + +Call `await runtime.connect(transport)` for each logical peer. The runtime opens the transport and serves it in the background. For message streams, the call returns when the transport is open, before MCP negotiation. For a dispatcher transport, it also waits for the dispatcher to signal readiness. Each peer has its own request-ID state; message streams negotiate their protocol era independently. + +| Option | Behavior | +| --- | --- | +| `server.serve(max_connections=100)` | Limits active connections. `connect()` waits for capacity before opening another transport. | +| `runtime.connect(..., transport_builder=...)` | Builds each inbound message's `TransportContext`, available as `ctx.transport` in handlers. | +| `runtime.connect(..., session_id=...)` | Supplies an optional identifier for a handshake-era connection. It is not authentication. | + +The default connection limit prevents an adapter from opening unlimited peers. Await admission in your listener instead of spawning unbounded tasks that wait for a slot. Message-size limits, broker queue limits, and per-peer request limits remain the adapter's responsibility. + +An error opening a transport reaches the caller of `connect()`. A later connection failure is logged and closes that peer without cancelling other peers. Exiting `server.serve()` stops admission, cancels active work, closes transports, and then exits application lifespan. Transport cleanup and lifespan cleanup each have a five-second cancellation deadline; cleanup code must cooperate with cancellation. Cleanup timeouts do not suppress an earlier listener or dispatcher startup failure. Dispatchers must join their handlers before returning; these deadlines do not permit closing application resources while a handler still uses them. Code that ignores cancellation can delay that join, so enforce hard process deadlines outside the SDK. Do not retain a runtime after its context exits. + +!!! warning "A peer label is not an identity" + The example attaches a label for demonstration. A real adapter must authenticate and authorize callers before binding identity to a request. Authenticating your server's broker connection does not authenticate every publisher. Validate reply destinations instead of forwarding messages to arbitrary client-supplied topics or queues. + +The [client transport contract](../client/transports.md#implement-a-message-transport) describes message types, resource ownership, and connection loss. For a native RPC binding, `runtime.connect()` also accepts an explicit [dispatcher transport](../client/transports.md#integrate-a-native-dispatcher). That entry serves modern per-request envelopes, not legacy handshakes. The built-in `run()` forms remain unchanged. + ## Server settings A couple of things about running are not about the transport. They are constructor arguments: diff --git a/docs_src/authorization/tutorial003.py b/docs_src/authorization/tutorial003.py new file mode 100644 index 0000000000..72867892d7 --- /dev/null +++ b/docs_src/authorization/tutorial003.py @@ -0,0 +1,27 @@ +import secrets +from dataclasses import dataclass + +from mcp.server import ServerRequestContext +from mcp.server.mcpserver import MCPServer +from mcp.server.request_state import RequestStateSecurity +from mcp.shared.transport import TransportContext + + +@dataclass(kw_only=True, frozen=True) +class VerifiedPeer(TransportContext): + principal: str + + +def principal(ctx: ServerRequestContext) -> str: + if not isinstance(ctx.transport, VerifiedPeer): + raise ValueError("Verified transport identity is required") + return ctx.transport.principal + + +mcp = MCPServer( + "broker-service", + request_state_security=RequestStateSecurity( + keys=[secrets.token_bytes(32)], + bind_principal=principal, + ), +) diff --git a/docs_src/client_transports/tutorial005.py b/docs_src/client_transports/tutorial005.py new file mode 100644 index 0000000000..796744ba4d --- /dev/null +++ b/docs_src/client_transports/tutorial005.py @@ -0,0 +1,57 @@ +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Any + +import anyio + +from mcp import Client +from mcp.server.mcpserver import Context, MCPServer +from mcp.server.runtime import ServerRuntime +from mcp.shared.memory import create_client_server_memory_streams +from mcp.shared.transport import MessageMetadata, TransportContext, TransportStreams + + +@dataclass(kw_only=True, frozen=True) +class PeerContext(TransportContext): + peer: str + + +server = MCPServer("Custom transport") + + +@server.tool() +async def identify(ctx: Context) -> str: + transport = ctx.transport + assert isinstance(transport, PeerContext) + return transport.peer + + +@asynccontextmanager +async def memory_client(runtime: ServerRuntime[Any], peer: str) -> AsyncIterator[TransportStreams]: + async with create_client_server_memory_streams() as (client_streams, server_streams): + + @asynccontextmanager + async def server_transport() -> AsyncIterator[TransportStreams]: + async with server_streams[0], server_streams[1]: + yield server_streams + + def build_context(metadata: MessageMetadata) -> PeerContext: + return PeerContext(kind="memory", can_send_request=True, peer=peer) + + await runtime.connect(server_transport(), transport_builder=build_context) + yield client_streams + + +async def main() -> None: + async with server.serve(max_connections=10) as runtime: + async with Client(memory_client(runtime, "alice")) as alice: + async with Client(memory_client(runtime, "bob")) as bob: + alice_result = await alice.call_tool("identify") + bob_result = await bob.call_tool("identify") + assert alice_result.structured_content == {"result": "alice"} + assert bob_result.structured_content == {"result": "bob"} + + +if __name__ == "__main__": + anyio.run(main) diff --git a/docs_src/client_transports/tutorial006.py b/docs_src/client_transports/tutorial006.py new file mode 100644 index 0000000000..f87cae91a0 --- /dev/null +++ b/docs_src/client_transports/tutorial006.py @@ -0,0 +1,54 @@ +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Any + +import anyio + +from mcp import Client +from mcp.server.mcpserver import Context, MCPServer +from mcp.server.runtime import ServerRuntime +from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair +from mcp.shared.dispatcher import Dispatcher +from mcp.shared.transport import DispatcherTransport, TransportContext + +server = MCPServer("Dispatcher transport") + + +@server.tool() +async def greet(name: str, ctx: Context) -> str: + assert ctx.transport is not None + assert not ctx.transport.can_send_request + return f"Hello, {name}!" + + +def direct_client(runtime: ServerRuntime[Any]) -> DispatcherTransport: + @asynccontextmanager + async def connection() -> AsyncIterator[Dispatcher[TransportContext]]: + client_dispatcher, server_dispatcher = create_direct_dispatcher_pair() + + @asynccontextmanager + async def server_connection() -> AsyncIterator[Dispatcher[TransportContext]]: + try: + yield server_dispatcher + finally: + server_dispatcher.close() + + try: + await runtime.connect(DispatcherTransport(server_connection())) + yield client_dispatcher + finally: + client_dispatcher.close() + server_dispatcher.close() + + return DispatcherTransport(connection()) + + +async def main() -> None: + async with server.serve() as runtime: + async with Client(direct_client(runtime)) as client: + result = await client.call_tool("greet", {"name": "Alice"}) + assert result.structured_content == {"result": "Hello, Alice!"} + + +if __name__ == "__main__": + anyio.run(main) diff --git a/src/mcp/client/_transport.py b/src/mcp/client/_transport.py index 0163fef950..f7ca1d30b8 100644 --- a/src/mcp/client/_transport.py +++ b/src/mcp/client/_transport.py @@ -1,21 +1,5 @@ """Transport protocol for MCP clients.""" -from __future__ import annotations - -from contextlib import AbstractAsyncContextManager -from typing import Protocol - -from mcp.shared._stream_protocols import ReadStream, WriteStream -from mcp.shared.message import SessionMessage +from mcp.shared.transport import ReadStream, Transport, TransportStreams, WriteStream __all__ = ["ReadStream", "WriteStream", "Transport", "TransportStreams"] - -TransportStreams = tuple[ReadStream[SessionMessage | Exception], WriteStream[SessionMessage]] - - -class Transport(AbstractAsyncContextManager[TransportStreams], Protocol): - """Protocol for MCP transports. - - A transport is an async context manager that yields read and write streams - for bidirectional communication with an MCP server. - """ diff --git a/src/mcp/client/client.py b/src/mcp/client/client.py index f921c7e30b..e229bf9c1f 100644 --- a/src/mcp/client/client.py +++ b/src/mcp/client/client.py @@ -71,6 +71,7 @@ from mcp.shared.extension import validate_extension_identifier from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher from mcp.shared.subscriptions import event_to_notification +from mcp.shared.transport import DispatcherTransport logger = logging.getLogger(__name__) @@ -90,10 +91,12 @@ ``__aenter__`` reads them for the handshake step.""" -def _connect_transport(transport: Transport) -> _Connector: - """Connector for the stream-backed paths (URL, user-supplied ``Transport``).""" +def _connect_transport(transport: Transport | DispatcherTransport) -> _Connector: + """Enter a message transport or an explicitly dispatcher-backed connection.""" async def connect(exit_stack: AsyncExitStack, _mode: ConnectMode, _raise_exceptions: bool) -> Dispatcher[Any]: + if isinstance(transport, DispatcherTransport): + return await exit_stack.enter_async_context(transport.connection) read_stream, write_stream = await exit_stack.enter_async_context(transport) return JSONRPCDispatcher(read_stream, write_stream) @@ -280,12 +283,13 @@ async def main(): ``` """ - server: Server[Any] | MCPServer | Transport | StdioServerParameters | str + server: Server[Any] | MCPServer | Transport | DispatcherTransport | StdioServerParameters | str """The MCP server to connect to. If the server is a URL string, it will be used as the URL for a `streamable_http_client` transport. If the server is a `StdioServerParameters`, the command is launched with `stdio_client`. If the server is a `Transport` instance, it will be used directly. + A `DispatcherTransport` explicitly supplies a dispatcher instead of streams. If the server is a `Server` or `MCPServer` instance, it will be connected in-process. """ diff --git a/src/mcp/server/context.py b/src/mcp/server/context.py index bfcb9c9ca4..a242228498 100644 --- a/src/mcp/server/context.py +++ b/src/mcp/server/context.py @@ -47,6 +47,8 @@ class ServerRequestContext(Generic[LifespanContextT, RequestT]): request: RequestT | None = None close_sse_stream: CloseSSEStreamCallback | None = None close_standalone_sse_stream: CloseSSEStreamCallback | None = None + transport: TransportContext | None = None + """Transport metadata supplied by the dispatcher; absent on manually constructed contexts unless provided.""" # Covariant: `lifespan` is exposed read-only, so a `Context[AppState]` passes as `Context[object]`. diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index 8a886dcc24..38f78f0038 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -64,6 +64,7 @@ async def main(): from mcp.server.context import HandlerResult, ServerMiddleware, ServerRequestContext from mcp.server.models import InitializationOptions from mcp.server.runner import serve_dual_era_loop +from mcp.server.runtime import ServerRuntime from mcp.server.streamable_http import EventStore from mcp.server.streamable_http_manager import ( DEFAULT_MAX_SESSIONS, @@ -75,6 +76,7 @@ async def main(): from mcp.shared._stream_protocols import ReadStream, WriteStream from mcp.shared.exceptions import MCPDeprecationWarning from mcp.shared.message import SessionMessage +from mcp.shared.transport import TransportContextBuilder logger = logging.getLogger(__name__) @@ -689,6 +691,15 @@ def session_manager(self) -> StreamableHTTPSessionManager: ) return self._session_manager + def serve(self, *, max_connections: int = 100) -> AbstractAsyncContextManager[ServerRuntime[LifespanResultT]]: + """Share one application lifespan across custom transport connections. + + Use `await runtime.connect(transport)` inside the context for each logical + peer. Admission waits at `max_connections`; exiting cancels active + connections and closes their transports before application cleanup. + """ + return ServerRuntime[LifespanResultT].open(self, max_connections=max_connections) + async def run( self, read_stream: ReadStream[SessionMessage | Exception], @@ -699,6 +710,8 @@ async def run( # but also make tracing exceptions much easier during testing and when using # in-process servers. raise_exceptions: bool = False, + *, + transport_builder: TransportContextBuilder | None = None, ) -> None: """Serve a single connection over the given streams until the read side closes. @@ -706,7 +719,9 @@ async def run( then drives the loop, serving the legacy handshake era and the modern per-request-envelope era (the client's first request decides which). Transports with their own lifespan owner (the streamable-HTTP manager) - call `serve_loop` directly instead. + call `serve_loop` directly instead. `transport_builder` converts each + inbound message's metadata to the `transport` exposed on its handler + context. Without it, the dispatcher supplies generic JSON-RPC metadata. """ async with self.lifespan(self) as lifespan_context: await serve_dual_era_loop( @@ -716,6 +731,7 @@ async def run( lifespan_state=lifespan_context, init_options=initialization_options, raise_exceptions=raise_exceptions, + transport_builder=transport_builder, ) def streamable_http_app( diff --git a/src/mcp/server/mcpserver/context.py b/src/mcp/server/mcpserver/context.py index 07c4799dc1..6265aece59 100644 --- a/src/mcp/server/mcpserver/context.py +++ b/src/mcp/server/mcpserver/context.py @@ -24,6 +24,7 @@ ResourceUpdated, ToolsListChanged, ) +from mcp.shared.transport_context import TransportContext if TYPE_CHECKING: from mcp.server.mcpserver.server import MCPServer @@ -278,6 +279,11 @@ async def log( related_request_id=self.request_id, ) + @property + def transport(self) -> TransportContext | None: + """Transport metadata for this request, when its context supplies it.""" + return self.request_context.transport + @property def headers(self) -> Mapping[str, str] | None: """Request headers carried by this message, when the transport has them. diff --git a/src/mcp/server/mcpserver/server.py b/src/mcp/server/mcpserver/server.py index fbd2c26dd8..486336d1d5 100644 --- a/src/mcp/server/mcpserver/server.py +++ b/src/mcp/server/mcpserver/server.py @@ -90,6 +90,7 @@ from mcp.server.mcpserver.utilities.context_injection import find_context_parameter from mcp.server.mcpserver.utilities.logging import configure_logging, get_logger from mcp.server.request_state import RequestStateBoundary, RequestStateSecurity +from mcp.server.runtime import ServerRuntime from mcp.server.sse import SseServerTransport from mcp.server.stdio import stdio_server from mcp.server.streamable_http import EventStore @@ -1062,6 +1063,15 @@ def decorator( return decorator + def serve(self, *, max_connections: int = 100) -> AbstractAsyncContextManager[ServerRuntime[LifespanResultT]]: + """Share one application lifespan across custom transport connections. + + Use `await runtime.connect(transport)` inside the context for each logical + peer. Admission waits at `max_connections`; exiting cancels active + connections and closes their transports before application cleanup. + """ + return self._lowlevel_server.serve(max_connections=max_connections) + async def run_stdio_async(self) -> None: """Run the server using stdio transport.""" async with stdio_server() as (read_stream, write_stream): diff --git a/src/mcp/server/runner.py b/src/mcp/server/runner.py index 26e8efbe57..6fb36b4521 100644 --- a/src/mcp/server/runner.py +++ b/src/mcp/server/runner.py @@ -67,6 +67,7 @@ from mcp.shared.inbound import InboundLadderRejection, classify_inbound_request from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher, handler_exception_to_error_data from mcp.shared.message import MessageMetadata, ServerMessageMetadata, SessionMessage +from mcp.shared.transport import TransportContextBuilder from mcp.shared.transport_context import TransportContext if TYPE_CHECKING: @@ -81,6 +82,7 @@ "serve_connection", "serve_dual_era_loop", "serve_loop", + "serve_modern_dispatcher", "serve_one", ] @@ -334,6 +336,7 @@ def _make_context( meta=meta, protocol_version=protocol_version, request=request, + transport=dctx.transport, close_sse_stream=close_sse_stream, close_standalone_sse_stream=close_standalone_sse_stream, ) @@ -476,6 +479,7 @@ async def serve_loop( session_id: str | None = None, init_options: InitializationOptions | None = None, raise_exceptions: bool = False, + transport_builder: TransportContextBuilder | None = None, ) -> None: """Drive ``server`` in handshake-only loop mode over a stream pair until the channel closes. @@ -490,6 +494,7 @@ async def serve_loop( read_stream, write_stream, raise_handler_exceptions=raise_exceptions, + transport_builder=transport_builder, # Handle `initialize` inline so a client that pipelines it with the # next request (spec: SHOULD NOT, not MUST NOT) sees the initialized # state instead of failing the init-gate. @@ -608,6 +613,7 @@ async def serve_dual_era_loop( session_id: str | None = None, init_options: InitializationOptions | None = None, raise_exceptions: bool = False, + transport_builder: TransportContextBuilder | None = None, ) -> None: """Drive `server` over a duplex stream pair, in the era the client opens with. @@ -630,7 +636,12 @@ async def serve_dual_era_loop( ) if opens_modern: await _serve_modern_stream( - server, replayed, write_stream, lifespan_state=lifespan_state, raise_exceptions=raise_exceptions + server, + replayed, + write_stream, + lifespan_state=lifespan_state, + raise_exceptions=raise_exceptions, + transport_builder=transport_builder, ) else: await _serve_legacy_stream( @@ -641,9 +652,11 @@ async def serve_dual_era_loop( session_id=session_id, init_options=init_options, raise_exceptions=raise_exceptions, + transport_builder=transport_builder, ) finally: - await write_stream.aclose() + with anyio.move_on_after(_EXIT_STACK_CLOSE_TIMEOUT, shield=True): + await write_stream.aclose() _PRE_REQUEST_REPLAY_LIMIT: int = 8 @@ -709,7 +722,8 @@ async def replay_then_relay() -> None: yield opening_request, replayed tg.cancel_scope.cancel() finally: - await read_stream.aclose() + with anyio.move_on_after(_EXIT_STACK_CLOSE_TIMEOUT, shield=True): + await read_stream.aclose() replay_send.close() replay_receive.close() @@ -723,12 +737,14 @@ async def _serve_legacy_stream( session_id: str | None, init_options: InitializationOptions | None, raise_exceptions: bool, + transport_builder: TransportContextBuilder | None, ) -> None: """Serve a 2025 handshake connection; enveloped requests are refused.""" dispatcher: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher( read_stream, write_stream, raise_handler_exceptions=raise_exceptions, + transport_builder=transport_builder, # `initialize` inline for the same pipelining reason as `serve_loop`. inline_methods=frozenset({"initialize"}), ) @@ -759,11 +775,30 @@ async def _serve_modern_stream( *, lifespan_state: LifespanT, raise_exceptions: bool, + transport_builder: TransportContextBuilder | None, ) -> None: """Serve a 2026-07-28 connection: every request carries its own envelope.""" dispatcher: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher( - read_stream, write_stream, raise_handler_exceptions=raise_exceptions + read_stream, write_stream, raise_handler_exceptions=raise_exceptions, transport_builder=transport_builder ) + await serve_modern_dispatcher(server, dispatcher, lifespan_state=lifespan_state, raise_exceptions=raise_exceptions) + + +async def serve_modern_dispatcher( + server: Server[LifespanT], + dispatcher: Dispatcher[TransportContext], + *, + lifespan_state: LifespanT, + raise_exceptions: bool = False, + task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED, +) -> None: + """Serve per-request-envelope MCP over a wire-independent dispatcher. + + Each request is classified before entering the shared handler pipeline. + Handshake-era initialization is rejected. The dispatcher owns request + scheduling and cancellation; the caller owns application lifespan and + transport resources. Prefer `ServerRuntime.connect()` for managed serving. + """ outbound = NotifyOnlyOutbound(dispatcher) async def on_request( @@ -809,7 +844,7 @@ async def on_notify(dctx: DispatchContext[TransportContext], method: str, params finally: await aclose_shielded(connection) - await dispatcher.run(on_request, on_notify) + await dispatcher.run(on_request, on_notify, task_status=task_status) async def serve_one( diff --git a/src/mcp/server/runtime.py b/src/mcp/server/runtime.py new file mode 100644 index 0000000000..9e67577eca --- /dev/null +++ b/src/mcp/server/runtime.py @@ -0,0 +1,159 @@ +"""Application lifespan and connection supervision for custom transports.""" + +from __future__ import annotations + +import logging +from collections.abc import AsyncIterator +from contextlib import AsyncExitStack, asynccontextmanager +from dataclasses import dataclass, field +from functools import partial +from typing import TYPE_CHECKING, Generic + +import anyio +import anyio.abc +from typing_extensions import TypeVar + +from mcp.server.runner import serve_dual_era_loop, serve_modern_dispatcher +from mcp.shared._compat import resync_tracer +from mcp.shared.transport import DispatcherTransport, Transport, TransportContextBuilder + +if TYPE_CHECKING: + from mcp.server.lowlevel.server import Server + +__all__ = ["ServerRuntime"] + +logger = logging.getLogger(__name__) +LifespanT = TypeVar("LifespanT") + + +@dataclass +class ServerRuntime(Generic[LifespanT]): + """An active server returned by `Server.serve()` or `MCPServer.serve()`. + + Each `connect()` call serves a peer with independent protocol and request-ID state. + """ + + _server: Server[LifespanT] + _lifespan_state: LifespanT + _task_group: anyio.abc.TaskGroup + _limiter: anyio.CapacityLimiter + _active: bool = field(default=True, init=False) + + @classmethod + @asynccontextmanager + async def open( + cls, server: Server[LifespanT], *, max_connections: int = 100 + ) -> AsyncIterator[ServerRuntime[LifespanT]]: + """Share one lifespan, closing connections before application cleanup. + + Cleanup has a five-second cancellation deadline per layer and must cooperate. + """ + if max_connections < 1: + raise ValueError("max_connections must be positive") + body_error: BaseException | None = None + with anyio.CancelScope() as lifespan_scope: + async with server.lifespan(server) as state: + try: + async with anyio.create_task_group() as tg: + runtime = cls(server, state, tg, anyio.CapacityLimiter(max_connections)) + try: + yield runtime + finally: + runtime._active = False + tg.cancel_scope.cancel() + except BaseException as exc: + body_error = exc + raise + finally: + lifespan_scope.shield = True + lifespan_scope.deadline = anyio.current_time() + 5 + if lifespan_scope.cancelled_caught: + logger.warning("Server lifespan cleanup exceeded five seconds") + if body_error is not None: + raise body_error + await resync_tracer() + + async def connect( + self, + transport: Transport | DispatcherTransport, + *, + session_id: str | None = None, + transport_builder: TransportContextBuilder | None = None, + ) -> None: + """Open and supervise one peer's transport until it disconnects. + + Waits for capacity, then owns the entered transport. Opening failures reach + the caller; later failures are logged and isolated. After return, caller + cancellation does not close the connection. + Dispatcher transports signal readiness through `Dispatcher.run()`; + message transports return before receiving the first MCP request. + + Args: + transport: An unopened message transport or `DispatcherTransport`. + session_id: Optional identity for a handshake-era connection. + transport_builder: Builds handler metadata for each inbound message. + + Raises: + RuntimeError: If this runtime has closed. + """ + if not self._active: + raise RuntimeError("Server runtime is closed") + if isinstance(transport, DispatcherTransport) and (session_id is not None or transport_builder is not None): + raise ValueError("Dispatcher transports supply their own context and do not use handshake-era sessions") + + async def serve(*, task_status: anyio.abc.TaskStatus[None]) -> None: + ready = False + run_error: BaseException | None = None + + class ReadyStatus: + def started(self, value: None = None) -> None: + nonlocal ready + task_status.started() + ready = True + + status = ReadyStatus() + try: + async with self._limiter: + with anyio.CancelScope() as cleanup_scope: + async with AsyncExitStack() as stack: + if isinstance(transport, DispatcherTransport): + dispatcher = await stack.enter_async_context(transport.connection) + run = partial( + serve_modern_dispatcher, + self._server, + dispatcher, + lifespan_state=self._lifespan_state, + task_status=status, + ) + else: + read, write = await stack.enter_async_context(transport) + run = partial( + serve_dual_era_loop, + self._server, + read, + write, + lifespan_state=self._lifespan_state, + session_id=session_id, + transport_builder=transport_builder, + ) + try: + if not isinstance(transport, DispatcherTransport): + status.started() + await run() + except BaseException as exc: + run_error = exc + raise + finally: + cleanup_scope.shield = True + cleanup_scope.deadline = anyio.current_time() + 5 + if cleanup_scope.cancelled_caught: + logger.warning("Transport cleanup exceeded five seconds") + if run_error is not None: + raise run_error + return + except Exception: + if not ready: + raise + logger.exception("Transport connection failed") + + await self._task_group.start(serve) diff --git a/src/mcp/server/sse.py b/src/mcp/server/sse.py index d71ef25004..7a35041d5e 100644 --- a/src/mcp/server/sse.py +++ b/src/mcp/server/sse.py @@ -59,6 +59,7 @@ async def handle_sse(request): ) from mcp.shared._context_streams import ContextSendStream, create_context_streams from mcp.shared.message import ServerMessageMetadata, SessionMessage +from mcp.shared.transport_context import TransportContext logger = logging.getLogger(__name__) @@ -279,7 +280,10 @@ async def _handle_post_message(self, scope: Scope, receive: Receive, send: Send) return # Pass the ASGI scope for framework-agnostic access to request data - metadata = ServerMessageMetadata(request_context=request) + metadata = ServerMessageMetadata( + request_context=request, + transport_context=TransportContext(kind="sse", can_send_request=True, headers=request.headers), + ) session_message = SessionMessage(message, metadata=metadata) logger.debug(f"Sending session message to writer: {session_message}") response = Response("Accepted", status_code=202) diff --git a/src/mcp/server/streamable_http.py b/src/mcp/server/streamable_http.py index 416dd9e2b4..b68fb97f5f 100644 --- a/src/mcp/server/streamable_http.py +++ b/src/mcp/server/streamable_http.py @@ -46,6 +46,7 @@ from mcp.shared._stream_protocols import ReadStream, WriteStream from mcp.shared.inbound import MCP_PROTOCOL_VERSION_HEADER from mcp.shared.message import CloseSSEStreamCallback, ServerMessageMetadata, SessionMessage +from mcp.shared.transport_context import TransportContext logger = logging.getLogger(__name__) @@ -252,6 +253,9 @@ def _message_metadata( close_standalone_sse_stream=close_standalone_sse_stream, on_request_unanswered=on_request_unanswered, can_send_request=not self.is_json_response_enabled, + transport_context=TransportContext( + kind="streamable-http", can_send_request=not self.is_json_response_enabled, headers=request.headers + ), ) def close_sse_stream(self, request_id: RequestId) -> None: diff --git a/src/mcp/server/streamable_http_manager.py b/src/mcp/server/streamable_http_manager.py index a7efac3fbc..0a9ae83886 100644 --- a/src/mcp/server/streamable_http_manager.py +++ b/src/mcp/server/streamable_http_manager.py @@ -6,6 +6,7 @@ import logging import math from collections.abc import AsyncIterator +from dataclasses import replace from typing import TYPE_CHECKING, Any, Final from uuid import uuid4 @@ -28,6 +29,7 @@ from mcp.shared._compat import resync_tracer from mcp.shared.inbound import MCP_PROTOCOL_VERSION_HEADER from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher +from mcp.shared.message import MessageMetadata, ServerMessageMetadata from mcp.shared.transport_context import TransportContext if TYPE_CHECKING: @@ -216,6 +218,11 @@ async def _handle_stateless_request( security_settings=self.security_settings, ) + def transport_context(metadata: MessageMetadata) -> TransportContext: + assert isinstance(metadata, ServerMessageMetadata) + assert metadata.transport_context is not None + return replace(metadata.transport_context, can_send_request=False) + # Start server in a new task async def run_stateless_server(*, task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED): async with http_transport.connect() as streams: @@ -230,7 +237,7 @@ async def run_stateless_server(*, task_status: TaskStatus[None] = anyio.TASK_STA # reply has nowhere to land — `can_send_request=False` # makes the per-request channel raise `NoBackChannelError` # for requests while still allowing notifications. - transport_builder=lambda _md: TransportContext(kind="streamable-http", can_send_request=False), + transport_builder=transport_context, ) # Born-ready, no standalone channel: the legacy stateless path # never opens a GET stream and need not see `initialize`. The diff --git a/src/mcp/shared/direct_dispatcher.py b/src/mcp/shared/direct_dispatcher.py index e17283afa2..b9e874eeac 100644 --- a/src/mcp/shared/direct_dispatcher.py +++ b/src/mcp/shared/direct_dispatcher.py @@ -18,7 +18,8 @@ from __future__ import annotations import logging -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from contextlib import asynccontextmanager from dataclasses import dataclass, field from typing import Any @@ -52,6 +53,10 @@ _Notify = Callable[[str, Mapping[str, Any] | None], Awaitable[None]] +class _DispatchClosed(Exception): + """A connection closed while an operation was running in its caller's task.""" + + @dataclass class _DirectDispatchContext: """`DispatchContext` for an inbound request on a `DirectDispatcher`. @@ -104,8 +109,11 @@ class DirectDispatcher: to have started, and once a side has closed - via `close()` or `run()` ending - `send_raw_request` raises `MCPError` (`CONNECTION_CLOSED`) and inbound requests fail the peer's call the same way instead of invoking the - handler. Notifications are fire-and-forget in both directions: after close - they are silently dropped. + handler. Closing either peer cancels active operations, including nested + back-channel calls. `run()` joins handler cleanup before returning so + application resources cannot close underneath it. Interrupted requests + fail with `CONNECTION_CLOSED`; interrupted notifications are dropped. + Notifications sent after close are also silently dropped. """ def __init__(self, transport_ctx: TransportContext, *, raise_handler_exceptions: bool = True): @@ -117,6 +125,7 @@ def __init__(self, transport_ctx: TransportContext, *, raise_handler_exceptions: self._on_notify_intercept: OnNotifyIntercept | None = None self._next_id = 0 self._in_flight_ids: set[RequestId] = set() + self._operations: dict[anyio.CancelScope, anyio.Event] = {} self._ready = anyio.Event() self._close_event = anyio.Event() self._running = False @@ -146,7 +155,11 @@ async def send_raw_request( raise MCPError(code=CONNECTION_CLOSED, message="Connection closed") if not self._running: raise RuntimeError("DirectDispatcher.send_raw_request called before run()") - return await self._peer._dispatch_request(method, params, opts) + try: + async with self._operation(self._peer): + return await self._peer._dispatch_request(method, params, opts) + except _DispatchClosed: + raise MCPError(code=CONNECTION_CLOSED, message="Connection closed") from None async def notify(self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None) -> None: """Send a notification by invoking the peer's `on_notify` directly. @@ -161,7 +174,11 @@ async def notify(self, method: str, params: Mapping[str, Any] | None, opts: Call if self._closed: logger.debug("dropped notification %r on closed DirectDispatcher", method) return - await self._peer._dispatch_notify(method, params) + try: + async with self._operation(self._peer): + await self._peer._dispatch_notify(method, params) + except _DispatchClosed: + logger.debug("dropped notification %r on closed DirectDispatcher", method) async def run( self, @@ -186,25 +203,40 @@ async def run( await self._close_event.wait() finally: self._running = False - self._closed = True - # run() may end via cancellation without close() ever being - # called; setting the event wakes `_wait_ready` waiters so they - # observe the closed state instead of parking forever. - self._close_event.set() + self.close() + with anyio.CancelScope(shield=True): + for finished in tuple(self._operations.values()): + await finished.wait() def close(self) -> None: + """Stop admitting work and cancel active calls; `run()` joins their cleanup.""" self._closed = True self._close_event.set() + for scope in tuple(self._operations): + scope.cancel() + + @asynccontextmanager + async def _operation(self, peer: DirectDispatcher) -> AsyncIterator[None]: + finished = anyio.Event() + with anyio.CancelScope() as scope: + self._operations[scope] = peer._operations[scope] = finished + try: + yield + finally: + self._operations.pop(scope) + peer._operations.pop(scope, None) + finished.set() + if scope.cancel_called: + raise _DispatchClosed def _make_context( self, on_progress: ProgressFnT | None = None, request_id: RequestId | None = None ) -> _DirectDispatchContext: assert self._peer is not None - peer = self._peer return _DirectDispatchContext( transport=self._transport_ctx, - _back_request=lambda m, p, o: peer._dispatch_request(m, p, o), - _back_notify=lambda m, p: peer._dispatch_notify(m, p), + _back_request=self.send_raw_request, + _back_notify=self.notify, request_id=request_id, _on_progress=on_progress, ) diff --git a/src/mcp/shared/dispatcher.py b/src/mcp/shared/dispatcher.py index f2ff96e7d5..15f0d4b3eb 100644 --- a/src/mcp/shared/dispatcher.py +++ b/src/mcp/shared/dispatcher.py @@ -211,9 +211,10 @@ def cancel_requested(self) -> anyio.Event: ... async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None: - """Report progress for the inbound request, if the peer supplied a progress token. + """Report progress for the inbound request when the peer opted in. - A no-op when no token was supplied. + JSON-RPC uses a progress token; direct and native bindings can carry + the callback opt-in separately. Without an opt-in this is a no-op. """ ... @@ -264,8 +265,11 @@ async def run( ) -> None: """Drive the receive loop until the underlying channel closes. - Each inbound request is dispatched to `on_request` in its own task; - the returned dict (or raised `MCPError`) is sent back as the response. + Dispatch each inbound request independently to `on_request`; the + returned dict (or raised `MCPError`) is sent back as the response. + On closure, cancel active operations and join their handler/callback + cleanup before returning. Application resources may close as soon as + this method exits; a shielded handler must not outlive that boundary. Implementations MUST offer every inbound notification to `on_notify_intercept` synchronously in receive order (via `run_notify_intercept`), handing only unconsumed ones to `on_notify`. diff --git a/src/mcp/shared/jsonrpc_dispatcher.py b/src/mcp/shared/jsonrpc_dispatcher.py index 87bdf31ceb..f4b41b4ea9 100644 --- a/src/mcp/shared/jsonrpc_dispatcher.py +++ b/src/mcp/shared/jsonrpc_dispatcher.py @@ -10,7 +10,7 @@ import contextvars import logging from collections.abc import Awaitable, Callable, Mapping -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from functools import partial from typing import Any, Generic, Literal, cast @@ -194,6 +194,8 @@ def _default_transport_builder(metadata: MessageMetadata) -> TransportContext: its streams. """ can_send_request = metadata.can_send_request if isinstance(metadata, ServerMessageMetadata) else True + if isinstance(metadata, ServerMessageMetadata) and metadata.transport_context is not None: + return replace(metadata.transport_context, can_send_request=can_send_request) return TransportContext(kind="jsonrpc", can_send_request=can_send_request) diff --git a/src/mcp/shared/message.py b/src/mcp/shared/message.py index 31e51e7128..ed34072626 100644 --- a/src/mcp/shared/message.py +++ b/src/mcp/shared/message.py @@ -5,11 +5,13 @@ """ from collections.abc import Awaitable, Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any from mcp_types import JSONRPCMessage, RequestId +from mcp.shared.transport_context import TransportContext + ResumptionToken = str ResumptionTokenUpdateCallback = Callable[[ResumptionToken], Awaitable[None]] @@ -50,6 +52,8 @@ class ServerMessageMetadata: # `TransportContext.can_send_request`); a transport that says nothing leaves # it True. can_send_request: bool = True + transport_context: TransportContext | None = field(default=None, kw_only=True, repr=False) + """Context supplied by the framing transport; omitted from repr to avoid logging request headers.""" MessageMetadata = ClientMessageMetadata | ServerMessageMetadata | None diff --git a/src/mcp/shared/transport.py b/src/mcp/shared/transport.py new file mode 100644 index 0000000000..cac6b12994 --- /dev/null +++ b/src/mcp/shared/transport.py @@ -0,0 +1,67 @@ +"""Public contracts for message transports on either side of an MCP connection.""" + +from __future__ import annotations + +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager +from dataclasses import dataclass +from typing import TypeAlias + +from typing_extensions import Protocol + +from mcp.shared._stream_protocols import ReadStream, WriteStream +from mcp.shared.dispatcher import Dispatcher +from mcp.shared.message import ClientMessageMetadata, MessageMetadata, ServerMessageMetadata, SessionMessage +from mcp.shared.transport_context import TransportContext + +__all__ = [ + "ClientMessageMetadata", + "DispatcherTransport", + "MessageMetadata", + "ReadStream", + "ServerMessageMetadata", + "SessionMessage", + "Transport", + "TransportContext", + "TransportContextBuilder", + "TransportStreams", + "WriteStream", +] + +TransportStreams: TypeAlias = tuple[ReadStream[SessionMessage | Exception], WriteStream[SessionMessage]] +TransportContextBuilder: TypeAlias = Callable[[MessageMetadata], TransportContext] + + +class Transport(AbstractAsyncContextManager[TransportStreams], Protocol): + """An async context manager yielding a logical peer's read and write streams. + + Entering opens the channel. Exiting closes owned resources and stops its + background tasks. Consumers may close the streams before context exit, so + stream closure must be idempotent. Borrowed network clients remain owned by + their caller. + + Each inbound item is a decoded `SessionMessage` or a recoverable exception. + End the read stream when the connection is lost; an exception item alone + does not fail pending requests. Writes must support cancellation and apply + backpressure instead of buffering indefinitely. + + A stream pair belongs to one logical peer, not an entire broker. The + adapter owns framing and routing; the SDK owns MCP protocol processing. + """ + + +@dataclass(frozen=True) +class DispatcherTransport: + """Explicitly opt into a dispatcher-backed connection instead of message streams. + + Pass this wrapper to `Client` or `ServerRuntime.connect()`. Entering + `connection` acquires the channel and yields an unstarted dispatcher; the + SDK owns its receive loop. Exiting releases the channel after the loop + stops. Native adapters can use their own framing without implementing MCP + negotiation, validation, callbacks, or a separate client-session API. + + Custom dispatcher implementations remain experimental until the lifecycle + contract has been validated against native network bindings. + """ + + connection: AbstractAsyncContextManager[Dispatcher[TransportContext]] diff --git a/tests/client/test_client.py b/tests/client/test_client.py index d7278e3a81..4982f8ed66 100644 --- a/tests/client/test_client.py +++ b/tests/client/test_client.py @@ -6,9 +6,11 @@ import sys from collections.abc import AsyncIterator, Iterator from contextlib import asynccontextmanager, contextmanager +from typing import Any from unittest.mock import patch import anyio +import anyio.abc import mcp_types as types import pytest from inline_snapshot import snapshot @@ -40,17 +42,40 @@ from mcp.client._memory import InMemoryTransport from mcp.client._transport import TransportStreams from mcp.client.client import Client -from mcp.client.session import ClientRequestContext +from mcp.client.session import ClientRequestContext, ClientSession from mcp.client.streamable_http import streamable_http_client from mcp.server import Server, ServerRequestContext from mcp.server.mcpserver import Context, MCPServer +from mcp.server.runtime import ServerRuntime +from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair +from mcp.shared.dispatcher import Dispatcher from mcp.shared.memory import MessageStream, create_client_server_memory_streams from mcp.shared.message import SessionMessage +from mcp.shared.transport import DispatcherTransport, TransportContext from tests.interaction._connect import BASE_URL, mounted_app pytestmark = pytest.mark.anyio +@asynccontextmanager +async def dispatcher_connection(runtime: ServerRuntime[Any]) -> AsyncIterator[Dispatcher[TransportContext]]: + client_dispatcher, server_dispatcher = create_direct_dispatcher_pair() + + @asynccontextmanager + async def server_connection() -> AsyncIterator[Dispatcher[TransportContext]]: + try: + yield server_dispatcher + finally: + server_dispatcher.close() + + try: + await runtime.connect(DispatcherTransport(server_connection())) + yield client_dispatcher + finally: + client_dispatcher.close() + server_dispatcher.close() + + @pytest.fixture def simple_server() -> Server: """Create a simple MCP server for testing.""" @@ -1002,3 +1027,179 @@ async def elicitation_callback( contents=[TextResourceContents(uri="memory://gated", text="unlocked")], ) ) + + +@pytest.mark.parametrize("mode", ["auto", "2026-07-28"]) +async def test_dispatcher_transport_preserves_custom_methods_payloads_and_progress(mode: str) -> None: + """The native entry keeps arbitrary method payloads and routes progress through the usual client API.""" + + class EchoParams(types.RequestParams): + value: dict[str, Any] + + class EchoResult(types.Result): + value: dict[str, Any] + + async def echo(ctx: ServerRequestContext, params: EchoParams) -> EchoResult: + assert ctx.method == "example/echo" + assert ctx.transport is not None + assert not ctx.transport.can_send_request + await ctx.session.report_progress(1, 2, "halfway") + return EchoResult(value=params.value) + + server = Server("native") + server.add_request_handler("example/echo", EchoParams, echo) + payload = {"large": 2**63 + 1, "vendor/field": [None, {"label": "café"}]} + progress_updates: list[tuple[float, float | None, str | None]] = [] + + async def progress(progress: float, total: float | None, message: str | None) -> None: + progress_updates.append((progress, total, message)) + + with anyio.fail_after(5): + async with server.serve() as runtime: + async with Client(DispatcherTransport(dispatcher_connection(runtime)), mode=mode) as client: + result = await client.session.send_request( + types.Request(method="example/echo", params=EchoParams(value=payload)), + EchoResult, + progress_callback=progress, + ) + assert result.value == payload + assert progress_updates == snapshot([(1, 2, "halfway")]) + + +async def test_dispatcher_transport_runs_multi_round_trip_callbacks() -> None: + """A native connection uses the existing client callback/retry driver rather than a separate session API.""" + + async def handler( + ctx: ServerRequestContext, params: types.ReadResourceRequestParams + ) -> ReadResourceResult | types.InputRequiredResult: + assert params.uri == "memory://native" + if params.input_responses: + answer = params.input_responses["ask"] + assert isinstance(answer, types.ElicitResult) + assert answer.content is not None + return ReadResourceResult(contents=[TextResourceContents(uri=params.uri, text=str(answer.content["name"]))]) + return types.InputRequiredResult(input_requests={"ask": _name_elicitation()}) + + server = Server("native", on_read_resource=handler) + + async def elicitation_callback( + context: ClientRequestContext, params: types.ElicitRequestParams + ) -> types.ElicitResult: + return types.ElicitResult(action="accept", content={"name": "Alice"}) + + with anyio.fail_after(5): + async with server.serve() as runtime: + async with Client( + DispatcherTransport(dispatcher_connection(runtime)), elicitation_callback=elicitation_callback + ) as client: + result = await client.read_resource("memory://native") + assert result.model_dump(by_alias=True, mode="json") == snapshot( + { + "_meta": {"io.modelcontextprotocol/serverInfo": {"name": "native", "version": ""}}, + "ttlMs": 0, + "cacheScope": "private", + "contents": [{"uri": "memory://native", "mimeType": None, "_meta": None, "text": "Alice"}], + "resultType": "complete", + } + ) + + +async def test_dispatcher_transport_rejects_legacy_handshake_without_stopping_runtime() -> None: + """The server entry is modern-only. ClientSession exposes initialize without Client's exception-group wrapping.""" + with anyio.fail_after(5): + async with Server("native").serve() as runtime: + async with dispatcher_connection(runtime) as dispatcher, ClientSession(dispatcher=dispatcher) as session: + with pytest.raises(MCPError) as exc: + await session.initialize() + assert exc.value.code == types.UNSUPPORTED_PROTOCOL_VERSION + assert exc.value.message == snapshot( + "connection is serving the 2026-07-28 protocol; the initialize handshake is not accepted" + ) + async with Client(DispatcherTransport(dispatcher_connection(runtime))) as client: + version = client.protocol_version + assert version == "2026-07-28" + + +async def test_inprocess_client_exit_joins_handler_before_closing_lifespan() -> None: + """An in-process handler runs in its caller's task, but its resources must remain alive through cleanup.""" + entered = anyio.Event() + cleaning = anyio.Event() + release = anyio.Event() + cleaned = anyio.Event() + stop = anyio.Event() + client_closed = anyio.Event() + lifespan_closed = anyio.Event() + call_finished = anyio.Event() + + @asynccontextmanager + async def lifespan(server: MCPServer[None]) -> AsyncIterator[None]: + try: + yield None + finally: + assert cleaned.is_set() + lifespan_closed.set() + + server = MCPServer("in-process shutdown", lifespan=lifespan) + + @server.tool() + async def wait() -> str: + entered.set() + try: + await anyio.sleep_forever() + finally: + with anyio.CancelScope(shield=True): + cleaning.set() + await release.wait() + assert not lifespan_closed.is_set() + cleaned.set() + raise NotImplementedError + + async def own_client(*, task_status: anyio.abc.TaskStatus[Client]) -> None: + async with Client(server) as client: + task_status.started(client) + await stop.wait() + client_closed.set() + + async def call(client: Client) -> None: + with pytest.raises(MCPError) as exc: + await client.call_tool("wait") + assert exc.value.code == types.CONNECTION_CLOSED + call_finished.set() + + with anyio.fail_after(5): + async with anyio.create_task_group() as tg: + client = await tg.start(own_client) + tg.start_soon(call, client) + try: + await entered.wait() + stop.set() + await cleaning.wait() + await anyio.wait_all_tasks_blocked() + assert not client_closed.is_set() + assert not lifespan_closed.is_set() + finally: + release.set() + await client_closed.wait() + await call_finished.wait() + assert lifespan_closed.is_set() + + +async def test_dispatcher_transport_propagates_connection_opening_failure() -> None: + """Client construction is lazy and a native connection opening failure reaches the caller unchanged.""" + failure = OSError("connection unavailable") + opened = False + + @asynccontextmanager + async def connection() -> AsyncIterator[Dispatcher[TransportContext]]: + nonlocal opened + opened = True + raise failure + yield + + client = Client(DispatcherTransport(connection())) + assert not opened + with pytest.raises(OSError) as exc: + async with client: + raise NotImplementedError + assert exc.value is failure + assert opened diff --git a/tests/docs_src/test_authorization.py b/tests/docs_src/test_authorization.py index 00c9adc81c..1da83f9e16 100644 --- a/tests/docs_src/test_authorization.py +++ b/tests/docs_src/test_authorization.py @@ -1,20 +1,88 @@ """`docs/run/authorization.md`: every claim the page makes, proved against the real SDK.""" +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +import anyio import httpx2 import pytest from inline_snapshot import snapshot -from mcp_types import TextContent +from mcp_types import ( + INVALID_PARAMS, + ElicitRequest, + ElicitRequestFormParams, + ElicitResult, + InputRequiredResult, + InputResponses, + TextContent, +) from starlette.routing import Route -from docs_src.authorization import tutorial001, tutorial002 -from mcp import Client +from docs_src.authorization import tutorial001, tutorial002, tutorial003 +from mcp import Client, MCPError from mcp.client.streamable_http import streamable_http_client from mcp.server import MCPServer +from mcp.server.mcpserver import Context +from mcp.shared.memory import create_client_server_memory_streams +from mcp.shared.transport import MessageMetadata, TransportContext, TransportStreams # See test_index.py for why this is a per-module mark and not a conftest hook. pytestmark = [pytest.mark.anyio, pytest.mark.filterwarnings("error::mcp.MCPDeprecationWarning")] +async def test_verified_transport_identity_binds_the_published_request_state_example() -> None: + """tutorial003: issued state accepts its verified peer and rejects an anonymous retry through the public runtime.""" + + @tutorial003.mcp.tool() + async def confirm(ctx: Context) -> str | InputRequiredResult: + if ctx.input_responses is not None: + assert isinstance(ctx.request_state, str) + return ctx.request_state + return InputRequiredResult( + input_requests={ + "confirm": ElicitRequest( + params=ElicitRequestFormParams( + message="Confirm?", + requested_schema={"type": "object", "properties": {}}, + ) + ) + }, + request_state="approved", + ) + + with anyio.fail_after(5): + async with tutorial003.mcp.serve() as runtime: + + @asynccontextmanager + async def connection(verified: bool) -> AsyncIterator[TransportStreams]: + async with create_client_server_memory_streams() as (client_streams, server_streams): + + def builder(metadata: MessageMetadata) -> TransportContext: + if verified: + return tutorial003.VerifiedPeer(kind="broker", can_send_request=False, principal="alice") + return TransportContext(kind="broker", can_send_request=False) + + @asynccontextmanager + async def transport() -> AsyncIterator[TransportStreams]: + yield server_streams + + await runtime.connect(transport(), transport_builder=builder) + yield client_streams + + async with Client(connection(True)) as alice, Client(connection(False)) as anonymous: + pending = await alice.session.call_tool("confirm", allow_input_required=True) + assert isinstance(pending, InputRequiredResult) + assert pending.request_state is not None + responses: InputResponses = {"confirm": ElicitResult(action="accept")} + with pytest.raises(MCPError) as exc: + await anonymous.call_tool("confirm", input_responses=responses, request_state=pending.request_state) + assert exc.value.code == INVALID_PARAMS + result = await alice.call_tool( + "confirm", input_responses=responses, request_state=pending.request_state + ) + assert result.structured_content == {"result": "approved"} + + async def test_the_in_memory_client_never_authenticates() -> None: """tutorial001: `Client(mcp)` connects to the server object directly, so no token is ever checked.""" async with Client(tutorial001.mcp) as client: diff --git a/tests/docs_src/test_client_transports.py b/tests/docs_src/test_client_transports.py index 914067c7a0..944370c5ae 100644 --- a/tests/docs_src/test_client_transports.py +++ b/tests/docs_src/test_client_transports.py @@ -4,7 +4,7 @@ import pytest -from docs_src.client_transports import tutorial001, tutorial004 +from docs_src.client_transports import tutorial001, tutorial004, tutorial005, tutorial006 from mcp import Client from mcp.client.stdio import get_default_environment from mcp.client.streamable_http import streamable_http_client @@ -19,6 +19,16 @@ async def test_the_in_memory_program_on_the_page_runs(capsys: pytest.CaptureFixt assert "Found 3 books matching 'dune'." in capsys.readouterr().out +async def test_custom_transport_example_serves_independent_peers() -> None: + """The public adapter example runs both clients through stream-backed server dispatch.""" + await tutorial005.main() + + +async def test_dispatcher_transport_example_uses_the_shared_mcp_pipeline() -> None: + """The explicit dispatcher wrapper drives a complete client/server call without message streams.""" + await tutorial006.main() + + async def test_in_memory_client_talks_to_the_server_object() -> None: """tutorial001: passing the server object connects in-process. No subprocess, no port.""" async with Client(tutorial001.mcp) as client: diff --git a/tests/server/test_runner.py b/tests/server/test_runner.py index 50e77f7134..2913abf7ec 100644 --- a/tests/server/test_runner.py +++ b/tests/server/test_runner.py @@ -13,7 +13,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass, field, replace from functools import partial -from typing import Any, cast +from typing import Any, Literal, cast import anyio import anyio.abc @@ -32,6 +32,7 @@ SERVER_INFO_META_KEY, UNSUPPORTED_PROTOCOL_VERSION, CallToolRequestParams, + CallToolResult, ClientCapabilities, EmptyResult, ErrorData, @@ -57,6 +58,7 @@ ) import mcp.server.runner +from mcp import Client from mcp.server.caching import CacheHint from mcp.server.connection import Connection, NotifyOnlyOutbound from mcp.server.context import ServerRequestContext @@ -79,8 +81,10 @@ from mcp.shared.dispatcher import CallOptions from mcp.shared.exceptions import MCPError, NoBackChannelError from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher +from mcp.shared.memory import create_client_server_memory_streams from mcp.shared.message import MessageMetadata, SessionMessage from mcp.shared.peer import dump_params +from mcp.shared.transport import TransportStreams from mcp.shared.transport_context import TransportContext from ..shared.conftest import jsonrpc_pair @@ -2109,3 +2113,57 @@ async def test_dual_era_client_propagates_body_exception_unwrapped(server: SrvT) with pytest.raises(RuntimeError, match="boom"): async with dual_era_client(server): raise RuntimeError("boom") + + +@pytest.mark.anyio +@pytest.mark.parametrize("mode", ["legacy", "auto", "2026-07-28"]) +async def test_run_delivers_custom_transport_context_without_overriding_protocol_rules( + mode: Literal["legacy", "auto", "2026-07-28"], +) -> None: + """The SDK carries adapter metadata to handlers; modern protocol rules still deny server requests.""" + + @dataclass(kw_only=True, frozen=True) + class BrokerContext(TransportContext): + peer: str + + transport_context = BrokerContext(kind="broker", can_send_request=True, peer="alice") + + def build_context(metadata: MessageMetadata) -> BrokerContext: + return transport_context + + async def inspect_context(ctx: Ctx, params: CallToolRequestParams) -> CallToolResult: + assert params.name == "inspect" + assert isinstance(ctx.transport, BrokerContext) + return CallToolResult( + content=[], + structured_content={"peer": ctx.transport.peer, "can_send_request": ctx.transport.can_send_request}, + ) + + async def list_tools(ctx: Ctx, params: PaginatedRequestParams | None) -> ListToolsResult: + return ListToolsResult(tools=[Tool(name="inspect", input_schema={"type": "object"})]) + + app = Server("custom-transport", on_call_tool=inspect_context, on_list_tools=list_tools) + + @asynccontextmanager + async def transport() -> AsyncIterator[TransportStreams]: + async with create_client_server_memory_streams() as (client_streams, server_streams): + async with anyio.create_task_group() as tg: + tg.start_soon( + partial( + app.run, + *server_streams, + app.create_initialization_options(), + transport_builder=build_context, + ) + ) + yield client_streams + tg.cancel_scope.cancel() + + with anyio.fail_after(5): + async with Client(transport(), mode=mode) as client: + result = await client.call_tool("inspect") + assert result.structured_content == { + "peer": transport_context.peer, + "can_send_request": mode == "legacy", + } + assert transport_context.can_send_request is True diff --git a/tests/server/test_runtime.py b/tests/server/test_runtime.py new file mode 100644 index 0000000000..a9fa9e7d96 --- /dev/null +++ b/tests/server/test_runtime.py @@ -0,0 +1,652 @@ +"""Public custom-transport hosting behavior, without a network broker.""" + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass +from functools import partial +from typing import Any, Literal + +import anyio +import anyio.abc +import anyio.lowlevel +import pytest +from inline_snapshot import snapshot +from mcp_types import ( + INVALID_PARAMS, + CallToolRequestParams, + CallToolResult, + ElicitRequest, + ElicitRequestFormParams, + ElicitResult, + InputRequiredResult, + ListToolsResult, + PaginatedRequestParams, + RequestId, + TextContent, + Tool, +) + +from mcp import Client, MCPError +from mcp.server import Server +from mcp.server.context import ServerRequestContext +from mcp.server.mcpserver import Context, MCPServer +from mcp.server.request_state import RequestStateSecurity +from mcp.server.runtime import ServerRuntime +from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair +from mcp.shared.dispatcher import Dispatcher, OnNotify, OnNotifyIntercept, OnRequest +from mcp.shared.memory import create_client_server_memory_streams +from mcp.shared.transport import ( + DispatcherTransport, + MessageMetadata, + SessionMessage, + TransportContext, + TransportContextBuilder, + TransportStreams, +) + +pytestmark = pytest.mark.anyio + + +@asynccontextmanager +async def connect( + host: ServerRuntime[Any], *, transport_builder: TransportContextBuilder | None = None +) -> AsyncIterator[TransportStreams]: + async with create_client_server_memory_streams() as (client_streams, server_streams): + + @asynccontextmanager + async def transport() -> AsyncIterator[TransportStreams]: + async with server_streams[0], server_streams[1]: + yield server_streams + + await host.connect(transport(), transport_builder=transport_builder) + yield client_streams + + +@pytest.mark.parametrize("highlevel", [False, True]) +@pytest.mark.parametrize("modes", [("legacy", "legacy"), ("legacy", "2026-07-28"), ("2026-07-28", "2026-07-28")]) +async def test_runtime_shares_lifespan_and_isolates_clients(highlevel: bool, modes: tuple[str, str]) -> None: + """SDK hosting runs one lifespan while peers independently negotiate and dispatch overlapping requests.""" + lifecycle: list[str] = [] + entered = {"alice": anyio.Event(), "bob": anyio.Event()} + request_ids: dict[str, RequestId | None] = {} + + @asynccontextmanager + async def lifespan(server: Server[str] | MCPServer[str]) -> AsyncIterator[str]: + lifecycle.append("startup") + try: + yield "shared-state" + finally: + lifecycle.append("shutdown") + + async def inspect(name: str, request_id: RequestId | None) -> str: + request_ids[name] = request_id + entered[name].set() + await entered["bob" if name == "alice" else "alice"].wait() + return name + + if highlevel: + app = MCPServer("peers", lifespan=lifespan) + + @app.tool() + async def echo(name: str, ctx: Context[str]) -> str: + assert ctx.request_context.lifespan_context == "shared-state" + return await inspect(name, ctx.request_context.request_id) + + else: + + async def echo_lowlevel(ctx: ServerRequestContext[str], params: CallToolRequestParams) -> CallToolResult: + assert params.name == "echo" + assert ctx.lifespan_context == "shared-state" + assert params.arguments is not None + name = params.arguments["name"] + assert isinstance(name, str) + return CallToolResult(content=[TextContent(text=await inspect(name, ctx.request_id))]) + + async def list_tools(ctx: ServerRequestContext[str], params: PaginatedRequestParams | None) -> ListToolsResult: + return ListToolsResult(tools=[Tool(name="echo", input_schema={"type": "object"})]) + + app = Server("peers", lifespan=lifespan, on_call_tool=echo_lowlevel, on_list_tools=list_tools) + + results: dict[str, str] = {} + + async def call(client: Client, name: str) -> None: + result = await client.call_tool("echo", {"name": name}) + content = result.content[0] + assert isinstance(content, TextContent) + results[name] = content.text + + with anyio.fail_after(5): + async with app.serve() as host: + async with Client(connect(host), mode=modes[0]) as alice: + async with Client(connect(host), mode=modes[1]) as bob: + async with anyio.create_task_group() as tg: + tg.start_soon(call, alice, "alice") + tg.start_soon(call, bob, "bob") + if modes[0] == modes[1]: + assert request_ids["alice"] == request_ids["bob"] + await call(alice, "alice") + assert lifecycle == ["startup"] + assert lifecycle == ["startup", "shutdown"] + assert results == {"alice": "alice", "bob": "bob"} + + +@pytest.mark.parametrize("mode", ["legacy", "auto", "2026-07-28"]) +async def test_host_exposes_adapter_metadata_in_highlevel_handlers( + mode: Literal["legacy", "auto", "2026-07-28"], +) -> None: + """Adapter context survives hosting; the modern protocol's back-channel denial remains authoritative.""" + + @dataclass(kw_only=True, frozen=True) + class BrokerContext(TransportContext): + peer: str + + metadata = BrokerContext(kind="broker", can_send_request=True, peer="alice") + + def context_builder(message_metadata: MessageMetadata) -> BrokerContext: + return metadata + + app = MCPServer("metadata") + + @app.tool() + async def inspect_context(ctx: Context) -> dict[str, str | bool]: + assert isinstance(ctx.transport, BrokerContext) + return {"peer": ctx.transport.peer, "can_send_request": ctx.transport.can_send_request} + + with anyio.fail_after(5): + async with app.serve() as host, Client(connect(host, transport_builder=context_builder), mode=mode) as client: + result = await client.call_tool("inspect_context") + assert result.structured_content == {"peer": metadata.peer, "can_send_request": mode == "legacy"} + assert metadata.can_send_request is True + + +async def test_host_releases_transport_before_application_lifespan(monkeypatch: pytest.MonkeyPatch) -> None: + """Host exit cancels a connected peer and lets its adapter clean up before the application does.""" + events: list[str] = [] + + @asynccontextmanager + async def lifespan(server: Server[None]) -> AsyncIterator[None]: + try: + yield None + finally: + await anyio.lowlevel.checkpoint() + events.append("lifespan") + + request_send, request_receive = anyio.create_memory_object_stream[SessionMessage | Exception]() + response_send, response_receive = anyio.create_memory_object_stream[SessionMessage]() + async with request_send, request_receive, response_send, response_receive: + server_streams = request_receive, response_send + read_close, write_close = request_receive.aclose, response_send.aclose + + async def close_read() -> None: + await anyio.lowlevel.checkpoint() + events.append("read") + await read_close() + + async def close_write() -> None: + await anyio.lowlevel.checkpoint() + events.append("write") + await write_close() + + monkeypatch.setattr(server_streams[0], "aclose", close_read) + monkeypatch.setattr(server_streams[1], "aclose", close_write) + + @asynccontextmanager + async def transport() -> AsyncIterator[TransportStreams]: + try: + yield server_streams + finally: + await anyio.lowlevel.checkpoint() + events.append("transport") + + with anyio.fail_after(5): + async with Server("shutdown", lifespan=lifespan).serve() as host: + await host.connect(transport()) + assert events == ["read", "write", "transport", "lifespan"] + with pytest.raises(anyio.EndOfStream): + await response_receive.receive() + + +async def test_host_survives_an_adapter_failure_after_opening(caplog: pytest.LogCaptureFixture) -> None: + """An arbitrary adapter failure after readiness is isolated to its peer and logged with a traceback.""" + crash = anyio.Event() + stopped = anyio.Event() + + async def fail() -> None: + await crash.wait() + raise OSError("broker disconnected") + + with anyio.fail_after(5): + async with Server("isolation").serve() as host: + async with create_client_server_memory_streams() as (client_streams, server_streams): + + @asynccontextmanager + async def transport() -> AsyncIterator[TransportStreams]: + try: + async with anyio.create_task_group() as tg: + tg.start_soon(fail) + yield server_streams + finally: + stopped.set() + + await host.connect(transport()) + crash.set() + await stopped.wait() + async with Client(connect(host)) as healthy: + await healthy.session.discover() + with pytest.raises(anyio.EndOfStream): + await client_streams[0].receive() + assert stopped.is_set() + errors = [record for record in caplog.records if record.name == "mcp.server.runtime"] + assert len(errors) == 1 + assert errors[0].exc_info is not None + + +async def test_host_propagates_opening_failure_without_poisoning_other_connections() -> None: + """The caller of connect receives the original opening failure and can continue using the host.""" + failure = OSError("broker unavailable") + + @asynccontextmanager + async def unavailable() -> AsyncIterator[TransportStreams]: + raise failure + yield + + with anyio.fail_after(5): + async with Server("startup").serve() as host: + with pytest.raises(OSError) as exc: + await host.connect(unavailable()) + async with Client(connect(host)) as client: + await client.session.discover() + assert exc.value is failure + + +async def test_host_admission_waits_for_a_connection_slot() -> None: + """The host opens no more than its configured number of logical peers, then admits a waiting peer on EOF.""" + opened = anyio.Event() + + with anyio.fail_after(5): + async with Server("capacity").serve(max_connections=1) as host: + async with ( + create_client_server_memory_streams() as (first_client, first_server), + create_client_server_memory_streams() as (second_client, second_server), + anyio.create_task_group() as tg, + ): + + @asynccontextmanager + async def first() -> AsyncIterator[TransportStreams]: + yield first_server + + @asynccontextmanager + async def second() -> AsyncIterator[TransportStreams]: + opened.set() + yield second_server + + await host.connect(first()) + tg.start_soon(host.connect, second()) + await anyio.wait_all_tasks_blocked() + assert not opened.is_set() + await first_client[1].aclose() + await opened.wait() + await second_client[1].aclose() + assert opened.is_set() + + +async def test_closed_host_rejects_connections_without_entering_the_transport() -> None: + """Holding a host after its context exits does not allow new peers to outlive application lifespan.""" + + @asynccontextmanager + async def transport() -> AsyncIterator[TransportStreams]: + raise NotImplementedError + yield + + with anyio.fail_after(5): + async with Server("closed").serve() as host: + pass + with pytest.raises(RuntimeError) as exc: + await host.connect(transport()) + assert str(exc.value) == snapshot("Server runtime is closed") + + +@pytest.mark.parametrize("limit", [0, -1]) +async def test_host_rejects_nonpositive_capacity_before_starting_lifespan(limit: int) -> None: + """Invalid admission limits fail before the server acquires application resources.""" + + @asynccontextmanager + async def lifespan(server: Server[None]) -> AsyncIterator[None]: + raise NotImplementedError + yield None + + with pytest.raises(ValueError) as exc: + async with Server("invalid", lifespan=lifespan).serve(max_connections=limit): + raise NotImplementedError + assert str(exc.value) == snapshot("max_connections must be positive") + + +async def test_host_shields_transport_and_lifespan_cleanup_from_parent_cancellation() -> None: + """Cancelling the host owner still lets both cleanup layers perform asynchronous resource release.""" + events: list[str] = [] + + @asynccontextmanager + async def lifespan(server: Server[None]) -> AsyncIterator[None]: + try: + yield None + finally: + await anyio.lowlevel.checkpoint() + events.append("lifespan") + + with anyio.fail_after(5): + async with create_client_server_memory_streams() as (_, server_streams): + + @asynccontextmanager + async def transport() -> AsyncIterator[TransportStreams]: + try: + yield server_streams + finally: + await anyio.lowlevel.checkpoint() + events.append("transport") + + with anyio.CancelScope() as scope: + async with Server("cancel", lifespan=lifespan).serve() as host: + await host.connect(transport()) + scope.cancel() + await anyio.sleep_forever() + assert scope.cancelled_caught + assert events == ["transport", "lifespan"] + + +@pytest.mark.parametrize("stall", ["transport", "lifespan"]) +async def test_host_abandons_unresponsive_cleanup(stall: str, caplog: pytest.LogCaptureFixture) -> None: + """SDK cleanup deadlines interrupt a stuck adapter or lifespan without parking shutdown forever.""" + interrupted = anyio.Event() + + async def cleanup(layer: str) -> None: + if layer == stall: + try: + await anyio.sleep_forever() + finally: + interrupted.set() + + @asynccontextmanager + async def lifespan(server: Server[None]) -> AsyncIterator[None]: + try: + yield None + finally: + await cleanup("lifespan") + + # The behavior under test includes the documented five-second cleanup grace. + with anyio.fail_after(10): + async with create_client_server_memory_streams() as (_, server_streams): + + @asynccontextmanager + async def transport() -> AsyncIterator[TransportStreams]: + try: + yield server_streams + finally: + await cleanup("transport") + + async with Server("stalled", lifespan=lifespan).serve() as host: + await host.connect(transport()) + assert interrupted.is_set() + records = [record for record in caplog.records if record.name == "mcp.server.runtime"] + assert len(records) == 1 + assert records[0].levelname == "WARNING" + + +async def test_runtime_bounds_transport_cleanup_after_normal_peer_eof(caplog: pytest.LogCaptureFixture) -> None: + """A peer closing its stream finishes dispatch before the adapter's bounded cleanup begins.""" + interrupted = anyio.Event() + # Normal EOF, followed by the documented five-second cleanup grace. + with anyio.fail_after(10): + async with create_client_server_memory_streams() as (client_streams, server_streams): + + @asynccontextmanager + async def connection() -> AsyncIterator[TransportStreams]: + try: + yield server_streams + finally: + try: + await anyio.sleep_forever() + finally: + interrupted.set() + + async with Server("eof").serve() as runtime: + await runtime.connect(connection()) + await client_streams[1].aclose() + await interrupted.wait() + assert interrupted.is_set() + warnings = [record for record in caplog.records if record.name == "mcp.server.runtime"] + assert len(warnings) == 1 + assert warnings[0].levelname == "WARNING" + + +async def test_runtime_preserves_body_failure_when_lifespan_cleanup_times_out() -> None: + """The cleanup deadline must not turn a failed listener into a successful context-manager exit.""" + failure = RuntimeError("listener failed") + + @asynccontextmanager + async def lifespan(server: Server[None]) -> AsyncIterator[None]: + try: + yield None + finally: + await anyio.sleep_forever() + + # This failure path includes the documented five-second cleanup grace. + with anyio.fail_after(10), pytest.RaisesGroup(RuntimeError) as exc: + async with Server("failed-listener", lifespan=lifespan).serve(): + raise failure + assert exc.value.exceptions == (failure,) + + +@pytest.mark.parametrize("option", ["session_id", "transport_builder"]) +async def test_native_runtime_rejects_stream_only_options_before_opening(option: str) -> None: + """Native dispatchers supply their contexts and do not acquire legacy sessions from stream-hosting options.""" + + @asynccontextmanager + async def connection() -> AsyncIterator[Dispatcher[TransportContext]]: + raise NotImplementedError + yield + + def builder(metadata: MessageMetadata) -> TransportContext: + raise NotImplementedError + + with anyio.fail_after(5): + async with Server("native").serve() as runtime: + with pytest.raises(ValueError) as exc: + await runtime.connect( + DispatcherTransport(connection()), + session_id="session" if option == "session_id" else None, + transport_builder=builder if option == "transport_builder" else None, + ) + assert str(exc.value) == snapshot( + "Dispatcher transports supply their own context and do not use handshake-era sessions" + ) + + +async def test_runtime_waits_for_native_dispatcher_readiness(monkeypatch: pytest.MonkeyPatch) -> None: + """A native adapter is not connected until its receive loop has installed the MCP callbacks.""" + _, dispatcher = create_direct_dispatcher_pair() + entered = anyio.Event() + release = anyio.Event() + connected = anyio.Event() + run = dispatcher.run + + async def delayed_run( + on_request: OnRequest, + on_notify: OnNotify, + on_notify_intercept: OnNotifyIntercept | None = None, + *, + task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED, + ) -> None: + entered.set() + await release.wait() + await run(on_request, on_notify, on_notify_intercept, task_status=task_status) + + monkeypatch.setattr(dispatcher, "run", delayed_run) + + @asynccontextmanager + async def connection() -> AsyncIterator[Dispatcher[TransportContext]]: + yield dispatcher + + with anyio.fail_after(5): + async with Server("readiness").serve() as runtime: + + async def connect_native() -> None: + await runtime.connect(DispatcherTransport(connection())) + connected.set() + + async with anyio.create_task_group() as tg: + tg.start_soon(connect_native) + await entered.wait() + await anyio.wait_all_tasks_blocked() + assert not connected.is_set() + release.set() + await connected.wait() + assert connected.is_set() + + +async def test_runtime_propagates_native_dispatcher_startup_failure(monkeypatch: pytest.MonkeyPatch) -> None: + """A native receive-loop startup failure reaches connect's caller rather than being logged after false readiness.""" + _, dispatcher = create_direct_dispatcher_pair() + failure = RuntimeError("could not install RPC handlers") + + async def failing_run( + on_request: OnRequest, + on_notify: OnNotify, + on_notify_intercept: OnNotifyIntercept | None = None, + *, + task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED, + ) -> None: + raise failure + + monkeypatch.setattr(dispatcher, "run", failing_run) + + @asynccontextmanager + async def connection() -> AsyncIterator[Dispatcher[TransportContext]]: + yield dispatcher + + with anyio.fail_after(5): + async with Server("startup").serve() as runtime: + with pytest.raises(RuntimeError) as exc: + await runtime.connect(DispatcherTransport(connection())) + assert exc.value is failure + + +async def test_runtime_preserves_startup_failure_when_transport_cleanup_times_out( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A stuck adapter cleanup must not replace its receive-loop startup failure with a false-readiness error.""" + _, dispatcher = create_direct_dispatcher_pair() + failure = RuntimeError("native dispatcher failed to start") + + async def failing_run( + on_request: OnRequest, + on_notify: OnNotify, + on_notify_intercept: OnNotifyIntercept | None = None, + *, + task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED, + ) -> None: + raise failure + + monkeypatch.setattr(dispatcher, "run", failing_run) + + @asynccontextmanager + async def connection() -> AsyncIterator[Dispatcher[TransportContext]]: + try: + yield dispatcher + finally: + await anyio.sleep_forever() + + # The failure includes the documented five-second adapter cleanup grace. + with anyio.fail_after(10): + async with Server("startup").serve() as runtime: + with pytest.raises(RuntimeError) as exc: + await runtime.connect(DispatcherTransport(connection())) + assert exc.value is failure + + +async def test_runtime_keeps_request_state_bound_to_verified_peer_metadata() -> None: + """The existing principal hook binds sealed state to adapter metadata, not caller-supplied claims. + + Both peers mint state concurrently, a cross-peer replay fails, and the original peer can still complete its request. + """ + + @dataclass(kw_only=True, frozen=True) + class VerifiedPeer(TransportContext): + principal: str + + def context_builder(metadata: MessageMetadata, *, principal: str) -> VerifiedPeer: + return VerifiedPeer(kind="broker", can_send_request=True, principal=principal) + + def bind_principal(ctx: ServerRequestContext) -> str: + if not isinstance(ctx.transport, VerifiedPeer): + raise ValueError("Verified transport identity is required") + return ctx.transport.principal + + server = MCPServer( + "principals", + request_state_security=RequestStateSecurity( + keys=[b"test-key-for-principal-binding-32"], bind_principal=bind_principal + ), + ) + entered = {"alice": anyio.Event(), "bob": anyio.Event()} + + @server.tool() + async def confirm(ctx: Context) -> str | InputRequiredResult: + if ctx.input_responses is not None: + assert ctx.request_state is not None + return ctx.request_state + assert isinstance(ctx.transport, VerifiedPeer) + principal = ctx.transport.principal + entered[principal].set() + await entered["bob" if principal == "alice" else "alice"].wait() + return InputRequiredResult( + input_requests={ + "confirm": ElicitRequest( + params=ElicitRequestFormParams( + message="Confirm?", requested_schema={"type": "object", "properties": {}} + ) + ) + }, + request_state=principal, + ) + + states: dict[str, str] = {} + + async def mint(client: Client, principal: str) -> None: + result = await client.session.call_tool("confirm", allow_input_required=True) + assert isinstance(result, InputRequiredResult) + assert result.request_state is not None + states[principal] = result.request_state + + with anyio.fail_after(5): + async with server.serve() as runtime: + async with ( + Client(connect(runtime, transport_builder=partial(context_builder, principal="alice"))) as alice, + Client(connect(runtime, transport_builder=partial(context_builder, principal="bob"))) as bob, + ): + async with anyio.create_task_group() as tg: + tg.start_soon(mint, alice, "alice") + tg.start_soon(mint, bob, "bob") + with pytest.raises(MCPError) as exc: + await bob.session.call_tool( + "confirm", + input_responses={"confirm": ElicitResult(action="accept")}, + request_state=states["alice"], + meta={"principal": "alice"}, + ) + assert exc.value.code == INVALID_PARAMS + assert exc.value.message == snapshot("Invalid or expired requestState") + async with Client(connect(runtime)) as anonymous: + with pytest.raises(MCPError) as missing_identity: + await anonymous.session.call_tool( + "confirm", + input_responses={"confirm": ElicitResult(action="accept")}, + request_state=states["alice"], + ) + assert missing_identity.value.code == INVALID_PARAMS + result = await alice.call_tool( + "confirm", + input_responses={"confirm": ElicitResult(action="accept")}, + request_state=states["alice"], + meta={"principal": "bob"}, + ) + assert result.structured_content == {"result": "alice"} diff --git a/tests/shared/test_dispatcher.py b/tests/shared/test_dispatcher.py index c6ebb401ff..b2ba3e1848 100644 --- a/tests/shared/test_dispatcher.py +++ b/tests/shared/test_dispatcher.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any import anyio +import anyio.abc import pytest from mcp_types import ( CONNECTION_CLOSED, @@ -573,6 +574,154 @@ def broken_intercept(method: str, params: Mapping[str, Any] | None) -> bool: assert [method for method, _ in crec.notifications] == ["notifications/survives"] +@pytest.mark.anyio +@pytest.mark.parametrize("closing_side", ["client", "server"]) +@pytest.mark.parametrize("operation", ["request", "notification"]) +@pytest.mark.parametrize("swallow_cancel", [False, True]) +async def test_direct_close_joins_in_flight_handler_cleanup( + closing_side: str, + operation: str, + swallow_cancel: bool, +) -> None: + """Closing either peer interrupts its conversation and keeps run alive until shielded handler cleanup finishes.""" + entered = anyio.Event() + cleaning = anyio.Event() + release = anyio.Event() + cleaned = anyio.Event() + finished = anyio.Event() + stopped = {"client": anyio.Event(), "server": anyio.Event()} + + async def handle() -> None: + entered.set() + try: + await anyio.sleep_forever() + except anyio.get_cancelled_exc_class(): + if not swallow_cancel: + raise + finally: + with anyio.CancelScope(shield=True): + cleaning.set() + await release.wait() + cleaned.set() + + async def request( + ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None + ) -> dict[str, Any]: + assert method == "work" + await handle() + return {} + + async def notify(ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None) -> None: + assert method == "work" + await handle() + + client, server = create_direct_dispatcher_pair() + + async def run(dispatcher: DirectDispatcher, side: str, *, task_status: anyio.abc.TaskStatus[None]) -> None: + await dispatcher.run(request, notify, task_status=task_status) + assert cleaned.is_set() + stopped[side].set() + + async def call() -> None: + if operation == "request": + with pytest.raises(MCPError) as exc: + await client.send_raw_request("work", None) + assert exc.value.code == CONNECTION_CLOSED + else: + await client.notify("work", None) + finished.set() + + with anyio.fail_after(5): + async with anyio.create_task_group() as tg: + await tg.start(run, client, "client") + await tg.start(run, server, "server") + tg.start_soon(call) + try: + await entered.wait() + (client if closing_side == "client" else server).close() + await cleaning.wait() + await anyio.wait_all_tasks_blocked() + assert not stopped[closing_side].is_set() + release.set() + await stopped[closing_side].wait() + await finished.wait() + finally: + release.set() + client.close() + server.close() + assert cleaned.is_set() + + +@pytest.mark.anyio +@pytest.mark.parametrize("closing_side", ["client", "server"]) +async def test_direct_close_cancels_nested_backchannel_requests(closing_side: str) -> None: + """Nested calls share the conversation lifetime even though both handlers execute in the originating task.""" + entered = anyio.Event() + cleaned: list[str] = [] + finished = anyio.Event() + + async def outer( + ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None + ) -> dict[str, Any]: + assert method == "outer" + try: + return await ctx.send_raw_request("inner", None) + finally: + cleaned.append("outer") + + async def inner( + ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None + ) -> dict[str, Any]: + assert method == "inner" + entered.set() + try: + await anyio.sleep_forever() + finally: + cleaned.append("inner") + raise NotImplementedError + + async def unused_notify( + ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None + ) -> None: + raise NotImplementedError + + client, server = create_direct_dispatcher_pair() + + async def call() -> None: + with pytest.raises(MCPError) as exc: + await client.send_raw_request("outer", None) + assert exc.value.code == CONNECTION_CLOSED + finished.set() + + with anyio.fail_after(5): + async with anyio.create_task_group() as tg: + await tg.start(client.run, inner, unused_notify) + await tg.start(server.run, outer, unused_notify) + tg.start_soon(call) + await entered.wait() + (client if closing_side == "client" else server).close() + await finished.wait() + client.close() + server.close() + assert cleaned == ["inner", "outer"] + + +@pytest.mark.anyio +async def test_direct_notification_handler_errors_are_not_mistaken_for_connection_shutdown() -> None: + """Shutdown drops interrupted notifications without swallowing an MCPError raised by a live handler.""" + failure = MCPError(code=CONNECTION_CLOSED, message="handler refusal") + + async def notify(ctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None) -> None: + assert method == "example/event" + raise failure + + with anyio.fail_after(5): + async with running_pair(direct_pair, server_on_notify=notify) as (client, *_): + with pytest.raises(MCPError) as exc: + await client.notify("example/event", None) + assert exc.value is failure + + if TYPE_CHECKING: _d: Dispatcher[TransportContext] = DirectDispatcher(TransportContext(kind="direct", can_send_request=True)) _o: Outbound = _d diff --git a/tests/shared/test_message.py b/tests/shared/test_message.py new file mode 100644 index 0000000000..8b6ac0b74a --- /dev/null +++ b/tests/shared/test_message.py @@ -0,0 +1,14 @@ +from mcp.shared.transport import ServerMessageMetadata, SessionMessage, TransportContext +from mcp.types import JSONRPCRequest + + +def test_transport_headers_are_not_exposed_by_message_representations() -> None: + """SDK-defined: framing metadata preserves header access without adding credentials to debug representations.""" + credential = "private-bearer-value" + context = TransportContext(kind="http", can_send_request=False, headers={"authorization": credential}) + metadata = ServerMessageMetadata(None, None, None, None, None, False, transport_context=context) + message = SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"), metadata) + assert metadata.transport_context is context + assert not metadata.can_send_request + assert credential not in repr(metadata) + assert credential not in repr(message) diff --git a/tests/shared/test_sse.py b/tests/shared/test_sse.py index 77d1b28a0a..f1b28c7848 100644 --- a/tests/shared/test_sse.py +++ b/tests/shared/test_sse.py @@ -241,7 +241,11 @@ async def test_sse_client_basic_connection_mounted_app() -> None: async def _handle_context_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult: assert params.name in ("echo_headers", "echo_context") assert ctx.request is not None - headers_info = dict(ctx.request.headers) + assert ctx.transport is not None + assert ctx.transport.kind == "sse" + assert ctx.transport.headers == ctx.request.headers + assert ctx.transport.headers is not None + headers_info = dict(ctx.transport.headers) if params.name == "echo_headers": return CallToolResult(content=[TextContent(type="text", text=json.dumps(headers_info))]) diff --git a/tests/shared/test_streamable_http.py b/tests/shared/test_streamable_http.py index 655d9941dc..e0d2e96c5f 100644 --- a/tests/shared/test_streamable_http.py +++ b/tests/shared/test_streamable_http.py @@ -1410,14 +1410,18 @@ async def _handle_context_call_tool( ) -> CallToolResult: assert params.name in ("echo_headers", "echo_context") assert isinstance(ctx.request, Request) + assert ctx.transport is not None + assert ctx.transport.kind == "streamable-http" + assert ctx.transport.headers == ctx.request.headers + assert ctx.transport.headers is not None if params.name == "echo_headers": - return CallToolResult(content=[TextContent(type="text", text=json.dumps(dict(ctx.request.headers)))]) + return CallToolResult(content=[TextContent(type="text", text=json.dumps(dict(ctx.transport.headers)))]) assert params.arguments is not None context_data: dict[str, Any] = { "request_id": params.arguments.get("request_id"), - "headers": dict(ctx.request.headers), + "headers": dict(ctx.transport.headers), "method": ctx.request.method, "path": ctx.request.url.path, "protocol_version": ctx.protocol_version,