[v2] ClientSession runs on JSONRPCDispatcher; BaseSession removed (#2838)
This commit is contained in:
Max
GitHub
parent
cf110e3214
commit
1012d60004
5 changed files with 2886 additions and 1734 deletions
|
|
@ -2,9 +2,8 @@ import pytest
|
|||
from pydantic import FileUrl
|
||||
|
||||
from mcp import Client
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.client import ClientRequestContext
|
||||
from mcp.server.mcpserver import Context, MCPServer
|
||||
from mcp.shared._context import RequestContext
|
||||
from mcp.types import ListRootsResult, Root, TextContent
|
||||
|
||||
|
||||
|
|
@ -20,7 +19,7 @@ async def test_list_roots_callback():
|
|||
)
|
||||
|
||||
async def list_roots_callback(
|
||||
context: RequestContext[ClientSession],
|
||||
context: ClientRequestContext,
|
||||
) -> ListRootsResult:
|
||||
return callback_return
|
||||
|
||||
|
|
|
|||
|
|
@ -1,66 +0,0 @@
|
|||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.session import BaseSession, RequestId, SendResultT
|
||||
from mcp.types import ClientNotification, ClientRequest, ClientResult, EmptyResult, ErrorData, PingRequest
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_send_request_stream_cleanup():
|
||||
"""Test that send_request properly cleans up streams when an exception occurs.
|
||||
|
||||
This test mocks out most of the session functionality to focus on stream cleanup.
|
||||
"""
|
||||
|
||||
# Create a mock session with the minimal required functionality
|
||||
class TestSession(BaseSession[ClientRequest, ClientNotification, ClientResult, Any, Any]):
|
||||
async def _send_response(
|
||||
self, request_id: RequestId, response: SendResultT | ErrorData
|
||||
) -> None: # pragma: no cover
|
||||
pass
|
||||
|
||||
@property
|
||||
def _receive_request_adapter(self) -> TypeAdapter[Any]:
|
||||
return TypeAdapter(object) # pragma: no cover
|
||||
|
||||
@property
|
||||
def _receive_notification_adapter(self) -> TypeAdapter[Any]:
|
||||
return TypeAdapter(object) # pragma: no cover
|
||||
|
||||
# Create streams
|
||||
write_stream_send, write_stream_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
read_stream_send, read_stream_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
|
||||
# Create the session
|
||||
session = TestSession(read_stream_receive, write_stream_send)
|
||||
|
||||
# Create a test request
|
||||
request = PingRequest()
|
||||
|
||||
# Patch the _write_stream.send method to raise an exception
|
||||
async def mock_send(*args: Any, **kwargs: Any):
|
||||
raise RuntimeError("Simulated network error")
|
||||
|
||||
# Record the response streams before the test
|
||||
initial_stream_count = len(session._response_streams)
|
||||
|
||||
# Run the test with the patched method
|
||||
with patch.object(session._write_stream, "send", mock_send):
|
||||
with pytest.raises(RuntimeError):
|
||||
await session.send_request(request, EmptyResult)
|
||||
|
||||
# Verify that no response streams were leaked
|
||||
assert len(session._response_streams) == initial_stream_count, (
|
||||
f"Expected {initial_stream_count} response streams after request, but found {len(session._response_streams)}"
|
||||
)
|
||||
|
||||
# Clean up
|
||||
await write_stream_send.aclose()
|
||||
await write_stream_receive.aclose()
|
||||
await read_stream_send.aclose()
|
||||
await read_stream_receive.aclose()
|
||||
|
|
@ -1,9 +1,8 @@
|
|||
import pytest
|
||||
|
||||
from mcp import Client
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.client import ClientRequestContext
|
||||
from mcp.server.mcpserver import Context, MCPServer
|
||||
from mcp.shared._context import RequestContext
|
||||
from mcp.types import (
|
||||
CreateMessageRequestParams,
|
||||
CreateMessageResult,
|
||||
|
|
@ -26,7 +25,7 @@ async def test_sampling_callback():
|
|||
)
|
||||
|
||||
async def sampling_callback(
|
||||
context: RequestContext[ClientSession],
|
||||
context: ClientRequestContext,
|
||||
params: CreateMessageRequestParams,
|
||||
) -> CreateMessageResult:
|
||||
return callback_return
|
||||
|
|
@ -71,7 +70,7 @@ async def test_create_message_backwards_compat_single_content():
|
|||
)
|
||||
|
||||
async def sampling_callback(
|
||||
context: RequestContext[ClientSession],
|
||||
context: ClientRequestContext,
|
||||
params: CreateMessageRequestParams,
|
||||
) -> CreateMessageResult:
|
||||
return callback_return
|
||||
|
|
|
|||
|
|
@ -1,22 +1,30 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
import anyio
|
||||
import anyio.abc
|
||||
import anyio.streams.memory
|
||||
import pytest
|
||||
from pydantic import FileUrl
|
||||
|
||||
from mcp import types
|
||||
from mcp import MCPError, types
|
||||
from mcp.client import ClientRequestContext
|
||||
from mcp.client.session import DEFAULT_CLIENT_INFO, ClientSession
|
||||
from mcp.shared._context import RequestContext
|
||||
from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair
|
||||
from mcp.shared.dispatcher import CallOptions, DispatchContext, OnNotify, OnRequest
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.session import RequestResponder
|
||||
from mcp.shared.transport_context import TransportContext
|
||||
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
|
||||
from mcp.types import (
|
||||
CONNECTION_CLOSED,
|
||||
INVALID_PARAMS,
|
||||
LATEST_PROTOCOL_VERSION,
|
||||
METHOD_NOT_FOUND,
|
||||
REQUEST_TIMEOUT,
|
||||
CallToolResult,
|
||||
Implementation,
|
||||
InitializedNotification,
|
||||
|
|
@ -420,7 +428,7 @@ async def test_client_capabilities_with_custom_callbacks():
|
|||
received_capabilities = None
|
||||
|
||||
async def custom_sampling_callback( # pragma: no cover
|
||||
context: RequestContext[ClientSession],
|
||||
context: ClientRequestContext,
|
||||
params: types.CreateMessageRequestParams,
|
||||
) -> types.CreateMessageResult | types.ErrorData:
|
||||
return types.CreateMessageResult(
|
||||
|
|
@ -430,7 +438,7 @@ async def test_client_capabilities_with_custom_callbacks():
|
|||
)
|
||||
|
||||
async def custom_list_roots_callback( # pragma: no cover
|
||||
context: RequestContext[ClientSession],
|
||||
context: ClientRequestContext,
|
||||
) -> types.ListRootsResult | types.ErrorData:
|
||||
return types.ListRootsResult(roots=[])
|
||||
|
||||
|
|
@ -504,7 +512,7 @@ async def test_client_capabilities_with_sampling_tools():
|
|||
received_capabilities = None
|
||||
|
||||
async def custom_sampling_callback( # pragma: no cover
|
||||
context: RequestContext[ClientSession],
|
||||
context: ClientRequestContext,
|
||||
params: types.CreateMessageRequestParams,
|
||||
) -> types.CreateMessageResult | types.ErrorData:
|
||||
return types.CreateMessageResult(
|
||||
|
|
@ -751,8 +759,32 @@ async def test_receive_loop_answers_malformed_inbound_request_with_invalid_param
|
|||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_receive_loop_answers_invalid_params_when_sampling_callback_raises():
|
||||
"""Same boundary catches exceptions from the request handler itself."""
|
||||
async def test_receive_loop_answers_unknown_request_method_with_method_not_found():
|
||||
"""An unknown request method is answered with METHOD_NOT_FOUND, not INVALID_PARAMS (spec-mandated)."""
|
||||
async with raw_client_session() as (_session, to_client, from_client):
|
||||
await to_client.send(SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=7, method="x/unknown")))
|
||||
out = await from_client.receive()
|
||||
assert isinstance(out.message, JSONRPCError)
|
||||
assert out.message.id == 7
|
||||
assert out.message.error == types.ErrorData(code=METHOD_NOT_FOUND, message="Method not found", data="x/unknown")
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_receive_loop_drops_unknown_notification_method_without_response():
|
||||
"""An unknown notification method is dropped silently: JSON-RPC forbids responses to notifications."""
|
||||
async with raw_client_session() as (_session, to_client, from_client):
|
||||
await to_client.send(SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="x/unknown")))
|
||||
# The answered follow-up ping proves no response was emitted and the loop survived.
|
||||
await to_client.send(SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")))
|
||||
out = await from_client.receive()
|
||||
assert isinstance(out.message, JSONRPCResponse)
|
||||
assert out.message.id == 1
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_raising_sampling_callback_answers_with_code_zero():
|
||||
"""A raising sampling callback is answered with code 0 and `str(exc)` (SDK-defined).
|
||||
Raw streams because the assertion is the outbound `JSONRPCError` envelope itself."""
|
||||
|
||||
async def boom(ctx: object, params: object) -> types.CreateMessageResult:
|
||||
raise RuntimeError("sampling boom")
|
||||
|
|
@ -767,12 +799,13 @@ async def test_receive_loop_answers_invalid_params_when_sampling_callback_raises
|
|||
)
|
||||
out = await from_client.receive()
|
||||
assert isinstance(out.message, JSONRPCError)
|
||||
assert out.message.error.code == INVALID_PARAMS
|
||||
assert out.message.error == types.ErrorData(code=0, message="sampling boom")
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_receive_loop_logs_and_drops_malformed_notification(caplog: pytest.LogCaptureFixture):
|
||||
"""A notification that fails ServerNotification validation is logged and dropped."""
|
||||
"""A malformed notification is logged and dropped without reaching `message_handler` (SDK-defined).
|
||||
Scripted peer: the typed API cannot emit a method outside the spec's notification union."""
|
||||
seen: list[object] = []
|
||||
delivered = anyio.Event()
|
||||
|
||||
|
|
@ -792,19 +825,54 @@ async def test_receive_loop_logs_and_drops_malformed_notification(caplog: pytest
|
|||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_receive_loop_forwards_transport_exception_to_message_handler():
|
||||
async def test_raising_message_handler_on_transport_exception_costs_the_delivery_not_the_connection(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
):
|
||||
"""A `message_handler` that raises on a transport-level `Exception` item is contained: the
|
||||
failure is logged and the receive loop keeps serving (SDK-defined). Raw streams because
|
||||
only a transport can put an `Exception` item on the read stream."""
|
||||
seen: list[object] = []
|
||||
delivered = anyio.Event()
|
||||
|
||||
async def handler(msg: object) -> None:
|
||||
seen.append(msg)
|
||||
delivered.set()
|
||||
# No checkpoint between set() and the containment log, so after wait() the log entry exists.
|
||||
raise RuntimeError("handler boom")
|
||||
|
||||
async with raw_client_session(message_handler=handler) as (_session, to_client, _):
|
||||
async with raw_client_session(message_handler=handler) as (_session, to_client, from_client):
|
||||
exc = ValueError("bad bytes")
|
||||
await to_client.send(exc)
|
||||
await delivered.wait()
|
||||
await to_client.send(SessionMessage(JSONRPCRequest(jsonrpc="2.0", id=9, method="ping")))
|
||||
out = await from_client.receive()
|
||||
assert seen == [exc]
|
||||
assert isinstance(out.message, JSONRPCResponse)
|
||||
assert out.message.id == 9
|
||||
assert "message_handler raised on transport exception" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_message_handler_awaiting_session_traffic_on_transport_exception_completes():
|
||||
"""A `message_handler` that awaits session traffic on a transport `Exception` item completes:
|
||||
fault deliveries are spawned into the task group, not run inline in the read loop (SDK-defined).
|
||||
Raw streams because only a transport can put an `Exception` item on the read stream."""
|
||||
ponged = anyio.Event()
|
||||
|
||||
# `session` resolves at call time, after the `as` clause binds it.
|
||||
async def handler(msg: object) -> None:
|
||||
assert isinstance(msg, Exception)
|
||||
await session.send_ping()
|
||||
ponged.set()
|
||||
|
||||
async with raw_client_session(message_handler=handler) as (session, to_client, from_client):
|
||||
await to_client.send(ValueError("bad bytes"))
|
||||
# Serve the handler's ping like a transport would; inline delivery would deadlock here.
|
||||
out = await from_client.receive()
|
||||
assert isinstance(out.message, JSONRPCRequest)
|
||||
assert out.message.method == "ping"
|
||||
await to_client.send(SessionMessage(JSONRPCResponse(jsonrpc="2.0", id=out.message.id, result={})))
|
||||
await ponged.wait()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
|
|
@ -814,6 +882,7 @@ async def test_receive_loop_consumes_server_cancelled_without_reaching_message_h
|
|||
The server dispatcher now emits this on sampling/elicitation timeout, but
|
||||
ClientSession has no in-flight tracking to act on it, so surfacing it would
|
||||
only break user handlers that exhaustively match ServerNotification.
|
||||
Scripted peer: the typed server API cannot emit a bare `notifications/cancelled`.
|
||||
"""
|
||||
seen: list[object] = []
|
||||
delivered = anyio.Event()
|
||||
|
|
@ -841,23 +910,344 @@ async def test_receive_loop_consumes_server_cancelled_without_reaching_message_h
|
|||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_receive_loop_swallows_progress_callback_exception(caplog: pytest.LogCaptureFixture):
|
||||
async def test_request_timeout_zero_overrides_session_timeout():
|
||||
"""`request_read_timeout_seconds=0` is a real per-request timeout (fail at the
|
||||
first checkpoint, `anyio.fail_after(0)` semantics), not a fall-through to the
|
||||
session-level timeout. The request is never answered, so falling back to the
|
||||
30s session timeout would trip the harness's 5s guard instead."""
|
||||
async with raw_client_session(read_timeout_seconds=30) as (session, _to_client, _from_client):
|
||||
with pytest.raises(MCPError) as exc_info:
|
||||
await session.send_request(types.PingRequest(), types.EmptyResult, request_read_timeout_seconds=0.0)
|
||||
assert exc_info.value.error.code == REQUEST_TIMEOUT
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_progress_notification_reaches_request_callback_and_message_handler():
|
||||
"""A `notifications/progress` for an in-flight request reaches both the `progress_callback` and
|
||||
`message_handler` (SDK-defined). Scripted peer: the progress token must echo the wire request id."""
|
||||
updates: list[tuple[float, float | None, str | None]] = []
|
||||
teed: list[types.ProgressNotification] = []
|
||||
request_id: types.RequestId | None = None
|
||||
progressed = anyio.Event()
|
||||
delivered = anyio.Event()
|
||||
|
||||
async def boom(progress: float, total: float | None, message: str | None) -> None:
|
||||
raise RuntimeError("progress boom")
|
||||
async def on_progress(progress: float, total: float | None, message: str | None) -> None:
|
||||
updates.append((progress, total, message))
|
||||
progressed.set()
|
||||
|
||||
async def handler(msg: object) -> None:
|
||||
# Only the progress notification is teed to the message handler here.
|
||||
assert isinstance(msg, types.ProgressNotification)
|
||||
teed.append(msg)
|
||||
delivered.set()
|
||||
|
||||
async with raw_client_session(message_handler=handler) as (session, to_client, _):
|
||||
# Register the callback under a known token without sending a request.
|
||||
session._progress_callbacks[42] = boom # pyright: ignore[reportPrivateUsage]
|
||||
params = {"progressToken": 42, "progress": 0.5}
|
||||
await to_client.send(
|
||||
SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/progress", params=params))
|
||||
)
|
||||
# The progress notification also reaches the message handler after the
|
||||
# callback runs, so this fires once the callback's exception is handled.
|
||||
await delivered.wait()
|
||||
assert "Progress callback raised an exception" in caplog.text
|
||||
async with raw_client_session(message_handler=handler) as (session, to_client, from_client):
|
||||
async with anyio.create_task_group() as tg:
|
||||
|
||||
async def call() -> None:
|
||||
await session.send_request(types.PingRequest(), types.EmptyResult, progress_callback=on_progress)
|
||||
|
||||
tg.start_soon(call)
|
||||
request = await from_client.receive()
|
||||
assert isinstance(request.message, JSONRPCRequest)
|
||||
request_id = request.message.id
|
||||
# The request id doubles as the progress token.
|
||||
params = {"progressToken": request_id, "progress": 0.5, "total": 1.0, "message": "halfway"}
|
||||
await to_client.send(
|
||||
SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/progress", params=params))
|
||||
)
|
||||
await progressed.wait()
|
||||
await delivered.wait()
|
||||
await to_client.send(SessionMessage(JSONRPCResponse(jsonrpc="2.0", id=request_id, result={})))
|
||||
assert updates == [(0.5, 1.0, "halfway")]
|
||||
assert request_id is not None
|
||||
assert len(teed) == 1
|
||||
assert teed[0].params == types.ProgressNotificationParams(
|
||||
progress_token=request_id, progress=0.5, total=1.0, message="halfway"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_dispatcher_keyword_runs_over_direct_dispatch():
|
||||
"""A session built with dispatcher= works without a stream pair (in-process embedding)."""
|
||||
client_side, server_side = create_direct_dispatcher_pair()
|
||||
|
||||
async def server_on_request(
|
||||
ctx: DispatchContext[TransportContext], method: str, params: dict[str, object] | None
|
||||
) -> dict[str, object]:
|
||||
assert method == "ping"
|
||||
return {}
|
||||
|
||||
notified: list[str] = []
|
||||
|
||||
async def server_on_notify(
|
||||
ctx: DispatchContext[TransportContext], method: str, params: dict[str, object] | None
|
||||
) -> None:
|
||||
notified.append(method)
|
||||
|
||||
session = ClientSession(dispatcher=client_side)
|
||||
results: list[types.EmptyResult] = []
|
||||
with anyio.fail_after(5):
|
||||
async with anyio.create_task_group() as tg:
|
||||
await tg.start(server_side.run, server_on_request, server_on_notify)
|
||||
async with session:
|
||||
results.append(await session.send_ping(meta=None))
|
||||
# Server-to-client: direct dispatch delivers ping with no params member (no _meta injection).
|
||||
assert await server_side.send_raw_request("ping", None) == {}
|
||||
await session.send_notification(types.RootsListChangedNotification())
|
||||
server_side.close()
|
||||
assert results == [types.EmptyResult()]
|
||||
assert notified == ["notifications/roots/list_changed"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_direct_dispatch_roots_list_reaches_callback_with_synthesized_request_id():
|
||||
"""A server-initiated roots/list over dispatcher= reaches the registered callback and round-trips
|
||||
the result; the callback context carries an int request_id (SDK-defined: DirectDispatcher
|
||||
synthesizes ids)."""
|
||||
client_side, server_side = create_direct_dispatcher_pair()
|
||||
contexts: list[ClientRequestContext] = []
|
||||
|
||||
async def list_roots(context: ClientRequestContext) -> types.ListRootsResult:
|
||||
contexts.append(context)
|
||||
return types.ListRootsResult(roots=[types.Root(uri=FileUrl("file:///workspace"))])
|
||||
|
||||
async def server_on_request(
|
||||
ctx: DispatchContext[TransportContext], method: str, params: dict[str, object] | None
|
||||
) -> dict[str, object]:
|
||||
raise NotImplementedError
|
||||
|
||||
async def server_on_notify(
|
||||
ctx: DispatchContext[TransportContext], method: str, params: dict[str, object] | None
|
||||
) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
session = ClientSession(dispatcher=client_side, list_roots_callback=list_roots)
|
||||
result: dict[str, Any] | None = None
|
||||
with anyio.fail_after(5):
|
||||
async with anyio.create_task_group() as tg:
|
||||
await tg.start(server_side.run, server_on_request, server_on_notify)
|
||||
async with session:
|
||||
result = await server_side.send_raw_request("roots/list", None)
|
||||
server_side.close()
|
||||
assert result == {"roots": [{"uri": "file:///workspace"}]}
|
||||
assert len(contexts) == 1
|
||||
assert isinstance(contexts[0].request_id, int)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_raising_notification_callbacks_over_direct_dispatch_cost_only_that_delivery(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
):
|
||||
"""A raising `logging_callback` or `message_handler` is contained in the session, so the
|
||||
in-process peer's notify() returns normally and the session keeps serving requests
|
||||
(SDK-defined: DirectDispatcher awaits notification handlers inline in the peer's call).
|
||||
A raising `logging_callback` skips the `message_handler` tee for that notification."""
|
||||
client_side, server_side = create_direct_dispatcher_pair()
|
||||
teed: list[types.ServerNotification] = []
|
||||
|
||||
async def logging_callback(params: types.LoggingMessageNotificationParams) -> None:
|
||||
raise ValueError("logging callback boom")
|
||||
|
||||
async def message_handler(
|
||||
message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
|
||||
) -> None:
|
||||
assert not isinstance(message, RequestResponder | Exception)
|
||||
teed.append(message)
|
||||
raise ValueError("message handler boom")
|
||||
|
||||
async def server_on_request(
|
||||
ctx: DispatchContext[TransportContext], method: str, params: dict[str, object] | None
|
||||
) -> dict[str, object]:
|
||||
assert method == "ping"
|
||||
return {}
|
||||
|
||||
async def server_on_notify(
|
||||
ctx: DispatchContext[TransportContext], method: str, params: dict[str, object] | None
|
||||
) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
session = ClientSession(dispatcher=client_side, logging_callback=logging_callback, message_handler=message_handler)
|
||||
with anyio.fail_after(5):
|
||||
async with anyio.create_task_group() as tg:
|
||||
await tg.start(server_side.run, server_on_request, server_on_notify)
|
||||
async with session:
|
||||
# logging_callback raises: notify() must return, and message_handler is skipped.
|
||||
await server_side.notify("notifications/message", {"level": "info", "data": "hello"})
|
||||
# message_handler raises: notify() must return.
|
||||
await server_side.notify("notifications/tools/list_changed", None)
|
||||
# The session still serves requests afterwards.
|
||||
assert await session.send_ping() == types.EmptyResult()
|
||||
server_side.close()
|
||||
assert [type(n) for n in teed] == [types.ToolListChangedNotification]
|
||||
assert caplog.text.count("notification callback for") == 2
|
||||
assert "notification callback for 'notifications/message' raised" in caplog.text
|
||||
assert "notification callback for 'notifications/tools/list_changed' raised" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_dispatcher_keyword_send_request_before_enter_raises_runtimeerror():
|
||||
"""The documented pre-enter RuntimeError holds for dispatcher= sessions too."""
|
||||
client_side, _server_side = create_direct_dispatcher_pair()
|
||||
session = ClientSession(dispatcher=client_side)
|
||||
with anyio.fail_after(5), pytest.raises(RuntimeError) as exc:
|
||||
await session.send_ping()
|
||||
assert str(exc.value) == "DirectDispatcher.send_raw_request called before run()"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_dispatcher_keyword_send_request_after_exit_raises_connection_closed():
|
||||
"""After __aexit__ a dispatcher= session raises MCPError(CONNECTION_CLOSED), matching the JSONRPC path."""
|
||||
client_side, server_side = create_direct_dispatcher_pair()
|
||||
|
||||
async def server_on_request(
|
||||
ctx: DispatchContext[TransportContext], method: str, params: dict[str, object] | None
|
||||
) -> dict[str, object]:
|
||||
assert method == "ping"
|
||||
return {}
|
||||
|
||||
async def server_on_notify(
|
||||
ctx: DispatchContext[TransportContext], method: str, params: dict[str, object] | None
|
||||
) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
session = ClientSession(dispatcher=client_side)
|
||||
with anyio.fail_after(5):
|
||||
async with anyio.create_task_group() as tg:
|
||||
await tg.start(server_side.run, server_on_request, server_on_notify)
|
||||
async with session:
|
||||
assert await session.send_ping() == types.EmptyResult()
|
||||
with pytest.raises(MCPError) as exc:
|
||||
await session.send_ping()
|
||||
assert exc.value.error.code == CONNECTION_CLOSED
|
||||
server_side.close()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_dispatcher_keyword_request_timeout_bounds_wait_for_never_run_peer():
|
||||
"""request_read_timeout_seconds fires even when the peer dispatcher never started running."""
|
||||
client_side, _server_side = create_direct_dispatcher_pair()
|
||||
session = ClientSession(dispatcher=client_side)
|
||||
with anyio.fail_after(5):
|
||||
async with session:
|
||||
with pytest.raises(MCPError) as exc:
|
||||
await session.send_request(types.PingRequest(), types.EmptyResult, request_read_timeout_seconds=0.01)
|
||||
assert exc.value.error.code == REQUEST_TIMEOUT
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_initialize_opts_out_of_cancel_on_abandon_while_other_requests_leave_it_unset():
|
||||
"""`send_request` passes `cancel_on_abandon=False` for `initialize` — the spec forbids
|
||||
cancelling it — and leaves the option unset for every other method."""
|
||||
|
||||
class RecordingDispatcher:
|
||||
"""Records `send_raw_request` opts and answers with canned results."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, CallOptions]] = []
|
||||
|
||||
async def run(
|
||||
self,
|
||||
on_request: OnRequest,
|
||||
on_notify: OnNotify,
|
||||
*,
|
||||
task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
|
||||
) -> None:
|
||||
task_status.started()
|
||||
await anyio.sleep_forever()
|
||||
|
||||
async def send_raw_request(
|
||||
self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None
|
||||
) -> dict[str, Any]:
|
||||
self.calls.append((method, opts or {}))
|
||||
if method == "initialize":
|
||||
return InitializeResult(
|
||||
protocol_version=LATEST_PROTOCOL_VERSION,
|
||||
capabilities=ServerCapabilities(),
|
||||
server_info=Implementation(name="mock-server", version="0.1.0"),
|
||||
).model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
return {}
|
||||
|
||||
async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
||||
pass
|
||||
|
||||
dispatcher = RecordingDispatcher()
|
||||
with anyio.fail_after(5):
|
||||
async with ClientSession(dispatcher=dispatcher) as session:
|
||||
await session.initialize()
|
||||
await session.send_ping()
|
||||
opts_by_method = dict(dispatcher.calls)
|
||||
assert opts_by_method["initialize"].get("cancel_on_abandon") is False
|
||||
assert "cancel_on_abandon" not in opts_by_method["ping"]
|
||||
|
||||
|
||||
def test_constructor_rejects_streams_and_dispatcher_together():
|
||||
client_side, _server_side = create_direct_dispatcher_pair()
|
||||
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
|
||||
with pytest.raises(ValueError, match="not both"):
|
||||
ClientSession(s2c_recv, dispatcher=client_side)
|
||||
s2c_send.close()
|
||||
s2c_recv.close()
|
||||
|
||||
|
||||
def test_constructor_requires_both_streams_without_dispatcher():
|
||||
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](1)
|
||||
with pytest.raises(ValueError, match="read_stream and write_stream are required"):
|
||||
ClientSession(s2c_recv)
|
||||
with pytest.raises(ValueError, match="read_stream and write_stream are required"):
|
||||
ClientSession()
|
||||
s2c_send.close()
|
||||
s2c_recv.close()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_aenter_cancelled_while_dispatcher_starts_unwinds_cleanly():
|
||||
"""Cancellation while `__aenter__` waits for the dispatcher to start unwinds the half-entered
|
||||
task group cleanly, not via anyio's "exited non-innermost cancel scope" RuntimeError (SDK-defined)."""
|
||||
|
||||
class NeverStartsDispatcher:
|
||||
"""`run()` parks without ever signalling `task_status.started()`."""
|
||||
|
||||
async def run(
|
||||
self,
|
||||
on_request: OnRequest,
|
||||
on_notify: OnNotify,
|
||||
*,
|
||||
task_status: anyio.abc.TaskStatus[None] = anyio.TASK_STATUS_IGNORED,
|
||||
) -> None:
|
||||
await anyio.sleep_forever()
|
||||
|
||||
async def send_raw_request(
|
||||
self, method: str, params: Mapping[str, Any] | None, opts: CallOptions | None = None
|
||||
) -> dict[str, Any]:
|
||||
raise NotImplementedError
|
||||
|
||||
async def notify(self, method: str, params: Mapping[str, Any] | None) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
session = ClientSession(dispatcher=NeverStartsDispatcher())
|
||||
async with AsyncExitStack() as stack:
|
||||
# `start()` is parked forever, so the deadline only ends the wait — any duration is non-racy.
|
||||
with anyio.move_on_after(0.01) as scope:
|
||||
await stack.enter_async_context(session)
|
||||
assert scope.cancelled_caught
|
||||
# The failed enter must not leave the session half-entered.
|
||||
assert session._task_group is None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_send_notification_after_close_is_dropped_silently():
|
||||
"""Post-close `send_notification` is fire-and-forget: the notification is dropped,
|
||||
not surfaced as a raw transport error (v1 leaked `anyio.ClosedResourceError`)."""
|
||||
s2c_send, s2c_recv = anyio.create_memory_object_stream[SessionMessage | Exception](4)
|
||||
c2s_send, c2s_recv = anyio.create_memory_object_stream[SessionMessage](4)
|
||||
try:
|
||||
async with ClientSession(s2c_recv, c2s_send) as session:
|
||||
pass
|
||||
with anyio.fail_after(5):
|
||||
await session.send_notification(types.RootsListChangedNotification())
|
||||
with pytest.raises(anyio.EndOfStream):
|
||||
c2s_recv.receive_nowait() # nothing reached the wire
|
||||
finally:
|
||||
for s in (s2c_send, s2c_recv, c2s_send, c2s_recv):
|
||||
s.close()
|
||||
|
|
|
|||
141
tests/client/test_session_concurrency.py
Normal file
141
tests/client/test_session_concurrency.py
Normal file
|
|
@ -0,0 +1,141 @@
|
|||
"""Concurrency over a single client session: multiple requests in flight at once, in both directions."""
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
from inline_snapshot import snapshot
|
||||
|
||||
from mcp import Client
|
||||
from mcp.client import ClientRequestContext
|
||||
from mcp.server.mcpserver import Context, MCPServer
|
||||
from mcp.types import (
|
||||
CallToolResult,
|
||||
CreateMessageRequestParams,
|
||||
CreateMessageResult,
|
||||
SamplingMessage,
|
||||
TextContent,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
async def test_concurrent_tool_calls_resolve_out_of_order_to_their_own_callers() -> None:
|
||||
"""Three tool calls in flight at once on one session each receive their own result, even though
|
||||
the responses come back in the reverse of the order the requests were sent.
|
||||
|
||||
SDK-defined contract: pins the client request machinery's support for concurrent in-flight
|
||||
calls with out-of-order response correlation. Each handler parks on its own release event
|
||||
after signalling it started; a session that serialized requests would never start the later
|
||||
handlers and the test would time out instead.
|
||||
"""
|
||||
send_order = ["a", "b", "c"]
|
||||
started = {tag: anyio.Event() for tag in send_order}
|
||||
release = {tag: anyio.Event() for tag in send_order}
|
||||
done = {tag: anyio.Event() for tag in send_order}
|
||||
completion_order: list[str] = []
|
||||
results: dict[str, CallToolResult] = {}
|
||||
|
||||
server = MCPServer("parking")
|
||||
|
||||
@server.tool()
|
||||
async def park(tag: str) -> str:
|
||||
started[tag].set()
|
||||
await release[tag].wait()
|
||||
return f"result:{tag}"
|
||||
|
||||
async with Client(server) as client:
|
||||
|
||||
async def call_and_record(tag: str) -> None:
|
||||
results[tag] = await client.call_tool("park", {"tag": tag})
|
||||
completion_order.append(tag)
|
||||
done[tag].set()
|
||||
|
||||
with anyio.fail_after(5):
|
||||
async with anyio.create_task_group() as task_group: # pragma: no branch
|
||||
# Waiting for each handler to start before issuing the next call fixes the send
|
||||
# order, and leaves all three parked in flight together once the loop finishes.
|
||||
for tag in send_order:
|
||||
task_group.start_soon(call_and_record, tag)
|
||||
await started[tag].wait()
|
||||
|
||||
# Nothing completed yet: all three calls are genuinely concurrent.
|
||||
assert completion_order == []
|
||||
|
||||
# Release in reverse, awaiting each completion so the finish order is forced.
|
||||
for tag in reversed(send_order):
|
||||
release[tag].set()
|
||||
await done[tag].wait()
|
||||
|
||||
assert completion_order == ["c", "b", "a"]
|
||||
assert results == snapshot(
|
||||
{
|
||||
"c": CallToolResult(content=[TextContent(text="result:c")], structured_content={"result": "result:c"}),
|
||||
"b": CallToolResult(content=[TextContent(text="result:b")], structured_content={"result": "result:b"}),
|
||||
"a": CallToolResult(content=[TextContent(text="result:a")], structured_content={"result": "result:a"}),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def test_overlapping_sampling_requests_are_serviced_concurrently_by_the_client() -> None:
|
||||
"""A server tool that fans out two sampling requests at once gets both echoes back: the client
|
||||
runs overlapping inbound `create_message` requests concurrently instead of serializing them in
|
||||
its receive loop.
|
||||
|
||||
Regression pin for https://github.com/modelcontextprotocol/python-sdk/issues/2489 -- v1's
|
||||
`BaseSession` awaited each inbound request handler inline, so the second sampling callback
|
||||
could not start until the first returned; here both rendezvous before either is released.
|
||||
"""
|
||||
sampling_started = {"x": anyio.Event(), "y": anyio.Event()}
|
||||
sampling_release = anyio.Event()
|
||||
tool_results: list[CallToolResult] = []
|
||||
|
||||
server = MCPServer("fan_out_server")
|
||||
|
||||
@server.tool()
|
||||
async def fan_out(ctx: Context) -> str:
|
||||
echoes: dict[str, str] = {}
|
||||
|
||||
async def sample(tag: str) -> None:
|
||||
result = await ctx.session.create_message(
|
||||
messages=[SamplingMessage(role="user", content=TextContent(text=tag))],
|
||||
max_tokens=10,
|
||||
)
|
||||
assert isinstance(result.content, TextContent)
|
||||
echoes[tag] = result.content.text
|
||||
|
||||
async with anyio.create_task_group() as sampler_group:
|
||||
sampler_group.start_soon(sample, "x")
|
||||
sampler_group.start_soon(sample, "y")
|
||||
return f"{echoes['x']} {echoes['y']}"
|
||||
|
||||
async def sampling_callback(
|
||||
context: ClientRequestContext, params: CreateMessageRequestParams
|
||||
) -> CreateMessageResult:
|
||||
content = params.messages[0].content
|
||||
assert isinstance(content, TextContent)
|
||||
sampling_started[content.text].set()
|
||||
await sampling_release.wait()
|
||||
return CreateMessageResult(
|
||||
role="assistant",
|
||||
content=TextContent(text=f"echo:{content.text}"),
|
||||
model="test-model",
|
||||
stop_reason="endTurn",
|
||||
)
|
||||
|
||||
async with Client(server, sampling_callback=sampling_callback) as client:
|
||||
with anyio.fail_after(5):
|
||||
async with anyio.create_task_group() as task_group: # pragma: no branch
|
||||
|
||||
async def invoke_fan_out() -> None:
|
||||
tool_results.append(await client.call_tool("fan_out", {}))
|
||||
|
||||
task_group.start_soon(invoke_fan_out)
|
||||
|
||||
# Both sampling callbacks are mid-flight before either may answer -- a client that
|
||||
# serialized inbound requests would never start the second one.
|
||||
await sampling_started["x"].wait()
|
||||
await sampling_started["y"].wait()
|
||||
sampling_release.set()
|
||||
|
||||
assert tool_results == snapshot(
|
||||
[CallToolResult(content=[TextContent(text="echo:x echo:y")], structured_content={"result": "echo:x echo:y"})]
|
||||
)
|
||||
Loading…
Reference in a new issue