diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 25421ab2b52..1551e575f04 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -37597,6 +37597,14 @@ "title": "MCPEnvVarScope", "type": "string" }, + "MCPExecutionMode": { + "enum": [ + "litellm", + "provider" + ], + "title": "MCPExecutionMode", + "type": "string" + }, "NewMCPServerRequest": { "properties": { "alias": { @@ -37788,10 +37796,8 @@ "title": "Env Vars" }, "execution_mode": { - "const": "litellm", - "default": "litellm", - "title": "Execution Mode", - "type": "string" + "$ref": "#/components/schemas/MCPExecutionMode", + "default": "litellm" }, "extra_headers": { "anyOf": [ @@ -38079,14 +38085,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": [ @@ -40641,10 +40654,8 @@ "title": "Env Vars" }, "execution_mode": { - "const": "litellm", - "default": "litellm", - "title": "Execution Mode", - "type": "string" + "$ref": "#/components/schemas/MCPExecutionMode", + "default": "litellm" }, "extra_headers": { "anyOf": [ @@ -40932,14 +40943,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": [ @@ -44709,10 +44727,8 @@ "title": "Env Vars" }, "execution_mode": { - "const": "litellm", - "default": "litellm", - "title": "Execution Mode", - "type": "string" + "$ref": "#/components/schemas/MCPExecutionMode", + "default": "litellm" }, "extra_headers": { "anyOf": [ @@ -45000,14 +45016,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": [ @@ -45523,14 +45546,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": [ @@ -47988,6 +48018,14 @@ "title": "MCPEnvVarScope", "type": "string" }, + "MCPExecutionMode": { + "enum": [ + "litellm", + "provider" + ], + "title": "MCPExecutionMode", + "type": "string" + }, "NewMCPServerRequest": { "properties": { "alias": { @@ -48179,10 +48217,8 @@ "title": "Env Vars" }, "execution_mode": { - "const": "litellm", - "default": "litellm", - "title": "Execution Mode", - "type": "string" + "$ref": "#/components/schemas/MCPExecutionMode", + "default": "litellm" }, "extra_headers": { "anyOf": [ @@ -48470,14 +48506,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 d97ccb9c99b..c11ae3deb78 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 @@ -1414,12 +1417,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 @@ -1480,6 +1483,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") @@ -1491,7 +1498,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") @@ -1523,11 +1538,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 @@ -1572,6 +1589,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: @@ -1581,7 +1620,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..042f334cade 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, @@ -332,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}). " @@ -340,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)}, ) @@ -1207,6 +1215,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, @@ -1342,13 +1367,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 +1995,31 @@ 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() + 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=provider_match, runtime=None, source="registry") + if provider_match is not None and provider_match.is_config + else await resolve_mcp_server( + 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"), + sid, + ), + id_client_ip=client_ip, + name_client_ip=client_ip, + match_name=True, + ) ) authorized: Final = await authorize_mcp_server( resolved, @@ -1992,18 +2049,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 +2147,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 +3256,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 +3368,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..4c44be6e4a4 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, + provider_native_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 provider_native_servers.items()}, + } + ), ) @@ -116,6 +161,8 @@ class CatalogSnapshots: self.manager = manager self._refresh_lock = asyncio.Lock() self._database_identity = "" + 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 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._provider_native_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._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) + return _snapshot(self.manager, self._database_identity, self._provider_native_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._provider_native_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)) + provider_native_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 provider_native_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._provider_native_servers = provider_native_servers def catalog_operation( diff --git a/litellm/proxy/mcp_server/db.py b/litellm/proxy/mcp_server/db.py index 27f8b90a586..cd856745549 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: @@ -88,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 @@ -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) @@ -631,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]", *, @@ -689,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, *, @@ -697,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) @@ -713,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 @@ -1076,16 +1122,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 +1186,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 @@ -1125,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 @@ -1246,35 +1324,106 @@ 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, + 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()) + 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 + + +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()) + 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=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, + 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)) + 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 +1437,8 @@ 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, + provider_url=provider_url, ) return await _mcp_server_write_if_identifier_free( prisma_client, @@ -1296,7 +1446,8 @@ 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, + provider_url=provider_url, ) @@ -1315,6 +1466,27 @@ 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, + provider_url=_canonical_provider_url_or_none(data.url) + if "url" in (fields_set or data.model_fields_set) + else None, + ) + 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 +1494,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 +1616,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 c1af5d6f980..f17dc96afa3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7669,17 +7669,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..56317b2de37 --- /dev/null +++ b/litellm/types/mcp_server/provider_registration.py @@ -0,0 +1,137 @@ +import json +import re +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) +# 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: + 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() + 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 + 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/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] 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..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( @@ -408,6 +436,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 +684,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 +804,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 +897,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 +2055,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 +8129,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_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_partial_update.py b/tests/unit/proxy/mcp_server/test_mcp_partial_update.py index 8e24e4fca80..42a046bff62 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,113 @@ 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( + "allowed_tools,descriptions,display_names", + [ + ([], {}, {}), + (["fetch"], {"fetch": "Fetch a page"}, {"fetch": "Fetch"}), + ], +) +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 _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", "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"}, + ) + 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"]) +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_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/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/__init__.py b/tests/unit/types/mcp_server/__init__.py new file mode 100644 index 00000000000..e69de29bb2d 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..aeb7f206d00 --- /dev/null +++ b/tests/unit/types/mcp_server/test_provider_registration.py @@ -0,0 +1,100 @@ +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"), + ("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: + 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", + "https://faß.example/mcp", + "https://fa%C3%9F.example/mcp", + "https://docs\uff0eexample/mcp", + ], +) +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 ee8c7c745ad..925cc4886d8 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -40121,12 +40121,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 */ @@ -40205,9 +40201,8 @@ export interface components { /** * Transport * @default sse - * @enum {string} */ - transport: "sse" | "http" | "stdio"; + transport: ("sse" | "http" | "stdio") | null; /** Url */ url?: string | null; }; @@ -50276,9 +50271,8 @@ export interface components { /** * Transport * @default sse - * @enum {string} */ - transport: "sse" | "http" | "stdio"; + transport: ("sse" | "http" | "stdio") | null; /** Url */ url?: string | null; };