Skip to content

Commit a80aae2

Browse files
committed
fix(client/auth): discover AS metadata before cold-start token refresh
On a cold start (stored refresh token reused before any 401) the eager pre-401 refresh built its URL from the urljoin(origin, "/token") fallback because authorization-server metadata had not been discovered yet. Servers whose token endpoint lives under a path returned 404, the client cleared its stored tokens, and headless clients were forced into an interactive re-auth they cannot perform (#3240, #3250). Run protected-resource + authorization-server metadata discovery before the eager refresh so it targets the discovered token endpoint, applying the same SEP-2352 issuer-binding checks as the 401 discovery path: when the stored credentials are bound to a different issuer they are dropped and the refresh is skipped, so credentials are never presented to an authorization server they are not bound to, and the subsequent 401 flow re-registers cleanly. Servers publishing no metadata keep the previous {origin}/token fallback behavior. Fixes #3240 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CjbXueCDdFNJK6imejCXgM
1 parent a4f4ccd commit a80aae2

2 files changed

Lines changed: 372 additions & 7 deletions

File tree

src/mcp/client/auth/oauth2.py

Lines changed: 97 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -577,6 +577,92 @@ 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 _refresh_with_discovery(self) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
581+
"""Refresh the token, discovering authorization-server metadata first when needed.
582+
583+
The token endpoint comes from the AS metadata. On a cold start (a stored refresh
584+
token reused before any 401) that metadata has not been discovered yet, so
585+
``_refresh_token`` would fall back to ``{origin}/token`` — dropping any issuer
586+
path and 404ing on servers whose token endpoint lives elsewhere. Discovery runs
587+
first, applying the same SEP-2352 issuer-binding checks as the 401 path so stored
588+
credentials are never sent to an authorization server they are not bound to: on a
589+
binding mismatch the credentials and tokens are dropped and the refresh is
590+
skipped, letting the subsequent 401 flow re-register against the new server.
591+
Yields the discovery and refresh requests so they run through the outer httpx
592+
auth flow rather than a side-channel client.
593+
"""
594+
if self.context.oauth_metadata is None:
595+
# Step 1: protected resource metadata -> authorization server URL (SEP-985).
596+
# Best-effort: a legacy server without PRM falls through to the origin
597+
# well-known fallback in the ASM step below. There is no 401 response at
598+
# this point, so no WWW-Authenticate resource_metadata hint is available.
599+
for url in build_protected_resource_metadata_discovery_urls(None, self.context.server_url):
600+
prm = await handle_protected_resource_response((yield create_oauth_metadata_request(url)))
601+
if prm:
602+
# Validate PRM resource matches server URL (RFC 8707)
603+
await self._validate_resource_match(prm)
604+
self.context.protected_resource_metadata = prm
605+
self.context.auth_server_url = str(prm.authorization_servers[0])
606+
break
607+
else:
608+
logger.debug(f"Protected resource metadata discovery failed: {url}")
609+
610+
# SEP-2352: stored credentials are bound to the issuer that registered them.
611+
# If the authorization server changed, drop them (and the old tokens) and skip
612+
# the refresh so the 401 flow re-registers instead of presenting another
613+
# server's credentials to the newly discovered one.
614+
if (
615+
self.context.client_info is not None
616+
and self.context.auth_server_url is not None
617+
and not credentials_match_issuer(
618+
self.context.client_info, self.context.auth_server_url, self.context.client_metadata_url
619+
)
620+
):
621+
logger.debug("Authorization server changed; discarding bound credentials and skipping refresh")
622+
self.context.client_info = None
623+
self.context.clear_tokens()
624+
return
625+
626+
# Step 2: authorization server metadata -> the token endpoint (with fallback
627+
# for legacy servers).
628+
for url in build_oauth_authorization_server_metadata_discovery_urls(
629+
self.context.auth_server_url, self.context.server_url
630+
):
631+
ok, asm = await handle_auth_metadata_response((yield create_oauth_metadata_request(url)))
632+
if not ok:
633+
break
634+
if asm:
635+
# SEP-2468: metadata issuer must match the discovery issuer
636+
if self.context.auth_server_url is not None:
637+
validate_metadata_issuer(asm, self.context.auth_server_url)
638+
self.context.oauth_metadata = asm
639+
break
640+
else:
641+
logger.debug(f"OAuth metadata discovery failed: {url}")
642+
643+
# SEP-2352: on the legacy no-PRM path the issuer is only known after ASM
644+
# discovery, so re-evaluate the binding here using the discovered metadata
645+
# issuer (mirroring the 401 path's post-ASM check).
646+
if (
647+
self.context.client_info is not None
648+
and self.context.auth_server_url is None
649+
and self.context.oauth_metadata is not None
650+
and not credentials_match_issuer(
651+
self.context.client_info,
652+
str(self.context.oauth_metadata.issuer),
653+
self.context.client_metadata_url,
654+
)
655+
):
656+
logger.debug("Authorization server changed; discarding bound credentials and skipping refresh")
657+
self.context.client_info = None
658+
self.context.clear_tokens()
659+
return
660+
661+
refresh_response = yield await self._refresh_token()
662+
if not await self._handle_refresh_response(refresh_response):
663+
# Refresh failed, need full re-authentication
664+
self._initialized = False
665+
580666
async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
581667
"""httpx2 auth flow integration."""
582668
async with self.context.lock:
@@ -587,13 +673,17 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
587673
self.context.protocol_version = request.headers.get(MCP_PROTOCOL_VERSION_HEADER)
588674

589675
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
676+
# Refresh the token, discovering authorization-server metadata first on a
677+
# cold start (see _refresh_with_discovery). Driven inline so its requests
678+
# run through this httpx auth flow, not a side-channel client.
679+
refresh_flow = self._refresh_with_discovery()
680+
refresh_request = await refresh_flow.__anext__()
681+
while True:
682+
refresh_response = yield refresh_request
683+
try:
684+
refresh_request = await refresh_flow.asend(refresh_response)
685+
except StopAsyncIteration:
686+
break
597687

598688
if self.context.is_token_valid():
599689
self._add_auth_header(request)

tests/client/test_auth.py

Lines changed: 275 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3253,3 +3253,278 @@ 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, mock_storage: MockTokenStorage, valid_tokens: OAuthToken
3261+
):
3262+
"""Regression for #3240/#3250: a cold-start eager refresh discovers the token endpoint.
3263+
3264+
On a restart with a stored (expired) token the pre-401 refresh used to POST to the
3265+
``{origin}/token`` fallback because authorization-server metadata had not been
3266+
discovered yet, 404ing on servers whose token endpoint lives under a path and
3267+
silently clearing the stored tokens. The refresh must run PRM + ASM discovery first
3268+
and target the discovered token endpoint.
3269+
"""
3270+
oauth_provider.context.current_tokens = valid_tokens
3271+
oauth_provider.context.token_expiry_time = time.time() - 100 # expired
3272+
oauth_provider.context.client_info = OAuthClientInformationFull(
3273+
client_id="test_client",
3274+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
3275+
token_endpoint_auth_method="none",
3276+
)
3277+
oauth_provider._initialized = True
3278+
assert oauth_provider.context.oauth_metadata is None
3279+
3280+
test_request = httpx2.Request("GET", "https://api.example.com/v1/mcp")
3281+
auth_flow = oauth_provider.async_auth_flow(test_request)
3282+
3283+
# 1) protected-resource metadata discovery (no WWW-Authenticate hint pre-401)
3284+
prm_request = await auth_flow.__anext__()
3285+
assert str(prm_request.url) == "https://api.example.com/.well-known/oauth-protected-resource/v1/mcp"
3286+
prm_response = httpx2.Response(
3287+
200,
3288+
content=(
3289+
b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}'
3290+
),
3291+
request=prm_request,
3292+
)
3293+
3294+
# 2) authorization-server metadata whose token endpoint is NOT {origin}/token
3295+
asm_request = await auth_flow.asend(prm_response)
3296+
assert str(asm_request.url) == "https://auth.example.com/.well-known/oauth-authorization-server"
3297+
asm_response = httpx2.Response(
3298+
200,
3299+
content=(
3300+
b'{"issuer": "https://auth.example.com", '
3301+
b'"authorization_endpoint": "https://auth.example.com/oauth2/authorize", '
3302+
b'"token_endpoint": "https://auth.example.com/oauth2/api/v1/token"}'
3303+
),
3304+
request=asm_request,
3305+
)
3306+
3307+
# 3) the refresh targets the discovered token endpoint, not the fallback
3308+
refresh_request = await auth_flow.asend(asm_response)
3309+
assert refresh_request.method == "POST"
3310+
assert str(refresh_request.url) == "https://auth.example.com/oauth2/api/v1/token"
3311+
assert "grant_type=refresh_token" in refresh_request.content.decode()
3312+
refresh_response = httpx2.Response(
3313+
200,
3314+
json={"access_token": "refreshed_token", "token_type": "Bearer", "expires_in": 3600},
3315+
request=refresh_request,
3316+
)
3317+
3318+
# 4) the original request goes out with the refreshed token
3319+
api_request = await auth_flow.asend(refresh_response)
3320+
assert str(api_request.url) == "https://api.example.com/v1/mcp"
3321+
assert api_request.headers["Authorization"] == "Bearer refreshed_token"
3322+
stored = await mock_storage.get_tokens()
3323+
assert stored is not None
3324+
assert stored.access_token == "refreshed_token"
3325+
3326+
with pytest.raises(StopAsyncIteration):
3327+
await auth_flow.asend(httpx2.Response(200, request=api_request))
3328+
3329+
3330+
@pytest.mark.anyio
3331+
async def test_eager_refresh_falls_back_to_origin_token_when_no_metadata_published(
3332+
oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
3333+
):
3334+
"""A legacy server publishing no metadata keeps the pre-existing ``{origin}/token`` fallback.
3335+
3336+
PRM discovery 404s at both well-known URLs and the legacy origin ASM fallback 404s too,
3337+
so the refresh still POSTs to ``{origin}/token`` exactly as before discovery-before-refresh
3338+
existed. A failed refresh then clears tokens and lets the request go out unauthenticated.
3339+
"""
3340+
oauth_provider.context.current_tokens = valid_tokens
3341+
oauth_provider.context.token_expiry_time = time.time() - 100 # expired
3342+
oauth_provider.context.client_info = OAuthClientInformationFull(
3343+
client_id="test_client",
3344+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
3345+
token_endpoint_auth_method="none",
3346+
)
3347+
oauth_provider._initialized = True
3348+
3349+
auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
3350+
3351+
# PRM discovery: path-based then root-based, both 404.
3352+
prm_request = await auth_flow.__anext__()
3353+
assert str(prm_request.url) == "https://api.example.com/.well-known/oauth-protected-resource/v1/mcp"
3354+
prm_request = await auth_flow.asend(httpx2.Response(404, request=prm_request))
3355+
assert str(prm_request.url) == "https://api.example.com/.well-known/oauth-protected-resource"
3356+
3357+
# ASM discovery: legacy origin fallback, 404 as well.
3358+
asm_request = await auth_flow.asend(httpx2.Response(404, request=prm_request))
3359+
assert str(asm_request.url) == "https://api.example.com/.well-known/oauth-authorization-server"
3360+
3361+
# Refresh falls back to {origin}/token (pre-existing legacy behavior).
3362+
refresh_request = await auth_flow.asend(httpx2.Response(404, request=asm_request))
3363+
assert refresh_request.method == "POST"
3364+
assert str(refresh_request.url) == "https://api.example.com/token"
3365+
3366+
# The refresh fails; tokens are cleared and the original request goes out unauthenticated.
3367+
api_request = await auth_flow.asend(httpx2.Response(401, request=refresh_request))
3368+
assert str(api_request.url) == "https://api.example.com/v1/mcp"
3369+
assert "Authorization" not in api_request.headers
3370+
assert oauth_provider.context.current_tokens is None
3371+
await auth_flow.aclose()
3372+
3373+
3374+
@pytest.mark.anyio
3375+
async def test_eager_refresh_stops_asm_discovery_on_server_error(
3376+
oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
3377+
):
3378+
"""A non-4XX ASM discovery error stops the fallback chain, mirroring the 401 path.
3379+
3380+
The refresh then proceeds against the ``{origin}/token`` fallback rather than
3381+
hammering further well-known URLs.
3382+
"""
3383+
oauth_provider.context.current_tokens = valid_tokens
3384+
oauth_provider.context.token_expiry_time = time.time() - 100 # expired
3385+
oauth_provider.context.client_info = OAuthClientInformationFull(
3386+
client_id="test_client",
3387+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
3388+
token_endpoint_auth_method="none",
3389+
)
3390+
oauth_provider._initialized = True
3391+
3392+
auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
3393+
3394+
# PRM discovery succeeds and points at the authorization server.
3395+
prm_request = await auth_flow.__anext__()
3396+
prm_response = httpx2.Response(
3397+
200,
3398+
content=(
3399+
b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}'
3400+
),
3401+
request=prm_request,
3402+
)
3403+
3404+
# ASM discovery hits a 500: stop trying further URLs.
3405+
asm_request = await auth_flow.asend(prm_response)
3406+
assert str(asm_request.url) == "https://auth.example.com/.well-known/oauth-authorization-server"
3407+
refresh_request = await auth_flow.asend(httpx2.Response(500, request=asm_request))
3408+
3409+
assert refresh_request.method == "POST"
3410+
assert str(refresh_request.url) == "https://api.example.com/token"
3411+
await auth_flow.aclose()
3412+
3413+
3414+
@pytest.mark.anyio
3415+
async def test_eager_refresh_skips_refresh_when_credentials_bound_to_different_issuer(
3416+
oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
3417+
):
3418+
"""SEP-2352: a cold-start refresh never sends credentials bound to another issuer.
3419+
3420+
When PRM discovery reveals an authorization server different from the one the stored
3421+
client credentials are bound to, the credentials and tokens are dropped and the
3422+
refresh is skipped, so the subsequent 401 flow re-registers against the new server
3423+
— mirroring the issuer-binding check on the 401 discovery path.
3424+
"""
3425+
oauth_provider.context.current_tokens = valid_tokens
3426+
oauth_provider.context.token_expiry_time = time.time() - 100 # expired
3427+
oauth_provider.context.client_info = OAuthClientInformationFull(
3428+
client_id="stale-client",
3429+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
3430+
issuer="https://old-as.example.com",
3431+
)
3432+
oauth_provider._initialized = True
3433+
3434+
auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
3435+
3436+
# PRM discovery points at auth.example.com, not the bound old-as.example.com.
3437+
prm_request = await auth_flow.__anext__()
3438+
prm_response = httpx2.Response(
3439+
200,
3440+
content=(
3441+
b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}'
3442+
),
3443+
request=prm_request,
3444+
)
3445+
3446+
# No refresh request: the next yield is the original request, unauthenticated.
3447+
api_request = await auth_flow.asend(prm_response)
3448+
assert str(api_request.url) == "https://api.example.com/v1/mcp"
3449+
assert "Authorization" not in api_request.headers
3450+
assert oauth_provider.context.client_info is None
3451+
assert oauth_provider.context.current_tokens is None
3452+
await auth_flow.aclose()
3453+
3454+
3455+
@pytest.mark.anyio
3456+
async def test_eager_refresh_legacy_path_rechecks_issuer_binding_after_asm(
3457+
oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
3458+
):
3459+
"""SEP-2352 on the legacy no-PRM path: the binding is checked against the ASM issuer.
3460+
3461+
PRM discovery fails so the issuer is only known once origin-fallback ASM discovery
3462+
succeeds; credentials bound to a different issuer are then dropped and the refresh is
3463+
skipped, exactly as on the 401 path's post-ASM re-evaluation.
3464+
"""
3465+
oauth_provider.context.current_tokens = valid_tokens
3466+
oauth_provider.context.token_expiry_time = time.time() - 100 # expired
3467+
oauth_provider.context.client_info = OAuthClientInformationFull(
3468+
client_id="stale-client",
3469+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
3470+
issuer="https://old-as.example.com",
3471+
)
3472+
oauth_provider._initialized = True
3473+
3474+
auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
3475+
3476+
# PRM discovery: both well-known URLs 404.
3477+
prm_request = await auth_flow.__anext__()
3478+
prm_request = await auth_flow.asend(httpx2.Response(404, request=prm_request))
3479+
3480+
# Origin-fallback ASM discovery succeeds with the resource origin as issuer.
3481+
asm_request = await auth_flow.asend(httpx2.Response(404, request=prm_request))
3482+
assert str(asm_request.url) == "https://api.example.com/.well-known/oauth-authorization-server"
3483+
asm_response = httpx2.Response(
3484+
200,
3485+
content=(
3486+
b'{"issuer": "https://api.example.com", '
3487+
b'"authorization_endpoint": "https://api.example.com/authorize", '
3488+
b'"token_endpoint": "https://api.example.com/token"}'
3489+
),
3490+
request=asm_request,
3491+
)
3492+
3493+
# No refresh request: the next yield is the original request, unauthenticated.
3494+
api_request = await auth_flow.asend(asm_response)
3495+
assert str(api_request.url) == "https://api.example.com/v1/mcp"
3496+
assert "Authorization" not in api_request.headers
3497+
assert oauth_provider.context.client_info is None
3498+
assert oauth_provider.context.current_tokens is None
3499+
# The just-discovered metadata is for the current server and is kept for the 401 flow.
3500+
assert oauth_provider.context.oauth_metadata is not None
3501+
await auth_flow.aclose()
3502+
3503+
3504+
@pytest.mark.anyio
3505+
async def test_eager_refresh_skips_discovery_when_metadata_already_known(
3506+
oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
3507+
):
3508+
"""With authorization-server metadata already discovered, the refresh is immediate."""
3509+
oauth_provider.context.current_tokens = valid_tokens
3510+
oauth_provider.context.token_expiry_time = time.time() - 100 # expired
3511+
oauth_provider.context.client_info = OAuthClientInformationFull(
3512+
client_id="test_client",
3513+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
3514+
token_endpoint_auth_method="none",
3515+
)
3516+
oauth_provider.context.oauth_metadata = OAuthMetadata.model_validate(
3517+
{
3518+
"issuer": "https://auth.example.com",
3519+
"authorization_endpoint": "https://auth.example.com/oauth2/authorize",
3520+
"token_endpoint": "https://auth.example.com/oauth2/api/v1/token",
3521+
}
3522+
)
3523+
oauth_provider._initialized = True
3524+
3525+
auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
3526+
3527+
refresh_request = await auth_flow.__anext__()
3528+
assert refresh_request.method == "POST"
3529+
assert str(refresh_request.url) == "https://auth.example.com/oauth2/api/v1/token"
3530+
await auth_flow.aclose()

0 commit comments

Comments
 (0)