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
27 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() event_source.response.raise_for_status()
logger.debug("SSE connection established") logger.debug("SSE connection established")
async def sse_reader( async def sse_reader(task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED):
task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED,
):
try: try:
async for sse in event_source.aiter_sse(): # pragma: no branch async for sse in event_source.aiter_sse(): # pragma: no branch
logger.debug(f"Received SSE event: {sse.event}") logger.debug(f"Received SSE event: {sse.event}")
@ -108,7 +106,7 @@ async def sse_client(
if not sse.data: if not sse.data:
continue continue
try: try:
message = types.JSONRPCMessage.model_validate_json( # noqa: E501 message = types.jsonrpc_message_adapter.validate_json(
sse.data, by_name=False sse.data, by_name=False
) )
logger.debug(f"Received server message: {message}") 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: for line in lines:
try: 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 except Exception as exc: # pragma: no cover
logger.exception("Failed to parse JSONRPC message from server") logger.exception("Failed to parse JSONRPC message from server")
await read_stream_writer.send(exc) await read_stream_writer.send(exc)

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

@ -25,6 +25,7 @@ from mcp.types import (
JSONRPCRequest, JSONRPCRequest,
JSONRPCResponse, JSONRPCResponse,
RequestId, RequestId,
jsonrpc_message_adapter,
) )
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -95,11 +96,11 @@ class StreamableHTTPTransport:
def _is_initialization_request(self, message: JSONRPCMessage) -> bool: def _is_initialization_request(self, message: JSONRPCMessage) -> bool:
"""Check if the message is an initialization request.""" """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: def _is_initialized_notification(self, message: JSONRPCMessage) -> bool:
"""Check if the message is an initialized notification.""" """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: def _maybe_extract_session_id_from_response(self, response: httpx.Response) -> None:
"""Extract and store session ID from response headers.""" """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: def _maybe_extract_protocol_version_from_message(self, message: JSONRPCMessage) -> None:
"""Extract protocol version from initialization response message.""" """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: try:
# Parse the result as InitializeResult for type safety # 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) self.protocol_version = str(init_result.protocol_version)
logger.info(f"Negotiated protocol version: {self.protocol_version}") logger.info(f"Negotiated protocol version: {self.protocol_version}")
except Exception: # pragma: no cover except Exception: # pragma: no cover
logger.warning("Failed to parse initialization response as InitializeResult", exc_info=True) 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( async def _handle_sse_event(
self, self,
@ -137,7 +138,7 @@ class StreamableHTTPTransport:
await resumption_callback(sse.id) await resumption_callback(sse.id)
return False return False
try: 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}") logger.debug(f"SSE message: {message}")
# Extract protocol version from initialization response # Extract protocol version from initialization response
@ -145,8 +146,8 @@ class StreamableHTTPTransport:
self._maybe_extract_protocol_version_from_message(message) self._maybe_extract_protocol_version_from_message(message)
# If this is a response and we have original_request_id, replace it # 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): if original_request_id is not None and isinstance(message, JSONRPCResponse | JSONRPCError):
message.root.id = original_request_id message.id = original_request_id
session_message = SessionMessage(message) session_message = SessionMessage(message)
await read_stream_writer.send(session_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 # If this is a response or error return True indicating completion
# Otherwise, return False to continue listening # Otherwise, return False to continue listening
return isinstance(message.root, JSONRPCResponse | JSONRPCError) return isinstance(message, JSONRPCResponse | JSONRPCError)
except Exception as exc: # pragma: no cover except Exception as exc: # pragma: no cover
logger.exception("Error parsing SSE message") logger.exception("Error parsing SSE message")
@ -222,8 +223,8 @@ class StreamableHTTPTransport:
# Extract original request ID to map responses # Extract original request ID to map responses
original_request_id = None original_request_id = None
if isinstance(ctx.session_message.message.root, JSONRPCRequest): # pragma: no branch if isinstance(ctx.session_message.message, JSONRPCRequest): # pragma: no branch
original_request_id = ctx.session_message.message.root.id original_request_id = ctx.session_message.message.id
async with aconnect_sse(ctx.client, "GET", self.url, headers=headers) as event_source: async with aconnect_sse(ctx.client, "GET", self.url, headers=headers) as event_source:
event_source.response.raise_for_status() event_source.response.raise_for_status()
@ -257,12 +258,9 @@ class StreamableHTTPTransport:
return return
if response.status_code == 404: # pragma: no branch if response.status_code == 404: # pragma: no branch
if isinstance(message.root, JSONRPCRequest): if isinstance(message, JSONRPCRequest): # pragma: no branch
await self._send_session_terminated_error( # pragma: no cover await self._send_session_terminated_error(ctx.read_stream_writer, message.id)
ctx.read_stream_writer, # pragma: no cover return
message.root.id, # pragma: no cover
) # pragma: no cover
return # pragma: no cover
response.raise_for_status() response.raise_for_status()
if is_initialization: if is_initialization:
@ -270,7 +268,7 @@ class StreamableHTTPTransport:
# Per https://modelcontextprotocol.io/specification/2025-06-18/basic#notifications: # Per https://modelcontextprotocol.io/specification/2025-06-18/basic#notifications:
# The server MUST NOT send a response to 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() content_type = response.headers.get("content-type", "").lower()
if content_type.startswith("application/json"): if content_type.startswith("application/json"):
await self._handle_json_response(response, ctx.read_stream_writer, is_initialization) await self._handle_json_response(response, ctx.read_stream_writer, is_initialization)
@ -291,7 +289,7 @@ class StreamableHTTPTransport:
"""Handle JSON response from the server.""" """Handle JSON response from the server."""
try: try:
content = await response.aread() 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 # Extract protocol version from initialization response
if is_initialization: if is_initialization:
@ -365,8 +363,8 @@ class StreamableHTTPTransport:
# Extract original request ID to map responses # Extract original request ID to map responses
original_request_id = None original_request_id = None
if isinstance(ctx.session_message.message.root, JSONRPCRequest): # pragma: no branch if isinstance(ctx.session_message.message, JSONRPCRequest): # pragma: no branch
original_request_id = ctx.session_message.message.root.id original_request_id = ctx.session_message.message.id
try: try:
async with aconnect_sse(ctx.client, "GET", self.url, headers=headers) as event_source: async with aconnect_sse(ctx.client, "GET", self.url, headers=headers) as event_source:
@ -416,7 +414,7 @@ class StreamableHTTPTransport:
id=request_id, id=request_id,
error=ErrorData(code=32600, message="Session terminated"), 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) await read_stream_writer.send(session_message)
async def post_writer( async def post_writer(
@ -463,7 +461,7 @@ class StreamableHTTPTransport:
await self._handle_post_request(ctx) await self._handle_post_request(ctx)
# If this is a request, start a new task to handle it # 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) tg.start_soon(handle_request_async)
else: else:
await handle_request_async() 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 with read_stream_writer:
async for raw_text in ws: async for raw_text in ws:
try: 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) session_message = SessionMessage(message)
await read_stream_writer.send(session_message) await read_stream_writer.send(session_message)
except ValidationError as exc: # pragma: no cover 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, ErrorData,
GetTaskPayloadRequest, GetTaskPayloadRequest,
GetTaskPayloadResult, GetTaskPayloadResult,
JSONRPCMessage,
RelatedTaskMetadata, RelatedTaskMetadata,
RequestId, RequestId,
) )
@ -107,12 +106,7 @@ class TaskResultHandler:
while True: while True:
task = await self._store.get_task(task_id) task = await self._store.get_task(task_id)
if task is None: if task is None:
raise McpError( raise McpError(ErrorData(code=INVALID_PARAMS, message=f"Task not found: {task_id}"))
ErrorData(
code=INVALID_PARAMS,
message=f"Task not found: {task_id}",
)
)
await self._deliver_queued_messages(task_id, session, request_id) await self._deliver_queued_messages(task_id, session, request_id)
@ -161,7 +155,7 @@ class TaskResultHandler:
# Send the message with relatedRequestId for routing # Send the message with relatedRequestId for routing
session_message = SessionMessage( session_message = SessionMessage(
message=JSONRPCMessage(message.message), message=message.message,
metadata=ServerMessageMetadata(related_request_id=request_id), metadata=ServerMessageMetadata(related_request_id=request_id),
) )
await self.send_message(session, session_message) 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}") logger.debug(f"Received JSON: {body}")
try: 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}") logger.debug(f"Validated client message: {message}")
except ValidationError as err: except ValidationError as err:
logger.exception("Failed to parse message") 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 with read_stream_writer:
async for line in stdin: async for line in stdin:
try: 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 except Exception as exc: # pragma: no cover
await read_stream_writer.send(exc) await read_stream_writer.send(exc)
continue continue

34
src/mcp/server/streamable_http.py
View file

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

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 with read_stream_writer:
async for msg in websocket.iter_text(): async for msg in websocket.iter_text():
try: 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: except ValidationError as exc:
await read_stream_writer.send(exc) await read_stream_writer.send(exc)
continue continue

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

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

18
tests/client/conftest.py
View file

@ -43,35 +43,33 @@ class StreamSpyCollection:
def get_client_requests(self, method: str | None = None) -> list[JSONRPCRequest]: # pragma: no cover def get_client_requests(self, method: str | None = None) -> list[JSONRPCRequest]: # pragma: no cover
"""Get client-sent requests, optionally filtered by method.""" """Get client-sent requests, optionally filtered by method."""
return [ return [
req.message.root req.message
for req in self.client.sent_messages for req in self.client.sent_messages
if isinstance(req.message.root, JSONRPCRequest) and (method is None or req.message.root.method == method) if isinstance(req.message, JSONRPCRequest) and (method is None or req.message.method == method)
] ]
def get_server_requests(self, method: str | None = None) -> list[JSONRPCRequest]: # pragma: no cover def get_server_requests(self, method: str | None = None) -> list[JSONRPCRequest]: # pragma: no cover
"""Get server-sent requests, optionally filtered by method.""" """Get server-sent requests, optionally filtered by method."""
return [ # pragma: no cover return [ # pragma: no cover
req.message.root req.message
for req in self.server.sent_messages for req in self.server.sent_messages
if isinstance(req.message.root, JSONRPCRequest) and (method is None or req.message.root.method == method) if isinstance(req.message, JSONRPCRequest) and (method is None or req.message.method == method)
] ]
def get_client_notifications(self, method: str | None = None) -> list[JSONRPCNotification]: # pragma: no cover def get_client_notifications(self, method: str | None = None) -> list[JSONRPCNotification]: # pragma: no cover
"""Get client-sent notifications, optionally filtered by method.""" """Get client-sent notifications, optionally filtered by method."""
return [ return [
notif.message.root notif.message
for notif in self.client.sent_messages for notif in self.client.sent_messages
if isinstance(notif.message.root, JSONRPCNotification) if isinstance(notif.message, JSONRPCNotification) and (method is None or notif.message.method == method)
and (method is None or notif.message.root.method == method)
] ]
def get_server_notifications(self, method: str | None = None) -> list[JSONRPCNotification]: # pragma: no cover def get_server_notifications(self, method: str | None = None) -> list[JSONRPCNotification]: # pragma: no cover
"""Get server-sent notifications, optionally filtered by method.""" """Get server-sent notifications, optionally filtered by method."""
return [ return [
notif.message.root notif.message
for notif in self.server.sent_messages for notif in self.server.sent_messages
if isinstance(notif.message.root, JSONRPCNotification) if isinstance(notif.message, JSONRPCNotification) and (method is None or notif.message.method == method)
and (method is None or notif.message.root.method == method)
] ]

182
tests/client/test_session.py
View file

@ -18,7 +18,6 @@ from mcp.types import (
InitializedNotification, InitializedNotification,
InitializeRequest, InitializeRequest,
InitializeResult, InitializeResult,
JSONRPCMessage,
JSONRPCNotification, JSONRPCNotification,
JSONRPCRequest, JSONRPCRequest,
JSONRPCResponse, JSONRPCResponse,
@ -41,7 +40,7 @@ async def test_client_session_initialize():
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -65,18 +64,16 @@ async def test_client_session_initialize():
async with server_to_client_send: async with server_to_client_send:
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
session_notification = await client_to_server_receive.receive() session_notification = await client_to_server_receive.receive()
jsonrpc_notification = session_notification.message jsonrpc_notification = session_notification.message
assert isinstance(jsonrpc_notification.root, JSONRPCNotification) assert isinstance(jsonrpc_notification, JSONRPCNotification)
initialized_notification = ClientNotification.model_validate( initialized_notification = ClientNotification.model_validate(
jsonrpc_notification.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_notification.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -128,7 +125,7 @@ async def test_client_session_custom_client_info():
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -146,12 +143,10 @@ async def test_client_session_custom_client_info():
async with server_to_client_send: async with server_to_client_send:
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
@ -189,7 +184,7 @@ async def test_client_session_default_client_info():
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -207,12 +202,10 @@ async def test_client_session_default_client_info():
async with server_to_client_send: async with server_to_client_send:
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
@ -220,10 +213,7 @@ async def test_client_session_default_client_info():
await client_to_server_receive.receive() await client_to_server_receive.receive()
async with ( async with (
ClientSession( ClientSession(server_to_client_receive, client_to_server_send) as session,
server_to_client_receive,
client_to_server_send,
) as session,
anyio.create_task_group() as tg, anyio.create_task_group() as tg,
client_to_server_send, client_to_server_send,
client_to_server_receive, client_to_server_receive,
@ -247,7 +237,7 @@ async def test_client_session_version_negotiation_success():
async def mock_server(): async def mock_server():
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -268,12 +258,10 @@ async def test_client_session_version_negotiation_success():
async with server_to_client_send: async with server_to_client_send:
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
@ -281,10 +269,7 @@ async def test_client_session_version_negotiation_success():
await client_to_server_receive.receive() await client_to_server_receive.receive()
async with ( async with (
ClientSession( ClientSession(server_to_client_receive, client_to_server_send) as session,
server_to_client_receive,
client_to_server_send,
) as session,
anyio.create_task_group() as tg, anyio.create_task_group() as tg,
client_to_server_send, client_to_server_send,
client_to_server_receive, client_to_server_receive,
@ -309,7 +294,7 @@ async def test_client_session_version_negotiation_failure():
async def mock_server(): async def mock_server():
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -327,21 +312,16 @@ async def test_client_session_version_negotiation_failure():
async with server_to_client_send: async with server_to_client_send:
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
async with ( async with (
ClientSession( ClientSession(server_to_client_receive, client_to_server_send) as session,
server_to_client_receive,
client_to_server_send,
) as session,
anyio.create_task_group() as tg, anyio.create_task_group() as tg,
client_to_server_send, client_to_server_send,
client_to_server_receive, client_to_server_receive,
@ -368,7 +348,7 @@ async def test_client_capabilities_default():
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -386,12 +366,10 @@ async def test_client_capabilities_default():
async with server_to_client_send: async with server_to_client_send:
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
@ -399,10 +377,7 @@ async def test_client_capabilities_default():
await client_to_server_receive.receive() await client_to_server_receive.receive()
async with ( async with (
ClientSession( ClientSession(server_to_client_receive, client_to_server_send) as session,
server_to_client_receive,
client_to_server_send,
) as session,
anyio.create_task_group() as tg, anyio.create_task_group() as tg,
client_to_server_send, client_to_server_send,
client_to_server_receive, client_to_server_receive,
@ -446,7 +421,7 @@ async def test_client_capabilities_with_custom_callbacks():
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -464,12 +439,10 @@ async def test_client_capabilities_with_custom_callbacks():
async with server_to_client_send: async with server_to_client_send:
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
@ -529,7 +502,7 @@ async def test_client_capabilities_with_sampling_tools():
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -547,12 +520,10 @@ async def test_client_capabilities_with_sampling_tools():
async with server_to_client_send: async with server_to_client_send:
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
@ -600,7 +571,7 @@ async def test_get_server_capabilities():
async def mock_server(): async def mock_server():
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -617,12 +588,10 @@ async def test_get_server_capabilities():
async with server_to_client_send: async with server_to_client_send:
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
@ -669,7 +638,7 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
# Receive initialization request from client # Receive initialization request from client
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -686,12 +655,10 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
# Answer initialization request # Answer initialization request
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
@ -702,14 +669,14 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
# Wait for the client to send a 'tools/call' request # Wait for the client to send a 'tools/call' request
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
assert jsonrpc_request.root.method == "tools/call" assert jsonrpc_request.method == "tools/call"
if meta is not None: if meta is not None:
assert jsonrpc_request.root.params assert jsonrpc_request.params
assert "_meta" in jsonrpc_request.root.params assert "_meta" in jsonrpc_request.params
assert jsonrpc_request.root.params["_meta"] == meta assert jsonrpc_request.params["_meta"] == meta
result = ServerResult( result = ServerResult(
CallToolResult(content=[TextContent(type="text", text="Called successfully")], is_error=False) CallToolResult(content=[TextContent(type="text", text="Called successfully")], is_error=False)
@ -718,12 +685,10 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
# Send the tools/call result # Send the tools/call result
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
@ -732,20 +697,18 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
# The client requires this step to validate the tool output schema # The client requires this step to validate the tool output schema
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
assert jsonrpc_request.root.method == "tools/list" assert jsonrpc_request.method == "tools/list"
result = types.ListToolsResult(tools=[mocked_tool]) result = types.ListToolsResult(tools=[mocked_tool])
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
@ -753,10 +716,7 @@ async def test_client_tool_call_with_meta(meta: dict[str, Any] | None):
server_to_client_send.close() server_to_client_send.close()
async with ( async with (
ClientSession( ClientSession(server_to_client_receive, client_to_server_send) as session,
server_to_client_receive,
client_to_server_send,
) as session,
anyio.create_task_group() as tg, anyio.create_task_group() as tg,
client_to_server_send, client_to_server_send,
client_to_server_receive, client_to_server_receive,

8
tests/client/test_stdio.py
View file

@ -47,8 +47,8 @@ async def test_stdio_client():
async with stdio_client(server_parameters) as (read_stream, write_stream): async with stdio_client(server_parameters) as (read_stream, write_stream):
# Test sending and receiving messages # Test sending and receiving messages
messages = [ messages = [
JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")), JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"),
JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})), JSONRPCResponse(jsonrpc="2.0", id=2, result={}),
] ]
async with write_stream: async with write_stream:
@ -67,8 +67,8 @@ async def test_stdio_client():
break break
assert len(read_messages) == 2 assert len(read_messages) == 2
assert read_messages[0] == JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")) assert read_messages[0] == JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
assert read_messages[1] == JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})) assert read_messages[1] == JSONRPCResponse(jsonrpc="2.0", id=2, result={})
@pytest.mark.anyio @pytest.mark.anyio

49
tests/experimental/tasks/client/test_capabilities.py
View file

@ -15,7 +15,6 @@ from mcp.types import (
Implementation, Implementation,
InitializeRequest, InitializeRequest,
InitializeResult, InitializeResult,
JSONRPCMessage,
JSONRPCRequest, JSONRPCRequest,
JSONRPCResponse, JSONRPCResponse,
ServerCapabilities, ServerCapabilities,
@ -36,7 +35,7 @@ async def test_client_capabilities_without_tasks():
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -54,12 +53,10 @@ async def test_client_capabilities_without_tasks():
async with server_to_client_send: async with server_to_client_send:
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
@ -110,7 +107,7 @@ async def test_client_capabilities_with_tasks():
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -128,12 +125,10 @@ async def test_client_capabilities_with_tasks():
async with server_to_client_send: async with server_to_client_send:
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
@ -194,7 +189,7 @@ async def test_client_capabilities_auto_built_from_handlers():
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -212,12 +207,10 @@ async def test_client_capabilities_auto_built_from_handlers():
async with server_to_client_send: async with server_to_client_send:
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )
@ -274,7 +267,7 @@ async def test_client_capabilities_with_task_augmented_handlers():
session_message = await client_to_server_receive.receive() session_message = await client_to_server_receive.receive()
jsonrpc_request = session_message.message jsonrpc_request = session_message.message
assert isinstance(jsonrpc_request.root, JSONRPCRequest) assert isinstance(jsonrpc_request, JSONRPCRequest)
request = ClientRequest.model_validate( request = ClientRequest.model_validate(
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
) )
@ -292,12 +285,10 @@ async def test_client_capabilities_with_task_augmented_handlers():
async with server_to_client_send: async with server_to_client_send:
await server_to_client_send.send( await server_to_client_send.send(
SessionMessage( SessionMessage(
JSONRPCMessage( JSONRPCResponse(
JSONRPCResponse( jsonrpc="2.0",
jsonrpc="2.0", id=jsonrpc_request.id,
id=jsonrpc_request.root.id, result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
)
) )
) )
) )

70
tests/experimental/tasks/client/test_handlers.py
View file

@ -151,15 +151,11 @@ async def test_client_handles_get_task_request(client_streams: ClientTestStreams
await client_ready.wait() await client_ready.wait()
typed_request = GetTaskRequest(params=GetTaskRequestParams(task_id="test-task-123")) typed_request = GetTaskRequest(params=GetTaskRequestParams(task_id="test-task-123"))
request = types.JSONRPCRequest( request = types.JSONRPCRequest(jsonrpc="2.0", id="req-1", **typed_request.model_dump(by_alias=True))
jsonrpc="2.0", await client_streams.server_send.send(SessionMessage(request))
id="req-1",
**typed_request.model_dump(by_alias=True),
)
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
response_msg = await client_streams.server_receive.receive() response_msg = await client_streams.server_receive.receive()
response = response_msg.message.root response = response_msg.message
assert isinstance(response, types.JSONRPCResponse) assert isinstance(response, types.JSONRPCResponse)
assert response.id == "req-1" assert response.id == "req-1"
@ -219,10 +215,10 @@ async def test_client_handles_get_task_result_request(client_streams: ClientTest
id="req-2", id="req-2",
**typed_request.model_dump(by_alias=True), **typed_request.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) await client_streams.server_send.send(SessionMessage(request))
response_msg = await client_streams.server_receive.receive() response_msg = await client_streams.server_receive.receive()
response = response_msg.message.root response = response_msg.message
assert isinstance(response, types.JSONRPCResponse) assert isinstance(response, types.JSONRPCResponse)
assert isinstance(response.result, dict) assert isinstance(response.result, dict)
@ -277,10 +273,10 @@ async def test_client_handles_list_tasks_request(client_streams: ClientTestStrea
id="req-3", id="req-3",
**typed_request.model_dump(by_alias=True), **typed_request.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) await client_streams.server_send.send(SessionMessage(request))
response_msg = await client_streams.server_receive.receive() response_msg = await client_streams.server_receive.receive()
response = response_msg.message.root response = response_msg.message
assert isinstance(response, types.JSONRPCResponse) assert isinstance(response, types.JSONRPCResponse)
result = ListTasksResult.model_validate(response.result) result = ListTasksResult.model_validate(response.result)
@ -340,10 +336,10 @@ async def test_client_handles_cancel_task_request(client_streams: ClientTestStre
id="req-4", id="req-4",
**typed_request.model_dump(by_alias=True), **typed_request.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) await client_streams.server_send.send(SessionMessage(request))
response_msg = await client_streams.server_receive.receive() response_msg = await client_streams.server_receive.receive()
response = response_msg.message.root response = response_msg.message
assert isinstance(response, types.JSONRPCResponse) assert isinstance(response, types.JSONRPCResponse)
result = CancelTaskResult.model_validate(response.result) result = CancelTaskResult.model_validate(response.result)
@ -448,11 +444,11 @@ async def test_client_task_augmented_sampling(client_streams: ClientTestStreams)
id="req-sampling", id="req-sampling",
**typed_request.model_dump(by_alias=True), **typed_request.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) await client_streams.server_send.send(SessionMessage(request))
# Step 2: Client responds with CreateTaskResult # Step 2: Client responds with CreateTaskResult
response_msg = await client_streams.server_receive.receive() response_msg = await client_streams.server_receive.receive()
response = response_msg.message.root response = response_msg.message
assert isinstance(response, types.JSONRPCResponse) assert isinstance(response, types.JSONRPCResponse)
task_result = CreateTaskResult.model_validate(response.result) task_result = CreateTaskResult.model_validate(response.result)
@ -469,10 +465,10 @@ async def test_client_task_augmented_sampling(client_streams: ClientTestStreams)
id="req-poll", id="req-poll",
**typed_poll.model_dump(by_alias=True), **typed_poll.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(poll_request))) await client_streams.server_send.send(SessionMessage(poll_request))
poll_response_msg = await client_streams.server_receive.receive() poll_response_msg = await client_streams.server_receive.receive()
poll_response = poll_response_msg.message.root poll_response = poll_response_msg.message
assert isinstance(poll_response, types.JSONRPCResponse) assert isinstance(poll_response, types.JSONRPCResponse)
status = GetTaskResult.model_validate(poll_response.result) status = GetTaskResult.model_validate(poll_response.result)
@ -485,10 +481,10 @@ async def test_client_task_augmented_sampling(client_streams: ClientTestStreams)
id="req-result", id="req-result",
**typed_result_req.model_dump(by_alias=True), **typed_result_req.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(result_request))) await client_streams.server_send.send(SessionMessage(result_request))
result_response_msg = await client_streams.server_receive.receive() result_response_msg = await client_streams.server_receive.receive()
result_response = result_response_msg.message.root result_response = result_response_msg.message
assert isinstance(result_response, types.JSONRPCResponse) assert isinstance(result_response, types.JSONRPCResponse)
assert isinstance(result_response.result, dict) assert isinstance(result_response.result, dict)
@ -588,11 +584,11 @@ async def test_client_task_augmented_elicitation(client_streams: ClientTestStrea
id="req-elicit", id="req-elicit",
**typed_request.model_dump(by_alias=True), **typed_request.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) await client_streams.server_send.send(SessionMessage(request))
# Step 2: Client responds with CreateTaskResult # Step 2: Client responds with CreateTaskResult
response_msg = await client_streams.server_receive.receive() response_msg = await client_streams.server_receive.receive()
response = response_msg.message.root response = response_msg.message
assert isinstance(response, types.JSONRPCResponse) assert isinstance(response, types.JSONRPCResponse)
task_result = CreateTaskResult.model_validate(response.result) task_result = CreateTaskResult.model_validate(response.result)
@ -609,10 +605,10 @@ async def test_client_task_augmented_elicitation(client_streams: ClientTestStrea
id="req-poll", id="req-poll",
**typed_poll.model_dump(by_alias=True), **typed_poll.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(poll_request))) await client_streams.server_send.send(SessionMessage(poll_request))
poll_response_msg = await client_streams.server_receive.receive() poll_response_msg = await client_streams.server_receive.receive()
poll_response = poll_response_msg.message.root poll_response = poll_response_msg.message
assert isinstance(poll_response, types.JSONRPCResponse) assert isinstance(poll_response, types.JSONRPCResponse)
status = GetTaskResult.model_validate(poll_response.result) status = GetTaskResult.model_validate(poll_response.result)
@ -625,10 +621,10 @@ async def test_client_task_augmented_elicitation(client_streams: ClientTestStrea
id="req-result", id="req-result",
**typed_result_req.model_dump(by_alias=True), **typed_result_req.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(result_request))) await client_streams.server_send.send(SessionMessage(result_request))
result_response_msg = await client_streams.server_receive.receive() result_response_msg = await client_streams.server_receive.receive()
result_response = result_response_msg.message.root result_response = result_response_msg.message
assert isinstance(result_response, types.JSONRPCResponse) assert isinstance(result_response, types.JSONRPCResponse)
# Verify the elicitation result # Verify the elicitation result
@ -667,10 +663,10 @@ async def test_client_returns_error_for_unhandled_task_request(client_streams: C
id="req-unhandled", id="req-unhandled",
**typed_request.model_dump(by_alias=True), **typed_request.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) await client_streams.server_send.send(SessionMessage(request))
response_msg = await client_streams.server_receive.receive() response_msg = await client_streams.server_receive.receive()
response = response_msg.message.root response = response_msg.message
assert isinstance(response, types.JSONRPCError) assert isinstance(response, types.JSONRPCError)
assert ( assert (
"not supported" in response.error.message.lower() "not supported" in response.error.message.lower()
@ -706,10 +702,10 @@ async def test_client_returns_error_for_unhandled_task_result_request(client_str
id="req-result", id="req-result",
**typed_request.model_dump(by_alias=True), **typed_request.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) await client_streams.server_send.send(SessionMessage(request))
response_msg = await client_streams.server_receive.receive() response_msg = await client_streams.server_receive.receive()
response = response_msg.message.root response = response_msg.message
assert isinstance(response, types.JSONRPCError) assert isinstance(response, types.JSONRPCError)
assert "not supported" in response.error.message.lower() assert "not supported" in response.error.message.lower()
@ -742,10 +738,10 @@ async def test_client_returns_error_for_unhandled_list_tasks_request(client_stre
id="req-list", id="req-list",
**typed_request.model_dump(by_alias=True), **typed_request.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) await client_streams.server_send.send(SessionMessage(request))
response_msg = await client_streams.server_receive.receive() response_msg = await client_streams.server_receive.receive()
response = response_msg.message.root response = response_msg.message
assert isinstance(response, types.JSONRPCError) assert isinstance(response, types.JSONRPCError)
assert "not supported" in response.error.message.lower() assert "not supported" in response.error.message.lower()
@ -778,10 +774,10 @@ async def test_client_returns_error_for_unhandled_cancel_task_request(client_str
id="req-cancel", id="req-cancel",
**typed_request.model_dump(by_alias=True), **typed_request.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) await client_streams.server_send.send(SessionMessage(request))
response_msg = await client_streams.server_receive.receive() response_msg = await client_streams.server_receive.receive()
response = response_msg.message.root response = response_msg.message
assert isinstance(response, types.JSONRPCError) assert isinstance(response, types.JSONRPCError)
assert "not supported" in response.error.message.lower() assert "not supported" in response.error.message.lower()
@ -822,10 +818,10 @@ async def test_client_returns_error_for_unhandled_task_augmented_sampling(client
id="req-sampling", id="req-sampling",
**typed_request.model_dump(by_alias=True), **typed_request.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) await client_streams.server_send.send(SessionMessage(request))
response_msg = await client_streams.server_receive.receive() response_msg = await client_streams.server_receive.receive()
response = response_msg.message.root response = response_msg.message
assert isinstance(response, types.JSONRPCError) assert isinstance(response, types.JSONRPCError)
assert "not supported" in response.error.message.lower() assert "not supported" in response.error.message.lower()
@ -868,10 +864,10 @@ async def test_client_returns_error_for_unhandled_task_augmented_elicitation(
id="req-elicit", id="req-elicit",
**typed_request.model_dump(by_alias=True), **typed_request.model_dump(by_alias=True),
) )
await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request))) await client_streams.server_send.send(SessionMessage(request))
response_msg = await client_streams.server_receive.receive() response_msg = await client_streams.server_receive.receive()
response = response_msg.message.root response = response_msg.message
assert isinstance(response, types.JSONRPCError) assert isinstance(response, types.JSONRPCError)
assert "not supported" in response.error.message.lower() assert "not supported" in response.error.message.lower()

15
tests/experimental/tasks/server/test_server.py
View file

@ -36,7 +36,6 @@ from mcp.types import (
GetTaskRequestParams, GetTaskRequestParams,
GetTaskResult, GetTaskResult,
JSONRPCError, JSONRPCError,
JSONRPCMessage,
JSONRPCNotification, JSONRPCNotification,
JSONRPCResponse, JSONRPCResponse,
ListTasksRequest, ListTasksRequest,
@ -724,7 +723,7 @@ async def test_send_message() -> None:
# Create a test message # Create a test message
notification = JSONRPCNotification(jsonrpc="2.0", method="test/notification") notification = JSONRPCNotification(jsonrpc="2.0", method="test/notification")
message = SessionMessage( message = SessionMessage(
message=JSONRPCMessage(notification), message=notification,
metadata=ServerMessageMetadata(related_request_id="test-req-1"), metadata=ServerMessageMetadata(related_request_id="test-req-1"),
) )
@ -733,8 +732,8 @@ async def test_send_message() -> None:
# Verify it was sent to the stream # Verify it was sent to the stream
received = await server_to_client_receive.receive() received = await server_to_client_receive.receive()
assert isinstance(received.message.root, JSONRPCNotification) assert isinstance(received.message, JSONRPCNotification)
assert received.message.root.method == "test/notification" assert received.message.method == "test/notification"
finally: # pragma: no cover finally: # pragma: no cover
await server_to_client_send.aclose() await server_to_client_send.aclose()
await server_to_client_receive.aclose() await server_to_client_receive.aclose()
@ -776,7 +775,7 @@ async def test_response_routing_success() -> None:
# Simulate receiving a response from client # Simulate receiving a response from client
response = JSONRPCResponse(jsonrpc="2.0", id="test-req-1", result={"status": "ok"}) response = JSONRPCResponse(jsonrpc="2.0", id="test-req-1", result={"status": "ok"})
message = SessionMessage(message=JSONRPCMessage(response)) message = SessionMessage(message=response)
# Send from "client" side # Send from "client" side
await client_to_server_send.send(message) await client_to_server_send.send(message)
@ -831,7 +830,7 @@ async def test_response_routing_error() -> None:
# Simulate receiving an error response from client # Simulate receiving an error response from client
error_data = ErrorData(code=INVALID_REQUEST, message="Test error") error_data = ErrorData(code=INVALID_REQUEST, message="Test error")
error_response = JSONRPCError(jsonrpc="2.0", id="test-req-2", error=error_data) error_response = JSONRPCError(jsonrpc="2.0", id="test-req-2", error=error_data)
message = SessionMessage(message=JSONRPCMessage(error_response)) message = SessionMessage(message=error_response)
# Send from "client" side # Send from "client" side
await client_to_server_send.send(message) await client_to_server_send.send(message)
@ -894,7 +893,7 @@ async def test_response_routing_skips_non_matching_routers() -> None:
# Send a response - should skip first router and be handled by second # Send a response - should skip first router and be handled by second
response = JSONRPCResponse(jsonrpc="2.0", id="test-req-1", result={"status": "ok"}) response = JSONRPCResponse(jsonrpc="2.0", id="test-req-1", result={"status": "ok"})
message = SessionMessage(message=JSONRPCMessage(response)) message = SessionMessage(message=response)
await client_to_server_send.send(message) await client_to_server_send.send(message)
with anyio.fail_after(5): with anyio.fail_after(5):
@ -953,7 +952,7 @@ async def test_error_routing_skips_non_matching_routers() -> None:
# Send an error - should skip first router and be handled by second # Send an error - should skip first router and be handled by second
error_data = ErrorData(code=INVALID_REQUEST, message="Test error") error_data = ErrorData(code=INVALID_REQUEST, message="Test error")
error_response = JSONRPCError(jsonrpc="2.0", id="test-req-2", error=error_data) error_response = JSONRPCError(jsonrpc="2.0", id="test-req-2", error=error_data)
message = SessionMessage(message=JSONRPCMessage(error_response)) message = SessionMessage(message=error_response)
await client_to_server_send.send(message) await client_to_server_send.send(message)
with anyio.fail_after(5): with anyio.fail_after(5):

10
tests/issues/test_192_request_id.py
View file

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

24
tests/issues/test_malformed_input.py
View file

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

86
tests/server/test_lifespan.py
View file

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

Some files were not shown because too many files have changed in this diff Show more