add clients

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"]

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()

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

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

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