1369 lines
62 KiB
Python
1369 lines
62 KiB
Python
"""JSON-RPC-specific Dispatcher tests.
|
|
|
|
Behaviors with no `DirectDispatcher` analog: request-id correlation, the
|
|
exception-to-wire boundary, peer-cancel handling, and shutdown fan-out.
|
|
The contract tests shared with `DirectDispatcher` live in
|
|
`test_dispatcher.py`.
|
|
"""
|
|
|
|
import contextvars
|
|
import json
|
|
import logging
|
|
from collections.abc import Mapping
|
|
from typing import Any
|
|
|
|
import anyio
|
|
import anyio.lowlevel
|
|
import pytest
|
|
|
|
from mcp import Client
|
|
from mcp.server import Server, ServerRequestContext
|
|
from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream
|
|
from mcp.shared.dispatcher import CallOptions, DispatchContext
|
|
from mcp.shared.exceptions import MCPError, NoBackChannelError
|
|
from mcp.shared.jsonrpc_dispatcher import ( # pyright: ignore[reportPrivateUsage]
|
|
JSONRPCDispatcher,
|
|
_coerce_id,
|
|
_outbound_metadata,
|
|
_Pending,
|
|
)
|
|
from mcp.shared.message import ClientMessageMetadata, MessageMetadata, ServerMessageMetadata, SessionMessage
|
|
from mcp.shared.transport_context import TransportContext
|
|
from mcp.types import (
|
|
CONNECTION_CLOSED,
|
|
INTERNAL_ERROR,
|
|
INVALID_PARAMS,
|
|
CallToolRequest,
|
|
CallToolRequestParams,
|
|
CallToolResult,
|
|
CancelledNotification,
|
|
CancelledNotificationParams,
|
|
ErrorData,
|
|
JSONRPCError,
|
|
JSONRPCNotification,
|
|
JSONRPCRequest,
|
|
JSONRPCResponse,
|
|
RequestId,
|
|
Tool,
|
|
)
|
|
|
|
from .conftest import jsonrpc_pair
|
|
from .test_dispatcher import Recorder, echo_handlers, running_pair
|
|
|
|
DCtx = DispatchContext[TransportContext]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_concurrent_send_raw_requests_correlate_by_id_when_responses_arrive_out_of_order():
|
|
release_first = anyio.Event()
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
if method == "first":
|
|
await release_first.wait()
|
|
return {"m": method}
|
|
|
|
async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, *_):
|
|
results: dict[str, dict[str, Any]] = {}
|
|
|
|
async def call(method: str) -> None:
|
|
results[method] = await client.send_raw_request(method, None)
|
|
|
|
with anyio.fail_after(5):
|
|
async with anyio.create_task_group() as tg: # pragma: no branch
|
|
tg.start_soon(call, "first")
|
|
await anyio.sleep(0)
|
|
tg.start_soon(call, "second")
|
|
await anyio.sleep(0)
|
|
# second resolves while first is still parked
|
|
assert "first" not in results
|
|
release_first.set()
|
|
assert results == {"first": {"m": "first"}, "second": {"m": "second"}}
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_handler_raising_exception_sends_code_zero_with_str_message():
|
|
"""Matches the existing server's `_handle_request`: code=0, message=str(e)."""
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
raise RuntimeError("kaboom")
|
|
|
|
async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, *_):
|
|
with anyio.fail_after(5), pytest.raises(MCPError) as exc:
|
|
await client.send_raw_request("tools/list", None)
|
|
assert exc.value.error.code == 0
|
|
assert exc.value.error.message == "kaboom"
|
|
assert exc.value.__cause__ is None # cause does not survive the wire
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_peer_cancel_interrupt_mode_writes_cancelled_error_response():
|
|
"""Matches the existing server: a peer-cancelled request is answered with code=0."""
|
|
handler_started = anyio.Event()
|
|
handler_exited = anyio.Event()
|
|
seen_ctx: list[DCtx] = []
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
seen_ctx.append(ctx)
|
|
handler_started.set()
|
|
try:
|
|
await anyio.sleep_forever()
|
|
finally:
|
|
handler_exited.set()
|
|
raise NotImplementedError
|
|
|
|
seen_error: list[ErrorData] = []
|
|
async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, *_):
|
|
with anyio.fail_after(5):
|
|
async with anyio.create_task_group() as tg: # pragma: no branch
|
|
|
|
async def call_then_record() -> None:
|
|
with pytest.raises(MCPError) as exc:
|
|
await client.send_raw_request("slow", None)
|
|
seen_error.append(exc.value.error)
|
|
|
|
tg.start_soon(call_then_record)
|
|
await handler_started.wait()
|
|
await client.notify("notifications/cancelled", {"requestId": 1})
|
|
await handler_exited.wait()
|
|
assert seen_ctx[0].cancel_requested.is_set()
|
|
assert seen_error == [ErrorData(code=0, message="Request cancelled")]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_peer_cancel_signal_mode_sets_event_but_handler_runs_to_completion():
|
|
handler_started = anyio.Event()
|
|
cancel_seen = anyio.Event()
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
handler_started.set()
|
|
await ctx.cancel_requested.wait()
|
|
cancel_seen.set()
|
|
return {"finished": True}
|
|
|
|
def factory(*, can_send_request: bool = True):
|
|
client, server, close = jsonrpc_pair(can_send_request=can_send_request)
|
|
# Reach in to set signal mode on the server side.
|
|
assert isinstance(server, JSONRPCDispatcher)
|
|
server._peer_cancel_mode = "signal" # pyright: ignore[reportPrivateUsage]
|
|
return client, server, close
|
|
|
|
result_box: list[dict[str, Any]] = []
|
|
async with running_pair(factory, server_on_request=server_on_request) as (client, *_):
|
|
with anyio.fail_after(5):
|
|
async with anyio.create_task_group() as tg: # pragma: no branch
|
|
|
|
async def call() -> None:
|
|
result_box.append(await client.send_raw_request("slow", None))
|
|
|
|
tg.start_soon(call)
|
|
await handler_started.wait()
|
|
await client.notify("notifications/cancelled", {"requestId": 1})
|
|
await cancel_seen.wait()
|
|
assert result_box == [{"finished": True}]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_send_raw_request_raises_connection_closed_when_read_stream_eofs_mid_await():
|
|
"""A blocked send_raw_request is woken with CONNECTION_CLOSED when run() exits."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
|
|
on_request, on_notify = echo_handlers(Recorder())
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(client.run, on_request, on_notify)
|
|
|
|
async def caller() -> None:
|
|
with pytest.raises(MCPError) as exc:
|
|
await client.send_raw_request("ping", None)
|
|
assert exc.value.error.code == CONNECTION_CLOSED
|
|
|
|
tg.start_soon(caller)
|
|
await anyio.sleep(0)
|
|
# No server: simulate the peer dropping by closing the read side.
|
|
s2c_send.close()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_run_returns_cleanly_when_read_stream_receive_end_is_closed():
|
|
"""Iterating a closed receive end raises ClosedResourceError; run() treats it as EOF.
|
|
|
|
Stateless SHTTP teardown closes the dispatcher's receive end after the
|
|
request is handled; the next loop iteration must not surface as a crash.
|
|
"""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
|
|
on_request, on_notify = echo_handlers(Recorder())
|
|
# Close the dispatcher's own receive end (not the send end) before run()
|
|
# iterates it: __anext__ on a closed stream raises ClosedResourceError.
|
|
c2s_recv.close()
|
|
with anyio.fail_after(5):
|
|
await server.run(on_request, on_notify)
|
|
for s in (c2s_send, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_run_cancels_in_flight_handlers_when_read_stream_eofs():
|
|
"""A handler that outlives its caller must not keep run() from returning.
|
|
|
|
Without the cancel-at-EOF, the task-group join would wait on this handler
|
|
forever (over SSE that leaks the handler task and the GET request hosting
|
|
the session).
|
|
"""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
|
|
handler_started = anyio.Event()
|
|
handler_cancelled = anyio.Event()
|
|
|
|
async def park(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
handler_started.set()
|
|
try:
|
|
await anyio.sleep_forever()
|
|
finally:
|
|
handler_cancelled.set()
|
|
raise NotImplementedError
|
|
|
|
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
|
|
raise NotImplementedError
|
|
|
|
run_returned = anyio.Event()
|
|
|
|
async def drive() -> None:
|
|
await server.run(park, on_notify)
|
|
run_returned.set()
|
|
|
|
async with anyio.create_task_group() as tg:
|
|
tg.start_soon(drive)
|
|
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="x", params=None)))
|
|
with anyio.fail_after(5):
|
|
await handler_started.wait()
|
|
c2s_send.close() # EOF the read side; run() must cancel the parked handler
|
|
await run_returned.wait()
|
|
assert handler_cancelled.is_set()
|
|
s2c_recv.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_run_closes_write_stream_on_exit():
|
|
"""run() enters both streams; the write end is released on EOF."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
|
|
on_request, on_notify = echo_handlers(Recorder())
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, on_request, on_notify)
|
|
c2s_send.close() # EOF the read side; run() exits
|
|
with anyio.fail_after(5), pytest.raises(anyio.EndOfStream): # pragma: no branch
|
|
# Write end was entered and released by run(); peer's receive sees EOF.
|
|
await s2c_recv.receive()
|
|
s2c_recv.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_late_response_after_timeout_is_dropped_without_crashing():
|
|
handler_started = anyio.Event()
|
|
proceed = anyio.Event()
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
handler_started.set()
|
|
await proceed.wait()
|
|
return {"late": True}
|
|
|
|
async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, *_):
|
|
with anyio.fail_after(5):
|
|
with pytest.raises(MCPError): # REQUEST_TIMEOUT
|
|
await client.send_raw_request("slow", None, {"timeout": 0})
|
|
# The server handler is still running; let it finish and write a
|
|
# response for an id the client has already discarded.
|
|
await handler_started.wait()
|
|
proceed.set()
|
|
# One more round-trip proves the dispatcher is still healthy.
|
|
assert await client.send_raw_request("ping", None) == {"late": True}
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_raise_handler_exceptions_true_propagates_out_of_run():
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
|
|
def builder(_meta: object) -> TransportContext:
|
|
return TransportContext(kind="jsonrpc", can_send_request=True)
|
|
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(
|
|
c2s_recv, s2c_send, transport_builder=builder, raise_handler_exceptions=True
|
|
)
|
|
|
|
async def boom(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
raise RuntimeError("propagate me")
|
|
|
|
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
|
|
raise NotImplementedError
|
|
|
|
try:
|
|
with pytest.raises(BaseException) as exc:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, boom, on_notify)
|
|
# Inject a request directly onto the server's read stream.
|
|
await c2s_send.send(
|
|
SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="x", params=None))
|
|
)
|
|
assert exc.group_contains(RuntimeError, match="propagate me")
|
|
# The error response was still written before re-raising.
|
|
sent = s2c_recv.receive_nowait()
|
|
assert isinstance(sent, SessionMessage)
|
|
assert isinstance(sent.message, JSONRPCError)
|
|
assert sent.message.error.code == 0
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ctx_send_raw_request_tags_outbound_with_server_message_metadata():
|
|
"""Server-to-client requests carry related_request_id for SHTTP routing."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
return await ctx.send_raw_request("sampling/createMessage", {"prompt": "hi"})
|
|
|
|
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
|
|
raise NotImplementedError
|
|
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, server_on_request, on_notify)
|
|
# Kick the server with an inbound request id=7.
|
|
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=7, method="t", params=None)))
|
|
with anyio.fail_after(5):
|
|
outbound = await s2c_recv.receive()
|
|
assert isinstance(outbound, SessionMessage)
|
|
assert isinstance(outbound.message, JSONRPCRequest)
|
|
assert isinstance(outbound.metadata, ServerMessageMetadata)
|
|
assert outbound.metadata.related_request_id == 7
|
|
# Reply so the handler completes cleanly.
|
|
await c2s_send.send(
|
|
SessionMessage(message=JSONRPCResponse(jsonrpc="2.0", id=outbound.message.id, result={"ok": True}))
|
|
)
|
|
with anyio.fail_after(5):
|
|
final = await s2c_recv.receive()
|
|
assert isinstance(final, SessionMessage)
|
|
assert isinstance(final.message, JSONRPCResponse)
|
|
assert final.message.id == 7
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_courtesy_cancel_on_timeout_tags_outbound_with_server_message_metadata():
|
|
"""The timeout-path `notifications/cancelled` carries the originating request id.
|
|
|
|
Streamable-HTTP's `message_router` keys on `ServerMessageMetadata.related_request_id`;
|
|
a cancel without it would fall through to the standalone GET stream and be dropped
|
|
when no GET stream is open, so the client never learns to stop work.
|
|
"""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
with pytest.raises(MCPError): # REQUEST_TIMEOUT
|
|
await ctx.send_raw_request("sampling/createMessage", None, {"timeout": 0})
|
|
return {"gave_up": True}
|
|
|
|
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
|
|
raise NotImplementedError
|
|
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, server_on_request, on_notify)
|
|
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=7, method="t", params=None)))
|
|
with anyio.fail_after(5):
|
|
outbound = await s2c_recv.receive()
|
|
assert isinstance(outbound, SessionMessage)
|
|
assert isinstance(outbound.message, JSONRPCRequest)
|
|
assert outbound.message.method == "sampling/createMessage"
|
|
sampling_id = outbound.message.id
|
|
# Don't respond; let the timeout fire. Next on the wire is the courtesy cancel.
|
|
with anyio.fail_after(5):
|
|
cancel = await s2c_recv.receive()
|
|
assert isinstance(cancel, SessionMessage)
|
|
assert isinstance(cancel.message, JSONRPCNotification)
|
|
assert cancel.message.method == "notifications/cancelled"
|
|
assert cancel.message.params == {"requestId": sampling_id, "reason": "timed out after 0s"}
|
|
assert isinstance(cancel.metadata, ServerMessageMetadata)
|
|
assert cancel.metadata.related_request_id == 7
|
|
with anyio.fail_after(5):
|
|
final = await s2c_recv.receive()
|
|
assert isinstance(final, SessionMessage)
|
|
assert isinstance(final.message, JSONRPCResponse)
|
|
assert final.message.result == {"gave_up": True}
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ctx_message_metadata_carries_inbound_request_metadata():
|
|
"""Transport-attached metadata (HTTP request, SSE close hooks) is readable off the dispatch context."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
|
|
metadata = ServerMessageMetadata(request_context="request-scoped-data")
|
|
seen: list[MessageMetadata] = []
|
|
|
|
async def on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
seen.append(ctx.message_metadata)
|
|
return {}
|
|
|
|
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
|
|
raise NotImplementedError
|
|
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, on_request, on_notify)
|
|
await c2s_send.send(
|
|
SessionMessage(
|
|
message=JSONRPCRequest(jsonrpc="2.0", id=1, method="tools/call", params=None),
|
|
metadata=metadata,
|
|
)
|
|
)
|
|
with anyio.fail_after(5):
|
|
await s2c_recv.receive() # response sent => the handler has run
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
assert len(seen) == 1
|
|
assert seen[0] is metadata # the exact object, passed through verbatim
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ctx_message_metadata_carries_inbound_notification_metadata():
|
|
"""Notifications get the same metadata pass-through as requests."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
|
|
metadata = ServerMessageMetadata(request_context="request-scoped-data")
|
|
seen: list[MessageMetadata] = []
|
|
notified = anyio.Event()
|
|
|
|
async def on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
raise NotImplementedError
|
|
|
|
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
|
|
seen.append(ctx.message_metadata)
|
|
notified.set()
|
|
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, on_request, on_notify)
|
|
await c2s_send.send(
|
|
SessionMessage(
|
|
message=JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized", params=None),
|
|
metadata=metadata,
|
|
)
|
|
)
|
|
with anyio.fail_after(5):
|
|
await notified.wait()
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
assert len(seen) == 1
|
|
assert seen[0] is metadata
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ctx_progress_with_only_progress_value_omits_total_and_message():
|
|
received: list[tuple[float, float | None, str | None]] = []
|
|
|
|
async def on_progress(progress: float, total: float | None, message: str | None) -> None:
|
|
received.append((progress, total, message))
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
await ctx.progress(0.25)
|
|
return {}
|
|
|
|
async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, *_):
|
|
with anyio.fail_after(5):
|
|
await client.send_raw_request("t", None, {"on_progress": on_progress})
|
|
assert received == [(0.25, None, None)]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ctx_after_handler_return_reports_closed_and_drops_backchannel_traffic():
|
|
"""Once `_handle_request` closes the dctx, the back-channel guard and ops agree.
|
|
|
|
Detached work that outlives the handler must see `can_send_request == False`,
|
|
get `NoBackChannelError` from `send_raw_request`, and have `notify`/`progress`
|
|
silently dropped rather than emitted with a stale `related_request_id`.
|
|
"""
|
|
captured: list[DCtx] = []
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
captured.append(ctx)
|
|
assert ctx.can_send_request is True
|
|
return {}
|
|
|
|
async def on_progress(progress: float, total: float | None, message: str | None) -> None:
|
|
raise NotImplementedError
|
|
|
|
async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, _server, crec, _srec):
|
|
with anyio.fail_after(5):
|
|
await client.send_raw_request("tools/call", None, {"on_progress": on_progress})
|
|
dctx = captured[0]
|
|
assert dctx.can_send_request is False
|
|
with pytest.raises(NoBackChannelError):
|
|
await dctx.send_raw_request("sampling/createMessage", None)
|
|
await dctx.notify("notifications/message", {"level": "info"})
|
|
await dctx.progress(0.9)
|
|
# A second round-trip flushes any notification the server might have
|
|
# written, so an empty client recorder afterwards proves the drop.
|
|
await client.send_raw_request("ping", None)
|
|
assert crec.notifications == []
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_progress_callback_exception_is_swallowed_and_logged(caplog: pytest.LogCaptureFixture):
|
|
"""A user progress callback raising must not crash the dispatcher."""
|
|
|
|
async def boom(progress: float, total: float | None, message: str | None) -> None:
|
|
raise RuntimeError("progress callback boom")
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
await ctx.progress(0.5)
|
|
return {"ok": True}
|
|
|
|
opts: CallOptions = {"on_progress": boom}
|
|
async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, *_):
|
|
with anyio.fail_after(5):
|
|
result = await client.send_raw_request("t", None, opts)
|
|
# Request still completes; the callback's crash was swallowed.
|
|
assert result == {"ok": True}
|
|
assert "progress callback raised" in caplog.text
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_inline_methods_are_handled_before_next_message_is_dequeued():
|
|
"""A method in `inline_methods` runs to completion before subsequent
|
|
messages are dispatched, so its side effects are visible to them."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(
|
|
c2s_recv, s2c_send, inline_methods=frozenset({"first"})
|
|
)
|
|
state = {"initialized": False}
|
|
seen: list[bool] = []
|
|
|
|
async def on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
if method == "first":
|
|
await anyio.lowlevel.checkpoint()
|
|
state["initialized"] = True
|
|
else:
|
|
seen.append(state["initialized"])
|
|
return {}
|
|
|
|
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
|
|
raise NotImplementedError
|
|
|
|
# Buffer both requests before run() reads anything.
|
|
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="first", params=None)))
|
|
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=2, method="second", params=None)))
|
|
c2s_send.close()
|
|
with anyio.fail_after(5):
|
|
await server.run(on_request, on_notify)
|
|
assert seen == [True]
|
|
s2c_recv.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_send_raw_request_always_carries_meta_on_the_wire():
|
|
"""Outbound requests always include `params._meta` (otel injection per SEP-414).
|
|
|
|
Caller-supplied `_meta` keys are preserved; the progress token is merged in.
|
|
"""
|
|
seen: list[Mapping[str, Any] | None] = []
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
seen.append(params)
|
|
return {}
|
|
|
|
async def noop_progress(progress: float, total: float | None, message: str | None) -> None:
|
|
raise NotImplementedError
|
|
|
|
opts: CallOptions = {"on_progress": noop_progress}
|
|
async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, *_):
|
|
with anyio.fail_after(5):
|
|
await client.send_raw_request("a", None)
|
|
await client.send_raw_request("b", {"x": 1, "_meta": {"k": "v"}}, opts)
|
|
# `_meta` is always present. Its contents depend on the active otel
|
|
# tracer (traceparent/tracestate may be injected), so assert presence
|
|
# and that anything beyond W3C keys is exactly what we expect.
|
|
w3c = {"traceparent", "tracestate"}
|
|
assert seen[0] is not None and seen[0].keys() == {"_meta"}
|
|
assert set(seen[0]["_meta"].keys()) <= w3c
|
|
assert seen[1] is not None and seen[1]["x"] == 1
|
|
assert set(seen[1]["_meta"].keys()) - w3c == {"k", "progressToken"}
|
|
assert seen[1]["_meta"]["k"] == "v"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_handler_raising_validation_error_sends_invalid_params():
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
Tool.model_validate({"name": 123}) # raises ValidationError
|
|
raise NotImplementedError
|
|
|
|
async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, *_):
|
|
with anyio.fail_after(5), pytest.raises(MCPError) as exc:
|
|
await client.send_raw_request("t", None)
|
|
assert exc.value.error == ErrorData(code=INVALID_PARAMS, message="Invalid request parameters", data="")
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_send_raw_request_before_run_raises_runtimeerror():
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
|
|
d: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
|
|
try:
|
|
with pytest.raises(RuntimeError, match="before run"):
|
|
await d.send_raw_request("ping", None)
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_transport_exception_in_read_stream_is_logged_and_dropped():
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
|
|
on_request, on_notify = echo_handlers(Recorder())
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, on_request, on_notify)
|
|
await c2s_send.send(ValueError("transport hiccup"))
|
|
# Dispatcher must remain healthy after the dropped exception.
|
|
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="t", params=None)))
|
|
with anyio.fail_after(5):
|
|
resp = await s2c_recv.receive()
|
|
assert isinstance(resp, SessionMessage)
|
|
assert isinstance(resp.message, JSONRPCResponse)
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_progress_notification_for_unknown_token_falls_through_to_on_notify():
|
|
async with running_pair(jsonrpc_pair) as (client, _server, _crec, srec):
|
|
with anyio.fail_after(5):
|
|
await client.notify("notifications/progress", {"progressToken": 999, "progress": 0.5})
|
|
await srec.notified.wait()
|
|
assert srec.notifications == [("notifications/progress", {"progressToken": 999, "progress": 0.5})]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_cancelled_notification_for_unknown_request_id_skips_cancel_but_reaches_on_notify():
|
|
async with running_pair(jsonrpc_pair) as (client, _server, _crec, srec):
|
|
with anyio.fail_after(5):
|
|
await client.notify("notifications/cancelled", {"requestId": 999})
|
|
await srec.notified.wait()
|
|
# No in-flight correlation; dispatcher remains healthy.
|
|
assert await client.send_raw_request("t", None) == {"echoed": "t", "params": {}}
|
|
# cancelled is teed to on_notify so middleware/handlers can observe it.
|
|
assert srec.notifications == [("notifications/cancelled", {"requestId": 999})]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_cancelled_notification_for_in_flight_request_is_teed_to_on_notify():
|
|
"""The dispatcher applies the cancellation itself AND forwards the notification."""
|
|
handler_started = anyio.Event()
|
|
handler_exited = anyio.Event()
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
handler_started.set()
|
|
try:
|
|
await anyio.sleep_forever()
|
|
finally:
|
|
handler_exited.set()
|
|
raise NotImplementedError
|
|
|
|
async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, _server, _crec, srec):
|
|
with anyio.fail_after(5):
|
|
async with anyio.create_task_group() as tg: # pragma: no branch
|
|
|
|
async def call() -> None:
|
|
with pytest.raises(MCPError):
|
|
await client.send_raw_request("slow", None)
|
|
|
|
tg.start_soon(call)
|
|
await handler_started.wait()
|
|
await client.notify("notifications/cancelled", {"requestId": 1})
|
|
await handler_exited.wait()
|
|
await srec.notified.wait()
|
|
assert srec.notifications == [("notifications/cancelled", {"requestId": 1})]
|
|
|
|
|
|
_probe: contextvars.ContextVar[str] = contextvars.ContextVar("probe", default="unset")
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize("inline", [frozenset[str](), frozenset({"t"})], ids=["spawned", "inline"])
|
|
async def test_handler_inherits_sender_contextvars(inline: frozenset[str]):
|
|
"""The handler task sees contextvars set by the task that wrote into the
|
|
read stream, on both the spawned and the inline-method dispatch paths."""
|
|
raw_send, raw_recv = anyio.create_memory_object_stream[tuple[contextvars.Context, SessionMessage | Exception]](4)
|
|
read_stream = ContextReceiveStream[SessionMessage | Exception](raw_recv)
|
|
write_send = ContextSendStream[SessionMessage | Exception](raw_send)
|
|
out_send, out_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(read_stream, out_send, inline_methods=inline)
|
|
|
|
seen: list[str] = []
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
seen.append(_probe.get())
|
|
return {}
|
|
|
|
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
|
|
raise NotImplementedError
|
|
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, server_on_request, on_notify)
|
|
|
|
async def sender() -> None:
|
|
_probe.set("from-sender")
|
|
await write_send.send(
|
|
SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="t", params=None))
|
|
)
|
|
|
|
tg.start_soon(sender)
|
|
with anyio.fail_after(5):
|
|
resp = await out_recv.receive()
|
|
assert isinstance(resp, SessionMessage)
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (raw_send, raw_recv, out_send, out_recv):
|
|
s.close()
|
|
assert seen == ["from-sender"]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_response_write_after_peer_drop_is_swallowed():
|
|
"""Handler completes after the write stream is closed; the dropped write doesn't crash run()."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
|
|
proceed = anyio.Event()
|
|
handlers_done = anyio.Event()
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
await proceed.wait()
|
|
if method == "raise":
|
|
handlers_done.set()
|
|
raise MCPError(code=INTERNAL_ERROR, message="x")
|
|
return {"ok": True}
|
|
|
|
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
|
|
raise NotImplementedError
|
|
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, server_on_request, on_notify)
|
|
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="ok", params=None)))
|
|
await c2s_send.send(
|
|
SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=2, method="raise", params=None))
|
|
)
|
|
await anyio.sleep(0)
|
|
# Peer drops: close the receive end so the server's writes hit BrokenResourceError.
|
|
s2c_recv.close()
|
|
proceed.set()
|
|
with anyio.fail_after(5):
|
|
await handlers_done.wait()
|
|
# run() must still be healthy - close the read side to let it exit cleanly.
|
|
c2s_send.close()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_cancel_outbound_after_write_stream_closed_is_swallowed():
|
|
"""Courtesy-cancel write hits a closed stream; the error is swallowed and cancellation propagates."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
|
|
client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
|
|
on_request, on_notify = echo_handlers(Recorder())
|
|
caller_done = anyio.Event()
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(client.run, on_request, on_notify)
|
|
caller_scope = anyio.CancelScope()
|
|
|
|
async def caller() -> None:
|
|
with caller_scope:
|
|
await client.send_raw_request("slow", None)
|
|
caller_done.set()
|
|
|
|
tg.start_soon(caller)
|
|
# Deterministic proof the request write completed: pull it off the wire.
|
|
with anyio.fail_after(5):
|
|
sent = await c2s_recv.receive()
|
|
assert isinstance(sent, SessionMessage)
|
|
assert isinstance(sent.message, JSONRPCRequest)
|
|
# Now safe: close the client's write end, then cancel the caller. The
|
|
# shielded `_cancel_outbound` write hits ClosedResourceError and is
|
|
# swallowed; cancellation propagates cleanly.
|
|
c2s_send.close()
|
|
caller_scope.cancel()
|
|
with anyio.fail_after(5):
|
|
await caller_done.wait()
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
def test_resolve_pending_drops_outcome_when_waiter_stream_already_closed():
|
|
"""White-box: a response for an id still in _pending but whose waiter has gone."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
|
|
d: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
|
|
send, recv = anyio.create_memory_object_stream[dict[str, Any] | ErrorData](1)
|
|
d._pending[1] = _Pending(send=send, receive=recv) # pyright: ignore[reportPrivateUsage]
|
|
recv.close() # waiter gone - send_nowait will raise BrokenResourceError
|
|
d._resolve_pending(1, {"late": True}) # pyright: ignore[reportPrivateUsage]
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv, send):
|
|
s.close()
|
|
|
|
|
|
def test_fan_out_closed_drops_signal_when_waiter_already_has_outcome():
|
|
"""White-box: the buffer=1 invariant - WouldBlock means waiter already has an outcome."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
|
|
d: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
|
|
send, recv = anyio.create_memory_object_stream[dict[str, Any] | ErrorData](1)
|
|
# Register a fake pending and pre-fill its single buffer slot.
|
|
d._pending[1] = _Pending(send=send, receive=recv) # pyright: ignore[reportPrivateUsage]
|
|
send.send_nowait({"real": "result"})
|
|
d._fan_out_closed() # pyright: ignore[reportPrivateUsage]
|
|
# The real result is still there; the close signal was dropped.
|
|
assert recv.receive_nowait() == {"real": "result"}
|
|
assert d._pending == {} # pyright: ignore[reportPrivateUsage]
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv, send, recv):
|
|
s.close()
|
|
|
|
|
|
def test_outbound_metadata_with_resumption_token_returns_client_metadata():
|
|
md = _outbound_metadata(None, {"resumption_token": "abc"})
|
|
assert isinstance(md, ClientMessageMetadata)
|
|
assert md.resumption_token == "abc"
|
|
assert _outbound_metadata(None, None) is None
|
|
assert _outbound_metadata(None, {}) is None
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_response_with_string_id_correlates_to_int_keyed_pending_request():
|
|
"""A peer that echoes the request ID as a JSON string still resolves the waiter."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
|
|
on_request, on_notify = echo_handlers(Recorder())
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(client.run, on_request, on_notify)
|
|
with anyio.fail_after(5):
|
|
|
|
async def respond_stringly() -> None:
|
|
out = await c2s_recv.receive()
|
|
assert isinstance(out, SessionMessage)
|
|
assert isinstance(out.message, JSONRPCRequest)
|
|
rid = out.message.id
|
|
assert isinstance(rid, int)
|
|
await s2c_send.send(
|
|
SessionMessage(message=JSONRPCResponse(jsonrpc="2.0", id=str(rid), result={"ok": True}))
|
|
)
|
|
|
|
tg.start_soon(respond_stringly)
|
|
result = await client.send_raw_request("ping", None)
|
|
assert result == {"ok": True}
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_progress_with_string_token_reaches_callback_for_int_keyed_request():
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
|
|
on_request, on_notify = echo_handlers(Recorder())
|
|
seen: list[float] = []
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(client.run, on_request, on_notify)
|
|
with anyio.fail_after(5):
|
|
|
|
async def respond_with_string_token_progress() -> None:
|
|
out = await c2s_recv.receive()
|
|
assert isinstance(out, SessionMessage)
|
|
assert isinstance(out.message, JSONRPCRequest)
|
|
rid = out.message.id
|
|
assert isinstance(rid, int)
|
|
await s2c_send.send(
|
|
SessionMessage(
|
|
message=JSONRPCNotification(
|
|
jsonrpc="2.0",
|
|
method="notifications/progress",
|
|
params={"progressToken": str(rid), "progress": 0.5},
|
|
)
|
|
)
|
|
)
|
|
await s2c_send.send(
|
|
SessionMessage(message=JSONRPCResponse(jsonrpc="2.0", id=rid, result={"ok": True}))
|
|
)
|
|
|
|
async def on_progress(progress: float, total: float | None, message: str | None) -> None:
|
|
seen.append(progress)
|
|
|
|
tg.start_soon(respond_with_string_token_progress)
|
|
result = await client.send_raw_request("ping", None, {"on_progress": on_progress})
|
|
assert result == {"ok": True}
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
assert seen == [0.5]
|
|
|
|
|
|
def test_coerce_id_passes_through_non_numeric_string_and_int():
|
|
assert _coerce_id("7") == 7
|
|
assert _coerce_id("not-an-int") == "not-an-int"
|
|
assert _coerce_id(42) == 42
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_jsonrpc_error_response_with_null_id_is_dropped():
|
|
"""Parse-error responses (id=null) have no waiter; they're logged and dropped."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
|
|
on_request, on_notify = echo_handlers(Recorder())
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(client.run, on_request, on_notify)
|
|
await s2c_send.send(
|
|
SessionMessage(message=JSONRPCError(jsonrpc="2.0", id=None, error=ErrorData(code=-32700, message="x")))
|
|
)
|
|
await anyio.sleep(0)
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_notify_without_params_omits_params_key_on_the_wire():
|
|
"""JSON-RPC 2.0 forbids `params: null`; the member must be absent.
|
|
|
|
Transports serialize with `exclude_unset=True`, so `notify` must leave
|
|
`params` unset on the model rather than passing an explicit None (strict
|
|
peers like the TypeScript SDK reject `"params": null`).
|
|
"""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
|
|
d: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
|
|
try:
|
|
await d.notify("notifications/tools/list_changed", None)
|
|
await d.notify("notifications/message", {"level": "info"})
|
|
bare = c2s_recv.receive_nowait()
|
|
with_params = c2s_recv.receive_nowait()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
assert isinstance(bare, SessionMessage)
|
|
wire = json.loads(bare.message.model_dump_json(by_alias=True, exclude_unset=True))
|
|
assert wire == {"jsonrpc": "2.0", "method": "notifications/tools/list_changed"}
|
|
assert isinstance(with_params, SessionMessage)
|
|
wire = json.loads(with_params.message.model_dump_json(by_alias=True, exclude_unset=True))
|
|
assert wire == {"jsonrpc": "2.0", "method": "notifications/message", "params": {"level": "info"}}
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_transport_builder_exception_on_request_is_answered_with_internal_error():
|
|
"""A raising builder costs only the one request, not the connection."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
calls = 0
|
|
|
|
def builder(_meta: MessageMetadata) -> TransportContext:
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls == 1:
|
|
raise RuntimeError("builder boom")
|
|
return TransportContext(kind="jsonrpc", can_send_request=True)
|
|
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send, transport_builder=builder)
|
|
on_request, on_notify = echo_handlers(Recorder())
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, on_request, on_notify)
|
|
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="t", params=None)))
|
|
with anyio.fail_after(5):
|
|
resp = await s2c_recv.receive()
|
|
assert isinstance(resp, SessionMessage)
|
|
assert isinstance(resp.message, JSONRPCError)
|
|
assert resp.message.id == 1
|
|
assert resp.message.error == ErrorData(code=INTERNAL_ERROR, message="transport context unavailable")
|
|
# The dispatcher stays healthy: the next request is served normally.
|
|
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=2, method="t", params=None)))
|
|
with anyio.fail_after(5):
|
|
resp2 = await s2c_recv.receive()
|
|
assert isinstance(resp2, SessionMessage)
|
|
assert isinstance(resp2.message, JSONRPCResponse)
|
|
assert resp2.message.id == 2
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_transport_builder_exception_on_notification_drops_only_that_notification():
|
|
"""A raising builder drops the one notification; the read loop survives."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
calls = 0
|
|
|
|
def builder(_meta: MessageMetadata) -> TransportContext:
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls == 1:
|
|
raise RuntimeError("builder boom")
|
|
return TransportContext(kind="jsonrpc", can_send_request=True)
|
|
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send, transport_builder=builder)
|
|
rec = Recorder()
|
|
on_request, on_notify = echo_handlers(rec)
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, on_request, on_notify)
|
|
await c2s_send.send(
|
|
SessionMessage(message=JSONRPCNotification(jsonrpc="2.0", method="notifications/x", params=None))
|
|
)
|
|
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=1, method="t", params=None)))
|
|
with anyio.fail_after(5):
|
|
resp = await s2c_recv.receive()
|
|
assert isinstance(resp, SessionMessage)
|
|
assert isinstance(resp.message, JSONRPCResponse)
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
assert rec.notifications == [] # the notification never reached on_notify
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_cancelled_with_bool_request_id_does_not_cancel_request_one():
|
|
"""`int()` match patterns accept bool, and `True == 1` would alias the
|
|
`_in_flight` lookup to request id 1; the bool guard must reject it."""
|
|
handler_started = anyio.Event()
|
|
handler_exited = anyio.Event()
|
|
|
|
async def server_on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
handler_started.set()
|
|
try:
|
|
await anyio.sleep_forever()
|
|
finally:
|
|
handler_exited.set()
|
|
raise NotImplementedError
|
|
|
|
async with running_pair(jsonrpc_pair, server_on_request=server_on_request) as (client, _server, _crec, srec):
|
|
with anyio.fail_after(5):
|
|
async with anyio.create_task_group() as tg: # pragma: no branch
|
|
|
|
async def call() -> None:
|
|
with pytest.raises(MCPError):
|
|
await client.send_raw_request("slow", None)
|
|
|
|
tg.start_soon(call)
|
|
await handler_started.wait()
|
|
await client.notify("notifications/cancelled", {"requestId": True})
|
|
# The malformed cancel is teed to on_notify; once observed, the
|
|
# correlation arm has already run - and must not have cancelled
|
|
# the request keyed by id 1.
|
|
await srec.notified.wait()
|
|
assert not handler_exited.is_set()
|
|
await client.notify("notifications/cancelled", {"requestId": 1})
|
|
await handler_exited.wait()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_progress_with_bool_token_or_bool_progress_does_not_fire_callback():
|
|
"""Bool `progressToken`/`progress` values are malformed; the callback must
|
|
not fire for the unrelated request keyed by id 1 (`True == 1`)."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
client: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(s2c_recv, c2s_send)
|
|
on_request, on_notify = echo_handlers(Recorder())
|
|
seen: list[float] = []
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(client.run, on_request, on_notify)
|
|
with anyio.fail_after(5):
|
|
|
|
async def respond_with_malformed_then_valid_progress() -> None:
|
|
out = await c2s_recv.receive()
|
|
assert isinstance(out, SessionMessage)
|
|
assert isinstance(out.message, JSONRPCRequest)
|
|
rid = out.message.id
|
|
for params in (
|
|
{"progressToken": True, "progress": 0.1}, # bool token
|
|
{"progressToken": rid, "progress": True}, # bool progress
|
|
{"progressToken": rid, "progress": 0.5}, # valid
|
|
):
|
|
await s2c_send.send(
|
|
SessionMessage(
|
|
message=JSONRPCNotification(
|
|
jsonrpc="2.0", method="notifications/progress", params=params
|
|
)
|
|
)
|
|
)
|
|
await s2c_send.send(
|
|
SessionMessage(message=JSONRPCResponse(jsonrpc="2.0", id=rid, result={"ok": True}))
|
|
)
|
|
|
|
async def on_progress(progress: float, total: float | None, message: str | None) -> None:
|
|
seen.append(progress)
|
|
|
|
tg.start_soon(respond_with_malformed_then_valid_progress)
|
|
result = await client.send_raw_request("ping", None, {"on_progress": on_progress})
|
|
assert result == {"ok": True}
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
assert seen == [0.5] # only the well-formed progress fired the callback
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_request_with_bool_meta_progress_token_is_not_adopted():
|
|
"""A bool `_meta.progressToken` is malformed; `ctx.progress()` must be a no-op
|
|
instead of emitting `progressToken: true` on the wire."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
|
|
|
|
async def on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
await ctx.progress(0.5)
|
|
return {"ok": True}
|
|
|
|
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
|
|
raise NotImplementedError
|
|
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, on_request, on_notify)
|
|
await c2s_send.send(
|
|
SessionMessage(
|
|
message=JSONRPCRequest(jsonrpc="2.0", id=1, method="t", params={"_meta": {"progressToken": True}})
|
|
)
|
|
)
|
|
with anyio.fail_after(5):
|
|
first = await s2c_recv.receive()
|
|
# No progress notification was emitted; the first wire message is the response.
|
|
assert isinstance(first, SessionMessage)
|
|
assert isinstance(first.message, JSONRPCResponse)
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize(
|
|
("request_id", "cancel_id"),
|
|
[(7, "7"), ("9", 9)],
|
|
ids=["string-cancel-for-int-request", "int-cancel-for-string-request"],
|
|
)
|
|
async def test_cancelled_correlates_across_string_and_int_request_id_forms(request_id: RequestId, cancel_id: object):
|
|
"""A peer that stringifies the id between request and cancel still cancels
|
|
(same `_coerce_id` treatment as the response path)."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
|
|
|
|
async def on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
await anyio.sleep_forever()
|
|
raise NotImplementedError
|
|
|
|
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
|
|
pass
|
|
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, on_request, on_notify)
|
|
await c2s_send.send(
|
|
SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=request_id, method="t", params=None))
|
|
)
|
|
await c2s_send.send(
|
|
SessionMessage(
|
|
message=JSONRPCNotification(
|
|
jsonrpc="2.0", method="notifications/cancelled", params={"requestId": cancel_id}
|
|
)
|
|
)
|
|
)
|
|
with anyio.fail_after(5):
|
|
resp = await s2c_recv.receive()
|
|
assert isinstance(resp, SessionMessage)
|
|
assert isinstance(resp.message, JSONRPCError)
|
|
assert resp.message.id == request_id # response echoes the peer's id form verbatim
|
|
assert resp.message.error == ErrorData(code=0, message="Request cancelled")
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_completed_handler_does_not_evict_reused_request_id_from_in_flight():
|
|
"""The awaited response write sits after the `_in_flight` pop; a second
|
|
request reusing the id during that window must keep its own entry (a
|
|
second, post-write pop would evict it and break its peer-cancellation)."""
|
|
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage | Exception](32)
|
|
# buffer=0: the first handler's response write parks until the test receives.
|
|
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](0)
|
|
server: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(c2s_recv, s2c_send)
|
|
calls = 0
|
|
second_started = anyio.Event()
|
|
second_exited = anyio.Event()
|
|
|
|
async def on_request(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls == 1:
|
|
return {"first": True}
|
|
second_started.set()
|
|
try:
|
|
await anyio.sleep_forever()
|
|
finally:
|
|
second_exited.set()
|
|
raise NotImplementedError
|
|
|
|
async def on_notify(ctx: DCtx, method: str, params: Mapping[str, Any] | None) -> None:
|
|
pass
|
|
|
|
try:
|
|
async with anyio.create_task_group() as tg:
|
|
await tg.start(server.run, on_request, on_notify)
|
|
with anyio.fail_after(5):
|
|
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=7, method="a")))
|
|
# First handler is now parked in `_write_result`; reuse its id.
|
|
await c2s_send.send(SessionMessage(message=JSONRPCRequest(jsonrpc="2.0", id=7, method="b")))
|
|
await second_started.wait()
|
|
resp1 = await s2c_recv.receive()
|
|
assert isinstance(resp1, SessionMessage)
|
|
assert isinstance(resp1.message, JSONRPCResponse)
|
|
assert resp1.message.result == {"first": True}
|
|
# Let the first handler task run to completion past the write.
|
|
await anyio.wait_all_tasks_blocked()
|
|
assert 7 in server._in_flight # pyright: ignore[reportPrivateUsage]
|
|
# The surviving entry must still be cancellable.
|
|
await c2s_send.send(
|
|
SessionMessage(
|
|
message=JSONRPCNotification(
|
|
jsonrpc="2.0", method="notifications/cancelled", params={"requestId": 7}
|
|
)
|
|
)
|
|
)
|
|
resp2 = await s2c_recv.receive()
|
|
assert isinstance(resp2, SessionMessage)
|
|
assert isinstance(resp2.message, JSONRPCError)
|
|
assert resp2.message.error == ErrorData(code=0, message="Request cancelled")
|
|
assert second_exited.is_set()
|
|
tg.cancel_scope.cancel()
|
|
finally:
|
|
for s in (c2s_send, c2s_recv, s2c_send, s2c_recv):
|
|
s.close()
|
|
|
|
|
|
def test_outbound_metadata_with_related_request_id_drops_resumption_hints_with_debug_log(
|
|
caplog: pytest.LogCaptureFixture,
|
|
):
|
|
"""`SessionMessage.metadata` carries one object; `related_request_id` wins
|
|
and resumption hints are dropped observably (debug log)."""
|
|
with caplog.at_level(logging.DEBUG, logger="mcp.shared.jsonrpc_dispatcher"):
|
|
md = _outbound_metadata(7, {"resumption_token": "abc"})
|
|
assert isinstance(md, ServerMessageMetadata)
|
|
assert md.related_request_id == 7
|
|
assert "dropping resumption hints" in caplog.text
|
|
caplog.clear()
|
|
with caplog.at_level(logging.DEBUG, logger="mcp.shared.jsonrpc_dispatcher"):
|
|
md = _outbound_metadata(7, {"timeout": 1.0})
|
|
assert isinstance(md, ServerMessageMetadata)
|
|
assert "dropping resumption hints" not in caplog.text
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_server_middleware_observes_cancelled_notification():
|
|
"""End-to-end over the JSON-RPC path: `Server.middleware` wraps every inbound
|
|
notification, including `notifications/cancelled` (the dispatcher applies
|
|
the cancellation itself, then forwards the notification)."""
|
|
handler_started = anyio.Event()
|
|
cancel_observed = anyio.Event()
|
|
observed: list[tuple[str, dict[str, Any]]] = []
|
|
request_id: RequestId | None = None
|
|
|
|
async def handle_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
|
|
nonlocal request_id
|
|
request_id = ctx.request_id
|
|
handler_started.set()
|
|
await anyio.sleep_forever()
|
|
raise NotImplementedError
|
|
|
|
async def observe(ctx: Any, method: str, params: Mapping[str, Any] | None, call_next: Any) -> Any:
|
|
if method == "notifications/cancelled":
|
|
observed.append((method, dict(params or {})))
|
|
cancel_observed.set()
|
|
return await call_next()
|
|
|
|
server = Server("test-server", on_call_tool=handle_call_tool)
|
|
server.middleware.append(observe)
|
|
|
|
async with Client(server) as client:
|
|
with anyio.fail_after(5):
|
|
async with anyio.create_task_group() as tg: # pragma: no branch
|
|
|
|
async def call() -> None:
|
|
with pytest.raises(MCPError):
|
|
await client.session.send_request(
|
|
CallToolRequest(params=CallToolRequestParams(name="t", arguments={})),
|
|
CallToolResult,
|
|
)
|
|
|
|
tg.start_soon(call)
|
|
await handler_started.wait()
|
|
assert request_id is not None
|
|
await client.session.send_notification(
|
|
CancelledNotification(
|
|
params=CancelledNotificationParams(request_id=request_id, reason="user clicked stop")
|
|
)
|
|
)
|
|
await cancel_observed.wait()
|
|
assert len(observed) == 1
|
|
assert observed[0][0] == "notifications/cancelled"
|
|
assert observed[0][1]["requestId"] == request_id
|
|
assert observed[0][1]["reason"] == "user clicked stop"
|