Handle the CIMD issuer change beside the SEP-2352 guard so stale metadata is dropped too

The CIMD token drop ran after authorization-server metadata rediscovery,
so when that rediscovery failed the previous server's cached metadata
survived under a record already re-stamped with the new issuer. Doing it
at the same point as the existing bound-credentials guard, right after PRM
names the issuer, drops the cached metadata as that guard does. A unit
test pins all three effects.

Author: maxisbey

src/mcp/client/auth/oauth2.py

@@ -652,6 +652,18 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
                         # 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
+                    elif (
+                        self.context.client_info is not None
+                        and self.context.client_info.client_id == self.context.client_metadata_url
+                        and self.context.auth_server_url is not None
+                        and self.context.client_info.issuer not in (None, self.context.auth_server_url)
+                    ):
+                        # A CIMD client_id is portable across authorization servers; the tokens issued
+                        # under it and the cached metadata are not. Keep the record, re-stamped.
+                        self.context.clear_tokens()
+                        self.context.oauth_metadata = None
+                        self.context.client_info.issuer = self.context.auth_server_url
+                        await self.context.storage.set_client_info(self.context.client_info)
 
                     asm_discovery_urls = build_oauth_authorization_server_metadata_discovery_urls(
                         self.context.auth_server_url, self.context.server_url
@@ -748,21 +760,6 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
                             # Held tokens belong to a previous client and cannot be refreshed by this one.
                             self.context.clear_tokens()
 
-                    # A CIMD client_id is portable across authorization servers (SEP-2352) but tokens
-                    # issued under it are not: on an issuer change keep the record, drop the tokens.
-                    client_info = self.context.client_info
-                    current_issuer = self.context.auth_server_url or (
-                        str(self.context.oauth_metadata.issuer) if self.context.oauth_metadata else None
-                    )
-                    if (
-                        client_info.client_id == self.context.client_metadata_url
-                        and current_issuer is not None
-                        and client_info.issuer not in (None, current_issuer)
-                    ):
-                        self.context.clear_tokens()
-                        client_info.issuer = current_issuer
-                        await self.context.storage.set_client_info(client_info)
-
                     # 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

tests/client/test_auth.py

@@ -3279,3 +3279,55 @@ async def test_expired_token_is_not_refreshed_ahead_of_the_request_before_metada
 
     with pytest.raises(StopAsyncIteration):
         await auth_flow.asend(httpx2.Response(200, request=request))
+
+
+@pytest.mark.anyio
+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
+) -> 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 the old issuer, tokens minted there, and
+    the old server's cached metadata. As soon as PRM names a different issuer, the tokens and the
+    cached metadata are dropped and the record is re-stamped and persisted, so a failed
+    rediscovery cannot leave the old endpoints in play and no refresh reaches the new 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="https://old-as.example.com"
+    )
+    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()