refactor: move MessageFrame class to types.py for better code organization

🤖 Generated with [Claude Code](https://claude.ai/code)
Co-Authored-By: Claude <[email protected]>

Author: dsp-ant

src/mcp/client/sse.py

@@ -10,12 +10,12 @@
 
 import mcp.types as types
 from mcp.shared.session import (
-    MessageFrame,
     ReadStream,
     ReadStreamWriter,
     WriteStream,
     WriteStreamReader,
 )
+from mcp.types import MessageFrame
 
 logger = logging.getLogger(__name__)
 
@@ -91,7 +91,7 @@ async def sse_reader(
                                     case "message":
                                         try:
                                             message = MessageFrame(
-                                                types.JSONRPCMessage.model_validate_json(  # noqa: E501
+                                                root=types.JSONRPCMessage.model_validate_json(  # noqa: E501
                                                     sse.data
                                                 ),
                                                 raw=sse,

src/mcp/server/sse.py

@@ -46,12 +46,12 @@ async def handle_sse(request):
 
 import mcp.types as types
 from mcp.shared.session import (
-    MessageFrame,
     ReadStream,
     ReadStreamWriter,
     WriteStream,
     WriteStreamReader,
 )
+from mcp.types import MessageFrame
 
 logger = logging.getLogger(__name__)
 
@@ -176,4 +176,4 @@ async def handle_post_message(
         logger.debug(f"Sending message to writer: {message}")
         response = Response("Accepted", status_code=202)
         await response(scope, receive, send)
-        await writer.send(MessageFrame(message, raw=request))
+        await writer.send(MessageFrame(root=message, raw=request))

src/mcp/server/stdio.py

@@ -27,12 +27,12 @@ async def run_server():
 
 import mcp.types as types
 from mcp.shared.session import (
-    MessageFrame,
     ReadStream,
     ReadStreamWriter,
     WriteStream,
     WriteStreamReader,
 )
+from mcp.types import MessageFrame
 
 
 @asynccontextmanager
@@ -72,14 +72,15 @@ async def stdin_reader():
                         await read_stream_writer.send(exc)
                         continue
 
-                    await read_stream_writer.send(MessageFrame(message, raw=line))
+                    await read_stream_writer.send(MessageFrame(root=message, raw=line))
         except anyio.ClosedResourceError:
             await anyio.lowlevel.checkpoint()
 
     async def stdout_writer():
         try:
             async with write_stream_reader:
                 async for message in write_stream_reader:
+                    # Extract the inner JSONRPCRequest/JSONRPCResponse from MessageFrame
                     json = message.model_dump_json(by_alias=True, exclude_none=True)
                     await stdout.write(json + "\n")
                     await stdout.flush()

src/mcp/server/websocket.py

@@ -7,12 +7,12 @@
 
 import mcp.types as types
 from mcp.shared.session import (
-    MessageFrame,
     ReadStream,
     ReadStreamWriter,
     WriteStream,
     WriteStreamReader,
 )
+from mcp.types import MessageFrame
 
 logger = logging.getLogger(__name__)
 
@@ -47,7 +47,7 @@ async def ws_reader():
                         continue
 
                     await read_stream_writer.send(
-                        MessageFrame(client_message, raw=message)
+                        MessageFrame(root=client_message, raw=message)
                     )
         except anyio.ClosedResourceError:
             await websocket.close()

src/mcp/shared/memory.py

@@ -11,7 +11,7 @@
 
 from mcp.client.session import ClientSession, ListRootsFnT, SamplingFnT
 from mcp.server import Server
-from mcp.shared.session import MessageFrame
+from mcp.types import MessageFrame
 
 MessageStream = tuple[
     MemoryObjectReceiveStream[MessageFrame | Exception],

src/mcp/shared/session.py

@@ -7,7 +7,7 @@
 import anyio.lowlevel
 import httpx
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-from pydantic import BaseModel, RootModel
+from pydantic import BaseModel
 
 from mcp.shared.exceptions import McpError
 from mcp.types import (
@@ -21,27 +21,17 @@
     JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
+    MessageFrame,
     RequestParams,
     ServerNotification,
     ServerRequest,
     ServerResult,
 )
 
-RawT = TypeVar("RawT")
-
-
-class MessageFrame(RootModel[JSONRPCMessage], Generic[RawT]):
-    root: JSONRPCMessage
-    raw: RawT | None = None
-
-    class Config:
-        arbitrary_types_allowed = True
-
-
-ReadStream = MemoryObjectReceiveStream[MessageFrame[RawT] | Exception]
-ReadStreamWriter = MemoryObjectSendStream[MessageFrame[RawT] | Exception]
-WriteStream = MemoryObjectSendStream[MessageFrame[RawT]]
-WriteStreamReader = MemoryObjectReceiveStream[MessageFrame[RawT]]
+ReadStream = MemoryObjectReceiveStream[MessageFrame | Exception]
+ReadStreamWriter = MemoryObjectSendStream[MessageFrame | Exception]
+WriteStream = MemoryObjectSendStream[MessageFrame]
+WriteStreamReader = MemoryObjectReceiveStream[MessageFrame]
 
 SendRequestT = TypeVar("SendRequestT", ClientRequest, ServerRequest)
 SendResultT = TypeVar("SendResultT", ClientResult, ServerResult)
@@ -242,7 +232,7 @@ async def send_request(
         # TODO: Support progress callbacks
 
         await self._write_stream.send(
-            MessageFrame(JSONRPCMessage(jsonrpc_request), None)
+            MessageFrame(root=JSONRPCMessage(jsonrpc_request), raw=None)
         )
 
         try:
@@ -280,15 +270,17 @@ async def send_notification(self, notification: SendNotificationT) -> None:
         )
 
         await self._write_stream.send(
-            MessageFrame(JSONRPCMessage(jsonrpc_notification))
+            MessageFrame(root=JSONRPCMessage(jsonrpc_notification), raw=None)
         )
 
     async def _send_response(
         self, request_id: RequestId, response: SendResultT | ErrorData
     ) -> None:
         if isinstance(response, ErrorData):
             jsonrpc_error = JSONRPCError(jsonrpc="2.0", id=request_id, error=response)
-            await self._write_stream.send(MessageFrame(JSONRPCMessage(jsonrpc_error)))
+            await self._write_stream.send(
+                MessageFrame(root=JSONRPCMessage(jsonrpc_error), raw=None)
+            )
         else:
             jsonrpc_response = JSONRPCResponse(
                 jsonrpc="2.0",
@@ -298,7 +290,7 @@ async def _send_response(
                 ),
             )
             await self._write_stream.send(
-                MessageFrame(JSONRPCMessage(jsonrpc_response))
+                MessageFrame(root=JSONRPCMessage(jsonrpc_response), raw=None)
             )
 
     async def _receive_loop(self) -> None:

src/mcp/types.py

@@ -180,6 +180,21 @@ class JSONRPCMessage(
     pass
 
 
+RawT = TypeVar("RawT")
+
+
+class MessageFrame(BaseModel, Generic[RawT]):
+    root: JSONRPCMessage
+    raw: RawT | None = None
+    model_config = ConfigDict(extra="allow")
+
+    def model_dump(self, *args, **kwargs):
+        return self.root.model_dump(*args, **kwargs)
+
+    def model_dump_json(self, *args, **kwargs):
+        return self.root.model_dump_json(*args, **kwargs)
+
+
 class EmptyResult(Result):
     """A response that indicates success but carries no data."""
 

tests/client/test_session.py

@@ -2,7 +2,6 @@
 import pytest
 
 from mcp.client.session import ClientSession
-from mcp.shared.session import MessageFrame
 from mcp.types import (
     LATEST_PROTOCOL_VERSION,
     ClientNotification,
@@ -12,8 +11,8 @@
     InitializeRequest,
     InitializeResult,
     JSONRPCMessage,
-    JSONRPCRequest,
     JSONRPCResponse,
+    MessageFrame,
     ServerCapabilities,
     ServerResult,
 )
@@ -34,7 +33,7 @@ async def mock_server():
         nonlocal initialized_notification
 
         jsonrpc_request = await client_to_server_receive.receive()
-        assert isinstance(jsonrpc_request.root, JSONRPCRequest)
+        assert isinstance(jsonrpc_request, MessageFrame)
         request = ClientRequest.model_validate(
             jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
         )
@@ -61,7 +60,7 @@ async def mock_server():
                     root=JSONRPCMessage(
                         JSONRPCResponse(
                             jsonrpc="2.0",
-                            id=jsonrpc_request.root.id,
+                            id=jsonrpc_request.root.root.id,
                             result=result.model_dump(
                                 by_alias=True, mode="json", exclude_none=True
                             ),
@@ -71,7 +70,7 @@ async def mock_server():
                 )
             )
             jsonrpc_notification = await client_to_server_receive.receive()
-            assert isinstance(jsonrpc_notification.root, MessageFrame)
+            assert isinstance(jsonrpc_notification.root, JSONRPCMessage)
             initialized_notification = ClientNotification.model_validate(
                 jsonrpc_notification.root.model_dump(
                     by_alias=True, mode="json", exclude_none=True

tests/issues/test_192_request_id.py

@@ -11,6 +11,7 @@
     JSONRPCMessage,
     JSONRPCNotification,
     JSONRPCRequest,
+    MessageFrame,
     NotificationParams,
 )
 
@@ -58,7 +59,9 @@ async def run_server():
             jsonrpc="2.0",
         )
 
-        await client_writer.send(JSONRPCMessage(root=init_req))
+        await client_writer.send(
+            MessageFrame(root=JSONRPCMessage(root=init_req), raw=None)
+        )
         await server_reader.receive()  # Get init response but don't need to check it
 
         # Send initialized notification
@@ -67,21 +70,25 @@ async def run_server():
             params=NotificationParams().model_dump(by_alias=True, exclude_none=True),
             jsonrpc="2.0",
         )
-        await client_writer.send(JSONRPCMessage(root=initialized_notification))
+        await client_writer.send(
+            MessageFrame(root=JSONRPCMessage(root=initialized_notification), raw=None)
+        )
 
         # Send ping request with custom ID
         ping_request = JSONRPCRequest(
             id=custom_request_id, method="ping", params={}, jsonrpc="2.0"
         )
 
-        await client_writer.send(JSONRPCMessage(root=ping_request))
+        await client_writer.send(
+            MessageFrame(root=JSONRPCMessage(root=ping_request), raw=None)
+        )
 
         # Read response
         response = await server_reader.receive()
 
         # Verify response ID matches request ID
         assert (
-            response.root.id == custom_request_id
+            response.root.root.id == custom_request_id
         ), "Response ID should match request ID"
 
         # Cancel server task