100 lines
3.5 KiB
Python
100 lines
3.5 KiB
Python
"""
|
|
In-memory transports
|
|
"""
|
|
|
|
from collections.abc import AsyncGenerator
|
|
from contextlib import asynccontextmanager
|
|
from datetime import timedelta
|
|
from typing import Any
|
|
|
|
import anyio
|
|
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
|
|
|
import mcp.types as types
|
|
from mcp.client.session import (
|
|
ClientSession,
|
|
ElicitationFnT,
|
|
ListRootsFnT,
|
|
LoggingFnT,
|
|
MessageHandlerFnT,
|
|
SamplingFnT,
|
|
)
|
|
from mcp.server import Server
|
|
from mcp.shared.message import SessionMessage
|
|
|
|
MessageStream = tuple[MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage]]
|
|
|
|
|
|
@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[SessionMessage | Exception](1)
|
|
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | 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[Any],
|
|
read_timeout_seconds: timedelta | None = None,
|
|
sampling_callback: SamplingFnT | None = None,
|
|
list_roots_callback: ListRootsFnT | None = None,
|
|
logging_callback: LoggingFnT | None = None,
|
|
message_handler: MessageHandlerFnT | None = None,
|
|
client_info: types.Implementation | None = None,
|
|
raise_exceptions: bool = False,
|
|
elicitation_callback: ElicitationFnT | None = None,
|
|
) -> 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,
|
|
sampling_callback=sampling_callback,
|
|
list_roots_callback=list_roots_callback,
|
|
logging_callback=logging_callback,
|
|
message_handler=message_handler,
|
|
client_info=client_info,
|
|
elicitation_callback=elicitation_callback,
|
|
) as client_session:
|
|
await client_session.initialize()
|
|
yield client_session
|
|
finally:
|
|
tg.cancel_scope.cancel()
|