mirror of
https://github.com/aljazceru/mcp-python-sdk.git
synced 2025-12-19 06:54:18 +01:00
9d0f2daddb
* refactor: improve typing with memory stream type aliases Move memory stream type definitions to models.py and use them throughout the codebase for better type safety and maintainability. GitHub-Issue:#201 * refactor: move streams to ParsedMessage * refactor: update test files to use ParsedMessage Updates test files to work with the ParsedMessage stream type aliases and fixes a line length issue in test_201_client_hangs_on_logging.py. Github-Issue:#201 * refactor: rename ParsedMessage to MessageFrame for clarity 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com> * refactor: move MessageFrame class to types.py for better code organization 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com> * fix pyright * refactor: update websocket client to use MessageFrame Modified the websocket client to work with the new MessageFrame type, preserving raw message text and properly extracting the root JSON-RPC message when sending. Github-Issue:#204 * fix: use NoneType instead of None for type parameters in MessageFrame 🤖 Generated with [Claude Code](https://claude.ai/code) * refactor: rename root to message
89 lines
2.9 KiB
Python
89 lines
2.9 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.client.session import ClientSession, ListRootsFnT, SamplingFnT
|
|
from mcp.server import Server
|
|
from mcp.types import MessageFrame
|
|
|
|
MessageStream = tuple[
|
|
MemoryObjectReceiveStream[MessageFrame | Exception],
|
|
MemoryObjectSendStream[MessageFrame],
|
|
]
|
|
|
|
|
|
@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[
|
|
MessageFrame | Exception
|
|
](1)
|
|
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
|
|
MessageFrame | 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,
|
|
sampling_callback: SamplingFnT | None = None,
|
|
list_roots_callback: ListRootsFnT | 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_read, client_write),
|
|
(server_read, server_write),
|
|
):
|
|
# 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,
|
|
) as client_session:
|
|
await client_session.initialize()
|
|
yield client_session
|
|
finally:
|
|
tg.cancel_scope.cancel()
|