diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 4e568a1b090..7302b74f15a 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1494,6 +1494,26 @@ async def upsert_mcp_server_oauth_client_credentials( ) +async def create_mcp_server_oauth_client_credentials_if_absent( + prisma_client: PrismaClient, server_id: str, credentials: MCPCredentials +) -> bool: + """Insert a server's DCR client into the server-scoped store only when no row exists yet, so two + first registrations racing each other cannot overwrite one another. Returns False when another + writer already holds the row, in which case the caller adopts that client.""" + from prisma.errors import UniqueViolationError # noqa: PLC0415 # prisma client is generated at build time + + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # noqa: PLC0415 # avoids circular import + + encrypted: Final = encrypt_credentials(credentials=MCPCredentials(**credentials), encryption_key=get_salt_key()) + try: + await _oauth_client_table_actions(prisma_client).create( + data={"server_id": server_id, "credentials": safe_dumps(encrypted)} + ) + except UniqueViolationError: + return False + return True + + def _reencrypt_mcp_credentials_blob( credentials: "str | Mapping[str, object] | None", new_master_key: str ) -> str | None: diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index a209bd88b9e..8876fd2dd87 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -84,6 +84,7 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( CIMD_METADATA_PATH, TOKEN_NO_CACHE_HEADERS, build_upstream_oauth2_token_request, + configured_proxy_base_url, get_cimd_client_id, get_cimd_document_url, get_request_base_url, @@ -752,7 +753,7 @@ def raise_if_not_oauth2(mcp_server: MCPServer) -> None: one against the upstream IdP (the admin UI's browser-only Authorize uses it). The minted token is upstream-audienced and held by the caller; the gateway persists nothing for these modes (``_persist_dcr_client_registration`` skips them unconditionally, so even the admin - Authorize path with ``persist_credentials`` enabled writes nothing to the server row). + Authorize path with ``replace_stale_client`` enabled writes nothing to the server row). """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import with mcp_server_manager at module load UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, @@ -1787,10 +1788,23 @@ async def _persisted_dcr_redirect_uri_is_stale(mcp_server: MCPServer, current_re DcrRegistrationPersistenceResult = Literal["persisted", "reused", "skipped", "failed"] +_UPSTREAM_CLIENT_MANAGEMENT_KEYS: Final = frozenset( + {"client_secret", "client_secret_expires_at", "registration_access_token", "registration_client_uri"} +) + +_UPSTREAM_REGISTRATION: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + + +def _without_upstream_client_management(registration: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in registration.items() if key not in _UPSTREAM_CLIENT_MANAGEMENT_KEYS} async def _persist_dcr_client_registration( - mcp_server: MCPServer, registration_response: object, current_redirect_uri: str + mcp_server: MCPServer, + registration_response: object, + current_redirect_uri: str, + *, + replace_stale_client: bool = True, ) -> DcrRegistrationPersistenceResult: """Persist the dynamically registered OAuth client (RFC 7591) to its single home: the server's ``LiteLLM_MCPServerTable`` row when it has one, otherwise the server-scoped store when the server @@ -1804,6 +1818,12 @@ async def _persist_dcr_client_registration( write that ``client_credentials`` and token exchange already use. Failures are logged, never raised here: the registration endpoint rejects a failed persistence result. + Whoever connects first mints the server's one shared client, admin or not, because the + ``refresh_token`` grant of every later user runs under it. Only a caller with + ``replace_stale_client`` may replace a client whose recorded ``redirect_uris`` miss the + current callback or overwrite the store; everyone else inserts only when the store is still + empty and otherwise adopts whatever a concurrent registration persisted first. + The client-forwarded token modes (``true_passthrough`` / ``oauth_delegate``) are skipped unconditionally: the caller holds the upstream token and the gateway must hold no OAuth client identity for these servers. Persisting here would stamp ``oauth2_flow`` and a @@ -1845,6 +1865,7 @@ async def _persist_dcr_client_registration( from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # avoids circular import McpIdentifierConflict, + create_mcp_server_oauth_client_credentials_if_absent, update_mcp_server, upsert_mcp_server_oauth_client_credentials, ) @@ -1872,6 +1893,11 @@ async def _persist_dcr_client_registration( or (stored.issuer and mcp_server.issuer and stored.issuer != mcp_server.issuer) ): return "failed" + if latest_credentials is not None and not replace_stale_client: + if not _apply_persisted_dcr_credentials(mcp_server, latest_credentials): + return "failed" + await _refresh_persisted_dcr_server(stored) + return "reused" if ( latest_credentials is not None and not _redirect_uri_not_registered(latest_credentials, current_redirect_uri) @@ -1909,11 +1935,19 @@ async def _persist_dcr_client_registration( else "failed" ) if global_mcp_server_manager.is_config_declared_server(mcp_server.server_id): - await upsert_mcp_server_oauth_client_credentials( + if replace_stale_client: + await upsert_mcp_server_oauth_client_credentials( + prisma_client=prisma_client, + server_id=mcp_server.server_id, + credentials=credentials, + ) + elif not await create_mcp_server_oauth_client_credentials_if_absent( prisma_client=prisma_client, server_id=mcp_server.server_id, credentials=credentials, - ) + ): + winner: Final = await _load_store_dcr_credentials(mcp_server, raise_on_error=True) + return "reused" if winner and _apply_persisted_dcr_credentials(mcp_server, winner) else "failed" _apply_persisted_dcr_credentials(mcp_server, _PersistedDcrCredentials.model_validate(credentials)) return "persisted" except Exception as exc: # noqa: BLE001 @@ -2109,7 +2143,7 @@ async def register_client_with_server( response_types: list | None, token_endpoint_auth_method: str | None, fallback_client_id: str | None = None, - persist_credentials: bool = False, + replace_stale_client: bool = False, client_redirect_uris: list[str] | None = None, client_application_type: Literal["native", "web"] | None = None, ): @@ -2131,7 +2165,7 @@ async def register_client_with_server( resolved_server.client_id and client_matches and not ( - persist_credentials + replace_stale_client and resolved_server.registration_url and await _persisted_dcr_redirect_uri_is_stale(resolved_server, current_redirect_uri) ) @@ -2140,7 +2174,7 @@ async def register_client_with_server( if await _reuse_persisted_dcr_client_if_available( resolved_server, - current_redirect_uri=current_redirect_uri if persist_credentials else None, + current_redirect_uri=current_redirect_uri if replace_stale_client else None, ): return dummy_return @@ -2179,18 +2213,36 @@ async def register_client_with_server( detail="redirect_uris is required to register a client with this server", ) - register_data: Final = { - **( - {"application_type": client_application_type} - if bridge_relay and client_application_type is not None - else {} - ), - "client_name": client_name, - "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 ""), - } + shared_client_persistable: Final = replace_stale_client or configured_proxy_base_url() is not None + if not bridge_relay and not shared_client_persistable: + verbose_logger.warning( + "register_client_with_server: PROXY_BASE_URL is not set, so the OAuth client minted for " + "server_id=%s is not kept as the server's shared client and token refresh will need a " + "re-authorization; set PROXY_BASE_URL or connect the server once as a proxy admin.", + resolved_server.server_id, + ) + register_data: Final[Mapping[str, object]] = ( + { + "client_name": "LiteLLM Proxy", + "redirect_uris": [current_redirect_uri], + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + "token_endpoint_auth_method": "none", + } + if shared_client_persistable and not replace_stale_client and not bridge_relay + else { + **( + {"application_type": client_application_type} + if bridge_relay and client_application_type is not None + else {} + ), + "client_name": client_name, + "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 ""), + } + ) response: Final = await _post_dcr_registration( registration_url=registration_url, register_data=register_data, @@ -2199,27 +2251,34 @@ async def register_client_with_server( token_response: Final = response.json() persistence_result: Final = ( - await _persist_dcr_client_registration(resolved_server, token_response, current_redirect_uri) - if persist_credentials and not bridge_relay + await _persist_dcr_client_registration( + resolved_server, token_response, current_redirect_uri, replace_stale_client=replace_stale_client + ) + if shared_client_persistable and not bridge_relay else None ) if persistence_result == "reused": return dummy_return if persistence_result == "failed": raise HTTPException(status_code=503, detail="OAuth client registration could not be saved; retry authorization") - bound_response: Final = ( + shared_client_response: Final[Mapping[str, object]] = ( + _without_upstream_client_management(_UPSTREAM_REGISTRATION.validate_python(token_response)) + if persistence_result == "persisted" and not replace_stale_client and isinstance(token_response, dict) + else token_response + ) + bound_response: Final[Mapping[str, object]] = ( { - **token_response, + **_UPSTREAM_REGISTRATION.validate_python(token_response), "dcr_issuer": resolved_server.issuer, "dcr_server_url": resolved_server.url, "dcr_redirect_uris": [current_redirect_uri], } - if persistence_result == "persisted" - else token_response + if persistence_result == "persisted" and replace_stale_client + else shared_client_response ) - client_response: Final = ( + client_response: Final[Mapping[str, object]] = ( {**bound_response, "redirect_uris": client_facing_redirect_uris} - if client_redirect_uris and not bridge_relay and isinstance(bound_response, dict) + if client_redirect_uris and not bridge_relay else bound_response ) return JSONResponse(client_response) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index b8675f749e6..b551e099aa5 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -132,6 +132,12 @@ def _resolve_proxy_base_url_env() -> str | None: return None +def configured_proxy_base_url() -> str | None: + """The operator-configured public origin (``PROXY_BASE_URL``), or None when the gateway only + knows itself through the current request's Host and forwarded headers.""" + return _resolve_proxy_base_url_env() + + CIMD_METADATA_PATH: Final = "/oauth/client-metadata.json" diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 923c82e5f63..f69d9db08a0 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2454,7 +2454,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), + replace_stale_client=_user_is_full_admin(user_api_key_dict), client_redirect_uris=client_redirect_uris, client_application_type=client_application_type, ) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 2590052f329..ab45ecb4fbc 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -1397,7 +1397,7 @@ async def test_register_client_persists_dcr_client_identity(): grant_types=["authorization_code", "refresh_token"], response_types=["code"], token_endpoint_auth_method="client_secret_basic", - persist_credentials=True, + replace_stale_client=True, ) import json @@ -1425,7 +1425,7 @@ async def test_register_client_persists_dcr_client_identity(): async def _register_persistence_attempted_for_auth_type(auth_type: MCPAuth) -> bool: - """Run register_client_with_server with persist_credentials=True for a server of ``auth_type`` + """Run register_client_with_server with replace_stale_client=True for a server of ``auth_type`` and report whether the DCR result was persisted onto the server row. The client-forwarded token modes must skip the persist even on the admin path: writing it stamps oauth2_flow and a client_id onto a server whose contract is that the gateway stores nothing, which makes a fresh @@ -1490,7 +1490,7 @@ async def _register_persistence_attempted_for_auth_type(auth_type: MCPAuth) -> b grant_types=["authorization_code", "refresh_token"], response_types=["code"], token_endpoint_auth_method="none", - persist_credentials=True, + replace_stale_client=True, ) expected_binding = { @@ -1505,7 +1505,7 @@ async def _register_persistence_attempted_for_auth_type(auth_type: MCPAuth) -> b @pytest.mark.asyncio @pytest.mark.parametrize("auth_type", [MCPAuth.true_passthrough, MCPAuth.oauth_delegate]) async def test_register_client_does_not_persist_for_client_forwarded_modes(auth_type): - """The admin Authorize path (persist_credentials=True) must not write the DCR client onto a + """The admin Authorize path (replace_stale_client=True) must not write the DCR client onto a true_passthrough / oauth_delegate server row: the browser still receives the registration, but the gateway keeps no OAuth client identity for these modes.""" assert await _register_persistence_attempted_for_auth_type(auth_type) is False @@ -1590,7 +1590,7 @@ async def test_register_client_persists_only_to_its_own_row_when_another_server_ grant_types=["authorization_code", "refresh_token"], response_types=["code"], token_endpoint_auth_method="none", - persist_credentials=True, + replace_stale_client=True, ) mock_async_client.post.assert_called_once() @@ -1668,7 +1668,7 @@ async def test_register_client_does_not_clobber_token_url_when_absent(): grant_types=["authorization_code", "refresh_token"], response_types=["code"], token_endpoint_auth_method="none", - persist_credentials=True, + replace_stale_client=True, ) mock_update.assert_called_once() @@ -1743,7 +1743,7 @@ async def test_register_client_reuses_persisted_client_id_for_non_admin_when_reg grant_types=["authorization_code", "refresh_token"], response_types=["code"], token_endpoint_auth_method="none", - persist_credentials=False, + replace_stale_client=False, ) assert response["client_id"] == "remote_server" @@ -1818,7 +1818,7 @@ async def test_register_client_reuse_refreshes_request_server_when_manager_updat grant_types=["authorization_code", "refresh_token"], response_types=["code"], token_endpoint_auth_method="none", - persist_credentials=False, + replace_stale_client=False, ) assert response["client_id"] == "remote_server" @@ -1901,7 +1901,7 @@ async def test_register_client_returns_reused_client_when_concurrent_persist_win grant_types=["authorization_code", "refresh_token"], response_types=["code"], token_endpoint_auth_method="none", - persist_credentials=True, + replace_stale_client=True, ) assert response["client_id"] == "remote_server" @@ -2003,7 +2003,7 @@ async def test_register_client_re_registers_when_persisted_redirect_uri_no_longe grant_types=["authorization_code", "refresh_token"], response_types=["code"], token_endpoint_auth_method="none", - persist_credentials=True, + replace_stale_client=True, ) mock_async_client.post.assert_called_once() @@ -2073,7 +2073,7 @@ async def test_register_client_grandfathers_persisted_client_without_recorded_re grant_types=["authorization_code", "refresh_token"], response_types=["code"], token_endpoint_auth_method="none", - persist_credentials=True, + replace_stale_client=True, ) mock_async_client.post.assert_not_called() @@ -2134,7 +2134,7 @@ async def test_register_client_keeps_persisted_client_when_recorded_redirect_uri grant_types=["authorization_code", "refresh_token"], response_types=["code"], token_endpoint_auth_method="none", - persist_credentials=True, + replace_stale_client=True, ) mock_async_client.post.assert_not_called() @@ -2144,10 +2144,10 @@ async def test_register_client_keeps_persisted_client_when_recorded_redirect_uri @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.""" + """Callers without replace authority (the public register routes and non-admin users) must + keep reusing the persisted client even when its recorded redirect_uris mismatch: a spoofable + forwarded origin must not let any caller re-register and orphan every user's refresh token. + Only the admin path re-registers.""" try: from fastapi import Request @@ -2196,7 +2196,7 @@ async def test_register_client_non_admin_reuses_persisted_client_despite_redirec grant_types=["authorization_code", "refresh_token"], response_types=["code"], token_endpoint_auth_method="none", - persist_credentials=False, + replace_stale_client=False, ) mock_async_client.post.assert_not_called() @@ -2204,6 +2204,101 @@ async def test_register_client_non_admin_reuses_persisted_client_despite_redirec assert response["client_secret"] == "dummy" +@pytest.mark.asyncio +async def test_register_client_without_replace_authority_adopts_concurrently_persisted_client(monkeypatch): + """Two users connect to a fresh config-declared server at once. Both miss the store at reuse + time and both mint a client upstream, but the loser's insert finds the row taken and must + adopt the winner's client (which the winner's browser is already authorizing under) instead of + overwriting it, and is answered with the gateway's dummy client.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + monkeypatch.setenv("PROXY_BASE_URL", "https://proxy.litellm.example") + config_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_response = MagicMock() + mock_response.json.return_value = {"client_id": "loser-client", "token_endpoint_auth_method": "none"} + mock_response.raise_for_status = MagicMock() + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + winner = {"client_id": "winner-client", "redirect_uris": ["https://proxy.litellm.example/callback"]} + mock_upsert = AsyncMock() + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ), + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=AsyncMock(return_value=None)), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=AsyncMock(side_effect=[None, winner]), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.create_mcp_server_oauth_client_credentials_if_absent", + new=AsyncMock(return_value=False), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", + new=mock_upsert, + ), + ): + response = await discoverable_endpoints.register_client_with_server( + request=mock_request, + mcp_server=config_server, + client_name="Claude Code", + grant_types=["authorization_code", "refresh_token"], + response_types=["code"], + token_endpoint_auth_method="none", + replace_stale_client=False, + ) + + mock_async_client.post.assert_called_once() + mock_upsert.assert_not_called() + assert config_server.client_id == "winner-client" + assert response["client_id"] == config_server.server_name + + +@pytest.mark.asyncio +async def test_persist_without_replace_authority_never_overwrites_a_row_client_it_cannot_apply(): + """A DB-declared server already holds a shared DCR client that this process cannot adopt (it + was bound to another server URL). A caller without replace authority must get a failure, not + fall through to update_mcp_server and overwrite the client every signed-in user depends on.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _persist_dcr_client_registration, + ) + + server = _dcr_redirect_test_server(client_id=None).model_copy(update={"url": "https://mine.example/mcp"}) + row = LiteLLM_MCPServerTable.model_validate(server.model_dump(exclude_none=True)) + row.credentials = {"client_id": "someone-elses-client", "dcr_server_url": "https://other.example/mcp"} + mock_update = AsyncMock() + + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=row)), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update), + ): + result = await _persist_dcr_client_registration( + server, {"client_id": "my-client"}, "https://proxy.litellm.example/callback", replace_stale_client=False + ) + + assert result == "failed" + mock_update.assert_not_called() + assert server.client_id is None + + @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 @@ -2275,64 +2370,68 @@ async def test_register_client_reuses_existing_client_id_without_re_dcr(): 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 +def _public_register_rig(minted): + from fastapi import Request + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + request_payload = { + "client_name": "attacker", + "grant_types": ["implicit"], + "response_types": ["token"], + "token_endpoint_auth_method": "client_secret_basic", + } + mock_response = MagicMock() + mock_response.json.return_value = minted + mock_response.raise_for_status = MagicMock() + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + return mock_request, request_payload, mock_async_client + + +GATEWAY_OWNED_DCR_BODY: Final = { + "client_name": "LiteLLM Proxy", + "redirect_uris": ["https://proxy.litellm.example/callback"], + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + "token_endpoint_auth_method": "none", +} + + +@pytest.mark.asyncio +async def test_public_register_route_persists_first_dcr_client_when_origin_is_configured(monkeypatch): + """A DB-declared oauth2 server with no client yet gets its first DCR client from whichever + client connects first, usually a non-admin user's coding tool on the unauthenticated + /{server}/register route. With PROXY_BASE_URL set the gateway knows its own callback, so that + client lands on the server row for every later user's refresh. The gateway owns what it + registers: the caller's grant types, response types, auth method and name are not sent + upstream, so an anonymous caller cannot shape or phish the shared client.""" + try: 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") + monkeypatch.setenv("PROXY_BASE_URL", "https://proxy.litellm.example") 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", - ) + oauth2_server = _dcr_redirect_test_server(client_id=None) 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()) + mock_request, request_payload, mock_async_client = _public_register_rig( + { + "client_id": "minted-client", + "client_secret": "minted-secret", + "registration_access_token": "manage-me", + "registration_client_uri": "https://provider.example/oauth/register/minted-client", + } + ) + persisted_row = LiteLLM_MCPServerTable.model_validate(oauth2_server.model_dump(exclude_none=True)) + mock_update = AsyncMock(return_value=persisted_row) + mock_update_server = AsyncMock() try: with ( @@ -2345,14 +2444,166 @@ async def test_public_register_route_does_not_persist_client_credentials(): 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.get_mcp_server", + new=AsyncMock(return_value=persisted_row), + ), + patch.object(global_mcp_server_manager, "update_server", new=mock_update_server), 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) + response = await register_client(request=mock_request, mcp_server_name=oauth2_server.server_name) finally: global_mcp_server_manager.registry.clear() + assert mock_async_client.post.call_args.kwargs["json"] == GATEWAY_OWNED_DCR_BODY + mock_update.assert_called_once() + assert mock_update.call_args.kwargs["data"].credentials["client_id"] == "minted-client" + mock_update_server.assert_called_once_with(persisted_row) + body = json.loads(response.body.decode("utf-8")) + assert body["client_id"] == "minted-client" + assert not {"client_secret", "registration_access_token", "registration_client_uri", "dcr_server_url"} & set(body) + + +@pytest.mark.asyncio +async def test_public_register_route_persists_first_dcr_client_for_config_server_and_serves_later_callers( + monkeypatch, +): + """The customer-visible shape of the bug: a config.yaml-declared DCR server (no + LiteLLM_MCPServerTable row) whose first connect comes from a non-admin user. The minted client + must reach the server-scoped OAuth client store, so a token refresh has a client_id and a + second connect reuses the same client instead of minting another.""" + try: + 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, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + monkeypatch.setenv("PROXY_BASE_URL", "https://proxy.litellm.example") + global_mcp_server_manager.registry.clear() + config_server = _dcr_redirect_test_server(client_id=None) + global_mcp_server_manager.registry[config_server.server_id] = config_server + mock_request, request_payload, mock_async_client = _public_register_rig( + {"client_id": "minted-client", "token_endpoint_auth_method": "none"} + ) + store = {} + + def insert_if_absent(**kwargs): + if "blob" in store: + return False + store["blob"] = kwargs["credentials"] + return True + + mock_create = AsyncMock(side_effect=insert_if_absent) + mock_upsert = 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, + ), + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=AsyncMock(return_value=None)), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=AsyncMock(side_effect=lambda **kwargs: store.get("blob")), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.create_mcp_server_oauth_client_credentials_if_absent", + new=mock_create, + ), + patch( + "litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", + new=mock_upsert, + ), + ): + first = await register_client(request=mock_request, mcp_server_name=config_server.server_name) + second = await register_client(request=mock_request, mcp_server_name=config_server.server_name) + finally: + global_mcp_server_manager.registry.clear() + + mock_create.assert_called_once() + assert mock_create.call_args.kwargs["server_id"] == config_server.server_id + assert mock_create.call_args.kwargs["credentials"]["client_id"] == "minted-client" + mock_upsert.assert_not_called() + mock_async_client.post.assert_called_once() + assert json.loads(first.body.decode("utf-8"))["client_id"] == "minted-client" + assert second["client_id"] == config_server.server_name + + +@pytest.mark.asyncio +async def test_public_register_route_without_configured_origin_does_not_pin_the_shared_client(monkeypatch): + """Without PROXY_BASE_URL the callback comes from the request's own Host, which an anonymous + caller can forge, so the public route keeps today's behavior: register for this caller with + the body it sent, echo the upstream response, and persist nothing. An admin's Connect (or + setting PROXY_BASE_URL) is what establishes the shared client.""" + try: + 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, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + global_mcp_server_manager.registry.clear() + config_server = _dcr_redirect_test_server(client_id=None) + global_mcp_server_manager.registry[config_server.server_id] = config_server + mock_request, request_payload, mock_async_client = _public_register_rig({"client_id": "caller-client"}) + mock_create = AsyncMock() + mock_upsert = AsyncMock() + mock_update = 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, + ), + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.create_mcp_server_oauth_client_credentials_if_absent", + new=mock_create, + ), + patch( + "litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", + new=mock_upsert, + ), + ): + response = await register_client(request=mock_request, mcp_server_name=config_server.server_name) + finally: + global_mcp_server_manager.registry.clear() + + posted = mock_async_client.post.call_args.kwargs["json"] + assert posted["grant_types"] == request_payload["grant_types"] + assert posted["client_name"] == request_payload["client_name"] + mock_create.assert_not_called() + mock_upsert.assert_not_called() mock_update.assert_not_called() + assert json.loads(response.body.decode("utf-8"))["client_id"] == "caller-client" @pytest.mark.asyncio @@ -5502,7 +5753,7 @@ def test_oauth_authorization_server_404_for_non_bridge_client_forwarded_server() assert exc.value.status_code == 404 -async def _bridge_register_response(server, request_payload, persist_credentials=False): +async def _bridge_register_response(server, request_payload, replace_stale_client=False): from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( register_client_with_server, ) @@ -5540,7 +5791,7 @@ async def _bridge_register_response(server, request_payload, persist_credentials grant_types=request_payload.get("grant_types"), response_types=request_payload.get("response_types"), token_endpoint_auth_method=request_payload.get("token_endpoint_auth_method"), - persist_credentials=persist_credentials, + replace_stale_client=replace_stale_client, client_redirect_uris=request_payload.get("redirect_uris"), ) return response, mock_async_client, mock_persist @@ -5680,7 +5931,7 @@ async def test_register_bridge_relay_never_persists(): _, _, mock_persist = await _bridge_register_response( _bridge_server(), {"client_name": "Claude", "redirect_uris": [_BRIDGE_CLIENT_REDIRECT]}, - persist_credentials=True, + replace_stale_client=True, ) mock_persist.assert_not_called() @@ -13179,14 +13430,16 @@ async def test_failed_registration_preserves_cached_credentials(monkeypatch): with pytest.raises(HTTPException) as error: await endpoints.register_client_with_server( request=request, mcp_server=server, client_name="app", grant_types=["authorization_code"], - response_types=["code"], token_endpoint_auth_method="none", persist_credentials=True, + response_types=["code"], token_endpoint_auth_method="none", replace_stale_client=True, ) assert error.value.status_code == 503 assert "could not be saved" in error.value.detail assert server.client_id == "old-client" assert server.client_secret == "old-secret" register.assert_awaited_once() - persist.assert_awaited_once_with(server, {"client_id": "new-client"}, "https://gateway.example/callback") + persist.assert_awaited_once_with( + server, {"client_id": "new-client"}, "https://gateway.example/callback", replace_stale_client=True + ) @pytest.mark.asyncio @@ -13206,7 +13459,7 @@ async def test_registration_without_database_keeps_client_in_temporary_server(mo request.headers = {} response = await endpoints.register_client_with_server( request=request, mcp_server=server, client_name="app", grant_types=["authorization_code"], - response_types=["code"], token_endpoint_auth_method="none", persist_credentials=True, + response_types=["code"], token_endpoint_auth_method="none", replace_stale_client=True, ) assert response.status_code == 200 assert json.loads(response.body)["client_id"] == "temporary-client" diff --git a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py index 21ff5534e03..941ecd32e14 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -3751,7 +3751,7 @@ class TestTemporaryMCPSessionEndpoints: response_types=["code"], token_endpoint_auth_method="client_secret_basic", fallback_client_id="server-1", - persist_credentials=True, + replace_stale_client=True, client_redirect_uris=None, client_application_type=None, ) @@ -3807,11 +3807,11 @@ class TestTemporaryMCPSessionEndpoints: assert register_mock.await_args.kwargs["client_redirect_uris"] == forwarded @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.""" + async def test_mcp_register_denies_non_admin_the_authority_to_replace_the_shared_client(self): + """A non-admin caller may mint the server's first shared DCR client (any user's connect + has to, or token refresh runs with no client_id), but must never replace an existing one: + register_client_with_server is invoked with replace_stale_client=False. Only a full + PROXY_ADMIN re-registers a client whose redirect_uris went stale.""" from litellm.proxy.management_endpoints.mcp_management_endpoints import ( mcp_register, ) @@ -3850,7 +3850,7 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is register_response - assert register_mock.await_args.kwargs["persist_credentials"] is False + assert register_mock.await_args.kwargs["replace_stale_client"] is False @pytest.mark.asyncio async def test_get_cached_temporary_mcp_server_falls_back_to_redis(self):