Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 34 additions & 6 deletions src/mcp/shared/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from contextlib import AsyncExitStack
from datetime import timedelta
from types import TracebackType
from typing import Any, Generic, Protocol, TypeVar
from typing import Any, Generic, Protocol, TypeVar, cast, get_args

import anyio
import httpx
Expand All @@ -17,6 +17,7 @@
from mcp.types import (
CONNECTION_CLOSED,
INVALID_PARAMS,
METHOD_NOT_FOUND,
CancelledNotification,
ClientNotification,
ClientRequest,
Expand Down Expand Up @@ -159,6 +160,31 @@ def cancelled(self) -> bool: # pragma: no cover
return self._cancel_scope.cancel_called


def _extract_known_request_methods(request_type: type[Any]) -> frozenset[str]:
"""Extract default method names from a Pydantic RootModel or Union of request models."""
try:
union_type = getattr(request_type, "__value__", None) or request_type
if hasattr(union_type, "model_fields") and "root" in union_type.model_fields:
union_type = union_type.model_fields["root"].annotation

methods: set[str] = set()

def _unpack_union(t: Any) -> None:
args = get_args(t)
if args:
for arg in args:
_unpack_union(arg)
elif hasattr(t, "model_fields") and "method" in t.model_fields:
m = cast(Any, t.model_fields["method"].default)
if isinstance(m, str):
methods.add(m)

_unpack_union(union_type)
return frozenset(methods)
except Exception:
return frozenset()


class BaseSession(
Generic[
SendRequestT,
Expand Down Expand Up @@ -197,6 +223,7 @@ def __init__(
self._request_id = 0
self._receive_request_type = receive_request_type
self._receive_notification_type = receive_notification_type
self._known_request_methods = _extract_known_request_methods(receive_request_type)
self._session_read_timeout_seconds = read_timeout_seconds
self._in_flight = {}
self._progress_callbacks = {}
Expand Down Expand Up @@ -348,6 +375,11 @@ async def _send_response(self, request_id: RequestId, response: SendResultT | Er
session_message = SessionMessage(message=JSONRPCMessage(jsonrpc_response))
await self._write_stream.send(session_message)

def _get_request_validation_error(self, request: JSONRPCRequest) -> ErrorData:
if self._known_request_methods and request.method not in self._known_request_methods:
return ErrorData(code=METHOD_NOT_FOUND, message="Method not found")
return ErrorData(code=INVALID_PARAMS, message="Invalid request parameters")

async def _receive_loop(self) -> None:
async with (
self._read_stream,
Expand Down Expand Up @@ -385,11 +417,7 @@ async def _receive_loop(self) -> None:
error_response = JSONRPCError(
jsonrpc="2.0",
id=message.message.root.id,
error=ErrorData(
code=INVALID_PARAMS,
message="Invalid request parameters",
data="",
),
error=self._get_request_validation_error(message.message.root),
)
session_message = SessionMessage(message=JSONRPCMessage(error_response))
await self._write_stream.send(session_message)
Expand Down
72 changes: 72 additions & 0 deletions tests/server/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -521,3 +521,75 @@ async def mock_client():

assert error_response_received
assert error_code == types.INVALID_PARAMS


@pytest.mark.anyio
async def test_unknown_method_returns_method_not_found():
"""Test that unknown request methods return METHOD_NOT_FOUND (-32601),
while malformed known methods return INVALID_PARAMS (-32602).
"""
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)

async def run_server():
async with ServerSession(
client_to_server_receive,
server_to_client_send,
InitializationOptions(
server_name="mcp",
server_version="0.1.0",
capabilities=ServerCapabilities(),
),
):
await anyio.sleep(0.2)

async def mock_client():
# 1. Send unknown method request
await client_to_server_send.send(
SessionMessage(
types.JSONRPCMessage(
types.JSONRPCRequest(
jsonrpc="2.0",
id=1,
method="totally/bogus",
params={},
)
)
)
)

resp1 = await server_to_client_receive.receive()
assert isinstance(resp1.message.root, types.JSONRPCError)
assert resp1.message.root.id == 1
assert resp1.message.root.error.code == types.METHOD_NOT_FOUND
assert resp1.message.root.error.message == "Method not found"

# 2. Send malformed known method request (initialize with invalid params shape)
await client_to_server_send.send(
SessionMessage(
types.JSONRPCMessage(
types.JSONRPCRequest(
jsonrpc="2.0",
id=2,
method="initialize",
params={"protocolVersion": 12345}, # invalid type for protocolVersion
)
)
)
)

resp2 = await server_to_client_receive.receive()
assert isinstance(resp2.message.root, types.JSONRPCError)
assert resp2.message.root.id == 2
assert resp2.message.root.error.code == types.INVALID_PARAMS
assert resp2.message.root.error.message == "Invalid request parameters"

async with (
client_to_server_send,
client_to_server_receive,
server_to_client_send,
server_to_client_receive,
anyio.create_task_group() as tg,
):
tg.start_soon(run_server)
tg.start_soon(mock_client)
56 changes: 56 additions & 0 deletions tests/shared/test_session.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from collections.abc import AsyncGenerator
from types import SimpleNamespace
from typing import Any

import anyio
Expand All @@ -10,6 +11,7 @@
from mcp.shared.exceptions import McpError
from mcp.shared.memory import create_client_server_memory_streams, create_connected_server_and_client_session
from mcp.shared.message import SessionMessage
from mcp.shared.session import _extract_known_request_methods
from mcp.types import (
CancelledNotification,
CancelledNotificationParams,
Expand Down Expand Up @@ -339,3 +341,57 @@ async def mock_server():
await ev_closed.wait()
with anyio.fail_after(1): # pragma: no cover
await ev_response.wait()


@pytest.mark.anyio
async def test_client_session_unknown_method_returns_method_not_found():
"""Test that ClientSession returns METHOD_NOT_FOUND (-32601) when receiving an unknown server request method."""
async with create_client_server_memory_streams() as (client_stream, server_stream):
client_read, client_write = client_stream
server_read, server_write = server_stream

async def mock_server():
# Send unknown request method to client
await server_write.send(
SessionMessage(
message=JSONRPCMessage(
JSONRPCRequest(
jsonrpc="2.0",
id=1,
method="invalid/server_method",
params={},
)
)
)
)
resp = await server_read.receive()
assert isinstance(resp, SessionMessage)
assert isinstance(resp.message.root, JSONRPCError)
assert resp.message.root.id == 1
assert resp.message.root.error.code == types.METHOD_NOT_FOUND
assert resp.message.root.error.message == "Method not found"

async with (
ClientSession(read_stream=client_read, write_stream=client_write),
anyio.create_task_group() as tg,
):
tg.start_soon(mock_server)


def test_extract_known_request_methods_ignores_non_string_defaults():
class RequestWithNonStringMethod:
model_fields = {"method": SimpleNamespace(default=None)}

assert _extract_known_request_methods(RequestWithNonStringMethod) == frozenset()


def test_extract_known_request_methods_fails_closed_on_schema_introspection_error():
class ExplodingMeta(type):
@property
def __value__(cls) -> type[Any]:
raise RuntimeError("broken schema")

class BrokenRequest(metaclass=ExplodingMeta):
pass

assert _extract_known_request_methods(BrokenRequest) == frozenset()
Loading