from __future__ import annotations import time import webbrowser from collections.abc import AsyncGenerator from contextlib import aclosing from typing import Any import anyio import httpx from key_value.aio.adapters.pydantic import PydanticAdapter from key_value.aio.protocols import AsyncKeyValue from key_value.aio.stores.memory import MemoryStore from mcp.client.auth import OAuthClientProvider, TokenStorage from mcp.shared._httpx_utils import McpHttpClientFactory from mcp.shared.auth import ( OAuthClientInformationFull, OAuthClientMetadata, OAuthToken, ) from pydantic import AnyHttpUrl from typing_extensions import override from uvicorn.server import Server from fastmcp.client.oauth_callback import ( OAuthCallbackResult, create_oauth_callback_server, ) from fastmcp.utilities.http import find_available_port from fastmcp.utilities.logging import get_logger __all__ = ["OAuth"] logger = get_logger(__name__) def _normalize_callback_host_for_bind(host: str) -> str: if host.startswith("[") and host.endswith("]"): return host[1:-1] return host def _format_callback_host_for_url(host: str) -> str: if ":" in host: return f"[{host}]" return host class ClientNotFoundError(Exception): """Raised when OAuth client credentials are not found on the server.""" 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: # noqa: SIM103 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 class TokenStorageAdapter(TokenStorage): _server_url: str _key_value_store: AsyncKeyValue _storage_oauth_token: PydanticAdapter[OAuthToken] _storage_client_info: PydanticAdapter[OAuthClientInformationFull] def __init__(self, async_key_value: AsyncKeyValue, server_url: str): self._server_url = server_url self._key_value_store = async_key_value self._storage_oauth_token = PydanticAdapter[OAuthToken]( default_collection="mcp-oauth-token", key_value=async_key_value, pydantic_model=OAuthToken, raise_on_validation_error=True, ) self._storage_client_info = PydanticAdapter[OAuthClientInformationFull]( default_collection="mcp-oauth-client-info", key_value=async_key_value, pydantic_model=OAuthClientInformationFull, raise_on_validation_error=True, ) def _get_token_cache_key(self) -> str: return f"{self._server_url}/tokens" def _get_client_info_cache_key(self) -> str: return f"{self._server_url}/client_info" def _get_token_expiry_cache_key(self) -> str: return f"{self._server_url}/token_expiry" async def clear(self) -> None: await self._storage_oauth_token.delete(key=self._get_token_cache_key()) await self._storage_client_info.delete(key=self._get_client_info_cache_key()) await self._key_value_store.delete( key=self._get_token_expiry_cache_key(), collection="mcp-oauth-token-expiry", ) @override async def get_tokens(self) -> OAuthToken | None: return await self._storage_oauth_token.get(key=self._get_token_cache_key()) @override async def set_tokens(self, tokens: OAuthToken) -> None: # Don't set TTL based on access token expiry - the refresh token may be # valid much longer. Use 1 year as a reasonable upper bound; the OAuth # provider handles actual token expiry/refresh logic. await self._storage_oauth_token.put( key=self._get_token_cache_key(), value=tokens, ttl=60 * 60 * 24 * 365, # 1 year ) # Store absolute expiry so reloads don't misinterpret the stale # relative expires_in value (#2862). if tokens.expires_in is not None: expires_at = time.time() + int(tokens.expires_in) await self._key_value_store.put( key=self._get_token_expiry_cache_key(), value={"expires_at": expires_at}, collection="mcp-oauth-token-expiry", ttl=60 * 60 * 24 * 365, ) async def get_token_expiry(self) -> float | None: raw = await self._key_value_store.get( key=self._get_token_expiry_cache_key(), collection="mcp-oauth-token-expiry", ) if raw is not None: return float(raw["expires_at"]) return None @override async def get_client_info(self) -> OAuthClientInformationFull | None: return await self._storage_client_info.get( key=self._get_client_info_cache_key() ) @override async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: ttl: int | None = None if client_info.client_secret_expires_at: ttl = client_info.client_secret_expires_at - int(time.time()) await self._storage_client_info.put( key=self._get_client_info_cache_key(), value=client_info, ttl=ttl, ) class OAuth(OAuthClientProvider): """ OAuth client provider for MCP servers with browser-based authentication. This class provides OAuth authentication for FastMCP clients by opening a browser for user authorization and running a local callback server. """ _bound: bool def __init__( self, mcp_url: str | None = None, scopes: str | list[str] | None = None, client_name: str = "FastMCP Client", token_storage: AsyncKeyValue | None = None, additional_client_metadata: dict[str, Any] | None = None, callback_port: int | None = None, callback_host: str = "localhost", callback_timeout: float = 300.0, httpx_client_factory: McpHttpClientFactory | None = None, # Alternative to dynamic client registration: # --- Clients host a static JSON document at an HTTPS URL --- client_metadata_url: str | None = None, # --- OR clients provide full client information --- client_id: str | None = None, client_secret: str | None = None, ): """ Initialize OAuth client provider for an MCP server. Args: mcp_url: Full URL to the MCP endpoint (e.g. "http://host/mcp/sse/"). Optional when OAuth is passed to Client(auth=...), which provides the URL automatically from the transport. 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: An AsyncKeyValue-compatible token store, tokens are stored in memory if not provided additional_client_metadata: Extra fields for OAuthClientMetadata callback_port: Fixed port for OAuth callback (default: random available port) callback_host: Hostname used for OAuth redirect URI and callback server. callback_timeout: Seconds to wait for OAuth callback before timing out. client_metadata_url: A CIMD (Client ID Metadata Document) URL. When provided, this URL is used as the client_id instead of performing Dynamic Client Registration. Must be an HTTPS URL with a non-root path (e.g. "https://myapp.example.com/oauth/client.json"). client_id: Pre-registered OAuth client ID. When provided, skips dynamic client registration and uses these static credentials instead. client_secret: OAuth client secret (optional, used with client_id) """ # Store config for deferred binding if mcp_url not yet known self._scopes = scopes self._client_name = client_name self._token_storage = token_storage self._additional_client_metadata = additional_client_metadata self._callback_port = callback_port self._callback_host = _normalize_callback_host_for_bind(callback_host) self._callback_timeout = callback_timeout self._client_metadata_url = client_metadata_url self._client_id = client_id self._client_secret = client_secret self._static_client_info = None self.httpx_client_factory = httpx_client_factory or httpx.AsyncClient self._bound = False if mcp_url is not None: self._bind(mcp_url) def _bind(self, mcp_url: str) -> None: """Bind this OAuth provider to a specific MCP server URL. Called automatically when mcp_url is provided to __init__, or by the transport when OAuth is used without an explicit URL. """ if self._bound: return mcp_url = mcp_url.rstrip("/") self.redirect_port = self._callback_port or find_available_port( host=self._callback_host ) redirect_host = _format_callback_host_for_url(self._callback_host) redirect_uri = f"http://{redirect_host}:{self.redirect_port}/callback" scopes_str: str if isinstance(self._scopes, list): scopes_str = " ".join(self._scopes) elif self._scopes is not None: scopes_str = str(self._scopes) else: scopes_str = "" client_metadata = OAuthClientMetadata( client_name=self._client_name, redirect_uris=[AnyHttpUrl(redirect_uri)], grant_types=["authorization_code", "refresh_token"], response_types=["code"], scope=scopes_str, **(self._additional_client_metadata or {}), ) if self._client_id: # Create the full static client info directly which will avoid DCR. # Spread client_metadata so redirect_uris, grant_types, response_types, # scope, etc. are included — servers may validate these fields. metadata = client_metadata.model_dump(exclude_none=True) # Default token_endpoint_auth_method based on whether a secret is # provided, unless the caller already set it via additional_client_metadata. if "token_endpoint_auth_method" not in metadata: metadata["token_endpoint_auth_method"] = ( "client_secret_post" if self._client_secret else "none" ) self._static_client_info = OAuthClientInformationFull( client_id=self._client_id, client_secret=self._client_secret, **metadata, ) token_storage = self._token_storage or MemoryStore() if isinstance(token_storage, MemoryStore): from warnings import warn warn( message="Using in-memory token storage -- tokens will be lost when the client restarts. " "For persistent storage across multiple MCP servers, provide an encrypted AsyncKeyValue backend. " "See https://gofastmcp.com/clients/auth/oauth#token-storage for details.", stacklevel=2, ) # Use full URL for token storage to properly separate tokens per MCP endpoint self.token_storage_adapter: TokenStorageAdapter = TokenStorageAdapter( async_key_value=token_storage, server_url=mcp_url ) self.mcp_url = mcp_url super().__init__( server_url=mcp_url, client_metadata=client_metadata, storage=self.token_storage_adapter, redirect_handler=self.redirect_handler, callback_handler=self.callback_handler, timeout=self._callback_timeout, client_metadata_url=self._client_metadata_url, ) self._bound = True async def _initialize(self) -> None: """Load stored tokens and client info, properly setting token expiry.""" await super()._initialize() if self._static_client_info is not None: self.context.client_info = self._static_client_info await self.token_storage_adapter.set_client_info(self._static_client_info) if self.context.current_tokens and self.context.current_tokens.expires_in: stored_expiry = await self.token_storage_adapter.get_token_expiry() if stored_expiry is not None: self.context.token_expiry_time = stored_expiry else: self.context.update_token_expiry(self.context.current_tokens) async def redirect_handler(self, authorization_url: str) -> None: """Open browser for authorization, with pre-flight check for invalid client.""" # Pre-flight check to detect invalid client_id before opening browser async with self.httpx_client_factory() as client: response = await client.get(authorization_url, follow_redirects=False) # Check for client not found error (400 typically means bad client_id) if response.status_code == 400: raise ClientNotFoundError( "OAuth client not found - cached credentials may be stale" ) # OAuth typically returns redirects, but some providers return 200 with HTML login pages if response.status_code not in (200, 302, 303, 307, 308): raise RuntimeError( f"Unexpected authorization response: {response.status_code}" ) logger.info(f"OAuth authorization URL: {authorization_url}") webbrowser.open(authorization_url) async def callback_handler(self) -> tuple[str, str | None]: """Handle OAuth callback and return (auth_code, state).""" # Create result container and event to capture the OAuth response result = OAuthCallbackResult() result_ready = anyio.Event() # Create server with result tracking server: Server = create_oauth_callback_server( port=self.redirect_port, host=self._callback_host, server_url=self.mcp_url, result_container=result, result_ready=result_ready, ) # 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://{self._callback_host}:{self.redirect_port}" ) try: with anyio.fail_after(self._callback_timeout): await result_ready.wait() if result.error: raise result.error return result.code, result.state # type: ignore except TimeoutError as e: raise TimeoutError( f"OAuth callback timed out after {self._callback_timeout} seconds" ) from e finally: server.should_exit = True await anyio.sleep(0.1) # Allow server to shut down gracefully tg.cancel_scope.cancel() raise RuntimeError("OAuth callback handler could not be started") async def async_auth_flow( self, request: httpx.Request ) -> AsyncGenerator[httpx.Request, httpx.Response]: """HTTPX auth flow with automatic retry on stale cached credentials. If the OAuth flow fails due to invalid/stale client credentials, clears the cache and retries once with fresh registration. """ if not self._bound: raise RuntimeError( "OAuth provider has no server URL. Either pass mcp_url to OAuth() " "or use it with Client(auth=...) which provides the URL automatically." ) try: # First attempt with potentially cached credentials async with aclosing(super().async_auth_flow(request)) as gen: response = None while True: try: # First iteration sends None, subsequent iterations send response yielded_request = await gen.asend(response) # ty: ignore[invalid-argument-type] response = yield yielded_request except StopAsyncIteration: break except ClientNotFoundError: # Static credentials are fixed — retrying won't help. Surface the # error so the user can correct their client_id / client_secret. if self._static_client_info is not None: raise ClientNotFoundError( "OAuth server rejected the static client credentials. " "Verify that the client_id (and client_secret, if provided) " "are correct and that the client is registered with the server." ) from None logger.debug( "OAuth client not found on server, clearing cache and retrying..." ) # Clear cached state and retry once self._initialized = False await self.token_storage_adapter.clear() # Retry with fresh registration async with aclosing(super().async_auth_flow(request)) as gen: response = None while True: try: yielded_request = await gen.asend(response) # ty: ignore[invalid-argument-type] response = yield yielded_request except StopAsyncIteration: break