From 1ba99633e40354d5a95d96c57b13d9643181bbe0 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Sat, 31 Oct 2026 12:57:01 -0700 Subject: [PATCH] feat: add provider MCP registration management --- docs/plans/provider-executed-mcp.md | 63 ++++++++-- litellm/proxy/_types.py | 24 +++- .../mcp_management_endpoints.py | 71 +++++++---- litellm/proxy/mcp_server/catalog.py | 33 +++++- litellm/proxy/mcp_server/db.py | 88 +++++++++++--- .../proxy/mcp_server/mcp_server_manager.py | 12 +- litellm/types/mcp.py | 5 +- .../types/mcp_server/mcp_server_manager.py | 3 +- .../types/mcp_server/provider_registration.py | 110 ++++++++++++++++++ .../test_mcp_management_endpoints.py | 6 + tests/unit/proxy/mcp_server/test_catalog.py | 40 +++++++ .../mcp_server/test_mcp_partial_update.py | 84 ++++++++++++- tests/unit/proxy/test__types.py | 18 ++- .../mcp_server/test_provider_registration.py | 67 +++++++++++ 14 files changed, 548 insertions(+), 76 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/docs/plans/provider-executed-mcp.md b/docs/plans/provider-executed-mcp.md index b8236df521b..a63e6559621 100644 --- a/docs/plans/provider-executed-mcp.md +++ b/docs/plans/provider-executed-mcp.md @@ -1,6 +1,6 @@ # Provider-executed MCP registration -Status: database slice approved and committed, runtime isolation slice prepared for local diff review +Status: database, isolation, and admin registration API slices approved. Scoped retrieval and provider config loading are intentionally deferred ## Implementation plan @@ -9,8 +9,8 @@ Work through these layers in order on `litellm_provider_mcp_registration`. Updat | Layer | Scope | Status | | --- | --- | --- | | 1 | Settle registration fields, values, storage, and validation | Approved | -| 2 | Implement persistence and execution-mode isolation before enabling provider registration | Database slice committed; typed mode and gateway isolation ready for review; persistence validation pending | -| 3 | Management UI/API, stored tool options, and existing key/team/toolset selectors | Pending | +| 2 | Implement persistence and execution-mode isolation before enabling provider registration | Database and isolation committed; API persistence validation approved; YAML provider inventory pending | +| 3 | Management UI/API, stored tool options, and existing key/team/toolset selectors | Admin CRUD and inventory API approved; UI and scoped selectors pending | | 4 | Authorize effective OpenAI Responses and Anthropic Messages MCP attachments through existing MCP permissions | Pending | | 5 | Live provider tests, streaming admission tests, UI tests, and gateway regressions | Pending | @@ -18,6 +18,26 @@ V1 supports URL-based OpenAI Responses MCP and both Anthropic Messages connector Response/SSE inspection, per-invocation argument policy inside a provider's tool loop, automatic MCP discovery for provider registrations, local MCP references, provider connector IDs, and additional API bridges are outside v1 +## Current review acceptance criteria + +This batch covers admin CRUD, input validation, cached management inventory, and execution isolation. Key/team-scoped retrieval, selectors, toolsets, and native-request admission remain later layers. Their absence is intentional for this batch + +Core provider fields remain `execution_mode`, `server_name`, `url`, `allowed_tools`, and optional `tool_name_to_description`, alongside the existing shared management fields. Neutral connection values remain accepted where the shared schema permits them, but never enable MCP execution. Active credentials, transports, OAuth settings, commands, and other connection configuration are rejected. No strict five-field schema rewrite is part of this revision + +| Acceptance area | Required proof in this batch | +| --- | --- | +| Admin CRUD | Create, list, detail, partial update, delete, missing/deleted IDs, duplicate URL, immutable mode | +| Input validation | Active connection settings reject; supported neutral settings succeed; rejected updates leave the row unchanged; omitted tools preserve scope and an explicit empty list stays enforced | +| Cached inventory | Provider registrations appear in admin inventory and refresh after updates and deletion without becoming executable | +| Execution isolation | No provider initialization, discovery, health probe, or tool call from management or gateway endpoints | +| Compatibility | Existing LiteLLM credential merging and update transaction rules remain in place | + +Provider updates keep read-merge-validation atomic through the existing writer. The gateway credential path is not moved into a new blanket transaction. One initial row lookup selects the immutable execution mode and is reused when the gateway credential flow needs stored values + +The later permissions layer will test matching list/detail visibility under key/team intersection, explicit denial, disjoint scopes, revocation after caching, and separation of read grants from mutation privileges. The admin-only live run is not evidence for that layer + +The next config slice should accept provider entries under the existing `mcp_servers` key, reuse registration validation, and publish read-only management records in the shared catalog without constructing gateway clients. Verify startup and reload, removal, stable IDs, conflicting config/database identities or URLs, and zero upstream activity. Existing LiteLLM config loading must remain unchanged + ## Current invocation routing On the Responses path, `type: "mcp"` alone does not tell LiteLLM to connect to the server. An external `server_url` such as `https://custom.example/mcp` stays in the provider-bound tools. On a provider/API supporting native MCP, the provider discovers and calls that endpoint, subject to its own connectivity, authentication, and approval requirements. An unsupported provider/API may reject the attachment. LiteLLM does not automatically become the MCP executor @@ -78,7 +98,7 @@ The model excerpt shows only the changed fields. All other columns remain in pla | `mcp_info.tool_allowlist_enforced` | Existing marker, persisted as `true` for provider mode so `[]` means deny all | | `tool_name_to_description` | Existing optional descriptions for names in `allowed_tools` | | `transport` | SQL `NULL`. LiteLLM neither chooses nor opens the provider's MCP transport | -| `auth_type`, `credentials` | SQL `NULL`. LiteLLM does not manage this registration's upstream authentication | +| `auth_type`, `credentials` | No upstream authentication settings or credentials. Existing empty storage representations are acceptable | | Other connection-only settings | Neutral existing defaults, with no active connection configuration accepted | | Object permissions and toolsets | Existing tables and existing server IDs/tool names | @@ -190,7 +210,7 @@ Existing `metadata.allowed_tools` remains an independent compatibility restricti | Area | Required proof | | --- | --- | -| Schema and compatibility | Old registrations load as `litellm`; provider records round-trip with null transport/credentials and exact allowed tools | +| Schema and compatibility | Old registrations load as `litellm`; provider records round-trip with null transport, no upstream credentials, and exact allowed tools | | Field validation | Mode/URL/list validation; partial updates; enforced empty list; duplicate endpoint rejection | | Management | API and UI create/edit/select the same server and tool IDs without live MCP discovery | | Toolsets | Mixed-mode toolsets remain assignable; grants cannot change execution mode or exceed server/team ceilings | @@ -199,14 +219,41 @@ Existing `metadata.allowed_tools` remains an independent compatibility restricti | API coverage | OpenAI Responses and both Anthropic Messages formats; request-supplied and model-configured attachments; streaming and non-streaming | | Live evidence | Real providers execute authorized read-only MCP tools; denied requests produce zero provider dispatches; existing gateway execution still works | -The approved database slice was committed as `fd0ac3ca76`. The next slice adds shared execution-mode types, stored-record transport validation, and gateway isolation. Stored records can represent both modes, while the connection-capable `MCPServer` type permits only `litellm` +The approved database slice was committed as `fd0ac3ca76` and the approved isolation slice as `42ea534932`. Stored records can represent both modes, while the connection-capable `MCPServer` type permits only `litellm`. The redundant named gateway-mode alias is replaced with a literal annotation, leaving one named execution-mode enum -Provider records are excluded during incremental add/update and database catalog reload, before cached gateway entries can be reused. Direct table-to-runtime conversion and both shared client-creation boundaries reject provider mode before OAuth discovery, credential resolution, or MCP connection. Config loading currently rejects provider mode before discovery. The create request schema currently accepts only `litellm`, so `execution_mode=provider` cannot be silently ignored. Update requests reject any supplied execution mode +Provider records remain excluded during incremental add/update and database catalog reload, before cached gateway entries can be reused. Direct table-to-runtime conversion and both shared client-creation boundaries reject provider mode before OAuth discovery, credential resolution, or MCP connection. YAML loading still rejects provider mode until its inventory path is implemented -This is an intermediate review slice. Management API/config creation of provider registrations remains disabled. Canonical URL and allowlist validation, merged-record update validation, provider inventory loading, and management tool selection are still pending. Later work will enable creation after those contracts are wired in. Native LLM admission authorization is unchanged +The current review slice enables provider registration on the existing create API, canonicalizes HTTPS endpoint identity, requires an explicit finite tool list, and rejects active gateway connection settings. Create persists null transport/auth fields and an enforced allowlist. Update validates the merged record under a row lock, preserves omitted fields, rejects null tool lists, and rejects changes to execution mode. The existing provider-only unique index rejects duplicate endpoints with HTTP 409 + +Admin management inventory includes active provider records. Detail reads return stored allowed tools without running health probes. The UI, scoped key/team selectors, toolset integration, YAML provider inventory, and native LLM admission authorization remain pending. Registration management alone does not yet enforce native-request permissions + +Management inventory and executable servers share one revision-based catalog snapshot. A refresh reads active rows once, retains provider registrations for management, and excludes them from runtime construction. Management listing does not issue a second database inventory query + +Create normalizes provider input in a dedicated request-model validator, then serializes once. Update locks and reads the row, merges the raw patch, validates through the same request model, then serializes and writes. Gateway credential merging retains its existing flow and transaction rules. No extra credential-cleanup SQL or null-preservation flag is needed + +For LiteLLM execution, an empty `allowed_tools` list keeps its existing unrestricted meaning unless `mcp_info.tool_allowlist_enforced=true` is set. Provider registration always persists that marker as true, making its empty list deny all Database-slice validation passed with Prisma schema validation, identical schema copies, the repository's no-data-rewriting migration check, and six checks on a disposable PostgreSQL 16 database built from the pre-change schema. Those checks cover existing rows and duplicate gateway URLs, new legacy defaults, provider records with null connection fields and exact tool lists, duplicate provider URL rejection, distinct provider URLs with empty tool lists, and invalid/null mode rejection. Application-level URL canonicalization, field validation, and authorization are not covered by this database-only proof Runtime isolation was exercised against a real local MCP server. The existing gateway successfully initialized, listed tools, and called an echo tool, producing 11 HTTP requests. Provider-mode table conversion, config loading, manager client creation, and direct upstream client creation were rejected with zero additional HTTP requests. Provider rows produced zero executable catalog entries. This proves the local execution boundary, not native-provider request authorization The runtime slice passes 369 focused model, request, catalog, and MCP utility tests, along with Ruff and the test lint rules. A type-check comparison of the six changed runtime files found no new diagnostics against the approved commit, and strict Ruff rule counts did not increase. These are focused slice checks, not a full repository CI run + +Registration API validation uses a real local proxy and disposable PostgreSQL database. The live checks cover create/read/admin inventory, partial updates, explicit empty lists, invalid/null tool scopes, immutable mode, duplicate endpoints, rejected connection settings, gateway-listing denial, and concurrent tool-list/description edits. The prior run verified null connection fields and zero outbound connections. Revised validation checks absence of actual credentials without requiring a particular JSON-null representation. Evidence is local to the review workspace and contains no customer credentials + +The registration slice passes 899 focused tests, including the CI directory-coverage check, and 16 live API/database checks. The new API tests retain legacy gateway update coverage alongside provider-specific validation + + +Revision review follows the code in this order + +1. `NewMCPServerRequest.normalize_registration` owns provider normalization, separately from transport validation +2. `update_mcp_server` reads the stored execution mode. Provider updates use the existing writer with a preparation callback under the row lock, validating the merged registration through `NewMCPServerRequest`. Only supplied fields plus normalized `mcp_info` are serialized for the patch. Explicit null metadata clears the value, while omitted tools remain untouched +3. `CatalogSnapshots._reload` reads active records once. The published snapshot contains both management inventory and gateway-only executable servers. `_resolve_accessible_mcp_servers` applies visibility to that cached inventory +4. Provider credentials use the existing empty representation. No additional SQL update enforces a distinction between SQL NULL and JSON null + +The revision passed 966 focused tests plus 132 catalog/manager tests. The live proxy and PostgreSQL run passed 20 checks including explicit metadata clearing, inventory refresh after updates and deletion, and concurrent edits. The provider endpoint received zero connections. Changed-file Ruff checks passed, strict rule counts did not increase, and a five-file type-check comparison introduced no diagnostics against the approved base. These revisions were reviewed and approved + + +The latest revision restores the gateway credential-update flow and its existing transaction rules. The shared writer accepts a preparation callback only for provider updates, keeping their stored-row validation under the lock. Neutral connection fields remain accepted without enabling upstream activity. No strict five-field rewrite or scoped-retrieval implementation was added + +Validation after this trim passed 967 focused tests, followed by 606 affected tests after the final adjustment. All 24 live proxy/PostgreSQL checks passed, including neutral fields on create/update, unchanged records after rejection, deleted-record errors, and concurrent edits. Provider connections remained zero. Type comparison found no new diagnostics and strict Ruff rule counts did not increase. The revision was reviewed in VS Code and approved for commit and push diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 43f514284a7..85b5c68e424 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -64,13 +64,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 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 @@ -1378,12 +1378,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 @@ -1444,6 +1444,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") @@ -1455,7 +1459,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") @@ -1491,7 +1503,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): 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 @@ -1545,7 +1557,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 f6b836bc4be..9f359af21d1 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 @@ -1342,10 +1343,24 @@ 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]] = {} + provider_servers: Final = [ + server for server in visible_inventory + if getattr(server, "execution_mode", MCPExecutionMode.litellm) == MCPExecutionMode.provider + ] + aggregated: Final[dict[str, LiteLLM_MCPServerTable]] = {server.server_id: server for server in provider_servers} 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) @@ -1992,18 +2007,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 +2105,10 @@ 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 Exception as e: verbose_proxy_logger.exception("Error creating mcp server: %s", e) raise HTTPException( @@ -3181,7 +3203,7 @@ if MCP_AVAILABLE: raise HTTPException( status_code=503, detail="Cannot validate MCP configuration while stored state is unavailable" ) - if stored is None: + if stored is None or stored.execution_mode == MCPExecutionMode.provider: return effective_transport: Final = payload.transport if "transport" in fields_set else stored.transport effective_info: Final = payload.mcp_info if "mcp_info" in fields_set else stored.mcp_info @@ -3302,12 +3324,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..06697bd9a62 100644 --- a/litellm/proxy/mcp_server/catalog.py +++ b/litellm/proxy/mcp_server/catalog.py @@ -23,6 +23,7 @@ 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 @@ -60,6 +61,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 +92,10 @@ 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 _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 @@ -108,6 +113,10 @@ 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: 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 +125,7 @@ class CatalogSnapshots: self.manager = manager self._refresh_lock = asyncio.Lock() self._database_identity = "" + self._registration_only_servers: Mapping[str, LiteLLM_MCPServerTable] = MappingProxyType({}) self._arrival_ticket = 0 self._completed_ticket = 0 self._shared_snapshot: CatalogSnapshot | None = None @@ -150,11 +160,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 +186,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 +222,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 +305,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 +507,13 @@ 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 + ) + }) previous_registry: Final = self.manager.registry new_registry: Final = await self._stage_servers(raw_rows, reuse_unchanged=reuse_unchanged) @@ -569,6 +593,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..aaff05fde2b 100644 --- a/litellm/proxy/mcp_server/db.py +++ b/litellm/proxy/mcp_server/db.py @@ -57,7 +57,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 +458,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) @@ -1082,6 +1090,11 @@ async def create_mcp_server( """ 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()) @@ -1246,35 +1259,70 @@ 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: set[str] | None = None, +) -> NewMCPServerRequest: + patch: Final = data.model_dump(include=fields_set, exclude_unset=True) + try: + registration: Final = LiteLLM_MCPServerTable.model_validate(stored.model_dump()) + merged: Final = {**registration.model_dump(), **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: set[str] | None, +) -> Mapping[str, object]: + from prisma import Json + + validated: Final = _validate_provider_registration_update(stored, data, fields_set) + supplied_fields: Final = fields_set if fields_set is not None else data.model_fields_set + normalized_patch: Final = data.model_copy(update=validated.model_dump(include=supplied_fields | {"mcp_info"})) + serialized: Final = _prepare_mcp_server_data(normalized_patch, exclude_unset=True, fields_set=supplied_fields) + return { + **serialized, + **({"credentials": Json(None)} if "credentials" in serialized and serialized["credentials"] is None else {}), + "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 +1336,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 +1344,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 +1363,21 @@ 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()) + ) + # 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 @@ -1330,7 +1393,7 @@ async def update_mcp_server( 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 = registration auth_type_changed: Final = bool( data.auth_type @@ -1442,9 +1505,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 b7513daed8c..6207dd7a40a 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -7868,16 +7868,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/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 ee7f6aeefe1..e1013828906 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, @@ -96,7 +95,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..bb078736458 --- /dev/null +++ b/litellm/types/mcp_server/provider_registration.py @@ -0,0 +1,110 @@ +import json +from collections.abc import Mapping +from typing import Final +from urllib.parse import urlsplit, urlunsplit + +from pydantic import AnyHttpUrl, TypeAdapter + +_STR: Final = TypeAdapter(str) +_TOOLS: Final = TypeAdapter(list[str]) +_MAP: Final = TypeAdapter(dict[str, object]) +_URL: Final = TypeAdapter(AnyHttpUrl) +_CONNECTION_FIELDS: Final = ( + "transport", + "auth_type", + "credentials", + "spec_path", + "extra_headers", + "static_headers", + "env_vars", + "command", + "args", + "env", + "issuer", + "authorization_url", + "token_url", + "registration_url", + "oauth2_flow", + "token_exchange_endpoint", + "audience", + "subject_token_type", + "token_exchange_profile", + "delegate_auth_to_upstream", + "oauth_passthrough", + "dcr_bridge", + "per_server_oauth_discovery", + "is_byok", + "byok_description", + "byok_api_key_help_url", + "timeout", + "max_concurrent_requests", + "rpm", + "pinned_tools", +) +_FLAG_FIELDS: Final = frozenset( + {"delegate_auth_to_upstream", "oauth_passthrough", "dcr_bridge", "per_server_oauth_discovery", "is_byok"} +) +_EMPTY_FIELDS: Final = frozenset( + {"credentials", "extra_headers", "static_headers", "env_vars", "args", "env", "byok_description", "pinned_tools"} +) +_JSON_FIELDS: Final = frozenset({"credentials", "static_headers", "env_vars", "env", "pinned_tools"}) + + +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 _active_connection_field(field: str, raw: object) -> bool: + value: Final = _decoded_json(raw) if field in _JSON_FIELDS else raw + return not ( + value is None or (field in _FLAG_FIELDS and value is False) or (field in _EMPTY_FIELDS and value in ({}, [])) + ) + + +def normalize_provider_registration(values: Mapping[str, object]) -> dict[str, object]: + unsupported: Final = next( + (field for field in _CONNECTION_FIELDS if _active_connection_field(field, values.get(field))), None + ) + if unsupported is not None: + raise ValueError(f"Provider-executed MCP registrations do not accept {unsupported}") + if values.get("available_on_public_internet", True) is not True: + raise ValueError("Provider-executed MCP registrations do not support gateway network restrictions") + if values.get("allowed_tools") is None: + raise ValueError("Provider registrations require an explicit allowed_tools list") + tools: Final = _TOOLS.validate_python(values.get("allowed_tools")) + 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") + info: Final = _MAP.validate_python(_decoded_json(values.get("mcp_info")) 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") + descriptions: Final = _MAP.validate_python(_decoded_json(values.get("tool_name_to_description")) or {}) + if not descriptions.keys() <= set(tools): + raise ValueError("Tool descriptions must reference allowed_tools") + return { + **values, + "url": canonical_provider_mcp_url(values.get("url")), + "allowed_tools": list(dict.fromkeys(tools)), + "mcp_info": {**info, "tool_allowlist_enforced": True}, + "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..bcacd9062f8 100644 --- a/tests/unit/proxy/mcp_server/test_catalog.py +++ b/tests/unit/proxy/mcp_server/test_catalog.py @@ -851,3 +851,43 @@ 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 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..9b0ff7a46e9 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"}}, + ], +) +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 written["transport"] is None + assert _credentials_cleared(written["credentials"]) + assert "allowed_tools" not in written + assert json.loads(written["mcp_info"])["tool_allowlist_enforced"] is True diff --git a/tests/unit/proxy/test__types.py b/tests/unit/proxy/test__types.py index 1d2cb7e4cfe..1d0db77d296 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..1d9ab05f4af --- /dev/null +++ b/tests/unit/types/mcp_server/test_provider_registration.py @@ -0,0 +1,67 @@ +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"}, + {"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