fix(client/auth): run discovery before the 403 step-up re-registers when no AS metadata is cached

After a restart _initialize() restores tokens and client info but never AS
metadata, so a 403 insufficient_scope step-up that needs to re-register (lapsed
secret, or no stored record) previously registered blind: it bypassed a
configured CIMD URL (should_use_client_metadata_url is False without metadata)
and POSTed the registration document to the resource origin's /register
fallback — the wrong server in the standard separate-AS topology, failing every
step-up until the access token expired.

The 401 flow's discovery sequence (PRM + SEP-2352 issuer checks + ASM) now
lives in a _discover_authorization_server_metadata sub-generator shared by both
branches; the 403 step-up runs it before the expiry discard whenever
oauth_metadata is None and re-registration is coming, so registration targets
the discovered endpoint (or resolves CIMD). extract_resource_metadata_from_
www_auth now honors the resource_metadata parameter on 403 challenges too
(RFC 9728 attaches it to any WWW-Authenticate challenge), so the step-up's
discovery is seeded from the challenge header exactly like the 401 flow's.

Covered by two restart-shaped step-up tests: DCR against the discovered
registration endpoint, and CIMD resolution once discovery restores the
advertised capability.

Author: claude

src/mcp/client/auth/oauth2.py

@@ -678,6 +678,101 @@ async def _complete_client_registration(self, response: httpx2.Response) -> None
         self.context.client_info = client_information
         await self.context.storage.set_client_info(client_information)
 
+    async def _discover_authorization_server_metadata(
+        self, response: httpx2.Response
+    ) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
+        """Discover protected resource and authorization server metadata.
+
+        Runs the discovery sequence shared by the 401 flow and the 403 step-up:
+        protected resource metadata (SEP-985, seeded from the challenge's
+        `resource_metadata` parameter when present), the SEP-2352 issuer checks on
+        stored credentials, and OAuth authorization server metadata. Yields each
+        discovery request; the caller must send the response back into the generator
+        (the httpx auth-flow protocol).
+        """
+        www_auth_resource_metadata_url = extract_resource_metadata_from_www_auth(response)
+
+        # Step 1: Discover protected resource metadata (SEP-985 with fallback support)
+        prm_discovery_urls = build_protected_resource_metadata_discovery_urls(
+            www_auth_resource_metadata_url, self.context.server_url
+        )
+
+        for url in prm_discovery_urls:  # pragma: no branch
+            discovery_request = create_oauth_metadata_request(url)
+
+            discovery_response = yield discovery_request  # sending request
+
+            prm = await handle_protected_resource_response(discovery_response)
+            if prm:
+                # Validate PRM resource matches server URL (RFC 8707)
+                await self._validate_resource_match(prm)
+                self.context.protected_resource_metadata = prm
+
+                # todo: try all authorization_servers to find the OASM
+                assert (
+                    len(prm.authorization_servers) > 0
+                )  # this is always true as authorization_servers has a min length of 1
+
+                self.context.auth_server_url = str(prm.authorization_servers[0])
+                break
+            else:
+                logger.debug(f"Protected resource metadata discovery failed: {url}")
+
+        # SEP-2352: stored credentials are bound to the issuer that registered them.
+        # If the authorization server changed, drop them (and the old tokens) so the
+        # flow re-registers instead of presenting another server's credentials.
+        if (
+            self.context.client_info is not None
+            and self.context.auth_server_url is not None
+            and not credentials_match_issuer(
+                self.context.client_info, self.context.auth_server_url, self.context.client_metadata_url
+            )
+        ):
+            logger.debug("Authorization server changed; discarding bound credentials and re-registering")
+            self.context.client_info = None
+            self.context.clear_tokens()
+            # Any cached AS metadata is for the old server; drop it so a failed
+            # rediscovery cannot leak the old registration/token endpoints into Step 4.
+            self.context.oauth_metadata = None
+
+        asm_discovery_urls = build_oauth_authorization_server_metadata_discovery_urls(
+            self.context.auth_server_url, self.context.server_url
+        )
+
+        # Step 2: Discover OAuth Authorization Server Metadata (OASM) (with fallback for legacy servers)
+        for url in asm_discovery_urls:  # pragma: no branch
+            oauth_metadata_request = create_oauth_metadata_request(url)
+            oauth_metadata_response = yield oauth_metadata_request
+
+            ok, asm = await handle_auth_metadata_response(oauth_metadata_response)
+            if not ok:
+                break
+            if ok and asm:
+                # SEP-2468: metadata issuer must match the discovery issuer
+                if self.context.auth_server_url is not None:
+                    validate_metadata_issuer(asm, self.context.auth_server_url)
+                self.context.oauth_metadata = asm
+                break
+            else:
+                logger.debug(f"OAuth metadata discovery failed: {url}")
+
+        # SEP-2352: on the legacy no-PRM path the issuer is only known after ASM
+        # discovery, so re-evaluate the binding here using the discovered metadata
+        # issuer (mirroring the bound_issuer fallback in Step 4).
+        if (
+            self.context.client_info is not None
+            and self.context.auth_server_url is None
+            and self.context.oauth_metadata is not None
+            and not credentials_match_issuer(
+                self.context.client_info,
+                str(self.context.oauth_metadata.issuer),
+                self.context.client_metadata_url,
+            )
+        ):
+            logger.debug("Authorization server changed; discarding bound credentials and re-registering")
+            self.context.client_info = None
+            self.context.clear_tokens()
+
     async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
         """httpx2 auth flow integration."""
         async with self.context.lock:
@@ -713,88 +808,19 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
                 try:
                     # OAuth flow must be inline due to generator constraints
 
-                    www_auth_resource_metadata_url = extract_resource_metadata_from_www_auth(response)
-
-                    # Step 1: Discover protected resource metadata (SEP-985 with fallback support)
-                    prm_discovery_urls = build_protected_resource_metadata_discovery_urls(
-                        www_auth_resource_metadata_url, self.context.server_url
-                    )
-
-                    for url in prm_discovery_urls:  # pragma: no branch
-                        discovery_request = create_oauth_metadata_request(url)
-
-                        discovery_response = yield discovery_request  # sending request
-
-                        prm = await handle_protected_resource_response(discovery_response)
-                        if prm:
-                            # Validate PRM resource matches server URL (RFC 8707)
-                            await self._validate_resource_match(prm)
-                            self.context.protected_resource_metadata = prm
-
-                            # todo: try all authorization_servers to find the OASM
-                            assert (
-                                len(prm.authorization_servers) > 0
-                            )  # this is always true as authorization_servers has a min length of 1
-
-                            self.context.auth_server_url = str(prm.authorization_servers[0])
-                            break
-                        else:
-                            logger.debug(f"Protected resource metadata discovery failed: {url}")
-
-                    # SEP-2352: stored credentials are bound to the issuer that registered them.
-                    # If the authorization server changed, drop them (and the old tokens) so the
-                    # flow re-registers instead of presenting another server's credentials.
-                    if (
-                        self.context.client_info is not None
-                        and self.context.auth_server_url is not None
-                        and not credentials_match_issuer(
-                            self.context.client_info, self.context.auth_server_url, self.context.client_metadata_url
-                        )
-                    ):
-                        logger.debug("Authorization server changed; discarding bound credentials and re-registering")
-                        self.context.client_info = None
-                        self.context.clear_tokens()
-                        # Any cached AS metadata is for the old server; drop it so a failed
-                        # rediscovery cannot leak the old registration/token endpoints into Step 4.
-                        self.context.oauth_metadata = None
-
-                    asm_discovery_urls = build_oauth_authorization_server_metadata_discovery_urls(
-                        self.context.auth_server_url, self.context.server_url
-                    )
-
-                    # Step 2: Discover OAuth Authorization Server Metadata (OASM) (with fallback for legacy servers)
-                    for url in asm_discovery_urls:  # pragma: no branch
-                        oauth_metadata_request = create_oauth_metadata_request(url)
-                        oauth_metadata_response = yield oauth_metadata_request
-
-                        ok, asm = await handle_auth_metadata_response(oauth_metadata_response)
-                        if not ok:
-                            break
-                        if ok and asm:
-                            # SEP-2468: metadata issuer must match the discovery issuer
-                            if self.context.auth_server_url is not None:
-                                validate_metadata_issuer(asm, self.context.auth_server_url)
-                            self.context.oauth_metadata = asm
-                            break
-                        else:
-                            logger.debug(f"OAuth metadata discovery failed: {url}")
-
-                    # SEP-2352: on the legacy no-PRM path the issuer is only known after ASM
-                    # discovery, so re-evaluate the binding here using the discovered metadata
-                    # issuer (mirroring the bound_issuer fallback in Step 4).
-                    if (
-                        self.context.client_info is not None
-                        and self.context.auth_server_url is None
-                        and self.context.oauth_metadata is not None
-                        and not credentials_match_issuer(
-                            self.context.client_info,
-                            str(self.context.oauth_metadata.issuer),
-                            self.context.client_metadata_url,
-                        )
-                    ):
-                        logger.debug("Authorization server changed; discarding bound credentials and re-registering")
-                        self.context.client_info = None
-                        self.context.clear_tokens()
+                    # Steps 1-2: Discover protected resource and authorization server
+                    # metadata, applying the SEP-2352 issuer checks along the way. The
+                    # sequence lives in a sub-generator shared with the 403 step-up;
+                    # its requests are relayed by hand (`yield from` cannot cross an
+                    # async generator).
+                    discovery = self._discover_authorization_server_metadata(response)
+                    try:
+                        discovery_request = await anext(discovery)
+                        while True:
+                            discovery_response = yield discovery_request
+                            discovery_request = await discovery.asend(discovery_response)
+                    except StopAsyncIteration:
+                        pass
 
                     # Step 3: Apply scope selection strategy
                     self.context.client_metadata.scope = get_client_metadata_scopes(
@@ -844,6 +870,29 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
                 # Step 2: Check if we need to step-up authorization
                 if error == "insufficient_scope":  # pragma: no branch
                     try:
+                        # After a restart the 403 step-up can be the first auth event:
+                        # `_initialize` restores tokens and client info but never AS
+                        # metadata, and the still-live access token keeps the 401 flow
+                        # (and its discovery) from running. When this step-up will need
+                        # to re-register — the stored secret lapsed, or no record is
+                        # stored at all — discover the metadata first instead of
+                        # registering blind: the blind fallback would bypass a
+                        # configured CIMD URL and POST the registration document to the
+                        # resource origin's `/register`, the wrong server in the
+                        # separate-AS topology. Running it before the expiry discard
+                        # lets the SEP-2352 issuer checks still see the record's stamp.
+                        if self.context.oauth_metadata is None and (
+                            self.context.registration_secret_expired() or self.context.client_info is None
+                        ):
+                            discovery = self._discover_authorization_server_metadata(response)
+                            try:
+                                discovery_request = await anext(discovery)
+                                while True:
+                                    discovery_response = yield discovery_request
+                                    discovery_request = await discovery.asend(discovery_response)
+                            except StopAsyncIteration:
+                                pass
+
                         # Step 2a: Union previously requested scopes with the newly challenged
                         # scopes (SEP-2350) so escalating one operation keeps the others' grants.
                         # Fold in the stored token's scope too: on a restart the token is reloaded

src/mcp/client/auth/utils.py

@@ -51,10 +51,13 @@ def extract_scope_from_www_auth(response: Response) -> str | None:
 def extract_resource_metadata_from_www_auth(response: Response) -> str | None:
     """Extract protected resource metadata URL from WWW-Authenticate header as per RFC 9728.
 
+    RFC 9728 attaches the `resource_metadata` parameter to any WWW-Authenticate
+    challenge, so both 401 and 403 (scope step-up) responses are honored.
+
     Returns:
         Resource metadata URL if found in WWW-Authenticate header, None otherwise
     """
-    if not response or response.status_code != 401:
+    if not response or response.status_code not in (401, 403):
         return None  # pragma: no cover
 
     return extract_field_from_www_auth(response, "resource_metadata")