fastmcp 2.5.2__py3-none-any.whl → 2.6.1__py3-none-any.whl

fastmcp/client/__init__.py

@@ -11,6 +11,7 @@ from .transports import (
     FastMCPTransport,
     StreamableHttpTransport,
 )
+from .auth import OAuth, BearerAuth
 
 __all__ = [
     "Client",
@@ -24,4 +25,6 @@ __all__ = [
     "NpxStdioTransport",
     "FastMCPTransport",
     "StreamableHttpTransport",
+    "OAuth",
+    "BearerAuth",
 ]

fastmcp/client/auth/__init__.py

@@ -0,0 +1,4 @@
+from .bearer import BearerAuth
+from .oauth import OAuth
+
+__all__ = ["BearerAuth", "OAuth"]

fastmcp/client/auth/bearer.py

@@ -0,0 +1,17 @@
+import httpx
+from pydantic import SecretStr
+
+from fastmcp.utilities.logging import get_logger
+
+__all__ = ["BearerAuth"]
+
+logger = get_logger(__name__)
+
+
+class BearerAuth(httpx.Auth):
+    def __init__(self, token: str):
+        self.token = SecretStr(token)
+
+    def auth_flow(self, request):
+        request.headers["Authorization"] = f"Bearer {self.token.get_secret_value()}"
+        yield request

fastmcp/client/auth/oauth.py

@@ -0,0 +1,391 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import webbrowser
+from pathlib import Path
+from typing import Any, Literal
+from urllib.parse import urljoin, urlparse
+
+import anyio
+import httpx
+from mcp.client.auth import OAuthClientProvider as _MCPOAuthClientProvider
+from mcp.client.auth import TokenStorage
+from mcp.shared.auth import (
+    OAuthClientInformationFull,
+    OAuthClientMetadata,
+)
+from mcp.shared.auth import (
+    OAuthMetadata as _MCPServerOAuthMetadata,
+)
+from mcp.shared.auth import (
+    OAuthToken as OAuthToken,
+)
+from pydantic import AnyHttpUrl, ValidationError
+
+from fastmcp.client.oauth_callback import (
+    create_oauth_callback_server,
+)
+from fastmcp.settings import settings as fastmcp_global_settings
+from fastmcp.utilities.http import find_available_port
+from fastmcp.utilities.logging import get_logger
+
+__all__ = ["OAuth"]
+
+logger = get_logger(__name__)
+
+
+def default_cache_dir() -> Path:
+    return fastmcp_global_settings.home / "oauth-mcp-client-cache"
+
+
+# Flexible OAuth models for real-world compatibility
+class ServerOAuthMetadata(_MCPServerOAuthMetadata):
+    """
+    More flexible OAuth metadata model that accepts broader ranges of values
+    than the restrictive MCP standard model.
+
+    This handles real-world OAuth servers like PayPal that may support
+    additional methods not in the MCP specification.
+    """
+
+    # Allow any code challenge methods, not just S256
+    code_challenge_methods_supported: list[str] | None = None
+
+    # Allow any token endpoint auth methods
+    token_endpoint_auth_methods_supported: list[str] | None = None
+
+    # Allow any grant types
+    grant_types_supported: list[str] | None = None
+
+    # Allow any response types
+    response_types_supported: list[str] = ["code"]
+
+    # Allow any response modes
+    response_modes_supported: list[str] | None = None
+
+
+class OAuthClientProvider(_MCPOAuthClientProvider):
+    """
+    OAuth client provider with more flexible OAuth metadata discovery.
+    """
+
+    async def _discover_oauth_metadata(
+        self, server_url: str
+    ) -> ServerOAuthMetadata | None:
+        """
+        Discover OAuth metadata with flexible validation.
+
+        This is nearly identical to the parent implementation but uses
+        ServerOAuthMetadata instead of the restrictive MCP OAuthMetadata.
+        """
+        # Extract base URL per MCP spec
+        auth_base_url = self._get_authorization_base_url(server_url)
+        url = urljoin(auth_base_url, "/.well-known/oauth-authorization-server")
+
+        from mcp.types import LATEST_PROTOCOL_VERSION
+
+        headers = {"MCP-Protocol-Version": LATEST_PROTOCOL_VERSION}
+
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.get(url, headers=headers)
+                if response.status_code == 404:
+                    return None
+                response.raise_for_status()
+                metadata_json = response.json()
+                logger.debug(f"OAuth metadata discovered: {metadata_json}")
+                return ServerOAuthMetadata.model_validate(metadata_json)
+            except Exception:
+                # Retry without MCP header for CORS compatibility
+                try:
+                    response = await client.get(url)
+                    if response.status_code == 404:
+                        return None
+                    response.raise_for_status()
+                    metadata_json = response.json()
+                    logger.debug(
+                        f"OAuth metadata discovered (no MCP header): {metadata_json}"
+                    )
+                    return ServerOAuthMetadata.model_validate(metadata_json)
+                except Exception:
+                    logger.exception("Failed to discover OAuth metadata")
+                    return None
+
+
+class FileTokenStorage(TokenStorage):
+    """
+    File-based token storage implementation for OAuth credentials and tokens.
+    Implements the mcp.client.auth.TokenStorage protocol.
+
+    Each instance is tied to a specific server URL for proper token isolation.
+    """
+
+    def __init__(self, server_url: str, cache_dir: Path | None = None):
+        """Initialize storage for a specific server URL."""
+        self.server_url = server_url
+        self.cache_dir = cache_dir or default_cache_dir()
+        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."""
+        parsed = urlparse(url)
+        return f"{parsed.scheme}://{parsed.netloc}"
+
+    def get_cache_key(self) -> str:
+        """Generate a safe filesystem key from the server's base URL."""
+        base_url = self.get_base_url(self.server_url)
+        return (
+            base_url.replace("://", "_")
+            .replace(".", "_")
+            .replace("/", "_")
+            .replace(":", "_")
+        )
+
+    def _get_file_path(self, file_type: Literal["client_info", "tokens"]) -> Path:
+        """Get the file path for the specified cache file type."""
+        key = self.get_cache_key()
+        return self.cache_dir / f"{key}_{file_type}.json"
+
+    async def get_tokens(self) -> OAuthToken | None:
+        """Load tokens from file storage."""
+        path = self._get_file_path("tokens")
+
+        try:
+            tokens = OAuthToken.model_validate_json(path.read_text())
+            # now = datetime.datetime.now(datetime.timezone.utc)
+            # if tokens.expires_at is not None and tokens.expires_at <= now:
+            #     logger.debug(f"Token expired for {self.get_base_url(self.server_url)}")
+            #     return None
+            return tokens
+        except (FileNotFoundError, json.JSONDecodeError, ValidationError) as e:
+            logger.debug(
+                f"Could not load tokens for {self.get_base_url(self.server_url)}: {e}"
+            )
+            return None
+
+    async def set_tokens(self, tokens: OAuthToken) -> None:
+        """Save tokens to file storage."""
+        path = self._get_file_path("tokens")
+        path.write_text(tokens.model_dump_json(indent=2))
+        logger.debug(f"Saved tokens for {self.get_base_url(self.server_url)}")
+
+    async def get_client_info(self) -> OAuthClientInformationFull | None:
+        """Load client information from file storage."""
+        path = self._get_file_path("client_info")
+        try:
+            client_info = OAuthClientInformationFull.model_validate_json(
+                path.read_text()
+            )
+            # Check if we have corresponding valid tokens
+            # If no tokens exist, the OAuth flow was incomplete and we should
+            # force a fresh client registration
+            tokens = await self.get_tokens()
+            if tokens is None:
+                logger.debug(
+                    f"No tokens found for client info at {self.get_base_url(self.server_url)}. "
+                    "OAuth flow may have been incomplete. Clearing client info to force fresh registration."
+                )
+                # Clear the incomplete client info
+                client_info_path = self._get_file_path("client_info")
+                client_info_path.unlink(missing_ok=True)
+                return None
+
+            return client_info
+        except (FileNotFoundError, json.JSONDecodeError, ValidationError) as e:
+            logger.debug(
+                f"Could not load client info for {self.get_base_url(self.server_url)}: {e}"
+            )
+            return None
+
+    async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
+        """Save client information to file storage."""
+        path = self._get_file_path("client_info")
+        path.write_text(client_info.model_dump_json(indent=2))
+        logger.debug(f"Saved client info for {self.get_base_url(self.server_url)}")
+
+    def clear(self) -> None:
+        """Clear all cached data for this server."""
+        file_types: list[Literal["client_info", "tokens"]] = ["client_info", "tokens"]
+        for file_type in file_types:
+            path = self._get_file_path(file_type)
+            path.unlink(missing_ok=True)
+        logger.info(f"Cleared OAuth cache for {self.get_base_url(self.server_url)}")
+
+    @classmethod
+    def clear_all(cls, cache_dir: Path | None = None) -> None:
+        """Clear all cached data for all servers."""
+        cache_dir = cache_dir or default_cache_dir()
+        if not cache_dir.exists():
+            return
+
+        file_types: list[Literal["client_info", "tokens"]] = ["client_info", "tokens"]
+        for file_type in file_types:
+            for file in cache_dir.glob(f"*_{file_type}.json"):
+                file.unlink(missing_ok=True)
+        logger.info("Cleared all OAuth client cache data.")
+
+
+async def discover_oauth_metadata(
+    server_base_url: str, httpx_kwargs: dict[str, Any] | None = None
+) -> _MCPServerOAuthMetadata | None:
+    """
+    Discover OAuth metadata from the server using RFC 8414 well-known endpoint.
+
+    Args:
+        server_base_url: Base URL of the OAuth server (e.g., "https://example.com")
+        httpx_kwargs: Additional kwargs for httpx client
+
+    Returns:
+        OAuth metadata if found, None otherwise
+    """
+    well_known_url = urljoin(server_base_url, "/.well-known/oauth-authorization-server")
+    logger.debug(f"Discovering OAuth metadata from: {well_known_url}")
+
+    async with httpx.AsyncClient(**(httpx_kwargs or {})) as client:
+        try:
+            response = await client.get(well_known_url, timeout=10.0)
+            if response.status_code == 200:
+                logger.debug("Successfully discovered OAuth metadata")
+                return _MCPServerOAuthMetadata.model_validate(response.json())
+            elif response.status_code == 404:
+                logger.debug(
+                    "OAuth metadata not found (404) - server may not require auth"
+                )
+                return None
+            else:
+                logger.warning(f"OAuth metadata request failed: {response.status_code}")
+                return None
+        except (httpx.RequestError, json.JSONDecodeError, ValidationError) as e:
+            logger.debug(f"OAuth metadata discovery failed: {e}")
+            return None
+
+
+async def check_if_auth_required(
+    mcp_url: str, httpx_kwargs: dict[str, Any] | None = None
+) -> bool:
+    """
+    Check if the MCP endpoint requires authentication by making a test request.
+
+    Returns:
+        True if auth appears to be required, False otherwise
+    """
+    async with httpx.AsyncClient(**(httpx_kwargs or {})) as client:
+        try:
+            # Try a simple request to the endpoint
+            response = await client.get(mcp_url, timeout=5.0)
+
+            # If we get 401/403, auth is likely required
+            if response.status_code in (401, 403):
+                return True
+
+            # Check for WWW-Authenticate header
+            if "WWW-Authenticate" in response.headers:
+                return True
+
+            # If we get a successful response, auth may not be required
+            return False
+
+        except httpx.RequestError:
+            # If we can't connect, assume auth might be required
+            return True
+
+
+def OAuth(
+    mcp_url: str,
+    scopes: str | list[str] | None = None,
+    client_name: str = "FastMCP Client",
+    token_storage_cache_dir: Path | None = None,
+    additional_client_metadata: dict[str, Any] | None = None,
+) -> _MCPOAuthClientProvider:
+    """
+    Create an OAuthClientProvider for an MCP server.
+
+    This is intended to be provided to the `auth` parameter of an
+    httpx.AsyncClient (or appropriate FastMCP client/transport instance)
+
+    Args:
+        mcp_url: Full URL to the MCP endpoint (e.g.,
+        "http://host/mcp/sse")
+        scopes: OAuth scopes to request. Can be a
+        space-separated string or a list of strings.
+        client_name: Name for this client during registration
+        token_storage_cache_dir: Directory for FileTokenStorage
+        additional_client_metadata: Extra fields for OAuthClientMetadata
+
+    Returns:
+        OAuthClientProvider
+    """
+    parsed_url = urlparse(mcp_url)
+    server_base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
+
+    # Setup OAuth client
+    redirect_port = find_available_port()
+    redirect_uri = f"http://127.0.0.1:{redirect_port}/callback"
+
+    if isinstance(scopes, list):
+        scopes = " ".join(scopes)
+
+    client_metadata = OAuthClientMetadata(
+        client_name=client_name,
+        redirect_uris=[AnyHttpUrl(redirect_uri)],
+        grant_types=["authorization_code", "refresh_token"],
+        response_types=["code"],
+        token_endpoint_auth_method="client_secret_post",
+        scope=scopes,
+        **(additional_client_metadata or {}),
+    )
+
+    # Create server-specific token storage
+    storage = FileTokenStorage(
+        server_url=server_base_url, cache_dir=token_storage_cache_dir
+    )
+
+    # Define OAuth handlers
+    async def redirect_handler(authorization_url: str) -> None:
+        """Open browser for authorization."""
+        logger.info(f"OAuth authorization URL: {authorization_url}")
+        webbrowser.open(authorization_url)
+
+    async def callback_handler() -> tuple[str, str | None]:
+        """Handle OAuth callback and return (auth_code, state)."""
+        # Create a future to capture the OAuth response
+        response_future = asyncio.get_running_loop().create_future()
+
+        # Create server with the future
+        server = create_oauth_callback_server(
+            port=redirect_port,
+            server_url=server_base_url,
+            response_future=response_future,
+        )
+
+        # Run server until response is received with timeout logic
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(server.serve)
+            logger.info(
+                f"🎧 OAuth callback server started on http://127.0.0.1:{redirect_port}"
+            )
+
+            TIMEOUT = 300.0  # 5 minute timeout
+            try:
+                with anyio.fail_after(TIMEOUT):
+                    auth_code, state = await response_future
+                    return auth_code, state
+            except TimeoutError:
+                raise TimeoutError(f"OAuth callback timed out after {TIMEOUT} seconds")
+            finally:
+                server.should_exit = True
+                await asyncio.sleep(0.1)  # Allow server to shutdown gracefully
+                tg.cancel_scope.cancel()
+
+    # Create OAuth provider
+    oauth_provider = OAuthClientProvider(
+        server_url=server_base_url,
+        client_metadata=client_metadata,
+        storage=storage,
+        redirect_handler=redirect_handler,
+        callback_handler=callback_handler,
+    )
+
+    return oauth_provider

fastmcp/client/client.py

@@ -1,9 +1,11 @@
+import asyncio
 import datetime
 from contextlib import AsyncExitStack, asynccontextmanager
 from pathlib import Path
-from typing import Any, Generic, cast, overload
+from typing import Any, Generic, Literal, cast, overload
 
 import anyio
+import httpx
 import mcp.types
 from exceptiongroup import catch
 from mcp import ClientSession
@@ -43,6 +45,7 @@ from .transports import (
 
 __all__ = [
     "Client",
+    "SessionKwargs",
     "RootsHandler",
     "RootsList",
     "LogHandler",
@@ -142,11 +145,11 @@ class Client(Generic[ClientTransportT]):
         progress_handler: ProgressHandler | None = None,
         timeout: datetime.timedelta | float | int | None = None,
         init_timeout: datetime.timedelta | float | int | None = None,
+        auth: httpx.Auth | Literal["oauth"] | str | None = None,
     ):
         self.transport = cast(ClientTransportT, infer_transport(transport))
-        self._session: ClientSession | None = None
-        self._exit_stack: AsyncExitStack | None = None
-        self._nesting_counter: int = 0
+        if auth is not None:
+            self.transport._set_auth(auth)
         self._initialize_result: mcp.types.InitializeResult | None = None
 
         if log_handler is None:
@@ -187,6 +190,15 @@ class Client(Generic[ClientTransportT]):
                 sampling_handler
             )
 
+        # session context management
+        self._session: ClientSession | None = None
+        self._exit_stack: AsyncExitStack | None = None
+        self._nesting_counter: int = 0
+        self._context_lock = anyio.Lock()
+        self._session_task: asyncio.Task | None = None
+        self._ready_event = anyio.Event()
+        self._stop_event = anyio.Event()
+
     @property
     def session(self) -> ClientSession:
         """Get the current active session. Raises RuntimeError if not connected."""
@@ -237,39 +249,75 @@ class Client(Generic[ClientTransportT]):
                 except TimeoutError:
                     raise RuntimeError("Failed to initialize server session")
                 finally:
-                    self._exit_stack = None
                     self._session = None
                     self._initialize_result = None
 
     async def __aenter__(self):
-        if self._nesting_counter == 0:
-            # Create exit stack to manage both context managers
-            stack = AsyncExitStack()
-            await stack.__aenter__()
-
-            await stack.enter_async_context(self._context_manager())
-
-            self._exit_stack = stack
-
-        self._nesting_counter += 1
-
+        await self._connect()
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        self._nesting_counter -= 1
+        await self._disconnect()
+
+    async def _connect(self):
+        # ensure only one session is running at a time to avoid race conditions
+        async with self._context_lock:
+            need_to_start = self._session_task is None or self._session_task.done()
+            if need_to_start:
+                self._stop_event = anyio.Event()
+                self._ready_event = anyio.Event()
+                self._session_task = asyncio.create_task(self._session_runner())
+            await self._ready_event.wait()
+            self._nesting_counter += 1
+        return self
 
-        if self._nesting_counter == 0:
-            # Exit the stack which will handle cleaning up the session
-            if self._exit_stack is not None:
-                try:
-                    await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
-                finally:
-                    self._exit_stack = None
+    async def _disconnect(self, force: bool = False):
+        # ensure only one session is running at a time to avoid race conditions
+        async with self._context_lock:
+            # if we are forcing a disconnect, reset the nesting counter
+            if force:
+                self._nesting_counter = 0
+
+            # otherwise decrement to check if we are done nesting
+            else:
+                self._nesting_counter = max(0, self._nesting_counter - 1)
+
+            # if we are still nested, return
+            if self._nesting_counter > 0:
+                return
+
+            # stop the active seesion
+            if self._session_task is None:
+                return
+            self._stop_event.set()
+            runner_task = self._session_task
+            self._session_task = None
+
+        # wait for the session to finish
+        if runner_task:
+            await runner_task
+
+        # Reset for future reconnects
+        self._stop_event = anyio.Event()
+        self._ready_event = anyio.Event()
+        self._session = None
+        self._initialize_result = None
+
+    async def _session_runner(self):
+        async with AsyncExitStack() as stack:
+            try:
+                await stack.enter_async_context(self._context_manager())
+                # Session/context is now ready
+                self._ready_event.set()
+                # Wait until disconnect/stop is requested
+                await self._stop_event.wait()
+            finally:
+                # On exit, ensure ready event is set (idempotent)
+                self._ready_event.set()
 
     async def close(self):
+        await self._disconnect(force=True)
         await self.transport.close()
-        self._session = None
-        self._initialize_result = None
 
     # --- MCP Client Methods ---