Skip to content

Commit f20d52e

Browse files
BabyChrist666claude
andcommitted
fix: restore eager OAuth discovery to avoid slow unauthenticated roundtrip (#1274)
When the client has no valid tokens, perform OAuth discovery and authorization BEFORE sending the MCP request. This restores the eager behavior from v1.11.0 that was removed in v1.12.0, eliminating the unnecessary unauthenticated roundtrip that servers like Notion handle slowly (~10s latency per operation). Both the eager (pre-request) and reactive (post-401) paths now share a single `_perform_oauth_discovery_and_auth()` helper, keeping the code DRY while preserving RFC 9728 WWW-Authenticate header support on the 401 path. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent be5bb7c commit f20d52e

File tree

3 files changed

+247
-99
lines changed

3 files changed

+247
-99
lines changed

src/mcp/client/auth/oauth2.py

Lines changed: 122 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)