@@ -498,6 +498,102 @@ async def _validate_resource_match(self, prm: ProtectedResourceMetadata) -> None
498498 if not check_resource_allowed (requested_resource = default_resource , configured_resource = prm_resource ):
499499 raise OAuthFlowError (f"Protected resource { prm_resource } does not match expected { default_resource } " )
500500
501+ async def _perform_oauth_discovery_and_auth (
502+ self ,
503+ www_auth_response : httpx .Response | None = None ,
504+ ) -> AsyncGenerator [httpx .Request , httpx .Response ]:
505+ """Perform the full OAuth discovery, registration, and authorization flow.
506+
507+ This is extracted as a helper to allow both eager (pre-request) and
508+ reactive (post-401) OAuth flows to share the same implementation.
509+
510+ Args:
511+ www_auth_response: Optional 401 response to extract WWW-Authenticate
512+ header from for RFC 9728 resource_metadata discovery. When None,
513+ falls back to well-known URL discovery.
514+ """
515+ www_auth_resource_metadata_url = (
516+ extract_resource_metadata_from_www_auth (www_auth_response ) if www_auth_response else None
517+ )
518+
519+ # Step 1: Discover protected resource metadata (SEP-985 with fallback support)
520+ prm_discovery_urls = build_protected_resource_metadata_discovery_urls (
521+ www_auth_resource_metadata_url , self .context .server_url
522+ )
523+
524+ for url in prm_discovery_urls : # pragma: no branch
525+ discovery_request = create_oauth_metadata_request (url )
526+
527+ discovery_response = yield discovery_request # sending request
528+
529+ prm = await handle_protected_resource_response (discovery_response )
530+ if prm :
531+ # Validate PRM resource matches server URL (RFC 8707)
532+ await self ._validate_resource_match (prm )
533+ self .context .protected_resource_metadata = prm
534+
535+ # todo: try all authorization_servers to find the OASM
536+ assert (
537+ len (prm .authorization_servers ) > 0
538+ ) # this is always true as authorization_servers has a min length of 1
539+
540+ self .context .auth_server_url = str (prm .authorization_servers [0 ])
541+ break
542+ else :
543+ logger .debug (f"Protected resource metadata discovery failed: { url } " )
544+
545+ asm_discovery_urls = build_oauth_authorization_server_metadata_discovery_urls (
546+ self .context .auth_server_url , self .context .server_url
547+ )
548+
549+ # Step 2: Discover OAuth Authorization Server Metadata (OASM)
550+ for url in asm_discovery_urls : # pragma: no branch
551+ oauth_metadata_request = create_oauth_metadata_request (url )
552+ oauth_metadata_response = yield oauth_metadata_request
553+
554+ ok , asm = await handle_auth_metadata_response (oauth_metadata_response )
555+ if not ok :
556+ break
557+ if ok and asm :
558+ self .context .oauth_metadata = asm
559+ break
560+ else :
561+ logger .debug (f"OAuth metadata discovery failed: { url } " )
562+
563+ # Step 3: Apply scope selection strategy
564+ self .context .client_metadata .scope = get_client_metadata_scopes (
565+ extract_scope_from_www_auth (www_auth_response ) if www_auth_response else None ,
566+ self .context .protected_resource_metadata ,
567+ self .context .oauth_metadata ,
568+ )
569+
570+ # Step 4: Register client or use URL-based client ID (CIMD)
571+ if not self .context .client_info :
572+ if should_use_client_metadata_url (self .context .oauth_metadata , self .context .client_metadata_url ):
573+ # Use URL-based client ID (CIMD)
574+ logger .debug (f"Using URL-based client ID (CIMD): { self .context .client_metadata_url } " )
575+ client_information = create_client_info_from_metadata_url (
576+ self .context .client_metadata_url , # type: ignore[arg-type]
577+ redirect_uris = self .context .client_metadata .redirect_uris ,
578+ )
579+ self .context .client_info = client_information
580+ await self .context .storage .set_client_info (client_information )
581+ else :
582+ # Fallback to Dynamic Client Registration
583+ registration_request = create_client_registration_request (
584+ self .context .oauth_metadata ,
585+ self .context .client_metadata ,
586+ self .context .get_authorization_base_url (self .context .server_url ),
587+ )
588+ registration_response = yield registration_request
589+ client_information = await handle_registration_response (registration_response )
590+ self .context .client_info = client_information
591+ await self .context .storage .set_client_info (client_information )
592+
593+ # Step 5: Perform authorization and complete token exchange
594+ token_response = yield await self ._perform_authorization ()
595+ await self ._handle_token_response (token_response )
596+
501597 async def async_auth_flow (self , request : httpx .Request ) -> AsyncGenerator [httpx .Request , httpx .Response ]:
502598 """HTTPX auth flow integration."""
503599 async with self .context .lock :
@@ -516,96 +612,48 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
516612 # Refresh failed, need full re-authentication
517613 self ._initialized = False
518614
615+ # Eager OAuth: if we have no valid token, can't refresh, AND we have
616+ # already been through the OAuth flow at least once (we have
617+ # client_info and discovery metadata), re-run the discovery/auth flow
618+ # BEFORE sending the MCP request. This avoids the unnecessary
619+ # unauthenticated round-trip that some servers (e.g. Notion) handle
620+ # slowly, causing ~10 s latency per request. See #1274.
621+ #
622+ # On the very first connection (no client_info), we skip the eager
623+ # flow and let the reactive 401 path handle discovery, because the
624+ # server's WWW-Authenticate header may carry routing information
625+ # (e.g. resource_metadata URL) that pure well-known discovery lacks.
626+ #
627+ # If the eager flow fails, we fall through gracefully and send the
628+ # MCP request without auth so the reactive 401 path can take over.
629+ if not self .context .is_token_valid () and self .context .client_info :
630+ try :
631+ oauth_gen = self ._perform_oauth_discovery_and_auth ()
632+ oauth_request = await oauth_gen .__anext__ ()
633+ while True :
634+ oauth_response = yield oauth_request
635+ oauth_request = await oauth_gen .asend (oauth_response )
636+ except StopAsyncIteration :
637+ pass
638+ except Exception :
639+ logger .debug ("Eager OAuth discovery failed, falling back to reactive 401 path" , exc_info = True )
640+
519641 if self .context .is_token_valid ():
520642 self ._add_auth_header (request )
521643
522644 response = yield request
523645
524646 if response .status_code == 401 :
525- # Perform full OAuth flow
647+ # Perform full OAuth flow (reactive path — uses WWW-Authenticate
648+ # header from the 401 response for RFC 9728 discovery)
526649 try :
527- # OAuth flow must be inline due to generator constraints
528- www_auth_resource_metadata_url = extract_resource_metadata_from_www_auth (response )
529-
530- # Step 1: Discover protected resource metadata (SEP-985 with fallback support)
531- prm_discovery_urls = build_protected_resource_metadata_discovery_urls (
532- www_auth_resource_metadata_url , self .context .server_url
533- )
534-
535- for url in prm_discovery_urls : # pragma: no branch
536- discovery_request = create_oauth_metadata_request (url )
537-
538- discovery_response = yield discovery_request # sending request
539-
540- prm = await handle_protected_resource_response (discovery_response )
541- if prm :
542- # Validate PRM resource matches server URL (RFC 8707)
543- await self ._validate_resource_match (prm )
544- self .context .protected_resource_metadata = prm
545-
546- # todo: try all authorization_servers to find the OASM
547- assert (
548- len (prm .authorization_servers ) > 0
549- ) # this is always true as authorization_servers has a min length of 1
550-
551- self .context .auth_server_url = str (prm .authorization_servers [0 ])
552- break
553- else :
554- logger .debug (f"Protected resource metadata discovery failed: { url } " )
555-
556- asm_discovery_urls = build_oauth_authorization_server_metadata_discovery_urls (
557- self .context .auth_server_url , self .context .server_url
558- )
559-
560- # Step 2: Discover OAuth Authorization Server Metadata (OASM) (with fallback for legacy servers)
561- for url in asm_discovery_urls : # pragma: no branch
562- oauth_metadata_request = create_oauth_metadata_request (url )
563- oauth_metadata_response = yield oauth_metadata_request
564-
565- ok , asm = await handle_auth_metadata_response (oauth_metadata_response )
566- if not ok :
567- break
568- if ok and asm :
569- self .context .oauth_metadata = asm
570- break
571- else :
572- logger .debug (f"OAuth metadata discovery failed: { url } " )
573-
574- # Step 3: Apply scope selection strategy
575- self .context .client_metadata .scope = get_client_metadata_scopes (
576- extract_scope_from_www_auth (response ),
577- self .context .protected_resource_metadata ,
578- self .context .oauth_metadata ,
579- )
580-
581- # Step 4: Register client or use URL-based client ID (CIMD)
582- if not self .context .client_info :
583- if should_use_client_metadata_url (
584- self .context .oauth_metadata , self .context .client_metadata_url
585- ):
586- # Use URL-based client ID (CIMD)
587- logger .debug (f"Using URL-based client ID (CIMD): { self .context .client_metadata_url } " )
588- client_information = create_client_info_from_metadata_url (
589- self .context .client_metadata_url , # type: ignore[arg-type]
590- redirect_uris = self .context .client_metadata .redirect_uris ,
591- )
592- self .context .client_info = client_information
593- await self .context .storage .set_client_info (client_information )
594- else :
595- # Fallback to Dynamic Client Registration
596- registration_request = create_client_registration_request (
597- self .context .oauth_metadata ,
598- self .context .client_metadata ,
599- self .context .get_authorization_base_url (self .context .server_url ),
600- )
601- registration_response = yield registration_request
602- client_information = await handle_registration_response (registration_response )
603- self .context .client_info = client_information
604- await self .context .storage .set_client_info (client_information )
605-
606- # Step 5: Perform authorization and complete token exchange
607- token_response = yield await self ._perform_authorization ()
608- await self ._handle_token_response (token_response )
650+ oauth_gen = self ._perform_oauth_discovery_and_auth (www_auth_response = response )
651+ oauth_request = await oauth_gen .__anext__ ()
652+ while True :
653+ oauth_response = yield oauth_request
654+ oauth_request = await oauth_gen .asend (oauth_response )
655+ except StopAsyncIteration :
656+ pass
609657 except Exception : # pragma: no cover
610658 logger .exception ("OAuth flow error" )
611659 raise
0 commit comments