step-up auth flow and fix tests

Author: dogacancolak

src/mcp/client/auth.py

@@ -97,7 +97,6 @@ class OAuthContext:
     oauth_metadata: OAuthMetadata | None = None
     auth_server_url: str | None = None
     protocol_version: str | None = None
-    www_authenticate_scope: str | None = None
 
     # Client registration
     client_info: OAuthClientInformationFull | None = None
@@ -212,9 +211,6 @@ def _extract_field_from_www_auth(self, init_response: httpx.Response, field_name
         Returns:
             Field value if found in WWW-Authenticate header, None otherwise
         """
-        if not init_response or init_response.status_code != 401:
-            return None
-
         www_auth_header = init_response.headers.get("WWW-Authenticate")
         if not www_auth_header:
             return None
@@ -236,6 +232,9 @@ def _extract_resource_metadata_from_www_auth(self, init_response: httpx.Response
         Returns:
             Resource metadata URL if found in WWW-Authenticate header, None otherwise
         """
+        if not init_response or init_response.status_code != 401:
+            return None
+
         return self._extract_field_from_www_auth(init_response, "resource_metadata")
 
     def _extract_scope_from_www_auth(self, init_response: httpx.Response) -> str | None:
@@ -268,26 +267,31 @@ async def _handle_protected_resource_response(self, response: httpx.Response) ->
                 if metadata.authorization_servers:
                     self.context.auth_server_url = str(metadata.authorization_servers[0])
 
-                # Per MCP spec, scope selection priority order:
-                # 1. Keep client scope if configured
-                # 2. Use scope from WWW-Authenticate header (if provided)
-                # 3. Use all scopes from PRM scopes_supported (if available)
-                # 4. Omit scope parameter if neither is available
-                #
-                # Priority 1: Don't touch if client scope is already configured
-                if self.context.client_metadata.scope is None:
-                    if self.context.www_authenticate_scope is not None:
-                        # Priority 2: WWW-Authenticate header scope
-                        self.context.client_metadata.scope = self.context.www_authenticate_scope
-                    elif self.context.protected_resource_metadata.scopes_supported is not None:
-                        # Priority 3: PRM scopes_supported
-                        self.context.client_metadata.scope = " ".join(
-                            self.context.protected_resource_metadata.scopes_supported
-                        )
-                    # Priority 4: Omit scope parameter
-
             except ValidationError:
                 pass
+        else:
+            raise OAuthFlowError(f"Protected Resource Metadata request failed: {response.status_code}")
+
+    def _configure_scope_selection(self, init_response: httpx.Response) -> None:
+        """Select scopes as outlined in the 'Scope Selection Strategy in the MCP spec."""
+        # Per MCP spec, scope selection priority order:
+        # 1. Use scope from WWW-Authenticate header (if provided)
+        # 2. Use all scopes from PRM scopes_supported (if available)
+        # 3. Omit scope parameter if neither is available
+        #
+        # Step 1: Extract scope from WWW-Authenticate header
+        www_authenticate_scope = self._extract_scope_from_www_auth(init_response)
+        if www_authenticate_scope is not None:
+            # Priority 1: WWW-Authenticate header scope
+            self.context.client_metadata.scope = www_authenticate_scope
+        elif self.context.protected_resource_metadata is not None and self.context.protected_resource_metadata.scopes_supported is not None:
+            # Priority 2: PRM scopes_supported
+            self.context.client_metadata.scope = " ".join(
+                self.context.protected_resource_metadata.scopes_supported
+            )
+        else:
+            # Priority 3: Omit scope parameter
+            self.context.client_metadata.scope = None
 
     def _get_discovery_urls(self) -> list[str]:
         """Generate ordered list of (url, type) tuples for discovery attempts."""
@@ -544,14 +548,14 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                 # Perform full OAuth flow
                 try:
                     # OAuth flow must be inline due to generator constraints
-                    # Step 1: Extract scope from WWW-Authenticate header
-                    self.context.www_authenticate_scope = self._extract_scope_from_www_auth(response)
-
-                    # Step 2: Discover protected resource metadata (RFC9728 with WWW-Authenticate support)
+                    # 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)
 
+                    # Step 2: Apply scope selection strategy
+                    self._configure_scope_selection(response)
+
                     # Step 3: Discover OAuth metadata (with fallback for legacy servers)
                     discovery_urls = self._get_discovery_urls()
                     for url in discovery_urls:
@@ -587,3 +591,27 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                 # Retry with new tokens
                 self._add_auth_header(request)
                 yield request
+            elif response.status_code == 403:
+                try:
+                    # Step 1: Extract error field from WWW-Authenticate header
+                    error = self._extract_field_from_www_auth(response, "error")
+
+                    # Step 2: Check if we need to step-up authorization
+                    if error == "insufficient_scope":
+                        # Step 2a: Update the required scopes
+                        self._configure_scope_selection(response)
+
+                        # Step 2b: Perform (re-)authorization
+                        auth_code, code_verifier = await self._perform_authorization()
+
+                        # Step 2c: Exchange authorization code for tokens
+                        token_request = await self._exchange_token(auth_code, code_verifier)
+                        token_response = yield token_request
+                        await self._handle_token_response(token_response)
+                except Exception:
+                    logger.exception("OAuth flow error")
+                    raise
+
+                # Retry with new tokens
+                self._add_auth_header(request)
+                yield request