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

This reverts commit 9d0f2daddb.

src/mcp/server/lowlevel/server.py

@@ -74,6 +74,7 @@ from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontext
 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
@@ -83,7 +84,7 @@ 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 ReadStream, RequestResponder, WriteStream
+from mcp.shared.session import RequestResponder
 
 logger = logging.getLogger(__name__)
 
@@ -473,8 +474,8 @@ class Server(Generic[LifespanResultT]):
 
     async def run(
         self,
-        read_stream: ReadStream,
-        write_stream: WriteStream,
+        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
+        write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
         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,7 +5,9 @@ and tools.
 
 from pydantic import BaseModel
 
-from mcp.types import ServerCapabilities
+from mcp.types import (
+    ServerCapabilities,
+)
 
 
 class InitializationOptions(BaseModel):

src/mcp/server/session.py

@@ -42,15 +42,14 @@ from typing import Any, TypeVar
 
 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,
 )
 
 
@@ -77,8 +76,8 @@ class ServerSession(
 
     def __init__(
         self,
-        read_stream: ReadStream,
-        write_stream: WriteStream,
+        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
+        write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
         init_options: InitializationOptions,
     ) -> None:
         super().__init__(

src/mcp/server/sse.py

@@ -38,6 +38,7 @@ from urllib.parse import quote
 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
@@ -45,13 +46,6 @@ 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__)
 
@@ -69,7 +63,9 @@ class SseServerTransport:
     """
 
     _endpoint: str
-    _read_stream_writers: dict[UUID, ReadStreamWriter]
+    _read_stream_writers: dict[
+        UUID, MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+    ]
 
     def __init__(self, endpoint: str) -> None:
         """
@@ -89,11 +85,11 @@ class SseServerTransport:
             raise ValueError("connect_sse can only handle HTTP requests")
 
         logger.debug("Setting up SSE connection")
-        read_stream: ReadStream
-        read_stream_writer: ReadStreamWriter
+        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
+        read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
 
-        write_stream: WriteStream
-        write_stream_reader: WriteStreamReader
+        write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
+        write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
 
         read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
         write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@@ -176,4 +172,4 @@ class SseServerTransport:
         logger.debug(f"Sending message to writer: {message}")
         response = Response("Accepted", status_code=202)
         await response(scope, receive, send)
-        await writer.send(MessageFrame(message=message, raw=request))
+        await writer.send(message)

src/mcp/server/stdio.py

@@ -24,15 +24,9 @@ from io import TextIOWrapper
 
 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
@@ -53,11 +47,11 @@ async def stdio_server(
     if not stdout:
         stdout = anyio.wrap_file(TextIOWrapper(sys.stdout.buffer, encoding="utf-8"))
 
-    read_stream: ReadStream
-    read_stream_writer: ReadStreamWriter
+    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
 
-    write_stream: WriteStream
-    write_stream_reader: WriteStreamReader
+    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
+    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
 
     read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@@ -72,9 +66,7 @@ async def stdio_server(
                         await read_stream_writer.send(exc)
                         continue
 
-                    await read_stream_writer.send(
-                        MessageFrame(message=message, raw=line)
-                    )
+                    await read_stream_writer.send(message)
         except anyio.ClosedResourceError:
             await anyio.lowlevel.checkpoint()
 
@@ -82,7 +74,6 @@ async def stdio_server(
         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,17 +2,11 @@ import logging
 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__)
 
@@ -27,11 +21,11 @@ async def websocket_server(scope: Scope, receive: Receive, send: Send):
     websocket = WebSocket(scope, receive, send)
     await websocket.accept(subprotocol="mcp")
 
-    read_stream: ReadStream
-    read_stream_writer: ReadStreamWriter
+    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
 
-    write_stream: WriteStream
-    write_stream_reader: WriteStreamReader
+    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
+    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
 
     read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@@ -46,9 +40,7 @@ async def websocket_server(scope: Scope, receive: Receive, send: Send):
                         await read_stream_writer.send(exc)
                         continue
 
-                    await read_stream_writer.send(
-                        MessageFrame(message=client_message, raw=message)
-                    )
+                    await read_stream_writer.send(client_message)
         except anyio.ClosedResourceError:
             await websocket.close()