modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Drop RootModel from JSONRPCMessage (#1908)

Browse source
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
11 changed files with 314 additions and 476 deletions
Download patch file Download diff file Expand all files Collapse all files

6
src/mcp/client/sse.py
View file

@ -73,9 +73,7 @@ async def sse_client(
event_source.response.raise_for_status()
logger.debug("SSE connection established")
async def sse_reader(
task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED,
):
async def sse_reader(task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED):
try:
async for sse in event_source.aiter_sse(): # pragma: no branch
logger.debug(f"Received SSE event: {sse.event}")
@ -108,7 +106,7 @@ async def sse_client(
if not sse.data:
continue
try:
message = types.JSONRPCMessage.model_validate_json( # noqa: E501
message = types.jsonrpc_message_adapter.validate_json(
sse.data, by_name=False
)
logger.debug(f"Received server message: {message}")

2
src/mcp/client/stdio/__init__.py
View file

@ -150,7 +150,7 @@ async def stdio_client(server: StdioServerParameters, errlog: TextIO = sys.stder
for line in lines:
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
logger.exception("Failed to parse JSONRPC message from server")
await read_stream_writer.send(exc)

44
src/mcp/client/streamable_http.py
View file

@ -25,6 +25,7 @@ from mcp.types import (
JSONRPCRequest,
JSONRPCResponse,
RequestId,
jsonrpc_message_adapter,
)
logger = logging.getLogger(__name__)
@ -95,11 +96,11 @@ class StreamableHTTPTransport:
def _is_initialization_request(self, message: JSONRPCMessage) -> bool:
"""Check if the message is an initialization request."""
return isinstance(message.root, JSONRPCRequest) and message.root.method == "initialize"
return isinstance(message, JSONRPCRequest) and message.method == "initialize"
def _is_initialized_notification(self, message: JSONRPCMessage) -> bool:
"""Check if the message is an initialized notification."""
return isinstance(message.root, JSONRPCNotification) and message.root.method == "notifications/initialized"
return isinstance(message, JSONRPCNotification) and message.method == "notifications/initialized"
def _maybe_extract_session_id_from_response(self, response: httpx.Response) -> None:
"""Extract and store session ID from response headers."""
@ -110,15 +111,15 @@ class StreamableHTTPTransport:
def _maybe_extract_protocol_version_from_message(self, message: JSONRPCMessage) -> None:
"""Extract protocol version from initialization response message."""
if isinstance(message.root, JSONRPCResponse) and message.root.result: # pragma: no branch
if isinstance(message, JSONRPCResponse) and message.result: # pragma: no branch
try:
# Parse the result as InitializeResult for type safety
init_result = InitializeResult.model_validate(message.root.result, by_name=False)
init_result = InitializeResult.model_validate(message.result, by_name=False)
self.protocol_version = str(init_result.protocol_version)
logger.info(f"Negotiated protocol version: {self.protocol_version}")
except Exception: # pragma: no cover
logger.warning("Failed to parse initialization response as InitializeResult", exc_info=True)
logger.warning(f"Raw result: {message.root.result}")
logger.warning(f"Raw result: {message.result}")
async def _handle_sse_event(
self,
@ -137,7 +138,7 @@ class StreamableHTTPTransport:
await resumption_callback(sse.id)
return False
try:
message = JSONRPCMessage.model_validate_json(sse.data, by_name=False)
message = jsonrpc_message_adapter.validate_json(sse.data, by_name=False)
logger.debug(f"SSE message: {message}")
# Extract protocol version from initialization response
@ -145,8 +146,8 @@ class StreamableHTTPTransport:
self._maybe_extract_protocol_version_from_message(message)
# If this is a response and we have original_request_id, replace it
if original_request_id is not None and isinstance(message.root, JSONRPCResponse | JSONRPCError):
message.root.id = original_request_id
if original_request_id is not None and isinstance(message, JSONRPCResponse | JSONRPCError):
message.id = original_request_id
session_message = SessionMessage(message)
await read_stream_writer.send(session_message)
@ -157,7 +158,7 @@ class StreamableHTTPTransport:
# If this is a response or error return True indicating completion
# Otherwise, return False to continue listening
return isinstance(message.root, JSONRPCResponse | JSONRPCError)
return isinstance(message, JSONRPCResponse | JSONRPCError)
except Exception as exc: # pragma: no cover
logger.exception("Error parsing SSE message")
@ -222,8 +223,8 @@ class StreamableHTTPTransport:
# Extract original request ID to map responses
original_request_id = None
if isinstance(ctx.session_message.message.root, JSONRPCRequest): # pragma: no branch
original_request_id = ctx.session_message.message.root.id
if isinstance(ctx.session_message.message, JSONRPCRequest): # pragma: no branch
original_request_id = ctx.session_message.message.id
async with aconnect_sse(ctx.client, "GET", self.url, headers=headers) as event_source:
event_source.response.raise_for_status()
@ -257,12 +258,9 @@ class StreamableHTTPTransport:
return
if response.status_code == 404: # pragma: no branch
if isinstance(message.root, JSONRPCRequest):
await self._send_session_terminated_error( # pragma: no cover
ctx.read_stream_writer, # pragma: no cover
message.root.id, # pragma: no cover
) # pragma: no cover
return # pragma: no cover
if isinstance(message, JSONRPCRequest): # pragma: no branch
await self._send_session_terminated_error(ctx.read_stream_writer, message.id)
return
response.raise_for_status()
if is_initialization:
@ -270,7 +268,7 @@ class StreamableHTTPTransport:
# Per https://modelcontextprotocol.io/specification/2025-06-18/basic#notifications:
# The server MUST NOT send a response to notifications.
if isinstance(message.root, JSONRPCRequest):
if isinstance(message, JSONRPCRequest):
content_type = response.headers.get("content-type", "").lower()
if content_type.startswith("application/json"):
await self._handle_json_response(response, ctx.read_stream_writer, is_initialization)
@ -291,7 +289,7 @@ class StreamableHTTPTransport:
"""Handle JSON response from the server."""
try:
content = await response.aread()
message = JSONRPCMessage.model_validate_json(content, by_name=False)
message = jsonrpc_message_adapter.validate_json(content, by_name=False)
# Extract protocol version from initialization response
if is_initialization:
@ -365,8 +363,8 @@ class StreamableHTTPTransport:
# Extract original request ID to map responses
original_request_id = None
if isinstance(ctx.session_message.message.root, JSONRPCRequest): # pragma: no branch
original_request_id = ctx.session_message.message.root.id
if isinstance(ctx.session_message.message, JSONRPCRequest): # pragma: no branch
original_request_id = ctx.session_message.message.id
try:
async with aconnect_sse(ctx.client, "GET", self.url, headers=headers) as event_source:
@ -416,7 +414,7 @@ class StreamableHTTPTransport:
id=request_id,
error=ErrorData(code=32600, message="Session terminated"),
)
session_message = SessionMessage(JSONRPCMessage(jsonrpc_error))
session_message = SessionMessage(jsonrpc_error)
await read_stream_writer.send(session_message)
async def post_writer(
@ -463,7 +461,7 @@ class StreamableHTTPTransport:
await self._handle_post_request(ctx)
# If this is a request, start a new task to handle it
if isinstance(message.root, JSONRPCRequest):
if isinstance(message, JSONRPCRequest):
tg.start_soon(handle_request_async)
else:
await handle_request_async()

2
src/mcp/client/websocket.py
View file

@ -51,7 +51,7 @@ async def websocket_client(
async with read_stream_writer:
async for raw_text in ws:
try:
message = types.JSONRPCMessage.model_validate_json(raw_text, by_name=False)
message = types.jsonrpc_message_adapter.validate_json(raw_text, by_name=False)
session_message = SessionMessage(message)
await read_stream_writer.send(session_message)
except ValidationError as exc: # pragma: no cover

10
src/mcp/server/experimental/task_result_handler.py
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)

2
src/mcp/server/sse.py
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")

2
src/mcp/server/stdio.py
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

36
src/mcp/server/streamable_http.py
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."""
)

2
src/mcp/server/websocket.py
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

44
src/mcp/shared/session.py
View file

@ -24,7 +24,6 @@ from mcp.types import (
ClientResult,
ErrorData,
JSONRPCError,
JSONRPCMessage,
JSONRPCNotification,
JSONRPCRequest,
JSONRPCResponse,
@ -271,7 +270,7 @@ class BaseSession(
**request_data,
)
await self._write_stream.send(SessionMessage(message=JSONRPCMessage(jsonrpc_request), metadata=metadata))
await self._write_stream.send(SessionMessage(message=jsonrpc_request, metadata=metadata))
# request read timeout takes precedence over session read timeout
timeout = None
@ -321,7 +320,7 @@ class BaseSession(
**notification.model_dump(by_alias=True, mode="json", exclude_none=True),
)
session_message = SessionMessage( # pragma: no cover
message=JSONRPCMessage(jsonrpc_notification),
message=jsonrpc_notification,
metadata=ServerMessageMetadata(related_request_id=related_request_id) if related_request_id else None,
)
await self._write_stream.send(session_message)
@ -329,7 +328,7 @@ class BaseSession(
async def _send_response(self, request_id: RequestId, response: SendResultT | ErrorData) -> None:
if isinstance(response, ErrorData):
jsonrpc_error = JSONRPCError(jsonrpc="2.0", id=request_id, error=response)
session_message = SessionMessage(message=JSONRPCMessage(jsonrpc_error))
session_message = SessionMessage(message=jsonrpc_error)
await self._write_stream.send(session_message)
else:
jsonrpc_response = JSONRPCResponse(
@ -337,7 +336,7 @@ class BaseSession(
id=request_id,
result=response.model_dump(by_alias=True, mode="json", exclude_none=True),
)
session_message = SessionMessage(message=JSONRPCMessage(jsonrpc_response))
session_message = SessionMessage(message=jsonrpc_response)
await self._write_stream.send(session_message)
async def _receive_loop(self) -> None:
@ -349,14 +348,14 @@ class BaseSession(
async for message in self._read_stream:
if isinstance(message, Exception): # pragma: no cover
await self._handle_incoming(message)
elif isinstance(message.message.root, JSONRPCRequest):
elif isinstance(message.message, JSONRPCRequest):
try:
validated_request = self._receive_request_type.model_validate(
message.message.root.model_dump(by_alias=True, mode="json", exclude_none=True),
message.message.model_dump(by_alias=True, mode="json", exclude_none=True),
by_name=False,
)
responder = RequestResponder(
request_id=message.message.root.id,
request_id=message.message.id,
request_meta=validated_request.root.params.meta
if validated_request.root.params
else None,
@ -374,23 +373,23 @@ class BaseSession(
# For request validation errors, send a proper JSON-RPC error
# response instead of crashing the server
logging.warning(f"Failed to validate request: {e}")
logging.debug(f"Message that failed validation: {message.message.root}")
logging.debug(f"Message that failed validation: {message.message}")
error_response = JSONRPCError(
jsonrpc="2.0",
id=message.message.root.id,
id=message.message.id,
error=ErrorData(
code=INVALID_PARAMS,
message="Invalid request parameters",
data="",
),
)
session_message = SessionMessage(message=JSONRPCMessage(error_response))
session_message = SessionMessage(message=error_response)
await self._write_stream.send(session_message)
elif isinstance(message.message.root, JSONRPCNotification):
elif isinstance(message.message, JSONRPCNotification):
try:
notification = self._receive_notification_type.model_validate(
message.message.root.model_dump(by_alias=True, mode="json", exclude_none=True),
message.message.model_dump(by_alias=True, mode="json", exclude_none=True),
by_name=False,
)
# Handle cancellation notifications
@ -419,10 +418,11 @@ class BaseSession(
)
await self._received_notification(notification)
await self._handle_incoming(notification)
except Exception as e: # pragma: no cover
except Exception: # pragma: no cover
# For other validation errors, log and continue
logging.warning(
f"Failed to validate notification: {e}. Message was: {message.message.root}"
f"Failed to validate notification:. Message was: {message.message}",
exc_info=True,
)
else: # Response or error
await self._handle_response(message)
@ -475,27 +475,25 @@ class BaseSession(
Checks response routers first (e.g., for task-related responses),
then falls back to the normal response stream mechanism.
"""
root = message.message.root
# This check is always true at runtime: the caller (_receive_loop) only invokes
# this method in the else branch after checking for JSONRPCRequest and
# JSONRPCNotification. However, the type checker can't infer this from the
# method signature, so we need this guard for type narrowing.
if not isinstance(root, JSONRPCResponse | JSONRPCError):
if not isinstance(message.message, JSONRPCResponse | JSONRPCError):
return # pragma: no cover
# Normalize response ID to handle type mismatches (e.g., "0" vs 0)
response_id = self._normalize_request_id(root.id)
response_id = self._normalize_request_id(message.message.id)
# First, check response routers (e.g., TaskResultHandler)
if isinstance(root, JSONRPCError):
if isinstance(message.message, JSONRPCError):
# Route error to routers
for router in self._response_routers:
if router.route_error(response_id, root.error):
if router.route_error(response_id, message.message.error):
return # Handled
else:
# Route success response to routers
response_data: dict[str, Any] = root.result or {}
response_data: dict[str, Any] = message.message.result or {}
for router in self._response_routers:
if router.route_response(response_id, response_data):
return # Handled
@ -503,7 +501,7 @@ class BaseSession(
# Fall back to normal response streams
stream = self._response_streams.pop(response_id, None)
if stream: # pragma: no cover
await stream.send(root)
await stream.send(message.message)
else: # pragma: no cover
await self._handle_incoming(RuntimeError(f"Received response with an unknown request ID: {message}"))

6
src/mcp/types.py
View file

@ -4,7 +4,7 @@ from collections.abc import Callable
from datetime import datetime
from typing import Annotated, Any, Final, Generic, Literal, TypeAlias, TypeVar
from pydantic import BaseModel, ConfigDict, Field, FileUrl, RootModel
from pydantic import BaseModel, ConfigDict, Field, FileUrl, RootModel, TypeAdapter
from pydantic.alias_generators import to_camel
LATEST_PROTOCOL_VERSION = "2025-11-25"
@ -197,8 +197,8 @@ class JSONRPCError(MCPModel):
error: ErrorData
class JSONRPCMessage(RootModel[JSONRPCRequest | JSONRPCNotification | JSONRPCResponse | JSONRPCError]):
pass
JSONRPCMessage = JSONRPCRequest | JSONRPCNotification | JSONRPCResponse | JSONRPCError
jsonrpc_message_adapter = TypeAdapter[JSONRPCMessage](JSONRPCMessage)
class EmptyResult(Result):