refactor: move streams to ParsedMessage

dsp-ant · dsp-ant · commit 3af37d61d0c8 · 2025-03-13T13:25:12.000Z
diff --git a/src/mcp/client/session.py b/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,
diff --git a/src/mcp/client/sse.py b/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 (
+    ParsedMessage,
+    ReadStream,
+    ReadStreamWriter,
+    WriteStream,
+    WriteStreamReader,
+)
 
 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 = ParsedMessage(
+                                                types.JSONRPCMessage.model_validate_json(  # noqa: E501
+                                                    sse.data
+                                                ),
+                                                raw=sse,
                                             )
                                             logger.debug(
                                                 f"Received server message: {message}"
diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py
@@ -78,12 +78,12 @@ async def main():
 
 import mcp.types as types
 from mcp.server.lowlevel.helper_types import ReadResourceContents
-from mcp.server.models import InitializationOptions, ReadStream, WriteStream
+from mcp.server.models import InitializationOptions
 from mcp.server.session import ServerSession
 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__)
 
diff --git a/src/mcp/server/models.py b/src/mcp/server/models.py
@@ -3,15 +3,9 @@
 and tools.
 """
 
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from pydantic import BaseModel
 
-from mcp.types import JSONRPCMessage, ServerCapabilities
-
-ReadStream = MemoryObjectReceiveStream[JSONRPCMessage | Exception]
-ReadStreamWriter = MemoryObjectSendStream[JSONRPCMessage | Exception]
-WriteStream = MemoryObjectSendStream[JSONRPCMessage]
-WriteStreamReader = MemoryObjectReceiveStream[JSONRPCMessage]
+from mcp.types import ServerCapabilities
 
 
 class InitializationOptions(BaseModel):
diff --git a/src/mcp/server/session.py b/src/mcp/server/session.py
@@ -45,10 +45,12 @@ async def handle_list_prompts(ctx: RequestContext) -> list[types.Prompt]:
 from pydantic import AnyUrl
 
 import mcp.types as types
-from mcp.server.models import InitializationOptions, ReadStream, WriteStream
+from mcp.server.models import InitializationOptions
 from mcp.shared.session import (
     BaseSession,
+    ReadStream,
     RequestResponder,
+    WriteStream,
 )
 
 
diff --git a/src/mcp/server/sse.py b/src/mcp/server/sse.py
@@ -45,7 +45,8 @@ async def handle_sse(request):
 from starlette.types import Receive, Scope, Send
 
 import mcp.types as types
-from mcp.server.models import (
+from mcp.shared.session import (
+    ParsedMessage,
     ReadStream,
     ReadStreamWriter,
     WriteStream,
@@ -175,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(ParsedMessage(message, raw=request))
diff --git a/src/mcp/server/stdio.py b/src/mcp/server/stdio.py
@@ -26,7 +26,8 @@ async def run_server():
 import anyio.lowlevel
 
 import mcp.types as types
-from mcp.server.models import (
+from mcp.shared.session import (
+    ParsedMessage,
     ReadStream,
     ReadStreamWriter,
     WriteStream,
@@ -71,7 +72,7 @@ async def stdin_reader():
                         await read_stream_writer.send(exc)
                         continue
 
-                    await read_stream_writer.send(message)
+                    await read_stream_writer.send(ParsedMessage(message, raw=line))
         except anyio.ClosedResourceError:
             await anyio.lowlevel.checkpoint()
 
diff --git a/src/mcp/server/websocket.py b/src/mcp/server/websocket.py
@@ -6,7 +6,8 @@
 from starlette.websockets import WebSocket
 
 import mcp.types as types
-from mcp.server.models import (
+from mcp.shared.session import (
+    ParsedMessage,
     ReadStream,
     ReadStreamWriter,
     WriteStream,
@@ -45,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(
+                        ParsedMessage(client_message, raw=message)
+                    )
         except anyio.ClosedResourceError:
             await websocket.close()
 
diff --git a/src/mcp/shared/memory.py b/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.shared.session import ParsedMessage
 
 MessageStream = tuple[
-    MemoryObjectReceiveStream[JSONRPCMessage | Exception],
-    MemoryObjectSendStream[JSONRPCMessage],
+    MemoryObjectReceiveStream[ParsedMessage | Exception],
+    MemoryObjectSendStream[ParsedMessage],
 ]
 
 
@@ -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
+        ParsedMessage | Exception
     ](1)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage | Exception
+        ParsedMessage | 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(
diff --git a/src/mcp/shared/session.py b/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
+from pydantic import BaseModel, RootModel
 from typing_extensions import Self
 
 from mcp.shared.exceptions import McpError
@@ -28,6 +28,22 @@
     ServerResult,
 )
 
+RawT = TypeVar("RawT")
+
+
+class ParsedMessage(RootModel[JSONRPCMessage], Generic[RawT]):
+    root: JSONRPCMessage
+    raw: RawT | None = None
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+ReadStream = MemoryObjectReceiveStream[ParsedMessage[RawT] | Exception]
+ReadStreamWriter = MemoryObjectSendStream[ParsedMessage[RawT] | Exception]
+WriteStream = MemoryObjectSendStream[ParsedMessage[RawT]]
+WriteStreamReader = MemoryObjectReceiveStream[ParsedMessage[RawT]]
+
 SendRequestT = TypeVar("SendRequestT", ClientRequest, ServerRequest)
 SendResultT = TypeVar("SendResultT", ClientResult, ServerResult)
 SendNotificationT = TypeVar("SendNotificationT", ClientNotification, ServerNotification)
@@ -165,8 +181,8 @@ class BaseSession(
 
     def __init__(
         self,
-        read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
-        write_stream: MemoryObjectSendStream[JSONRPCMessage],
+        read_stream: ReadStream,
+        write_stream: WriteStream,
         receive_request_type: type[ReceiveRequestT],
         receive_notification_type: type[ReceiveNotificationT],
         # If none, reading will never time out
@@ -242,7 +258,9 @@ async def send_request(
 
         # TODO: Support progress callbacks
 
-        await self._write_stream.send(JSONRPCMessage(jsonrpc_request))
+        await self._write_stream.send(
+            ParsedMessage(JSONRPCMessage(jsonrpc_request), None)
+        )
 
         try:
             with anyio.fail_after(
@@ -278,14 +296,16 @@ async def send_notification(self, notification: SendNotificationT) -> None:
             **notification.model_dump(by_alias=True, mode="json", exclude_none=True),
         )
 
-        await self._write_stream.send(JSONRPCMessage(jsonrpc_notification))
+        await self._write_stream.send(
+            ParsedMessage(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(JSONRPCMessage(jsonrpc_error))
+            await self._write_stream.send(ParsedMessage(JSONRPCMessage(jsonrpc_error)))
         else:
             jsonrpc_response = JSONRPCResponse(
                 jsonrpc="2.0",
@@ -294,18 +314,23 @@ async def _send_response(
                     by_alias=True, mode="json", exclude_none=True
                 ),
             )
-            await self._write_stream.send(JSONRPCMessage(jsonrpc_response))
+            await self._write_stream.send(
+                ParsedMessage(JSONRPCMessage(jsonrpc_response))
+            )
 
     async def _receive_loop(self) -> None:
         async with (
             self._read_stream,
             self._write_stream,
             self._incoming_message_stream_writer,
         ):
-            async for message in self._read_stream:
-                if isinstance(message, Exception):
-                    await self._incoming_message_stream_writer.send(message)
-                elif isinstance(message.root, JSONRPCRequest):
+            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.root
+                if isinstance(message.root, JSONRPCRequest):
                     validated_request = self._receive_request_type.model_validate(
                         message.root.model_dump(
                             by_alias=True, mode="json", exclude_none=True
