OAuth client: refresh before re-authorizing, and discover before refreshing

docs/client/oauth-clients.md

@@ -81,7 +81,7 @@ The first time `Client` sends a request, the server answers `401`. The provider
 3. **Authorization.** It generates the PKCE pair and a `state`, builds the authorization URL, awaits your `redirect_handler`, then awaits your `callback_handler` for the code.
 4. **Exchange.** It trades the code for an `OAuthToken`, stores it, and replays your original request with `Authorization: Bearer ...`.
 
-After that it is quiet. Tokens come out of storage, an expired access token is refreshed with the refresh token, and only when none of that works does it run the flow again.
+After that it is quiet. Tokens come out of storage, an expired access token is refreshed with the refresh token, and only when none of that works does it run the flow again. That holds across restarts: a new process that finds a refresh token in storage answers the first `401` by rediscovering the authorization server and refreshing, not by sending anyone back to the browser.
 
 You wrote none of it. Two keyword arguments remain (`client_metadata_url` and `validate_resource_url`), and this file needs neither. `client_metadata_url` is the one worth knowing about; it gets its own section below.
 

src/mcp/client/auth/oauth2.py

@@ -374,7 +374,7 @@
         if self.context.client_metadata.redirect_uris is None:
             raise OAuthFlowError("No redirect URIs provided for authorization code grant")  # pragma: no cover
         if not self.context.redirect_handler:
-            raise OAuthFlowError("No redirect handler provided for authorization code grant")  # pragma: no cover
+            raise OAuthFlowError("No redirect handler provided for authorization code grant")
         if not self.context.callback_handler:
             raise OAuthFlowError("No callback handler provided for authorization code grant")  # pragma: no cover
 
@@ -521,6 +521,8 @@
         if response.status_code != 200:
             logger.warning(f"Token refresh failed: {response.status_code}")
             self.context.clear_tokens()
+            # Re-read storage on the next request: the failure may have been transient.
+            self._initialized = False
             return False
 
         try:
@@ -545,8 +547,32 @@
         except ValidationError:  # pragma: no cover
             logger.exception("Invalid refresh response")
             self.context.clear_tokens()
+            self._initialized = False
             return False
 
+    async def _apply_issuer_binding(self, issuer: str) -> bool:
+        """Apply SEP-2352 to the held registration now that the authorization server's issuer is known.
+
+        Credentials bound to another issuer are discarded with their tokens so the flow re-registers.
+        A CIMD record is portable, so it is kept and re-stamped, but tokens it carried over from
+        another issuer (or of unknown origin, when the record is unstamped) are dropped. Returns
+        True when the held state was for a different issuer.
+        """
+        client_info = self.context.client_info
+        if client_info is None:
+            return False
+        if not credentials_match_issuer(client_info, 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()
+            return True
+        if client_info.client_id == self.context.client_metadata_url and client_info.issuer != issuer:
+            self.context.clear_tokens()
+            client_info.issuer = issuer
+            await self.context.storage.set_client_info(client_info)
+            return True
+        return False
+
     async def _initialize(self) -> None:
         """Load stored tokens and client info."""
         self.context.current_tokens = await self.context.storage.get_tokens()
@@ -586,14 +612,15 @@
             # 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
+            # Refresh ahead of the request only when the token endpoint is already known; on a cold
+            # start the request goes out and the 401 branch discovers, then refreshes.
+            if (
+                not self.context.is_token_valid()
+                and self.context.can_refresh_token()
+                and self.context.oauth_metadata is not None
+            ):
+                refresh_response = yield await self._refresh_token()
+                await self._handle_refresh_response(refresh_response)
 
             if self.context.is_token_valid():
                 self._add_auth_header(request)
@@ -632,21 +659,12 @@
                         else:
                             logger.debug(f"Protected resource metadata discovery failed: {url}")
 
-                    # SEP-2352: stored credentials are bound to the issuer that registered them.
-                    # If the authorization server changed, 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 self.context.auth_server_url is not None
-                        and not credentials_match_issuer(
-                            self.context.client_info, self.context.auth_server_url, self.context.client_metadata_url
-                        )
+                    # SEP-2352: stored credentials and tokens belong to the issuer they came from.
+                    if self.context.auth_server_url is not None and await self._apply_issuer_binding(
+                        self.context.auth_server_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.
+                        # rediscovery cannot leak the old endpoints into Steps 4-5.
                         self.context.oauth_metadata = None
 
                     asm_discovery_urls = build_oauth_authorization_server_metadata_discovery_urls(
@@ -671,21 +689,9 @@
                             logger.debug(f"OAuth metadata discovery failed: {url}")
 
                     # SEP-2352: on the legacy no-PRM path the issuer is only known after ASM
-                    # discovery, so re-evaluate the binding here using the discovered metadata
-                    # issuer (mirroring the bound_issuer fallback in Step 4).
-                    if (
-                        self.context.client_info is not None
-                        and self.context.auth_server_url is None
-                        and self.context.oauth_metadata is not None
-                        and not credentials_match_issuer(
-                            self.context.client_info,
-                            str(self.context.oauth_metadata.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()
+                    # discovery (mirroring the bound_issuer fallback in Step 4).
+                    if self.context.auth_server_url is None and self.context.oauth_metadata is not None:
+                        await self._apply_issuer_binding(str(self.context.oauth_metadata.issuer))
 
                     # Step 3: Apply scope selection strategy
                     self.context.client_metadata.scope = get_client_metadata_scopes(
@@ -741,10 +747,18 @@
                                 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)
+                            # Held tokens belong to a previous client and cannot be refreshed by this one.
+                            self.context.clear_tokens()
+
+                    # Step 5: Refresh with the stored refresh token first (RFC 6749 §6); run the full
+                    # authorization only when there is none or the server rejects it.
+                    refreshed = False
+                    if self.context.can_refresh_token():
+                        refresh_response = yield await self._refresh_token()
+                        refreshed = await self._handle_refresh_response(refresh_response)
+                    if not refreshed:
+                        token_response = yield await self._perform_authorization()
+                        await self._handle_token_response(token_response)
                 except Exception:
                     logger.exception("OAuth flow error")
                     raise

tests/client/test_auth.py

@@ -4,7 +4,7 @@
 import json
 import time
 from unittest import mock
-from urllib.parse import parse_qs, quote, unquote, urlparse
+from urllib.parse import parse_qs, parse_qsl, quote, unquote, urlparse
 
 import httpx2
 import pytest
@@ -3253,3 +3253,139 @@ async def echo_callback() -> AuthorizationCodeResult:
         await auth_flow.asend(httpx2.Response(200, request=final_req))
     except StopAsyncIteration:
         pass
+
+
+@pytest.mark.anyio
+async def test_expired_token_is_not_refreshed_ahead_of_the_request_before_metadata_is_discovered(
+    oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
+) -> None:
+    """With no authorization-server metadata yet, an expired token is not refreshed at a guessed endpoint.
+
+    The request goes out unauthenticated instead, so the 401 branch discovers the real token
+    endpoint before the refresh token is presented anywhere (#3240).
+    """
+    oauth_provider.context.current_tokens = valid_tokens
+    oauth_provider.context.token_expiry_time = time.time() - 60
+    oauth_provider.context.client_info = OAuthClientInformationFull(client_id="c", redirect_uris=None)
+    oauth_provider.context.oauth_metadata = None
+    oauth_provider._initialized = True
+
+    request = httpx2.Request("POST", "https://api.example.com/v1/mcp")
+    auth_flow = oauth_provider.async_auth_flow(request)
+    first = await auth_flow.__anext__()
+
+    assert first is request
+    assert "Authorization" not in first.headers
+
+    with pytest.raises(StopAsyncIteration):
+        await auth_flow.asend(httpx2.Response(200, request=request))
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("stamped_issuer", ["https://old-as.example.com", None], ids=["stamped-elsewhere", "unstamped"])
+async def test_cimd_record_is_restamped_and_its_tokens_and_cached_metadata_dropped_when_prm_names_a_new_issuer(
+    client_metadata: OAuthClientMetadata,
+    mock_storage: MockTokenStorage,
+    valid_tokens: OAuthToken,
+    stamped_issuer: str | None,
+) -> None:
+    """SEP-2352 for CIMD: the URL client_id survives an authorization-server change, nothing else does.
+
+    A long-lived provider holds a CIMD record stamped with another issuer (or, from an older store,
+    not stamped at all), tokens of matching provenance, and cached metadata. As soon as PRM names
+    the issuer in use, the tokens and the cached metadata are dropped and the record is re-stamped
+    and persisted, so a failed rediscovery cannot leave old endpoints in play and no refresh token
+    of unconfirmed origin reaches the named server.
+    """
+    cimd_url = "https://client.example.com/.well-known/mcp-client"
+    provider = OAuthClientProvider(
+        server_url="https://api.example.com/v1/mcp",
+        client_metadata=client_metadata,
+        storage=mock_storage,
+        client_metadata_url=cimd_url,
+    )
+    provider.context.client_info = OAuthClientInformationFull(
+        client_id=cimd_url, token_endpoint_auth_method="none", issuer=stamped_issuer
+    )
+    provider.context.current_tokens = valid_tokens
+    provider.context.token_expiry_time = time.time() + 1800
+    provider.context.oauth_metadata = OAuthMetadata(
+        issuer=AnyHttpUrl("https://old-as.example.com"),
+        authorization_endpoint=AnyHttpUrl("https://old-as.example.com/authorize"),
+        token_endpoint=AnyHttpUrl("https://old-as.example.com/token"),
+    )
+    provider._initialized = True
+
+    auth_flow = provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
+    request = await auth_flow.__anext__()
+    prm_req = await auth_flow.asend(httpx2.Response(401, request=request))
+    prm_response = httpx2.Response(
+        200,
+        content=b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://new-as.example.com"]}',
+        request=prm_req,
+    )
+    asm_req = await auth_flow.asend(prm_response)
+
+    assert str(asm_req.url) == "https://new-as.example.com/.well-known/oauth-authorization-server"
+    assert provider.context.current_tokens is None
+    assert provider.context.oauth_metadata is None
+    assert provider.context.client_info is not None
+    assert (provider.context.client_info.client_id, provider.context.client_info.issuer) == (
+        cimd_url,
+        "https://new-as.example.com",
+    )
+    assert mock_storage._client_info is provider.context.client_info
+    await auth_flow.aclose()
+
+
+@pytest.mark.anyio
+async def test_cimd_record_is_restamped_and_its_tokens_dropped_when_only_asm_reveals_a_new_issuer(
+    client_metadata: OAuthClientMetadata, mock_storage: MockTokenStorage, valid_tokens: OAuthToken
+) -> None:
+    """The CIMD rebinding also applies on the legacy no-PRM path, where the issuer is learned from AS metadata.
+
+    PRM discovery 404s, so the issuer only becomes known from the root well-known metadata; it
+    differs from the record's stamp, so the tokens are dropped and the record re-stamped before
+    any refresh could be attempted, and the flow proceeds to authorize rather than refresh.
+    """
+    cimd_url = "https://client.example.com/.well-known/mcp-client"
+    provider = OAuthClientProvider(
+        server_url="https://api.example.com/v1/mcp",
+        client_metadata=client_metadata,
+        storage=mock_storage,
+        client_metadata_url=cimd_url,
+    )
+    provider.context.client_info = OAuthClientInformationFull(
+        client_id=cimd_url, token_endpoint_auth_method="none", issuer="https://old-as.example.com"
+    )
+    provider.context.current_tokens = valid_tokens
+    provider.context.token_expiry_time = time.time() + 1800
+    provider._initialized = True
+    provider._perform_authorization_code_grant = mock.AsyncMock(return_value=("auth-code", "verifier"))
+
+    auth_flow = provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
+    request = await auth_flow.__anext__()
+    prm_req = await auth_flow.asend(httpx2.Response(401, request=request))
+    prm_req = await auth_flow.asend(httpx2.Response(404, request=prm_req))
+    asm_req = await auth_flow.asend(httpx2.Response(404, request=prm_req))
+    assert str(asm_req.url) == "https://api.example.com/.well-known/oauth-authorization-server"
+    asm_response = httpx2.Response(
+        200,
+        content=(
+            b'{"issuer": "https://api.example.com", '
+            b'"authorization_endpoint": "https://api.example.com/authorize", '
+            b'"token_endpoint": "https://api.example.com/token", '
+            b'"client_id_metadata_document_supported": true}'
+        ),
+        request=asm_req,
+    )
+    next_req = await auth_flow.asend(asm_response)
+
+    assert dict(parse_qsl(next_req.content.decode()))["grant_type"] == "authorization_code"
+    assert provider.context.client_info is not None
+    assert (provider.context.client_info.client_id, provider.context.client_info.issuer) == (
+        cimd_url,
+        "https://api.example.com",
+    )
+    assert mock_storage._client_info is provider.context.client_info
+    await auth_flow.aclose()

tests/interaction/_requirements.py

@@ -3883,6 +3883,29 @@ def __post_init__(self) -> None:
         transports=("streamable-http",),
         note="OAuth is HTTP-only.",
     ),
+    "client-auth:refresh:on-401": Requirement(
+        source="issue:#3250",
+        behavior=(
+            "A 401 received while a refresh token is held is answered, after rediscovery, with a "
+            "refresh_token grant before any interactive authorization, so a client constructed over "
+            "persisted tokens and client registration recovers from an expired access token headlessly."
+        ),
+        transports=("streamable-http",),
+        note="OAuth is HTTP-only. RFC 6749 §1.5 (E)-(H); matches the TypeScript, C# and Rust SDKs.",
+    ),
+    "client-auth:refresh:discovered-endpoint": Requirement(
+        source="issue:#3240",
+        behavior=(
+            "A refresh in a process that has not yet discovered the authorization server happens only after "
+            "protected-resource and authorization-server metadata discovery and posts to the token endpoint "
+            "that metadata advertises."
+        ),
+        transports=("streamable-http",),
+        note=(
+            "OAuth is HTTP-only. When discovery yields no AS metadata at all, the 2025-03-26 origin-derived "
+            "fallback endpoint is still used, as it is for the authorization itself."
+        ),
+    ),
     "client-auth:resource-parameter": Requirement(
         source=f"{SPEC_BASE_URL}/basic/authorization#resource-parameter-implementation",
         behavior=(