Add in-memory transport (#25)

## Goal

Support running an MCP server in the same process as the client, while preserving MCP abstractions.

## Details

1. **(core change)** Adds a new `memory` transport module that enables in-process client-server communication.
This includes:
   - `create_client_server_memory_streams()` to create bidirectional memory streams
   - `create_connected_server_and_client_session()` to establish an in-process client-server connection

3. (minor) Enhances error handling and timeout support:
   - Adds configurable read timeouts to sessions via `read_timeout_seconds` parameter
   - Improves exception handling in the server with a new `raise_exceptions` flag to control whether errors are returned to clients or raised directly
   - Ensures proper cleanup of request context tokens in error cases

4. (minor) Makes server improvements:
   - Adds built-in ping handler support
This commit is contained in:
Nick Merrill 2024-11-05 18:42:41 -05:00 • committed by GitHub
commit 60e9c7a0d7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1,8 +1,10 @@
from contextlib import AbstractAsyncContextManager
from datetime import timedelta
from typing import Generic, TypeVar
import anyio
import anyio.lowlevel
import httpx
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic import BaseModel
@ -87,6 +89,8 @@ class BaseSession(
write_stream: MemoryObjectSendStream[JSONRPCMessage],
receive_request_type: type[ReceiveRequestT],
receive_notification_type: type[ReceiveNotificationT],
# If none, reading will never time out
read_timeout_seconds: timedelta | None = None,
) -> None:
self._read_stream = read_stream
self._write_stream = write_stream
@ -94,6 +98,7 @@ class BaseSession(
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 +152,25 @@ class BaseSession(
await self._write_stream.send(JSONRPCMessage(jsonrpc_request))
response_or_error = await response_stream_reader.receive()
try:
with anyio.fail_after(
None if self._read_timeout_seconds is None
else self._read_timeout_seconds.total_seconds()
):
response_or_error = await response_stream_reader.receive()
except TimeoutError:
raise McpError(
ErrorData(
code=httpx.codes.REQUEST_TIMEOUT,
message=(
f"Timed out while waiting for response to "
f"{request.__class__.__name__}. Waited "
f"{self._read_timeout_seconds} seconds."
),
)
)
if isinstance(response_or_error, JSONRPCError):
raise McpError(response_or_error.error)
else: