diff --git a/src/fastmcp/client/__init__.py b/src/fastmcp/client/__init__.py new file mode 100644 index 000000000..ceb3f12a0 --- /dev/null +++ b/src/fastmcp/client/__init__.py @@ -0,0 +1,5 @@ +from .websocket import WebSocketClient +from .sse import SSEClient +from .stdio import StdioClient + +__all__ = ["StdioClient", "SSEClient", "WebSocketClient"] diff --git a/src/fastmcp/client/base.py b/src/fastmcp/client/base.py new file mode 100644 index 000000000..53316f706 --- /dev/null +++ b/src/fastmcp/client/base.py @@ -0,0 +1,190 @@ +import abc +import contextlib +import datetime +from typing import Any, AsyncContextManager, Optional + +import mcp.types +from mcp import ClientSession +from mcp.client.session import ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT +from mcp.shared.context import LifespanContextT, RequestContext +from pydantic import AnyUrl + + +def _get_roots_callback(roots: list[mcp.types.Root]) -> ListRootsFnT | None: + async def _roots_callback( + context: RequestContext[ClientSession, LifespanContextT], + ) -> mcp.types.ListRootsResult: + return mcp.types.ListRootsResult(roots=roots) + + return _roots_callback + + +class BaseClient(abc.ABC): + def __init__( + self, + roots: list[mcp.types.Root] | None = None, + sampling_callback: SamplingFnT | None = None, + list_roots_callback: ListRootsFnT | None = None, + logging_callback: LoggingFnT | None = None, + message_handler: MessageHandlerFnT | None = None, + read_timeout_seconds: datetime.timedelta | None = None, + ): + self._transport: Any = None + self._session: Optional[ClientSession] = None + self._cm: Optional[AsyncContextManager] = None + + if roots is not None: + if list_roots_callback is not None: + raise ValueError( + "Cannot provide both `roots` and `list_roots_callback`. " + "Either provide a list of roots or a callback to list roots." + ) + else: + list_roots_callback = _get_roots_callback(roots) + + self._sampling_callback = sampling_callback + self._list_roots_callback = list_roots_callback + self._logging_callback = logging_callback + self._message_handler = message_handler + self._read_timeout_seconds = read_timeout_seconds + + @property + def transport(self): + """Get the current transport connection""" + if self._transport is None: + raise RuntimeError( + "Client is not connected. Use 'async with client:' context manager first." + ) + return self._transport + + @property + def session(self): + """Get the current session""" + if self._session is None: + raise RuntimeError( + "Client is not connected. Use 'async with client:' context manager first." + ) + return self._session + + def is_connected(self): + """Check if the client is currently connected""" + return self._session is not None + + @abc.abstractmethod + def _connect( + self, + sampling_callback: SamplingFnT | None = None, + list_roots_callback: ListRootsFnT | None = None, + logging_callback: LoggingFnT | None = None, + message_handler: MessageHandlerFnT | None = None, + ) -> AsyncContextManager: + """Return an async context manager that handles connection lifecycle. + This will be called by __aenter__ to establish the connection.""" + raise NotImplementedError("Subclasses must implement this method") + + @contextlib.asynccontextmanager + async def _create_connection_context(self): + """Create and manage the connection context if not already connected. + This handles both creating a new connection or reusing an existing one.""" + created_connection = False + try: + if not self.is_connected(): + # Only create a new connection if not already connected + self._cm = self._connect() + await self._cm.__aenter__() + created_connection = True + yield + finally: + if created_connection and self._cm is not None: + # Only close if we created the connection in this context + await self._cm.__aexit__(None, None, None) + self._transport = None + self._session = None + self._cm = None + + @contextlib.asynccontextmanager + async def _set_session(self, transport: Any, session: ClientSession): + self._transport = transport + self._session = session + try: + await self._session.initialize() + yield + finally: + self._transport = None + self._session = None + + async def __aenter__(self): + self._connection_ctx = self._create_connection_context() + await self._connection_ctx.__aenter__() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self._connection_ctx.__aexit__(exc_type, exc_val, exc_tb) + + # --- MCP Client Methods --- + + async def ping(self) -> None: + """Send a ping request.""" + await self.session.send_ping() + + async def progress( + self, progress_token: str | int, progress: float, total: float | None = None + ) -> None: + """Send a progress notification.""" + await self.session.send_progress_notification(progress_token, progress, total) + + async def set_logging_level(self, level: mcp.types.LoggingLevel) -> None: + """Send a logging/setLevel request.""" + await self.session.set_logging_level(level) + + async def list_resources(self) -> mcp.types.ListResourcesResult: + """Send a resources/list request.""" + return await self.session.list_resources() + + async def list_resource_templates(self) -> mcp.types.ListResourceTemplatesResult: + """Send a resources/listResourceTemplates request.""" + return await self.session.list_resource_templates() + + async def read_resource(self, uri: AnyUrl) -> mcp.types.ReadResourceResult: + """Send a resources/read request.""" + return await self.session.read_resource(uri) + + async def subscribe_resource(self, uri: AnyUrl) -> None: + """Send a resources/subscribe request.""" + await self.session.subscribe_resource(uri) + + async def unsubscribe_resource(self, uri: AnyUrl) -> None: + """Send a resources/unsubscribe request.""" + await self.session.unsubscribe_resource(uri) + + async def list_prompts(self) -> mcp.types.ListPromptsResult: + """Send a prompts/list request.""" + return await self.session.list_prompts() + + async def get_prompt( + self, name: str, arguments: dict[str, str] | None = None + ) -> mcp.types.GetPromptResult: + """Send a prompts/get request.""" + return await self.session.get_prompt(name, arguments) + + async def complete( + self, + ref: mcp.types.ResourceReference | mcp.types.PromptReference, + argument: dict[str, str], + ) -> mcp.types.CompleteResult: + """Send a completion/complete request.""" + return await self.session.complete(ref, argument) + + async def list_tools(self) -> mcp.types.ListToolsResult: + """Send a tools/list request.""" + return await self.session.list_tools() + + async def call_tool( + self, name: str, arguments: dict[str, Any] | None = None + ) -> mcp.types.CallToolResult: + """Send a tools/call request.""" + return await self.session.call_tool(name, arguments) + + async def send_roots_list_changed(self) -> None: + """Send a roots/list_changed notification.""" + await self.session.send_roots_list_changed() diff --git a/src/fastmcp/client/sse.py b/src/fastmcp/client/sse.py new file mode 100644 index 000000000..de01fe0fb --- /dev/null +++ b/src/fastmcp/client/sse.py @@ -0,0 +1,55 @@ +import contextlib +import datetime + +import mcp.types +from mcp import ClientSession +from mcp.client.sse import sse_client + +from fastmcp.client.base import ( + BaseClient, + ListRootsFnT, + LoggingFnT, + MessageHandlerFnT, + SamplingFnT, +) + + +class SSEClient(BaseClient): + def __init__( + self, + url: str, + headers: dict[str, str] | None = None, + roots: list[mcp.types.Root] | None = None, + sampling_callback: SamplingFnT | None = None, + list_roots_callback: ListRootsFnT | None = None, + logging_callback: LoggingFnT | None = None, + message_handler: MessageHandlerFnT | None = None, + read_timeout_seconds: datetime.timedelta | None = None, + ): + super().__init__( + roots=roots, + sampling_callback=sampling_callback, + list_roots_callback=list_roots_callback, + logging_callback=logging_callback, + message_handler=message_handler, + read_timeout_seconds=read_timeout_seconds, + ) + self.url = url + self.headers = headers or {} + + @contextlib.asynccontextmanager + async def _connect(self): + """Set up SSE connection and session""" + async with sse_client(self.url, headers=self.headers) as transport: + read_stream, write_stream = transport + async with ClientSession( + read_stream=read_stream, + write_stream=write_stream, + sampling_callback=self._sampling_callback, + list_roots_callback=self._list_roots_callback, + logging_callback=self._logging_callback, + message_handler=self._message_handler, + read_timeout_seconds=self._read_timeout_seconds, + ) as session: + async with self._set_session(transport, session): + yield self diff --git a/src/fastmcp/client/stdio.py b/src/fastmcp/client/stdio.py new file mode 100644 index 000000000..9c5f882e4 --- /dev/null +++ b/src/fastmcp/client/stdio.py @@ -0,0 +1,64 @@ +import contextlib +import datetime + +import mcp.types +from mcp import ClientSession, StdioServerParameters +from mcp.client.stdio import stdio_client + +from fastmcp.client.base import ( + BaseClient, + ListRootsFnT, + LoggingFnT, + MessageHandlerFnT, + SamplingFnT, +) + + +class StdioClient(BaseClient): + def __init__( + self, + server_script_path: str, + roots: list[mcp.types.Root] | None = None, + sampling_callback: SamplingFnT | None = None, + list_roots_callback: ListRootsFnT | None = None, + logging_callback: LoggingFnT | None = None, + message_handler: MessageHandlerFnT | None = None, + read_timeout_seconds: datetime.timedelta | None = None, + ): + super().__init__( + roots=roots, + sampling_callback=sampling_callback, + list_roots_callback=list_roots_callback, + logging_callback=logging_callback, + message_handler=message_handler, + read_timeout_seconds=read_timeout_seconds, + ) + self.server_script_path = server_script_path + + @contextlib.asynccontextmanager + async def _connect(self): + """Set up stdio connection and session""" + is_python = self.server_script_path.endswith(".py") + is_js = self.server_script_path.endswith(".js") + if not (is_python or is_js): + raise ValueError("Server script must be a .py or .js file") + + command = "python" if is_python else "node" + server_params = StdioServerParameters( + command=command, args=[self.server_script_path], env=None + ) + + async with stdio_client(server_params) as transport: + stdio, write = transport + + async with ClientSession( + read_stream=stdio, + write_stream=write, + sampling_callback=self._sampling_callback, + list_roots_callback=self._list_roots_callback, + logging_callback=self._logging_callback, + message_handler=self._message_handler, + read_timeout_seconds=self._read_timeout_seconds, + ) as session: + async with self._set_session(transport, session): + yield self diff --git a/src/fastmcp/client/websocket.py b/src/fastmcp/client/websocket.py new file mode 100644 index 000000000..c66173ec4 --- /dev/null +++ b/src/fastmcp/client/websocket.py @@ -0,0 +1,54 @@ +import contextlib +import datetime + +import mcp.types +from mcp import ClientSession +from mcp.client.websocket import websocket_client + +from fastmcp.client.base import ( + BaseClient, + ListRootsFnT, + LoggingFnT, + MessageHandlerFnT, + SamplingFnT, +) + + +class WebSocketClient(BaseClient): + def __init__( + self, + url: str, + roots: list[mcp.types.Root] | None = None, + sampling_callback: SamplingFnT | None = None, + list_roots_callback: ListRootsFnT | None = None, + logging_callback: LoggingFnT | None = None, + message_handler: MessageHandlerFnT | None = None, + read_timeout_seconds: datetime.timedelta | None = None, + ): + super().__init__( + roots=roots, + sampling_callback=sampling_callback, + list_roots_callback=list_roots_callback, + logging_callback=logging_callback, + message_handler=message_handler, + read_timeout_seconds=read_timeout_seconds, + ) + self.url = url + + @contextlib.asynccontextmanager + async def _connect(self): + """Set up WebSocket connection and session""" + async with websocket_client(self.url) as transport: + read_stream, write_stream = transport + + async with ClientSession( + read_stream=read_stream, + write_stream=write_stream, + sampling_callback=self._sampling_callback, + list_roots_callback=self._list_roots_callback, + logging_callback=self._logging_callback, + message_handler=self._message_handler, + read_timeout_seconds=self._read_timeout_seconds, + ) as session: + async with self._set_session(transport, session): + yield self