From 394ed982bb34ce8f39827db6fc05188d7c9ff02e Mon Sep 17 00:00:00 2001 From: ychampion Date: Wed, 8 Jul 2026 19:38:38 +0000 Subject: [PATCH] Reject changed duplicate initialize requests Keep duplicate initialize requests idempotent when the client sends the same parameters, but return INVALID_REQUEST if a stateful server is reinitialized with different parameters after negotiation. --- src/mcp/server/session.py | 49 +++++++++++------- tests/server/test_session.py | 97 ++++++++++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 17 deletions(-) diff --git a/src/mcp/server/session.py b/src/mcp/server/session.py index f80971b012..e75e2434cb 100644 --- a/src/mcp/server/session.py +++ b/src/mcp/server/session.py @@ -94,6 +94,7 @@ def __init__( stateless: bool = False, ) -> None: super().__init__(read_stream, write_stream, types.ClientRequest, types.ClientNotification) + self._stateless = stateless self._initialization_state = ( InitializationState.Initialized if stateless else InitializationState.NotInitialized ) @@ -172,30 +173,44 @@ async def _receive_loop(self) -> None: async with self._incoming_message_stream_writer: await super()._receive_loop() + def _initialize_result(self, requested_version: str | int) -> types.InitializeResult: + return types.InitializeResult( + protocolVersion=requested_version + if requested_version in SUPPORTED_PROTOCOL_VERSIONS + else types.LATEST_PROTOCOL_VERSION, + capabilities=self._init_options.capabilities, + serverInfo=types.Implementation( + name=self._init_options.server_name, + version=self._init_options.server_version, + websiteUrl=self._init_options.website_url, + icons=self._init_options.icons, + ), + instructions=self._init_options.instructions, + ) + async def _received_request(self, responder: RequestResponder[types.ClientRequest, types.ServerResult]): match responder.request.root: case types.InitializeRequest(params=params): + if not self._stateless and self._initialization_state == InitializationState.Initialized: + if params == self._client_params: + with responder: + await responder.respond(types.ServerResult(self._initialize_result(params.protocolVersion))) + return + + with responder: + await responder.respond( + types.ErrorData( + code=types.INVALID_REQUEST, + message="Server is already initialized", + ) + ) + return + requested_version = params.protocolVersion self._initialization_state = InitializationState.Initializing self._client_params = params with responder: - await responder.respond( - types.ServerResult( - types.InitializeResult( - protocolVersion=requested_version - if requested_version in SUPPORTED_PROTOCOL_VERSIONS - else types.LATEST_PROTOCOL_VERSION, - capabilities=self._init_options.capabilities, - serverInfo=types.Implementation( - name=self._init_options.server_name, - version=self._init_options.server_version, - websiteUrl=self._init_options.website_url, - icons=self._init_options.icons, - ), - instructions=self._init_options.instructions, - ) - ) - ) + await responder.respond(types.ServerResult(self._initialize_result(requested_version))) self._initialization_state = InitializationState.Initialized case types.PingRequest(): # Ping requests are allowed at any time diff --git a/tests/server/test_session.py b/tests/server/test_session.py index ba1b44126d..bdae5c278d 100644 --- a/tests/server/test_session.py +++ b/tests/server/test_session.py @@ -220,6 +220,103 @@ async def mock_client(): assert received_protocol_version == "2024-11-05" +@pytest.mark.anyio +async def test_duplicate_initialize_after_initialized_preserves_negotiated_params(): + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](10) + first_capabilities = types.ClientCapabilities(sampling=types.SamplingCapability()) + final_client_params: types.InitializeRequestParams | None = None + + def initialize_message( + request_id: int, + client_name: str, + capabilities: types.ClientCapabilities, + protocol_version: str = types.LATEST_PROTOCOL_VERSION, + ) -> SessionMessage: + return SessionMessage( + types.JSONRPCMessage( + types.JSONRPCRequest( + jsonrpc="2.0", + id=request_id, + method="initialize", + params=types.InitializeRequestParams( + protocolVersion=protocol_version, + capabilities=capabilities, + clientInfo=types.Implementation(name=client_name, version="1.0.0"), + ).model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) + ) + + async def run_server(): + nonlocal final_client_params + + async with ServerSession( + client_to_server_receive, + server_to_client_send, + InitializationOptions( + server_name="mcp", + server_version="0.1.0", + capabilities=ServerCapabilities(), + ), + ) as server_session: + async for message in server_session.incoming_messages: + if isinstance(message, Exception): # pragma: no cover + raise message + + final_client_params = server_session.client_params + assert server_session.check_client_capability(first_capabilities) + + async def mock_client(): + await client_to_server_send.send(initialize_message(1, "client-a", first_capabilities)) + + with anyio.fail_after(5): + first_response = await server_to_client_receive.receive() + assert isinstance(first_response.message.root, types.JSONRPCResponse) + + await client_to_server_send.send( + SessionMessage( + types.JSONRPCMessage( + types.JSONRPCNotification( + jsonrpc="2.0", + method="notifications/initialized", + ) + ) + ) + ) + + await client_to_server_send.send(initialize_message(2, "client-a", first_capabilities)) + + with anyio.fail_after(5): + second_response = await server_to_client_receive.receive() + assert isinstance(second_response.message.root, types.JSONRPCResponse) + + await client_to_server_send.send( + initialize_message(3, "client-b", types.ClientCapabilities(), protocol_version="2024-11-05") + ) + + with anyio.fail_after(5): + third_response = await server_to_client_receive.receive() + assert isinstance(third_response.message.root, types.JSONRPCError) + assert third_response.message.root.error.code == types.INVALID_REQUEST + + await client_to_server_send.aclose() + + 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) + + assert final_client_params is not None + assert final_client_params.clientInfo.name == "client-a" + assert final_client_params.protocolVersion == types.LATEST_PROTOCOL_VERSION + + @pytest.mark.anyio async def test_ping_request_before_initialization(): """Test that ping requests are allowed before initialization is complete."""