From 418914e909ff41d0e2e22d289449cc1673c13609 Mon Sep 17 00:00:00 2001 From: mfadul24 Date: Sun, 2 Aug 2026 23:41:10 -0700 Subject: [PATCH] Discover auth-server metadata before eager token refresh MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On a cold start with a cached-but-expired token (reusing a stored refresh token before any 401), async_auth_flow refreshed before discovery ran, so oauth_metadata was None and _refresh_token fell back to {origin}/token. That drops any issuer path and 404s on servers whose token endpoint lives under a path (e.g. .../oauth2/api/v1/token) — the refresh fails, tokens are cleared, and the client is forced into interactive re-auth it may not be able to complete. Discover AS metadata before the eager refresh. Adds _discover_oauth_metadata (pure discovery, driven through the auth flow) and _refresh_with_discovery, plus a regression test. --- src/mcp/client/auth/oauth2.py | 76 +++++++++++++++++++++++++++++++---- tests/client/test_auth.py | 57 ++++++++++++++++++++++++++ 2 files changed, 126 insertions(+), 7 deletions(-) diff --git a/src/mcp/client/auth/oauth2.py b/src/mcp/client/auth/oauth2.py index 7dc62b52b9..e64f32c617 100644 --- a/src/mcp/client/auth/oauth2.py +++ b/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) diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index be96cc8eec..f800ccfca1 100644 --- a/tests/client/test_auth.py +++ b/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()