Simplify

Signed-off-by: Sid Murching <sid.murching@databricks.com>

Author: smurching

src/mcp/client/auth.py

@@ -108,10 +108,6 @@ 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)
@@ -228,23 +224,23 @@ def _extract_resource_metadata_from_www_auth(self, init_response: httpx.Response
 
         return None
 
-    async def _discover_protected_resource(self, init_response: httpx.Response) -> httpx.Request:
-        # RFC9728: Try to extract resource_metadata URL from WWW-Authenticate header of the initial response
-        url = self._extract_resource_metadata_from_www_auth(init_response)
-
-        if not url:
-            # Fallback to well-known discovery with path component included
-            parsed = urlparse(self.context.server_url)
-            auth_base_url = f"{parsed.scheme}://{parsed.netloc}"
+    def _get_protected_resource_discovery_urls(self) -> list[str]:
+        """Generate ordered list of URLs for protected resource discovery attempts."""
+        urls: list[str] = []
+        parsed = urlparse(self.context.server_url)
+        base_url = f"{parsed.scheme}://{parsed.netloc}"
 
-            if parsed.path and parsed.path != "/":
-                # Include path component in the well-known URL
-                path_component = parsed.path.rstrip("/")
-                url = urljoin(auth_base_url, f"/.well-known/oauth-protected-resource{path_component}")
-            else:
-                url = urljoin(auth_base_url, "/.well-known/oauth-protected-resource")
+        if parsed.path and parsed.path != "/":
+            # Try path-specific endpoint first
+            path_component = parsed.path.rstrip("/")
+            urls.append(urljoin(base_url, f"/.well-known/oauth-protected-resource{path_component}"))
+            # Then fallback to base endpoint
+            urls.append(urljoin(base_url, "/.well-known/oauth-protected-resource"))
+        else:
+            # No path, just use base endpoint
+            urls.append(urljoin(base_url, "/.well-known/oauth-protected-resource"))
 
-        return httpx.Request("GET", url, headers={MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION})
+        return urls
 
     async def _handle_protected_resource_response(self, response: httpx.Response) -> None:
         """Handle discovery response."""
@@ -517,9 +513,28 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                 try:
                     # OAuth flow must be inline due to generator constraints
                     # Step 1: Discover protected resource metadata (RFC9728 with WWW-Authenticate support)
-                    discovery_request = await self._discover_protected_resource(response)
-                    discovery_response = yield discovery_request
-                    await self._handle_protected_resource_response(discovery_response)
+                    # Check if WWW-Authenticate provides resource_metadata URL first
+                    www_auth_url = self._extract_resource_metadata_from_www_auth(response)
+                    if www_auth_url:
+                        discovery_request = httpx.Request(
+                            "GET", www_auth_url, headers={MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION}
+                        )
+                        discovery_response = yield discovery_request
+                        await self._handle_protected_resource_response(discovery_response)
+                    else:
+                        # Try well-known discovery URLs with fallback
+                        discovery_urls = self._get_protected_resource_discovery_urls()
+                        for url in discovery_urls:
+                            discovery_request = httpx.Request(
+                                "GET", url, headers={MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION}
+                            )
+                            discovery_response = yield discovery_request
+
+                            if discovery_response.status_code == 200:
+                                await self._handle_protected_resource_response(discovery_response)
+                                break  # Success, stop trying other URLs
+                            elif discovery_response.status_code != 404:
+                                break  # Non-404 error, stop trying
 
                     # Step 2: Discover OAuth metadata (with fallback for legacy servers)
                     discovery_urls = self._get_discovery_urls()

tests/client/test_auth.py

@@ -198,15 +198,16 @@ class TestOAuthFlow:
     """Test OAuth flow methods."""
 
     @pytest.mark.anyio
-    async def test_discover_protected_resource_request(self, client_metadata, mock_storage):
-        """Test protected resource discovery request building maintains backward compatibility."""
+    async def test_protected_resource_discovery_urls(self, client_metadata, mock_storage):
+        """Test protected resource discovery URL generation with fallback."""
 
         async def redirect_handler(url: str) -> None:
             pass
 
         async def callback_handler() -> tuple[str, str | None]:
             return "test_auth_code", "test_state"
 
+        # Test with path component
         provider = OAuthClientProvider(
             server_url="https://api.example.com/api/2.0/mcp",
             client_metadata=client_metadata,
@@ -215,25 +216,23 @@ async def callback_handler() -> tuple[str, str | None]:
             callback_handler=callback_handler,
         )
 
-        # Test without WWW-Authenticate (fallback)
-        init_response = httpx.Response(
-            status_code=401, headers={}, request=httpx.Request("GET", "https://request-api.example.com")
-        )
-
-        request = await provider._discover_protected_resource(init_response)
-        assert request.method == "GET"
-        assert str(request.url) == "https://api.example.com/.well-known/oauth-protected-resource/api/2.0/mcp"
-        assert "mcp-protocol-version" in request.headers
+        urls = provider._get_protected_resource_discovery_urls()
+        assert urls == [
+            "https://api.example.com/.well-known/oauth-protected-resource/api/2.0/mcp",
+            "https://api.example.com/.well-known/oauth-protected-resource",
+        ]
 
-        # Test with WWW-Authenticate header
-        init_response.headers["WWW-Authenticate"] = (
-            'Bearer resource_metadata="https://prm.example.com/.well-known/oauth-protected-resource/path"'
+        # Test without path component
+        provider = OAuthClientProvider(
+            server_url="https://api.example.com",
+            client_metadata=client_metadata,
+            storage=mock_storage,
+            redirect_handler=redirect_handler,
+            callback_handler=callback_handler,
         )
 
-        request = await provider._discover_protected_resource(init_response)
-        assert request.method == "GET"
-        assert str(request.url) == "https://prm.example.com/.well-known/oauth-protected-resource/path"
-        assert "mcp-protocol-version" in request.headers
+        urls = provider._get_protected_resource_discovery_urls()
+        assert urls == ["https://api.example.com/.well-known/oauth-protected-resource"]
 
     @pytest.mark.anyio
     def test_create_oauth_metadata_request(self, oauth_provider):