merge

Author: SoldierSacha

src/mcp/client/auth/oauth2.py

@@ -19,23 +19,20 @@
 import httpx
 from pydantic import BaseModel, Field, ValidationError
 
-from mcp.client.auth import OAuthFlowError, OAuthTokenError
+from mcp.client.auth import OAuthFlowError, OAuthRegistrationError, OAuthTokenError
 from mcp.client.auth.utils import (
     build_oauth_authorization_server_metadata_discovery_urls,
     build_protected_resource_metadata_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,
     handle_auth_metadata_response,
     handle_protected_resource_response,
-    handle_registration_response,
     handle_token_response_scopes,
 )
 from mcp.client.streamable_http import MCP_PROTOCOL_VERSION
-from mcp.types import LATEST_PROTOCOL_VERSION
 from mcp.shared.auth import (
     OAuthClientInformationFull,
     OAuthClientMetadata,
@@ -48,6 +45,7 @@
     check_resource_allowed,
     resource_url_from_server_url,
 )
+from mcp.types import LATEST_PROTOCOL_VERSION
 
 logger = logging.getLogger(__name__)
 
@@ -251,9 +249,7 @@ def _create_registration_request(self, metadata: OAuthMetadata | None = None) ->
             headers={"Content-Type": "application/json"},
         )
 
-    async def _handle_registration_response(
-        self, response: httpx.Response
-    ) -> OAuthClientInformationFull:
+    async def _handle_registration_response(self, response: httpx.Response) -> OAuthClientInformationFull:
         if response.status_code not in (200, 201):
             await response.aread()
             raise OAuthRegistrationError(f"Registration failed: {response.status_code} {response.text}")
@@ -323,15 +319,11 @@ def __init__(
 
     def _build_protected_resource_discovery_urls(self, resource_metadata_url: str | None) -> list[str]:
         """Build the list of PRM discovery URLs with legacy fallbacks."""
-        return build_protected_resource_metadata_discovery_urls(
-            resource_metadata_url, self.context.server_url
-        )
+        return build_protected_resource_metadata_discovery_urls(resource_metadata_url, self.context.server_url)
 
     def _get_discovery_urls(self, server_url: str | None = None) -> list[str]:
         """Build OAuth authorization server discovery URLs with legacy fallbacks."""
-        return build_oauth_authorization_server_metadata_discovery_urls(
-            server_url, self.context.server_url
-        )
+        return build_oauth_authorization_server_metadata_discovery_urls(server_url, self.context.server_url)
 
     async def _handle_protected_resource_response(self, response: httpx.Response) -> bool:
         """
@@ -356,9 +348,7 @@ async def _handle_protected_resource_response(self, response: httpx.Response) ->
         )
         return False
 
-    async def _handle_oauth_metadata_response(
-        self, response: httpx.Response
-    ) -> tuple[bool, OAuthMetadata | None]:
+    async def _handle_oauth_metadata_response(self, response: httpx.Response) -> tuple[bool, OAuthMetadata | None]:
         ok, asm = await handle_auth_metadata_response(response)
         if asm:
             self.context.oauth_metadata = asm
@@ -580,9 +570,7 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                     self._metadata = None
 
                     # Step 1: Discover protected resource metadata (SEP-985 with fallback support)
-                    prm_discovery_urls = self._build_protected_resource_discovery_urls(
-                        www_auth_resource_metadata_url
-                    )
+                    prm_discovery_urls = self._build_protected_resource_discovery_urls(www_auth_resource_metadata_url)
 
                     for url in prm_discovery_urls:  # pragma: no branch
                         discovery_request = self._create_oauth_metadata_request(url)

tests/client/test_auth.py

@@ -1364,9 +1364,7 @@ async def callback_handler() -> tuple[str, str | None]:
         )
 
         # Mock authorization
-        provider._perform_authorization_code_grant = AsyncMock(
-            return_value=("test_auth_code", "test_code_verifier")
-        )
+        provider._perform_authorization_code_grant = AsyncMock(return_value=("test_auth_code", "test_code_verifier"))
 
         # Next should be token exchange
         token_request = await auth_flow.asend(oauth_metadata_response)
@@ -1470,9 +1468,7 @@ async def callback_handler() -> tuple[str, str | None]:
             request=oauth_metadata_request,
         )
 
-        provider._perform_authorization_code_grant = AsyncMock(
-            return_value=("test_auth_code", "test_code_verifier")
-        )
+        provider._perform_authorization_code_grant = AsyncMock(return_value=("test_auth_code", "test_code_verifier"))
 
         token_request = await auth_flow.asend(oauth_metadata_response)
         assert str(token_request.url) == "https://api.example.com/token"