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:
parent
1a60e1b7c7
commit
60e9c7a0d7
1 changed files with 210 additions and 10 deletions
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Reference in a new issue