add fallback

Author: ihrpr

src/mcp/client/auth.py

@@ -106,6 +106,10 @@ class OAuthContext:
     # State
     lock: anyio.Lock = field(default_factory=anyio.Lock)
 
+    # Discovery state for fallback support
+    discovery_base_url: str | None = None
+    discovery_pathname: str | None = None
+
     def get_authorization_base_url(self, server_url: str) -> str:
         """Extract base URL by removing path component."""
         parsed = urlparse(server_url)
@@ -197,26 +201,53 @@ async def _handle_protected_resource_response(self, response: httpx.Response) ->
             except ValidationError:
                 pass
 
+    def _build_well_known_path(self, pathname: str) -> str:
+        """Construct well-known path for OAuth metadata discovery."""
+        well_known_path = f"/.well-known/oauth-authorization-server{pathname}"
+        if pathname.endswith("/"):
+            # Strip trailing slash from pathname to avoid double slashes
+            well_known_path = well_known_path[:-1]
+        return well_known_path
+
+    def _should_attempt_fallback(self, response_status: int, pathname: str) -> bool:
+        """Determine if fallback to root discovery should be attempted."""
+        return response_status == 404 and pathname != "/"
+
+    async def _try_metadata_discovery(self, url: str) -> httpx.Request:
+        """Build metadata discovery request for a specific URL."""
+        return httpx.Request("GET", url, headers={MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION})
+
     async def _discover_oauth_metadata(self) -> httpx.Request:
-        """Build OAuth metadata discovery request."""
+        """Build OAuth metadata discovery request with fallback support."""
         if self.context.auth_server_url:
             auth_server_url = self.context.auth_server_url
         else:
             auth_server_url = self.context.server_url
 
-        # Per RFC 8414, preserve the path component when constructing discovery URL
+        # Per RFC 8414, try path-aware discovery first
         parsed = urlparse(auth_server_url)
-        well_known_path = f"/.well-known/oauth-authorization-server{parsed.path}"
-        if parsed.path.endswith("/"):
-            # Strip trailing slash from pathname
-            well_known_path = well_known_path[:-1]
-
+        well_known_path = self._build_well_known_path(parsed.path)
         base_url = f"{parsed.scheme}://{parsed.netloc}"
         url = urljoin(base_url, well_known_path)
-        return httpx.Request("GET", url, headers={MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION})
 
-    async def _handle_oauth_metadata_response(self, response: httpx.Response) -> None:
-        """Handle OAuth metadata response."""
+        # Store fallback info for use in response handler
+        self.context.discovery_base_url = base_url
+        self.context.discovery_pathname = parsed.path
+
+        return await self._try_metadata_discovery(url)
+
+    async def _discover_oauth_metadata_fallback(self) -> httpx.Request:
+        """Build fallback OAuth metadata discovery request for legacy servers."""
+        base_url = getattr(self.context, "discovery_base_url", "")
+        if not base_url:
+            raise OAuthFlowError("No base URL available for fallback discovery")
+
+        # Fallback to root discovery for legacy servers
+        url = urljoin(base_url, "/.well-known/oauth-authorization-server")
+        return await self._try_metadata_discovery(url)
+
+    async def _handle_oauth_metadata_response(self, response: httpx.Response, is_fallback: bool = False) -> bool:
+        """Handle OAuth metadata response. Returns True if handled successfully."""
         if response.status_code == 200:
             try:
                 content = await response.aread()
@@ -225,9 +256,18 @@ async def _handle_oauth_metadata_response(self, response: httpx.Response) -> Non
                 # Apply default scope if none specified
                 if self.context.client_metadata.scope is None and metadata.scopes_supported is not None:
                     self.context.client_metadata.scope = " ".join(metadata.scopes_supported)
+                return True
             except ValidationError:
                 pass
 
+        # Check if we should attempt fallback (404 on path-aware discovery)
+        if not is_fallback and self._should_attempt_fallback(
+            response.status_code, getattr(self.context, "discovery_pathname", "/")
+        ):
+            return False  # Signal that fallback should be attempted
+
+        return True  # Signal no fallback needed (either success or non-404 error)
+
     async def _register_client(self) -> httpx.Request | None:
         """Build registration request or skip if already registered."""
         if self.context.client_info:
@@ -426,10 +466,16 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                     discovery_response = yield discovery_request
                     await self._handle_protected_resource_response(discovery_response)
 
-                    # Step 2: Discover OAuth metadata
+                    # Step 2: Discover OAuth metadata (with fallback for legacy servers)
                     oauth_request = await self._discover_oauth_metadata()
                     oauth_response = yield oauth_request
-                    await self._handle_oauth_metadata_response(oauth_response)
+                    handled = await self._handle_oauth_metadata_response(oauth_response, is_fallback=False)
+
+                    # If path-aware discovery failed with 404, try fallback to root
+                    if not handled:
+                        fallback_request = await self._discover_oauth_metadata_fallback()
+                        fallback_response = yield fallback_request
+                        await self._handle_oauth_metadata_response(fallback_response, is_fallback=True)
 
                     # Step 3: Register client if needed
                     registration_request = await self._register_client()
@@ -472,10 +518,16 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                     discovery_response = yield discovery_request
                     await self._handle_protected_resource_response(discovery_response)
 
-                    # Step 2: Discover OAuth metadata
+                    # Step 2: Discover OAuth metadata (with fallback for legacy servers)
                     oauth_request = await self._discover_oauth_metadata()
                     oauth_response = yield oauth_request
-                    await self._handle_oauth_metadata_response(oauth_response)
+                    handled = await self._handle_oauth_metadata_response(oauth_response, is_fallback=False)
+
+                    # If path-aware discovery failed with 404, try fallback to root
+                    if not handled:
+                        fallback_request = await self._discover_oauth_metadata_fallback()
+                        fallback_response = yield fallback_request
+                        await self._handle_oauth_metadata_response(fallback_response, is_fallback=True)
 
                     # Step 3: Register client if needed
                     registration_request = await self._register_client()

tests/client/test_auth.py

@@ -260,6 +260,107 @@ async def callback_handler() -> tuple[str, str | None]:
         assert str(request.url) == "https://api.example.com/.well-known/oauth-authorization-server/v1/mcp"
         assert "mcp-protocol-version" in request.headers
 
+
+class TestOAuthFallback:
+    """Test OAuth discovery fallback behavior for legacy (act as AS not RS) servers."""
+
+    @pytest.mark.anyio
+    async def test_fallback_discovery_request(self, client_metadata, mock_storage):
+        """Test fallback discovery request building."""
+
+        async def redirect_handler(url: str) -> None:
+            pass
+
+        async def callback_handler() -> tuple[str, str | None]:
+            return "test_auth_code", "test_state"
+
+        provider = OAuthClientProvider(
+            server_url="https://api.example.com/v1/mcp",
+            client_metadata=client_metadata,
+            storage=mock_storage,
+            redirect_handler=redirect_handler,
+            callback_handler=callback_handler,
+        )
+
+        # Set up discovery state manually as if path-aware discovery was attempted
+        provider.context.discovery_base_url = "https://api.example.com"
+        provider.context.discovery_pathname = "/v1/mcp"
+
+        # Test fallback request building
+        request = await provider._discover_oauth_metadata_fallback()
+
+        assert request.method == "GET"
+        assert str(request.url) == "https://api.example.com/.well-known/oauth-authorization-server"
+        assert "mcp-protocol-version" in request.headers
+
+    @pytest.mark.anyio
+    async def test_should_attempt_fallback(self, oauth_provider):
+        """Test fallback decision logic."""
+        # Should attempt fallback on 404 with non-root path
+        assert oauth_provider._should_attempt_fallback(404, "/v1/mcp")
+
+        # Should NOT attempt fallback on 404 with root path
+        assert not oauth_provider._should_attempt_fallback(404, "/")
+
+        # Should NOT attempt fallback on other status codes
+        assert not oauth_provider._should_attempt_fallback(200, "/v1/mcp")
+        assert not oauth_provider._should_attempt_fallback(500, "/v1/mcp")
+
+    @pytest.mark.anyio
+    async def test_handle_metadata_response_success(self, oauth_provider):
+        """Test successful metadata response handling."""
+        # Create minimal valid OAuth metadata
+        content = b"""{
+            "issuer": "https://auth.example.com",
+            "authorization_endpoint": "https://auth.example.com/authorize", 
+            "token_endpoint": "https://auth.example.com/token"
+        }"""
+        response = httpx.Response(200, content=content)
+
+        # Should return True (success) and set metadata
+        result = await oauth_provider._handle_oauth_metadata_response(response, is_fallback=False)
+        assert result is True
+        assert oauth_provider.context.oauth_metadata is not None
+        assert str(oauth_provider.context.oauth_metadata.issuer) == "https://auth.example.com/"
+
+    @pytest.mark.anyio
+    async def test_handle_metadata_response_404_needs_fallback(self, oauth_provider):
+        """Test 404 response handling that should trigger fallback."""
+        # Set up discovery state for non-root path
+        oauth_provider.context.discovery_base_url = "https://api.example.com"
+        oauth_provider.context.discovery_pathname = "/v1/mcp"
+
+        # Mock 404 response
+        response = httpx.Response(404)
+
+        # Should return False (needs fallback)
+        result = await oauth_provider._handle_oauth_metadata_response(response, is_fallback=False)
+        assert result is False
+
+    @pytest.mark.anyio
+    async def test_handle_metadata_response_404_no_fallback_needed(self, oauth_provider):
+        """Test 404 response handling when no fallback is needed."""
+        # Set up discovery state for root path
+        oauth_provider.context.discovery_base_url = "https://api.example.com"
+        oauth_provider.context.discovery_pathname = "/"
+
+        # Mock 404 response
+        response = httpx.Response(404)
+
+        # Should return True (no fallback needed)
+        result = await oauth_provider._handle_oauth_metadata_response(response, is_fallback=False)
+        assert result is True
+
+    @pytest.mark.anyio
+    async def test_handle_metadata_response_404_fallback_attempt(self, oauth_provider):
+        """Test 404 response handling during fallback attempt."""
+        # Mock 404 response during fallback
+        response = httpx.Response(404)
+
+        # Should return True (fallback attempt complete, no further action needed)
+        result = await oauth_provider._handle_oauth_metadata_response(response, is_fallback=True)
+        assert result is True
+
     @pytest.mark.anyio
     async def test_register_client_request(self, oauth_provider):
         """Test client registration request building."""