add in memory client

This commit is contained in:
Jeremiah Lowin 2025-04-07 11:23:29 -04:00
commit 4f6aefed7a
2 changed files with 63 additions and 2 deletions

View file

@ -1,5 +1,12 @@
from .websocket import WebSocketClient
from .sse import SSEClient
from .stdio import StdioClient
from .stdio import StdioClient, UvxClient
from .memory import InMemoryClient
__all__ = ["StdioClient", "SSEClient", "WebSocketClient"]
__all__ = [
"StdioClient",
"SSEClient",
"WebSocketClient",
"UvxClient",
"InMemoryClient",
]

View file

@ -0,0 +1,54 @@
import contextlib
from typing import Any, TypeVar
from mcp.server import Server
from mcp.shared.memory import create_connected_server_and_client_session
from typing_extensions import Unpack
from fastmcp.client.base import BaseClient, ClientKwargs
T = TypeVar("T")
class InMemoryClient(BaseClient):
"""Client that connects to an in-memory MCP server.
This client creates and manages an in-memory connection to a server,
without using any external processes or network connections.
"""
def __init__(
self,
server: Server[Any],
raise_exceptions: bool = False,
**kwargs: Unpack[ClientKwargs],
):
"""Initialize an InMemoryClient that connects to an in-memory MCP server.
Args:
server: The MCP server instance to connect to
raise_exceptions: Whether to raise exceptions from the server
**kwargs: Additional arguments for BaseClient
"""
super().__init__(**kwargs)
self.server = server
self.raise_exceptions = raise_exceptions
self._cm_session = None
@contextlib.asynccontextmanager
async def _connect(self):
"""Set up in-memory connection and session"""
self._cm_session = create_connected_server_and_client_session(
server=self.server,
read_timeout_seconds=self._read_timeout_seconds,
sampling_callback=self._sampling_callback,
list_roots_callback=self._list_roots_callback,
logging_callback=self._logging_callback,
message_handler=self._message_handler,
raise_exceptions=self.raise_exceptions,
)
async with self._cm_session as session:
# No need to call initialize as create_connected_server_and_client_session already does
async with self._set_session((None, None), session):
yield self