Skip to content

Commit 418914e

Browse files
committed
Discover auth-server metadata before eager token refresh
On a cold start with a cached-but-expired token (reusing a stored refresh token before any 401), async_auth_flow refreshed before discovery ran, so oauth_metadata was None and _refresh_token fell back to {origin}/token. That drops any issuer path and 404s on servers whose token endpoint lives under a path (e.g. .../oauth2/api/v1/token) — the refresh fails, tokens are cleared, and the client is forced into interactive re-auth it may not be able to complete. Discover AS metadata before the eager refresh. Adds _discover_oauth_metadata (pure discovery, driven through the auth flow) and _refresh_with_discovery, plus a regression test.
1 parent a4f4ccd commit 418914e

2 files changed

Lines changed: 126 additions & 7 deletions

File tree

src/mcp/client/auth/oauth2.py

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

tests/client/test_auth.py

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3253,3 +3253,60 @@ async def echo_callback() -> AuthorizationCodeResult:
32533253
await auth_flow.asend(httpx2.Response(200, request=final_req))
32543254
except StopAsyncIteration:
32553255
pass
3256+
3257+
3258+
@pytest.mark.anyio
3259+
async def test_eager_refresh_discovers_token_endpoint_before_refreshing(
3260+
oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
3261+
):
3262+
"""Regression: on a cold start (cached expired token, no prior discovery) the
3263+
eager refresh must discover authorization-server metadata first, so it targets
3264+
the real token endpoint instead of the ``{origin}/token`` fallback. That
3265+
fallback drops any issuer path and 404s on servers whose token endpoint lives
3266+
under a path, which silently clears tokens and forces interactive re-auth.
3267+
"""
3268+
oauth_provider.context.current_tokens = valid_tokens
3269+
oauth_provider.context.token_expiry_time = time.time() - 100 # expired
3270+
oauth_provider.context.client_info = OAuthClientInformationFull(
3271+
client_id="test_client",
3272+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
3273+
token_endpoint_auth_method="none",
3274+
)
3275+
oauth_provider._initialized = True
3276+
assert oauth_provider.context.oauth_metadata is None
3277+
3278+
test_request = httpx2.Request("GET", "https://api.example.com/v1/mcp")
3279+
auth_flow = oauth_provider.async_auth_flow(test_request)
3280+
3281+
# 1) protected-resource metadata discovery
3282+
prm_request = await auth_flow.__anext__()
3283+
assert "oauth-protected-resource" in str(prm_request.url)
3284+
prm_response = httpx2.Response(
3285+
200,
3286+
content=(
3287+
b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}'
3288+
),
3289+
request=prm_request,
3290+
)
3291+
3292+
# 2) authorization-server metadata whose token endpoint is NOT {origin}/token
3293+
asm_request = await auth_flow.asend(prm_response)
3294+
assert "oauth-authorization-server" in str(asm_request.url)
3295+
asm_response = httpx2.Response(
3296+
200,
3297+
content=(
3298+
b'{"issuer": "https://auth.example.com", '
3299+
b'"authorization_endpoint": "https://auth.example.com/oauth2/authorize", '
3300+
b'"token_endpoint": "https://auth.example.com/oauth2/api/v1/token"}'
3301+
),
3302+
request=asm_request,
3303+
)
3304+
3305+
# 3) the refresh must target the discovered token endpoint, not the fallback
3306+
refresh_request = await auth_flow.asend(asm_response)
3307+
assert refresh_request.method == "POST"
3308+
assert str(refresh_request.url) == "https://auth.example.com/oauth2/api/v1/token"
3309+
assert str(refresh_request.url) != "https://api.example.com/token"
3310+
assert "grant_type=refresh_token" in refresh_request.content.decode()
3311+
3312+
await auth_flow.aclose()

0 commit comments

Comments
 (0)