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

src/mcp/client/auth/oauth2.py

@@ -577,6 +577,92 @@
         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, applying the same SEP-2352 issuer-binding checks as the 401 path so stored
+        credentials are never sent to an authorization server they are not bound to: on a
+        binding mismatch the credentials and tokens are dropped and the refresh is
+        skipped, letting the subsequent 401 flow re-register against the new server.
+        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:
+            # Step 1: protected resource metadata -> authorization server URL (SEP-985).
+            # Best-effort: a legacy server without PRM falls through to the origin
+            # well-known fallback in the ASM step below. There is no 401 response at
+            # this point, so no WWW-Authenticate resource_metadata hint is available.
+            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:
+                    # Validate PRM resource matches server URL (RFC 8707)
+                    await self._validate_resource_match(prm)
+                    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.
+            # If the authorization server changed, drop them (and the old tokens) and skip
+            # the refresh so the 401 flow re-registers instead of presenting another
+            # server's credentials to the newly discovered one.
+            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 skipping refresh")
+                self.context.client_info = None
+                self.context.clear_tokens()
+                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:
+                    # 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 401 path's post-ASM check).
+            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 skipping refresh")
+                self.context.client_info = None
+                self.context.clear_tokens()
+                return
+
+        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:
@@ -584,16 +670,20 @@
                 await self._initialize()
 
             # Capture protocol version from request headers
             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.
+                refresh_flow = self._refresh_with_discovery()
+                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)

tests/client/test_auth.py

@@ -3253,3 +3253,278 @@ async def echo_callback() -> AuthorizationCodeResult:
         await auth_flow.asend(httpx2.Response(200, request=final_req))
     except StopAsyncIteration:
         pass
+
+
+@pytest.mark.anyio
+async def test_eager_refresh_discovers_token_endpoint_before_refreshing(
+    oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, valid_tokens: OAuthToken
+):
+    """Regression for #3240/#3250: a cold-start eager refresh discovers the token endpoint.
+
+    On a restart with a stored (expired) token the pre-401 refresh used to POST to the
+    ``{origin}/token`` fallback because authorization-server metadata had not been
+    discovered yet, 404ing on servers whose token endpoint lives under a path and
+    silently clearing the stored tokens. The refresh must run PRM + ASM discovery first
+    and target the discovered token endpoint.
+    """
+    oauth_provider.context.current_tokens = valid_tokens
+    oauth_provider.context.token_expiry_time = time.time() - 100  # expired
+    oauth_provider.context.client_info = OAuthClientInformationFull(
+        client_id="test_client",
+        redirect_uris=[AnyUrl("http://localhost:3030/callback")],
+        token_endpoint_auth_method="none",
+    )
+    oauth_provider._initialized = True
+    assert oauth_provider.context.oauth_metadata is None
+
+    test_request = httpx2.Request("GET", "https://api.example.com/v1/mcp")
+    auth_flow = oauth_provider.async_auth_flow(test_request)
+
+    # 1) protected-resource metadata discovery (no WWW-Authenticate hint pre-401)
+    prm_request = await auth_flow.__anext__()
+    assert str(prm_request.url) == "https://api.example.com/.well-known/oauth-protected-resource/v1/mcp"
+    prm_response = httpx2.Response(
+        200,
+        content=(
+            b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}'
+        ),
+        request=prm_request,
+    )
+
+    # 2) authorization-server metadata whose token endpoint is NOT {origin}/token
+    asm_request = await auth_flow.asend(prm_response)
+    assert str(asm_request.url) == "https://auth.example.com/.well-known/oauth-authorization-server"
+    asm_response = httpx2.Response(
+        200,
+        content=(
+            b'{"issuer": "https://auth.example.com", '
+            b'"authorization_endpoint": "https://auth.example.com/oauth2/authorize", '
+            b'"token_endpoint": "https://auth.example.com/oauth2/api/v1/token"}'
+        ),
+        request=asm_request,
+    )
+
+    # 3) the refresh targets the discovered token endpoint, not the fallback
+    refresh_request = await auth_flow.asend(asm_response)
+    assert refresh_request.method == "POST"
+    assert str(refresh_request.url) == "https://auth.example.com/oauth2/api/v1/token"
+    assert "grant_type=refresh_token" in refresh_request.content.decode()
+    refresh_response = httpx2.Response(
+        200,
+        json={"access_token": "refreshed_token", "token_type": "Bearer", "expires_in": 3600},
+        request=refresh_request,
+    )
+
+    # 4) the original request goes out with the refreshed token
+    api_request = await auth_flow.asend(refresh_response)
+    assert str(api_request.url) == "https://api.example.com/v1/mcp"
+    assert api_request.headers["Authorization"] == "Bearer refreshed_token"
+    stored = await mock_storage.get_tokens()
+    assert stored is not None
+    assert stored.access_token == "refreshed_token"
+
+    with pytest.raises(StopAsyncIteration):
+        await auth_flow.asend(httpx2.Response(200, request=api_request))
+
+
+@pytest.mark.anyio
+async def test_eager_refresh_falls_back_to_origin_token_when_no_metadata_published(
+    oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
+):
+    """A legacy server publishing no metadata keeps the pre-existing ``{origin}/token`` fallback.
+
+    PRM discovery 404s at both well-known URLs and the legacy origin ASM fallback 404s too,
+    so the refresh still POSTs to ``{origin}/token`` exactly as before discovery-before-refresh
+    existed. A failed refresh then clears tokens and lets the request go out unauthenticated.
+    """
+    oauth_provider.context.current_tokens = valid_tokens
+    oauth_provider.context.token_expiry_time = time.time() - 100  # expired
+    oauth_provider.context.client_info = OAuthClientInformationFull(
+        client_id="test_client",
+        redirect_uris=[AnyUrl("http://localhost:3030/callback")],
+        token_endpoint_auth_method="none",
+    )
+    oauth_provider._initialized = True
+
+    auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
+
+    # PRM discovery: path-based then root-based, both 404.
+    prm_request = await auth_flow.__anext__()
+    assert str(prm_request.url) == "https://api.example.com/.well-known/oauth-protected-resource/v1/mcp"
+    prm_request = await auth_flow.asend(httpx2.Response(404, request=prm_request))
+    assert str(prm_request.url) == "https://api.example.com/.well-known/oauth-protected-resource"
+
+    # ASM discovery: legacy origin fallback, 404 as well.
+    asm_request = await auth_flow.asend(httpx2.Response(404, request=prm_request))
+    assert str(asm_request.url) == "https://api.example.com/.well-known/oauth-authorization-server"
+
+    # Refresh falls back to {origin}/token (pre-existing legacy behavior).
+    refresh_request = await auth_flow.asend(httpx2.Response(404, request=asm_request))
+    assert refresh_request.method == "POST"
+    assert str(refresh_request.url) == "https://api.example.com/token"
+
+    # The refresh fails; tokens are cleared and the original request goes out unauthenticated.
+    api_request = await auth_flow.asend(httpx2.Response(401, request=refresh_request))
+    assert str(api_request.url) == "https://api.example.com/v1/mcp"
+    assert "Authorization" not in api_request.headers
+    assert oauth_provider.context.current_tokens is None
+    await auth_flow.aclose()
+
+
+@pytest.mark.anyio
+async def test_eager_refresh_stops_asm_discovery_on_server_error(
+    oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
+):
+    """A non-4XX ASM discovery error stops the fallback chain, mirroring the 401 path.
+
+    The refresh then proceeds against the ``{origin}/token`` fallback rather than
+    hammering further well-known URLs.
+    """
+    oauth_provider.context.current_tokens = valid_tokens
+    oauth_provider.context.token_expiry_time = time.time() - 100  # expired
+    oauth_provider.context.client_info = OAuthClientInformationFull(
+        client_id="test_client",
+        redirect_uris=[AnyUrl("http://localhost:3030/callback")],
+        token_endpoint_auth_method="none",
+    )
+    oauth_provider._initialized = True
+
+    auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
+
+    # PRM discovery succeeds and points at the authorization server.
+    prm_request = await auth_flow.__anext__()
+    prm_response = httpx2.Response(
+        200,
+        content=(
+            b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}'
+        ),
+        request=prm_request,
+    )
+
+    # ASM discovery hits a 500: stop trying further URLs.
+    asm_request = await auth_flow.asend(prm_response)
+    assert str(asm_request.url) == "https://auth.example.com/.well-known/oauth-authorization-server"
+    refresh_request = await auth_flow.asend(httpx2.Response(500, request=asm_request))
+
+    assert refresh_request.method == "POST"
+    assert str(refresh_request.url) == "https://api.example.com/token"
+    await auth_flow.aclose()
+
+
+@pytest.mark.anyio
+async def test_eager_refresh_skips_refresh_when_credentials_bound_to_different_issuer(
+    oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
+):
+    """SEP-2352: a cold-start refresh never sends credentials bound to another issuer.
+
+    When PRM discovery reveals an authorization server different from the one the stored
+    client credentials are bound to, the credentials and tokens are dropped and the
+    refresh is skipped, so the subsequent 401 flow re-registers against the new server
+    — mirroring the issuer-binding check on the 401 discovery path.
+    """
+    oauth_provider.context.current_tokens = valid_tokens
+    oauth_provider.context.token_expiry_time = time.time() - 100  # expired
+    oauth_provider.context.client_info = OAuthClientInformationFull(
+        client_id="stale-client",
+        redirect_uris=[AnyUrl("http://localhost:3030/callback")],
+        issuer="https://old-as.example.com",
+    )
+    oauth_provider._initialized = True
+
+    auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
+
+    # PRM discovery points at auth.example.com, not the bound old-as.example.com.
+    prm_request = await auth_flow.__anext__()
+    prm_response = httpx2.Response(
+        200,
+        content=(
+            b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}'
+        ),
+        request=prm_request,
+    )
+
+    # No refresh request: the next yield is the original request, unauthenticated.
+    api_request = await auth_flow.asend(prm_response)
+    assert str(api_request.url) == "https://api.example.com/v1/mcp"
+    assert "Authorization" not in api_request.headers
+    assert oauth_provider.context.client_info is None
+    assert oauth_provider.context.current_tokens is None
+    await auth_flow.aclose()
+
+
+@pytest.mark.anyio
+async def test_eager_refresh_legacy_path_rechecks_issuer_binding_after_asm(
+    oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
+):
+    """SEP-2352 on the legacy no-PRM path: the binding is checked against the ASM issuer.
+
+    PRM discovery fails so the issuer is only known once origin-fallback ASM discovery
+    succeeds; credentials bound to a different issuer are then dropped and the refresh is
+    skipped, exactly as on the 401 path's post-ASM re-evaluation.
+    """
+    oauth_provider.context.current_tokens = valid_tokens
+    oauth_provider.context.token_expiry_time = time.time() - 100  # expired
+    oauth_provider.context.client_info = OAuthClientInformationFull(
+        client_id="stale-client",
+        redirect_uris=[AnyUrl("http://localhost:3030/callback")],
+        issuer="https://old-as.example.com",
+    )
+    oauth_provider._initialized = True
+
+    auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
+
+    # PRM discovery: both well-known URLs 404.
+    prm_request = await auth_flow.__anext__()
+    prm_request = await auth_flow.asend(httpx2.Response(404, request=prm_request))
+
+    # Origin-fallback ASM discovery succeeds with the resource origin as issuer.
+    asm_request = await auth_flow.asend(httpx2.Response(404, request=prm_request))
+    assert str(asm_request.url) == "https://api.example.com/.well-known/oauth-authorization-server"
+    asm_response = httpx2.Response(
+        200,
+        content=(
+            b'{"issuer": "https://api.example.com", '
+            b'"authorization_endpoint": "https://api.example.com/authorize", '
+            b'"token_endpoint": "https://api.example.com/token"}'
+        ),
+        request=asm_request,
+    )
+
+    # No refresh request: the next yield is the original request, unauthenticated.
+    api_request = await auth_flow.asend(asm_response)
+    assert str(api_request.url) == "https://api.example.com/v1/mcp"
+    assert "Authorization" not in api_request.headers
+    assert oauth_provider.context.client_info is None
+    assert oauth_provider.context.current_tokens is None
+    # The just-discovered metadata is for the current server and is kept for the 401 flow.
+    assert oauth_provider.context.oauth_metadata is not None
+    await auth_flow.aclose()
+
+
+@pytest.mark.anyio
+async def test_eager_refresh_skips_discovery_when_metadata_already_known(
+    oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
+):
+    """With authorization-server metadata already discovered, the refresh is immediate."""
+    oauth_provider.context.current_tokens = valid_tokens
+    oauth_provider.context.token_expiry_time = time.time() - 100  # expired
+    oauth_provider.context.client_info = OAuthClientInformationFull(
+        client_id="test_client",
+        redirect_uris=[AnyUrl("http://localhost:3030/callback")],
+        token_endpoint_auth_method="none",
+    )
+    oauth_provider.context.oauth_metadata = OAuthMetadata.model_validate(
+        {
+            "issuer": "https://auth.example.com",
+            "authorization_endpoint": "https://auth.example.com/oauth2/authorize",
+            "token_endpoint": "https://auth.example.com/oauth2/api/v1/token",
+        }
+    )
+    oauth_provider._initialized = True
+
+    auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
+
+    refresh_request = await auth_flow.__anext__()
+    assert refresh_request.method == "POST"
+    assert str(refresh_request.url) == "https://auth.example.com/oauth2/api/v1/token"
+    await auth_flow.aclose()