fix(session): return METHOD_NOT_FOUND (-32601) for unknown request methods

Author: Ranga-Prasath-22

src/mcp/shared/session.py

@@ -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, get_args
 
 import anyio
 import httpx
@@ -17,6 +17,7 @@
 from mcp.types import (
     CONNECTION_CLOSED,
     INVALID_PARAMS,
+    METHOD_NOT_FOUND,
     CancelledNotification,
     ClientNotification,
     ClientRequest,
@@ -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()
+
+        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 = 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,
@@ -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 = {}
@@ -348,6 +375,13 @@ 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 (
+            not isinstance(request.method, str) or 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,
@@ -385,11 +419,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)

tests/server/test_session.py

@@ -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)

tests/shared/test_session.py

@@ -339,3 +339,37 @@ 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.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)