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

Author: dsp-ant

src/mcp/server/lowlevel/server.py

@@ -72,12 +72,11 @@ 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
 from mcp.server.lowlevel.helper_types import ReadResourceContents
-from mcp.server.models import InitializationOptions
+from mcp.server.models import InitializationOptions, ReadStream, WriteStream
 from mcp.server.session import ServerSession
 from mcp.server.stdio import stdio_server as stdio_server
 from mcp.shared.context import RequestContext
@@ -472,8 +471,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

@@ -3,11 +3,15 @@
 and tools.
 """
 
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from pydantic import BaseModel
 
-from mcp.types import (
-    ServerCapabilities,
-)
+from mcp.types import JSONRPCMessage, ServerCapabilities
+
+ReadStream = MemoryObjectReceiveStream[JSONRPCMessage | Exception]
+ReadStreamWriter = MemoryObjectSendStream[JSONRPCMessage | Exception]
+WriteStream = MemoryObjectSendStream[JSONRPCMessage]
+WriteStreamReader = MemoryObjectReceiveStream[JSONRPCMessage]
 
 
 class InitializationOptions(BaseModel):

src/mcp/server/session.py

@@ -42,11 +42,10 @@ 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.server.models import InitializationOptions, ReadStream, WriteStream
 from mcp.shared.session import (
     BaseSession,
     RequestResponder,
@@ -73,8 +72,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,19 @@ 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.server.models import (
+    ReadStream,
+    ReadStreamWriter,
+    WriteStream,
+    WriteStreamReader,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -63,9 +68,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 +88,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)

src/mcp/server/stdio.py

@@ -24,9 +24,14 @@ async def run_server():
 
 import anyio
 import anyio.lowlevel
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 
 import mcp.types as types
+from mcp.server.models import (
+    ReadStream,
+    ReadStreamWriter,
+    WriteStream,
+    WriteStreamReader,
+)
 
 
 @asynccontextmanager
@@ -47,11 +52,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)

src/mcp/server/websocket.py

@@ -2,11 +2,16 @@
 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.server.models import (
+    ReadStream,
+    ReadStreamWriter,
+    WriteStream,
+    WriteStreamReader,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -21,11 +26,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)