Add in-memory transport

Author: nick-merrill

mcp_python/client/session.py

@@ -36,8 +36,9 @@ def __init__(
         self,
         read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
         write_stream: MemoryObjectSendStream[JSONRPCMessage],
+        read_timeout_seconds: int | float | None = None,
     ) -> None:
-        super().__init__(read_stream, write_stream, ServerRequest, ServerNotification)
+        super().__init__(read_stream, write_stream, ServerRequest, ServerNotification, read_timeout_seconds=read_timeout_seconds)
 
     async def initialize(self) -> InitializeResult:
         from mcp_python.types import (

mcp_python/server/__init__.py

@@ -63,9 +63,13 @@ def pkg_version(package: str) -> str:
             try:
                 from importlib.metadata import version
 
-                return version(package)
+                v = version(package)
+                if v is not None:
+                    return v
             except Exception:
-                return "unknown"
+                pass
+
+            return "unknown"
 
         return types.InitializationOptions(
             server_name=self.name,
@@ -330,6 +334,11 @@ async def run(
         read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
         write_stream: MemoryObjectSendStream[JSONRPCMessage],
         initialization_options: types.InitializationOptions,
+        # When True, exceptions are returned as messages to the client.
+        # When False, exceptions are raised, which will cause the server to shut down
+        # but also make tracing exceptions much easier during testing and when using
+        # in-process servers.
+        raise_exceptions: bool = False,
     ):
         with warnings.catch_warnings(record=True) as w:
             async with ServerSession(
@@ -349,6 +358,7 @@ async def run(
                                     f"Dispatching request of type {type(req).__name__}"
                                 )
 
+                                token = None
                                 try:
                                     # Set our global state that can be retrieved via
                                     # app.get_request_context()
@@ -360,20 +370,24 @@ async def run(
                                         )
                                     )
                                     response = await handler(req)
-                                    # Reset the global state after we are done
-                                    request_ctx.reset(token)
                                 except Exception as err:
+                                    if raise_exceptions:
+                                        raise err
                                     response = ErrorData(
                                         code=0, message=str(err), data=None
                                     )
+                                finally:
+                                    # Reset the global state after we are done
+                                    if token is not None:
+                                        request_ctx.reset(token)
 
                                 await message.respond(response)
                             else:
                                 await message.respond(
                                     ErrorData(
                                         code=METHOD_NOT_FOUND,
                                         message="Method not found",
-                                    )
+                                    ),
                                 )
 
                             logger.debug("Response sent")

mcp_python/server/memory.py

@@ -0,0 +1,71 @@
+"""
+In-memory transports
+"""
+
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator, Tuple
+
+import anyio
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+
+from mcp_python.client.session import ClientSession
+from mcp_python.server import Server
+from mcp_python.server.session import ServerSession
+from mcp_python.types import ErrorData, JSONRPCMessage
+
+@asynccontextmanager
+async def create_client_server_memory_streams() -> AsyncGenerator[ Tuple[
+    Tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]],
+    Tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]
+], None]:
+    """
+    Creates a pair of bidirectional memory streams for client-server communication.
+
+    Returns:
+        A tuple of (client_streams, server_streams) where each is a tuple of
+        (read_stream, write_stream)
+    """
+    # Create streams for both directions
+    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[JSONRPCMessage | Exception](1)
+    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[JSONRPCMessage | Exception](1)
+
+    # Return streams grouped by client/server
+    client_streams = (server_to_client_receive, client_to_server_send)
+    server_streams = (client_to_server_receive, server_to_client_send)
+
+    async with (server_to_client_receive, client_to_server_send,
+                client_to_server_receive, server_to_client_send):
+        yield client_streams, server_streams
+
+
+@asynccontextmanager
+async def create_connected_server_and_client_session(server: Server) -> AsyncGenerator[ClientSession, None]:
+    """Creates a ServerSession that is connected to the `server`."""
+    async with create_client_server_memory_streams() as (client_streams, server_streams):
+        # Unpack the streams
+        client_read, client_write = client_streams
+        server_read, server_write = server_streams
+        print("stream-1")
+
+        # Create a cancel scope for the server task
+        async with anyio.create_task_group() as tg:
+
+            tg.start_soon(
+                server.run,
+                server_read,
+                server_write,
+                server.create_initialization_options()
+            )
+
+            print("stream2")
+
+            try:
+                # Client session could be created here using client_read and client_write
+                # This would allow testing the server with a client in the same process
+                async with ClientSession(
+                    read_stream=client_read, write_stream=client_write
+                ) as client_session:
+                    await client_session.initialize()
+                    yield client_session
+            finally:
+                tg.cancel_scope.cancel()

mcp_python/shared/session.py

@@ -87,13 +87,16 @@ def __init__(
         write_stream: MemoryObjectSendStream[JSONRPCMessage],
         receive_request_type: type[ReceiveRequestT],
         receive_notification_type: type[ReceiveNotificationT],
+        # If none, reading will never time out
+        read_timeout_seconds: int | float | None = None,
     ) -> None:
         self._read_stream = read_stream
         self._write_stream = write_stream
         self._response_streams = {}
         self._request_id = 0
         self._receive_request_type = receive_request_type
         self._receive_notification_type = receive_notification_type
+        self._read_timeout_seconds = read_timeout_seconds
 
         self._incoming_message_stream_writer, self._incoming_message_stream_reader = (
             anyio.create_memory_object_stream[
@@ -147,7 +150,13 @@ async def send_request(
 
         await self._write_stream.send(JSONRPCMessage(jsonrpc_request))
 
-        response_or_error = await response_stream_reader.receive()
+        try:
+            with anyio.fail_after(self._read_timeout_seconds):
+                response_or_error = await response_stream_reader.receive()
+        except TimeoutError:
+            # TODO: make sure this response comes back correctly to the client
+            raise McpError(ErrorData(code=408, message=f"Timed out while waiting for response to {request.__class__.__name__}. Waited {READ_TIMEOUT_SECONDS} seconds."))
+
         if isinstance(response_or_error, JSONRPCError):
             raise McpError(response_or_error.error)
         else:

mcp_python/types.py

@@ -1,6 +1,7 @@
 from typing import Any, Generic, Literal, TypeVar
 
 from pydantic import BaseModel, ConfigDict, RootModel
+from pydantic.fields import Field
 from pydantic.networks import AnyUrl
 
 """
@@ -141,16 +142,19 @@ class ErrorData(BaseModel):
 
     code: int
     """The error type that occurred."""
+
     message: str
     """
     A short description of the error. The message SHOULD be limited to a concise single
     sentence.
     """
+
     data: Any | None = None
     """
     Additional information about the error. The value of this member is defined by the
     sender (e.g. detailed error information, nested errors etc.).
     """
+
     model_config = ConfigDict(extra="allow")
 
 

tests/conftest.py

@@ -0,0 +1,30 @@
+import pytest
+from pydantic import HttpUrl
+
+from mcp_python.server import Server
+from mcp_python.server.types import InitializationOptions
+from mcp_python.types import Resource, ServerCapabilities
+
+
+TEST_INITIALIZATION_OPTIONS = InitializationOptions(
+    server_name="my_mcp_server",
+    server_version="0.1.0",
+    capabilities=ServerCapabilities(),
+)
+
+@pytest.fixture
+def mcp_server() -> Server:
+    server = Server(name="test_server")
+
+    # Add a simple resource for testing
+    @server.list_resources()
+    async def handle_list_resources():
+        return [
+            Resource(
+                uri=HttpUrl("memory://test"),
+                name="Test Resource",
+                description="A test resource"
+            )
+        ]
+
+    return server

tests/test_memory.py

@@ -0,0 +1,37 @@
+from typing_extensions import AsyncGenerator
+import anyio
+from pydantic import HttpUrl
+from pydantic_core import Url
+import pytest
+
+from mcp_python.client.session import ClientSession
+from mcp_python.server import Server
+from mcp_python.server.memory import create_client_server_memory_streams, create_connected_server_and_client_session
+from mcp_python.server.session import ServerSession
+from mcp_python.server.types import InitializationOptions
+from mcp_python.types import (
+    ClientNotification,
+    EmptyResult,
+    InitializedNotification,
+    JSONRPCMessage,
+    Resource,
+    ServerCapabilities,
+)
+
+
+@pytest.fixture
+async def client_connected_to_server(mcp_server: Server) -> AsyncGenerator[ClientSession, None]:
+    print('11111')
+    async with create_connected_server_and_client_session(mcp_server) as client_session:
+        print('2222k')
+        yield client_session
+        print('33')
+
+
+@pytest.mark.anyio
+async def test_memory_server_and_client_connection(client_connected_to_server: ClientSession):
+    """Shows how a client and server can communicate over memory streams."""
+    response = await client_connected_to_server.send_ping()
+    print('foo')
+    assert isinstance(response, EmptyResult)
+    print('bar')