diff --git a/src/fastmcp/client/auth/oauth.py b/src/fastmcp/client/auth/oauth.py index b589c2873..c703950d4 100644 --- a/src/fastmcp/client/auth/oauth.py +++ b/src/fastmcp/client/auth/oauth.py @@ -4,6 +4,7 @@ import asyncio import json import webbrowser from asyncio import Future +from collections.abc import AsyncGenerator from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any, Literal @@ -34,6 +35,12 @@ __all__ = ["OAuth"] logger = get_logger(__name__) +class ClientNotFoundError(Exception): + """Raised when OAuth client credentials are not found on the server.""" + + pass + + class StoredToken(BaseModel): """Token storage format with absolute expiry time.""" @@ -173,7 +180,7 @@ class FileTokenStorage(TokenStorage): 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)}") + logger.debug(f"Cleared OAuth cache for {self.get_base_url(self.server_url)}") @classmethod def clear_all(cls, cache_dir: Path | None = None) -> None: @@ -300,7 +307,23 @@ class OAuth(OAuthClientProvider): self.context.update_token_expiry(self.context.current_tokens) async def redirect_handler(self, authorization_url: str) -> None: - """Open browser for authorization.""" + """Open browser for authorization, with pre-flight check for invalid client.""" + # Pre-flight check to detect invalid client_id before opening browser + async with httpx.AsyncClient() 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" + ) + + # For any non-redirect response, something is wrong + if response.status_code not in (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) @@ -336,3 +359,56 @@ class OAuth(OAuthClientProvider): 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. + """ + try: + # First attempt with potentially cached credentials + gen = super().async_auth_flow(request) + response = None + while True: + try: + yielded_request = await gen.asend(response) + response = yield yielded_request + except StopAsyncIteration: + break + + except ClientNotFoundError: + logger.debug( + "OAuth client not found on server, clearing cache and retrying..." + ) + + # Clear cached state and retry once + self._initialized = False + + # Try to clear storage if it supports it + if hasattr(self.context.storage, "clear"): + try: + self.context.storage.clear() + except Exception as e: + logger.warning(f"Failed to clear OAuth storage cache: {e}") + # Can't retry without clearing cache, re-raise original error + raise ClientNotFoundError( + "OAuth client not found and cache could not be cleared" + ) from e + else: + logger.warning( + "Storage does not support clear() - cannot retry with fresh credentials" + ) + # Can't retry without clearing cache, re-raise original error + raise + + gen = super().async_auth_flow(request) + response = None + while True: + try: + yielded_request = await gen.asend(response) + response = yield yielded_request + except StopAsyncIteration: + break diff --git a/src/fastmcp/server/auth/oauth_proxy.py b/src/fastmcp/server/auth/oauth_proxy.py index f78f7389e..a75f8b690 100644 --- a/src/fastmcp/server/auth/oauth_proxy.py +++ b/src/fastmcp/server/auth/oauth_proxy.py @@ -172,7 +172,6 @@ class OAuthProxy(OAuthProvider): 1. Client Registration (DCR): - Accept any client registration request - Store ProxyDCRClient that accepts dynamic redirect URIs - - Return shared upstream credentials to all clients 2. Authorization: - Store transaction mapping client details to proxy flow @@ -323,67 +322,28 @@ class OAuthProxy(OAuthProvider): # ------------------------------------------------------------------------- async def get_client(self, client_id: str) -> OAuthClientInformationFull | None: - """Get client information by ID. + """Get client information by ID. This is generally the random ID + provided to the DCR client during registration, not the upstream client ID. - For unregistered clients, returns a ProxyDCRClient that accepts - any localhost redirect URI for DCR clients. - - Even registered clients use ProxyDCRClient to ensure they can - authenticate with different dynamic ports on reconnection. This - handles the case where a client with cached tokens reconnects - on a different port. + For unregistered clients, returns None (which will raise an error in the SDK). """ client = self._clients.get(client_id) - if client is None: - # For unregistered DCR clients, create a permissive client - # that will accept any localhost redirect URI - # We need at least one URI for Pydantic validation, but our custom - # validate_redirect_uri will accept any localhost URI - client = ProxyDCRClient( - client_id=client_id, - client_secret=None, - redirect_uris=[ - AnyUrl("http://localhost") - ], # Placeholder, validation uses allowed_patterns - grant_types=["authorization_code", "refresh_token"], - scope=self._default_scope_str, - token_endpoint_auth_method="none", - allowed_redirect_uri_patterns=self._allowed_client_redirect_uris, - ) - logger.debug("Created ProxyDCRClient for unregistered client %s", client_id) - return client async def register_client(self, client_info: OAuthClientInformationFull) -> None: - """Register a client locally using fixed upstream credentials. + """Register a client locally - This implementation always uses the upstream client_id and client_secret - regardless of what the client requests. It modifies the client_info object - in place since the MCP framework ignores return values. - - This ensures all clients use the same credentials that are registered - with the upstream server. - - Implementation Detail: - We store a ProxyDCRClient (not the original client_info) to ensure - the client can reconnect with different dynamic redirect URIs. This is - essential for cached token scenarios where the client port changes. - - The flow: - 1. Client provides its desired redirect URIs (dynamic localhost ports) - 2. We create a ProxyDCRClient that will accept ANY localhost URI - 3. We store this flexible client for future authentications - 4. When client reconnects with a different port, ProxyDCRClient accepts it + When a client registers, we create a ProxyDCRClient that is more + forgiving about validating redirect URIs, since the DCR client's + redirect URI will likely be localhost or unknown to the proxied IDP. The + proxied IDP only knows about this server's fixed redirect URI. """ - # Always use the upstream credentials - upstream_id = self._upstream_client_id - upstream_secret = self._upstream_client_secret.get_secret_value() # Create a ProxyDCRClient with configured redirect URI validation proxy_client = ProxyDCRClient( - client_id=upstream_id, - client_secret=upstream_secret, + client_id=client_info.client_id, + client_secret=client_info.client_secret, redirect_uris=client_info.redirect_uris or [AnyUrl("http://localhost")], grant_types=client_info.grant_types or ["authorization_code", "refresh_token"], @@ -392,8 +352,8 @@ class OAuthProxy(OAuthProvider): allowed_redirect_uri_patterns=self._allowed_client_redirect_uris, ) - # Store the ProxyDCRClient using the upstream ID - self._clients[upstream_id] = proxy_client + # Store the ProxyDCRClient + self._clients[client_info.client_id] = proxy_client # Log redirect URIs to help users discover what patterns they might need if client_info.redirect_uris: @@ -406,7 +366,7 @@ class OAuthProxy(OAuthProvider): logger.debug( "Registered client %s with %d redirect URIs", - upstream_id, + client_info.client_id, len(proxy_client.redirect_uris), )