234 lines
9.9 KiB
Python
234 lines
9.9 KiB
Python
"""Cancellation interactions against the low-level Server, driven through the public Client API.
|
|
|
|
There is no client-side cancellation API: cancelling means sending a CancelledNotification
|
|
carrying the request id, which only the server-side handler can observe (`ctx.request_id`), so
|
|
these tests capture the id from inside the blocked handler before cancelling. The handler blocks
|
|
on an Event rather than a sleep, and every wait is bounded by `anyio.fail_after`.
|
|
"""
|
|
|
|
import anyio
|
|
import pytest
|
|
from inline_snapshot import snapshot
|
|
|
|
from mcp import MCPError, types
|
|
from mcp.client import 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 (
|
|
CallToolResult,
|
|
EmptyResult,
|
|
ErrorData,
|
|
Implementation,
|
|
InitializeResult,
|
|
JSONRPCNotification,
|
|
JSONRPCRequest,
|
|
JSONRPCResponse,
|
|
PingRequest,
|
|
ServerCapabilities,
|
|
TextContent,
|
|
)
|
|
from tests.interaction._connect import Connect
|
|
from tests.interaction._helpers import IncomingMessage
|
|
from tests.interaction._requirements import requirement
|
|
|
|
pytestmark = pytest.mark.anyio
|
|
|
|
|
|
@requirement("protocol:cancel:in-flight")
|
|
@requirement("protocol:cancel:handler-abort-propagates")
|
|
async def test_cancellation_stops_in_flight_handler(connect: Connect) -> None:
|
|
"""Cancelling an in-flight request interrupts its handler and fails the pending call.
|
|
|
|
The server answers the cancelled request with an error response (the spec says it should
|
|
not respond at all; see the divergence note on the requirement), so the caller's pending
|
|
request raises rather than hanging.
|
|
"""
|
|
started = anyio.Event()
|
|
handler_cancelled = anyio.Event()
|
|
request_ids: list[types.RequestId] = []
|
|
errors: list[ErrorData] = []
|
|
|
|
async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult:
|
|
assert params.name == "block"
|
|
assert ctx.request_id is not None
|
|
request_ids.append(ctx.request_id)
|
|
started.set()
|
|
try:
|
|
await anyio.Event().wait() # blocks until cancelled; nothing ever sets this event
|
|
except anyio.get_cancelled_exc_class():
|
|
handler_cancelled.set()
|
|
raise
|
|
raise NotImplementedError # unreachable: the wait above never completes normally
|
|
|
|
server = Server("blocker", on_call_tool=call_tool)
|
|
|
|
async with connect(server) as client:
|
|
with anyio.fail_after(5):
|
|
async with anyio.create_task_group() as task_group:
|
|
|
|
async def call_and_capture_error() -> None:
|
|
with pytest.raises(MCPError) as exc_info:
|
|
await client.call_tool("block", {})
|
|
errors.append(exc_info.value.error)
|
|
|
|
task_group.start_soon(call_and_capture_error)
|
|
await started.wait()
|
|
await client.session.send_notification(
|
|
types.CancelledNotification(
|
|
params=types.CancelledNotificationParams(request_id=request_ids[0], reason="user aborted")
|
|
)
|
|
)
|
|
|
|
await handler_cancelled.wait()
|
|
|
|
assert errors == snapshot([ErrorData(code=0, message="Request cancelled")])
|
|
|
|
|
|
@requirement("protocol:cancel:server-survives")
|
|
async def test_session_serves_requests_after_cancellation(connect: Connect) -> None:
|
|
"""A request cancelled mid-flight does not poison the session: the next request succeeds."""
|
|
started = anyio.Event()
|
|
request_ids: list[types.RequestId] = []
|
|
|
|
async def list_tools(
|
|
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
|
|
) -> types.ListToolsResult:
|
|
return types.ListToolsResult(
|
|
tools=[
|
|
types.Tool(name="block", input_schema={"type": "object"}),
|
|
types.Tool(name="echo", input_schema={"type": "object"}),
|
|
]
|
|
)
|
|
|
|
async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult:
|
|
if params.name == "echo":
|
|
return CallToolResult(content=[TextContent(text="still alive")])
|
|
assert ctx.request_id is not None
|
|
request_ids.append(ctx.request_id)
|
|
started.set()
|
|
await anyio.Event().wait() # blocks until cancelled
|
|
raise NotImplementedError # unreachable
|
|
|
|
server = Server("blocker", on_list_tools=list_tools, on_call_tool=call_tool)
|
|
|
|
async with connect(server) as client:
|
|
with anyio.fail_after(5):
|
|
async with anyio.create_task_group() as task_group:
|
|
|
|
async def call_and_swallow_cancellation_error() -> None:
|
|
with pytest.raises(MCPError):
|
|
await client.call_tool("block", {})
|
|
|
|
task_group.start_soon(call_and_swallow_cancellation_error)
|
|
await started.wait()
|
|
await client.session.send_notification(
|
|
types.CancelledNotification(params=types.CancelledNotificationParams(request_id=request_ids[0]))
|
|
)
|
|
|
|
result = await client.call_tool("echo", {})
|
|
|
|
assert result == snapshot(CallToolResult(content=[TextContent(text="still alive")]))
|
|
|
|
|
|
@requirement("protocol:cancel:unknown-id-ignored")
|
|
async def test_cancellation_for_unknown_request_is_ignored(connect: Connect) -> None:
|
|
"""A cancellation referencing a request id that is not in flight is ignored without error."""
|
|
|
|
async def list_tools(
|
|
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
|
|
) -> types.ListToolsResult:
|
|
return types.ListToolsResult(tools=[types.Tool(name="echo", input_schema={"type": "object"})])
|
|
|
|
async def call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> CallToolResult:
|
|
assert params.name == "echo"
|
|
return CallToolResult(content=[TextContent(text="unbothered")])
|
|
|
|
server = Server("calm", on_list_tools=list_tools, on_call_tool=call_tool)
|
|
|
|
async with connect(server) as client:
|
|
await client.session.send_notification(
|
|
types.CancelledNotification(params=types.CancelledNotificationParams(request_id=9999))
|
|
)
|
|
result = await client.call_tool("echo", {})
|
|
|
|
assert result == snapshot(CallToolResult(content=[TextContent(text="unbothered")]))
|
|
|
|
|
|
@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.
|
|
|
|
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.
|
|
|
|
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
|
|
other tests in this file run over the transport matrix; this one is in-memory only because the
|
|
scripted-peer mechanism is the in-memory stream pair, not because the behaviour is
|
|
transport-specific.
|
|
"""
|
|
|
|
async def scripted_server(streams: MessageStream) -> None:
|
|
server_read, server_write = streams
|
|
|
|
def respond(request_id: types.RequestId, result: types.Result) -> SessionMessage:
|
|
return SessionMessage(
|
|
JSONRPCResponse(
|
|
jsonrpc="2.0",
|
|
id=request_id,
|
|
# Serialized exactly as a real server serializes results onto the wire.
|
|
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
|
|
)
|
|
)
|
|
|
|
init = await server_read.receive()
|
|
assert isinstance(init, SessionMessage)
|
|
assert isinstance(init.message, JSONRPCRequest)
|
|
assert init.message.method == "initialize"
|
|
await server_write.send(
|
|
respond(
|
|
init.message.id,
|
|
InitializeResult(
|
|
protocol_version="2025-11-25",
|
|
capabilities=ServerCapabilities(),
|
|
server_info=Implementation(name="scripted", version="0.0.1"),
|
|
),
|
|
)
|
|
)
|
|
|
|
initialized = await server_read.receive()
|
|
assert isinstance(initialized, SessionMessage)
|
|
assert isinstance(initialized.message, JSONRPCNotification)
|
|
assert initialized.message.method == "notifications/initialized"
|
|
|
|
ping = await server_read.receive()
|
|
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.
|
|
await server_write.send(respond(9999, EmptyResult()))
|
|
await server_write.send(respond(ping.message.id, EmptyResult()))
|
|
|
|
incoming: list[IncomingMessage] = []
|
|
|
|
async def message_handler(message: IncomingMessage) -> None:
|
|
incoming.append(message)
|
|
|
|
async with (
|
|
create_client_server_memory_streams() as ((client_read, client_write), server_streams),
|
|
anyio.create_task_group() as task_group,
|
|
ClientSession(client_read, client_write, message_handler=message_handler) as session,
|
|
):
|
|
task_group.start_soon(scripted_server, server_streams)
|
|
with anyio.fail_after(5):
|
|
await session.initialize()
|
|
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:")
|