refactor: reorganize message handling for better type safety and clarity (#239)

* refactor: improve typing with memory stream type aliases

Move memory stream type definitions to models.py and use them throughout
the codebase for better type safety and maintainability.

GitHub-Issue:#201

* refactor: move streams to ParsedMessage

* refactor: update test files to use ParsedMessage

Updates test files to work with the ParsedMessage stream type aliases
and fixes a line length issue in test_201_client_hangs_on_logging.py.

Github-Issue:#201

* refactor: rename ParsedMessage to MessageFrame for clarity

🤖 Generated with [Claude Code](https://claude.ai/code)
Co-Authored-By: Claude <noreply@anthropic.com>

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

🤖 Generated with [Claude Code](https://claude.ai/code)
Co-Authored-By: Claude <noreply@anthropic.com>

* fix pyright

* refactor: update websocket client to use MessageFrame

Modified the websocket client to work with the new MessageFrame type,
preserving raw message text and properly extracting the root JSON-RPC
message when sending.

Github-Issue:#204

* fix: use NoneType instead of None for type parameters in MessageFrame

🤖 Generated with [Claude Code](https://claude.ai/code)

* refactor: rename root to message

tests/server/test_lifespan.py

@@ -17,6 +17,7 @@ from mcp.types import (
     JSONRPCMessage,
     JSONRPCNotification,
     JSONRPCRequest,
+    MessageFrame,
 )
 
 
@@ -64,7 +65,7 @@ async def test_lowlevel_server_lifespan():
                 send_stream2,
                 InitializationOptions(
                     server_name="test",
-                    server_version="0.1.0",
+                    server_version="1.0.0",
                     capabilities=server.get_capabilities(
                         notification_options=NotificationOptions(),
                         experimental_capabilities={},
@@ -82,42 +83,51 @@ async def test_lowlevel_server_lifespan():
             clientInfo=Implementation(name="test-client", version="0.1.0"),
         )
         await send_stream1.send(
-            JSONRPCMessage(
-                root=JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id=1,
-                    method="initialize",
-                    params=TypeAdapter(InitializeRequestParams).dump_python(params),
-                )
+            MessageFrame(
+                message=JSONRPCMessage(
+                    root=JSONRPCRequest(
+                        jsonrpc="2.0",
+                        id=1,
+                        method="initialize",
+                        params=TypeAdapter(InitializeRequestParams).dump_python(params),
+                    )
+                ),
+                raw=None,
             )
         )
         response = await receive_stream2.receive()
 
         # Send initialized notification
         await send_stream1.send(
-            JSONRPCMessage(
-                root=JSONRPCNotification(
-                    jsonrpc="2.0",
-                    method="notifications/initialized",
-                )
+            MessageFrame(
+                message=JSONRPCMessage(
+                    root=JSONRPCNotification(
+                        jsonrpc="2.0",
+                        method="notifications/initialized",
+                    )
+                ),
+                raw=None,
             )
         )
 
         # Call the tool to verify lifespan context
         await send_stream1.send(
-            JSONRPCMessage(
-                root=JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id=2,
-                    method="tools/call",
-                    params={"name": "check_lifespan", "arguments": {}},
-                )
+            MessageFrame(
+                message=JSONRPCMessage(
+                    root=JSONRPCRequest(
+                        jsonrpc="2.0",
+                        id=2,
+                        method="tools/call",
+                        params={"name": "check_lifespan", "arguments": {}},
+                    )
+                ),
+                raw=None,
             )
         )
 
         # Get response and verify
         response = await receive_stream2.receive()
-        assert response.root.result["content"][0]["text"] == "true"
+        assert response.message.root.result["content"][0]["text"] == "true"
 
         # Cancel server task
         tg.cancel_scope.cancel()
@@ -178,42 +188,51 @@ async def test_fastmcp_server_lifespan():
             clientInfo=Implementation(name="test-client", version="0.1.0"),
         )
         await send_stream1.send(
-            JSONRPCMessage(
-                root=JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id=1,
-                    method="initialize",
-                    params=TypeAdapter(InitializeRequestParams).dump_python(params),
-                )
+            MessageFrame(
+                message=JSONRPCMessage(
+                    root=JSONRPCRequest(
+                        jsonrpc="2.0",
+                        id=1,
+                        method="initialize",
+                        params=TypeAdapter(InitializeRequestParams).dump_python(params),
+                    )
+                ),
+                raw=None,
             )
         )
         response = await receive_stream2.receive()
 
         # Send initialized notification
         await send_stream1.send(
-            JSONRPCMessage(
-                root=JSONRPCNotification(
-                    jsonrpc="2.0",
-                    method="notifications/initialized",
-                )
+            MessageFrame(
+                message=JSONRPCMessage(
+                    root=JSONRPCNotification(
+                        jsonrpc="2.0",
+                        method="notifications/initialized",
+                    )
+                ),
+                raw=None,
             )
         )
 
         # Call the tool to verify lifespan context
         await send_stream1.send(
-            JSONRPCMessage(
-                root=JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id=2,
-                    method="tools/call",
-                    params={"name": "check_lifespan", "arguments": {}},
-                )
+            MessageFrame(
+                message=JSONRPCMessage(
+                    root=JSONRPCRequest(
+                        jsonrpc="2.0",
+                        id=2,
+                        method="tools/call",
+                        params={"name": "check_lifespan", "arguments": {}},
+                    )
+                ),
+                raw=None,
             )
         )
 
         # Get response and verify
         response = await receive_stream2.receive()
-        assert response.root.result["content"][0]["text"] == "true"
+        assert response.message.root.result["content"][0]["text"] == "true"
 
         # Cancel server task
         tg.cancel_scope.cancel()

tests/server/test_session.py

@@ -9,7 +9,7 @@ from mcp.server.session import ServerSession
 from mcp.types import (
     ClientNotification,
     InitializedNotification,
-    JSONRPCMessage,
+    MessageFrame,
     PromptsCapability,
     ResourcesCapability,
     ServerCapabilities,
@@ -19,10 +19,10 @@ from mcp.types import (
 @pytest.mark.anyio
 async def test_server_session_initialize():
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        MessageFrame[None]
     ](1)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        MessageFrame[None]
     ](1)
 
     async def run_client(client: ClientSession):

tests/server/test_stdio.py

@@ -4,7 +4,7 @@ import anyio
 import pytest
 
 from mcp.server.stdio import stdio_server
-from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
+from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse, MessageFrame
 
 
 @pytest.mark.anyio
@@ -13,8 +13,8 @@ async def test_stdio_server():
     stdout = io.StringIO()
 
     messages = [
-        JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")),
-        JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})),
+        JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"),
+        JSONRPCResponse(jsonrpc="2.0", id=2, result={}),
     ]
 
     for message in messages:
@@ -35,17 +35,29 @@ async def test_stdio_server():
 
         # Verify received messages
         assert len(received_messages) == 2
-        assert received_messages[0] == JSONRPCMessage(
-            root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
-        )
-        assert received_messages[1] == JSONRPCMessage(
-            root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})
-        )
+        assert isinstance(received_messages[0].message, JSONRPCMessage)
+        assert isinstance(received_messages[0].message.root, JSONRPCRequest)
+        assert received_messages[0].message.root.id == 1
+        assert received_messages[0].message.root.method == "ping"
+
+        assert isinstance(received_messages[1].message, JSONRPCMessage)
+        assert isinstance(received_messages[1].message.root, JSONRPCResponse)
+        assert received_messages[1].message.root.id == 2
 
         # Test sending responses from the server
         responses = [
-            JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")),
-            JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=4, result={})),
+            MessageFrame(
+                message=JSONRPCMessage(
+                    root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")
+                ),
+                raw=None,
+            ),
+            MessageFrame(
+                message=JSONRPCMessage(
+                    root=JSONRPCResponse(jsonrpc="2.0", id=4, result={})
+                ),
+                raw=None,
+            ),
         ]
 
         async with write_stream:
@@ -56,13 +68,10 @@ async def test_stdio_server():
     output_lines = stdout.readlines()
     assert len(output_lines) == 2
 
-    received_responses = [
-        JSONRPCMessage.model_validate_json(line.strip()) for line in output_lines
-    ]
-    assert len(received_responses) == 2
-    assert received_responses[0] == JSONRPCMessage(
-        root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")
-    )
-    assert received_responses[1] == JSONRPCMessage(
-        root=JSONRPCResponse(jsonrpc="2.0", id=4, result={})
-    )
+    # Parse and verify the JSON responses directly
+    request_json = JSONRPCRequest.model_validate_json(output_lines[0].strip())
+    response_json = JSONRPCResponse.model_validate_json(output_lines[1].strip())
+
+    assert request_json.id == 3
+    assert request_json.method == "ping"
+    assert response_json.id == 4