mirror of
https://github.com/aljazceru/mcp-python-sdk.git
synced 2025-12-19 14:54:24 +01:00
60e9c7a0d7
## 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
88 lines
2.8 KiB
Python
88 lines
2.8 KiB
Python
"""
|
|
In-memory transports
|
|
"""
|
|
|
|
from contextlib import asynccontextmanager
|
|
from datetime import timedelta
|
|
from typing import AsyncGenerator
|
|
|
|
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.types import JSONRPCMessage
|
|
|
|
MessageStream = tuple[
|
|
MemoryObjectReceiveStream[JSONRPCMessage | Exception],
|
|
MemoryObjectSendStream[JSONRPCMessage]
|
|
]
|
|
|
|
@asynccontextmanager
|
|
async def create_client_server_memory_streams() -> AsyncGenerator[
|
|
tuple[MessageStream, MessageStream],
|
|
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)
|
|
|
|
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,
|
|
read_timeout_seconds: timedelta | None = None,
|
|
raise_exceptions: bool = False,
|
|
) -> AsyncGenerator[ClientSession, None]:
|
|
"""Creates a ClientSession that is connected to a running MCP server."""
|
|
async with create_client_server_memory_streams() as (
|
|
client_streams,
|
|
server_streams,
|
|
):
|
|
client_read, client_write = client_streams
|
|
server_read, server_write = server_streams
|
|
|
|
# Create a cancel scope for the server task
|
|
async with anyio.create_task_group() as tg:
|
|
tg.start_soon(
|
|
lambda: server.run(
|
|
server_read,
|
|
server_write,
|
|
server.create_initialization_options(),
|
|
raise_exceptions=raise_exceptions,
|
|
)
|
|
)
|
|
|
|
try:
|
|
async with ClientSession(
|
|
read_stream=client_read,
|
|
write_stream=client_write,
|
|
read_timeout_seconds=read_timeout_seconds,
|
|
) as client_session:
|
|
await client_session.initialize()
|
|
yield client_session
|
|
finally:
|
|
tg.cancel_scope.cancel()
|