convert to Pydantic models

akash329d · akash329d · commit 30b475bea93a · 2025-05-05T13:43:02.000-07:00
diff --git a/examples/servers/simple-prompt/mcp_simple_prompt/server.py b/examples/servers/simple-prompt/mcp_simple_prompt/server.py
@@ -88,13 +88,12 @@ async def get_prompt(
         )
 
     if transport == "sse":
+        from mcp.server.message_queue.redis import RedisMessageDispatch
         from mcp.server.sse import SseServerTransport
         from starlette.applications import Starlette
         from starlette.responses import Response
         from starlette.routing import Mount, Route
 
-        from mcp.server.message_queue.redis import RedisMessageDispatch
-
         message_dispatch = RedisMessageDispatch("redis://localhost:6379/0")
 
         sse = SseServerTransport("/messages/", message_dispatch=message_dispatch)
diff --git a/src/mcp/client/sse.py b/src/mcp/client/sse.py
@@ -98,7 +98,7 @@ async def sse_reader(
                                             await read_stream_writer.send(exc)
                                             continue
 
-                                        session_message = SessionMessage(message)
+                                        session_message = SessionMessage(message=message)
                                         await read_stream_writer.send(session_message)
                                     case _:
                                         logger.warning(
diff --git a/src/mcp/client/stdio/__init__.py b/src/mcp/client/stdio/__init__.py
@@ -144,7 +144,7 @@ async def stdout_reader():
                             await read_stream_writer.send(exc)
                             continue
 
-                        session_message = SessionMessage(message)
+                        session_message = SessionMessage(message=message)
                         await read_stream_writer.send(session_message)
         except anyio.ClosedResourceError:
             await anyio.lowlevel.checkpoint()
diff --git a/src/mcp/client/streamable_http.py b/src/mcp/client/streamable_http.py
@@ -153,7 +153,7 @@ async def _handle_sse_event(
                 ):
                     message.root.id = original_request_id
 
-                session_message = SessionMessage(message)
+                session_message = SessionMessage(message=message)
                 await read_stream_writer.send(session_message)
 
                 # Call resumption token callback if we have an ID
@@ -286,7 +286,7 @@ async def _handle_json_response(
         try:
             content = await response.aread()
             message = JSONRPCMessage.model_validate_json(content)
-            session_message = SessionMessage(message)
+            session_message = SessionMessage(message=message)
             await read_stream_writer.send(session_message)
         except Exception as exc:
             logger.error(f"Error parsing JSON response: {exc}")
@@ -333,7 +333,7 @@ async def _send_session_terminated_error(
             id=request_id,
             error=ErrorData(code=32600, message="Session terminated"),
         )
-        session_message = SessionMessage(JSONRPCMessage(jsonrpc_error))
+        session_message = SessionMessage(message=JSONRPCMessage(jsonrpc_error))
         await read_stream_writer.send(session_message)
 
     async def post_writer(
diff --git a/src/mcp/client/websocket.py b/src/mcp/client/websocket.py
@@ -60,7 +60,7 @@ async def ws_reader():
                 async for raw_text in ws:
                     try:
                         message = types.JSONRPCMessage.model_validate_json(raw_text)
-                        session_message = SessionMessage(message)
+                        session_message = SessionMessage(message=message)
                         await read_stream_writer.send(session_message)
                     except ValidationError as exc:
                         # If JSON parse or model validation fails, send the exception
diff --git a/src/mcp/server/message_queue/redis.py b/src/mcp/server/message_queue/redis.py
@@ -8,8 +8,8 @@
 from anyio.abc import TaskGroup
 from pydantic import ValidationError
 
-import mcp.types as types
 from mcp.server.message_queue.base import MessageCallback
+from mcp.shared.message import SessionMessage
 
 try:
     import redis.asyncio as redis
@@ -165,7 +165,7 @@ async def _handle_message(self, session_id: UUID, data: str) -> None:
             # Parse message or pass validation error to callback
             msg_or_error = None
             try:
-                msg_or_error = types.JSONRPCMessage.model_validate_json(data)
+                msg_or_error = SessionMessage.model_validate_json(data)
             except ValidationError as exc:
                 msg_or_error = exc
 
@@ -174,7 +174,7 @@ async def _handle_message(self, session_id: UUID, data: str) -> None:
             logger.error(f"Error in message handler for {session_id}: {e}")
 
     async def publish_message(
-        self, session_id: UUID, message: types.JSONRPCMessage | str
+        self, session_id: UUID, message: SessionMessage | str
     ) -> bool:
         """Publish a message for the specified session."""
         if not await self.session_exists(session_id):
diff --git a/src/mcp/server/stdio.py b/src/mcp/server/stdio.py
@@ -67,7 +67,7 @@ async def stdin_reader():
                         await read_stream_writer.send(exc)
                         continue
 
-                    session_message = SessionMessage(message)
+                    session_message = SessionMessage(message=message)
                     await read_stream_writer.send(session_message)
         except anyio.ClosedResourceError:
             await anyio.lowlevel.checkpoint()
diff --git a/src/mcp/server/streamable_http.py b/src/mcp/server/streamable_http.py
@@ -398,7 +398,7 @@ async def _handle_post_request(
                 await response(scope, receive, send)
 
                 # Process the message after sending the response
-                session_message = SessionMessage(message)
+                session_message = SessionMessage(message=message)
                 await writer.send(session_message)
 
                 return
@@ -413,7 +413,7 @@ async def _handle_post_request(
 
             if self.is_json_response_enabled:
                 # Process the message
-                session_message = SessionMessage(message)
+                session_message = SessionMessage(message=message)
                 await writer.send(session_message)
                 try:
                     # Process messages from the request-specific stream
@@ -512,7 +512,7 @@ async def sse_writer():
                     async with anyio.create_task_group() as tg:
                         tg.start_soon(response, scope, receive, send)
                         # Then send the message to be processed by the server
-                        session_message = SessionMessage(message)
+                        session_message = SessionMessage(message=message)
                         await writer.send(session_message)
                 except Exception:
                     logger.exception("SSE response error")
diff --git a/src/mcp/server/websocket.py b/src/mcp/server/websocket.py
@@ -42,7 +42,7 @@ async def ws_reader():
                         await read_stream_writer.send(exc)
                         continue
 
-                    session_message = SessionMessage(client_message)
+                    session_message = SessionMessage(message=client_message)
                     await read_stream_writer.send(session_message)
         except anyio.ClosedResourceError:
             await websocket.close()
diff --git a/src/mcp/shared/message.py b/src/mcp/shared/message.py
@@ -38,4 +38,4 @@ class SessionMessage(BaseModel):
     """A message with specific metadata for transport-specific features."""
 
     message: JSONRPCMessage
-    metadata: MessageMetadata = None
+    metadata: MessageMetadata | None = None
diff --git a/tests/client/test_session.py b/tests/client/test_session.py
@@ -62,7 +62,7 @@ async def mock_server():
         async with server_to_client_send:
             await server_to_client_send.send(
                 SessionMessage(
-                    JSONRPCMessage(
+                    message=JSONRPCMessage(
                         JSONRPCResponse(
                             jsonrpc="2.0",
                             id=jsonrpc_request.root.id,
@@ -153,7 +153,7 @@ async def mock_server():
         async with server_to_client_send:
             await server_to_client_send.send(
                 SessionMessage(
-                    JSONRPCMessage(
+                    message=JSONRPCMessage(
                         JSONRPCResponse(
                             jsonrpc="2.0",
                             id=jsonrpc_request.root.id,
@@ -220,7 +220,7 @@ async def mock_server():
         async with server_to_client_send:
             await server_to_client_send.send(
                 SessionMessage(
-                    JSONRPCMessage(
+                    message=JSONRPCMessage(
                         JSONRPCResponse(
                             jsonrpc="2.0",
                             id=jsonrpc_request.root.id,
diff --git a/tests/client/test_stdio.py b/tests/client/test_stdio.py
@@ -23,7 +23,7 @@ async def test_stdio_client():
 
         async with write_stream:
             for message in messages:
-                session_message = SessionMessage(message)
+                session_message = SessionMessage(message=message)
                 await write_stream.send(session_message)
 
         read_messages = []
diff --git a/tests/issues/test_192_request_id.py b/tests/issues/test_192_request_id.py
@@ -65,7 +65,7 @@ async def run_server():
             jsonrpc="2.0",
         )
 
-        await client_writer.send(SessionMessage(JSONRPCMessage(root=init_req)))
+        await client_writer.send(SessionMessage(message=JSONRPCMessage(root=init_req)))
         response = (
             await server_reader.receive()
         )  # Get init response but don't need to check it
@@ -77,15 +77,15 @@ async def run_server():
             jsonrpc="2.0",
         )
         await client_writer.send(
-            SessionMessage(JSONRPCMessage(root=initialized_notification))
+            SessionMessage(message=JSONRPCMessage(root=initialized_notification))
         )
 
         # Send ping request with custom ID
         ping_request = JSONRPCRequest(
             id=custom_request_id, method="ping", params={}, jsonrpc="2.0"
         )
 
-        await client_writer.send(SessionMessage(JSONRPCMessage(root=ping_request)))
+        await client_writer.send(SessionMessage(message=JSONRPCMessage(root=ping_request)))
 
         # Read response
         response = await server_reader.receive()
diff --git a/tests/server/message_queue/test_redis.py b/tests/server/message_queue/test_redis.py
@@ -2,10 +2,12 @@
 from uuid import uuid4
 
 import anyio
-from pydantic import ValidationError
 import pytest
+from pydantic import ValidationError
 
 import mcp.types as types
+from mcp.server.message_queue.redis import RedisMessageDispatch
+from mcp.shared.message import SessionMessage
 
 # Set up fakeredis for testing
 try:
@@ -97,7 +99,7 @@ async def test_subscribe_unsubscribe(redis_dispatch):
 
 
 @pytest.mark.anyio
-async def test_publish_message_valid_json(redis_dispatch):
+async def test_publish_message_valid_json(redis_dispatch: RedisMessageDispatch):
     """Test publishing a valid JSON-RPC message."""
     session_id = uuid4()
     callback = AsyncMock()
@@ -108,7 +110,9 @@ async def test_publish_message_valid_json(redis_dispatch):
     # Subscribe to messages
     async with redis_dispatch.subscribe(session_id, callback):
         # Publish message
-        published = await redis_dispatch.publish_message(session_id, message)
+        published = await redis_dispatch.publish_message(
+            session_id, SessionMessage(message=message)
+        )
         assert published
 
         # Give some time for the message to be processed
@@ -117,9 +121,11 @@ async def test_publish_message_valid_json(redis_dispatch):
         # Callback should have been called with the message
         callback.assert_called_once()
         call_args = callback.call_args[0][0]
-        assert isinstance(call_args, types.JSONRPCMessage)
-        assert isinstance(call_args.root, types.JSONRPCRequest)
-        assert call_args.root.method == "test"  # Access method through root attribute
+        assert isinstance(call_args, SessionMessage)
+        assert isinstance(call_args.message.root, types.JSONRPCRequest)
+        assert (
+            call_args.message.root.method == "test"
+        )  # Access method through root attribute
 
 
 @pytest.mark.anyio
@@ -177,7 +183,7 @@ async def test_extract_session_id(redis_dispatch):
 
 
 @pytest.mark.anyio
-async def test_multiple_sessions(redis_dispatch):
+async def test_multiple_sessions(redis_dispatch: RedisMessageDispatch):
     """Test handling multiple concurrent sessions."""
     session1 = uuid4()
     session2 = uuid4()
@@ -194,13 +200,17 @@ async def test_multiple_sessions(redis_dispatch):
             message1 = types.JSONRPCMessage.model_validate(
                 {"jsonrpc": "2.0", "method": "test1", "params": {}, "id": 1}
             )
-            await redis_dispatch.publish_message(session1, message1)
+            await redis_dispatch.publish_message(
+                session1, SessionMessage(message=message1)
+            )
 
             # Publish to session2
             message2 = types.JSONRPCMessage.model_validate(
                 {"jsonrpc": "2.0", "method": "test2", "params": {}, "id": 2}
             )
-            await redis_dispatch.publish_message(session2, message2)
+            await redis_dispatch.publish_message(
+                session2, SessionMessage(message=message2)
+            )
 
             # Give some time for messages to be processed
             await anyio.sleep(0.1)
@@ -210,10 +220,12 @@ async def test_multiple_sessions(redis_dispatch):
             callback2.assert_called_once()
 
             call1_args = callback1.call_args[0][0]
-            assert call1_args.root.method == "test1"
+            assert isinstance(call1_args, SessionMessage)
+            assert call1_args.message.root.method == "test1"  # type: ignore
 
             call2_args = callback2.call_args[0][0]
-            assert call2_args.root.method == "test2"
+            assert isinstance(call2_args, SessionMessage)
+            assert call2_args.message.root.method == "test2"  # type: ignore
 
 
 @pytest.mark.anyio
@@ -306,7 +318,10 @@ async def session1_runner():
 
 @pytest.mark.anyio
 async def test_listener_task_handoff_on_cancellation(redis_dispatch):
-    """Test that the single listening task is properly handed off when a session is cancelled."""
+    """
+    Test that the single listening task is properly
+    handed off when a session is cancelled.
+    """
     session1 = uuid4()
     session2 = uuid4()
 
diff --git a/tests/server/message_queue/test_redis_integration.py b/tests/server/message_queue/test_redis_integration.py
@@ -168,13 +168,13 @@ async def test_redis_integration_tool_call(server: None, server_url: str) -> Non
 
             # Call a tool
             result = await session.call_tool("test_tool", {})
-            assert result.content[0].text == "Called test_tool"
+            assert result.content[0].text == "Called test_tool"  # type: ignore
 
 
 @pytest.mark.anyio
 async def test_redis_integration_session_lifecycle() -> None:
-    """Test that sessions are properly added to and removed from Redis using direct Redis access"""
-    # Create a fresh Redis instance with decode_responses=True to get str instead of bytes
+    """Test that sessions are properly added to and
+    removed from Redis using direct Redis access"""
     mock_redis = fake_redis.FakeRedis(decode_responses=True)
     active_sessions_key = "mcp:pubsub:active_sessions"
 
@@ -221,14 +221,14 @@ async def mock_callback(message):
 @pytest.mark.anyio
 async def test_redis_integration_message_publishing_direct() -> None:
     """Test message publishing through Redis channels using direct Redis access"""
-    # Create a fresh Redis instance with decode_responses=True to get str instead of bytes
     mock_redis = fake_redis.FakeRedis(decode_responses=True)
 
     # Mock Redis in RedisMessageDispatch
     with patch(
         "mcp.server.message_queue.redis.redis.from_url", return_value=mock_redis
     ):
         from mcp.server.message_queue.redis import RedisMessageDispatch
+        from mcp.shared.message import SessionMessage
         from mcp.types import JSONRPCMessage, JSONRPCRequest
 
         # Create Redis message dispatch with our specific mock redis instance
@@ -258,7 +258,9 @@ async def message_callback(message):
             )
 
             # Publish the message
-            success = await message_dispatch.publish_message(session_id, test_message)
+            success = await message_dispatch.publish_message(
+                session_id, SessionMessage(message=test_message)
+            )
             assert success
 
             # Give some time for the message to be processed
@@ -270,6 +272,6 @@ async def message_callback(message):
                 len(messages_received) > 0
             ), "No messages were received through the callback"
             received_message = messages_received[0]
-            assert isinstance(received_message, JSONRPCMessage)
-            assert received_message.root.method == "test_method"
-            assert received_message.root.id == 1
+            assert isinstance(received_message, SessionMessage)
+            assert received_message.message.root.method == "test_method"
+            assert received_message.message.root.id == 1
diff --git a/tests/server/test_lifespan.py b/tests/server/test_lifespan.py
diff --git a/tests/server/test_stdio.py b/tests/server/test_stdio.py
