From 20cf6ba70bbb3ef90731899d0543da628bc653c1 Mon Sep 17 00:00:00 2001 From: mrinal Date: Tue, 3 Nov 2026 16:36:57 +0000 Subject: [PATCH 1/7] feat(mcp): add provider registration configuration and management Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 145 ++++++++----- litellm/proxy/_types.py | 51 ++++- .../mcp_management_endpoints.py | 126 ++++++++--- litellm/proxy/mcp_server/catalog.py | 99 ++++++++- litellm/proxy/mcp_server/db.py | 129 +++++++++-- .../proxy/mcp_server/mcp_server_manager.py | 88 ++++++-- litellm/proxy/mcp_server/rest_endpoints.py | 3 + litellm/proxy/proxy_server.py | 19 +- litellm/types/mcp.py | 5 +- .../types/mcp_server/mcp_server_manager.py | 3 +- .../types/mcp_server/provider_registration.py | 129 +++++++++++ .../test_mcp_management_endpoints.py | 6 + tests/unit/proxy/mcp_server/test_catalog.py | 201 ++++++++++++++++++ .../mcp_server/test_mcp_partial_update.py | 84 +++++++- .../proxy/mcp_server/test_rest_endpoints.py | 13 ++ .../proxy/proxy_server/test_proxy_config.py | 31 +++ tests/unit/proxy/test__types.py | 18 +- .../mcp_server/test_provider_registration.py | 96 +++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 14 +- 19 files changed, 1106 insertions(+), 154 deletions(-) create mode 100644 litellm/types/mcp_server/provider_registration.py create mode 100644 tests/unit/types/mcp_server/test_provider_registration.py diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 42bcf3eb61c..3d9265f1761 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -37315,6 +37315,14 @@ "title": "MCPEnvVarScope", "type": "string" }, + "MCPExecutionMode": { + "enum": [ + "litellm", + "provider" + ], + "title": "MCPExecutionMode", + "type": "string" + }, "NewMCPServerRequest": { "properties": { "alias": { @@ -37506,10 +37514,8 @@ "title": "Env Vars" }, "execution_mode": { - "const": "litellm", - "default": "litellm", - "title": "Execution Mode", - "type": "string" + "$ref": "#/components/schemas/MCPExecutionMode", + "default": "litellm" }, "extra_headers": { "anyOf": [ @@ -37797,14 +37803,21 @@ "title": "Tool Name To Display Name" }, "transport": { - "default": "sse", - "enum": [ - "sse", - "http", - "stdio" + "anyOf": [ + { + "enum": [ + "sse", + "http", + "stdio" + ], + "type": "string" + }, + { + "type": "null" + } ], - "title": "Transport", - "type": "string" + "default": "sse", + "title": "Transport" }, "url": { "anyOf": [ @@ -40359,10 +40372,8 @@ "title": "Env Vars" }, "execution_mode": { - "const": "litellm", - "default": "litellm", - "title": "Execution Mode", - "type": "string" + "$ref": "#/components/schemas/MCPExecutionMode", + "default": "litellm" }, "extra_headers": { "anyOf": [ @@ -40650,14 +40661,21 @@ "title": "Tool Name To Display Name" }, "transport": { - "default": "sse", - "enum": [ - "sse", - "http", - "stdio" + "anyOf": [ + { + "enum": [ + "sse", + "http", + "stdio" + ], + "type": "string" + }, + { + "type": "null" + } ], - "title": "Transport", - "type": "string" + "default": "sse", + "title": "Transport" }, "url": { "anyOf": [ @@ -44427,10 +44445,8 @@ "title": "Env Vars" }, "execution_mode": { - "const": "litellm", - "default": "litellm", - "title": "Execution Mode", - "type": "string" + "$ref": "#/components/schemas/MCPExecutionMode", + "default": "litellm" }, "extra_headers": { "anyOf": [ @@ -44718,14 +44734,21 @@ "title": "Tool Name To Display Name" }, "transport": { - "default": "sse", - "enum": [ - "sse", - "http", - "stdio" + "anyOf": [ + { + "enum": [ + "sse", + "http", + "stdio" + ], + "type": "string" + }, + { + "type": "null" + } ], - "title": "Transport", - "type": "string" + "default": "sse", + "title": "Transport" }, "url": { "anyOf": [ @@ -45241,14 +45264,21 @@ "title": "Tool Name To Display Name" }, "transport": { - "default": "sse", - "enum": [ - "sse", - "http", - "stdio" + "anyOf": [ + { + "enum": [ + "sse", + "http", + "stdio" + ], + "type": "string" + }, + { + "type": "null" + } ], - "title": "Transport", - "type": "string" + "default": "sse", + "title": "Transport" }, "url": { "anyOf": [ @@ -47706,6 +47736,14 @@ "title": "MCPEnvVarScope", "type": "string" }, + "MCPExecutionMode": { + "enum": [ + "litellm", + "provider" + ], + "title": "MCPExecutionMode", + "type": "string" + }, "NewMCPServerRequest": { "properties": { "alias": { @@ -47897,10 +47935,8 @@ "title": "Env Vars" }, "execution_mode": { - "const": "litellm", - "default": "litellm", - "title": "Execution Mode", - "type": "string" + "$ref": "#/components/schemas/MCPExecutionMode", + "default": "litellm" }, "extra_headers": { "anyOf": [ @@ -48188,14 +48224,21 @@ "title": "Tool Name To Display Name" }, "transport": { - "default": "sse", - "enum": [ - "sse", - "http", - "stdio" + "anyOf": [ + { + "enum": [ + "sse", + "http", + "stdio" + ], + "type": "string" + }, + { + "type": "null" + } ], - "title": "Transport", - "type": "string" + "default": "sse", + "title": "Transport" }, "url": { "anyOf": [ diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 88b06ee1a97..c1d4eff0363 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2,6 +2,7 @@ import enum import json import os from collections.abc import Callable, Mapping, Sequence +from collections.abc import Set as AbstractSet from datetime import datetime from types import MappingProxyType as MappingProxyType # noqa: PLC0414 # public re-export from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple, TypeAlias @@ -15,7 +16,9 @@ from pydantic import ( Field, Json, JsonValue, + ModelWrapValidatorHandler, PositiveInt, + PrivateAttr, TypeAdapter, field_validator, model_validator, @@ -65,13 +68,13 @@ from litellm.types.mcp import ( MCPCredentials, MCPDisabledCapabilities, MCPExecutionMode, - MCPGatewayExecutionMode, MCPTransport, MCPTransportType, MCPUpstreamProtocol, validate_mcp_protocol_transport, ) from litellm.types.mcp_server.mcp_server_manager import MCPInfo +from litellm.types.mcp_server.provider_registration import active_registration_extras, normalize_provider_registration from litellm.types.proxy.agent_identity import ManagedAgentContext from litellm.types.proxy.auth.special_headers import ( SpecialHeaders as SpecialHeaders, # noqa: PLC0414 # public re-export @@ -1388,12 +1391,12 @@ def _validate_mcp_transport_fields(values: object) -> None: class NewMCPServerRequest(LiteLLMPydanticObjectBase): - execution_mode: MCPGatewayExecutionMode = MCPExecutionMode.litellm + execution_mode: MCPExecutionMode = MCPExecutionMode.litellm server_id: str | None = None server_name: str | None = None alias: str | None = None description: str | None = None - transport: MCPTransportType = MCPTransport.sse + transport: MCPTransportType | None = MCPTransport.sse auth_type: MCPAuthType | None = None credentials: MCPCredentials | None = None url: str | None = None @@ -1454,6 +1457,10 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): @model_validator(mode="after") def validate_protocol_transport(self) -> "NewMCPServerRequest": + if self.execution_mode == MCPExecutionMode.provider: + return self + if self.transport is None: + raise ValueError("LiteLLM-executed MCP registrations require a transport") validate_mcp_protocol_transport( TypeAdapter[MCPUpstreamProtocol](MCPUpstreamProtocol).validate_python( (self.mcp_info or {}).get("protocol_version", "auto") @@ -1465,7 +1472,15 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): @model_validator(mode="before") @classmethod def validate_transport_fields(cls, values): - _validate_mcp_transport_fields(values) + if not isinstance(values, dict) or values.get("execution_mode") != MCPExecutionMode.provider: + _validate_mcp_transport_fields(values) + return values + + @model_validator(mode="before") + @classmethod + def normalize_registration(cls, values: object) -> object: + if isinstance(values, dict) and values.get("execution_mode") == MCPExecutionMode.provider: + return normalize_provider_registration(TypeAdapter(dict[str, object]).validate_python(values)) return values @model_validator(mode="before") @@ -1497,11 +1512,13 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): + _provider_extra_fields: tuple[str, ...] = PrivateAttr(default=()) + server_id: str server_name: str | None = None alias: str | None = None description: str | None = None - transport: MCPTransportType = MCPTransport.sse + transport: MCPTransportType | None = MCPTransport.sse auth_type: MCPAuthType | None = None credentials: MCPCredentials | None = None url: str | None = None @@ -1546,6 +1563,28 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): max_concurrent_requests: int | None = None rpm: int | None = Field(default=None, ge=0) + @model_validator(mode="wrap") + @classmethod + def retain_extra_field_names( + cls, + values: object, + handler: ModelWrapValidatorHandler["UpdateMCPServerRequest"], + ) -> "UpdateMCPServerRequest": + validated: Final = handler(values) + if isinstance(values, Mapping): + validated._provider_extra_fields = active_registration_extras( + TypeAdapter(dict[str, object]).validate_python(values), + cls.model_fields.keys(), + ) + return validated + + def provider_registration_patch(self, fields_set: AbstractSet[str] | None = None) -> Mapping[str, object]: + if self._provider_extra_fields: + raise ValueError( + f"Provider-executed MCP registrations do not accept {', '.join(self._provider_extra_fields)}" + ) + return self.model_dump(include=set(fields_set) if fields_set is not None else None, exclude_unset=True) + @model_validator(mode="before") @classmethod def reject_execution_mode_changes(cls, values: object) -> object: @@ -1555,7 +1594,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): @model_validator(mode="after") def validate_protocol_transport(self) -> "UpdateMCPServerRequest": - if not {"transport", "mcp_info"}.issubset(self.model_fields_set): + if self.transport is None or not {"transport", "mcp_info"}.issubset(self.model_fields_set): return self validate_mcp_protocol_transport( TypeAdapter[MCPUpstreamProtocol](MCPUpstreamProtocol).validate_python( diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index a99cc47ecaf..a323c277b52 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -50,6 +50,7 @@ from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm.proxy.mcp_server.catalog import catalog_operation, public_catalog_operation +from litellm.types.mcp import MCPExecutionMode try: from prisma.errors import RecordNotFoundError, UniqueViolationError @@ -257,6 +258,7 @@ if MCP_AVAILABLE: ) from litellm.proxy.mcp_server.server_resolution import ( MCPServerTargetCatalog, + ResolvedMCPServer, authorize_mcp_server, resolve_mcp_server, ) @@ -272,6 +274,7 @@ if MCP_AVAILABLE: MCPCredentials, MCPGatewaySessionsResponse, MCPGatewaySessionsTerminateResponse, + MCPTransportType, MCPUpstreamProtocol, normalize_upstream_header_name, validate_mcp_protocol_transport, @@ -1342,13 +1345,34 @@ if MCP_AVAILABLE: aligned with the cards actually rendered: an admin in view_all mode sees every server even when their key carries no per-server MCP grant. """ - if _get_user_mcp_management_mode() == "view_all" and not _is_restricted_virtual_key_request(user_api_key_dict): - return await global_mcp_server_manager.get_all_mcp_servers_unfiltered() + unrestricted: Final = not _is_restricted_virtual_key_request(user_api_key_dict) + view_all: Final = unrestricted and _get_user_mcp_management_mode() == "view_all" + view_providers: Final = unrestricted and user_api_key_has_admin_view(user_api_key_dict) + inventory: Final = ( + await global_mcp_server_manager.get_all_mcp_servers_unfiltered() if view_all or view_providers else [] + ) + visible_inventory: Final = [ + server + for server in inventory + if view_providers + or getattr(server, "execution_mode", MCPExecutionMode.litellm) != MCPExecutionMode.provider + ] + if view_all: + return visible_inventory - aggregated: Final[dict[str, LiteLLM_MCPServerTable]] = {} - for auth_context in await build_effective_auth_contexts(user_api_key_dict): - for server in await global_mcp_server_manager.get_all_allowed_mcp_servers(user_api_key_auth=auth_context): - aggregated.setdefault(server.server_id, server) + provider_servers: Final = [ + server + for server in visible_inventory + if getattr(server, "execution_mode", MCPExecutionMode.litellm) == MCPExecutionMode.provider + ] + allowed_by_context: Final = [ + await global_mcp_server_manager.get_all_allowed_mcp_servers(user_api_key_auth=auth_context) + for auth_context in await build_effective_auth_contexts(user_api_key_dict) + ] + aggregated: Final = { + server.server_id: server + for server in (*chain.from_iterable(reversed(allowed_by_context)), *provider_servers) + } return list(aggregated.values()) @catalog_operation(lambda: global_mcp_server_manager) @@ -1949,20 +1973,40 @@ if MCP_AVAILABLE: --header 'Authorization: Bearer your_api_key_here' ``` """ - prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") - from litellm.proxy.auth.ip_address_utils import IPAddressUtils client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) is_admin_view: Final = user_api_key_has_admin_view(user_api_key_dict) is_restricted_virtual_key: Final = _is_restricted_virtual_key_request(user_api_key_dict) - resolved: Final = await resolve_mcp_server( - server_id, - manager=global_mcp_server_manager, - db_lookup=lambda sid: get_mcp_server(prisma_client, sid), - id_client_ip=client_ip, - name_client_ip=client_ip, - match_name=True, + snapshot: Final = global_mcp_server_manager.catalog.current() + config_match: Final = ( + next( + ( + server + for server in snapshot.management_servers.values() + if server.is_config + and server.execution_mode == MCPExecutionMode.provider + and server_id in (server.server_id, server.server_name, server.alias) + ), + None, + ) + if snapshot is not None + else None + ) + resolved: Final = ( + ResolvedMCPServer(table=config_match, runtime=None, source="registry") + if config_match is not None + else await resolve_mcp_server( + server_id, + manager=global_mcp_server_manager, + db_lookup=lambda sid: get_mcp_server( + get_prisma_client_or_throw("Database not connected. Connect a database to your proxy"), + sid, + ), + id_client_ip=client_ip, + name_client_ip=client_ip, + match_name=True, + ) ) authorized: Final = await authorize_mcp_server( resolved, @@ -1992,18 +2036,21 @@ if MCP_AVAILABLE: if from_db: await global_mcp_server_manager.add_server(mcp_server) - # Perform health check on the server using server manager - try: - health_result: Final = await global_mcp_server_manager.health_check_server(server_id) - # Update the server object with health check results - mcp_server.status = _mcp_health_status_for_response(health_result.status, include_reachability) or "unknown" - mcp_server.last_health_check = health_result.last_health_check - mcp_server.health_check_error = health_result.health_check_error - except Exception as e: - verbose_proxy_logger.debug("Error performing health check on server %s: %s", server_id, e) - mcp_server.status = "unknown" - mcp_server.last_health_check = datetime.now() - mcp_server.health_check_error = str(e) + if mcp_server.execution_mode == MCPExecutionMode.litellm: + # Perform health check on the server using server manager + try: + health_result: Final = await global_mcp_server_manager.health_check_server(server_id) + # Update the server object with health check results + mcp_server.status = ( + _mcp_health_status_for_response(health_result.status, include_reachability) or "unknown" + ) + mcp_server.last_health_check = health_result.last_health_check + mcp_server.health_check_error = health_result.health_check_error + except Exception as e: + verbose_proxy_logger.debug("Error performing health check on server %s: %s", server_id, e) + mcp_server.status = "unknown" + mcp_server.last_health_check = datetime.now() + mcp_server.health_check_error = str(e) redacted: Final = _redact_mcp_credentials(mcp_server) if is_restricted_virtual_key: @@ -2087,6 +2134,12 @@ if MCP_AVAILABLE: payload, touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, ) + except UniqueViolationError as exc: + raise HTTPException( + status_code=409, detail="An MCP registration already uses this identity or provider endpoint" + ) from exc + except HTTPException: + raise except Exception as e: verbose_proxy_logger.exception("Error creating mcp server: %s", e) raise HTTPException( @@ -3190,7 +3243,7 @@ if MCP_AVAILABLE: TypeAdapter[MCPUpstreamProtocol](MCPUpstreamProtocol).validate_python( (effective_info or {}).get("protocol_version", "auto") ), - effective_transport, + TypeAdapter[MCPTransportType](MCPTransportType).validate_python(effective_transport), ) except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc @@ -3302,12 +3355,17 @@ if MCP_AVAILABLE: ) # try to update the mcp server - mcp_server_record_updated: Final = await update_mcp_server( - prisma_client, - payload, - touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, - fields_set=payload_fields_set, - ) + try: + mcp_server_record_updated: Final = await update_mcp_server( + prisma_client, + payload, + touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, + fields_set=payload_fields_set, + ) + except UniqueViolationError as exc: + raise HTTPException( + status_code=409, detail="An MCP registration already uses this provider endpoint" + ) from exc if isinstance(mcp_server_record_updated, McpIdentifierConflict): raise_mcp_identifier_conflict(mcp_server_record_updated) diff --git a/litellm/proxy/mcp_server/catalog.py b/litellm/proxy/mcp_server/catalog.py index f8899fb948d..01ec5ea41af 100644 --- a/litellm/proxy/mcp_server/catalog.py +++ b/litellm/proxy/mcp_server/catalog.py @@ -23,8 +23,10 @@ from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache +from litellm.models.mcp_server import LiteLLM_MCPServerTable from litellm.proxy.mcp_server.result_conversion import age_freshness, aggregate_freshness from litellm.types.mcp import MCPExecutionMode +from litellm.types.mcp_server.provider_registration import canonical_provider_mcp_url if TYPE_CHECKING: from mcp.types import ListToolsResult, PaginatedRequestParams, PaginatedResult @@ -60,6 +62,7 @@ class CatalogSnapshot: identity: str tools: Mapping[str, MCPTool] routing: MutableMapping[str, str] + management_servers: Mapping[str, LiteLLM_MCPServerTable] def _configuration_identity(server: MCPServer) -> str: @@ -90,7 +93,41 @@ def _check_oauth_revision(selected: MCPServer, candidate: MCPServer | None) -> N raise HTTPException(status_code=503, detail="OAuth metadata discovery changed repeatedly; retry shortly") -def _snapshot(manager: MCPServerManager, database_identity: str) -> CatalogSnapshot: +def _provider_registration_conflicts(first: LiteLLM_MCPServerTable, second: LiteLLM_MCPServerTable) -> bool: + if MCPExecutionMode.provider not in (first.execution_mode, second.execution_mode): + return False + identifiers: Final = frozenset( + value.lower() for value in (first.server_id, first.server_name, first.alias) if value + ) + other_identifiers: Final = frozenset( + value.lower() for value in (second.server_id, second.server_name, second.alias) if value + ) + same_provider_url: Final = ( + first.execution_mode == second.execution_mode == MCPExecutionMode.provider + and first.url is not None + and second.url is not None + and canonical_provider_mcp_url(first.url) == canonical_provider_mcp_url(second.url) + ) + return bool(identifiers & other_identifiers) or same_provider_url + + +def validate_provider_registration_conflicts( + config_servers: Sequence[LiteLLM_MCPServerTable], + database_servers: Sequence[LiteLLM_MCPServerTable] = (), +) -> None: + for index, config_server in enumerate(config_servers): + for other in (*config_servers[index + 1 :], *database_servers): + if _provider_registration_conflicts(config_server, other): + raise ValueError( + f"MCP config registration '{config_server.server_name}' conflicts with registration '{other.server_name}'" + ) + + +def _snapshot( + manager: MCPServerManager, + database_identity: str, + registration_only_servers: Mapping[str, LiteLLM_MCPServerTable], +) -> CatalogSnapshot: from litellm.proxy.mcp_server.tool_registry import global_mcp_tool_registry servers: Final = manager.config_mcp_servers | manager.registry @@ -98,6 +135,7 @@ def _snapshot(manager: MCPServerManager, database_identity: str) -> CatalogSnaps serialized: Final = json.dumps( ( database_identity, + tuple(server.model_dump(mode="json") for server in manager.catalog.config_registrations.values()), tuple(sorted(manager.registry)), tuple((key, _configuration_identity(value)) for key, value in sorted(manager.config_mcp_servers.items())), ), @@ -108,6 +146,13 @@ def _snapshot(manager: MCPServerManager, database_identity: str) -> CatalogSnaps hashlib.sha256(serialized.encode()).hexdigest(), MappingProxyType(dict(global_mcp_tool_registry.published_tools)), dict(manager.published_tool_routes), + MappingProxyType( + { + **{key: value.model_copy(deep=True) for key, value in manager.catalog.config_registrations.items()}, + **{key: manager.build_mcp_server_table(value) for key, value in servers.items()}, + **{key: value.model_copy(deep=True) for key, value in registration_only_servers.items()}, + } + ), ) @@ -116,6 +161,8 @@ class CatalogSnapshots: self.manager = manager self._refresh_lock = asyncio.Lock() self._database_identity = "" + self._registration_only_servers: Mapping[str, LiteLLM_MCPServerTable] = MappingProxyType({}) + self.config_registrations: Mapping[str, LiteLLM_MCPServerTable] = MappingProxyType({}) self._arrival_ticket = 0 self._completed_ticket = 0 self._shared_snapshot: CatalogSnapshot | None = None @@ -129,6 +176,24 @@ class CatalogSnapshots: "mcp_catalog_routing", default=None ) + def validate_config_inventory(self, servers: Sequence[LiteLLM_MCPServerTable]) -> None: + database_servers: Final = ( + *(self.manager.build_mcp_server_table(server) for server in self.manager.registry.values()), + *self._registration_only_servers.values(), + ) + validate_provider_registration_conflicts(servers, database_servers) + + def publish_config_registrations(self, servers: Mapping[str, LiteLLM_MCPServerTable]) -> None: + self.config_registrations = MappingProxyType(dict(servers)) + self._shared_snapshot = None + + def validate_database_registration(self, server: LiteLLM_MCPServerTable) -> None: + config_servers: Final = ( + *self.config_registrations.values(), + *(self.manager.build_mcp_server_table(value) for value in self.manager.config_mcp_servers.values()), + ) + validate_provider_registration_conflicts(config_servers, (server,)) + def current(self) -> CatalogSnapshot | None: scoped: Final = self._operation.get() return scoped[0] if scoped is not None and not scoped[1].is_set() else None @@ -150,11 +215,11 @@ class CatalogSnapshots: from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - return _snapshot(self.manager, self._database_identity) + return _snapshot(self.manager, self._database_identity, self._registration_only_servers) from litellm.proxy.proxy_server import should_load_db_object if not should_load_db_object("mcp"): - return _snapshot(self.manager, self._database_identity) + return _snapshot(self.manager, self._database_identity, self._registration_only_servers) from litellm.proxy.mcp_server.db import get_mcp_catalog_revision try: @@ -176,6 +241,10 @@ class CatalogSnapshots: async with self.operation() as snapshot: return snapshot.servers + async def management_inventory(self) -> Mapping[str, LiteLLM_MCPServerTable]: + async with self.operation() as snapshot: + return snapshot.management_servers + def assert_current(self, server: MCPServer) -> None: from fastapi import HTTPException @@ -208,6 +277,9 @@ class CatalogSnapshots: shared, servers=MappingProxyType({key: value.model_copy(deep=True) for key, value in shared.servers.items()}), routing=routing, + management_servers=MappingProxyType( + {key: value.model_copy(deep=True) for key, value in shared.management_servers.items()} + ), ) closed: Final = asyncio.Event() token: Final = self._operation.set((snapshot, closed)) @@ -288,7 +360,7 @@ class CatalogSnapshots: self._completed_ticket = covered raise else: - self._shared_snapshot = _snapshot(self.manager, self._database_identity) + self._shared_snapshot = _snapshot(self.manager, self._database_identity, self._registration_only_servers) self._applied_revision = revision self._completed_ticket = covered finally: @@ -490,6 +562,24 @@ class CatalogSnapshots: ).hexdigest() verbose_logger.info("Found %s MCP servers in database", len(raw_rows)) + registration_only_servers: Final = MappingProxyType( + { + server.server_id: server + for server in ( + LiteLLM_MCPServerTable.model_validate(row.model_dump()) + for row in raw_rows + if row.model_dump().get("execution_mode") == MCPExecutionMode.provider + ) + } + ) + config_servers: Final = ( + *self.config_registrations.values(), + *(self.manager.build_mcp_server_table(value) for value in self.manager.config_mcp_servers.values()), + ) + if self.config_registrations or registration_only_servers: + validate_provider_registration_conflicts( + config_servers, tuple(LiteLLM_MCPServerTable.model_validate(row.model_dump()) for row in raw_rows) + ) previous_registry: Final = self.manager.registry new_registry: Final = await self._stage_servers(raw_rows, reuse_unchanged=reuse_unchanged) @@ -569,6 +659,7 @@ class CatalogSnapshots: ", ".join(sorted(capturing_config_server_ids)), ) self._warned_capturing_config_server_ids = capturing_config_server_ids + self._registration_only_servers = registration_only_servers def catalog_operation( diff --git a/litellm/proxy/mcp_server/db.py b/litellm/proxy/mcp_server/db.py index 27f8b90a586..ebcaa606831 100644 --- a/litellm/proxy/mcp_server/db.py +++ b/litellm/proxy/mcp_server/db.py @@ -3,6 +3,7 @@ import binascii import hashlib import json from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence +from collections.abc import Set as AbstractSet from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast @@ -57,7 +58,13 @@ from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) from litellm.types.llms.custom_http import httpxSpecialProvider -from litellm.types.mcp import MCPCredentials, MCPTransportType, MCPUpstreamProtocol, validate_mcp_protocol_transport +from litellm.types.mcp import ( + MCPCredentials, + MCPExecutionMode, + MCPTransportType, + MCPUpstreamProtocol, + validate_mcp_protocol_transport, +) from litellm.types.mcp_server.mcp_server_manager import MCPInfo, PinnedMCPTool if TYPE_CHECKING: @@ -452,6 +459,8 @@ def _prepare_mcp_server_data( data_dict[json_map_field] = {} else: data_dict = data.model_dump(exclude_none=True) + if isinstance(data, NewMCPServerRequest) and data.execution_mode == MCPExecutionMode.provider: + data_dict["transport"] = data.transport # Ensure alias is always present in the dict (even if None) if "alias" not in data_dict: data_dict["alias"] = getattr(data, "alias", None) @@ -1076,16 +1085,46 @@ async def delete_mcp_server( return deleted_server # pyright: ignore[reportReturnType] # prisma row, not domain LiteLLM_MCPServerTable +def _validate_mcp_config_registration_write(values: Mapping[str, object]) -> None: + from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager + + catalog: Final = global_mcp_server_manager.catalog + if not catalog.config_registrations and values.get("execution_mode") != MCPExecutionMode.provider: + return + identity: Final = { + key: values[key] + for key in ( + "server_id", + "server_name", + "alias", + "url", + "transport", + "execution_mode", + ) + if key in values + } + try: + catalog.validate_database_registration(LiteLLM_MCPServerTable.model_validate(identity)) + except ValueError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc + + async def create_mcp_server( prisma_client: PrismaClient, data: NewMCPServerRequest, touched_by: str ) -> LiteLLM_MCPServerTable: """ Create a new mcp server record in the db """ + if data.execution_mode == MCPExecutionMode.provider: + provider_record: Final = await create_mcp_server_if_identifier_free(prisma_client, data, touched_by) + if isinstance(provider_record, McpIdentifierConflict): + raise HTTPException(status_code=409, detail="An MCP registration already uses this identity") + return provider_record if data.server_id is None: data.server_id = str(uuid.uuid4()) # Use helper to prepare data with proper JSON serialization + _validate_mcp_config_registration_write(data.model_dump()) data_dict: Final = _prepare_mcp_server_data(data) # Add audit fields @@ -1110,6 +1149,7 @@ async def create_mcp_server_if_identifier_free( if data.server_id is None: data.server_id = str(uuid.uuid4()) + _validate_mcp_config_registration_write(data.model_dump()) data_dict: Final = _prepare_mcp_server_data(data) data_dict["created_by"] = touched_by data_dict["updated_by"] = touched_by @@ -1246,35 +1286,75 @@ def _validate_mcp_protocol_write( raise HTTPException(status_code=400, detail=str(exc)) from exc +def _validate_provider_registration_update( + stored: "prisma_db_models.LiteLLM_MCPServerTable", + data: UpdateMCPServerRequest, + fields_set: AbstractSet[str] | None = None, +) -> NewMCPServerRequest: + try: + patch: Final = data.provider_registration_patch(fields_set) + registration: Final = LiteLLM_MCPServerTable.model_validate(stored.model_dump()) + merged: Final = {**registration.model_dump(include=set(NewMCPServerRequest.model_fields)), **patch} + return NewMCPServerRequest.model_validate(merged) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + +def _prepare_provider_registration_update( + stored: "prisma_db_models.LiteLLM_MCPServerTable", + data: UpdateMCPServerRequest, + touched_by: str, + fields_set: AbstractSet[str] | None, +) -> Mapping[str, object]: + from litellm.types.mcp_server.provider_registration import ProviderMCPRegistration + + validated: Final = _validate_provider_registration_update(stored, data, fields_set) + _validate_mcp_config_registration_write(validated.model_dump()) + supplied_fields: Final = fields_set if fields_set is not None else data.model_fields_set + normalized_patch: Final = UpdateMCPServerRequest.model_validate( + validated.model_dump( + include=set(supplied_fields & ProviderMCPRegistration.model_fields.keys()) | {"server_id", "mcp_info"} + ) + ) + serialized: Final = _prepare_mcp_server_data(normalized_patch, exclude_unset=True, fields_set=set(supplied_fields)) + return { + **serialized, + "updated_by": touched_by, + } + + async def _update_mcp_server_row( prisma_client: PrismaClient, *, server_id: str, data_dict: Mapping[str, object], + prepare_update: Callable[["prisma_db_models.LiteLLM_MCPServerTable"], Mapping[str, object]] | None = None, expected_updated_at: datetime | None = None, ) -> "prisma_db_models.LiteLLM_MCPServerTable | McpIdentifierConflict | None": identifier_write: Final = any(field in data_dict for field in ("server_name", "alias")) protocol_write: Final = bool({"transport", "mcp_info"}.intersection(data_dict)) + needs_stored_row: Final = protocol_write or prepare_update is not None async def _update( table: "TableActions[prisma_db_models.LiteLLM_MCPServerTable]", ) -> "prisma_db_models.LiteLLM_MCPServerTable | None": - if protocol_write: - stored: Final = await table.find_unique(where={"server_id": server_id}) - if stored is None: - return None + stored: Final = await table.find_unique(where={"server_id": server_id}) if needs_stored_row else None + if needs_stored_row and stored is None: + return None + if protocol_write and prepare_update is None and stored is not None: _validate_mcp_protocol_write(stored, data_dict) + write_data: Final = prepare_update(stored) if prepare_update is not None and stored is not None else data_dict if expected_updated_at is not None: changed: Final = await table.update_many( - where={"server_id": server_id, "updated_at": expected_updated_at}, data=data_dict + where={"server_id": server_id, "updated_at": expected_updated_at}, data=write_data ) return await table.find_unique(where={"server_id": server_id}) if changed else None return await table.update( where={"server_id": server_id}, - data=data_dict, + data=write_data, ) - if not identifier_write and not protocol_write: + if not identifier_write and not needs_stored_row: return await _update(_mcp_server_table_actions(prisma_client)) if "alias" in data_dict and not data_dict["alias"] and "server_name" not in data_dict: # Clearing the alias drops the prefix to the stored server_name, which @@ -1288,7 +1368,7 @@ async def _update_mcp_server_row( alias=None, exclude_server_id=server_id, write=_update, - lock_server_id=server_id if protocol_write else None, + lock_server_id=server_id if needs_stored_row else None, ) return await _mcp_server_write_if_identifier_free( prisma_client, @@ -1296,7 +1376,7 @@ async def _update_mcp_server_row( alias=_identifier_field(data_dict, "alias"), exclude_server_id=server_id, write=_update, - lock_server_id=server_id if protocol_write else None, + lock_server_id=server_id if needs_stored_row else None, ) @@ -1315,6 +1395,24 @@ async def update_mcp_server( """ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + registration: Final = await _db_find_mcp_server_row(prisma_client, data.server_id) + if getattr(registration, "execution_mode", MCPExecutionMode.litellm) == MCPExecutionMode.provider: + updated_provider: Final = await _update_mcp_server_row( + prisma_client, + server_id=data.server_id, + data_dict=data.model_dump(include=fields_set, exclude_unset=True), + prepare_update=lambda stored: _prepare_provider_registration_update(stored, data, touched_by, fields_set), + expected_updated_at=expected_updated_at, + ) + return ( + updated_provider + if updated_provider is None or isinstance(updated_provider, McpIdentifierConflict) + else LiteLLM_MCPServerTable.model_validate(updated_provider.model_dump()) + ) + + if registration is not None: + _validate_mcp_config_registration_write({**registration.model_dump(), **data.model_dump(exclude_unset=True)}) + # Use helper to prepare data with proper JSON serialization. # exclude_unset=True makes this a true partial update: fields the caller did # not provide are not written, so they keep their existing DB value instead @@ -1322,15 +1420,17 @@ async def update_mcp_server( prepared_data: Final = _prepare_mcp_server_data(data, exclude_unset=True, fields_set=fields_set) # Pre-fetch existing record once if we need it for auth_type, url, or credential logic - existing = None has_credentials: Final = "credentials" in prepared_data and prepared_data["credentials"] is not None # An explicit token-exchange column write (set or clear) also migrates the # legacy blob copies below, so the existing row is needed for those updates. explicit_te_write: Final = bool(_TOKEN_EXCHANGE_COLUMN_FIELDS & prepared_data.keys()) url_provided: Final = "url" in prepared_data and prepared_data["url"] is not None issuer_provided: Final = "issuer" in prepared_data - if data.auth_type or has_credentials or explicit_te_write or url_provided or issuer_provided: - existing = await _db_find_mcp_server_row(prisma_client, data.server_id) + existing: Final = ( + registration + if data.auth_type or has_credentials or explicit_te_write or url_provided or issuer_provided + else None + ) auth_type_changed: Final = bool( data.auth_type @@ -1442,9 +1542,6 @@ async def update_mcp_server( # Add audit fields data_dict["updated_by"] = touched_by - # prisma-python rejects a raw ``None`` for a ``Json?`` field ("value is required but not set"); the - # clear paths above use ``None`` as the merge-skip sentinel, so translate it here to ``Json(None)``, - # which writes SQL null and reads back as ``None``. Done at the edge so the merge guards stay simple. if "credentials" in data_dict and data_dict["credentials"] is None: from prisma import Json # noqa: PLC0415 # local import: prisma may be ungenerated at module load in some tools diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index 25c32952f7d..ff012d94c83 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -89,6 +89,7 @@ from litellm.proxy._types import ( MCPEnvVar, MCPTransport, MCPTransportType, + NewMCPServerRequest, SpecialMCPServerNames, UserAPIKeyAuth, is_per_server_oauth_discovery_eligible, @@ -408,6 +409,7 @@ class MCPServerConfig(TypedDict, total=False): server_id: ReadOnly[str] execution_mode: ReadOnly[MCPExecutionMode] + tool_name_to_description: ReadOnly[Mapping[str, str]] protocol_version: ReadOnly[MCPUpstreamProtocol] alias: str description: str @@ -2558,13 +2560,83 @@ class MCPServerManager: store; a rowless server that is NOT config-declared is a throwaway temp/session server whose client must not be persisted. This never overrides the row-existence check: a server that has a LiteLLM_MCPServerTable row is always resolved to that row first.""" - return server_id in self.config_mcp_servers + return server_id in self.config_mcp_servers or server_id in self.catalog.config_registrations + + def _config_registration( + self, + name: str, + values: MCPServerConfig, + aliases: Mapping[str, str] | None, + ) -> LiteLLM_MCPServerTable: + alias: Final = _assigned_alias(name, values, aliases) + pinned_id: Final = _pinned_config_server_id(values.get("server_id"), name) + if values.get("execution_mode", MCPExecutionMode.litellm) == MCPExecutionMode.provider: + request: Final = NewMCPServerRequest.model_validate( + { + **{field: value for field, value in values.items() if field != "access_groups"}, + "server_name": name, + "alias": alias, + "mcp_access_groups": values.get("access_groups", []), + } + ) + stable_id: Final = hashlib.sha256(f"provider|{name}|{request.url}|{alias or ''}".encode()).hexdigest()[:32] + return LiteLLM_MCPServerTable.model_validate( + { + **request.model_dump(exclude_none=True), + "server_id": pinned_id or stable_id, + "transport": request.transport, + "is_config": True, + } + ) + validate_mcp_gateway_execution(values.get("execution_mode", MCPExecutionMode.litellm)) + return LiteLLM_MCPServerTable( + server_id=pinned_id + or self._generate_stable_server_id( + name, + values.get("url") or "", + values.get("transport", MCPTransport.http), + values.get("auth_type"), + alias, + ), + server_name=name, + alias=alias, + url=values.get("url"), + transport=values.get("transport", MCPTransport.http), + is_config=True, + ) async def load_servers_from_config( self, mcp_servers_config: dict[str, MCPServerConfig], mcp_aliases: dict[str, str] | None = None, - ): + ) -> None: + _validate_config_server_names(mcp_servers_config) + registrations: Final = tuple( + self._config_registration(name, config, mcp_aliases) for name, config in mcp_servers_config.items() + ) + self.catalog.validate_config_inventory(registrations) + providers: Final = { + server.server_id: server for server in registrations if server.execution_mode == MCPExecutionMode.provider + } + provider_names: Final = frozenset(server.server_name for server in providers.values()) + for server in tuple(self.config_mcp_servers.values()): + if server.server_name in provider_names or server.server_id in providers: + self._cleanup_server_tool_routing_artifacts(server) + self.invalidate_oauth_discovery_state(server.server_id) + self.config_mcp_servers.pop(server.server_id, None) + gateway_config: Final = { + name: config + for name, config in mcp_servers_config.items() + if config.get("execution_mode", MCPExecutionMode.litellm) != MCPExecutionMode.provider + } + await self._load_gateway_servers_from_config(gateway_config, mcp_aliases) + self.catalog.publish_config_registrations(providers) + + async def _load_gateway_servers_from_config( + self, + mcp_servers_config: Mapping[str, MCPServerConfig], + mcp_aliases: Mapping[str, str] | None = None, + ) -> None: """ Load the MCP Servers from the config @@ -7869,16 +7941,8 @@ class MCPServerManager: _build_mcp_server_table = build_mcp_server_table async def get_all_mcp_servers_unfiltered(self) -> list[LiteLLM_MCPServerTable]: - """Return all MCP servers from registry without applying access controls.""" - - registry: Final = self.get_registry() - if not registry: - return [] - - servers: Final[list[LiteLLM_MCPServerTable]] = [] - for server in registry.values(): - servers.append(self.build_mcp_server_table(server)) - return servers + """Return the cached management inventory without applying access controls.""" + return list((await self.catalog.management_inventory()).values()) async def get_all_mcp_servers_with_health_unfiltered( self, server_ids: list[str] | None = None diff --git a/litellm/proxy/mcp_server/rest_endpoints.py b/litellm/proxy/mcp_server/rest_endpoints.py index 2737b567985..7df863d2676 100644 --- a/litellm/proxy/mcp_server/rest_endpoints.py +++ b/litellm/proxy/mcp_server/rest_endpoints.py @@ -1671,6 +1671,9 @@ if MCP_AVAILABLE: Returns: The dict returned by *operation*, or an error dict on failure. """ + if request.transport is None: + return {"error": "Provider-executed MCP registrations do not support gateway connections"} + try: client_id, client_secret, scopes = _extract_credentials(request) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 87bc78730b5..7f26667d35d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7612,17 +7612,18 @@ class ProxyConfig: global_agent_registry.load_agents_from_config(agent_config) + from litellm.proxy.mcp_server.utils import is_mcp_available + mcp_servers_config: Final = config.get("mcp_servers", None) - if mcp_servers_config: - from litellm.proxy.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) + if mcp_servers_config or is_mcp_available(): + from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager - # Get mcp_aliases from litellm_settings if available - litellm_settings: Final = config.get("litellm_settings", {}) - mcp_aliases: Final = litellm_settings.get("mcp_aliases", None) - - await global_mcp_server_manager.load_servers_from_config(mcp_servers_config, mcp_aliases) + if mcp_servers_config: + litellm_settings: Final = config.get("litellm_settings", {}) + mcp_aliases: Final = litellm_settings.get("mcp_aliases", None) + await global_mcp_server_manager.load_servers_from_config(mcp_servers_config, mcp_aliases) + elif global_mcp_server_manager.catalog.config_registrations: + global_mcp_server_manager.catalog.publish_config_registrations({}) ## VECTOR STORES vector_store_registry_config: Final = config.get("vector_store_registry", None) diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 3556198723e..13218dad84c 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -29,11 +29,10 @@ class MCPExecutionMode(str, enum.Enum): provider = "provider" -MCPGatewayExecutionMode: TypeAlias = Literal[MCPExecutionMode.litellm] -_GATEWAY_EXECUTION_MODE: Final = TypeAdapter[MCPGatewayExecutionMode](MCPGatewayExecutionMode) +_GATEWAY_EXECUTION_MODE: Final = TypeAdapter[Literal[MCPExecutionMode.litellm]](Literal[MCPExecutionMode.litellm]) -def validate_mcp_gateway_execution(value: object) -> MCPGatewayExecutionMode: +def validate_mcp_gateway_execution(value: object) -> Literal[MCPExecutionMode.litellm]: return _GATEWAY_EXECUTION_MODE.validate_python(value) diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 6080bbcabec..8a5cf750df6 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -12,7 +12,6 @@ from litellm.types.mcp import ( MCPAuth, MCPAuthType, MCPExecutionMode, - MCPGatewayExecutionMode, MCPTokenEndpointAuthMethod, MCPTransportType, MCPUpstreamProtocol, @@ -98,7 +97,7 @@ def parse_pinned_tools(value: object) -> dict[str, PinnedMCPTool] | None: class MCPServer(LiteLLMBaseModel): server_id: str name: str - execution_mode: MCPGatewayExecutionMode = MCPExecutionMode.litellm + execution_mode: Literal[MCPExecutionMode.litellm] = MCPExecutionMode.litellm alias: str | None = None server_name: str | None = None url: str | None = None diff --git a/litellm/types/mcp_server/provider_registration.py b/litellm/types/mcp_server/provider_registration.py new file mode 100644 index 00000000000..cf11b0914eb --- /dev/null +++ b/litellm/types/mcp_server/provider_registration.py @@ -0,0 +1,129 @@ +import json +from collections.abc import Collection, Mapping, Sequence +from datetime import datetime +from typing import Final, Literal +from urllib.parse import urlsplit, urlunsplit + +from pydantic import AnyHttpUrl, BaseModel, ConfigDict, Field, TypeAdapter, field_validator, model_validator +from typing_extensions import Self + +from litellm.types.mcp import MCPExecutionMode + +_STR: Final = TypeAdapter(str) +_MAP: Final = TypeAdapter(Mapping[str, object]) +_URL: Final = TypeAdapter(AnyHttpUrl) + + +def _is_neutral_registration_value(value: object) -> bool: + if isinstance(value, str): + return value.strip() in ("", "null", "{}", "[]") + return value is None or value is False or value in ({}, []) + + +def active_registration_extras(values: Mapping[str, object], known_fields: Collection[str]) -> tuple[str, ...]: + return tuple( + field + for field, value in values.items() + if field not in known_fields and not _is_neutral_registration_value(value) + ) + + +class ProviderMCPRegistration(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + execution_mode: Literal[MCPExecutionMode.provider] = MCPExecutionMode.provider + server_id: str | None = None + server_name: str | None = None + alias: str | None = None + description: str | None = None + url: str + allowed_tools: Sequence[str] + mcp_info: Mapping[str, object] = Field(default_factory=dict, validate_default=True) + mcp_access_groups: Sequence[str] = () + tool_name_to_display_name: Mapping[str, str] | None = None + tool_name_to_description: Mapping[str, str] | None = None + instructions: str | None = None + source_url: str | None = None + allow_all_keys: bool = False + available_on_public_internet: Literal[True] = True + protocol_version: Literal["auto"] = "auto" + approval_status: str | None = None + submitted_by: str | None = None + submitted_at: datetime | None = None + + @model_validator(mode="before") + @classmethod + def discard_neutral_extras(cls, values: object) -> object: + if not isinstance(values, Mapping): + return values + fields: Final = _MAP.validate_python(values) + return { + field: value + for field, value in fields.items() + if field in cls.model_fields or not _is_neutral_registration_value(value) + } + + @field_validator("url", mode="before") + @classmethod + def normalize_url(cls, value: object) -> str: + return canonical_provider_mcp_url(value) + + @field_validator("allowed_tools") + @classmethod + def normalize_allowed_tools(cls, tools: Sequence[str]) -> tuple[str, ...]: + if any(not tool.strip() or any(char in tool for char in "*?[]") for tool in tools): + raise ValueError("allowed_tools must contain nonblank exact tool names without wildcards") + return tuple(dict.fromkeys(tools)) + + @field_validator("mcp_info", mode="before") + @classmethod + def normalize_metadata(cls, value: object) -> Mapping[str, object]: + info: Final = _MAP.validate_python(_decoded_json(value) or {}) + if "tool_allowlist_enforced" in info and info["tool_allowlist_enforced"] is not True: + raise ValueError("Provider registrations require mcp_info.tool_allowlist_enforced=true") + if info.get("protocol_version", "auto") != "auto": + raise ValueError("Provider registrations do not configure the upstream MCP protocol") + return {**info, "tool_allowlist_enforced": True} + + @field_validator("tool_name_to_description", "tool_name_to_display_name", mode="before") + @classmethod + def decode_stored_metadata(cls, value: object) -> object: + return _decoded_json(value) + + @model_validator(mode="after") + def validate_tool_descriptions(self) -> Self: + if not (self.tool_name_to_description or {}).keys() <= set(self.allowed_tools): + raise ValueError("Tool descriptions must reference allowed_tools") + return self + + +def canonical_provider_mcp_url(value: object) -> str: + if value is None: + raise ValueError("Provider MCP url is required") + url: Final = _STR.validate_python(value) + if any(char.isspace() or ord(char) < 32 for char in url) or "\\" in url: + raise ValueError("Provider MCP url must not contain whitespace, control characters, or backslashes") + _URL.validate_python(url) + parsed: Final = urlsplit(url) + if parsed.scheme.lower() != "https" or not parsed.hostname or parsed.username is not None or "#" in url: + raise ValueError("Provider MCP url must use HTTPS without userinfo or a fragment") + host: Final = parsed.hostname.lower().encode("idna").decode("ascii") + hostname: Final = f"[{host}]" if ":" in host else host + port: Final = parsed.port + authority: Final = f"{hostname}:{port}" if port is not None and port != 443 else hostname + return urlunsplit(("https", authority, parsed.path, parsed.query, "")) + + +def _decoded_json(value: object) -> object: + return json.loads(value) if isinstance(value, str) else value + + +def normalize_provider_registration(values: Mapping[str, object]) -> Mapping[str, object]: + registration: Final = ProviderMCPRegistration.model_validate(values) + return { + **registration.model_dump(mode="json", exclude_unset=True), + "mcp_info": registration.mcp_info, + "transport": None, + "auth_type": None, + "credentials": None, + } 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 4b1415c8879..6b2a1a3ffc4 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -408,6 +408,7 @@ class TestListMCPServers: ), ] mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=mock_servers) + mock_manager.get_all_mcp_servers_unfiltered = AsyncMock(return_value=[]) for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} @@ -655,6 +656,7 @@ class TestListMCPServers: ), ] mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=mock_servers) + mock_manager.get_all_mcp_servers_unfiltered = AsyncMock(return_value=[]) for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} @@ -774,6 +776,7 @@ class TestListMCPServers: ), ] mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=mock_servers) + mock_manager.get_all_mcp_servers_unfiltered = AsyncMock(return_value=[]) for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} @@ -866,6 +869,7 @@ class TestListMCPServers: # Mock manager mock_manager = MagicMock() mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=[server_1, server_2]) + mock_manager.get_all_mcp_servers_unfiltered = AsyncMock(return_value=[server_1, server_2]) with ( patch( @@ -2023,6 +2027,7 @@ class TestFetchAllMCPServersOrdering: resolved: Final = list(reversed(servers) if reverse else servers) mock_manager: Final = MagicMock() mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=resolved) + mock_manager.get_all_mcp_servers_unfiltered = AsyncMock(return_value=[]) with ( patch( # test-quality-ok: the route reads a module-global manager with no injection seam "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", @@ -8096,6 +8101,7 @@ class TestConnectedAppViewAnnotation: mock_manager = MagicMock() mock_manager.get_all_allowed_mcp_servers = AsyncMock(side_effect=per_context_servers) + mock_manager.get_all_mcp_servers_unfiltered = AsyncMock(return_value=[]) mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=["server-1", "server-2"]) with ( diff --git a/tests/unit/proxy/mcp_server/test_catalog.py b/tests/unit/proxy/mcp_server/test_catalog.py index 384a412c4f1..10629672bd9 100644 --- a/tests/unit/proxy/mcp_server/test_catalog.py +++ b/tests/unit/proxy/mcp_server/test_catalog.py @@ -851,3 +851,204 @@ async def test_provider_mode_is_rejected_before_upstream_auth_or_connection(boun ) with pytest.raises(ValidationError): await attempt + + +@pytest.mark.asyncio +async def test_management_and_runtime_inventory_share_refresh_and_request_snapshot(monkeypatch: pytest.MonkeyPatch) -> None: + from fastapi import HTTPException + + from litellm.models.mcp_server import LiteLLM_MCPServerTable + from litellm.proxy import proxy_server + from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager + + provider: Final = LiteLLM_MCPServerTable( + server_id="provider", execution_mode="provider", transport=None, + url="https://provider.example/mcp", allowed_tools=["search"], + mcp_info={"tool_allowlist_enforced": True}, + ) + updated: Final = provider.model_copy(update={"allowed_tools": []}) + read_rows: Final = AsyncMock(side_effect=([provider], [updated], RuntimeError("offline"), [])) + read_revision: Final = AsyncMock(return_value=SimpleNamespace(reload_revision=1)) + tables: Final = SimpleNamespace( + litellm_mcpservertable=SimpleNamespace(find_many=read_rows), + litellm_config=SimpleNamespace(find_unique=read_revision), + ) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables)) + monkeypatch.setattr(proxy_server, "general_settings", {}) + manager: Final = MCPServerManager() + async with manager.catalog.operation(): + assert (await manager.get_all_mcp_servers_unfiltered())[0].allowed_tools == ["search"] + assert await manager.catalog.list() == {} + read_revision.return_value = SimpleNamespace(reload_revision=2) + assert (await manager.get_all_mcp_servers_unfiltered())[0].allowed_tools == ["search"] + assert (await manager.get_all_mcp_servers_unfiltered())[0].allowed_tools == [] + assert await manager.catalog.list() == {} + assert read_rows.await_count == 2 + read_revision.return_value = SimpleNamespace(reload_revision=3) + with pytest.raises(HTTPException) as unavailable: + await manager.get_all_mcp_servers_unfiltered() + assert unavailable.value.status_code == 503 + assert await manager.get_all_mcp_servers_unfiltered() == [] + assert await manager.catalog.list() == {} + assert read_rows.await_count == 4 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("pinned_id", [None, "configured-docs"]) +async def test_provider_config_inventory_reloads_without_entering_gateway( + monkeypatch: pytest.MonkeyPatch, pinned_id: str | None, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager + + monkeypatch.setattr(proxy_server, "prisma_client", None) + manager: Final = MCPServerManager() + config: Final = { + "docs": { + "execution_mode": "provider", "url": "https://DOCS.example:443/mcp", "allowed_tools": ["search"], + "tool_name_to_description": {"search": "Search documentation"}, + **({"server_id": pinned_id} if pinned_id else {}), + }, + } + await manager.load_servers_from_config(config) + first: Final = (await manager.get_all_mcp_servers_unfiltered())[0] + assert first.is_config and first.execution_mode == "provider" + assert first.url == "https://docs.example/mcp" + assert first.allowed_tools == ["search"] + assert first.tool_name_to_description == {"search": "Search documentation"} + assert first.transport is None and first.credentials is None + assert first.mcp_info["tool_allowlist_enforced"] is True + assert not manager.get_registry() + async with manager.catalog.operation(): + await manager.load_servers_from_config({"docs": { + **config["docs"], "url": first.url, "allowed_tools": [], "tool_name_to_description": {}, + }}) + assert (await manager.get_all_mcp_servers_unfiltered())[0].allowed_tools == ["search"] + updated: Final = (await manager.get_all_mcp_servers_unfiltered())[0] + assert updated.server_id == first.server_id + assert updated.allowed_tools == [] + await manager.load_servers_from_config({}) + assert await manager.get_all_mcp_servers_unfiltered() == [] + assert not manager.get_registry() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("field,value", [ + ("authentication_token", "synthetic"), ("client_secret", "synthetic"), + ("aws_access_key_id", "synthetic"), ("client_private_key", "synthetic"), + ("allow_sampling", True), ("protocol_version", "2025-03-26"), + ("gateway_custom_option", {"enabled": True}), +]) +async def test_provider_config_rejects_active_config_only_connection_fields( + monkeypatch: pytest.MonkeyPatch, field: str, value: object, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager + + monkeypatch.setattr(proxy_server, "prisma_client", None) + manager: Final = MCPServerManager() + config: Final = {"execution_mode": "provider", "url": "https://docs.example/mcp", "allowed_tools": ["search"]} + await manager.load_servers_from_config({"docs": config}) + previous: Final = await manager.get_all_mcp_servers_unfiltered() + with pytest.raises(ValueError, match=field): + await manager.load_servers_from_config({"docs": {**config, field: value}}) + assert await manager.get_all_mcp_servers_unfiltered() == previous + assert not manager.get_registry() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("other", [ + {"execution_mode": "provider", "url": "https://DOCS.example:443/mcp", "allowed_tools": []}, + {"execution_mode": "provider", "url": "https://other.example/mcp", "allowed_tools": [], "alias": "docs"}, + {"url": "https://gateway.example/mcp", "transport": "http", "server_id": "docs"}, +]) +async def test_provider_config_rejects_ambiguous_config_identities( + monkeypatch: pytest.MonkeyPatch, other: dict[str, object], +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager + + monkeypatch.setattr(proxy_server, "prisma_client", None) + manager: Final = MCPServerManager() + with pytest.raises(ValueError, match="conflicts"): + await manager.load_servers_from_config({ + "docs": {"execution_mode": "provider", "url": "https://docs.example/mcp", "allowed_tools": []}, + "other": other, + }) + assert not manager.get_registry() + assert await manager.get_all_mcp_servers_unfiltered() == [] + + +@pytest.mark.asyncio +async def test_provider_config_replaces_gateway_entry_without_retaining_execution(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager + + monkeypatch.setattr(proxy_server, "prisma_client", None) + manager: Final = MCPServerManager() + await manager.load_servers_from_config({"docs": {"transport": "http", "url": "https://docs.example/mcp"}}) + assert manager.get_mcp_server_by_name("docs") is not None + await manager.load_servers_from_config({"docs": { + "execution_mode": "provider", "url": "https://docs.example/mcp", "allowed_tools": [], + }}) + assert manager.get_mcp_server_by_name("docs") is None + assert not manager.get_registry() + assert (await manager.get_all_mcp_servers_unfiltered())[0].execution_mode == "provider" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("database_first", [False, True]) +async def test_provider_config_rejects_database_url_conflict_in_either_loading_order( + monkeypatch: pytest.MonkeyPatch, database_first: bool, +) -> None: + from fastapi import HTTPException + + from litellm.models.mcp_server import LiteLLM_MCPServerTable + from litellm.proxy import proxy_server + from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager + + row: Final = LiteLLM_MCPServerTable( + server_id="database-docs", server_name="database-docs", execution_mode="provider", transport=None, + url="https://docs.example/mcp", allowed_tools=["search"], + ) + tables: Final = SimpleNamespace( + litellm_mcpservertable=SimpleNamespace(find_many=AsyncMock(return_value=[row])), + litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=SimpleNamespace(reload_revision=1))), + ) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables)) + monkeypatch.setattr(proxy_server, "general_settings", {}) + manager: Final = MCPServerManager() + config: Final = {"docs": {"execution_mode": "provider", "url": row.url, "allowed_tools": []}} + if database_first: + await manager.get_all_mcp_servers_unfiltered() + with pytest.raises(ValueError, match="conflicts"): + await manager.load_servers_from_config(config) + else: + await manager.load_servers_from_config(config) + with pytest.raises(HTTPException) as unavailable: + await manager.get_all_mcp_servers_unfiltered() + assert unavailable.value.status_code == 503 + assert not manager.get_registry() + + +@pytest.mark.asyncio +async def test_provider_config_changes_invalidate_cached_inventory_without_database_revision_change( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager + + tables: Final = SimpleNamespace( + litellm_mcpservertable=SimpleNamespace(find_many=AsyncMock(return_value=[])), + litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=SimpleNamespace(reload_revision=1))), + ) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables)) + monkeypatch.setattr(proxy_server, "general_settings", {}) + manager: Final = MCPServerManager() + initial: Final = {"execution_mode": "provider", "url": "https://docs.example/mcp", "allowed_tools": ["search"]} + await manager.load_servers_from_config({"docs": initial}) + assert (await manager.get_all_mcp_servers_unfiltered())[0].allowed_tools == ["search"] + await manager.load_servers_from_config({"docs": {**initial, "allowed_tools": []}}) + assert (await manager.get_all_mcp_servers_unfiltered())[0].allowed_tools == [] + await manager.load_servers_from_config({}) + assert await manager.get_all_mcp_servers_unfiltered() == [] diff --git a/tests/unit/proxy/mcp_server/test_mcp_partial_update.py b/tests/unit/proxy/mcp_server/test_mcp_partial_update.py index 8e24e4fca80..83b14feb1f6 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_partial_update.py +++ b/tests/unit/proxy/mcp_server/test_mcp_partial_update.py @@ -26,7 +26,7 @@ from litellm.types.mcp_server.mcp_server_manager import PinnedMCPTool def _credentials_cleared(value) -> bool: - """The clear sentinel after the edge translation: prisma Json(None) (SQL null) or a bare None.""" + """The clear sentinel after the edge translation: prisma Json(None) or a bare None.""" return value is None or (isinstance(value, Json) and getattr(value, "data", "x") is None) @@ -1366,3 +1366,85 @@ async def test_issuer_edit_preserves_replacement_with_same_client_id(replacement assert credentials["dcr_server_url"] == "https://new.example/mcp" assert credentials.get("client_secret") == replacement.get("client_secret") assert credentials.get("token_endpoint_auth_method") == replacement.get("token_endpoint_auth_method") + + +@pytest.mark.parametrize( + "patch", + [ + {"allowed_tools": None}, + {"mcp_info": {"tool_allowlist_enforced": False}}, + {"tool_name_to_description": {"write": "Not approved"}}, + {"credentials": {"auth_value": "synthetic"}}, + {"gateway_custom_option": {"enabled": True}}, + ], +) +def test_provider_update_validates_the_merged_stored_registration(patch: dict[str, object]) -> None: + from litellm.proxy.mcp_server.db import _validate_provider_registration_update + + stored = models.LiteLLM_MCPServerTable.model_construct( + server_id="provider", execution_mode="provider", url="https://docs.example/mcp", transport=None, + allowed_tools=["search"], mcp_info={"tool_allowlist_enforced": True}, + ) + with pytest.raises(HTTPException) as error: + _validate_provider_registration_update(stored, UpdateMCPServerRequest(server_id="provider", **patch)) + assert error.value.status_code == 400 + + +def test_provider_partial_update_preserves_tools_and_empty_list_stays_enforced() -> None: + from litellm.proxy.mcp_server.db import _validate_provider_registration_update + + stored = models.LiteLLM_MCPServerTable.model_construct( + server_id="provider", execution_mode="provider", url="https://docs.example/mcp", transport=None, + allowed_tools=["search"], mcp_info={"tool_allowlist_enforced": True, "description": "Docs"}, + ) + renamed = _validate_provider_registration_update(stored, UpdateMCPServerRequest(server_id="provider", description="New")) + assert renamed.allowed_tools == ["search"] + assert renamed.mcp_info == {"tool_allowlist_enforced": True, "description": "Docs"} + cleared = _validate_provider_registration_update(stored, UpdateMCPServerRequest(server_id="provider", allowed_tools=[])) + assert cleared.allowed_tools == [] + assert cleared.mcp_info["tool_allowlist_enforced"] is True + + +@pytest.mark.parametrize("field", ["description", "instructions", "source_url", "tool_name_to_description"]) +def test_provider_update_clears_explicit_null_metadata_without_writing_omitted_fields(field: str) -> None: + from litellm.proxy.mcp_server.db import _prepare_provider_registration_update + + stored = models.LiteLLM_MCPServerTable.model_construct( + server_id="provider", execution_mode="provider", url="https://docs.example/mcp", transport=None, + allowed_tools=["search"], mcp_info={"tool_allowlist_enforced": True}, + description="Docs", instructions="Search", source_url="https://docs.example", + tool_name_to_description={"search": "Search"}, + ) + patch = UpdateMCPServerRequest(server_id="provider", **{field: None}) + serialized = _prepare_provider_registration_update(stored, patch, "admin", None) + assert serialized[field] == ("{}" if field == "tool_name_to_description" else None) + assert "allowed_tools" not in serialized + assert "url" not in serialized + assert json.loads(serialized["mcp_info"])["tool_allowlist_enforced"] is True + + +@pytest.mark.asyncio +async def test_provider_accepts_neutral_connection_settings_without_credentials_or_gateway_transport(): + neutral = { + "transport": None, "auth_type": None, "credentials": {}, "static_headers": {}, + "extra_headers": [], "env_vars": [], "is_byok": False, "command": None, + } + created = NewMCPServerRequest( + server_name="provider", execution_mode="provider", url="https://docs.example/mcp", + allowed_tools=["search"], **neutral, + ) + assert created.transport is None + assert created.credentials is None + prisma = _mock_prisma() + stored = models.LiteLLM_MCPServerTable.model_construct( + server_id="provider", execution_mode="provider", url=created.url, transport=None, + allowed_tools=["search"], mcp_info={"tool_allowlist_enforced": True}, + ) + prisma.db.litellm_mcpservertable.find_unique.return_value = stored + prisma.db.litellm_mcpservertable.update.return_value = stored + updated = await update_mcp_server(prisma, UpdateMCPServerRequest(server_id="provider", **neutral), "admin") + assert updated.allowed_tools == ["search"] + written = prisma.db.litellm_mcpservertable.update.call_args.kwargs["data"] + assert not neutral.keys() & written.keys() + assert "allowed_tools" not in written + assert json.loads(written["mcp_info"])["tool_allowlist_enforced"] is True diff --git a/tests/unit/proxy/mcp_server/test_rest_endpoints.py b/tests/unit/proxy/mcp_server/test_rest_endpoints.py index fdb1ff3cd09..26b55580ad4 100644 --- a/tests/unit/proxy/mcp_server/test_rest_endpoints.py +++ b/tests/unit/proxy/mcp_server/test_rest_endpoints.py @@ -5106,3 +5106,16 @@ async def test_rest_guardrail_bearer_excludes_admission_key( assert (authorization.removeprefix("Bearer ") or None) == expected finally: global_mcp_tool_registry.unregister_tools_with_prefix("guarded-") + + +@pytest.mark.asyncio +async def test_provider_registration_does_not_open_gateway_connection() -> None: + request: Final = NewMCPServerRequest( + execution_mode="provider", url="https://mcp.example.com/mcp", allowed_tools=["read"] + ) + + async def unexpected_operation(_client: object) -> dict[str, object]: + pytest.fail("Provider registrations must not open a gateway client") + + result: Final = await rest_endpoints._execute_with_mcp_client(request, unexpected_operation) + assert result == {"error": "Provider-executed MCP registrations do not support gateway connections"} diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 712b4134c05..0f1aa2243e7 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -6593,3 +6593,34 @@ async def test_proxy_config_validates_advertised_mcp_versions_at_load(tmp_path, return with pytest.raises(ValidationError): await ProxyConfig().load_config(router=None, config_file_path=str(config)) + + +@pytest.mark.asyncio +async def test_provider_config_management_and_removal_without_database(monkeypatch): + from starlette.requests import Request + + from litellm.proxy import proxy_server + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.mcp_management_endpoints import fetch_all_mcp_servers, fetch_mcp_server + from litellm.proxy.mcp_server import mcp_server_manager + from litellm.proxy.management_endpoints import mcp_management_endpoints + + manager = mcp_server_manager.MCPServerManager() + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) + monkeypatch.setattr(mcp_management_endpoints, "global_mcp_server_manager", manager) + monkeypatch.setattr(proxy_server, "prisma_client", None) + config = proxy_server.ProxyConfig() + await config._init_non_llm_configs({"mcp_servers": {"config_docs": { + "server_id": "config-docs", "execution_mode": "provider", "url": "https://docs.example/mcp", + "allowed_tools": ["search"], + }}}) + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + inventory = await fetch_all_mcp_servers(admin, team_id=None, user_id=None) + assert len(inventory) == 1 and inventory[0].is_config + request = Request({"type": "http", "headers": [], "client": ("127.0.0.1", 12345)}) + detail = await fetch_mcp_server(request, "config_docs", admin, include_reachability=True) + assert detail.server_id == "config-docs" and detail.allowed_tools == ["search"] + assert detail.last_health_check is None + await config._init_non_llm_configs({}) + assert await fetch_all_mcp_servers(admin, team_id=None, user_id=None) == [] + assert not manager.get_registry() diff --git a/tests/unit/proxy/test__types.py b/tests/unit/proxy/test__types.py index 7d57f246b1e..1ccf5d734ee 100644 --- a/tests/unit/proxy/test__types.py +++ b/tests/unit/proxy/test__types.py @@ -50,12 +50,18 @@ SERVER_ONLY_MARKERS = ( ) -@pytest.mark.parametrize("request_type", [NewMCPServerRequest, UpdateMCPServerRequest]) -def test_gateway_management_requests_do_not_ignore_provider_execution_mode( - request_type: type[NewMCPServerRequest] | type[UpdateMCPServerRequest], -) -> None: - with pytest.raises(ValidationError, match="execution_mode"): - request_type(server_id="docs", url="https://docs.example/mcp", execution_mode="provider") +def test_provider_management_create_normalizes_the_registration() -> None: + request: Final = NewMCPServerRequest( + server_id="docs", + url="https://DOCS.example:443/mcp", + execution_mode="provider", + allowed_tools=["search", "search"], + ) + assert request.url == "https://docs.example/mcp" + assert request.transport is None + assert request.credentials is None + assert request.allowed_tools == ["search"] + assert request.mcp_info == {"tool_allowlist_enforced": True} @pytest.mark.parametrize("mode", ["litellm", "provider", None]) diff --git a/tests/unit/types/mcp_server/test_provider_registration.py b/tests/unit/types/mcp_server/test_provider_registration.py new file mode 100644 index 00000000000..df038b08476 --- /dev/null +++ b/tests/unit/types/mcp_server/test_provider_registration.py @@ -0,0 +1,96 @@ +from typing import Final + +import pytest + +from litellm.proxy._types import NewMCPServerRequest +from litellm.proxy.mcp_server.db import _prepare_mcp_server_data +from litellm.types.mcp_server.provider_registration import canonical_provider_mcp_url, normalize_provider_registration + + +@pytest.mark.parametrize( + "url,expected", + [ + ("https://DOCS.example:443/mcp/tenant?scope=a", "https://docs.example/mcp/tenant?scope=a"), + ("https://docs.example:8443/mcp/", "https://docs.example:8443/mcp/"), + ("https://[::1]:443/mcp", "https://[::1]/mcp"), + ], +) +def test_endpoint_identity_preserves_path_query_and_nondefault_port(url: str, expected: str) -> None: + assert canonical_provider_mcp_url(url) == expected + + +@pytest.mark.parametrize( + "url", + [ + "http://docs.example/mcp", + "https://u:p@docs.example/mcp", + "https://docs.example/mcp#x", + "https://docs.example/mcp\n", + ], +) +def test_unsafe_endpoint_identity_is_rejected(url: str) -> None: + with pytest.raises(ValueError, match="Provider MCP url"): + canonical_provider_mcp_url(url) + + +@pytest.mark.parametrize( + "patch", + [ + {"allowed_tools": None}, + {"allowed_tools": ["*"]}, + {"allowed_tools": [" "]}, + {"transport": "http"}, + {"auth_type": "none"}, + {"credentials": {"auth_value": "synthetic"}}, + {"static_headers": {"Authorization": "synthetic"}}, + {"command": "python"}, + {"gateway_custom_option": {"enabled": True}}, + {"timeout": 0}, + {"mcp_info": {"tool_allowlist_enforced": False}}, + {"mcp_info": {"tool_allowlist_enforced": None}}, + {"tool_name_to_description": {"not_allowed": "Not in the list"}}, + ], +) +def test_provider_registration_rejects_unbounded_tools_and_connection_settings(patch: dict[str, object]) -> None: + with pytest.raises(ValueError, match=r"allowed_tools|valid list|Provider|Tool descriptions"): + normalize_provider_registration({"url": "https://docs.example/mcp", "allowed_tools": ["search"], **patch}) + + +def test_empty_tool_list_is_enforced_without_transport_or_stored_credentials() -> None: + request: Final = NewMCPServerRequest( + execution_mode="provider", url="https://docs.example/mcp", allowed_tools=[], mcp_info={"description": "Docs"} + ) + prepared: Final = _prepare_mcp_server_data(request) + assert prepared["allowed_tools"] == [] + assert request.mcp_info == {"description": "Docs", "tool_allowlist_enforced": True} + assert prepared["transport"] is None + assert not prepared.get("credentials") + assert prepared.get("auth_type") is None + + +@pytest.mark.parametrize("neutral", [None, False, "", {}, []]) +def test_provider_registration_ignores_neutral_unknown_settings(neutral: object) -> None: + request: Final = NewMCPServerRequest( + execution_mode="provider", url="https://docs.example/mcp", allowed_tools=["search"], + gateway_custom_option=neutral, + ) + assert request.allowed_tools == ["search"] + assert request.transport is None + assert "gateway_custom_option" not in request.model_dump() + + +def test_provider_registration_preserves_supported_management_fields() -> None: + request: Final = NewMCPServerRequest( + execution_mode="provider", server_name="docs", alias="documentation", + url="https://docs.example/mcp", allowed_tools=["search"], description="Documentation", + instructions="Search docs", source_url="https://docs.example", mcp_access_groups=["research"], + allow_all_keys=True, tool_name_to_description={"search": "Search docs"}, + tool_name_to_display_name={"search": "Search"}, mcp_info={"custom_metadata": "retained"}, + ) + assert request.server_name == "docs" and request.alias == "documentation" + assert request.description == "Documentation" and request.instructions == "Search docs" + assert request.source_url == "https://docs.example" and request.mcp_access_groups == ["research"] + assert request.allow_all_keys + assert request.tool_name_to_description == {"search": "Search docs"} + assert request.tool_name_to_display_name == {"search": "Search"} + assert request.mcp_info == {"custom_metadata": "retained", "tool_allowlist_enforced": True} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 111ada46919..6ecb9772a90 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -39540,12 +39540,8 @@ export interface components { }; /** Env Vars */ env_vars?: components["schemas"]["MCPEnvVar"][] | null; - /** - * Execution Mode - * @default litellm - * @constant - */ - execution_mode: "litellm"; + /** @default litellm */ + execution_mode: components["schemas"]["MCPExecutionMode"]; /** Extra Headers */ extra_headers?: string[] | null; /** Instructions */ @@ -39624,9 +39620,8 @@ export interface components { /** * Transport * @default sse - * @enum {string} */ - transport: "sse" | "http" | "stdio"; + transport: ("sse" | "http" | "stdio") | null; /** Url */ url?: string | null; }; @@ -49300,9 +49295,8 @@ export interface components { /** * Transport * @default sse - * @enum {string} */ - transport: "sse" | "http" | "stdio"; + transport: ("sse" | "http" | "stdio") | null; /** Url */ url?: string | null; }; From fab4eb29a30047d3b956d04fbfe9bb22f6f1282e Mon Sep 17 00:00:00 2001 From: mrinal Date: Tue, 3 Nov 2026 17:53:38 +0000 Subject: [PATCH 2/7] test(mcp): complete provider registration database doubles Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/proxy/mcp_server/test_db_credentials.py | 4 ++++ tests/unit/proxy/mcp_server/test_mcp_env_vars.py | 4 ++++ tests/unit/proxy/mcp_server/test_mcp_sigv4_auth.py | 1 + tests/unit/types/mcp_server/__init__.py | 0 4 files changed, 9 insertions(+) create mode 100644 tests/unit/types/mcp_server/__init__.py diff --git a/tests/unit/proxy/mcp_server/test_db_credentials.py b/tests/unit/proxy/mcp_server/test_db_credentials.py index 5be76def8d8..1ba18199a3c 100644 --- a/tests/unit/proxy/mcp_server/test_db_credentials.py +++ b/tests/unit/proxy/mcp_server/test_db_credentials.py @@ -510,6 +510,10 @@ class _MapTable: async def find_many(self, where: object = None) -> list[PrismaMCPServer]: return [_prisma_map_row(row, self.quoted) for row in self.rows.values()] + async def find_unique(self, *, where: dict[str, str]) -> PrismaMCPServer | None: + row: Final = self.rows.get(where["server_id"]) + return None if row is None else _prisma_map_row(row, self.quoted) + @pytest.mark.asyncio @pytest.mark.parametrize("field", ["static_headers", "env"]) diff --git a/tests/unit/proxy/mcp_server/test_mcp_env_vars.py b/tests/unit/proxy/mcp_server/test_mcp_env_vars.py index 87bd388124b..f03216e399b 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_env_vars.py +++ b/tests/unit/proxy/mcp_server/test_mcp_env_vars.py @@ -1562,6 +1562,7 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri mock_prisma.db.litellm_mcpservertable.create = AsyncMock( return_value=_prisma_row_with_json_string_env_vars() ) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None) created = await create_mcp_server( mock_prisma, @@ -1581,6 +1582,9 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri mock_prisma_upd.db.litellm_mcpservertable.update = AsyncMock( return_value=_prisma_row_with_json_string_env_vars() ) + mock_prisma_upd.db.litellm_mcpservertable.find_unique = AsyncMock( + return_value=_prisma_row_with_json_string_env_vars() + ) updated = await update_mcp_server( mock_prisma_upd, UpdateMCPServerRequest(server_id="srv-update"), diff --git a/tests/unit/proxy/mcp_server/test_mcp_sigv4_auth.py b/tests/unit/proxy/mcp_server/test_mcp_sigv4_auth.py index 28f96c3bdd2..dbb3ee84285 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_sigv4_auth.py +++ b/tests/unit/proxy/mcp_server/test_mcp_sigv4_auth.py @@ -678,6 +678,7 @@ class TestCredentialMergeOnUpdate: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row()) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=_updated_row()) data = UpdateMCPServerRequest( server_id="test-server", diff --git a/tests/unit/types/mcp_server/__init__.py b/tests/unit/types/mcp_server/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 931f73b09a83ae3cb3fa76992e11da07e00c24e9 Mon Sep 17 00:00:00 2001 From: mrinal Date: Tue, 3 Nov 2026 20:54:00 +0000 Subject: [PATCH 3/7] refactor(mcp): name provider-only catalog records explicitly Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/mcp_server/catalog.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/mcp_server/catalog.py b/litellm/proxy/mcp_server/catalog.py index 01ec5ea41af..4c44be6e4a4 100644 --- a/litellm/proxy/mcp_server/catalog.py +++ b/litellm/proxy/mcp_server/catalog.py @@ -126,7 +126,7 @@ def validate_provider_registration_conflicts( def _snapshot( manager: MCPServerManager, database_identity: str, - registration_only_servers: Mapping[str, LiteLLM_MCPServerTable], + provider_native_servers: Mapping[str, LiteLLM_MCPServerTable], ) -> CatalogSnapshot: from litellm.proxy.mcp_server.tool_registry import global_mcp_tool_registry @@ -150,7 +150,7 @@ def _snapshot( { **{key: value.model_copy(deep=True) for key, value in manager.catalog.config_registrations.items()}, **{key: manager.build_mcp_server_table(value) for key, value in servers.items()}, - **{key: value.model_copy(deep=True) for key, value in registration_only_servers.items()}, + **{key: value.model_copy(deep=True) for key, value in provider_native_servers.items()}, } ), ) @@ -161,7 +161,7 @@ class CatalogSnapshots: self.manager = manager self._refresh_lock = asyncio.Lock() self._database_identity = "" - self._registration_only_servers: Mapping[str, LiteLLM_MCPServerTable] = MappingProxyType({}) + self._provider_native_servers: Mapping[str, LiteLLM_MCPServerTable] = MappingProxyType({}) self.config_registrations: Mapping[str, LiteLLM_MCPServerTable] = MappingProxyType({}) self._arrival_ticket = 0 self._completed_ticket = 0 @@ -179,7 +179,7 @@ class CatalogSnapshots: def validate_config_inventory(self, servers: Sequence[LiteLLM_MCPServerTable]) -> None: database_servers: Final = ( *(self.manager.build_mcp_server_table(server) for server in self.manager.registry.values()), - *self._registration_only_servers.values(), + *self._provider_native_servers.values(), ) validate_provider_registration_conflicts(servers, database_servers) @@ -215,11 +215,11 @@ class CatalogSnapshots: from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - return _snapshot(self.manager, self._database_identity, self._registration_only_servers) + return _snapshot(self.manager, self._database_identity, self._provider_native_servers) from litellm.proxy.proxy_server import should_load_db_object if not should_load_db_object("mcp"): - return _snapshot(self.manager, self._database_identity, self._registration_only_servers) + return _snapshot(self.manager, self._database_identity, self._provider_native_servers) from litellm.proxy.mcp_server.db import get_mcp_catalog_revision try: @@ -360,7 +360,7 @@ class CatalogSnapshots: self._completed_ticket = covered raise else: - self._shared_snapshot = _snapshot(self.manager, self._database_identity, self._registration_only_servers) + self._shared_snapshot = _snapshot(self.manager, self._database_identity, self._provider_native_servers) self._applied_revision = revision self._completed_ticket = covered finally: @@ -562,7 +562,7 @@ class CatalogSnapshots: ).hexdigest() verbose_logger.info("Found %s MCP servers in database", len(raw_rows)) - registration_only_servers: Final = MappingProxyType( + provider_native_servers: Final = MappingProxyType( { server.server_id: server for server in ( @@ -576,7 +576,7 @@ class CatalogSnapshots: *self.config_registrations.values(), *(self.manager.build_mcp_server_table(value) for value in self.manager.config_mcp_servers.values()), ) - if self.config_registrations or registration_only_servers: + if self.config_registrations or provider_native_servers: validate_provider_registration_conflicts( config_servers, tuple(LiteLLM_MCPServerTable.model_validate(row.model_dump()) for row in raw_rows) ) @@ -659,7 +659,7 @@ class CatalogSnapshots: ", ".join(sorted(capturing_config_server_ids)), ) self._warned_capturing_config_server_ids = capturing_config_server_ids - self._registration_only_servers = registration_only_servers + self._provider_native_servers = provider_native_servers def catalog_operation( From 5f4d3f59208ba94e2504ff5578f2a9731d114a8d Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Tue, 3 Nov 2026 14:16:22 -0800 Subject: [PATCH 4/7] fix(mcp): reject non-ASCII provider MCP hostnames MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit IDNA2003, UTS-46 and WHATWG parse Unicode and percent-encoded hosts differently (faß.example -> fass.example vs xn--fa-hia.example), so the host matched against a registration could differ from the host the provider connects to. Require plain ASCII hosts; punycode stays allowed. --- litellm/types/mcp_server/provider_registration.py | 10 +++++++++- .../types/mcp_server/test_provider_registration.py | 4 ++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/litellm/types/mcp_server/provider_registration.py b/litellm/types/mcp_server/provider_registration.py index cf11b0914eb..56317b2de37 100644 --- a/litellm/types/mcp_server/provider_registration.py +++ b/litellm/types/mcp_server/provider_registration.py @@ -1,4 +1,5 @@ import json +import re from collections.abc import Collection, Mapping, Sequence from datetime import datetime from typing import Final, Literal @@ -12,6 +13,9 @@ from litellm.types.mcp import MCPExecutionMode _STR: Final = TypeAdapter(str) _MAP: Final = TypeAdapter(Mapping[str, object]) _URL: Final = TypeAdapter(AnyHttpUrl) +# Plain ASCII hosts only: IDNA2003, UTS-46 and WHATWG disagree on Unicode and percent-encoded +# hostnames, so the host we authorize could differ from the host the provider connects to. +_ASCII_HOST: Final = re.compile(r"[a-z0-9.\-:]+") def _is_neutral_registration_value(value: object) -> bool: @@ -107,7 +111,11 @@ def canonical_provider_mcp_url(value: object) -> str: parsed: Final = urlsplit(url) if parsed.scheme.lower() != "https" or not parsed.hostname or parsed.username is not None or "#" in url: raise ValueError("Provider MCP url must use HTTPS without userinfo or a fragment") - host: Final = parsed.hostname.lower().encode("idna").decode("ascii") + host: Final = parsed.hostname.lower() + if not _ASCII_HOST.fullmatch(host): + raise ValueError( + "Provider MCP url hostname must be plain ASCII; register internationalized hosts in punycode (xn--) form" + ) hostname: Final = f"[{host}]" if ":" in host else host port: Final = parsed.port authority: Final = f"{hostname}:{port}" if port is not None and port != 443 else hostname diff --git a/tests/unit/types/mcp_server/test_provider_registration.py b/tests/unit/types/mcp_server/test_provider_registration.py index df038b08476..aeb7f206d00 100644 --- a/tests/unit/types/mcp_server/test_provider_registration.py +++ b/tests/unit/types/mcp_server/test_provider_registration.py @@ -13,6 +13,7 @@ from litellm.types.mcp_server.provider_registration import canonical_provider_mc ("https://DOCS.example:443/mcp/tenant?scope=a", "https://docs.example/mcp/tenant?scope=a"), ("https://docs.example:8443/mcp/", "https://docs.example:8443/mcp/"), ("https://[::1]:443/mcp", "https://[::1]/mcp"), + ("https://xn--fa-hia.example/mcp", "https://xn--fa-hia.example/mcp"), ], ) def test_endpoint_identity_preserves_path_query_and_nondefault_port(url: str, expected: str) -> None: @@ -26,6 +27,9 @@ def test_endpoint_identity_preserves_path_query_and_nondefault_port(url: str, ex "https://u:p@docs.example/mcp", "https://docs.example/mcp#x", "https://docs.example/mcp\n", + "https://faß.example/mcp", + "https://fa%C3%9F.example/mcp", + "https://docs\uff0eexample/mcp", ], ) def test_unsafe_endpoint_identity_is_rejected(url: str) -> None: From 3375dc410b17dbca443aebcf85a72e92b6c123eb Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Tue, 3 Nov 2026 15:01:03 -0800 Subject: [PATCH 5/7] fix(mcp): resolve provider details by exact ID and prune metadata for removed tools GET /v1/mcp/server/{id} now finds database provider registrations by name or alias, and an exact server_id always wins over another row's name or alias. Changing a provider's allowed_tools without resending descriptions or display names drops entries for removed tools instead of rejecting the update. --- .../mcp_management_endpoints.py | 36 +++++++++++-------- litellm/proxy/mcp_server/db.py | 24 ++++++++++++- .../test_mcp_management_endpoints.py | 28 +++++++++++++++ .../mcp_server/test_mcp_partial_update.py | 26 ++++++++++++++ 4 files changed, 99 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index a323c277b52..67d816f5a4b 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1210,6 +1210,23 @@ if MCP_AVAILABLE: def _mcp_server_display_order(server: LiteLLM_MCPServerTable) -> tuple[str, str]: return ((server.server_name or server.alias or server.server_id).lower(), server.server_id) + def _provider_registration_for_identifier( + servers: Sequence[LiteLLM_MCPServerTable], identifier: str + ) -> LiteLLM_MCPServerTable | None: + """An exact server_id wins over every name or alias, so another row's alias cannot shadow it.""" + exact: Final = next((server for server in servers if server.server_id == identifier), None) + if exact is not None: + return exact if exact.execution_mode == MCPExecutionMode.provider else None + return next( + ( + server + for server in servers + if server.execution_mode == MCPExecutionMode.provider + and identifier in (server.server_name, server.alias) + ), + None, + ) + def _get_user_mcp_management_mode() -> UserMCPManagementMode: from litellm.proxy.proxy_server import ( general_settings as proxy_general_settings, @@ -1979,25 +1996,16 @@ if MCP_AVAILABLE: is_admin_view: Final = user_api_key_has_admin_view(user_api_key_dict) is_restricted_virtual_key: Final = _is_restricted_virtual_key_request(user_api_key_dict) snapshot: Final = global_mcp_server_manager.catalog.current() - config_match: Final = ( - next( - ( - server - for server in snapshot.management_servers.values() - if server.is_config - and server.execution_mode == MCPExecutionMode.provider - and server_id in (server.server_id, server.server_name, server.alias) - ), - None, - ) + provider_match: Final = ( + _provider_registration_for_identifier(tuple(snapshot.management_servers.values()), server_id) if snapshot is not None else None ) resolved: Final = ( - ResolvedMCPServer(table=config_match, runtime=None, source="registry") - if config_match is not None + ResolvedMCPServer(table=provider_match, runtime=None, source="registry") + if provider_match is not None and provider_match.is_config else await resolve_mcp_server( - server_id, + provider_match.server_id if provider_match is not None else server_id, manager=global_mcp_server_manager, db_lookup=lambda sid: get_mcp_server( get_prisma_client_or_throw("Database not connected. Connect a database to your proxy"), diff --git a/litellm/proxy/mcp_server/db.py b/litellm/proxy/mcp_server/db.py index ebcaa606831..be4d4a0f6b2 100644 --- a/litellm/proxy/mcp_server/db.py +++ b/litellm/proxy/mcp_server/db.py @@ -1286,6 +1286,27 @@ def _validate_mcp_protocol_write( raise HTTPException(status_code=400, detail=str(exc)) from exc +_TOOL_METADATA_FIELDS: Final = ("tool_name_to_description", "tool_name_to_display_name") + + +def _tool_metadata_for_new_allowlist( + current: Mapping[str, object], patch: Mapping[str, object] +) -> Mapping[str, object]: + allowed: Final = patch.get("allowed_tools") + if not isinstance(allowed, Sequence) or isinstance(allowed, str): + return {} + kept: Final = frozenset(TypeAdapter(tuple[str, ...]).validate_python(allowed)) + return { + field: { + tool: value + for tool, value in TypeAdapter(dict[str, object]).validate_python(stored).items() + if tool in kept + } + for field in _TOOL_METADATA_FIELDS + if field not in patch and isinstance(stored := current.get(field), Mapping) + } + + def _validate_provider_registration_update( stored: "prisma_db_models.LiteLLM_MCPServerTable", data: UpdateMCPServerRequest, @@ -1294,7 +1315,8 @@ def _validate_provider_registration_update( try: patch: Final = data.provider_registration_patch(fields_set) registration: Final = LiteLLM_MCPServerTable.model_validate(stored.model_dump()) - merged: Final = {**registration.model_dump(include=set(NewMCPServerRequest.model_fields)), **patch} + current: Final = registration.model_dump(include=set(NewMCPServerRequest.model_fields)) + merged: Final = {**current, **_tool_metadata_for_new_allowlist(current, patch), **patch} return NewMCPServerRequest.model_validate(merged) except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc 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 6b2a1a3ffc4..47dcd94ea44 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -218,6 +218,34 @@ async def test_mcp_publication_list_and_detail_derive_current_status( assert record.mcp_info == original_metadata +def _provider_row(server_id: str, **fields: object) -> LiteLLM_MCPServerTable: + return LiteLLM_MCPServerTable( + server_id=server_id, execution_mode="provider", transport=None, + url=f"https://{server_id}.example/mcp", allowed_tools=["search"], **fields, + ) + + +@pytest.mark.parametrize( + "identifier,expected", + ( + ("native-id", "native-id"), + ("docs", "native-id"), + ("docs_alias", "native-id"), + ("gateway-id", None), + ("unknown", None), + ), +) +def test_provider_detail_lookup_prefers_exact_ids_over_other_rows_aliases(identifier: str, expected: str | None) -> None: + rows: Final = ( + _provider_row("shadow-id", alias="native-id"), + _provider_row("gateway-shadow", alias="gateway-id"), + _provider_row("native-id", server_name="docs", alias="docs_alias"), + LiteLLM_MCPServerTable(server_id="gateway-id", transport=MCPTransport.http, url="https://gw.example/mcp"), + ) + match: Final = mgmt_endpoints._provider_registration_for_identifier(rows, identifier) + assert (match.server_id if match is not None else None) == expected + + @pytest.mark.parametrize("approval_status", ("pending_review", "rejected", "draft", "active")) @pytest.mark.parametrize("strict", (False, True)) def test_mcp_publication_projection_excludes_unregistered_lifecycle_records( diff --git a/tests/unit/proxy/mcp_server/test_mcp_partial_update.py b/tests/unit/proxy/mcp_server/test_mcp_partial_update.py index 83b14feb1f6..ea104f581dd 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_partial_update.py +++ b/tests/unit/proxy/mcp_server/test_mcp_partial_update.py @@ -1405,6 +1405,32 @@ def test_provider_partial_update_preserves_tools_and_empty_list_stays_enforced() assert cleared.mcp_info["tool_allowlist_enforced"] is True +@pytest.mark.parametrize( + "allowed_tools,descriptions,display_names", + [ + ([], {}, {}), + (["fetch"], {"fetch": "Fetch a page"}, {"fetch": "Fetch"}), + ], +) +def test_provider_allowlist_change_drops_metadata_for_removed_tools( + allowed_tools: list[str], descriptions: dict[str, str], display_names: dict[str, str] +) -> None: + from litellm.proxy.mcp_server.db import _validate_provider_registration_update + + stored = models.LiteLLM_MCPServerTable.model_construct( + server_id="provider", execution_mode="provider", url="https://docs.example/mcp", transport=None, + allowed_tools=["search", "fetch"], mcp_info={"tool_allowlist_enforced": True}, + tool_name_to_description={"search": "Search docs", "fetch": "Fetch a page"}, + tool_name_to_display_name={"search": "Search", "fetch": "Fetch"}, + ) + updated = _validate_provider_registration_update( + stored, UpdateMCPServerRequest(server_id="provider", allowed_tools=allowed_tools) + ) + assert (updated.allowed_tools, updated.tool_name_to_description, updated.tool_name_to_display_name) == ( + allowed_tools, descriptions, display_names, + ) + + @pytest.mark.parametrize("field", ["description", "instructions", "source_url", "tool_name_to_description"]) def test_provider_update_clears_explicit_null_metadata_without_writing_omitted_fields(field: str) -> None: from litellm.proxy.mcp_server.db import _prepare_provider_registration_update From 7c27414bcca4a898edc2737997f92cd09c765a70 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Tue, 3 Nov 2026 15:19:32 -0800 Subject: [PATCH 6/7] fix(mcp): persist pruned tool metadata when a provider allowlist changes The update validated the merged registration with metadata for removed tools dropped, but wrote only the requested fields, so the stored descriptions and display names kept entries for tools no longer allowed. Write the stored metadata fields along with allowed_tools. --- litellm/proxy/mcp_server/db.py | 12 ++++++++++-- .../proxy/mcp_server/test_mcp_partial_update.py | 16 +++++++++------- 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/mcp_server/db.py b/litellm/proxy/mcp_server/db.py index be4d4a0f6b2..b09123b2f1a 100644 --- a/litellm/proxy/mcp_server/db.py +++ b/litellm/proxy/mcp_server/db.py @@ -1332,13 +1332,21 @@ def _prepare_provider_registration_update( validated: Final = _validate_provider_registration_update(stored, data, fields_set) _validate_mcp_config_registration_write(validated.model_dump()) - supplied_fields: Final = fields_set if fields_set is not None else data.model_fields_set + requested_fields: Final = fields_set if fields_set is not None else data.model_fields_set + # A changed allowlist rewrites stored tool metadata too, so entries for removed tools are dropped. + pruned_metadata: Final[frozenset[str]] = ( + frozenset(field for field in _TOOL_METADATA_FIELDS if getattr(stored, field, None) is not None) + - requested_fields + if "allowed_tools" in requested_fields + else frozenset() + ) + supplied_fields: Final[set[str]] = {*requested_fields, *pruned_metadata} normalized_patch: Final = UpdateMCPServerRequest.model_validate( validated.model_dump( include=set(supplied_fields & ProviderMCPRegistration.model_fields.keys()) | {"server_id", "mcp_info"} ) ) - serialized: Final = _prepare_mcp_server_data(normalized_patch, exclude_unset=True, fields_set=set(supplied_fields)) + serialized: Final = _prepare_mcp_server_data(normalized_patch, exclude_unset=True, fields_set=supplied_fields) return { **serialized, "updated_by": touched_by, diff --git a/tests/unit/proxy/mcp_server/test_mcp_partial_update.py b/tests/unit/proxy/mcp_server/test_mcp_partial_update.py index ea104f581dd..42a046bff62 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_partial_update.py +++ b/tests/unit/proxy/mcp_server/test_mcp_partial_update.py @@ -1412,10 +1412,10 @@ def test_provider_partial_update_preserves_tools_and_empty_list_stays_enforced() (["fetch"], {"fetch": "Fetch a page"}, {"fetch": "Fetch"}), ], ) -def test_provider_allowlist_change_drops_metadata_for_removed_tools( +def test_provider_allowlist_change_writes_metadata_without_removed_tools( allowed_tools: list[str], descriptions: dict[str, str], display_names: dict[str, str] ) -> None: - from litellm.proxy.mcp_server.db import _validate_provider_registration_update + from litellm.proxy.mcp_server.db import _prepare_provider_registration_update stored = models.LiteLLM_MCPServerTable.model_construct( server_id="provider", execution_mode="provider", url="https://docs.example/mcp", transport=None, @@ -1423,12 +1423,14 @@ def test_provider_allowlist_change_drops_metadata_for_removed_tools( tool_name_to_description={"search": "Search docs", "fetch": "Fetch a page"}, tool_name_to_display_name={"search": "Search", "fetch": "Fetch"}, ) - updated = _validate_provider_registration_update( - stored, UpdateMCPServerRequest(server_id="provider", allowed_tools=allowed_tools) - ) - assert (updated.allowed_tools, updated.tool_name_to_description, updated.tool_name_to_display_name) == ( - allowed_tools, descriptions, display_names, + written = _prepare_provider_registration_update( + stored, UpdateMCPServerRequest(server_id="provider", allowed_tools=allowed_tools), "admin", None ) + assert ( + written["allowed_tools"], + json.loads(written["tool_name_to_description"]), + json.loads(written["tool_name_to_display_name"]), + ) == (allowed_tools, descriptions, display_names) @pytest.mark.parametrize("field", ["description", "instructions", "source_url", "tool_name_to_description"]) From ecbb53dbc99a5b32f456e839780ea53eea76a184 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Tue, 3 Nov 2026 17:35:14 -0800 Subject: [PATCH 7/7] fix(mcp): reject duplicate provider URLs without relying on the index Provider create and update now check for another provider registration with the same canonical URL inside the existing identifier transaction, under an advisory lock on that URL, and return 409. Setups that skip migration SQL (prisma db push, adopted databases) have no partial unique index, so they previously accepted duplicates. --- .../mcp_management_endpoints.py | 7 ++- litellm/proxy/mcp_server/db.py | 54 +++++++++++++++++-- tests/integration/mcp/test_mcp_management.py | 38 ++++++++++++- 3 files changed, 92 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 67d816f5a4b..042f334cade 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -335,6 +335,11 @@ if MCP_AVAILABLE: _validate_upstream_token_header(payload) def mcp_identifier_conflict_message(conflict: McpIdentifierConflict) -> str: + if conflict.field == "url": + return ( + f"Provider MCP registration {conflict.server_id} already uses {conflict.value}. " + "Each provider endpoint can be registered once." + ) return ( f"An MCP server with {conflict.field} '{conflict.value}' already exists " f"(server_id={conflict.server_id}). " @@ -343,7 +348,7 @@ if MCP_AVAILABLE: def raise_mcp_identifier_conflict(conflict: McpIdentifierConflict) -> NoReturn: raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, + status_code=status.HTTP_409_CONFLICT if conflict.field == "url" else status.HTTP_400_BAD_REQUEST, detail={"error": mcp_identifier_conflict_message(conflict)}, ) diff --git a/litellm/proxy/mcp_server/db.py b/litellm/proxy/mcp_server/db.py index b09123b2f1a..cd856745549 100644 --- a/litellm/proxy/mcp_server/db.py +++ b/litellm/proxy/mcp_server/db.py @@ -95,7 +95,7 @@ class McpIdentifierConflict: string, and ``server_id`` the existing row that owns it. """ - field: Literal["server_name", "alias"] + field: Literal["server_name", "alias", "url"] value: str server_id: str @@ -640,6 +640,16 @@ def _identifier_field(data_dict: "Mapping[str, object]", field: str) -> str | No return value if isinstance(value, str) else None +def _canonical_provider_url_or_none(url: object) -> str | None: + """The identity a provider URL update must not share; None leaves an invalid URL to update validation.""" + from litellm.types.mcp_server.provider_registration import canonical_provider_mcp_url + + try: + return canonical_provider_mcp_url(url) + except ValueError: + return None + + async def _find_mcp_server_identifier_conflict( table: "TableActions[prisma_db_models.LiteLLM_MCPServerTable]", *, @@ -698,6 +708,26 @@ def _mcp_identifier_lock_keys(*identifiers: str | None) -> tuple[int, ...]: ) +async def _find_provider_url_conflict( + table: "TableActions[prisma_db_models.LiteLLM_MCPServerTable]", + *, + provider_url: str | None, + exclude_server_id: str | None, +) -> McpIdentifierConflict | None: + if not provider_url: + return None + from prisma.enums import MCPExecutionMode as StoredExecutionMode + + own_row_guard: Final[tuple[prisma_db_types.LiteLLM_MCPServerTableWhereInput, ...]] = ( + ({"NOT": [{"server_id": exclude_server_id}]},) if exclude_server_id is not None else () + ) + same_url: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = { + "AND": [{"execution_mode": StoredExecutionMode.provider}, {"url": provider_url}, *own_row_guard] + } + row: Final = await table.find_first(where=same_url) + return None if row is None else McpIdentifierConflict(field="url", value=provider_url, server_id=row.server_id) + + async def _mcp_server_write_if_identifier_free( prisma_client: PrismaClient, *, @@ -706,14 +736,19 @@ async def _mcp_server_write_if_identifier_free( exclude_server_id: str | None, write: "Callable[[TableActions[prisma_db_models.LiteLLM_MCPServerTable]], Awaitable[prisma_db_models.LiteLLM_MCPServerTable | None]]", lock_server_id: str | None = None, + provider_url: str | None = None, ) -> "prisma_db_models.LiteLLM_MCPServerTable | McpIdentifierConflict | None": - """Run ``write`` only when no other live row owns ``server_name``/``alias``. + """Run ``write`` only when no other live row owns ``server_name``/``alias``, nor, + for a provider registration, its canonical ``provider_url``. The conflict check and the write share a transaction guarded by per-identifier - advisory locks, so two concurrent requests for the same name cannot both - pass the check and both insert. + advisory locks, so two concurrent requests for the same name or provider URL + cannot both pass the check and both write. This does not depend on the partial + unique index, which setups that skip migration SQL never create. """ - lock_keys: Final = _mcp_identifier_lock_keys(server_name, alias) + lock_keys: Final = _mcp_identifier_lock_keys( + server_name, alias, f"provider_url:{provider_url}" if provider_url else None + ) async with _db_transaction_manager(prisma_client) as tx: for lock_key in lock_keys: await tx.execute_raw("SELECT pg_advisory_xact_lock($1::bigint)", lock_key) @@ -722,6 +757,8 @@ async def _mcp_server_write_if_identifier_free( server_name=server_name, alias=alias, exclude_server_id=exclude_server_id, + ) or await _find_provider_url_conflict( + tx.litellm_mcpservertable, provider_url=provider_url, exclude_server_id=exclude_server_id ) if conflict is not None: return conflict @@ -1165,6 +1202,7 @@ async def create_mcp_server_if_identifier_free( alias=_identifier_field(data_dict, "alias"), exclude_server_id=None, write=_create, + provider_url=_identifier_field(data_dict, "url") if data.execution_mode == MCPExecutionMode.provider else None, ) if isinstance(written, McpIdentifierConflict): return written @@ -1360,6 +1398,7 @@ async def _update_mcp_server_row( data_dict: Mapping[str, object], prepare_update: Callable[["prisma_db_models.LiteLLM_MCPServerTable"], Mapping[str, object]] | None = None, expected_updated_at: datetime | None = None, + provider_url: str | None = None, ) -> "prisma_db_models.LiteLLM_MCPServerTable | McpIdentifierConflict | None": identifier_write: Final = any(field in data_dict for field in ("server_name", "alias")) protocol_write: Final = bool({"transport", "mcp_info"}.intersection(data_dict)) @@ -1399,6 +1438,7 @@ async def _update_mcp_server_row( exclude_server_id=server_id, write=_update, lock_server_id=server_id if needs_stored_row else None, + provider_url=provider_url, ) return await _mcp_server_write_if_identifier_free( prisma_client, @@ -1407,6 +1447,7 @@ async def _update_mcp_server_row( exclude_server_id=server_id, write=_update, lock_server_id=server_id if needs_stored_row else None, + provider_url=provider_url, ) @@ -1433,6 +1474,9 @@ async def update_mcp_server( data_dict=data.model_dump(include=fields_set, exclude_unset=True), prepare_update=lambda stored: _prepare_provider_registration_update(stored, data, touched_by, fields_set), expected_updated_at=expected_updated_at, + provider_url=_canonical_provider_url_or_none(data.url) + if "url" in (fields_set or data.model_fields_set) + else None, ) return ( updated_provider diff --git a/tests/integration/mcp/test_mcp_management.py b/tests/integration/mcp/test_mcp_management.py index 90435f49b34..bcd8f3001ef 100644 --- a/tests/integration/mcp/test_mcp_management.py +++ b/tests/integration/mcp/test_mcp_management.py @@ -7,7 +7,7 @@ from typing import Final import pytest import yaml from integration._support.client import JSON_OBJECT, Gateway, eventually -from integration._support.database import read_rows +from integration._support.database import read_rows, scratch_database from integration._support.mcp import ( McpCaller, McpPeer, @@ -209,6 +209,42 @@ def _openapi_server_lists_and_calls_only_its_own_tools( assert [(item["method"], item["path"]) for item in peer.drain()] == [("GET", "/pets/7")], identity +def test_duplicate_provider_url_is_rejected_where_db_push_skipped_the_unique_index( + gateway: Gateway, tmp_path: Path +) -> None: + import concurrent.futures + + config: Final = tmp_path / "pushed.yaml" + config.write_text(yaml.safe_dump({"model_list": [], "general_settings": {"master_key": gateway.key}})) + with scratch_database("provider_url") as database_url, owned_proxy( + gateway, tmp_path, {"DATABASE_URL": database_url}, config=config + ) as pushed: + index: Final = read_rows( + "SELECT to_regclass('\"LiteLLM_MCPServerTable_provider_url_key\"')::text AS name", (), database_url=database_url + ) + assert index == [{"name": None}], "db push should leave the SQL-only partial index out" + url: Final = f"https://{uuid.uuid4().hex[:12]}.example/mcp" + + def register(name: str, endpoint: str) -> int: + return pushed.request( + "POST", + "/v1/mcp/server", + {"server_id": name, "server_name": name, "execution_mode": "provider", "url": endpoint, "allowed_tools": ["search"]}, + ).status_code + + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool: + statuses: Final = sorted(pool.map(partial(register, endpoint=url), (f"prov{n}" for n in range(8)))) + assert statuses == [201] + [409] * 7, statuses + stored: Final = read_rows( + 'SELECT server_id FROM "LiteLLM_MCPServerTable" WHERE url = %s', (url,), database_url=database_url + ) + assert len(stored) == 1, stored + + assert register("prov_other", f"https://other-{uuid.uuid4().hex[:8]}.example/mcp") == 201 + moved: Final = pushed.request("PUT", "/v1/mcp/server", {"server_id": "prov_other", "url": url}) + assert moved.status_code == 409, moved.text + + def test_openapi_listing_is_scoped_to_the_exact_alias_when_aliases_overlap(gateway: Gateway) -> None: with openapi_peer() as short, openapi_peer() as long, gateway.scenario() as scenario: stem: Final = "pet" + uuid.uuid4().hex[:8]