Skip to content

Commit 33a7c1c

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 7ba41dc commit 33a7c1c

File tree

3 files changed

+303
-99
lines changed

3 files changed

+303
-99
lines changed

src/mcp/client/auth/oauth2.py

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

Comments
 (0)