From fe9fb0b81a7ad199a8b890d2669b19207400593b Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 3 Jul 2026 10:42:55 -0700 Subject: [PATCH] fix(mcp): persist DCR client_id so interactive OAuth token refresh works (#31912) * fix(mcp): persist DCR client_id so interactive OAuth token refresh works Interactive authorization_code MCP servers register an OAuth client via Dynamic Client Registration (RFC 7591) during the authorize flow, but the minted client_id and the discovered token_url were returned to the caller and never written to the server row. The autonomous refresh_token grant reads client_id, client_secret and token_url off the server, so an expired access token could not be refreshed; the user was bounced back to re-authorize and tools/list returned zero tools Persist the DCR client_id (plus client_secret and token_endpoint_auth_method when the registration returns them) and the discovered token_url onto the server row, reusing the encrypt_credentials write that client_credentials and token exchange already use, then refresh the in-memory registry so the value is live at refresh time. Both the v1 refresher and the v2 AuthorizationCodeRefresher read those same fields, so egress needs no change * fix: reuse persisted MCP DCR clients * fix: reuse persisted MCP DCR clients --------- Co-authored-by: Cursor Agent (cherry picked from commit 15ff389eb488831c1e639d61776d7375f41d826c) --- .../mcp_server/discoverable_endpoints.py | 191 +++++- .../mcp_management_endpoints.py | 1 + .../mcp_server/test_discoverable_endpoints.py | 542 ++++++++++++++++++ .../test_mcp_management_endpoints.py | 47 ++ 4 files changed, 778 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index f7f7e29365b..c9d1d09b4f0 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -2,12 +2,13 @@ import asyncio import html as _html import json import time -from typing import Any, Dict, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import httpx from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse +from pydantic import BaseModel, ValidationError from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( @@ -30,9 +31,12 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.utils import get_server_root_path -from litellm.types.mcp import MCPAuth +from litellm.types.mcp import MCPAuth, MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer +if TYPE_CHECKING: + from litellm.proxy._types import LiteLLM_MCPServerTable + # TTL cache for upstream OAuth metadata fetched from pass-through MCP servers. # Keeps us from hammering the upstream IdP on each discovery request. # Keyed by (server_id, resource_url) → (expires_at_epoch, payload). @@ -509,6 +513,178 @@ async def exchange_token_with_server( return JSONResponse(result, headers=TOKEN_NO_CACHE_HEADERS) +class _DcrClientRegistration(BaseModel): + """RFC 7591 dynamic client registration response, narrowed to the fields the gateway + must persist to authenticate later token-endpoint calls. Extra members are ignored.""" + + client_id: str + client_secret: Optional[str] = None + token_endpoint_auth_method: Optional[str] = None + + +class _PersistedDcrCredentials(BaseModel): + client_id: Optional[str] = None + client_secret: Optional[str] = None + token_endpoint_auth_method: Optional[str] = None + + +def _get_persisted_dcr_credentials(credentials: object) -> Optional[_PersistedDcrCredentials]: + if not credentials: + return None + try: + return ( + _PersistedDcrCredentials.model_validate_json(credentials) + if isinstance(credentials, str) + else _PersistedDcrCredentials.model_validate(credentials) + ) + except ValidationError: + return None + + +def _decrypt_persisted_dcr_credential(value: Optional[str], key: str) -> Optional[str]: + if value is None: + return None + return decrypt_value_helper( + value=value, + key=key, + exception_type="debug", + return_original_value=True, + ) + + +def _apply_persisted_dcr_credentials(mcp_server: MCPServer, credentials: _PersistedDcrCredentials) -> bool: + client_id = _decrypt_persisted_dcr_credential(credentials.client_id, "client_id") + if not client_id: + return False + mcp_server.client_id = client_id + mcp_server.client_secret = _decrypt_persisted_dcr_credential(credentials.client_secret, "client_secret") + mcp_server.token_endpoint_auth_method = credentials.token_endpoint_auth_method + return True + + +async def _get_persisted_mcp_server_with_dcr_client_id( + mcp_server: MCPServer, +) -> Optional[tuple["LiteLLM_MCPServerTable", _PersistedDcrCredentials]]: + from litellm.proxy._experimental.mcp_server.db import get_mcp_server # noqa: PLC0415 + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 + + try: + prisma_client = get_prisma_client_or_throw("Database not connected. Cannot read MCP OAuth client registration.") + persisted_mcp_server = await get_mcp_server( + prisma_client=prisma_client, + server_id=mcp_server.server_id, + ) + except Exception as exc: # noqa: BLE001 + verbose_logger.debug( + "register_client_with_server: failed to read persisted DCR client registration for server_id=%s: %s", + mcp_server.server_id, + exc, + ) + return None + + if persisted_mcp_server is None: + return None + + credentials = _get_persisted_dcr_credentials(persisted_mcp_server.credentials) + if credentials is None or not credentials.client_id: + return None + + return persisted_mcp_server, credentials + + +async def _reuse_persisted_dcr_client_if_available(mcp_server: MCPServer) -> 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 not _apply_persisted_dcr_credentials(mcp_server, credentials): + return False + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 + global_mcp_server_manager, + ) + + try: + await global_mcp_server_manager.update_server(persisted_mcp_server) + except Exception as exc: # noqa: BLE001 + verbose_logger.warning( + "register_client_with_server: failed to refresh persisted DCR client registration for server_id=%s: %s", + mcp_server.server_id, + exc, + ) + return bool(mcp_server.client_id) + + +DcrRegistrationPersistenceResult = Literal["persisted", "reused", "failed"] + + +async def _persist_dcr_client_registration( + mcp_server: MCPServer, registration_response: object +) -> DcrRegistrationPersistenceResult: + """Persist the dynamically registered OAuth client (RFC 7591) onto the MCP server row. + + The interactive authorization_code flow mints a ``client_id`` via Dynamic Client + Registration that discovery cannot re-derive; without persisting it the autonomous + ``refresh_token`` grant has no client identity, so an expired access token forces a + full re-authorization instead of a silent refresh. Mirrors the ``encrypt_credentials`` + write that ``client_credentials`` and token exchange already use. Failures are logged, + never raised: registration still returns to the caller even when persistence fails. + """ + try: + registration = _DcrClientRegistration.model_validate(registration_response) + except ValidationError as exc: + verbose_logger.warning( + "register_client_with_server: DCR response has no usable client_id for server_id=%s; " + "client registration not persisted (%s)", + mcp_server.server_id, + exc, + ) + return "failed" + + if await _reuse_persisted_dcr_client_if_available(mcp_server): + 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 {} + ), + } + + from litellm.proxy._experimental.mcp_server.db import update_mcp_server # noqa: PLC0415 + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 + global_mcp_server_manager, + ) + from litellm.proxy._types import UpdateMCPServerRequest # noqa: PLC0415 + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 + + try: + prisma_client = get_prisma_client_or_throw( + "Database not connected. Cannot persist MCP OAuth client registration." + ) + updated_row = await update_mcp_server( + prisma_client=prisma_client, + data=UpdateMCPServerRequest( + server_id=mcp_server.server_id, + credentials=credentials, + **({"token_url": mcp_server.token_url} if mcp_server.token_url else {}), + ), + touched_by="mcp_oauth_dcr", + ) + await global_mcp_server_manager.update_server(updated_row) + return "persisted" + except Exception as exc: # noqa: BLE001 + verbose_logger.warning( + "register_client_with_server: failed to persist DCR client registration for server_id=%s: %s", + mcp_server.server_id, + exc, + ) + return "failed" + + async def register_client_with_server( request: Request, mcp_server: MCPServer, @@ -517,6 +693,7 @@ async def register_client_with_server( response_types: Optional[list], token_endpoint_auth_method: Optional[str], fallback_client_id: Optional[str] = None, + persist_credentials: bool = False, ): request_base_url = get_request_base_url(request) dummy_return = { @@ -525,7 +702,10 @@ async def register_client_with_server( "redirect_uris": [f"{request_base_url}/callback"], } - if mcp_server.client_id and mcp_server.client_secret: + if mcp_server.client_id: + return dummy_return + + if await _reuse_persisted_dcr_client_if_available(mcp_server): return dummy_return if mcp_server.authorization_url is None: @@ -561,6 +741,11 @@ async def register_client_with_server( token_response = response.json() + if persist_credentials: + persistence_result = await _persist_dcr_client_registration(mcp_server, token_response) + if persistence_result == "reused": + return dummy_return + return JSONResponse(token_response) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index ab9d04a4eb4..9e973f7922f 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1714,6 +1714,7 @@ if MCP_AVAILABLE: response_types=data.get("response_types", []), token_endpoint_auth_method=data.get("token_endpoint_auth_method", ""), fallback_client_id=server_id, + persist_credentials=_user_is_full_admin(user_api_key_dict), ) @router.delete( 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 153d43c24ed..3e08d6f35e6 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 @@ -534,6 +534,548 @@ async def test_register_client_remote_registration_success(): ) +@pytest.mark.asyncio +async def test_register_client_persists_dcr_client_identity(): + """A dynamic client registration (RFC 7591) must persist the issued client_id / + client_secret / token_endpoint_auth_method and the token_url onto the server row so + autonomous refresh can authenticate as the registered client. Without persistence the + minted client_id is discarded and the refresh_token grant has no client 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, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + oauth2_server = MCPServer( + server_id="remote_server", + name="remote_server", + server_name="remote_server", + alias="remote_server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + + 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": "generated-client", + "client_secret": "generated-secret", + "token_endpoint_auth_method": "client_secret_basic", + } + mock_response.raise_for_status = MagicMock() + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + mock_update = 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.update_mcp_server", new=mock_update), + 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="client_secret_basic", + persist_credentials=True, + ) + + import json + + assert response.status_code == 200 + assert json.loads(response.body.decode("utf-8")) == mock_response.json.return_value + + mock_update.assert_called_once() + update_data = mock_update.call_args.kwargs["data"] + assert update_data.server_id == "remote_server" + assert update_data.token_url == "https://provider.example/oauth/token" + 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" + + mock_update_server.assert_called_once() + + +@pytest.mark.asyncio +async def test_register_client_does_not_clobber_token_url_when_absent(): + """When the in-memory server has no token_url, the DCR persist must omit it from the + partial update rather than passing None, so exclude_unset leaves the token_url column + untouched instead of overwriting an existing value with NULL.""" + 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, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + oauth2_server = MCPServer( + server_id="remote_server", + name="remote_server", + server_name="remote_server", + alias="remote_server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url=None, + registration_url="https://provider.example/oauth/register", + ) + + 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": "generated-client"} + mock_response.raise_for_status = MagicMock() + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + mock_update = 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.update_mcp_server", new=mock_update), + patch.object(global_mcp_server_manager, "update_server", new=mock_update_server), + ): + 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_update.assert_called_once() + update_data = mock_update.call_args.kwargs["data"] + assert update_data.credentials["client_id"] == "generated-client" + assert "token_url" not in update_data.model_fields_set + + +@pytest.mark.asyncio +async def test_register_client_reuses_persisted_client_id_for_non_admin_when_registry_is_stale(): + 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, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + oauth2_server = MCPServer( + server_id="remote_server", + name="remote_server", + server_name="remote_server", + alias="remote_server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + + persisted_server = MagicMock() + persisted_server.credentials = {"client_id": "persisted-client"} + mock_get_mcp_server = AsyncMock(return_value=persisted_server) + mock_update_mcp_server = AsyncMock() + mock_update_server = AsyncMock() + mock_async_client = MagicMock() + mock_async_client.post = 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=False, + ) + + assert response["client_id"] == "remote_server" + assert oauth2_server.client_id == "persisted-client" + mock_async_client.post.assert_not_called() + mock_update_mcp_server.assert_not_called() + mock_update_server.assert_called_once_with(persisted_server) + + +@pytest.mark.asyncio +async def test_register_client_reuse_refreshes_request_server_when_manager_update_fails(): + 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, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + oauth2_server = MCPServer( + server_id="remote_server", + name="remote_server", + server_name="remote_server", + alias="remote_server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + + persisted_server = MagicMock() + persisted_server.credentials = { + "client_id": "persisted-client", + "client_secret": "persisted-secret", + "token_endpoint_auth_method": "client_secret_basic", + } + mock_get_mcp_server = AsyncMock(return_value=persisted_server) + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock() + mock_update_server = AsyncMock(side_effect=RuntimeError("registry update failed")) + + 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, + ) + + assert response["client_id"] == "remote_server" + assert oauth2_server.client_id == "persisted-client" + assert oauth2_server.client_secret == "persisted-secret" + assert oauth2_server.token_endpoint_auth_method == "client_secret_basic" + mock_async_client.post.assert_not_called() + mock_update_server.assert_called_once_with(persisted_server) + + +@pytest.mark.asyncio +async def test_register_client_returns_reused_client_when_concurrent_persist_wins(): + 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, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + oauth2_server = MCPServer( + server_id="remote_server", + name="remote_server", + server_name="remote_server", + alias="remote_server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + + 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": "generated-client"} + 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": "persisted-client"} + mock_get_mcp_server = AsyncMock(side_effect=[None, persisted_server]) + mock_update_mcp_server = AsyncMock() + 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, + ) + + assert response["client_id"] == "remote_server" + assert oauth2_server.client_id == "persisted-client" + mock_async_client.post.assert_called_once() + mock_update_mcp_server.assert_not_called() + mock_update_server.assert_called_once_with(persisted_server) + + +@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 + reused, not re-registered, even without a client_secret. A client_id is one-per-application + in OAuth and shared across users; re-minting per authorize would orphan other users' refresh + tokens by overwriting the server's client_id.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + oauth2_server = MCPServer( + server_id="remote_server", + name="remote_server", + server_name="remote_server", + alias="remote_server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="existing-shared-client", + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + + request_payload = { + "client_name": "Litellm Proxy", + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + "token_endpoint_auth_method": "none", + } + + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock() + + try: + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", + new=AsyncMock(return_value=request_payload), + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ), + ): + response = await register_client(request=mock_request, mcp_server_name=oauth2_server.server_name) + finally: + global_mcp_server_manager.registry.clear() + + mock_async_client.post.assert_not_called() + body = response if isinstance(response, dict) else json.loads(response.body.decode("utf-8")) + assert body["client_secret"] == "dummy" + + +@pytest.mark.asyncio +async def test_public_register_route_does_not_persist_client_credentials(): + """The unauthenticated root /register route must not persist the DCR result onto the + server row; only the authenticated management path passes persist_credentials=True. An + external caller could otherwise bind a caller-controlled client (and leak its secret) to + a server that has no client yet.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + oauth2_server = MCPServer( + server_id="remote_server", + name="remote_server", + server_name="remote_server", + alias="remote_server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + + request_payload = { + "client_name": "attacker", + "grant_types": ["authorization_code"], + "response_types": ["code"], + "token_endpoint_auth_method": "none", + } + + mock_response = MagicMock() + mock_response.json.return_value = { + "client_id": "attacker-client", + "client_secret": "attacker-secret", + } + mock_response.raise_for_status = MagicMock() + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + mock_update = AsyncMock(return_value=MagicMock()) + + try: + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", + new=AsyncMock(return_value=request_payload), + ), + 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.object(global_mcp_server_manager, "update_server", new=AsyncMock()), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update), + ): + await register_client(request=mock_request, mcp_server_name=oauth2_server.server_name) + finally: + global_mcp_server_manager.registry.clear() + + mock_update.assert_not_called() + + @pytest.mark.asyncio @pytest.mark.usefixtures("trust_xff") async def test_authorize_endpoint_respects_x_forwarded_proto(): diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index f40904e234d..f87fc1031c0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -2264,8 +2264,55 @@ class TestTemporaryMCPSessionEndpoints: response_types=["code"], token_endpoint_auth_method="client_secret_basic", fallback_client_id="server-1", + persist_credentials=True, ) + @pytest.mark.asyncio + async def test_mcp_register_does_not_persist_for_non_admin(self): + """A non-admin caller (who may have access to a real server) must not persist the DCR + result onto the shared server row. register_client_with_server is invoked with + persist_credentials=False, so user-side registration returns the DCR response without + writing shared client credentials. Only a full PROXY_ADMIN establishes the shared client.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + mcp_register, + ) + + request = MagicMock() + server = generate_mock_mcp_server_config_record(server_id="server-1") + register_response = {"client_id": "generated"} + request_body = { + "client_name": "LiteLLM", + "grant_types": ["authorization_code"], + "response_types": ["code"], + "token_endpoint_auth_method": "client_secret_basic", + } + non_admin_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", + return_value=server, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._read_request_body", + AsyncMock(return_value=request_body), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.register_client_with_server", + AsyncMock(return_value=register_response), + ) as register_mock, + ): + result = await mcp_register( + request=request, + server_id="server-1", + user_api_key_dict=non_admin_auth, + ) + + assert result is register_response + assert register_mock.await_args.kwargs["persist_credentials"] is False + @pytest.mark.asyncio async def test_get_cached_temporary_mcp_server_falls_back_to_redis(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import (