From ecbb53dbc99a5b32f456e839780ea53eea76a184 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Tue, 3 Nov 2026 17:35:14 -0800 Subject: [PATCH] fix(mcp): reject duplicate provider URLs without relying on the index Provider create and update now check for another provider registration with the same canonical URL inside the existing identifier transaction, under an advisory lock on that URL, and return 409. Setups that skip migration SQL (prisma db push, adopted databases) have no partial unique index, so they previously accepted duplicates. --- .../mcp_management_endpoints.py | 7 ++- litellm/proxy/mcp_server/db.py | 54 +++++++++++++++++-- tests/integration/mcp/test_mcp_management.py | 38 ++++++++++++- 3 files changed, 92 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 67d816f5a4b..042f334cade 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -335,6 +335,11 @@ if MCP_AVAILABLE: _validate_upstream_token_header(payload) def mcp_identifier_conflict_message(conflict: McpIdentifierConflict) -> str: + if conflict.field == "url": + return ( + f"Provider MCP registration {conflict.server_id} already uses {conflict.value}. " + "Each provider endpoint can be registered once." + ) return ( f"An MCP server with {conflict.field} '{conflict.value}' already exists " f"(server_id={conflict.server_id}). " @@ -343,7 +348,7 @@ if MCP_AVAILABLE: def raise_mcp_identifier_conflict(conflict: McpIdentifierConflict) -> NoReturn: raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, + status_code=status.HTTP_409_CONFLICT if conflict.field == "url" else status.HTTP_400_BAD_REQUEST, detail={"error": mcp_identifier_conflict_message(conflict)}, ) diff --git a/litellm/proxy/mcp_server/db.py b/litellm/proxy/mcp_server/db.py index b09123b2f1a..cd856745549 100644 --- a/litellm/proxy/mcp_server/db.py +++ b/litellm/proxy/mcp_server/db.py @@ -95,7 +95,7 @@ class McpIdentifierConflict: string, and ``server_id`` the existing row that owns it. """ - field: Literal["server_name", "alias"] + field: Literal["server_name", "alias", "url"] value: str server_id: str @@ -640,6 +640,16 @@ def _identifier_field(data_dict: "Mapping[str, object]", field: str) -> str | No return value if isinstance(value, str) else None +def _canonical_provider_url_or_none(url: object) -> str | None: + """The identity a provider URL update must not share; None leaves an invalid URL to update validation.""" + from litellm.types.mcp_server.provider_registration import canonical_provider_mcp_url + + try: + return canonical_provider_mcp_url(url) + except ValueError: + return None + + async def _find_mcp_server_identifier_conflict( table: "TableActions[prisma_db_models.LiteLLM_MCPServerTable]", *, @@ -698,6 +708,26 @@ def _mcp_identifier_lock_keys(*identifiers: str | None) -> tuple[int, ...]: ) +async def _find_provider_url_conflict( + table: "TableActions[prisma_db_models.LiteLLM_MCPServerTable]", + *, + provider_url: str | None, + exclude_server_id: str | None, +) -> McpIdentifierConflict | None: + if not provider_url: + return None + from prisma.enums import MCPExecutionMode as StoredExecutionMode + + own_row_guard: Final[tuple[prisma_db_types.LiteLLM_MCPServerTableWhereInput, ...]] = ( + ({"NOT": [{"server_id": exclude_server_id}]},) if exclude_server_id is not None else () + ) + same_url: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = { + "AND": [{"execution_mode": StoredExecutionMode.provider}, {"url": provider_url}, *own_row_guard] + } + row: Final = await table.find_first(where=same_url) + return None if row is None else McpIdentifierConflict(field="url", value=provider_url, server_id=row.server_id) + + async def _mcp_server_write_if_identifier_free( prisma_client: PrismaClient, *, @@ -706,14 +736,19 @@ async def _mcp_server_write_if_identifier_free( exclude_server_id: str | None, write: "Callable[[TableActions[prisma_db_models.LiteLLM_MCPServerTable]], Awaitable[prisma_db_models.LiteLLM_MCPServerTable | None]]", lock_server_id: str | None = None, + provider_url: str | None = None, ) -> "prisma_db_models.LiteLLM_MCPServerTable | McpIdentifierConflict | None": - """Run ``write`` only when no other live row owns ``server_name``/``alias``. + """Run ``write`` only when no other live row owns ``server_name``/``alias``, nor, + for a provider registration, its canonical ``provider_url``. The conflict check and the write share a transaction guarded by per-identifier - advisory locks, so two concurrent requests for the same name cannot both - pass the check and both insert. + advisory locks, so two concurrent requests for the same name or provider URL + cannot both pass the check and both write. This does not depend on the partial + unique index, which setups that skip migration SQL never create. """ - lock_keys: Final = _mcp_identifier_lock_keys(server_name, alias) + lock_keys: Final = _mcp_identifier_lock_keys( + server_name, alias, f"provider_url:{provider_url}" if provider_url else None + ) async with _db_transaction_manager(prisma_client) as tx: for lock_key in lock_keys: await tx.execute_raw("SELECT pg_advisory_xact_lock($1::bigint)", lock_key) @@ -722,6 +757,8 @@ async def _mcp_server_write_if_identifier_free( server_name=server_name, alias=alias, exclude_server_id=exclude_server_id, + ) or await _find_provider_url_conflict( + tx.litellm_mcpservertable, provider_url=provider_url, exclude_server_id=exclude_server_id ) if conflict is not None: return conflict @@ -1165,6 +1202,7 @@ async def create_mcp_server_if_identifier_free( alias=_identifier_field(data_dict, "alias"), exclude_server_id=None, write=_create, + provider_url=_identifier_field(data_dict, "url") if data.execution_mode == MCPExecutionMode.provider else None, ) if isinstance(written, McpIdentifierConflict): return written @@ -1360,6 +1398,7 @@ async def _update_mcp_server_row( data_dict: Mapping[str, object], prepare_update: Callable[["prisma_db_models.LiteLLM_MCPServerTable"], Mapping[str, object]] | None = None, expected_updated_at: datetime | None = None, + provider_url: str | None = None, ) -> "prisma_db_models.LiteLLM_MCPServerTable | McpIdentifierConflict | None": identifier_write: Final = any(field in data_dict for field in ("server_name", "alias")) protocol_write: Final = bool({"transport", "mcp_info"}.intersection(data_dict)) @@ -1399,6 +1438,7 @@ async def _update_mcp_server_row( exclude_server_id=server_id, write=_update, lock_server_id=server_id if needs_stored_row else None, + provider_url=provider_url, ) return await _mcp_server_write_if_identifier_free( prisma_client, @@ -1407,6 +1447,7 @@ async def _update_mcp_server_row( exclude_server_id=server_id, write=_update, lock_server_id=server_id if needs_stored_row else None, + provider_url=provider_url, ) @@ -1433,6 +1474,9 @@ async def update_mcp_server( data_dict=data.model_dump(include=fields_set, exclude_unset=True), prepare_update=lambda stored: _prepare_provider_registration_update(stored, data, touched_by, fields_set), expected_updated_at=expected_updated_at, + provider_url=_canonical_provider_url_or_none(data.url) + if "url" in (fields_set or data.model_fields_set) + else None, ) return ( updated_provider diff --git a/tests/integration/mcp/test_mcp_management.py b/tests/integration/mcp/test_mcp_management.py index 90435f49b34..bcd8f3001ef 100644 --- a/tests/integration/mcp/test_mcp_management.py +++ b/tests/integration/mcp/test_mcp_management.py @@ -7,7 +7,7 @@ from typing import Final import pytest import yaml from integration._support.client import JSON_OBJECT, Gateway, eventually -from integration._support.database import read_rows +from integration._support.database import read_rows, scratch_database from integration._support.mcp import ( McpCaller, McpPeer, @@ -209,6 +209,42 @@ def _openapi_server_lists_and_calls_only_its_own_tools( assert [(item["method"], item["path"]) for item in peer.drain()] == [("GET", "/pets/7")], identity +def test_duplicate_provider_url_is_rejected_where_db_push_skipped_the_unique_index( + gateway: Gateway, tmp_path: Path +) -> None: + import concurrent.futures + + config: Final = tmp_path / "pushed.yaml" + config.write_text(yaml.safe_dump({"model_list": [], "general_settings": {"master_key": gateway.key}})) + with scratch_database("provider_url") as database_url, owned_proxy( + gateway, tmp_path, {"DATABASE_URL": database_url}, config=config + ) as pushed: + index: Final = read_rows( + "SELECT to_regclass('\"LiteLLM_MCPServerTable_provider_url_key\"')::text AS name", (), database_url=database_url + ) + assert index == [{"name": None}], "db push should leave the SQL-only partial index out" + url: Final = f"https://{uuid.uuid4().hex[:12]}.example/mcp" + + def register(name: str, endpoint: str) -> int: + return pushed.request( + "POST", + "/v1/mcp/server", + {"server_id": name, "server_name": name, "execution_mode": "provider", "url": endpoint, "allowed_tools": ["search"]}, + ).status_code + + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool: + statuses: Final = sorted(pool.map(partial(register, endpoint=url), (f"prov{n}" for n in range(8)))) + assert statuses == [201] + [409] * 7, statuses + stored: Final = read_rows( + 'SELECT server_id FROM "LiteLLM_MCPServerTable" WHERE url = %s', (url,), database_url=database_url + ) + assert len(stored) == 1, stored + + assert register("prov_other", f"https://other-{uuid.uuid4().hex[:8]}.example/mcp") == 201 + moved: Final = pushed.request("PUT", "/v1/mcp/server", {"server_id": "prov_other", "url": url}) + assert moved.status_code == 409, moved.text + + def test_openapi_listing_is_scoped_to_the_exact_alias_when_aliases_overlap(gateway: Gateway) -> None: with openapi_peer() as short, openapi_peer() as long, gateway.scenario() as scenario: stem: Final = "pet" + uuid.uuid4().hex[:8]