diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 9689a7cb047..67018915107 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -820,6 +820,22 @@ class _PersistedDcrCredentials(BaseModel): client_id: Optional[str] = None client_secret: Optional[str] = None token_endpoint_auth_method: Optional[str] = None + redirect_uris: Optional[list[str]] = None + + +def _redirect_uri_not_registered(credentials: _PersistedDcrCredentials, current_redirect_uri: str) -> bool: + """Whether a persisted DCR client is positively known NOT to cover the current callback. + + A DCR client is bound to the redirect_uris it was registered with; if the proxy's + resolved public origin has since changed, every authorize built for it will be + rejected by the IdP. Clients persisted before ``redirect_uris`` was recorded (and + admin-configured clients, which never get a recording) return False so they are + grandfathered rather than re-registered, because re-minting a client_id orphans + every user's refresh tokens for that server.""" + recorded = credentials.redirect_uris + if not recorded: + return False + return current_redirect_uri not in recorded def _get_persisted_dcr_credentials(credentials: object) -> Optional[_PersistedDcrCredentials]: @@ -886,11 +902,22 @@ async def _get_persisted_mcp_server_with_dcr_client_id( return persisted_mcp_server, credentials -async def _reuse_persisted_dcr_client_if_available(mcp_server: MCPServer) -> bool: +async def _reuse_persisted_dcr_client_if_available( + mcp_server: MCPServer, current_redirect_uri: Optional[str] = None +) -> bool: persisted = await _get_persisted_mcp_server_with_dcr_client_id(mcp_server) if persisted is None: return False persisted_mcp_server, credentials = persisted + if current_redirect_uri is not None and _redirect_uri_not_registered(credentials, current_redirect_uri): + verbose_logger.warning( + "register_client_with_server: not reusing persisted DCR client for server_id=%s; its registered " + "redirect_uris=%s do not include the current callback %s", + mcp_server.server_id, + credentials.redirect_uris, + current_redirect_uri, + ) + return False if not _apply_persisted_dcr_credentials(mcp_server, credentials): return False @@ -909,11 +936,36 @@ async def _reuse_persisted_dcr_client_if_available(mcp_server: MCPServer) -> boo return bool(mcp_server.client_id) +async def _persisted_dcr_redirect_uri_is_stale(mcp_server: MCPServer, current_redirect_uri: str) -> bool: + """Whether the server's persisted DCR client is bound to redirect_uris that no longer + cover the current proxy callback, meaning authorize is guaranteed to fail IdP-side. + + Consulted when the in-memory server already carries a hydrated client_id, which + otherwise short-circuits registration before any redirect check can run. Servers + without a persisted DCR recording (admin-configured client_id, or registered before + redirect_uris were recorded) are never reported stale.""" + persisted = await _get_persisted_mcp_server_with_dcr_client_id(mcp_server) + if persisted is None: + return False + _, credentials = persisted + if not _redirect_uri_not_registered(credentials, current_redirect_uri): + return False + verbose_logger.warning( + "register_client_with_server: persisted DCR client for server_id=%s is registered with redirect_uris=%s " + "which do not include the current callback %s (proxy origin changed); registering a replacement client. " + "Users previously signed in to this server will need to re-authenticate.", + mcp_server.server_id, + credentials.redirect_uris, + current_redirect_uri, + ) + return True + + DcrRegistrationPersistenceResult = Literal["persisted", "reused", "skipped", "failed"] async def _persist_dcr_client_registration( - mcp_server: MCPServer, registration_response: object + mcp_server: MCPServer, registration_response: object, current_redirect_uri: str ) -> DcrRegistrationPersistenceResult: """Persist the dynamically registered OAuth client (RFC 7591) onto the MCP server row. @@ -929,6 +981,13 @@ async def _persist_dcr_client_registration( client identity for these servers. Persisting here would stamp ``oauth2_flow`` and a ``client_id`` onto a server whose mode promises the gateway stores nothing, making a fresh pass-through server read as gateway-authorized. + + ``redirect_uris`` records what the client is bound to so a later origin change can be + detected as a positive mismatch and trigger re-registration instead of stranding the + server on IdP-side redirect_uri rejections. ``client_secret`` and + ``token_endpoint_auth_method`` are written explicitly (None when absent) because + ``update_mcp_server`` merges credential blobs: a re-registered public client must not + inherit the previous client's secret or auth method. """ if mcp_server.is_true_passthrough or mcp_server.is_oauth_delegate: return "skipped" @@ -944,17 +1003,16 @@ async def _persist_dcr_client_registration( ) return "failed" - if await _reuse_persisted_dcr_client_if_available(mcp_server): + if await _reuse_persisted_dcr_client_if_available(mcp_server, current_redirect_uri=current_redirect_uri): return "reused" credentials: MCPCredentials = { "client_id": registration.client_id, - **({"client_secret": registration.client_secret} if registration.client_secret is not None else {}), - **( - {"token_endpoint_auth_method": "client_secret_basic"} - if registration.token_endpoint_auth_method == "client_secret_basic" - else {} + "client_secret": registration.client_secret, + "token_endpoint_auth_method": ( + "client_secret_basic" if registration.token_endpoint_auth_method == "client_secret_basic" else None ), + "redirect_uris": [current_redirect_uri], } from litellm.proxy._experimental.mcp_server.db import update_mcp_server # noqa: PLC0415 @@ -1017,16 +1075,24 @@ async def register_client_with_server( ): _raise_if_not_oauth2(mcp_server) request_base_url = get_request_base_url(request) + current_redirect_uri = f"{request_base_url}/callback" dummy_return = { "client_id": fallback_client_id or mcp_server.server_name, "client_secret": "dummy", - "redirect_uris": [f"{request_base_url}/callback"], + "redirect_uris": [current_redirect_uri], } - if mcp_server.client_id: + if mcp_server.client_id and not ( + persist_credentials + and mcp_server.registration_url + and await _persisted_dcr_redirect_uri_is_stale(mcp_server, current_redirect_uri) + ): return dummy_return - if await _reuse_persisted_dcr_client_if_available(mcp_server): + if await _reuse_persisted_dcr_client_if_available( + mcp_server, + current_redirect_uri=current_redirect_uri if persist_credentials else None, + ): return dummy_return if mcp_server.authorization_url is None: @@ -1044,7 +1110,7 @@ async def register_client_with_server( register_data = { "client_name": client_name, - "redirect_uris": client_redirect_uris if bridge_relay else [f"{request_base_url}/callback"], + "redirect_uris": client_redirect_uris if bridge_relay else [current_redirect_uri], "grant_types": grant_types or (["authorization_code", "refresh_token"] if bridge_relay else []), "response_types": response_types or (["code"] if bridge_relay else []), "token_endpoint_auth_method": token_endpoint_auth_method or ("none" if bridge_relay else ""), @@ -1072,7 +1138,7 @@ async def register_client_with_server( token_response = response.json() if persist_credentials and not bridge_relay: - persistence_result = await _persist_dcr_client_registration(mcp_server, token_response) + persistence_result = await _persist_dcr_client_registration(mcp_server, token_response, current_redirect_uri) if persistence_result == "reused": return dummy_return diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index d273e8ec4db..ac411ad9d9a 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -165,6 +165,15 @@ class MCPCredentials(TypedDict, total=False): sends HTTP Basic; defaults to "client_secret_post" when unset. """ + redirect_uris: Optional[List[str]] + """ + The redirect URIs a dynamically registered (RFC 7591) OAuth client was bound to at + registration time. Lets a later registration detect that the proxy's public origin no + longer matches the registered callback and re-register instead of reusing a client the + IdP will reject. Absent for admin-configured clients and for clients registered before + this field existed. Not a secret; stored unencrypted. + """ + token_exchange_profile: Optional[str] """ Token exchange wire dialect: "rfc8693" (default, the standard token-exchange grant) or diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 11a6fea0cea..bf450cf0cfc 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -660,6 +660,7 @@ async def test_register_client_persists_dcr_client_identity(): assert update_data.credentials["client_id"] == "generated-client" assert update_data.credentials["client_secret"] == "generated-secret" assert update_data.credentials["token_endpoint_auth_method"] == "client_secret_basic" + assert update_data.credentials["redirect_uris"] == ["https://proxy.litellm.example/callback"] assert update_data.oauth2_flow == "authorization_code" mock_update_server.assert_called_once() @@ -1141,6 +1142,296 @@ async def test_register_client_returns_reused_client_when_concurrent_persist_win mock_update_server.assert_called_once_with(persisted_server) +def _dcr_redirect_test_server(client_id): + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + return MCPServer( + server_id="remote_server", + name="remote_server", + server_name="remote_server", + alias="remote_server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=client_id, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + + +@pytest.mark.asyncio +async def test_register_client_re_registers_when_persisted_redirect_uri_no_longer_matches_origin(): + """A persisted DCR client is bound to the redirect_uri it was registered with. When the + proxy's resolved public origin changes, every authorize built for the reused client is + rejected IdP-side and the server is permanently stranded (GH #32473). A positive mismatch + between the recorded redirect_uris and the current callback must therefore re-register on + the admin path and persist the replacement client, with the new binding recorded and the + old client's secret/auth method cleared rather than merged into the new identity.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client_with_server, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + oauth2_server = _dcr_redirect_test_server(client_id="stale-client") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + + mock_response = MagicMock() + mock_response.json.return_value = { + "client_id": "fresh-client", + "redirect_uris": ["https://proxy.litellm.example/callback"], + } + mock_response.raise_for_status = MagicMock() + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + persisted_server = MagicMock() + persisted_server.credentials = { + "client_id": "stale-client", + "client_secret": "stale-secret", + "token_endpoint_auth_method": "client_secret_basic", + "redirect_uris": ["https://old-origin.example/callback"], + } + mock_get_mcp_server = AsyncMock(return_value=persisted_server) + mock_update_mcp_server = AsyncMock(return_value=MagicMock()) + mock_update_server = AsyncMock() + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=mock_get_mcp_server, + ), + patch( + "litellm.proxy._experimental.mcp_server.db.update_mcp_server", + new=mock_update_mcp_server, + ), + patch.object(global_mcp_server_manager, "update_server", new=mock_update_server), + ): + response = await register_client_with_server( + request=mock_request, + mcp_server=oauth2_server, + client_name="Litellm Proxy", + grant_types=["authorization_code", "refresh_token"], + response_types=["code"], + token_endpoint_auth_method="none", + persist_credentials=True, + ) + + mock_async_client.post.assert_called_once() + register_payload = mock_async_client.post.call_args.kwargs["json"] + assert register_payload["redirect_uris"] == ["https://proxy.litellm.example/callback"] + + mock_update_mcp_server.assert_called_once() + update_data = mock_update_mcp_server.call_args.kwargs["data"] + assert update_data.credentials["client_id"] == "fresh-client" + assert update_data.credentials["redirect_uris"] == ["https://proxy.litellm.example/callback"] + assert update_data.credentials["client_secret"] is None + assert update_data.credentials["token_endpoint_auth_method"] is None + + assert response.status_code == 200 + assert json.loads(response.body.decode("utf-8"))["client_id"] == "fresh-client" + + +@pytest.mark.asyncio +async def test_register_client_grandfathers_persisted_client_without_recorded_redirect_uris(): + """Clients persisted before redirect_uris were recorded (and admin-configured clients, + which never get a recording) have nothing to compare against; treating that as a mismatch + would re-mint a client_id for every existing install on upgrade and orphan all users' + refresh tokens for those servers. A missing recording must read as a match: no DCR call, + no persistence write, existing client returned.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client_with_server, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + oauth2_server = _dcr_redirect_test_server(client_id="legacy-client") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock() + + persisted_server = MagicMock() + persisted_server.credentials = {"client_id": "legacy-client"} + mock_get_mcp_server = AsyncMock(return_value=persisted_server) + mock_update_mcp_server = AsyncMock() + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=mock_get_mcp_server, + ), + patch( + "litellm.proxy._experimental.mcp_server.db.update_mcp_server", + new=mock_update_mcp_server, + ), + ): + response = await register_client_with_server( + request=mock_request, + mcp_server=oauth2_server, + client_name="Litellm Proxy", + grant_types=["authorization_code", "refresh_token"], + response_types=["code"], + token_endpoint_auth_method="none", + persist_credentials=True, + ) + + mock_async_client.post.assert_not_called() + mock_update_mcp_server.assert_not_called() + assert response["client_secret"] == "dummy" + assert oauth2_server.client_id == "legacy-client" + + +@pytest.mark.asyncio +async def test_register_client_keeps_persisted_client_when_recorded_redirect_uri_matches_origin(): + """When the recorded redirect_uris still cover the current callback the persisted client + is valid; re-registering would orphan refresh tokens for no reason.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client_with_server, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + oauth2_server = _dcr_redirect_test_server(client_id="kept-client") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock() + + persisted_server = MagicMock() + persisted_server.credentials = { + "client_id": "kept-client", + "redirect_uris": ["https://proxy.litellm.example/callback"], + } + mock_get_mcp_server = AsyncMock(return_value=persisted_server) + mock_update_mcp_server = AsyncMock() + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=mock_get_mcp_server, + ), + patch( + "litellm.proxy._experimental.mcp_server.db.update_mcp_server", + new=mock_update_mcp_server, + ), + ): + response = await register_client_with_server( + request=mock_request, + mcp_server=oauth2_server, + client_name="Litellm Proxy", + grant_types=["authorization_code", "refresh_token"], + response_types=["code"], + token_endpoint_auth_method="none", + persist_credentials=True, + ) + + mock_async_client.post.assert_not_called() + mock_update_mcp_server.assert_not_called() + assert response["client_secret"] == "dummy" + + +@pytest.mark.asyncio +async def test_register_client_non_admin_reuses_persisted_client_despite_redirect_mismatch(): + """Non-persisting callers (the public register routes and non-admin users) must keep + today's reuse behavior even when the recorded redirect_uris mismatch: re-registering + without persistence would mint an orphan upstream client on every connect while the + stored client keeps being used at authorize time. Only the admin path re-registers.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client_with_server, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + oauth2_server = _dcr_redirect_test_server(client_id=None) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock() + + persisted_server = MagicMock() + persisted_server.credentials = { + "client_id": "persisted-client", + "redirect_uris": ["https://old-origin.example/callback"], + } + mock_get_mcp_server = AsyncMock(return_value=persisted_server) + mock_update_server = AsyncMock() + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=mock_get_mcp_server, + ), + patch.object(global_mcp_server_manager, "update_server", new=mock_update_server), + ): + response = await register_client_with_server( + request=mock_request, + mcp_server=oauth2_server, + client_name="Litellm Proxy", + grant_types=["authorization_code", "refresh_token"], + response_types=["code"], + token_endpoint_auth_method="none", + persist_credentials=False, + ) + + mock_async_client.post.assert_not_called() + assert oauth2_server.client_id == "persisted-client" + assert response["client_secret"] == "dummy" + + @pytest.mark.asyncio async def test_register_client_reuses_existing_client_id_without_re_dcr(): """A server that already has a client_id (admin-configured or previously DCR'd) must be