mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(mcp): reject duplicate provider URLs without relying on the index
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
LiteLLM Rust DB / db-queries (push) Has been cancelled
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
LiteLLM Rust DB / db-queries (push) Has been cancelled
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
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.
This commit is contained in:
parent
7c27414bcc
commit
ecbb53dbc9
3 changed files with 92 additions and 7 deletions
|
|
@ -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)},
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue