test step-up auth flow

Author: dogacancolak

src/mcp/client/auth.py

@@ -284,11 +284,12 @@ def _configure_scope_selection(self, init_response: httpx.Response) -> None:
         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:
+        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
-            )
+            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
@@ -592,12 +593,12 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                 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 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 2: Check if we need to step-up authorization
+                if error == "insufficient_scope":
+                    try:
                         # Step 2a: Update the required scopes
                         self._configure_scope_selection(response)
 
@@ -608,10 +609,10 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                         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
+                    except Exception:
+                        logger.exception("OAuth flow error")
+                        raise
 
-                # Retry with new tokens
-                self._add_auth_header(request)
-                yield request
+                    # Retry with new tokens
+                    self._add_auth_header(request)
+                    yield request

tests/client/test_auth.py

@@ -858,6 +858,98 @@ async def test_auth_flow_no_unnecessary_retry_after_oauth(
         # Verify exactly one request was yielded (no double-sending)
         assert request_yields == 1, f"Expected 1 request yield, got {request_yields}"
 
+    @pytest.mark.anyio
+    async def test_403_insufficient_scope_updates_scope_from_header(
+        self,
+        oauth_provider: OAuthClientProvider,
+        mock_storage: MockTokenStorage,
+        valid_tokens: OAuthToken,
+    ):
+        """Test that 403 response correctly updates scope from WWW-Authenticate header."""
+        # Pre-store valid tokens and client info
+        client_info = OAuthClientInformationFull(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            redirect_uris=[AnyUrl("http://localhost:3030/callback")],
+        )
+        await mock_storage.set_tokens(valid_tokens)
+        await mock_storage.set_client_info(client_info)
+        oauth_provider.context.current_tokens = valid_tokens
+        oauth_provider.context.token_expiry_time = time.time() + 1800
+        oauth_provider.context.client_info = client_info
+        oauth_provider._initialized = True
+
+        # Original scope
+        assert oauth_provider.context.client_metadata.scope == "read write"
+
+        redirect_captured = False
+        captured_state = None
+
+        async def capture_redirect(url: str) -> None:
+            nonlocal redirect_captured, captured_state
+            redirect_captured = True
+            # Verify the new scope is included in authorization URL
+            assert "scope=admin%3Awrite+admin%3Adelete" in url or "scope=admin:write+admin:delete" in url.replace(
+                "%3A", ":"
+            ).replace("+", " ")
+            # Extract state from redirect URL
+            from urllib.parse import parse_qs, urlparse
+
+            parsed = urlparse(url)
+            params = parse_qs(parsed.query)
+            captured_state = params.get("state", [None])[0]
+
+        oauth_provider.context.redirect_handler = capture_redirect
+
+        # Mock callback
+        async def mock_callback() -> tuple[str, str | None]:
+            return "auth_code", captured_state
+
+        oauth_provider.context.callback_handler = mock_callback
+
+        test_request = httpx.Request("GET", "https://api.example.com/mcp")
+        auth_flow = oauth_provider.async_auth_flow(test_request)
+
+        # First request
+        request = await auth_flow.__anext__()
+
+        # Send 403 with new scope requirement
+        response_403 = httpx.Response(
+            403,
+            headers={"WWW-Authenticate": 'Bearer error="insufficient_scope", scope="admin:write admin:delete"'},
+            request=request,
+        )
+
+        # Trigger step-up - should get token exchange request
+        token_exchange_request = await auth_flow.asend(response_403)
+
+        # Verify scope was updated
+        assert oauth_provider.context.client_metadata.scope == "admin:write admin:delete"
+        assert redirect_captured
+
+        # Complete the flow with successful token response
+        token_response = httpx.Response(
+            200,
+            json={
+                "access_token": "new_token_with_new_scope",
+                "token_type": "Bearer",
+                "expires_in": 3600,
+                "scope": "admin:write admin:delete",
+            },
+            request=token_exchange_request,
+        )
+
+        # Should get final retry request
+        final_request = await auth_flow.asend(token_response)
+
+        # Send success response - flow should complete
+        success_response = httpx.Response(200, request=final_request)
+        try:
+            await auth_flow.asend(success_response)
+            pytest.fail("Should have stopped after successful response")
+        except StopAsyncIteration:
+            pass  # Expected
+
 
 @pytest.mark.parametrize(
     (