diff --git a/src/fastmcp/client/base.py b/src/fastmcp/client/auth/__init__.py similarity index 100% rename from src/fastmcp/client/base.py rename to src/fastmcp/client/auth/__init__.py diff --git a/src/fastmcp/client/auth/httpx_client.py b/src/fastmcp/client/auth/httpx_client.py new file mode 100644 index 000000000..693bd0e3e --- /dev/null +++ b/src/fastmcp/client/auth/httpx_client.py @@ -0,0 +1,424 @@ +from __future__ import annotations + +import asyncio +import socket +import time +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager, contextmanager +from contextvars import ContextVar +from typing import Any +from urllib.parse import urljoin + +import anyio +import httpx +import mcp.client.sse +import mcp.client.streamable_http +import mcp.shared._httpx_utils +from authlib.integrations.httpx_client import AsyncOAuth2Client +from starlette.applications import Starlette +from starlette.responses import PlainTextResponse +from starlette.routing import Route +from uvicorn import Config, Server + +from fastmcp.client.auth.oauth_cache import oauth_cache +from fastmcp.utilities.logging import get_logger + +_current_mcp_endpoint: ContextVar[str | None] = ContextVar("mcp_endpoint", default=None) + + +logger = get_logger(__name__) + + +def create_mcp_http_client( + headers: dict[str, Any] | None = None, + timeout: httpx.Timeout | None = None, + **kwargs: Any, +) -> httpx.AsyncClient: + # re-implements logic from mcp.shared._httpx_utils.create_mcp_http_client, but with **kwargs support + kwargs.setdefault("follow_redirects", True) + if timeout is None: + timeout = httpx.Timeout(30.0) + + return httpx.AsyncClient(headers=headers, timeout=timeout, **kwargs) + + +def find_available_port() -> int: + """Find an available port by letting the OS assign one.""" + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +async def _get_redirect( + port: int, path: str = "/callback", timeout: float = 100.0 +) -> str: + """ + Start a temporary server to handle OAuth redirect and get the full redirect URL. + + Args: + port: The port to run the server on + path: The path to listen for redirects on + timeout: Number of seconds to wait before timing out + + Returns: + The full redirect URL from the browser + + Raises: + TimeoutError: If no redirect is received within the timeout period + """ + fut = asyncio.get_running_loop().create_future() + + async def cb(request): + if not fut.done(): + fut.set_result(str(request.url)) # full redirect URL + return PlainTextResponse( + "✅ FastMCP login complete! You can close this tab now." + ) + + server = Server( + Config( + app=Starlette(routes=[Route(path, cb)]), + host="127.0.0.1", + port=port, + lifespan="off", + log_level="error", + ) + ) + + async with anyio.create_task_group() as tg: + tg.start_soon(server.serve) # background task for server + + try: + # Use anyio.fail_after to implement timeout + with anyio.fail_after(timeout): + redirect_url = await fut # wait for browser hit or timeout + return redirect_url + finally: + server.should_exit = True # stop the server loop + tg.cancel_scope.cancel() # tear down immediately + + +class OAuthBearerAuth(httpx.Auth): + """Auth handler that adds the OAuth bearer token to requests.""" + + def __init__(self, client: AsyncOAuth2Client) -> None: + self._client = client + + async def async_auth_flow(self, request: httpx.Request): + # Ensure token is loaded in the OAuth client + if ( + not self._client.token + or self._client.token.get("expires_at") + and self._client.token["expires_at"] < time.time() + ): + # We'll refresh or reauthorize in create_mcp_oauth_client + pass + + if self._client.token: + request.headers["Authorization"] = ( + f"Bearer {self._client.token['access_token']}" + ) + yield request + + +async def discover_oauth_metadata(base_url: str) -> dict[str, Any] | None: + """ + Discover OAuth metadata from the server according to RFC 8414. + + Returns None if the server appears to not require authentication. + """ + # First, try the well-known URL + well_known_url = urljoin(base_url, "/.well-known/oauth-authorization-server") + logger.debug(f"Attempting OAuth metadata discovery from: {well_known_url}") + + async with httpx.AsyncClient() as client: + # First try the well-known URL + try: + response = await client.get(well_known_url, timeout=10) + if response.status_code == 200: + logger.debug("Successfully discovered OAuth metadata") + return response.json() + except httpx.RequestError as e: + logger.debug(f"Failed to fetch OAuth metadata: {e}") + + # If well-known discovery fails, check WWW-Authenticate header + try: + response = await client.get(base_url, timeout=10) + + # If the base URL request succeeds without a 401/403 and has no WWW-Authenticate header, + # the server likely doesn't require authentication + if ( + response.status_code < 400 + and "WWW-Authenticate" not in response.headers + ): + logger.debug("Server appears to not require authentication") + return None + + auth_header = response.headers.get("WWW-Authenticate") + if auth_header and "resource_metadata" in auth_header: + # Extract metadata URL from header + import re + + metadata_match = re.search(r'resource_metadata="([^"]+)"', auth_header) + if metadata_match: + metadata_url = metadata_match.group(1) + metadata_response = await client.get(metadata_url, timeout=10) + if metadata_response.status_code == 200: + logger.debug( + "Successfully discovered OAuth metadata from WWW-Authenticate header" + ) + return metadata_response.json() + except httpx.RequestError as e: + logger.debug(f"Failed to fetch OAuth metadata from WWW-Authenticate: {e}") + + # Fallback to default endpoints based on the base URL + logger.debug("Falling back to default OAuth endpoints") + return { + "issuer": base_url, + "authorization_endpoint": urljoin(base_url, "/authorize"), + "token_endpoint": urljoin(base_url, "/token"), + "registration_endpoint": urljoin(base_url, "/register"), + "response_types_supported": ["code"], + "response_modes_supported": ["query"], + "grant_types_supported": ["authorization_code", "refresh_token"], + "token_endpoint_auth_methods_supported": [ + "client_secret_basic", + "client_secret_post", + "none", + ], + "code_challenge_methods_supported": ["S256"], + } + + +async def register_client( + registration_endpoint: str, redirect_uri: str +) -> dict[str, Any]: + """ + Register an OAuth client using RFC 7591 dynamic registration. + + May raise httpx.HTTPStatusError if registration fails. + """ + logger.debug(f"Registering client at: {registration_endpoint}") + + payload = { + "client_name": "FastMCP Client", + "redirect_uris": [redirect_uri], + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + "token_endpoint_auth_method": "none", # public PKCE client + } + + async with httpx.AsyncClient() as client: + response = await client.post(registration_endpoint, json=payload, timeout=10) + # Allow HTTPStatusError to propagate to the caller + response.raise_for_status() + logger.debug("Client registration successful") + return response.json() + + +@asynccontextmanager +async def create_mcp_oauth_client( + mcp_endpoint: str, + redirect_uri: str | None = None, + scope: list[str] | None = None, + headers: dict[str, Any] | None = None, + timeout: httpx.Timeout | None = None, + **httpx_kwargs: Any, +) -> AsyncIterator[httpx.AsyncClient]: + """ + Create an authenticated OAuth client for an MCP server from an endpoint URL. + + This function handles: + 1. OAuth metadata discovery + 2. Dynamic client registration if needed + 3. Authorization code flow with PKCE + 4. Token refreshing + 5. Token persistence + + If the server doesn't require authentication, a regular client will be returned. + + Args: + mcp_endpoint: Full URL to an MCP endpoint (e.g., + https://mcp.example.com/sse). This will be used to discover the OAuth + configuration. + redirect_uri: OAuth redirect URI for the authorization flow. If None, + a server will be started on an available port. + scope: OAuth scopes to request + headers: Additional headers to include in the requests + timeout: Timeout for the requests + **httpx_kwargs: Additional arguments for the httpx client + + Returns: + An httpx.AsyncClient that handles authentication automatically + """ + # Extract base URL for OAuth discovery + base_url = oauth_cache.get_base_url(mcp_endpoint) + logger.debug(f"MCP Endpoint: {mcp_endpoint}") + logger.debug(f"Base URL for OAuth: {base_url}") + + # Discover OAuth metadata + metadata = await discover_oauth_metadata(base_url) + + # If metadata is None, the server doesn't require authentication + if metadata is None: + logger.info("Server doesn't require authentication, creating regular client") + async with create_mcp_http_client( + headers=headers, + timeout=timeout, + **httpx_kwargs, + ) as client: + yield client + return + + logger.debug(f"Using OAuth endpoints: {metadata}") + + # Use dynamic redirect URI if none provided + # Generate port only once and reuse it for all operations + port = None + if redirect_uri is None: + port = find_available_port() + redirect_uri = f"http://127.0.0.1:{port}/callback" + logger.debug(f"Using dynamic redirect URI: {redirect_uri}") + + # Load or register client - check if we need to update registration due to new redirect URI + creds = oauth_cache.load(mcp_endpoint, "client") + if creds and redirect_uri not in creds.get("redirect_uris", []): + logger.debug("Redirect URI not in registered URIs, re-registering client") + creds = None # Force re-registration + + # Register if needed + if not creds: + try: + creds = await register_client( + metadata["registration_endpoint"], redirect_uri + ) + oauth_cache.save(mcp_endpoint, creds, "client") + logger.debug(f"Client registered with redirect URI: {redirect_uri}") + except httpx.HTTPStatusError as e: + if e.response.status_code == 404: + # If registration endpoint returns 404, server likely doesn't support OAuth + logger.info("Registration endpoint not found, creating regular client") + async with create_mcp_http_client( + headers=headers, + timeout=timeout, + **httpx_kwargs, + ) as client: + yield client + return + else: + # Other HTTP errors should be propagated + raise + + # Create the OAuth client + oauth_client = AsyncOAuth2Client( + client_id=creds["client_id"], + client_secret=creds.get("client_secret"), # "public" clients omit secret + scope=scope or ["openid", "profile", "email"], + redirect_uri=redirect_uri, + **httpx_kwargs, + ) + + # Load token if exists - passing mcp_endpoint directly + token = oauth_cache.load(mcp_endpoint, "token") + if token: + oauth_client.token = token + + try: + # Ensure we have a valid token + if ( + not oauth_client.token + or not oauth_client.token.get("expires_at") + or oauth_client.token["expires_at"] < time.time() + ): + # Try to refresh if possible + if oauth_client.token and oauth_client.token.get("refresh_token"): + logger.debug("Refreshing token") + try: + # ignore type because refresh_token is awaitable but not typed as such + token = await oauth_client.refresh_token( # type: ignore[await-expr] + url=metadata["token_endpoint"], + refresh_token=oauth_client.token["refresh_token"], + ) + except Exception as e: + logger.warning(f"Failed to refresh token: {e}") + token = None + else: + token = None + + # If token is still not available, start authorization flow + if not token: + # Start authorization flow with PKCE + logger.info("Starting authorization flow") + uri, _ = oauth_client.create_authorization_url( + metadata["authorization_endpoint"], + redirect_uri=redirect_uri, + code_challenge_method="S256", + ) + import webbrowser + + webbrowser.open(uri) + + # Wait for redirect after user approval - reuse the same port + try: + # We need to ensure port is not None for the _get_redirect function + redirect_port = port if port is not None else find_available_port() + redirect_url = await _get_redirect(port=redirect_port) + logger.info("Received redirect, fetching token") + + # ignore type because fetch_token is awaitable but not typed as such + token = await oauth_client.fetch_token( # type: ignore[await-expr] + url=metadata["token_endpoint"], + authorization_response=redirect_url, + timeout=15, # seconds + ) + except Exception as e: + logger.warning(f"Failed to fetch token: {e}") + token = None + + # Save token for future use - passing mcp_endpoint directly + if token is not None: + oauth_cache.save(mcp_endpoint, token, "token") + logger.debug("Token saved successfully") + + # Create a standard httpx client with the OAuth bearer auth + async with create_mcp_http_client( + auth=OAuthBearerAuth(oauth_client), + headers=headers, + timeout=timeout, + **httpx_kwargs, + ) as client: + # Yield the authenticated client + yield client + finally: + await oauth_client.aclose() + + +@contextmanager +def patch_mcp_httpx_client(mcp_endpoint: str): + """ + This context manager can be used to monkeypatch the low-level function that + returns an httpx.AsyncClient. It replaces it with a function that returns an + MCP OAuth-aware client. + + This is ugly, but it lets us reuse the low-level SDK without maintaining a fork. + """ + original_shttp_client_fn = mcp.client.streamable_http.create_mcp_http_client # type: ignore + original_sse_client_fn = mcp.client.sse.create_mcp_http_client # type: ignore + + # use tokens to manage context across concurrent requests + token = _current_mcp_endpoint.set(mcp_endpoint) + + def patched_mcp_client(**kwargs): + url = _current_mcp_endpoint.get() + if url is None: + return mcp.shared._httpx_utils.create_mcp_http_client(**kwargs) + return create_mcp_oauth_client(mcp_endpoint=url, **kwargs) + + try: + mcp.client.streamable_http.create_mcp_http_client = patched_mcp_client # type: ignore + mcp.client.sse.create_mcp_http_client = patched_mcp_client # type: ignore + yield + finally: + _current_mcp_endpoint.reset(token) + mcp.client.streamable_http.create_mcp_http_client = original_shttp_client_fn # type: ignore + mcp.client.sse.create_mcp_http_client = original_sse_client_fn # type: ignore diff --git a/src/fastmcp/client/auth/oauth_cache.py b/src/fastmcp/client/auth/oauth_cache.py new file mode 100644 index 000000000..928b9fa86 --- /dev/null +++ b/src/fastmcp/client/auth/oauth_cache.py @@ -0,0 +1,122 @@ +import json +import time +from pathlib import Path +from typing import Any, ClassVar, Literal +from urllib.parse import urlparse + +from fastmcp.client.auth.httpx_client import logger +from fastmcp.settings import settings + + +class OAuthCache: + """Manages OAuth credentials and tokens caching.""" + + # Class variables + CACHE_DIR: ClassVar[Path] = settings.home / "oauth-cache" + + def __init__(self): + """Initialize the cache directory.""" + self.CACHE_DIR.mkdir(exist_ok=True, parents=True) + + @staticmethod + def get_base_url(url: str) -> str: + """Extract the base URL (scheme + host) from a URL with a path.""" + parsed = urlparse(url) + return f"{parsed.scheme}://{parsed.netloc}" + + def get_cache_key(self, url: str) -> str: + """Generate a safe filesystem key from a URL, automatically extracting the base URL.""" + base_url = self.get_base_url(url) + # Replace scheme:// and non-alphanumeric characters with _ for safety + return base_url.replace("://", "_").replace(".", "_").replace("/", "_") + + def get_file_path(self, url: str, file_type: Literal["client", "token"]) -> Path: + """Get the file path for the specified cache file type and URL.""" + key = self.get_cache_key(url) + return self.CACHE_DIR / f"{key}_{file_type}.json" + + def save( + self, url: str, data: dict[str, Any], file_type: Literal["client", "token"] + ) -> None: + """Save data to the cache file using the base URL extracted from url.""" + path = self.get_file_path(url, file_type) + path.write_text(json.dumps(data)) + base_url = self.get_base_url(url) + logger.debug(f"Saved {file_type} data for {base_url}") + + def load( + self, url: str, file_type: Literal["client", "token"] + ) -> dict[str, Any] | None: + """Load data from the cache file using the base URL extracted from url.""" + path = self.get_file_path(url, file_type) + try: + return json.loads(path.read_text()) + except (FileNotFoundError, json.JSONDecodeError): + base_url = self.get_base_url(url) + logger.debug(f"No valid {file_type} cache found for {base_url}") + return None + + def has_valid_token(self, url: str) -> bool: + """Check if there's a valid non-expired token for the given URL.""" + token = self.load(url, "token") + if not token: + return False + + # Check expiration + expires_at = token.get("expires_at") + if not expires_at or expires_at < time.time(): + return False + + return True + + def list_cached_endpoints(self) -> list[str]: + """List all base URLs with cached credentials or tokens.""" + endpoints = set() + + file_types = ["client", "token"] + for file_type in file_types: + for file in self.CACHE_DIR.glob(f"*_{file_type}.json"): + key = file.name.replace(f"_{file_type}.json", "") + # This is a simplified conversion back to URL format + # May need enhancement for complex URLs + url = key.replace("_", "://", 1) + # Attempt to reconstruct the URL in a basic way + parts = url.split("_") + if len(parts) > 1: + # Reconstruct with dots and slashes + reconstructed = parts[0] + for part in parts[1:]: + if part: + reconstructed += f".{part}" + endpoints.add(reconstructed) + else: + endpoints.add(url) + + return sorted(list(endpoints)) + + def clear(self, url: str | None = None) -> None: + """ + Clear the OAuth cache for a specific URL or all cached data. + + Args: + url: The URL to clear cache for. If None, clears all cache. + """ + if url is None: + # Clear all files in the cache directory + file_types = ["client", "token"] + for file_type in file_types: + for file in self.CACHE_DIR.glob(f"*_{file_type}.json"): + file.unlink(missing_ok=True) + logger.info("Cleared all OAuth cache data") + else: + # Clear only files for the specific URL + path = self.get_file_path(url, "client") + path.unlink(missing_ok=True) + path = self.get_file_path(url, "token") + path.unlink(missing_ok=True) + base_url = self.get_base_url(url) + logger.info(f"Cleared OAuth cache for {base_url}") + + +# Initialize global cache instance +oauth_cache = OAuthCache() diff --git a/src/fastmcp/client/client.py b/src/fastmcp/client/client.py index 9f5a111f9..97d4205c0 100644 --- a/src/fastmcp/client/client.py +++ b/src/fastmcp/client/client.py @@ -1,12 +1,22 @@ +import abc +import contextlib import datetime +from collections.abc import AsyncIterator from contextlib import AsyncExitStack from pathlib import Path -from typing import Any, cast +from typing import Any, TypedDict, cast import mcp.types from exceptiongroup import catch from mcp import ClientSession +from mcp.client.session import ( + ListRootsFnT, + LoggingFnT, + MessageHandlerFnT, + SamplingFnT, +) from pydantic import AnyUrl +from typing_extensions import Unpack from fastmcp.client.logging import LogHandler, MessageHandler from fastmcp.client.roots import ( @@ -19,10 +29,10 @@ from fastmcp.exceptions import ToolError from fastmcp.server import FastMCP from fastmcp.utilities.exceptions import get_catch_handlers -from .transports import ClientTransport, SessionKwargs, infer_transport - __all__ = [ "Client", + "ClientTransport", + "SessionKwargs", "RootsHandler", "RootsList", "LogHandler", @@ -31,6 +41,51 @@ __all__ = [ ] +class SessionKwargs(TypedDict, total=False): + """Keyword arguments for the MCP ClientSession constructor.""" + + sampling_callback: SamplingFnT | None + list_roots_callback: ListRootsFnT | None + logging_callback: LoggingFnT | None + message_handler: MessageHandlerFnT | None + read_timeout_seconds: datetime.timedelta | None + + +class ClientTransport(abc.ABC): + """ + Abstract base class for different MCP client transport mechanisms. + + A Transport is responsible for establishing and managing connections + to an MCP server, and providing a ClientSession within an async context. + """ + + @abc.abstractmethod + @contextlib.asynccontextmanager + async def connect_session( + self, **session_kwargs: Unpack[SessionKwargs] + ) -> AsyncIterator[ClientSession]: + """ + Establishes a connection and yields an active, initialized ClientSession. + + The session is guaranteed to be valid only within the scope of the + async context manager. Connection setup and teardown are handled + within this context. + + Args: + **session_kwargs: Keyword arguments to pass to the ClientSession + constructor (e.g., callbacks, timeouts). + + Yields: + An initialized mcp.ClientSession instance. + """ + raise NotImplementedError + yield None # type: ignore + + def __repr__(self) -> str: + # Basic representation for subclasses + return f"<{self.__class__.__name__}>" + + class Client: """ MCP client that delegates connection management to a Transport instance. @@ -76,6 +131,8 @@ class Client: message_handler: MessageHandler | None = None, timeout: datetime.timedelta | float | int | None = None, ): + from fastmcp.client.transports import infer_transport + self.transport = infer_transport(transport) self._session: ClientSession | None = None self._exit_stack: AsyncExitStack | None = None diff --git a/src/fastmcp/client/sse.py b/src/fastmcp/client/sse.py new file mode 100644 index 000000000..4e17c1dc9 --- /dev/null +++ b/src/fastmcp/client/sse.py @@ -0,0 +1,65 @@ +import contextlib +import datetime +import logging +from collections.abc import AsyncIterator +from typing import cast + +from mcp import ClientSession +from mcp.client.sse import sse_client +from pydantic import AnyUrl +from typing_extensions import Unpack + +from fastmcp.client.auth.httpx_client import patch_mcp_httpx_client +from fastmcp.client.client import ClientTransport, SessionKwargs + +logger = logging.getLogger(__name__) + + +class SSETransport(ClientTransport): + """Transport implementation that connects to an MCP server via Server-Sent Events.""" + + def __init__( + self, + url: str | AnyUrl, + headers: dict[str, str] | None = None, + sse_read_timeout: datetime.timedelta | float | int | None = None, + ): + if isinstance(url, AnyUrl): + url = str(url) + if not isinstance(url, str) or not url.startswith("http"): + raise ValueError("Invalid HTTP/S URL provided for SSE.") + self.url = url + self.headers = headers or {} + + if isinstance(sse_read_timeout, int | float): + sse_read_timeout = datetime.timedelta(seconds=sse_read_timeout) + self.sse_read_timeout = sse_read_timeout + + @contextlib.asynccontextmanager + async def connect_session( + self, **session_kwargs: Unpack[SessionKwargs] + ) -> AsyncIterator[ClientSession]: + client_kwargs = {} + # sse_read_timeout has a default value set, so we can't pass None without overriding it + # instead we simply leave the kwarg out if it's not provided + if self.sse_read_timeout is not None: + client_kwargs["sse_read_timeout"] = self.sse_read_timeout.total_seconds() + if session_kwargs.get("read_timeout_seconds", None) is not None: + read_timeout_seconds = cast( + datetime.timedelta, session_kwargs.get("read_timeout_seconds") + ) + client_kwargs["timeout"] = read_timeout_seconds.total_seconds() + + with patch_mcp_httpx_client(self.url): + async with sse_client( + self.url, headers=self.headers, **client_kwargs + ) as transport: + read_stream, write_stream = transport + async with ClientSession( + read_stream, write_stream, **session_kwargs + ) as session: + await session.initialize() + yield session + + def __repr__(self) -> str: + return f"" diff --git a/src/fastmcp/client/streamable_http.py b/src/fastmcp/client/streamable_http.py new file mode 100644 index 000000000..a5d7b8e03 --- /dev/null +++ b/src/fastmcp/client/streamable_http.py @@ -0,0 +1,61 @@ +import contextlib +import datetime +from collections.abc import AsyncIterator + +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client +from pydantic import AnyUrl +from typing_extensions import Unpack + +from fastmcp.client.auth.httpx_client import patch_mcp_httpx_client +from fastmcp.client.client import ClientTransport, SessionKwargs +from fastmcp.utilities.logging import get_logger + +logger = get_logger(__name__) + + +class StreamableHttpTransport(ClientTransport): + """Transport implementation that connects to an MCP server via Streamable HTTP Requests.""" + + def __init__( + self, + url: str | AnyUrl, + headers: dict[str, str] | None = None, + sse_read_timeout: datetime.timedelta | float | int | None = None, + ): + if isinstance(url, AnyUrl): + url = str(url) + if not isinstance(url, str) or not url.startswith("http"): + raise ValueError("Invalid HTTP/S URL provided for Streamable HTTP.") + self.url = url + self.headers = headers or {} + + if isinstance(sse_read_timeout, int | float): + sse_read_timeout = datetime.timedelta(seconds=sse_read_timeout) + self.sse_read_timeout = sse_read_timeout + + @contextlib.asynccontextmanager + async def connect_session( + self, **session_kwargs: Unpack[SessionKwargs] + ) -> AsyncIterator[ClientSession]: + client_kwargs = {} + # sse_read_timeout has a default value set, so we can't pass None without overriding it + # instead we simply leave the kwarg out if it's not provided + if self.sse_read_timeout is not None: + client_kwargs["sse_read_timeout"] = self.sse_read_timeout + if session_kwargs.get("read_timeout_seconds", None) is not None: + client_kwargs["timeout"] = session_kwargs.get("read_timeout_seconds") + + with patch_mcp_httpx_client(self.url): + async with streamablehttp_client( + self.url, headers=self.headers, **client_kwargs + ) as transport: + read_stream, write_stream, _ = transport + async with ClientSession( + read_stream, write_stream, **session_kwargs + ) as session: + await session.initialize() + yield session + + def __repr__(self) -> str: + return f"" diff --git a/src/fastmcp/client/transports.py b/src/fastmcp/client/transports.py index 7faeab613..fc33af335 100644 --- a/src/fastmcp/client/transports.py +++ b/src/fastmcp/client/transports.py @@ -1,76 +1,38 @@ -import abc import contextlib -import datetime -import inspect import os import shutil import sys -import warnings from collections.abc import AsyncIterator from pathlib import Path -from typing import Any, TypedDict, cast +from typing import Any from mcp import ClientSession, StdioServerParameters -from mcp.client.session import ( - ListRootsFnT, - LoggingFnT, - MessageHandlerFnT, - SamplingFnT, -) -from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client -from mcp.client.streamable_http import streamablehttp_client from mcp.client.websocket import websocket_client from mcp.shared.memory import create_connected_server_and_client_session from pydantic import AnyUrl from typing_extensions import Unpack +from fastmcp.client.client import ClientTransport, SessionKwargs +from fastmcp.client.sse import SSETransport +from fastmcp.client.streamable_http import StreamableHttpTransport from fastmcp.server import FastMCP as FastMCPServer - -class SessionKwargs(TypedDict, total=False): - """Keyword arguments for the MCP ClientSession constructor.""" - - sampling_callback: SamplingFnT | None - list_roots_callback: ListRootsFnT | None - logging_callback: LoggingFnT | None - message_handler: MessageHandlerFnT | None - read_timeout_seconds: datetime.timedelta | None - - -class ClientTransport(abc.ABC): - """ - Abstract base class for different MCP client transport mechanisms. - - A Transport is responsible for establishing and managing connections - to an MCP server, and providing a ClientSession within an async context. - """ - - @abc.abstractmethod - @contextlib.asynccontextmanager - async def connect_session( - self, **session_kwargs: Unpack[SessionKwargs] - ) -> AsyncIterator[ClientSession]: - """ - Establishes a connection and yields an active, initialized ClientSession. - - The session is guaranteed to be valid only within the scope of the - async context manager. Connection setup and teardown are handled - within this context. - - Args: - **session_kwargs: Keyword arguments to pass to the ClientSession - constructor (e.g., callbacks, timeouts). - - Yields: - An initialized mcp.ClientSession instance. - """ - raise NotImplementedError - yield None # type: ignore - - def __repr__(self) -> str: - # Basic representation for subclasses - return f"<{self.__class__.__name__}>" +__all__ = [ + "ClientTransport", + "SSETransport", + "StreamableHttpTransport", + "FastMCPServer", + "WSTransport", + "StdioTransport", + "PythonStdioTransport", + "FastMCPStdioTransport", + "NodeStdioTransport", + "UvxStdioTransport", + "NpxStdioTransport", + "FastMCPTransport", + "infer_transport", +] class WSTransport(ClientTransport): @@ -99,101 +61,6 @@ class WSTransport(ClientTransport): return f"" -class SSETransport(ClientTransport): - """Transport implementation that connects to an MCP server via Server-Sent Events.""" - - def __init__( - self, - url: str | AnyUrl, - headers: dict[str, str] | None = None, - sse_read_timeout: datetime.timedelta | float | int | None = None, - ): - if isinstance(url, AnyUrl): - url = str(url) - if not isinstance(url, str) or not url.startswith("http"): - raise ValueError("Invalid HTTP/S URL provided for SSE.") - self.url = url - self.headers = headers or {} - - if isinstance(sse_read_timeout, int | float): - sse_read_timeout = datetime.timedelta(seconds=sse_read_timeout) - self.sse_read_timeout = sse_read_timeout - - @contextlib.asynccontextmanager - async def connect_session( - self, **session_kwargs: Unpack[SessionKwargs] - ) -> AsyncIterator[ClientSession]: - client_kwargs = {} - # sse_read_timeout has a default value set, so we can't pass None without overriding it - # instead we simply leave the kwarg out if it's not provided - if self.sse_read_timeout is not None: - client_kwargs["sse_read_timeout"] = self.sse_read_timeout.total_seconds() - if session_kwargs.get("read_timeout_seconds", None) is not None: - read_timeout_seconds = cast( - datetime.timedelta, session_kwargs.get("read_timeout_seconds") - ) - client_kwargs["timeout"] = read_timeout_seconds.total_seconds() - - async with sse_client( - self.url, headers=self.headers, **client_kwargs - ) as transport: - read_stream, write_stream = transport - async with ClientSession( - read_stream, write_stream, **session_kwargs - ) as session: - await session.initialize() - yield session - - def __repr__(self) -> str: - return f"" - - -class StreamableHttpTransport(ClientTransport): - """Transport implementation that connects to an MCP server via Streamable HTTP Requests.""" - - def __init__( - self, - url: str | AnyUrl, - headers: dict[str, str] | None = None, - sse_read_timeout: datetime.timedelta | float | int | None = None, - ): - if isinstance(url, AnyUrl): - url = str(url) - if not isinstance(url, str) or not url.startswith("http"): - raise ValueError("Invalid HTTP/S URL provided for Streamable HTTP.") - self.url = url - self.headers = headers or {} - - if isinstance(sse_read_timeout, int | float): - sse_read_timeout = datetime.timedelta(seconds=sse_read_timeout) - self.sse_read_timeout = sse_read_timeout - - @contextlib.asynccontextmanager - async def connect_session( - self, **session_kwargs: Unpack[SessionKwargs] - ) -> AsyncIterator[ClientSession]: - client_kwargs = {} - # sse_read_timeout has a default value set, so we can't pass None without overriding it - # instead we simply leave the kwarg out if it's not provided - if self.sse_read_timeout is not None: - client_kwargs["sse_read_timeout"] = self.sse_read_timeout - if session_kwargs.get("read_timeout_seconds", None) is not None: - client_kwargs["timeout"] = session_kwargs.get("read_timeout_seconds") - - async with streamablehttp_client( - self.url, headers=self.headers, **client_kwargs - ) as transport: - read_stream, write_stream, _ = transport - async with ClientSession( - read_stream, write_stream, **session_kwargs - ) as session: - await session.initialize() - yield session - - def __repr__(self) -> str: - return f"" - - class StdioTransport(ClientTransport): """ Base transport for connecting to an MCP server via subprocess with stdio. @@ -500,18 +367,9 @@ def infer_transport( # the transport is an http(s) URL elif isinstance(transport, AnyUrl | str) and str(transport).startswith("http"): if str(transport).rstrip("/").endswith("/sse"): - warnings.warn( - inspect.cleandoc( - """ - As of FastMCP 2.3.0, HTTP URLs are inferred to use Streamable HTTP. - The provided URL ends in `/sse`, so you may encounter unexpected behavior. - If you intended to use SSE, please use the `SSETransport` class directly. - """ - ), - category=UserWarning, - stacklevel=2, - ) - return StreamableHttpTransport(url=transport) + return SSETransport(url=transport) + else: + return StreamableHttpTransport(url=transport) # the transport is a websocket URL elif isinstance(transport, AnyUrl | str) and str(transport).startswith("ws"): diff --git a/src/fastmcp/settings.py b/src/fastmcp/settings.py index 34209f5b0..104550025 100644 --- a/src/fastmcp/settings.py +++ b/src/fastmcp/settings.py @@ -1,6 +1,7 @@ from __future__ import annotations as _annotations import inspect +from pathlib import Path from typing import TYPE_CHECKING, Annotated, Literal from mcp.server.auth.settings import AuthSettings @@ -27,6 +28,8 @@ class Settings(BaseSettings): nested_model_default_partial_update=True, ) + home: Path = Path.home() / ".fastmcp" + test_mode: bool = False log_level: LOG_LEVEL = "INFO" client_raise_first_exceptiongroup_error: Annotated[