diff --git a/src/mcp/client/auth/oauth2.py b/src/mcp/client/auth/oauth2.py index c513a99b46..67b5780d5a 100644 --- a/src/mcp/client/auth/oauth2.py +++ b/src/mcp/client/auth/oauth2.py @@ -172,7 +172,17 @@ class OAuthContext: token_expiry_time: float | None = None # State + # + # `lock` guards short-lived reads and writes of provider state. A request + # that already has a usable token yields outside the lock, so a long-poll + # GET SSE does not block other calls. A request with no usable token holds + # the lock until that challenge finishes: the call that met the 401 is the + # one whose timeout `waiting_on_a_person` pauses, and peers wait on the lock. + # + # `refresh_lock` single-flights token refresh: one refresh runs, and waiters + # re-check the token cache instead of starting another. lock: anyio.Lock = field(default_factory=anyio.Lock) + refresh_lock: anyio.Lock = field(default_factory=anyio.Lock) def get_authorization_base_url(self, server_url: str) -> str: """Extract base URL by removing path component.""" @@ -595,199 +605,268 @@ def _select_authorization_server(self, advertised: list[str]) -> str: """Which of the servers listed in protected resource metadata to use: the first (the list is never empty).""" return advertised[0] + def _needs_token_refresh(self) -> bool: + return (not self.context.is_token_valid()) and self.context.can_refresh_token() + + def _unauthorized_token_changed(self, response: httpx2.Response, sent_access_token: str | None) -> bool: + """True when a 401 is stale because another request already replaced the access token.""" + if response.status_code != 401: + return False + current = self.context.current_tokens.access_token if self.context.current_tokens else None + return current is not None and current != sent_access_token + def _expected_issuer(self) -> str: """The issuer that authorization server metadata and client credentials must belong to: the PRM-advertised server, or on the legacy no-PRM path the resource server's origin, which is what the 2025-03-26 well-known URL is built from (RFC 8414 §3.3).""" return self.context.auth_server_url or _origin_issuer(self.context.server_url) + async def _unauthorized_followups( + self, + request: httpx2.Request, + response: httpx2.Response, + sent_access_token: str | None, + ) -> AsyncGenerator[httpx2.Request, httpx2.Response]: + """Run discovery and login after a 401 or scope step-up. + + Caller holds ``context.lock``. Each yield is a follow-up HTTP request. + """ + step_up = response.status_code == 403 and extract_field_from_www_auth(response, "error") == "insufficient_scope" + if self._unauthorized_token_changed(response, sent_access_token): + # Another request refreshed or re-authorized while this one was in flight. + self._add_auth_header(request) + yield request + return + + # Perform full OAuth flow + try: + # Read before discovery, which may clear the tokens: on a restart the stored + # token's scope is the only record of what was granted (see Step 3). + granted_scope = self.context.current_tokens.scope if self.context.current_tokens else None + + # OAuth flow must be inline due to generator constraints. + # Steps 1-2 run on every 401. A scope step-up reuses the metadata discovered earlier + # in this process, and discovers it first when none is held yet (for example when + # tokens were loaded from storage), so re-authorization targets the right server. + if response.status_code == 401 or self.context.oauth_metadata is None: + www_auth_resource_metadata_url = extract_resource_metadata_from_www_auth(response) + + # Step 1: Discover protected resource metadata (SEP-985 with fallback support) + prm_discovery_urls = build_protected_resource_metadata_discovery_urls( + www_auth_resource_metadata_url, self.context.server_url + ) + + prm_request_failed: int | None = None + for url in prm_discovery_urls: + discovery_request = create_oauth_metadata_request(url) + + discovery_response = yield discovery_request # sending request + + if discovery_response.status_code >= 500 or discovery_response.status_code == 429: + prm_request_failed = discovery_response.status_code + prm = await handle_protected_resource_response(discovery_response) + if prm: + # Validate PRM resource matches server URL (RFC 8707) + await self._validate_resource_match(prm) + self.context.protected_resource_metadata = prm + + self.context.auth_server_url = self._select_authorization_server( + [str(url) for url in prm.authorization_servers] + ) + break + else: + logger.debug(f"Protected resource metadata discovery failed: {url}") + else: + if prm_request_failed is not None: + # A server error says nothing about whether the resource publishes + # metadata, so it must not send the flow down the legacy path. + raise OAuthFlowError(f"Protected resource metadata request failed: HTTP {prm_request_failed}") + + expected_issuer = self._expected_issuer() + + # SEP-2352: stored credentials are bound to the issuer that registered them. + # Decided before any metadata is fetched: if the expected issuer is a different + # server, drop them (and the old tokens) so the flow re-registers instead of + # presenting another server's credentials. + if self.context.client_info is not None and not credentials_match_issuer( + self.context.client_info, expected_issuer, self.context.client_metadata_url + ): + logger.debug("Authorization server changed; discarding bound credentials and re-registering") + self.context.client_info = None + self.context.clear_tokens() + # Any cached AS metadata is for the old server; drop it so a failed + # rediscovery cannot leak the old registration/token endpoints into Step 4. + self.context.oauth_metadata = None + + asm_discovery_urls = build_oauth_authorization_server_metadata_discovery_urls( + self.context.auth_server_url, self.context.server_url + ) + + # Step 2: Discover OAuth Authorization Server Metadata (OASM) (with fallback for legacy servers) + for url in asm_discovery_urls: # pragma: no branch + oauth_metadata_request = create_oauth_metadata_request(url) + oauth_metadata_response = yield oauth_metadata_request + + ok, asm = await handle_auth_metadata_response(oauth_metadata_response) + if not ok: + break + if ok and asm: + # SEP-2468 / RFC 8414 §3.3: the metadata must name the expected issuer. + # On the legacy path a root issuer rendered with its trailing slash + # names the same origin. + if self.context.auth_server_url is None and issuers_match(str(asm.issuer), expected_issuer): + expected_issuer = str(asm.issuer) + validate_metadata_issuer(asm, expected_issuer) + self.context.oauth_metadata = asm + break + else: + logger.debug(f"OAuth metadata discovery failed: {url}") + + # Step 3: Apply scope selection strategy + challenged_scope = get_client_metadata_scopes( + extract_scope_from_www_auth(response), + self.context.protected_resource_metadata, + self.context.oauth_metadata, + self.context.client_metadata.grant_types, + ) + if step_up: + # SEP-2350: union previously requested scopes with the newly challenged ones so + # escalating one operation keeps the others' grants, folding in the granted + # scope read above since client_metadata.scope is not reloaded on a restart. + prior_scope = union_scopes(self.context.client_metadata.scope, granted_scope) + self.context.client_metadata.scope = union_scopes(prior_scope, challenged_scope) + else: + self.context.client_metadata.scope = challenged_scope + + # Step 4: Register client or use URL-based client ID (CIMD) + if not self.context.client_info: + # SEP-2352: the issuer to bind these credentials to, once metadata for it + # was actually found. + discovered_issuer = self._expected_issuer() if self.context.oauth_metadata is not None else None + + if should_use_client_metadata_url(self.context.oauth_metadata, self.context.client_metadata_url): + # Use URL-based client ID (CIMD). CIMD records are portable across + # authorization servers, so the issuer stamp is informational. + logger.debug(f"Using URL-based client ID (CIMD): {self.context.client_metadata_url}") + client_information = create_client_info_from_metadata_url( + self.context.client_metadata_url, # type: ignore[arg-type] + redirect_uris=self.context.client_metadata.redirect_uris, + ) + client_information.issuer = discovered_issuer + self.context.client_info = client_information + await self.context.storage.set_client_info(client_information) + else: + # Fallback to Dynamic Client Registration + fallback_base = self.context.get_authorization_base_url(self.context.server_url) + registration_request = create_client_registration_request( + self.context.oauth_metadata, self.context.client_metadata, fallback_base + ) + registration_response = yield registration_request + client_information = await handle_registration_response(registration_response) + check_registration_usable(client_information) + # Only record the issuer when the registration above actually targeted + # the discovered AS — either via its published registration_endpoint, + # or because the resource-origin /register fallback is on the issuer's + # own host (legacy same-origin embedded AS). Otherwise the fallback hit + # a different server and recording a binding to the PRM-advertised AS + # would persist a binding that was never established. + if ( + self.context.oauth_metadata is not None + and discovered_issuer is not None + and ( + self.context.oauth_metadata.registration_endpoint is not None + or self.context.get_authorization_base_url(discovered_issuer) == fallback_base + ) + ): + client_information.issuer = discovered_issuer + self.context.client_info = client_information + await self.context.storage.set_client_info(client_information) + + # Step 5: Perform authorization and complete token exchange + token_response = yield await self._perform_authorization() + await self._handle_token_response(token_response) + except Exception: + logger.exception("OAuth flow error") + raise + + # Retry with new tokens + self._add_auth_header(request) + yield request + async def _auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: - """The OAuth flow proper; `async_auth_flow` drives it (see `RedirectAwareAuth`).""" + """The OAuth flow proper; `async_auth_flow` drives it (see `RedirectAwareAuth`). + + ``context.lock`` is held only while reading or mutating provider state, + except while a request that has no usable token is in flight. That + challenge holds the lock through discovery and login so a peer waits + here instead of starting a second login. A request that already carries + a token yields outside the lock. ``refresh_lock`` single-flights refresh. + """ + needs_refresh = False async with self.context.lock: if not self._initialized: await self._initialize() # Capture protocol version from request headers self.context.protocol_version = request.headers.get(MCP_PROTOCOL_VERSION_HEADER) - - if not self.context.is_token_valid() and self.context.can_refresh_token(): - # Try to refresh token - refresh_request = await self._refresh_token() - refresh_response = yield refresh_request - - if not await self._handle_refresh_response(refresh_response): - # Refresh failed, need full re-authentication - self._initialized = False - + needs_refresh = self._needs_token_refresh() + + if needs_refresh: + async with self.context.refresh_lock: + # Re-check outside the inner `async with` so a peer refresh that + # landed while this caller waited is a normal branch, not one + # hidden in the lock's exit. + still_needs_refresh = False + async with self.context.lock: + still_needs_refresh = self._needs_token_refresh() + refresh_request = None + if still_needs_refresh: + async with self.context.lock: + refresh_request = await self._refresh_token() + if refresh_request is not None: + refresh_response = yield refresh_request + refresh_ok = False + async with self.context.lock: + refresh_ok = await self._handle_refresh_response(refresh_response) + if not refresh_ok: + async with self.context.lock: + self._initialized = False + + attached = False + async with self.context.lock: if self.context.is_token_valid(): self._add_auth_header(request) + attached = True + sent_access_token = self.context.current_tokens.access_token if self.context.current_tokens else None + # No token yet: hold the lock across this yield and the login it starts. + # A peer blocks on the lock, so the challenged request is the one whose + # read timeout `waiting_on_a_person` pauses. + holding = False + try: + if not attached: + await self.context.lock.acquire() + holding = True response = yield request step_up = ( response.status_code == 403 and extract_field_from_www_auth(response, "error") == "insufficient_scope" ) - if response.status_code == 401 or step_up: - # Perform full OAuth flow + if not holding: + await self.context.lock.acquire() + holding = True + followups = self._unauthorized_followups(request, response, sent_access_token) try: - # Read before discovery, which may clear the tokens: on a restart the stored - # token's scope is the only record of what was granted (see Step 3). - granted_scope = self.context.current_tokens.scope if self.context.current_tokens else None - - # OAuth flow must be inline due to generator constraints. - # Steps 1-2 run on every 401. A scope step-up reuses the metadata discovered earlier - # in this process, and discovers it first when none is held yet (for example when - # tokens were loaded from storage), so re-authorization targets the right server. - if response.status_code == 401 or self.context.oauth_metadata is None: - www_auth_resource_metadata_url = extract_resource_metadata_from_www_auth(response) - - # Step 1: Discover protected resource metadata (SEP-985 with fallback support) - prm_discovery_urls = build_protected_resource_metadata_discovery_urls( - www_auth_resource_metadata_url, self.context.server_url - ) - - prm_request_failed: int | None = None - for url in prm_discovery_urls: - discovery_request = create_oauth_metadata_request(url) - - discovery_response = yield discovery_request # sending request - - if discovery_response.status_code >= 500 or discovery_response.status_code == 429: - prm_request_failed = discovery_response.status_code - prm = await handle_protected_resource_response(discovery_response) - if prm: - # Validate PRM resource matches server URL (RFC 8707) - await self._validate_resource_match(prm) - self.context.protected_resource_metadata = prm - - self.context.auth_server_url = self._select_authorization_server( - [str(url) for url in prm.authorization_servers] - ) - break - else: - logger.debug(f"Protected resource metadata discovery failed: {url}") - else: - if prm_request_failed is not None: - # A server error says nothing about whether the resource publishes - # metadata, so it must not send the flow down the legacy path. - raise OAuthFlowError( - f"Protected resource metadata request failed: HTTP {prm_request_failed}" - ) - - expected_issuer = self._expected_issuer() - - # SEP-2352: stored credentials are bound to the issuer that registered them. - # Decided before any metadata is fetched: if the expected issuer is a different - # server, drop them (and the old tokens) so the flow re-registers instead of - # presenting another server's credentials. - if self.context.client_info is not None and not credentials_match_issuer( - self.context.client_info, expected_issuer, self.context.client_metadata_url - ): - logger.debug( - "Authorization server changed; discarding bound credentials and re-registering" - ) - self.context.client_info = None - self.context.clear_tokens() - # Any cached AS metadata is for the old server; drop it so a failed - # rediscovery cannot leak the old registration/token endpoints into Step 4. - self.context.oauth_metadata = None - - asm_discovery_urls = build_oauth_authorization_server_metadata_discovery_urls( - self.context.auth_server_url, self.context.server_url - ) - - # Step 2: Discover OAuth Authorization Server Metadata (OASM) (with fallback for legacy servers) - for url in asm_discovery_urls: # pragma: no branch - oauth_metadata_request = create_oauth_metadata_request(url) - oauth_metadata_response = yield oauth_metadata_request - - ok, asm = await handle_auth_metadata_response(oauth_metadata_response) - if not ok: - break - if ok and asm: - # SEP-2468 / RFC 8414 §3.3: the metadata must name the expected issuer. - # On the legacy path a root issuer rendered with its trailing slash - # names the same origin. - if self.context.auth_server_url is None and issuers_match( - str(asm.issuer), expected_issuer - ): - expected_issuer = str(asm.issuer) - validate_metadata_issuer(asm, expected_issuer) - self.context.oauth_metadata = asm - break - else: - logger.debug(f"OAuth metadata discovery failed: {url}") - - # Step 3: Apply scope selection strategy - challenged_scope = get_client_metadata_scopes( - extract_scope_from_www_auth(response), - self.context.protected_resource_metadata, - self.context.oauth_metadata, - self.context.client_metadata.grant_types, - ) - if step_up: - # SEP-2350: union previously requested scopes with the newly challenged ones so - # escalating one operation keeps the others' grants, folding in the granted - # scope read above since client_metadata.scope is not reloaded on a restart. - prior_scope = union_scopes(self.context.client_metadata.scope, granted_scope) - self.context.client_metadata.scope = union_scopes(prior_scope, challenged_scope) - else: - self.context.client_metadata.scope = challenged_scope - - # Step 4: Register client or use URL-based client ID (CIMD) - if not self.context.client_info: - # SEP-2352: the issuer to bind these credentials to, once metadata for it - # was actually found. - discovered_issuer = self._expected_issuer() if self.context.oauth_metadata is not None else None - - if should_use_client_metadata_url( - self.context.oauth_metadata, self.context.client_metadata_url - ): - # Use URL-based client ID (CIMD). CIMD records are portable across - # authorization servers, so the issuer stamp is informational. - logger.debug(f"Using URL-based client ID (CIMD): {self.context.client_metadata_url}") - client_information = create_client_info_from_metadata_url( - self.context.client_metadata_url, # type: ignore[arg-type] - redirect_uris=self.context.client_metadata.redirect_uris, - ) - client_information.issuer = discovered_issuer - self.context.client_info = client_information - await self.context.storage.set_client_info(client_information) - else: - # Fallback to Dynamic Client Registration - fallback_base = self.context.get_authorization_base_url(self.context.server_url) - registration_request = create_client_registration_request( - self.context.oauth_metadata, self.context.client_metadata, fallback_base - ) - registration_response = yield registration_request - client_information = await handle_registration_response(registration_response) - check_registration_usable(client_information) - # Only record the issuer when the registration above actually targeted - # the discovered AS — either via its published registration_endpoint, - # or because the resource-origin /register fallback is on the issuer's - # own host (legacy same-origin embedded AS). Otherwise the fallback hit - # a different server and recording a binding to the PRM-advertised AS - # would persist a binding that was never established. - if ( - self.context.oauth_metadata is not None - and discovered_issuer is not None - and ( - self.context.oauth_metadata.registration_endpoint is not None - or self.context.get_authorization_base_url(discovered_issuer) == fallback_base - ) - ): - client_information.issuer = discovered_issuer - self.context.client_info = client_information - await self.context.storage.set_client_info(client_information) - - # Step 5: Perform authorization and complete token exchange - token_response = yield await self._perform_authorization() - await self._handle_token_response(token_response) - except Exception: - logger.exception("OAuth flow error") - raise - - # Retry with new tokens - self._add_auth_header(request) - yield request + outgoing = await followups.__anext__() + while True: + incoming = yield outgoing + try: + outgoing = await followups.asend(incoming) + except StopAsyncIteration: + break + finally: + await followups.aclose() + finally: + if holding: + self.context.lock.release() diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index 18a1566705..f31d47f27a 100644 --- a/tests/client/test_auth.py +++ b/tests/client/test_auth.py @@ -7,6 +7,7 @@ from unittest import mock from urllib.parse import parse_qs, quote, unquote, urlparse +import anyio import httpx2 import pytest from inline_snapshot import Is, snapshot @@ -3667,3 +3668,160 @@ async def echo_callback() -> AuthorizationCodeResult: await auth_flow.asend(httpx2.Response(200, request=final_req)) except StopAsyncIteration: pass + + +def _expired_refreshable(oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken) -> None: + """Token is present and refreshable, but past its expiry. Discovery is already done.""" + oauth_provider._initialized = True + oauth_provider.context.current_tokens = valid_tokens + oauth_provider.context.token_expiry_time = time.time() - 30 + oauth_provider.context.client_info = OAuthClientInformationFull( + client_id="test_client", + client_secret="test_secret", + redirect_uris=[AnyUrl("http://localhost:3030/callback")], + token_endpoint_auth_method="client_secret_post", + ) + + +@pytest.mark.anyio +async def test_expired_token_is_refreshed_then_the_original_request_is_sent( + oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken, mock_storage: MockTokenStorage +): + """A successful refresh stores the new token and the original request goes out with it. + + SDK behaviour: refresh runs outside ``context.lock``, and a 200 refresh response + leaves the provider initialized. + """ + _expired_refreshable(oauth_provider, valid_tokens) + + auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp")) + refresh_request = await auth_flow.__anext__() + + assert refresh_request.method == "POST" + assert "grant_type=refresh_token" in refresh_request.content.decode() + refresh_response = httpx2.Response( + 200, + json={"access_token": "refreshed-access", "token_type": "Bearer", "expires_in": 3600, "refresh_token": "kept"}, + request=refresh_request, + ) + primary = await auth_flow.asend(refresh_response) + + assert primary.method == "GET" + assert primary.headers["Authorization"] == "Bearer refreshed-access" + assert oauth_provider._initialized is True + stored = await mock_storage.get_tokens() + assert stored is not None + assert stored.access_token == "refreshed-access" + try: + await auth_flow.asend(httpx2.Response(200, request=primary)) + except StopAsyncIteration: + pass + + +@pytest.mark.anyio +async def test_failed_refresh_clears_the_token_and_sends_the_request_unauthenticated( + oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken +): + """A non-200 refresh drops the stored token and the original request goes out without it. + + SDK behaviour: refresh failure marks the provider uninitialized so the next call rediscovers. + """ + _expired_refreshable(oauth_provider, valid_tokens) + + auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp")) + refresh_request = await auth_flow.__anext__() + assert refresh_request.method == "POST" + primary = await auth_flow.asend(httpx2.Response(400, json={"error": "invalid_grant"}, request=refresh_request)) + + assert primary.method == "GET" + assert "Authorization" not in primary.headers + assert oauth_provider._initialized is False + assert oauth_provider.context.current_tokens is None + try: + await auth_flow.asend(httpx2.Response(200, request=primary)) + except StopAsyncIteration: + pass + + +@pytest.mark.anyio +async def test_refresh_waiter_sends_the_token_stored_while_it_waited( + oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken +): + """A caller that finds an expired token does not refresh if a peer stored a valid one while it waited. + + The re-check is under ``refresh_lock``. This drives ``async_auth_flow`` and holds that lock + so the wait is real; the public client has no other way to observe the single-flight. + """ + _expired_refreshable(oauth_provider, valid_tokens) + + class HoldableLock: + """Stand-in for ``refresh_lock`` that signals when a waiter arrives.""" + + def __init__(self) -> None: + self._inner = anyio.Lock() + self.waiting = anyio.Event() + self.arm = False + + async def __aenter__(self): + if self.arm: + self.waiting.set() + await self._inner.acquire() + return self + + async def __aexit__(self, *_exc: object) -> None: + self._inner.release() + + lock = HoldableLock() + # anyio.Lock.acquire is read-only, so the stand-in cannot be a subclass patch. + # The field is typed as anyio.Lock; setattr keeps the assignment off the checker. + object.__setattr__(oauth_provider.context, "refresh_lock", lock) + # First entry is before the waiter is armed, so both arms of `__aenter__` run. + await lock.__aenter__() + lock.arm = True + sent: list[httpx2.Request] = [] + + async def drive() -> None: + auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp")) + request = await auth_flow.__anext__() + sent.append(request) + try: + await auth_flow.asend(httpx2.Response(200, request=request)) + except StopAsyncIteration: + pass + + async with anyio.create_task_group() as task_group: + task_group.start_soon(drive) + with anyio.fail_after(5): + await lock.waiting.wait() + oauth_provider.context.token_expiry_time = time.time() + 3600 + await lock.__aexit__(None, None, None) + + assert len(sent) == 1 + assert sent[0].method == "GET" + assert sent[0].headers["Authorization"] == "Bearer test_access_token" + + +@pytest.mark.anyio +async def test_stale_401_retries_once_with_the_token_stored_during_the_request( + oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken +): + """A 401 whose access token was replaced while the request was in flight retries once. + + SDK behaviour: that retry uses the stored token and does not start a second login. + """ + oauth_provider._initialized = True + oauth_provider.context.current_tokens = valid_tokens + + auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp")) + primary = await auth_flow.__anext__() + assert primary.headers["Authorization"] == "Bearer test_access_token" + + valid_tokens.access_token = "replaced-access" + retry = await auth_flow.asend(httpx2.Response(401, request=primary)) + + assert retry.method == "GET" + assert retry.headers["Authorization"] == "Bearer replaced-access" + try: + await auth_flow.asend(httpx2.Response(200, request=retry)) + except StopAsyncIteration: + pass