@@ -577,6 +577,64 @@ async def _validate_resource_match(self, prm: ProtectedResourceMetadata) -> None
577577 if not check_resource_allowed (requested_resource = default_resource , configured_resource = prm_resource ):
578578 raise OAuthFlowError (f"Protected resource { prm_resource } does not match expected { default_resource } " )
579579
580+ async def _discover_oauth_metadata (self ) -> AsyncGenerator [httpx2 .Request , httpx2 .Response ]:
581+ """Discover authorization server metadata and populate the context.
582+
583+ Yields the discovery requests so they run through the outer httpx auth flow
584+ (no side-channel client). This is pure discovery: it fills in
585+ ``protected_resource_metadata`` / ``auth_server_url`` / ``oauth_metadata`` and
586+ does not register clients or mutate stored credentials. Used to populate the
587+ token endpoint before an eager refresh, and available for the 401 path.
588+ """
589+ # Protected resource metadata -> authorization server URL. Best-effort: legacy
590+ # servers without PRM fall through to the origin well-known in the ASM step.
591+ if self .context .auth_server_url is None :
592+ for url in build_protected_resource_metadata_discovery_urls (None , self .context .server_url ):
593+ prm = await handle_protected_resource_response ((yield create_oauth_metadata_request (url )))
594+ if prm :
595+ await self ._validate_resource_match (prm )
596+ self .context .protected_resource_metadata = prm
597+ self .context .auth_server_url = str (prm .authorization_servers [0 ])
598+ break
599+
600+ # Authorization server metadata -> token / authorization / registration endpoints.
601+ for url in build_oauth_authorization_server_metadata_discovery_urls (
602+ self .context .auth_server_url , self .context .server_url
603+ ):
604+ ok , asm = await handle_auth_metadata_response ((yield create_oauth_metadata_request (url )))
605+ if not ok :
606+ break
607+ if asm :
608+ if self .context .auth_server_url is not None :
609+ validate_metadata_issuer (asm , self .context .auth_server_url )
610+ self .context .oauth_metadata = asm
611+ break
612+
613+ async def _refresh_with_discovery (self ) -> AsyncGenerator [httpx2 .Request , httpx2 .Response ]:
614+ """Eager token refresh that discovers authorization-server metadata first when
615+ it is not yet known.
616+
617+ The token endpoint comes from the AS metadata. On a cold start (e.g. reusing a
618+ stored refresh token before any 401) that metadata has not been discovered, so
619+ ``_refresh_token`` would fall back to ``{origin}/token`` — dropping any issuer
620+ path and 404ing on servers whose token endpoint lives under a path. Yields the
621+ discovery and refresh requests so they run through the outer httpx auth flow.
622+ """
623+ if self .context .oauth_metadata is None :
624+ discovery = self ._discover_oauth_metadata ()
625+ discovery_request = await discovery .asend (None )
626+ while True :
627+ discovery_response = yield discovery_request
628+ try :
629+ discovery_request = await discovery .asend (discovery_response )
630+ except StopAsyncIteration :
631+ break
632+
633+ refresh_response = yield await self ._refresh_token ()
634+ if not await self ._handle_refresh_response (refresh_response ):
635+ # Refresh failed, need full re-authentication
636+ self ._initialized = False
637+
580638 async def async_auth_flow (self , request : httpx2 .Request ) -> AsyncGenerator [httpx2 .Request , httpx2 .Response ]:
581639 """httpx2 auth flow integration."""
582640 async with self .context .lock :
@@ -587,13 +645,17 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
587645 self .context .protocol_version = request .headers .get (MCP_PROTOCOL_VERSION_HEADER )
588646
589647 if not self .context .is_token_valid () and self .context .can_refresh_token ():
590- # Try to refresh token
591- refresh_request = await self ._refresh_token ()
592- refresh_response = yield refresh_request
593-
594- if not await self ._handle_refresh_response (refresh_response ):
595- # Refresh failed, need full re-authentication
596- self ._initialized = False
648+ # Refresh the token, discovering authorization-server metadata first when
649+ # it is not yet known (see _refresh_with_discovery). Driven here so its
650+ # requests run through this httpx auth flow, not a side-channel client.
651+ refresh_flow = self ._refresh_with_discovery ()
652+ refresh_request = await refresh_flow .asend (None )
653+ while True :
654+ refresh_response = yield refresh_request
655+ try :
656+ refresh_request = await refresh_flow .asend (refresh_response )
657+ except StopAsyncIteration :
658+ break
597659
598660 if self .context .is_token_valid ():
599661 self ._add_auth_header (request )
0 commit comments