Merge branch 'main' into quickstart

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

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