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
2 changed files with 314 additions and 476 deletions

View file

@ -66,7 +66,7 @@ async def test_request_id_match() -> None:
jsonrpc="2.0",
)
await client_writer.send(SessionMessage(JSONRPCMessage(root=init_req)))
await client_writer.send(SessionMessage(init_req))
response = await server_reader.receive() # Get init response but don't need to check it
# Send initialized notification
@ -75,12 +75,12 @@ async def test_request_id_match() -> None:
params=NotificationParams().model_dump(by_alias=True, exclude_none=True),
jsonrpc="2.0",
)
await client_writer.send(SessionMessage(JSONRPCMessage(root=initialized_notification)))
await client_writer.send(SessionMessage(initialized_notification))
# Send ping request with custom ID
ping_request = JSONRPCRequest(id=custom_request_id, method="ping", params={}, jsonrpc="2.0")
await client_writer.send(SessionMessage(JSONRPCMessage(root=ping_request)))
await client_writer.send(SessionMessage(ping_request))
# Read response
response = await server_reader.receive()
@ -88,8 +88,8 @@ async def test_request_id_match() -> None:
# Verify response ID matches request ID
assert isinstance(response, SessionMessage)
assert isinstance(response.message, JSONRPCMessage)
assert isinstance(response.message.root, JSONRPCResponse)
assert response.message.root.id == custom_request_id, "Response ID should match request ID"
assert isinstance(response.message, JSONRPCResponse)
assert response.message.id == custom_request_id, "Response ID should match request ID"
# Cancel server task
tg.cancel_scope.cancel()

View file

@ -1,21 +1,13 @@
# Claude Debug
"""Test for HackerOne vulnerability report #3156202 - malformed input DOS."""
from typing import Any
import anyio
import pytest
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
from mcp.shared.message import SessionMessage
from mcp.types import (
INVALID_PARAMS,
JSONRPCError,
JSONRPCMessage,
JSONRPCRequest,
ServerCapabilities,
)
from mcp.types import INVALID_PARAMS, JSONRPCError, JSONRPCMessage, JSONRPCRequest, ServerCapabilities
@pytest.mark.anyio
@ -37,7 +29,7 @@ async def test_malformed_initialize_request_does_not_crash_server():
)
# Wrap in session message
request_message = SessionMessage(message=JSONRPCMessage(malformed_request))
request_message = SessionMessage(message=malformed_request)
# Start a server session
async with ServerSession(
@ -58,7 +50,7 @@ async def test_malformed_initialize_request_does_not_crash_server():
# Check that we received an error response instead of a crash
try:
response_message = write_receive_stream.receive_nowait()
response = response_message.message.root
response = response_message.message
# Verify it's a proper JSON-RPC error response
assert isinstance(response, JSONRPCError)
@ -75,14 +67,14 @@ async def test_malformed_initialize_request_does_not_crash_server():
method="tools/call",
# params=None # Missing required params
)
another_request_message = SessionMessage(message=JSONRPCMessage(another_malformed_request))
another_request_message = SessionMessage(message=another_malformed_request)
await read_send_stream.send(another_request_message)
await anyio.sleep(0.1)
# Should get another error response, not a crash
second_response_message = write_receive_stream.receive_nowait()
second_response = second_response_message.message.root
second_response = second_response_message.message
assert isinstance(second_response, JSONRPCError)
assert second_response.id == "test_id_2"
@ -125,7 +117,7 @@ async def test_multiple_concurrent_malformed_requests():
method="initialize",
# params=None # Missing required params
)
request_message = SessionMessage(message=JSONRPCMessage(malformed_request))
request_message = SessionMessage(message=malformed_request)
malformed_requests.append(request_message)
# Send all requests
@ -136,11 +128,11 @@ async def test_multiple_concurrent_malformed_requests():
await anyio.sleep(0.2)
# Verify we get error responses for all requests
error_responses: list[Any] = []
error_responses: list[JSONRPCMessage] = []
try:
while True:
response_message = write_receive_stream.receive_nowait()
error_responses.append(response_message.message.root)
error_responses.append(response_message.message)
except anyio.WouldBlock:
pass # No more messages