Discover auth-server metadata before eager token refresh

src/mcp/client/auth/oauth2.py

@@ -577,6 +577,64 @@ 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 _discover_oauth_metadata(self) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
+        """Discover authorization server metadata and populate the context.
+
+        Yields the discovery requests so they run through the outer httpx auth flow
+        (no side-channel client). This is pure discovery: it fills in
+        ``protected_resource_metadata`` / ``auth_server_url`` / ``oauth_metadata`` and
+        does not register clients or mutate stored credentials. Used to populate the
+        token endpoint before an eager refresh, and available for the 401 path.
+        """
+        # Protected resource metadata -> authorization server URL. Best-effort: legacy
+        # servers without PRM fall through to the origin well-known in the ASM step.
+        if self.context.auth_server_url is None:
+            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:
+                    await self._validate_resource_match(prm)
+                    self.context.protected_resource_metadata = prm
+                    self.context.auth_server_url = str(prm.authorization_servers[0])
+                    break
+
+        # Authorization server metadata -> token / authorization / registration endpoints.
+        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:
+                    validate_metadata_issuer(asm, self.context.auth_server_url)
+                self.context.oauth_metadata = asm
+                break
+
+    async def _refresh_with_discovery(self) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
+        """Eager token refresh that discovers authorization-server metadata first when
+        it is not yet known.
+
+        The token endpoint comes from the AS metadata. On a cold start (e.g. reusing a
+        stored refresh token before any 401) that metadata has not been discovered, so
+        ``_refresh_token`` would fall back to ``{origin}/token`` — dropping any issuer
+        path and 404ing on servers whose token endpoint lives under a path. Yields the
+        discovery and refresh requests so they run through the outer httpx auth flow.
+        """
+        if self.context.oauth_metadata is None:
+            discovery = self._discover_oauth_metadata()
+            discovery_request = await discovery.asend(None)
+            while True:
+                discovery_response = yield discovery_request
+                try:
+                    discovery_request = await discovery.asend(discovery_response)
+                except StopAsyncIteration:
+                    break
+
+        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 +645,17 @@ 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 when
+                # it is not yet known (see _refresh_with_discovery). Driven here 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.asend(None)
+                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,60 @@ 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, valid_tokens: OAuthToken
+):
+    """Regression: on a cold start (cached expired token, no prior discovery) the
+    eager refresh must discover authorization-server metadata first, so it targets
+    the real token endpoint instead of the ``{origin}/token`` fallback. That
+    fallback drops any issuer path and 404s on servers whose token endpoint lives
+    under a path, which silently clears tokens and forces interactive re-auth.
+    """
+    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
+    prm_request = await auth_flow.__anext__()
+    assert "oauth-protected-resource" in str(prm_request.url)
+    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 "oauth-authorization-server" in str(asm_request.url)
+    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 must target 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 str(refresh_request.url) != "https://api.example.com/token"
+    assert "grant_type=refresh_token" in refresh_request.content.decode()
+
+    await auth_flow.aclose()