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

View file

@ -26,7 +26,6 @@ from mcp.types import (
ErrorData,
GetTaskPayloadRequest,
GetTaskPayloadResult,
JSONRPCMessage,
RelatedTaskMetadata,
RequestId,
)
@ -107,12 +106,7 @@ class TaskResultHandler:
while True:
task = await self._store.get_task(task_id)
if task is None:
raise McpError(
ErrorData(
code=INVALID_PARAMS,
message=f"Task not found: {task_id}",
)
)
raise McpError(ErrorData(code=INVALID_PARAMS, message=f"Task not found: {task_id}"))
await self._deliver_queued_messages(task_id, session, request_id)
@ -161,7 +155,7 @@ class TaskResultHandler:
# Send the message with relatedRequestId for routing
session_message = SessionMessage(
message=JSONRPCMessage(message.message),
message=message.message,
metadata=ServerMessageMetadata(related_request_id=request_id),
)
await self.send_message(session, session_message)

View file

@ -227,7 +227,7 @@ class SseServerTransport:
logger.debug(f"Received JSON: {body}")
try:
message = types.JSONRPCMessage.model_validate_json(body, by_name=False)
message = types.jsonrpc_message_adapter.validate_json(body, by_name=False)
logger.debug(f"Validated client message: {message}")
except ValidationError as err:
logger.exception("Failed to parse message")

View file

@ -60,7 +60,7 @@ async def stdio_server(
async with read_stream_writer:
async for line in stdin:
try:
message = types.JSONRPCMessage.model_validate_json(line, by_name=False)
message = types.jsonrpc_message_adapter.validate_json(line, by_name=False)
except Exception as exc: # pragma: no cover
await read_stream_writer.send(exc)
continue

View file

@ -42,6 +42,7 @@ from mcp.types import (
JSONRPCRequest,
JSONRPCResponse,
RequestId,
jsonrpc_message_adapter,
)
logger = logging.getLogger(__name__)
@ -301,10 +302,7 @@ class StreamableHTTPServerTransport:
error_response = JSONRPCError(
jsonrpc="2.0",
id="server-error", # We don't have a request ID for general errors
error=ErrorData(
code=error_code,
message=error_message,
),
error=ErrorData(code=error_code, message=error_message),
)
return Response(
@ -455,6 +453,7 @@ class StreamableHTTPServerTransport:
body = await request.body()
try:
# TODO(Marcelo): Replace `json.loads` with `pydantic_core.from_json`.
raw_message = json.loads(body)
except json.JSONDecodeError as e:
response = self._create_error_response(f"Parse error: {str(e)}", HTTPStatus.BAD_REQUEST, PARSE_ERROR)
@ -462,7 +461,7 @@ class StreamableHTTPServerTransport:
return
try: # pragma: no cover
message = JSONRPCMessage.model_validate(raw_message, by_name=False)
message = jsonrpc_message_adapter.validate_python(raw_message, by_name=False)
except ValidationError as e: # pragma: no cover
response = self._create_error_response(
f"Validation error: {str(e)}",
@ -473,9 +472,7 @@ class StreamableHTTPServerTransport:
return
# Check if this is an initialization request
is_initialization_request = (
isinstance(message.root, JSONRPCRequest) and message.root.method == "initialize"
) # pragma: no cover
is_initialization_request = isinstance(message, JSONRPCRequest) and message.method == "initialize"
if is_initialization_request: # pragma: no cover
# Check if the server already has an established session
@ -495,7 +492,7 @@ class StreamableHTTPServerTransport:
return
# For notifications and responses only, return 202 Accepted
if not isinstance(message.root, JSONRPCRequest): # pragma: no cover
if not isinstance(message, JSONRPCRequest): # pragma: no cover
# Create response object and send it
response = self._create_json_response(
None,
@ -514,13 +511,13 @@ class StreamableHTTPServerTransport:
# For initialize requests, get from request params.
# For other requests, get from header (already validated).
protocol_version = (
str(message.root.params.get("protocolVersion", DEFAULT_NEGOTIATED_VERSION))
if is_initialization_request and message.root.params
str(message.params.get("protocolVersion", DEFAULT_NEGOTIATED_VERSION))
if is_initialization_request and message.params
else request.headers.get(MCP_PROTOCOL_VERSION_HEADER, DEFAULT_NEGOTIATED_VERSION)
)
# Extract the request ID outside the try block for proper scope
request_id = str(message.root.id) # pragma: no cover
request_id = str(message.id) # pragma: no cover
# Register this stream for the request ID
self._request_streams[request_id] = anyio.create_memory_object_stream[EventMessage](0) # pragma: no cover
request_stream_reader = self._request_streams[request_id][1] # pragma: no cover
@ -538,12 +535,12 @@ class StreamableHTTPServerTransport:
# Use similar approach to SSE writer for consistency
async for event_message in request_stream_reader:
# If it's a response, this is what we're waiting for
if isinstance(event_message.message.root, JSONRPCResponse | JSONRPCError):
if isinstance(event_message.message, JSONRPCResponse | JSONRPCError):
response_message = event_message.message
break
# For notifications and request, keep waiting
else:
logger.debug(f"received: {event_message.message.root.method}")
logger.debug(f"received: {event_message.message.method}")
# At this point we should have a response
if response_message:
@ -589,10 +586,7 @@ class StreamableHTTPServerTransport:
await sse_stream_writer.send(event_data)
# If response, remove from pending streams and close
if isinstance(
event_message.message.root,
JSONRPCResponse | JSONRPCError,
):
if isinstance(event_message.message, JSONRPCResponse | JSONRPCError):
break
except anyio.ClosedResourceError:
# Expected when close_sse_stream() is called
@ -984,8 +978,8 @@ class StreamableHTTPServerTransport:
message = session_message.message
target_request_id = None
# Check if this is a response
if isinstance(message.root, JSONRPCResponse | JSONRPCError):
response_id = str(message.root.id)
if isinstance(message, JSONRPCResponse | JSONRPCError):
response_id = str(message.id)
# If this response is for an existing request stream,
# send it there
target_request_id = response_id
@ -1022,7 +1016,7 @@ class StreamableHTTPServerTransport:
self._request_streams.pop(request_stream_id, None)
else:
logger.debug(
f"""Request stream {request_stream_id} not found
f"""Request stream {request_stream_id} not found
for message. Still processing message as the client
might reconnect and replay."""
)

View file

@ -36,7 +36,7 @@ async def websocket_server(scope: Scope, receive: Receive, send: Send):
async with read_stream_writer:
async for msg in websocket.iter_text():
try:
client_message = types.JSONRPCMessage.model_validate_json(msg, by_name=False)
client_message = types.jsonrpc_message_adapter.validate_json(msg, by_name=False)
except ValidationError as exc:
await read_stream_writer.send(exc)
continue