fix(client/auth): discover AS metadata before cold-start token refresh, with issuer-binding check

src/mcp/client/auth/oauth2.py

@@ -10,6 +10,7 @@
 import string
 import time
 from collections.abc import AsyncGenerator, Awaitable, Callable
+from contextlib import aclosing
 from dataclasses import dataclass, field
 from typing import Any, Protocol, get_args
 from urllib.parse import quote, urlencode, urljoin, urlparse
@@ -160,6 +161,9 @@ class OAuthContext:
     oauth_metadata: OAuthMetadata | None = None
     auth_server_url: str | None = None
     protocol_version: str | None = None
+    # Whether the eager (pre-401) refresh already ran its blind discovery probes, so a
+    # server that publishes no metadata is not re-probed on every in-process refresh.
+    eager_discovery_attempted: bool = False
 
     # Client registration
     client_info: OAuthClientInformationFull | None = None
@@ -577,6 +581,126 @@ async def _validate_resource_match(self, prm: ProtectedResourceMetadata) -> None
         if not check_resource_allowed(requested_resource=default_resource, configured_resource=prm_resource):
             raise OAuthFlowError(f"Protected resource {prm_resource} does not match expected {default_resource}")
 
+    async def _refresh_with_discovery(self) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
+        """Refresh the token, discovering authorization-server metadata first when needed.
+
+        The token endpoint comes from the AS metadata. On a cold start (a stored refresh
+        token reused before any 401) that metadata has not been discovered yet, so
+        ``_refresh_token`` would fall back to ``{origin}/token`` — dropping any issuer
+        path and 404ing on servers whose token endpoint lives elsewhere. Discovery runs
+        first so the refresh targets the discovered token endpoint.
+
+        Unlike the 401 path, this discovery is unanchored: there is no WWW-Authenticate
+        ``resource_metadata`` hint, only blind well-known probes, so a co-hosted origin
+        can legitimately serve some *other* resource's documents. Results are therefore
+        treated as best-effort, never authoritative: a resource-mismatched PRM or an
+        issuer-mismatched ASM counts as a failed discovery rather than an error, and a
+        SEP-2352 issuer-binding mismatch
+        skips the eager refresh (so stored credentials are never presented to an
+        unvalidated authorization server) while leaving the credentials themselves for
+        the anchored 401 path to judge — that path re-discovers with the server's hint
+        and drops/re-registers only on a confirmed change. Servers publishing no
+        metadata at all keep the pre-existing ``{origin}/token`` fallback, and the
+        probes run only once per context (``eager_discovery_attempted``). Yields the
+        discovery and refresh requests so they run through the outer httpx auth flow
+        rather than a side-channel client.
+        """
+        if self.context.oauth_metadata is None and not self.context.eager_discovery_attempted:
+            # Step 1: protected resource metadata -> authorization server URL (SEP-985).
+            # Best-effort: a PRM that fails resource validation is some other co-hosted
+            # resource's document, not ours — skip it; a legacy server without PRM falls
+            # through to the origin well-known fallback in the ASM step below.
+            for url in build_protected_resource_metadata_discovery_urls(None, self.context.server_url):
+                prm = await handle_protected_resource_response((yield create_oauth_metadata_request(url)))
+                if prm:
+                    try:
+                        # Validate PRM resource matches server URL (RFC 8707)
+                        await self._validate_resource_match(prm)
+                    except OAuthFlowError:
+                        logger.debug(f"Ignoring protected resource metadata for a different resource: {url}")
+                        continue
+                    self.context.protected_resource_metadata = prm
+                    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.
+            # A mismatch here may mean the AS changed — or merely that the blind probe
+            # found a different co-hosted resource's PRM. Skip the eager refresh so the
+            # credentials are never presented to an unvalidated AS, discard the
+            # unanchored discovery results, and let the 401 path decide with the
+            # server's own hint whether to drop the credentials and re-register.
+            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(
+                    "Eagerly discovered authorization server does not match stored credential binding; "
+                    "skipping refresh and deferring to 401 discovery"
+                )
+                self.context.protected_resource_metadata = None
+                self.context.auth_server_url = None
+                self.context.eager_discovery_attempted = True
+                return
+
+            # Step 2: authorization server metadata -> the token endpoint (with fallback
+            # for legacy servers).
+            for url in build_oauth_authorization_server_metadata_discovery_urls(
+                self.context.auth_server_url, self.context.server_url
+            ):
+                ok, asm = await handle_auth_metadata_response((yield create_oauth_metadata_request(url)))
+                if not ok:
+                    break
+                if asm:
+                    if self.context.auth_server_url is not None:
+                        try:
+                            # SEP-2468: metadata issuer must match the discovery issuer
+                            validate_metadata_issuer(asm, self.context.auth_server_url)
+                        except OAuthFlowError:
+                            logger.debug(f"Ignoring authorization server metadata with mismatched issuer: {url}")
+                            continue
+                    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 (mirroring the 401 path's
+            # post-ASM check). As above, skip the refresh and discard the unanchored
+            # metadata rather than acting on it — keeping it would let a rejected
+            # issuer's endpoints leak into a later 401 flow's registration step.
+            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(
+                    "Eagerly discovered authorization server does not match stored credential binding; "
+                    "skipping refresh and deferring to 401 discovery"
+                )
+                self.context.oauth_metadata = None
+                self.context.eager_discovery_attempted = True
+                return
+
+            # Mark completion only now: an interrupted probe sequence (the transport
+            # failing mid-discovery closes this generator) is retried on the next
+            # refresh instead of being recorded as done.
+            self.context.eager_discovery_attempted = True
+
+        refresh_response = yield await self._refresh_token()
+        if not await self._handle_refresh_response(refresh_response):
+            # Refresh failed, need full re-authentication
+            self._initialized = False
+
     async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
         """httpx2 auth flow integration."""
         async with self.context.lock:
@@ -587,13 +711,18 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
             self.context.protocol_version = request.headers.get(MCP_PROTOCOL_VERSION_HEADER)
 
             if not self.context.is_token_valid() and self.context.can_refresh_token():
-                # Try to refresh token
-                refresh_request = await self._refresh_token()
-                refresh_response = yield refresh_request
-
-                if not await self._handle_refresh_response(refresh_response):
-                    # Refresh failed, need full re-authentication
-                    self._initialized = False
+                # Refresh the token, discovering authorization-server metadata first on a
+                # cold start (see _refresh_with_discovery). Driven inline so its requests
+                # run through this httpx auth flow, not a side-channel client; aclosing
+                # finalizes the inner generator when httpx closes this flow mid-refresh.
+                async with aclosing(self._refresh_with_discovery()) as refresh_flow:
+                    refresh_request = await refresh_flow.__anext__()
+                    while True:
+                        refresh_response = yield refresh_request
+                        try:
+                            refresh_request = await refresh_flow.asend(refresh_response)
+                        except StopAsyncIteration:
+                            break
 
             if self.context.is_token_valid():
                 self._add_auth_header(request)