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