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

Author: dsp-ant

src/mcp/client/session.py

@@ -1,12 +1,11 @@
 from datetime import timedelta
 from typing import Any, Protocol
 
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from pydantic import AnyUrl, TypeAdapter
 
 import mcp.types as types
 from mcp.shared.context import RequestContext
-from mcp.shared.session import BaseSession, RequestResponder
+from mcp.shared.session import BaseSession, ReadStream, RequestResponder, WriteStream
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 
 
@@ -59,8 +58,8 @@ class ClientSession(
 ):
     def __init__(
         self,
-        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
-        write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
+        read_stream: ReadStream,
+        write_stream: WriteStream,
         read_timeout_seconds: timedelta | None = None,
         sampling_callback: SamplingFnT | None = None,
         list_roots_callback: ListRootsFnT | None = None,

src/mcp/client/sse.py

@@ -6,10 +6,16 @@
 import anyio
 import httpx
 from anyio.abc import TaskStatus
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from httpx_sse import aconnect_sse
 
 import mcp.types as types
+from mcp.shared.session import (
+    ReadStream,
+    ReadStreamWriter,
+    WriteStream,
+    WriteStreamReader,
+)
+from mcp.types import MessageFrame
 
 logger = logging.getLogger(__name__)
 
@@ -31,11 +37,11 @@ async def sse_client(
     `sse_read_timeout` determines how long (in seconds) the client will wait for a new
     event before disconnecting. All other HTTP operations are controlled by `timeout`.
     """
-    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
-    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+    read_stream: ReadStream
+    read_stream_writer: ReadStreamWriter
 
-    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
-    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+    write_stream: WriteStream
+    write_stream_reader: WriteStreamReader
 
     read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@@ -84,8 +90,11 @@ async def sse_reader(
 
                                     case "message":
                                         try:
-                                            message = types.JSONRPCMessage.model_validate_json(  # noqa: E501
-                                                sse.data
+                                            message = MessageFrame(
+                                                message=types.JSONRPCMessage.model_validate_json(  # noqa: E501
+                                                    sse.data
+                                                ),
+                                                raw=sse,
                                             )
                                             logger.debug(
                                                 f"Received server message: {message}"

src/mcp/client/websocket.py

@@ -1,7 +1,7 @@
 import json
 import logging
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator
+from typing import Any, AsyncGenerator
 
 import anyio
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
@@ -10,6 +10,7 @@
 from websockets.typing import Subprotocol
 
 import mcp.types as types
+from mcp.types import MessageFrame
 
 logger = logging.getLogger(__name__)
 
@@ -19,8 +20,8 @@ async def websocket_client(
     url: str,
 ) -> AsyncGenerator[
     tuple[
-        MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
-        MemoryObjectSendStream[types.JSONRPCMessage],
+        MemoryObjectReceiveStream[MessageFrame[Any] | Exception],
+        MemoryObjectSendStream[MessageFrame[Any]],
     ],
     None,
 ]:
@@ -53,7 +54,11 @@ async def ws_reader():
             async with read_stream_writer:
                 async for raw_text in ws:
                     try:
-                        message = types.JSONRPCMessage.model_validate_json(raw_text)
+                        json_message = types.JSONRPCMessage.model_validate_json(
+                            raw_text
+                        )
+                        # Create MessageFrame with JSON message as root
+                        message = MessageFrame(message=json_message, raw=raw_text)
                         await read_stream_writer.send(message)
                     except ValidationError as exc:
                         # If JSON parse or model validation fails, send the exception
@@ -66,8 +71,8 @@ async def ws_writer():
             """
             async with write_stream_reader:
                 async for message in write_stream_reader:
-                    # Convert to a dict, then to JSON
-                    msg_dict = message.model_dump(
+                    # Extract the JSON-RPC message from MessageFrame and convert to JSON
+                    msg_dict = message.message.model_dump(
                         by_alias=True, mode="json", exclude_none=True
                     )
                     await ws.send(json.dumps(msg_dict))

src/mcp/server/lowlevel/server.py

@@ -74,7 +74,6 @@ async def main():
 from typing import Any, AsyncIterator, Generic, TypeVar
 
 import anyio
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from pydantic import AnyUrl
 
 import mcp.types as types
@@ -84,7 +83,7 @@ async def main():
 from mcp.server.stdio import stdio_server as stdio_server
 from mcp.shared.context import RequestContext
 from mcp.shared.exceptions import McpError
-from mcp.shared.session import RequestResponder
+from mcp.shared.session import ReadStream, RequestResponder, WriteStream
 
 logger = logging.getLogger(__name__)
 
@@ -474,8 +473,8 @@ async def handler(req: types.CompleteRequest):
 
     async def run(
         self,
-        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
-        write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
+        read_stream: ReadStream,
+        write_stream: WriteStream,
         initialization_options: InitializationOptions,
         # When False, exceptions are returned as messages to the client.
         # When True, exceptions are raised, which will cause the server to shut down

src/mcp/server/models.py

@@ -5,9 +5,7 @@
 
 from pydantic import BaseModel
 
-from mcp.types import (
-    ServerCapabilities,
-)
+from mcp.types import ServerCapabilities
 
 
 class InitializationOptions(BaseModel):

src/mcp/server/session.py

@@ -42,14 +42,15 @@ async def handle_list_prompts(ctx: RequestContext) -> list[types.Prompt]:
 
 import anyio
 import anyio.lowlevel
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from pydantic import AnyUrl
 
 import mcp.types as types
 from mcp.server.models import InitializationOptions
 from mcp.shared.session import (
     BaseSession,
+    ReadStream,
     RequestResponder,
+    WriteStream,
 )
 
 
@@ -76,8 +77,8 @@ class ServerSession(
 
     def __init__(
         self,
-        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
-        write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
+        read_stream: ReadStream,
+        write_stream: WriteStream,
         init_options: InitializationOptions,
     ) -> None:
         super().__init__(

src/mcp/server/sse.py

@@ -38,14 +38,20 @@ async def handle_sse(request):
 from uuid import UUID, uuid4
 
 import anyio
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from pydantic import ValidationError
 from sse_starlette import EventSourceResponse
 from starlette.requests import Request
 from starlette.responses import Response
 from starlette.types import Receive, Scope, Send
 
 import mcp.types as types
+from mcp.shared.session import (
+    ReadStream,
+    ReadStreamWriter,
+    WriteStream,
+    WriteStreamReader,
+)
+from mcp.types import MessageFrame
 
 logger = logging.getLogger(__name__)
 
@@ -63,9 +69,7 @@ class SseServerTransport:
     """
 
     _endpoint: str
-    _read_stream_writers: dict[
-        UUID, MemoryObjectSendStream[types.JSONRPCMessage | Exception]
-    ]
+    _read_stream_writers: dict[UUID, ReadStreamWriter]
 
     def __init__(self, endpoint: str) -> None:
         """
@@ -85,11 +89,11 @@ async def connect_sse(self, scope: Scope, receive: Receive, send: Send):
             raise ValueError("connect_sse can only handle HTTP requests")
 
         logger.debug("Setting up SSE connection")
-        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
-        read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+        read_stream: ReadStream
+        read_stream_writer: ReadStreamWriter
 
-        write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
-        write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+        write_stream: WriteStream
+        write_stream_reader: WriteStreamReader
 
         read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
         write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@@ -172,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(message)
+        await writer.send(MessageFrame(message=message, raw=request))

src/mcp/server/stdio.py

@@ -24,9 +24,15 @@ async def run_server():
 
 import anyio
 import anyio.lowlevel
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 
 import mcp.types as types
+from mcp.shared.session import (
+    ReadStream,
+    ReadStreamWriter,
+    WriteStream,
+    WriteStreamReader,
+)
+from mcp.types import MessageFrame
 
 
 @asynccontextmanager
@@ -47,11 +53,11 @@ async def stdio_server(
     if not stdout:
         stdout = anyio.wrap_file(TextIOWrapper(sys.stdout.buffer, encoding="utf-8"))
 
-    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
-    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+    read_stream: ReadStream
+    read_stream_writer: ReadStreamWriter
 
-    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
-    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+    write_stream: WriteStream
+    write_stream_reader: WriteStreamReader
 
     read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@@ -66,14 +72,17 @@ async def stdin_reader():
                         await read_stream_writer.send(exc)
                         continue
 
-                    await read_stream_writer.send(message)
+                    await read_stream_writer.send(
+                        MessageFrame(message=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

@@ -2,11 +2,17 @@
 from contextlib import asynccontextmanager
 
 import anyio
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from starlette.types import Receive, Scope, Send
 from starlette.websockets import WebSocket
 
 import mcp.types as types
+from mcp.shared.session import (
+    ReadStream,
+    ReadStreamWriter,
+    WriteStream,
+    WriteStreamReader,
+)
+from mcp.types import MessageFrame
 
 logger = logging.getLogger(__name__)
 
@@ -21,11 +27,11 @@ async def websocket_server(scope: Scope, receive: Receive, send: Send):
     websocket = WebSocket(scope, receive, send)
     await websocket.accept(subprotocol="mcp")
 
-    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
-    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+    read_stream: ReadStream
+    read_stream_writer: ReadStreamWriter
 
-    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
-    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+    write_stream: WriteStream
+    write_stream_reader: WriteStreamReader
 
     read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@@ -40,7 +46,9 @@ async def ws_reader():
                         await read_stream_writer.send(exc)
                         continue
 
-                    await read_stream_writer.send(client_message)
+                    await read_stream_writer.send(
+                        MessageFrame(message=client_message, raw=message)
+                    )
         except anyio.ClosedResourceError:
             await websocket.close()
 

src/mcp/shared/memory.py

@@ -11,11 +11,11 @@
 
 from mcp.client.session import ClientSession, ListRootsFnT, SamplingFnT
 from mcp.server import Server
-from mcp.types import JSONRPCMessage
+from mcp.types import MessageFrame
 
 MessageStream = tuple[
-    MemoryObjectReceiveStream[JSONRPCMessage | Exception],
-    MemoryObjectSendStream[JSONRPCMessage],
+    MemoryObjectReceiveStream[MessageFrame | Exception],
+    MemoryObjectSendStream[MessageFrame],
 ]
 
 
@@ -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[
-        JSONRPCMessage | Exception
+        MessageFrame | Exception
     ](1)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage | Exception
+        MessageFrame | Exception
     ](1)
 
     client_streams = (server_to_client_receive, client_to_server_send)
@@ -60,12 +60,9 @@ 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_streams,
-        server_streams,
+        (client_read, client_write),
+        (server_read, server_write),
     ):
-        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(