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
49 changes: 32 additions & 17 deletions src/mcp/server/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down Expand Up @@ -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
Expand Down
97 changes: 97 additions & 0 deletions tests/server/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
Loading