Revert "refactor: reorganize message handling for better type safety and clar…" (#282)

This reverts commit 9d0f2daddb.

tests/client/test_session.py

@@ -1,5 +1,3 @@
-from types import NoneType
-
 import anyio
 import pytest
 
@@ -13,9 +11,9 @@ from mcp.types import (
     InitializeRequest,
     InitializeResult,
     JSONRPCMessage,
+    JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
-    MessageFrame,
     ServerCapabilities,
     ServerResult,
 )
@@ -24,10 +22,10 @@ from mcp.types import (
 @pytest.mark.anyio
 async def test_client_session_initialize():
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
-        MessageFrame[NoneType]
+        JSONRPCMessage
     ](1)
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
-        MessageFrame[NoneType]
+        JSONRPCMessage
     ](1)
 
     initialized_notification = None
@@ -36,7 +34,7 @@ async def test_client_session_initialize():
         nonlocal initialized_notification
 
         jsonrpc_request = await client_to_server_receive.receive()
-        assert isinstance(jsonrpc_request, MessageFrame)
+        assert isinstance(jsonrpc_request.root, JSONRPCRequest)
         request = ClientRequest.model_validate(
             jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
         )
@@ -58,25 +56,21 @@ async def test_client_session_initialize():
         )
 
         async with server_to_client_send:
-            assert isinstance(jsonrpc_request.message.root, JSONRPCRequest)
             await server_to_client_send.send(
-                MessageFrame(
-                    message=JSONRPCMessage(
-                        JSONRPCResponse(
-                            jsonrpc="2.0",
-                            id=jsonrpc_request.message.root.id,
-                            result=result.model_dump(
-                                by_alias=True, mode="json", exclude_none=True
-                            ),
-                        )
-                    ),
-                    raw=None,
+                JSONRPCMessage(
+                    JSONRPCResponse(
+                        jsonrpc="2.0",
+                        id=jsonrpc_request.root.id,
+                        result=result.model_dump(
+                            by_alias=True, mode="json", exclude_none=True
+                        ),
+                    )
                 )
             )
             jsonrpc_notification = await client_to_server_receive.receive()
-            assert isinstance(jsonrpc_notification.message, JSONRPCMessage)
+            assert isinstance(jsonrpc_notification.root, JSONRPCNotification)
             initialized_notification = ClientNotification.model_validate(
-                jsonrpc_notification.message.model_dump(
+                jsonrpc_notification.model_dump(
                     by_alias=True, mode="json", exclude_none=True
                 )
             )