refactor: pull more oauth helpers out

Author: maxisbey

src/mcp/client/auth/__init__.py

@@ -4,11 +4,9 @@
 Implements authorization code flow with PKCE and automatic token refresh.
 """
 
+from mcp.client.auth.exceptions import OAuthFlowError, OAuthRegistrationError, OAuthTokenError
 from mcp.client.auth.oauth2 import (
     OAuthClientProvider,
-    OAuthFlowError,
-    OAuthRegistrationError,
-    OAuthTokenError,
     PKCEParameters,
     TokenStorage,
 )

src/mcp/client/auth/exceptions.py

@@ -0,0 +1,10 @@
+class OAuthFlowError(Exception):
+    """Base exception for OAuth flow errors."""
+
+
+class OAuthTokenError(OAuthFlowError):
+    """Raised when token operations fail."""
+
+
+class OAuthRegistrationError(OAuthFlowError):
+    """Raised when client registration fails."""

src/mcp/client/auth/oauth2.py

@@ -19,14 +19,19 @@
 import httpx
 from pydantic import BaseModel, Field, ValidationError
 
+from mcp.client.auth import OAuthFlowError, OAuthTokenError
 from mcp.client.auth.utils import (
     build_protected_resource_discovery_urls,
+    create_client_registration_request,
+    create_oauth_metadata_request,
     extract_field_from_www_auth,
     extract_resource_metadata_from_www_auth,
     extract_scope_from_www_auth,
     get_client_metadata_scopes,
     get_discovery_urls,
+    handle_auth_metadata_response,
     handle_protected_resource_response,
+    handle_registration_response,
 )
 from mcp.client.streamable_http import MCP_PROTOCOL_VERSION
 from mcp.shared.auth import (
@@ -37,23 +42,10 @@
     ProtectedResourceMetadata,
 )
 from mcp.shared.auth_utils import check_resource_allowed, resource_url_from_server_url
-from mcp.types import LATEST_PROTOCOL_VERSION
 
 logger = logging.getLogger(__name__)
 
 
-class OAuthFlowError(Exception):
-    """Base exception for OAuth flow errors."""
-
-
-class OAuthTokenError(OAuthFlowError):
-    """Raised when token operations fail."""
-
-
-class OAuthRegistrationError(OAuthFlowError):
-    """Raised when client registration fails."""
-
-
 class PKCEParameters(BaseModel):
     """PKCE (Proof Key for Code Exchange) parameters."""
 
@@ -255,19 +247,19 @@ async def _register_client(self) -> httpx.Request | None:
             "POST", registration_url, json=registration_data, headers={"Content-Type": "application/json"}
         )
 
-    async def _handle_registration_response(self, response: httpx.Response) -> None:
-        """Handle registration response."""
-        if response.status_code not in (200, 201):
-            await response.aread()
-            raise OAuthRegistrationError(f"Registration failed: {response.status_code} {response.text}")
-
-        try:
-            content = await response.aread()
-            client_info = OAuthClientInformationFull.model_validate_json(content)
-            self.context.client_info = client_info
-            await self.context.storage.set_client_info(client_info)
-        except ValidationError as e:
-            raise OAuthRegistrationError(f"Invalid registration response: {e}")
+    # async def _handle_registration_response(self, response: httpx.Response) -> None:
+    #     """Handle registration response."""
+    #     if response.status_code not in (200, 201):
+    #         await response.aread()
+    #         raise OAuthRegistrationError(f"Registration failed: {response.status_code} {response.text}")
+    #
+    #     try:
+    #         content = await response.aread()
+    #         client_info = OAuthClientInformationFull.model_validate_json(content)
+    #         self.context.client_info = client_info
+    #         await self.context.storage.set_client_info(client_info)
+    #     except ValidationError as e:
+    #         raise OAuthRegistrationError(f"Invalid registration response: {e}")
 
     async def _perform_authorization(self) -> httpx.Request:
         """Perform the authorization flow."""
@@ -456,9 +448,6 @@ def _add_auth_header(self, request: httpx.Request) -> None:
         if self.context.current_tokens and self.context.current_tokens.access_token:
             request.headers["Authorization"] = f"Bearer {self.context.current_tokens.access_token}"
 
-    def _create_oauth_metadata_request(self, url: str) -> httpx.Request:
-        return httpx.Request("GET", url, headers={MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION})
-
     async def _handle_oauth_metadata_response(self, response: httpx.Response) -> None:
         content = await response.aread()
         metadata = OAuthMetadata.model_validate_json(content)
@@ -494,54 +483,63 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                     www_auth_resource_metadata_url = extract_resource_metadata_from_www_auth(response)
 
                     # Step 1: Discover protected resource metadata (SEP-985 with fallback support)
-                    discovery_urls = build_protected_resource_discovery_urls(
+                    prm_discovery_urls = build_protected_resource_discovery_urls(
                         www_auth_resource_metadata_url, self.context.server_url
                     )
-                    discovery_success = False
-                    for url in discovery_urls:
-                        discovery_request = httpx.Request(
-                            "GET", url, headers={MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION}
-                        )
-                        discovery_response = yield discovery_request
-                        discovery_success, prm, auth_server_url = await handle_protected_resource_response(
-                            discovery_response
-                        )
-                        if discovery_success:
-                            assert prm is not None
-                            assert auth_server_url is not None
+                    prm_discovery_success = False
+                    for url in prm_discovery_urls:
+                        discovery_request = create_oauth_metadata_request(url)
+
+                        discovery_response = yield discovery_request  # sending request
+
+                        prm = await handle_protected_resource_response(discovery_response)
+                        if prm:
+                            prm_discovery_success = True
+
+                            # saving the response metadata
                             self.context.protected_resource_metadata = prm
                             if prm.authorization_servers:
-                                self.context.auth_server_url = auth_server_url
-                            break
+                                self.context.auth_server_url = str(prm.authorization_servers[0])
 
-                    if not discovery_success:
+                            break
+                        else:
+                            logger.debug(f"Protected resource metadata discovery failed: {url}")
+                    if not prm_discovery_success:
                         raise OAuthFlowError("Protected resource metadata discovery failed: no valid metadata found")
 
-                    # Step 2: Apply scope selection strategy
-                    self.context.client_metadata.scope = get_client_metadata_scopes(
-                        www_auth_resource_metadata_url, self.context.protected_resource_metadata
-                    )
-
-                    # Step 3: Discover OAuth metadata (with fallback for legacy servers)
-                    discovery_urls = get_discovery_urls(self.context.auth_server_url or self.context.server_url)
-                    for url in discovery_urls:
-                        oauth_metadata_request = self._create_oauth_metadata_request(url)
+                    # Step 2: Discover OAuth metadata (with fallback for legacy servers)
+                    asm_discovery_urls = get_discovery_urls(self.context.auth_server_url or self.context.server_url)
+                    for url in asm_discovery_urls:
+                        oauth_metadata_request = create_oauth_metadata_request(url)
                         oauth_metadata_response = yield oauth_metadata_request
 
-                        if oauth_metadata_response.status_code == 200:
-                            try:
-                                await self._handle_oauth_metadata_response(oauth_metadata_response)
-                                break
-                            except ValidationError:
-                                continue
-                        elif oauth_metadata_response.status_code < 400 or oauth_metadata_response.status_code >= 500:
-                            break  # Non-4XX error, stop trying
+                        ok, asm = await handle_auth_metadata_response(oauth_metadata_response)
+                        if not ok:
+                            break
+                        if ok and asm:
+                            self.context.oauth_metadata = asm
+                            break
+                        else:
+                            logger.debug(f"OAuth metadata discovery failed: {url}")
+
+                    # Step 3: Apply scope selection strategy
+                    self.context.client_metadata.scope = get_client_metadata_scopes(
+                        www_auth_resource_metadata_url,
+                        self.context.protected_resource_metadata,
+                        self.context.oauth_metadata,
+                    )
 
                     # Step 4: Register client if needed
-                    registration_request = await self._register_client()
-                    if registration_request:
+                    registration_request = create_client_registration_request(
+                        self.context.oauth_metadata,
+                        self.context.client_metadata,
+                        self.context.get_authorization_base_url(self.context.server_url),
+                    )
+                    if not self.context.client_info:
                         registration_response = yield registration_request
-                        await self._handle_registration_response(registration_response)
+                        client_information = await handle_registration_response(registration_response)
+                        self.context.client_info = client_information
+                        await self.context.storage.set_client_info(client_information)
 
                     # Step 5: Perform authorization and complete token exchange
                     token_response = yield await self._perform_authorization()

src/mcp/client/auth/utils.py

@@ -2,11 +2,13 @@
 import re
 from urllib.parse import urljoin, urlparse
 
-from httpx import Response
+from httpx import Request, Response
 from pydantic import ValidationError
 
-from mcp.client.auth import OAuthFlowError
-from mcp.shared.auth import ProtectedResourceMetadata
+from mcp.client.auth import OAuthRegistrationError
+from mcp.client.streamable_http import MCP_PROTOCOL_VERSION
+from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthMetadata, ProtectedResourceMetadata
+from mcp.types import LATEST_PROTOCOL_VERSION
 
 logger = logging.getLogger(__name__)
 
@@ -95,7 +97,9 @@ def build_protected_resource_discovery_urls(www_auth_url: str | None, server_url
 
 
 def get_client_metadata_scopes(
-    www_authenticate_scope: str | None, protected_resource_metadata: ProtectedResourceMetadata | None
+    www_authenticate_scope: str | None,
+    protected_resource_metadata: ProtectedResourceMetadata | None,
+    authorization_server_metadata: OAuthMetadata | None = None,
 ) -> str | None:
     """Select scopes as outlined in the 'Scope Selection Strategy' in the MCP spec."""
     # Per MCP spec, scope selection priority order:
@@ -109,6 +113,8 @@ def get_client_metadata_scopes(
     elif protected_resource_metadata is not None and protected_resource_metadata.scopes_supported is not None:
         # Priority 2: PRM scopes_supported
         return " ".join(protected_resource_metadata.scopes_supported)
+    elif authorization_server_metadata is not None and authorization_server_metadata.scopes_supported is not None:
+        return " ".join(authorization_server_metadata.scopes_supported)
     else:
         # Priority 3: Omit scope parameter
         return None
@@ -143,7 +149,7 @@ def get_discovery_urls(auth_server_url: str) -> list[str]:
 
 async def handle_protected_resource_response(
     response: Response,
-) -> tuple[bool, ProtectedResourceMetadata | None, str | None]:
+) -> ProtectedResourceMetadata | None:
     """
     Handle protected resource metadata discovery response.
 
@@ -156,22 +162,96 @@ async def handle_protected_resource_response(
         try:
             content = await response.aread()
             metadata = ProtectedResourceMetadata.model_validate_json(content)
-            auth_server_url: str | None = None
-            if metadata.authorization_servers:
-                auth_server_url = str(metadata.authorization_servers[0])
-            return True, metadata, auth_server_url
+            return metadata
 
         except ValidationError:
             # Invalid metadata - try next URL
-            logger.warning(f"Invalid protected resource metadata at {response.request.url}")
-            return False, None, None
-    elif response.status_code == 404:
+            return None
+    else:
         # Not found - try next URL in fallback chain
-        logger.debug(f"Protected resource metadata not found at {response.request.url}, trying next URL")
-        return False, None, None
+        return None
+
+
+async def handle_auth_metadata_response(response: Response) -> tuple[bool, OAuthMetadata | None]:
+    if response.status_code == 200:
+        try:
+            content = await response.aread()
+            asm = OAuthMetadata.model_validate_json(content)
+            return True, asm
+        except ValidationError:
+            return True, None
+    elif response.status_code < 400 or response.status_code >= 500:
+        return False, None  # Non-4XX error, stop trying
+    return True, None
+
+
+def create_oauth_metadata_request(url: str) -> Request:
+    return Request("GET", url, headers={MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION})
+
+
+def create_client_registration_request(
+    auth_server_metadata: OAuthMetadata | None, client_metadata: OAuthClientMetadata, auth_base_url: str
+) -> Request:
+    """Build registration request or skip if already registered."""
+
+    if auth_server_metadata and auth_server_metadata.registration_endpoint:
+        registration_url = str(auth_server_metadata.registration_endpoint)
     else:
-        # Other error - fail immediately
-        raise OAuthFlowError(f"Protected Resource Metadata request failed: {response.status_code}")
+        registration_url = urljoin(auth_base_url, "/register")
+
+    registration_data = client_metadata.model_dump(by_alias=True, mode="json", exclude_none=True)
+
+    return Request("POST", registration_url, json=registration_data, headers={"Content-Type": "application/json"})
+
+
+async def handle_registration_response(response: Response) -> OAuthClientInformationFull:
+    """Handle registration response."""
+    if response.status_code not in (200, 201):
+        await response.aread()
+        raise OAuthRegistrationError(f"Registration failed: {response.status_code} {response.text}")
+
+    try:
+        content = await response.aread()
+        client_info = OAuthClientInformationFull.model_validate_json(content)
+        return client_info
+        # self.context.client_info = client_info
+        # await self.context.storage.set_client_info(client_info)
+    except ValidationError as e:
+        raise OAuthRegistrationError(f"Invalid registration response: {e}")
+
+
+# async def prm_discovery(
+#     server_url: str,
+#     www_auth_resource_metadata_url: str | None,
+# ) -> ProtectedResourceMetadata | None:
+#     # Step 1: Discover protected resource metadata (SEP-985 with fallback support)
+#     prm_discovery_urls = build_protected_resource_discovery_urls(www_auth_resource_metadata_url, server_url)
+#     prm_discovery_success = False
+#     for url in prm_discovery_urls:
+#         discovery_request = create_oauth_metadata_request(url)
+#
+#         discovery_response = yield discovery_request  # sending request
+#
+#         prm = await handle_protected_resource_response(discovery_response)
+#         if prm:
+#             prm_discovery_success = True
+#
+#             # saving the response metadata
+#             self.context.protected_resource_metadata = prm
+#             if prm.authorization_servers:
+#                 self.context.auth_server_url = str(prm.authorization_servers[0])
+#
+#             break
+#         else:
+#             logger.debug(f"Protected resource metadata discovery failed: {url}")
+#     if not prm_discovery_success:
+#         raise OAuthFlowError("Protected resource metadata discovery failed: no valid metadata found")
+
+
+# class OAuthHandler:
+#     async def make_request(self, request: Request) -> Response: ...
+#
+#     async def store_prm(self, prm: ProtectedResourceMetadata) -> None: ...
 
 
 # async def discovery_process(discovery_urls: list[str]) -> AsyncGenerator[Request, Response]: