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)