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

This reverts commit 9d0f2daddb.

src/mcp/shared/memory.py

@@ -11,11 +11,11 @@ from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStre
 
 from mcp.client.session import ClientSession, ListRootsFnT, SamplingFnT
 from mcp.server import Server
-from mcp.types import MessageFrame
+from mcp.types import JSONRPCMessage
 
 MessageStream = tuple[
-    MemoryObjectReceiveStream[MessageFrame | Exception],
-    MemoryObjectSendStream[MessageFrame],
+    MemoryObjectReceiveStream[JSONRPCMessage | Exception],
+    MemoryObjectSendStream[JSONRPCMessage],
 ]
 
 
@@ -32,10 +32,10 @@ async def create_client_server_memory_streams() -> (
     """
     # Create streams for both directions
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
-        MessageFrame | Exception
+        JSONRPCMessage | Exception
     ](1)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
-        MessageFrame | Exception
+        JSONRPCMessage | Exception
     ](1)
 
     client_streams = (server_to_client_receive, client_to_server_send)
@@ -60,9 +60,12 @@ async def create_connected_server_and_client_session(
 ) -> AsyncGenerator[ClientSession, None]:
     """Creates a ClientSession that is connected to a running MCP server."""
     async with create_client_server_memory_streams() as (
-        (client_read, client_write),
-        (server_read, server_write),
+        client_streams,
+        server_streams,
     ):
+        client_read, client_write = client_streams
+        server_read, server_write = server_streams
+
         # Create a cancel scope for the server task
         async with anyio.create_task_group() as tg:
             tg.start_soon(

src/mcp/shared/session.py

@@ -22,18 +22,12 @@ from mcp.types import (
     JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
-    MessageFrame,
     RequestParams,
     ServerNotification,
     ServerRequest,
     ServerResult,
 )
 
-ReadStream = MemoryObjectReceiveStream[MessageFrame | Exception]
-ReadStreamWriter = MemoryObjectSendStream[MessageFrame | Exception]
-WriteStream = MemoryObjectSendStream[MessageFrame]
-WriteStreamReader = MemoryObjectReceiveStream[MessageFrame]
-
 SendRequestT = TypeVar("SendRequestT", ClientRequest, ServerRequest)
 SendResultT = TypeVar("SendResultT", ClientResult, ServerResult)
 SendNotificationT = TypeVar("SendNotificationT", ClientNotification, ServerNotification)
@@ -171,8 +165,8 @@ class BaseSession(
 
     def __init__(
         self,
-        read_stream: ReadStream,
-        write_stream: WriteStream,
+        read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
+        write_stream: MemoryObjectSendStream[JSONRPCMessage],
         receive_request_type: type[ReceiveRequestT],
         receive_notification_type: type[ReceiveNotificationT],
         # If none, reading will never time out
@@ -248,9 +242,7 @@ class BaseSession(
 
         # TODO: Support progress callbacks
 
-        await self._write_stream.send(
-            MessageFrame(message=JSONRPCMessage(jsonrpc_request), raw=None)
-        )
+        await self._write_stream.send(JSONRPCMessage(jsonrpc_request))
 
         try:
             with anyio.fail_after(
@@ -286,18 +278,14 @@ class BaseSession(
             **notification.model_dump(by_alias=True, mode="json", exclude_none=True),
         )
 
-        await self._write_stream.send(
-            MessageFrame(message=JSONRPCMessage(jsonrpc_notification), raw=None)
-        )
+        await self._write_stream.send(JSONRPCMessage(jsonrpc_notification))
 
     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(message=JSONRPCMessage(jsonrpc_error), raw=None)
-            )
+            await self._write_stream.send(JSONRPCMessage(jsonrpc_error))
         else:
             jsonrpc_response = JSONRPCResponse(
                 jsonrpc="2.0",
@@ -306,9 +294,7 @@ class BaseSession(
                     by_alias=True, mode="json", exclude_none=True
                 ),
             )
-            await self._write_stream.send(
-                MessageFrame(message=JSONRPCMessage(jsonrpc_response), raw=None)
-            )
+            await self._write_stream.send(JSONRPCMessage(jsonrpc_response))
 
     async def _receive_loop(self) -> None:
         async with (
@@ -316,13 +302,10 @@ class BaseSession(
             self._write_stream,
             self._incoming_message_stream_writer,
         ):
-            async for raw_message in self._read_stream:
-                if isinstance(raw_message, Exception):
-                    await self._incoming_message_stream_writer.send(raw_message)
-                    continue
-
-                message = raw_message.message
-                if isinstance(message.root, JSONRPCRequest):
+            async for message in self._read_stream:
+                if isinstance(message, Exception):
+                    await self._incoming_message_stream_writer.send(message)
+                elif isinstance(message.root, JSONRPCRequest):
                     validated_request = self._receive_request_type.model_validate(
                         message.root.model_dump(
                             by_alias=True, mode="json", exclude_none=True