Drop RootModel from JSONRPCMessage (#1908)

This commit is contained in:
Marcelo Trylesinski 2026-01-19 14:04:15 +01:00 • committed by GitHub
commit f4672c5084
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 314 additions and 476 deletions

View file

@ -82,13 +82,11 @@ async def test_lowlevel_server_lifespan():
)
await send_stream1.send(
SessionMessage(
JSONRPCMessage(
root=JSONRPCRequest(
jsonrpc="2.0",
id=1,
method="initialize",
params=TypeAdapter(InitializeRequestParams).dump_python(params),
)
JSONRPCRequest(
jsonrpc="2.0",
id=1,
method="initialize",
params=TypeAdapter(InitializeRequestParams).dump_python(params),
)
)
)
@ -96,27 +94,16 @@ async def test_lowlevel_server_lifespan():
response = response.message
# Send initialized notification
await send_stream1.send(
SessionMessage(
JSONRPCMessage(
root=JSONRPCNotification(
jsonrpc="2.0",
method="notifications/initialized",
)
)
)
)
await send_stream1.send(SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized")))
# Call the tool to verify lifespan context
await send_stream1.send(
SessionMessage(
JSONRPCMessage(
root=JSONRPCRequest(
jsonrpc="2.0",
id=2,
method="tools/call",
params={"name": "check_lifespan", "arguments": {}},
)
JSONRPCRequest(
jsonrpc="2.0",
id=2,
method="tools/call",
params={"name": "check_lifespan", "arguments": {}},
)
)
)
@ -125,8 +112,8 @@ async def test_lowlevel_server_lifespan():
response = await receive_stream2.receive()
response = response.message
assert isinstance(response, JSONRPCMessage)
assert isinstance(response.root, JSONRPCResponse)
assert response.root.result["content"][0]["text"] == "true"
assert isinstance(response, JSONRPCResponse)
assert response.result["content"][0]["text"] == "true"
# Cancel server task
tg.cancel_scope.cancel()
@ -162,13 +149,7 @@ async def test_fastmcp_server_lifespan():
return True
# Run server in background task
async with (
anyio.create_task_group() as tg,
send_stream1,
receive_stream1,
send_stream2,
receive_stream2,
):
async with anyio.create_task_group() as tg, send_stream1, receive_stream1, send_stream2, receive_stream2:
async def run_server():
await server._mcp_server.run(
@ -188,13 +169,11 @@ async def test_fastmcp_server_lifespan():
)
await send_stream1.send(
SessionMessage(
JSONRPCMessage(
root=JSONRPCRequest(
jsonrpc="2.0",
id=1,
method="initialize",
params=TypeAdapter(InitializeRequestParams).dump_python(params),
)
JSONRPCRequest(
jsonrpc="2.0",
id=1,
method="initialize",
params=TypeAdapter(InitializeRequestParams).dump_python(params),
)
)
)
@ -202,27 +181,16 @@ async def test_fastmcp_server_lifespan():
response = response.message
# Send initialized notification
await send_stream1.send(
SessionMessage(
JSONRPCMessage(
root=JSONRPCNotification(
jsonrpc="2.0",
method="notifications/initialized",
)
)
)
)
await send_stream1.send(SessionMessage(JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized")))
# Call the tool to verify lifespan context
await send_stream1.send(
SessionMessage(
JSONRPCMessage(
root=JSONRPCRequest(
jsonrpc="2.0",
id=2,
method="tools/call",
params={"name": "check_lifespan", "arguments": {}},
)
JSONRPCRequest(
jsonrpc="2.0",
id=2,
method="tools/call",
params={"name": "check_lifespan", "arguments": {}},
)
)
)
@ -231,8 +199,8 @@ async def test_fastmcp_server_lifespan():
response = await receive_stream2.receive()
response = response.message
assert isinstance(response, JSONRPCMessage)
assert isinstance(response.root, JSONRPCResponse)
assert response.root.result["content"][0]["text"] == "true"
assert isinstance(response, JSONRPCResponse)
assert response.result["content"][0]["text"] == "true"
# Cancel server task
tg.cancel_scope.cancel()

View file

@ -169,25 +169,23 @@ async def test_server_session_initialize_with_older_protocol_version():
# Send initialization request with older protocol version (2024-11-05)
await client_to_server_send.send(
SessionMessage(
types.JSONRPCMessage(
types.JSONRPCRequest(
jsonrpc="2.0",
id=1,
method="initialize",
params=types.InitializeRequestParams(
protocol_version="2024-11-05",
capabilities=types.ClientCapabilities(),
client_info=types.Implementation(name="test-client", version="1.0.0"),
).model_dump(by_alias=True, mode="json", exclude_none=True),
)
types.JSONRPCRequest(
jsonrpc="2.0",
id=1,
method="initialize",
params=types.InitializeRequestParams(
protocol_version="2024-11-05",
capabilities=types.ClientCapabilities(),
client_info=types.Implementation(name="test-client", version="1.0.0"),
).model_dump(by_alias=True, mode="json", exclude_none=True),
)
)
)
# Wait for the initialize response
init_response_message = await server_to_client_receive.receive()
assert isinstance(init_response_message.message.root, types.JSONRPCResponse)
result_data = init_response_message.message.root.result
assert isinstance(init_response_message.message, types.JSONRPCResponse)
result_data = init_response_message.message.result
init_result = types.InitializeResult.model_validate(result_data)
# Check that the server responded with the requested protocol version
@ -196,14 +194,7 @@ async def test_server_session_initialize_with_older_protocol_version():
# Send initialized notification
await client_to_server_send.send(
SessionMessage(
types.JSONRPCMessage(
types.JSONRPCNotification(
jsonrpc="2.0",
method="notifications/initialized",
)
)
)
SessionMessage(types.JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized"))
)
async with (
@ -256,24 +247,14 @@ async def test_ping_request_before_initialization():
nonlocal ping_response_received, ping_response_id
# Send ping request before any initialization
await client_to_server_send.send(
SessionMessage(
types.JSONRPCMessage(
types.JSONRPCRequest(
jsonrpc="2.0",
id=42,
method="ping",
)
)
)
)
await client_to_server_send.send(SessionMessage(types.JSONRPCRequest(jsonrpc="2.0", id=42, method="ping")))
# Wait for the ping response
ping_response_message = await server_to_client_receive.receive()
assert isinstance(ping_response_message.message.root, types.JSONRPCResponse)
assert isinstance(ping_response_message.message, types.JSONRPCResponse)
ping_response_received = True
ping_response_id = ping_response_message.message.root.id
ping_response_id = ping_response_message.message.id
async with (
client_to_server_send,
@ -493,22 +474,14 @@ async def test_other_requests_blocked_before_initialization():
# Try to send a non-ping request before initialization
await client_to_server_send.send(
SessionMessage(
types.JSONRPCMessage(
types.JSONRPCRequest(
jsonrpc="2.0",
id=1,
method="prompts/list",
)
)
)
SessionMessage(types.JSONRPCRequest(jsonrpc="2.0", id=1, method="prompts/list"))
)
# Wait for the error response
error_message = await server_to_client_receive.receive()
if isinstance(error_message.message.root, types.JSONRPCError): # pragma: no branch
if isinstance(error_message.message, types.JSONRPCError): # pragma: no branch
error_response_received = True
error_code = error_message.message.root.error.code
error_code = error_message.message.error.code
async with (
client_to_server_send,

View file

@ -87,54 +87,35 @@ async def test_request_immediately_after_initialize_response():
# Step 1: Send InitializeRequest
await client_to_server_send.send(
SessionMessage(
types.JSONRPCMessage(
types.JSONRPCRequest(
jsonrpc="2.0",
id=1,
method="initialize",
params=types.InitializeRequestParams(
protocol_version=types.LATEST_PROTOCOL_VERSION,
capabilities=types.ClientCapabilities(),
client_info=types.Implementation(name="test-client", version="1.0.0"),
).model_dump(by_alias=True, mode="json", exclude_none=True),
)
types.JSONRPCRequest(
jsonrpc="2.0",
id=1,
method="initialize",
params=types.InitializeRequestParams(
protocol_version=types.LATEST_PROTOCOL_VERSION,
capabilities=types.ClientCapabilities(),
client_info=types.Implementation(name="test-client", version="1.0.0"),
).model_dump(by_alias=True, mode="json", exclude_none=True),
)
)
)
# Step 2: Wait for InitializeResult
init_msg = await server_to_client_receive.receive()
assert isinstance(init_msg.message.root, types.JSONRPCResponse)
assert isinstance(init_msg.message, types.JSONRPCResponse)
# Step 3: Immediately send tools/list BEFORE InitializedNotification
# This is the race condition scenario
await client_to_server_send.send(
SessionMessage(
types.JSONRPCMessage(
types.JSONRPCRequest(
jsonrpc="2.0",
id=2,
method="tools/list",
)
)
)
)
await client_to_server_send.send(SessionMessage(types.JSONRPCRequest(jsonrpc="2.0", id=2, method="tools/list")))
# Step 4: Check the response
tools_msg = await server_to_client_receive.receive()
if isinstance(tools_msg.message.root, types.JSONRPCError): # pragma: no cover
error_received = tools_msg.message.root.error.message
if isinstance(tools_msg.message, types.JSONRPCError): # pragma: no cover
error_received = tools_msg.message.error.message
# Step 5: Send InitializedNotification
await client_to_server_send.send(
SessionMessage(
types.JSONRPCMessage(
types.JSONRPCNotification(
jsonrpc="2.0",
method="notifications/initialized",
)
)
)
SessionMessage(types.JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized"))
)
async with (

View file

@ -5,7 +5,7 @@ import pytest
from mcp.server.stdio import stdio_server
from mcp.shared.message import SessionMessage
from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse, jsonrpc_message_adapter
@pytest.mark.anyio
@ -14,8 +14,8 @@ async def test_stdio_server():
stdout = io.StringIO()
messages = [
JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")),
JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})),
JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"),
JSONRPCResponse(jsonrpc="2.0", id=2, result={}),
]
for message in messages:
@ -37,13 +37,13 @@ async def test_stdio_server():
# Verify received messages
assert len(received_messages) == 2
assert received_messages[0] == JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"))
assert received_messages[1] == JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={}))
assert received_messages[0] == JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
assert received_messages[1] == JSONRPCResponse(jsonrpc="2.0", id=2, result={})
# Test sending responses from the server
responses = [
JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")),
JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=4, result={})),
JSONRPCRequest(jsonrpc="2.0", id=3, method="ping"),
JSONRPCResponse(jsonrpc="2.0", id=4, result={}),
]
async with write_stream:
@ -55,7 +55,7 @@ async def test_stdio_server():
output_lines = stdout.readlines()
assert len(output_lines) == 2
received_responses = [JSONRPCMessage.model_validate_json(line.strip()) for line in output_lines]
received_responses = [jsonrpc_message_adapter.validate_json(line.strip()) for line in output_lines]
assert len(received_responses) == 2
assert received_responses[0] == JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping"))
assert received_responses[1] == JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=4, result={}))
assert received_responses[0] == JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")
assert received_responses[1] == JSONRPCResponse(jsonrpc="2.0", id=4, result={})