[v2] ClientSession runs on JSONRPCDispatcher; BaseSession removed (#2838)

This commit is contained in:
Max 2026-06-15 14:46:34 +01:00 • committed by GitHub
commit 1012d60004
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 2886 additions and 1734 deletions

View file

@ -193,11 +193,13 @@ many requirements at once; if the assertions would be separate, write separate t
### Notifications and concurrency
The client's receive loop dispatches each incoming message to completion before reading the next,
and the in-memory transport delivers everything on one ordered stream. Together these guarantee
that every notification a server handler emits before its response reaches the client callback
before the originating request returns — so tests collect notifications into a plain list and
assert after the call, with no synchronisation. The exceptions:
The client's dispatcher starts a task per incoming notification in arrival order but does not
await it before reading the next message, so completion order is not structural. What still
holds: the in-memory transport delivers everything on one ordered stream, and a callback that
records synchronously (no `await` before the append) finishes its scheduling slice before the
awaited request's waiter — woken strictly later — resumes. So tests whose callbacks are plain
appends may still collect into a list and assert after the call. A callback that awaits before
recording loses that ordering and must synchronise. The other exceptions:
- a notification not triggered by a request the test is awaiting needs an `anyio.Event` set in
the receiving handler and awaited under `anyio.fail_after(5)`;
@ -220,9 +222,8 @@ but still inside an outer `async with`, and no restructure can avoid it.
A handful of `# pragma: lax no cover` markers in `src/` cover teardown exception handlers whose
execution is timing-dependent under the in-process HTTP bridge — the POST-stream and
stateless-session `except Exception` handlers in `server/streamable_http*.py`, the `_terminated`
check in `message_router`, and the response-stream double-close guard in
`BaseSession._receive_loop`. `strict-no-cover` does not check `lax` lines; do not promote them to
strict `no cover` without first making the teardown ordering deterministic. The suite also relies
on a one-line `src/mcp/server/sse.py` fix (`sse_stream_reader.aclose()`) that closes a stream the
SSE leg would otherwise leak.
stateless-session `except Exception` handlers in `server/streamable_http*.py` and the
`_terminated` check in `message_router`. `strict-no-cover` does not check `lax` lines; do not
promote them to strict `no cover` without first making the teardown ordering deterministic. The
suite also relies on a one-line `src/mcp/server/sse.py` fix (`sse_stream_reader.aclose()`) that
closes a stream the SSE leg would otherwise leak.

View file

@ -67,8 +67,9 @@ class _RecordingWriteStream:
self._log = log
async def send(self, item: SessionMessage, /) -> None:
self._log.append(item)
# Record only after the inner send returns: a failed or cancelled send never reached the transport.
await self._inner.send(item)
self._log.append(item)
async def aclose(self) -> None:
await self._inner.aclose()

View file

@ -268,18 +268,15 @@ REQUIREMENTS: dict[str, Requirement] = {
divergence=Divergence(
note=(
"The spec says receivers of a cancellation SHOULD NOT send a response for the cancelled "
"request; the server sends an error response (code 0, 'Request cancelled'), which is what "
"unblocks the SDK client's pending call."
"request; both seats send an error response (code 0, 'Request cancelled') instead — the "
"server for cancelled client requests, and the client for cancelled server-initiated "
"requests — which is what unblocks the sender's pending call."
),
),
),
"protocol:cancel:initialize-not-cancellable": Requirement(
source=f"{SPEC_BASE_URL}/basic/utilities/cancellation#behavior-requirements",
behavior="The client never sends notifications/cancelled for the initialize request.",
deferred=(
"Not implemented in the SDK: the client has no public cancellation API at all, so no pathway "
"exists that could cancel initialize; there is no distinct behaviour to pin beyond that absence."
),
),
"protocol:cancel:late-response-ignored": Requirement(
source=f"{SPEC_BASE_URL}/basic/utilities/cancellation#behavior-requirements",
@ -287,14 +284,6 @@ REQUIREMENTS: dict[str, Requirement] = {
"A response that arrives after the sender issued notifications/cancelled is ignored; the "
"request stays failed and no error is raised."
),
divergence=Divergence(
note=(
"A response whose id matches no in-flight request is delivered to the message handler "
"as a RuntimeError rather than being silently ignored. The post-cancellation case is the "
"same code path; tested in its unknown-id form because that is deterministic without the "
"client-side cancellation API the SDK does not yet provide."
),
),
),
"protocol:cancel:server-survives": Requirement(
source="sdk",
@ -306,19 +295,6 @@ REQUIREMENTS: dict[str, Requirement] = {
"A server that abandons an in-flight server-initiated request (sampling, elicitation, roots) "
"cancels it, and the client stops processing the cancelled request."
),
divergence=Divergence(
note=(
"Abandoning a server-side send_request emits no cancellation notification, and the client "
"could not act on one anyway: client callbacks run inline in the receive loop, so a "
"cancellation is not even read until the callback has finished."
),
),
deferred=(
"Not implemented in the SDK: abandoning a server-side send_request emits no cancellation "
"notification (the same sender-side gap recorded on protocol:timeout:sends-cancellation), and "
"the client could not act on one anyway because client callbacks run inline in the receive "
"loop, so a cancellation would not even be read until the callback had already finished."
),
),
"protocol:cancel:unknown-id-ignored": Requirement(
source=f"{SPEC_BASE_URL}/basic/utilities/cancellation#error-handling",
@ -363,6 +339,27 @@ REQUIREMENTS: dict[str, Requirement] = {
source=f"{SPEC_BASE_URL}/basic#responses",
behavior="A request whose method has no registered handler is answered with a METHOD_NOT_FOUND error.",
),
"protocol:error:null-id": Requirement(
source="sdk",
behavior=(
"An error response carrying a null id — the JSON-RPC shape for a peer reporting a failure it "
"could not attribute to a request, such as a parse error — is surfaced to the application "
"rather than silently discarded."
),
divergence=Divergence(
note=(
"The dispatcher drops null-id error responses with a debug log; in v1, JSONRPCError.id was "
"non-nullable, so a null-id error response failed transport validation and the resulting "
"ValidationError was surfaced to message_handler as an exception. A typed fault channel "
"restoring visibility is planned before v2 stable."
),
),
deferred=(
"Not yet covered here: the current drop is pinned at the dispatcher level by "
"tests/shared/test_jsonrpc_dispatcher.py; an interaction-level test waits on the planned "
"fault channel."
),
),
"protocol:meta:related-task": Requirement(
source=f"{SPEC_BASE_URL}/basic/utilities/tasks#related-task-metadata",
behavior="Messages may carry related-task _meta associating them with a task.",
@ -466,11 +463,6 @@ REQUIREMENTS: dict[str, Requirement] = {
"When a request times out, the sender issues notifications/cancelled for that request before "
"failing the local call."
),
divergence=Divergence(
note=(
"The client only raises locally and sends nothing on timeout, so the server keeps running the handler."
),
),
),
"protocol:timeout:session-survives": Requirement(
source=f"{SPEC_BASE_URL}/basic/lifecycle#timeouts",

View file

@ -11,11 +11,12 @@ import pytest
from inline_snapshot import snapshot
from mcp import MCPError, types
from mcp.client import ClientSession
from mcp.client import ClientRequestContext, ClientSession
from mcp.server import Server, ServerRequestContext
from mcp.shared.memory import MessageStream, create_client_server_memory_streams
from mcp.shared.message import SessionMessage
from mcp.types import (
REQUEST_TIMEOUT,
CallToolResult,
EmptyResult,
ErrorData,
@ -155,14 +156,71 @@ async def test_cancellation_for_unknown_request_is_ignored(connect: Connect) ->
assert result == snapshot(CallToolResult(content=[TextContent(text="unbothered")]))
@requirement("protocol:cancel:server-to-client")
async def test_abandoned_server_request_cancels_the_client_callback(connect: Connect) -> None:
"""A server that abandons a sampling request cancels it, interrupting the client's callback mid-await."""
callback_started = anyio.Event()
callback_cancelled = anyio.Event()
async def sampling_callback(
context: ClientRequestContext, params: types.CreateMessageRequestParams
) -> types.CreateMessageResult:
callback_started.set()
try:
await anyio.Event().wait() # blocks until the cancellation interrupts it
except anyio.get_cancelled_exc_class():
callback_cancelled.set()
raise
raise NotImplementedError # unreachable
async def list_tools(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> types.ListToolsResult:
return types.ListToolsResult(tools=[types.Tool(name="impatient", input_schema={"type": "object"})])
async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult:
assert params.name == "impatient"
request = types.CreateMessageRequest(
params=types.CreateMessageRequestParams(
messages=[types.SamplingMessage(role="user", content=TextContent(text="Say hello."))],
max_tokens=8,
)
)
async with anyio.create_task_group() as abandon_scope:
async def sample() -> None:
await ctx.session.send_request(request, types.CreateMessageResult)
raise NotImplementedError # unreachable: the scope is cancelled
abandon_scope.start_soon(sample)
with anyio.fail_after(5):
await callback_started.wait()
abandon_scope.cancel_scope.cancel()
with anyio.fail_after(5):
await callback_cancelled.wait()
return CallToolResult(content=[TextContent(text="abandoned")])
server = Server("abandoner", on_list_tools=list_tools, on_call_tool=call_tool)
async with connect(server, sampling_callback=sampling_callback) as client:
result = await client.call_tool("impatient", {})
assert result == snapshot(CallToolResult(content=[TextContent(text="abandoned")]))
assert callback_cancelled.is_set()
@requirement("protocol:cancel:late-response-ignored")
async def test_a_response_for_an_unknown_request_id_surfaces_to_the_message_handler() -> None:
"""A response whose id matches no in-flight request is surfaced to the message handler as a RuntimeError.
async def test_a_response_for_an_unknown_request_id_is_ignored() -> None:
"""A response whose id matches no in-flight request is ignored, as the spec asks.
The spec says a sender SHOULD ignore a response that arrives after it issued a cancellation;
that is the same client-side code path as any response with an unknown id, and that form is
deterministic to test without depending on the cancellation API the SDK does not yet provide.
See the divergence note on the requirement.
deterministic to test without a client-side cancellation API.
"Ignored" is proved in two halves: the pong round-trip proves the read loop survived the
fabricated response (the ordered in-memory stream routed it first), and `surfaced` holding
only the control notification proves the fabricated response was never delivered to
`message_handler` (v1 surfaced it there as a RuntimeError).
A real Server cannot be made to answer with a fabricated id, so the test plays the server's
side of the wire by hand. Reserve this pattern for behaviour no real server can produce. The
@ -208,14 +266,18 @@ async def test_a_response_for_an_unknown_request_id_surfaces_to_the_message_hand
assert isinstance(ping, SessionMessage)
assert isinstance(ping.message, JSONRPCRequest)
assert ping.message.method == "ping"
# First answer with a fabricated id that matches nothing in flight, then the real id.
# First a fabricated id that matches nothing in flight, then a control notification that
# is surfaced to message_handler (proving the handler is live), then the real id.
await server_write.send(respond(9999, EmptyResult()))
await server_write.send(
SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/tools/list_changed"))
)
await server_write.send(respond(ping.message.id, EmptyResult()))
incoming: list[IncomingMessage] = []
surfaced: list[IncomingMessage] = []
async def message_handler(message: IncomingMessage) -> None:
incoming.append(message)
surfaced.append(message)
async with (
create_client_server_memory_streams() as ((client_read, client_write), server_streams),
@ -228,7 +290,56 @@ async def test_a_response_for_an_unknown_request_id_surfaces_to_the_message_hand
pong = await session.send_request(PingRequest(), EmptyResult)
assert pong == snapshot(EmptyResult())
assert len(incoming) == 1
assert isinstance(incoming[0], RuntimeError)
# The full message embeds the response object's repr; only the prefix is stable.
assert str(incoming[0]).startswith("Received response with an unknown request ID:")
# The stream is ordered, so the fabricated response was routed before the control
# notification: only the control surfaced, so the unknown-id response was dropped.
assert surfaced == snapshot([types.ToolListChangedNotification()])
@requirement("protocol:cancel:initialize-not-cancellable")
async def test_timed_out_initialize_sends_no_cancellation() -> None:
"""An abandoned initialize is not followed by notifications/cancelled on the wire (spec-mandated).
A real Server always answers initialize, so the test plays a stalling server by hand.
"""
received_methods: list[str] = []
async def scripted_server(streams: MessageStream) -> None:
server_read, server_write = streams
# Hold the initialize request unanswered until the client's read timeout fires.
init = await server_read.receive()
assert isinstance(init, SessionMessage)
assert isinstance(init.message, JSONRPCRequest)
received_methods.append(init.message.method)
follow_up = await server_read.receive()
assert isinstance(follow_up, SessionMessage)
assert isinstance(follow_up.message, JSONRPCRequest)
received_methods.append(follow_up.message.method)
await server_write.send(
SessionMessage(
JSONRPCResponse(
jsonrpc="2.0",
id=follow_up.message.id,
result=EmptyResult().model_dump(by_alias=True, mode="json", exclude_none=True),
)
)
)
async with (
create_client_server_memory_streams() as ((client_read, client_write), server_streams),
anyio.create_task_group() as task_group,
# The session-level read timeout is the only public pathway that abandons initialize.
ClientSession(client_read, client_write, read_timeout_seconds=0.000001) as session,
):
task_group.start_soon(scripted_server, server_streams)
with anyio.fail_after(5):
with pytest.raises(MCPError) as exc_info:
await session.initialize()
assert exc_info.value.error.code == REQUEST_TIMEOUT
# Override the session-level timeout: this ping must round-trip normally.
pong = await session.send_request(PingRequest(), EmptyResult, request_read_timeout_seconds=5)
assert pong == snapshot(EmptyResult())
# The stream is ordered, so a courtesy cancel would have arrived ahead of the ping.
assert received_methods == snapshot(["initialize", "ping"])

View file

@ -1,12 +1,8 @@
"""Logging interactions against the low-level Server, driven through the public Client API.
Notification ordering: the in-memory transport delivers every server-to-client message on one
ordered stream, and the client's receive loop dispatches each incoming message to completion
before reading the next one. Over streamable HTTP that ordered single-stream guarantee holds
only for messages that carry a ``related_request_id`` (they ride the originating request's POST
stream); without it the message routes to the standalone GET stream and may arrive after the
response. These tests pass ``related_request_id`` so they can collect into a plain list and
assert after the request completes on every transport leg -- no events, no waiting.
Notification ordering: await-free callbacks finish in arrival order, and passing
``related_request_id`` keeps each notification on the originating request's POST stream over
streamable HTTP, so plain-list collection is deterministic on every transport leg.
"""
import pytest

View file

@ -87,8 +87,8 @@ async def test_progress_token_visible_to_handler(connect: Connect) -> None:
async with connect(server) as client:
result = await client.call_tool("inspect", {}, progress_callback=ignore)
# The token is the request id of the tools/call request itself (initialize is request 0).
assert result == snapshot(CallToolResult(content=[TextContent(text="1")]))
# The token is the request id of the tools/call request itself (initialize is request 1).
assert result == snapshot(CallToolResult(content=[TextContent(text="2")]))
@requirement("protocol:progress:no-token")

View file

@ -3,8 +3,9 @@
The handler blocks on an event that is never set, so the awaited response can never arrive and
any positive timeout fires deterministically on the next event-loop pass. Per-request timeouts are
set to an effectively-zero duration; the session-level test runs on trio's virtual clock instead
(see the comment there). Either way the tests add no wall-clock time to the suite. (Zero itself
cannot be used: a falsy read_timeout_seconds is silently treated as "no timeout".)
(see the comment there). Either way the tests add no wall-clock time to the suite. (Zero would
also time out immediately, but a tiny positive value keeps the duration visible in the
cancellation reason these tests snapshot.)
"""
import anyio
@ -13,9 +14,13 @@ from inline_snapshot import snapshot
from trio.testing import MockClock
from mcp import MCPError, types
from mcp.client import ClientRequestContext
from mcp.client._memory import InMemoryTransport
from mcp.client.client import Client
from mcp.server import Server, ServerRequestContext
from mcp.types import REQUEST_TIMEOUT, CallToolResult, ErrorData, TextContent
from mcp.shared.message import SessionMessage
from mcp.types import REQUEST_TIMEOUT, CallToolResult, ErrorData, JSONRPCNotification, TextContent
from tests.interaction._helpers import RecordingTransport
from tests.interaction._requirements import requirement
pytestmark = pytest.mark.anyio
@ -26,16 +31,19 @@ pytestmark = pytest.mark.anyio
async def test_request_timeout_fails_the_pending_call() -> None:
"""A request whose response does not arrive within its read timeout fails with a timeout error.
No cancellation is sent to the server (see the divergence note on the requirement): the handler
starts and is still running after the caller has already given up. The test waits for the
handler to have started only after the timeout has fired, so the timeout itself races nothing.
The timeout is followed by notifications/cancelled, which interrupts the server's handler.
"""
handler_started = anyio.Event()
handler_cancelled = anyio.Event()
async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult:
assert params.name == "block"
handler_started.set()
await anyio.Event().wait() # blocks until the session is torn down
try:
await anyio.Event().wait() # blocks until the courtesy cancellation interrupts it
except anyio.get_cancelled_exc_class():
handler_cancelled.set()
raise
raise NotImplementedError # unreachable
server = Server("blocker", on_call_tool=call_tool)
@ -44,18 +52,79 @@ async def test_request_timeout_fails_the_pending_call() -> None:
with pytest.raises(MCPError) as exc_info:
await client.call_tool("block", {}, read_timeout_seconds=0.000001)
# The request was already on the wire: the handler still runs even though the caller gave up.
# The request was already on the wire: the handler started and was then cancelled.
with anyio.fail_after(5):
await handler_started.wait()
await handler_cancelled.wait()
assert exc_info.value.error == snapshot(
ErrorData(
code=REQUEST_TIMEOUT,
message="Timed out while waiting for response to CallToolRequest. Waited 1e-06 seconds.",
message="Request 'tools/call' timed out",
)
)
@requirement("protocol:timeout:basic")
@requirement("protocol:timeout:sends-cancellation")
async def test_server_request_timeout_sends_cancellation_to_the_client() -> None:
"""A server-initiated request that times out fails server-side and cancels the client's work.
The sampling callback answers only after the server gave up; the late response is discarded.
"""
release = anyio.Event()
callback_started = anyio.Event()
errors: list[ErrorData] = []
async def list_tools(
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> types.ListToolsResult:
return types.ListToolsResult(tools=[types.Tool(name="impatient", input_schema={"type": "object"})])
async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult:
assert params.name == "impatient"
request = types.CreateMessageRequest(
params=types.CreateMessageRequestParams(
messages=[types.SamplingMessage(role="user", content=TextContent(text="Say hello."))],
max_tokens=8,
)
)
with pytest.raises(MCPError) as exc_info:
await ctx.session.send_request(request, types.CreateMessageResult, request_read_timeout_seconds=0.000001)
errors.append(exc_info.value.error)
release.set()
return CallToolResult(content=[TextContent(text="gave up")])
server = Server("impatient", on_list_tools=list_tools, on_call_tool=call_tool)
recording = RecordingTransport(InMemoryTransport(server))
async def sampling_callback(
context: ClientRequestContext, params: types.CreateMessageRequestParams
) -> types.CreateMessageResult:
callback_started.set()
with anyio.fail_after(5):
await release.wait()
return types.CreateMessageResult(role="assistant", content=TextContent(text="too late"), model="test-model")
async with Client(recording, sampling_callback=sampling_callback) as client:
result = await client.call_tool("impatient", {})
assert result == snapshot(CallToolResult(content=[TextContent(text="gave up")]))
assert callback_started.is_set()
assert errors == snapshot([ErrorData(code=REQUEST_TIMEOUT, message="Request 'sampling/createMessage' timed out")])
cancellations = [
item.message
for item in recording.received
if isinstance(item, SessionMessage)
and isinstance(item.message, JSONRPCNotification)
and item.message.method == "notifications/cancelled"
]
# requestId 1 is the sampling request, the server's first outbound request.
assert [notification.params for notification in cancellations] == snapshot(
[{"requestId": 1, "reason": "timed out after 1e-06s"}]
)
@requirement("protocol:timeout:session-survives")
async def test_session_serves_requests_after_timeout() -> None:
"""A timed-out request does not poison the session: the next request succeeds."""
@ -73,7 +142,7 @@ async def test_session_serves_requests_after_timeout() -> None:
async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult:
if params.name == "echo":
return CallToolResult(content=[TextContent(text="still alive")])
await anyio.Event().wait() # blocks until the session is torn down
await anyio.Event().wait() # blocks until the courtesy cancellation interrupts it
raise NotImplementedError # unreachable
server = Server("blocker", on_list_tools=list_tools, on_call_tool=call_tool)
@ -105,7 +174,7 @@ async def test_session_level_timeout_applies_to_every_request() -> None:
async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult:
assert params.name == "block"
await anyio.Event().wait() # blocks until the session is torn down
await anyio.Event().wait() # blocks until the courtesy cancellation interrupts it
raise NotImplementedError # unreachable
server = Server("blocker", on_call_tool=call_tool)
@ -117,6 +186,6 @@ async def test_session_level_timeout_applies_to_every_request() -> None:
assert exc_info.value.error == snapshot(
ErrorData(
code=REQUEST_TIMEOUT,
message="Timed out while waiting for response to CallToolRequest. Waited 0.05 seconds.",
message="Request 'tools/call' timed out",
)
)

View file

@ -61,7 +61,7 @@ def _echo_server() -> Server:
async def test_request_ids_are_unique_and_never_null() -> None:
"""Every request the client sends carries a distinct, non-null id.
The id sequence is pinned: sequential integers from zero, in send order.
The id sequence is pinned: sequential integers from one, in send order.
"""
recording = RecordingTransport(InMemoryTransport(_echo_server()))
@ -77,7 +77,7 @@ async def test_request_ids_are_unique_and_never_null() -> None:
assert len(request_ids) == len(set(request_ids))
# initialize, tools/list, tools/call, tools/call, ping -- the client does not issue a
# schema-cache refresh here because the explicit tools/list already populated the cache.
assert request_ids == snapshot([0, 1, 2, 3, 4])
assert request_ids == snapshot([1, 2, 3, 4, 5])
@requirement("protocol:notifications:no-response")