mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
fix(mcp): let config.yaml MCP servers pin server_id (#39286)
* fix(mcp): let config.yaml MCP servers pin server_id A config-defined MCP server's id is a hash of server_name|url|transport| auth_type|alias, recomputed on every config load, so editing any of those fields mints a new id. Every key and team granted the old id via object_permission.mcp_servers keeps pointing at an id that no longer exists, and the server disappears from tools/list for them with nothing logged. load_servers_from_config now uses an explicit server_id from the server's config entry when present and falls back to the existing hash otherwise, so grants survive url/name/alias edits. Rejected at config load: a blank or non-string server_id, two entries claiming the same id, a pinned id already held by a database-backed server, and a pinned id that is another entry's server_name or alias (expand_permission_list matches ids before names, so that one would capture the other server's grants). Because the database registry loads after the config on startup, a database row that lands on a pinned config id is reported as a warning from the database reload instead, where it is decidable; the warning is latched on the shadowed set so the config-reload timer does not reprint it every interval. Deployments that do not set server_id keep the exact id they have today. * fix(mcp): close two more pinned-id capture paths A pinned server_id equal to an alias supplied through litellm_settings mcp_aliases was accepted, because the collision index only held the entry's own alias field. expand_permission_list matches ids before names, so grants written for the aliased server resolved to the pinning one. mcp_aliases keys whose target is a config server are now reserved the same way. A pinned server_id equal to a database-backed server's name, server_name or alias had the same effect against the database side, and could not be rejected at config load because the database registry is not loaded yet. The database reload now warns about it, latched like the existing shadow warning. * fix(mcp): reserve only the aliases the loader actually assigns Reserving every mcp_aliases key targeting a config server was too broad in two ways: the mapping is ignored when the entry sets its own alias, and only the first mapping for a server is ever applied. Both cases made a pinned server_id that could never have collided abort proxy startup. Reserve only the name load_servers_from_config will really assign. The database capture warning also fired for a database server whose own id is the config server_id. There the database row wins the id outright through get_registry precedence, so the shadow warning above it is the accurate one and the capture message contradicted it. Skip those rows. Also mark the two litellm-internal patches in the reload test helper, which the test-quality gate counts; the database reload has no other seam. * fix(mcp): match the loader's alias check exactly, is None not falsiness load_servers_from_config consults mcp_aliases only when the entry has no alias key at all, so an entry setting alias: "" gets no mapped alias. The collision index used falsiness and reserved the mapped name anyway, which failed startup on a pinned server_id that could never have collided with it. * fix(mcp): skip one identifier, not the whole database row A database row can shadow one config server_id by id and capture another by name at the same time. Skipping the entire row when its id shadowed a config entry dropped the second warning, leaving the operator with half a diagnosis. Skip only the identifier equal to the row's own id. * fix(mcp): reject conflicting self-pinned server ids * fix(mcp): validate config server names before building the identifier index The collision check reads every entry's body up front, so a malformed entry under an invalid name surfaced as an AttributeError instead of the name validation error the loader gave before this change.
This commit is contained in:
parent
aea5358c48
commit
fafd294878
2 changed files with 782 additions and 4 deletions
|
|
@ -13,9 +13,19 @@ import json
|
|||
import os
|
||||
import re
|
||||
import time
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
|
||||
from collections.abc import (
|
||||
AsyncIterator,
|
||||
Awaitable,
|
||||
Callable,
|
||||
Container,
|
||||
Iterable,
|
||||
Mapping,
|
||||
MutableMapping,
|
||||
Sequence,
|
||||
)
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass, replace
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast
|
||||
from urllib.parse import ParseResult, urlparse
|
||||
|
||||
|
|
@ -307,6 +317,7 @@ class MCPServerConfig(TypedDict, total=False):
|
|||
:meth:`MCPServerManager.load_servers_from_config`. Every key is optional: YAML supplies
|
||||
whatever the admin wrote, and each read applies its own default."""
|
||||
|
||||
server_id: ReadOnly[str]
|
||||
alias: str
|
||||
description: str
|
||||
mcp_info: MCPInfo
|
||||
|
|
@ -400,6 +411,164 @@ def _blank_to_none(value: str | None) -> str | None:
|
|||
return value.strip() or None
|
||||
|
||||
|
||||
def _pinned_config_server_id(raw_server_id: object, server_name: str) -> str | None:
|
||||
"""Return the ``server_id`` an admin pinned for this config.yaml server, or ``None`` when absent.
|
||||
|
||||
Without a pin the id is derived by hashing ``server_name|url|transport|auth_type|alias``, so
|
||||
editing any of those fields mints a new id and every ``object_permission.mcp_servers`` grant
|
||||
holding the old one silently stops matching. A pinned id is used verbatim and survives those
|
||||
edits. Blank and non-string values are rejected rather than silently falling back to the hash,
|
||||
because a config that pins an id and still churns is the failure this field exists to prevent.
|
||||
|
||||
Under ``LITELLM_USE_SHORT_MCP_TOOL_PREFIX`` the tool prefix is derived from the server_id, so
|
||||
pinning an id other than the one already in use renames every tool that server exposes.
|
||||
"""
|
||||
if raw_server_id is None:
|
||||
return None
|
||||
if not isinstance(raw_server_id, str) or not raw_server_id.strip():
|
||||
raise ValueError(
|
||||
f"Invalid config for MCP server '{server_name}': server_id must be a non-empty string "
|
||||
f"(got {raw_server_id!r})."
|
||||
)
|
||||
return raw_server_id.strip()
|
||||
|
||||
|
||||
def _first_mapped_alias(server_name: str, mcp_aliases: Mapping[str, str] | None) -> str | None:
|
||||
"""The ``mcp_aliases`` name ``load_servers_from_config`` will assign to this server, if any.
|
||||
|
||||
Mirrors that loop, which takes the first mapping pointing at the server and stops. A later
|
||||
mapping for the same server is never applied, so it stays free for another entry to pin.
|
||||
"""
|
||||
if mcp_aliases is None:
|
||||
return None
|
||||
return next(
|
||||
(alias_name for alias_name, target_server_name in mcp_aliases.items() if target_server_name == server_name),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _assigned_alias(
|
||||
server_name: str, server_config: MCPServerConfig, mcp_aliases: Mapping[str, str] | None
|
||||
) -> str | None:
|
||||
"""The alias ``load_servers_from_config`` will give this entry: its own, else the first mapping.
|
||||
|
||||
``is None``, not falsiness: the loader only consults the mapping when the key is absent, so an
|
||||
entry that sets ``alias: ""`` gets no mapped alias and reserves nothing.
|
||||
"""
|
||||
alias: Final = server_config.get("alias")
|
||||
return _first_mapped_alias(server_name, mcp_aliases) if alias is None else alias
|
||||
|
||||
|
||||
def _validate_config_server_names(mcp_servers_config: Mapping[str, MCPServerConfig]) -> None:
|
||||
"""Reject bad server names before ``_config_identifier_owners`` reads any entry's body.
|
||||
|
||||
The identifier index walks every entry up front, so without this pass a malformed entry under
|
||||
a bad name would surface as an ``AttributeError`` from the index instead of the name error.
|
||||
"""
|
||||
for server_name in mcp_servers_config:
|
||||
validate_mcp_server_name(server_name)
|
||||
|
||||
|
||||
def _config_identifier_owners(
|
||||
mcp_servers_config: Mapping[str, MCPServerConfig],
|
||||
mcp_aliases: Mapping[str, str] | None,
|
||||
) -> Mapping[str, frozenset[str]]:
|
||||
"""Map every server_name and alias in the config to the entries that own it.
|
||||
|
||||
``expand_permission_list`` resolves a grant against the registry keys before it falls back to
|
||||
matching alias and server_name, so an id equal to another entry's name or alias captures that
|
||||
entry's grants. Derived ids are hashes and never collide with a name, so this only matters once
|
||||
an id is pinned.
|
||||
|
||||
An alias is either set on the entry or mapped to it from ``litellm_settings.mcp_aliases``. Only
|
||||
a name the loader below will really assign is reserved: the mapping is ignored for an entry that
|
||||
sets its own ``alias``, and only the first mapping wins for one that does not, so reserving every
|
||||
mapping would fail startup on a pin that was never going to collide.
|
||||
|
||||
One identifier can have several owners when an entry's alias equals another entry's name. All of
|
||||
them are kept: a grant naming that identifier resolves to every match while no id is pinned, and
|
||||
a pin equal to it would narrow the grant to the pinning entry alone, even when that entry is one
|
||||
of the owners.
|
||||
"""
|
||||
claims: Final = tuple(
|
||||
(identifier, server_name)
|
||||
for server_name, server_config in mcp_servers_config.items()
|
||||
for identifier in (server_name, _assigned_alias(server_name, server_config, mcp_aliases))
|
||||
if identifier
|
||||
)
|
||||
return MappingProxyType(
|
||||
{identifier: frozenset(owner for claimed, owner in claims if claimed == identifier) for identifier, _ in claims}
|
||||
)
|
||||
|
||||
|
||||
def _config_ids_capturing_db_identifiers(
|
||||
config_server_ids: Container[str],
|
||||
db_servers: Iterable[MCPServer],
|
||||
) -> frozenset[str]:
|
||||
"""Config server ids that are a database-backed server's name, server_name or alias.
|
||||
|
||||
``expand_permission_list`` matches a grant against the registry keys before it matches names, so
|
||||
such an id answers every grant written for the database server, and the database server itself
|
||||
stops being reachable by name. The config load cannot catch this because the database registry
|
||||
is not loaded yet, so it is reported from the reload that does have both halves.
|
||||
|
||||
An identifier equal to the database server's own id is skipped: ``get_registry`` is
|
||||
``config_mcp_servers | registry``, so there the database server wins the id outright and the
|
||||
shadow warning above is the accurate one. Reporting both would contradict. The skip is per
|
||||
identifier rather than per server, so a row that shadows one config id and captures another
|
||||
still reports the capture.
|
||||
"""
|
||||
return frozenset(
|
||||
identifier
|
||||
for server in db_servers
|
||||
for identifier in (server.name, server.server_name, server.alias)
|
||||
if identifier and identifier != server.server_id and identifier in config_server_ids
|
||||
)
|
||||
|
||||
|
||||
def _reject_config_server_id_collision(
|
||||
assigned_server_ids: Mapping[str, str],
|
||||
server_id: str,
|
||||
server_name: str,
|
||||
pinned: bool,
|
||||
db_backed_server_ids: Mapping[str, object],
|
||||
identifier_owners: Mapping[str, frozenset[str]],
|
||||
) -> None:
|
||||
"""Raise when ``server_id`` is already taken, either by an earlier config entry or by the database.
|
||||
|
||||
Two config entries sharing an id would silently overwrite each other in ``config_mcp_servers``,
|
||||
and an id already held by a database-backed server is hidden by it, because ``get_registry`` is
|
||||
``config_mcp_servers | registry`` and the right operand wins. A pinned id that is another
|
||||
entry's server_name or alias captures that entry's permission grants the same way. Derived ids
|
||||
cannot collide (the unique config key is part of the hash input), so all three only happen once
|
||||
an id is pinned.
|
||||
|
||||
Pinning an identifier this entry itself owns is allowed, because a grant naming it already
|
||||
resolved here, but only when no other entry owns it too. An entry whose alias is this entry's
|
||||
server_name shares the identifier, and pinning it would take that entry's grants.
|
||||
"""
|
||||
claimed_by = assigned_server_ids.get(server_id)
|
||||
if claimed_by is not None:
|
||||
raise ValueError(
|
||||
f"Invalid config for MCP server '{server_name}': server_id '{server_id}' is already "
|
||||
f"used by MCP server '{claimed_by}'. Each mcp_servers entry needs its own id."
|
||||
)
|
||||
if pinned and server_id in db_backed_server_ids:
|
||||
raise ValueError(
|
||||
f"Invalid config for MCP server '{server_name}': server_id '{server_id}' belongs to a "
|
||||
"database-backed MCP server. The database entry takes precedence over config.yaml, so "
|
||||
"this server would never be reachable."
|
||||
)
|
||||
other_owners: Final = identifier_owners.get(server_id, frozenset()) - frozenset((server_name,))
|
||||
if pinned and other_owners:
|
||||
owner_names: Final = "', '".join(sorted(other_owners))
|
||||
raise ValueError(
|
||||
f"Invalid config for MCP server '{server_name}': server_id '{server_id}' is the "
|
||||
f"server_name or alias of MCP server '{owner_names}'. Permission entries naming "
|
||||
f"'{server_id}' would resolve to '{server_name}' alone and no longer reach '{owner_names}'."
|
||||
)
|
||||
|
||||
|
||||
def _uses_issuer_anchor(manual_issuer: str | None, is_discovery_auth_type: bool) -> bool:
|
||||
"""Whether the endpoints are authoritatively anchored to an admin-pinned issuer (RFC 8414 §3.3).
|
||||
|
||||
|
|
@ -1565,6 +1734,11 @@ class MCPServerManager:
|
|||
# empty result, or failure). Used to throttle re-probes for servers that do
|
||||
# not return instructions, and to apply a short cooldown after failures.
|
||||
self._upstream_initialize_instructions_probed_at: dict[str, float] = {}
|
||||
# Last set of config server ids found shadowed by database rows. reload_servers_from_database
|
||||
# runs on the config-reload timer, so this keeps a standing misconfiguration from re-logging
|
||||
# the same warning every interval; a change in the set logs again.
|
||||
self._warned_shadowed_config_server_ids: frozenset[str] = frozenset()
|
||||
self._warned_capturing_config_server_ids: frozenset[str] = frozenset()
|
||||
self._oauth_discovery_on_startup = _mcp_oauth_discovery_on_startup_enabled()
|
||||
self._oauth_discovery_generation_counter = 0
|
||||
self._oauth_discovery_slots: tuple[_OAuthDiscoverySlot, ...] = ()
|
||||
|
|
@ -1958,10 +2132,14 @@ class MCPServerManager:
|
|||
|
||||
# Track which aliases have been used to ensure only first occurrence is used
|
||||
used_aliases: Final = set()
|
||||
# server_id -> the config server_name that claimed it, so a pinned id cannot silently
|
||||
# overwrite another server's entry in self.config_mcp_servers.
|
||||
assigned_server_ids: MutableMapping[str, str] = {} # mutable-ok: per-load collision index
|
||||
_validate_config_server_names(mcp_servers_config)
|
||||
identifier_owners: Final = _config_identifier_owners(mcp_servers_config, mcp_aliases)
|
||||
|
||||
for server_name, raw_server_config in mcp_servers_config.items():
|
||||
server_config: MCPServerConfig = raw_server_config
|
||||
validate_mcp_server_name(server_name)
|
||||
_mcp_info: MCPInfo = server_config.get("mcp_info", None) or {}
|
||||
# Preserve all custom fields from config while setting defaults for core fields
|
||||
mcp_info: MCPInfo = _mcp_info.copy()
|
||||
|
|
@ -1994,14 +2172,24 @@ class MCPServerManager:
|
|||
name_for_prefix = get_server_prefix(temp_server)
|
||||
|
||||
server_url = server_config.get("url", None) or ""
|
||||
# Generate stable server ID based on parameters
|
||||
server_id = self._generate_stable_server_id(
|
||||
# An explicitly pinned server_id wins; otherwise derive one from the parameters.
|
||||
pinned_server_id = _pinned_config_server_id(server_config.get("server_id"), server_name)
|
||||
server_id = pinned_server_id or self._generate_stable_server_id(
|
||||
server_name=server_name,
|
||||
url=server_url,
|
||||
transport=server_config.get("transport", MCPTransport.http),
|
||||
auth_type=server_config.get("auth_type", None),
|
||||
alias=alias,
|
||||
)
|
||||
_reject_config_server_id_collision(
|
||||
assigned_server_ids,
|
||||
server_id,
|
||||
server_name,
|
||||
pinned=pinned_server_id is not None,
|
||||
db_backed_server_ids=self.registry,
|
||||
identifier_owners=identifier_owners,
|
||||
)
|
||||
assigned_server_ids[server_id] = server_name
|
||||
|
||||
_warn_on_server_name_fields(
|
||||
server_id=server_id,
|
||||
|
|
@ -6123,6 +6311,33 @@ class MCPServerManager:
|
|||
|
||||
verbose_logger.debug("MCP registry refreshed (%s servers in registry)", len(registered_registry))
|
||||
|
||||
# get_registry() is ``config_mcp_servers | registry``, so a database row sharing an id with a
|
||||
# config.yaml server hides that server everywhere. Only reachable once an operator pins
|
||||
# ``server_id`` in config.yaml; say so rather than letting the server disappear silently.
|
||||
shadowed_config_server_ids: Final = frozenset(self.config_mcp_servers.keys() & registered_registry.keys())
|
||||
if shadowed_config_server_ids and shadowed_config_server_ids != self._warned_shadowed_config_server_ids:
|
||||
verbose_logger.warning(
|
||||
"config.yaml MCP server_id(s) %s are also database-backed MCP servers. The database "
|
||||
"entry takes precedence, so the config.yaml server is unreachable. Give the config "
|
||||
"entry a different server_id.",
|
||||
", ".join(sorted(shadowed_config_server_ids)),
|
||||
)
|
||||
self._warned_shadowed_config_server_ids = shadowed_config_server_ids
|
||||
|
||||
# The mirror image of the block above: a config server_id that is a database server's name
|
||||
# answers that server's grants instead, because ids are matched before names.
|
||||
capturing_config_server_ids: Final = _config_ids_capturing_db_identifiers(
|
||||
self.config_mcp_servers.keys(), registered_registry.values()
|
||||
)
|
||||
if capturing_config_server_ids and capturing_config_server_ids != self._warned_capturing_config_server_ids:
|
||||
verbose_logger.warning(
|
||||
"config.yaml MCP server_id(s) %s are the name or alias of a database-backed MCP "
|
||||
"server. Permission entries naming them resolve to the config.yaml server, not the "
|
||||
"database one. Give the config entry a different server_id.",
|
||||
", ".join(sorted(capturing_config_server_ids)),
|
||||
)
|
||||
self._warned_capturing_config_server_ids = capturing_config_server_ids
|
||||
|
||||
await self._hydrate_config_servers_dcr_clients()
|
||||
|
||||
def get_mcp_servers_from_ids(self, server_ids: list[str]) -> list[MCPServer]:
|
||||
|
|
|
|||
|
|
@ -11428,6 +11428,569 @@ class TestOpenApiHandlerRelaysUpstreamAuth:
|
|||
assert "upstream returned HTTP 503" in result.content[0].text
|
||||
|
||||
|
||||
class TestConfigServerIdPinning:
|
||||
"""config.yaml servers may pin ``server_id`` so permission grants survive connection edits."""
|
||||
|
||||
@staticmethod
|
||||
def _config(**overrides: object) -> dict[str, dict[str, object]]:
|
||||
return {
|
||||
"docs_server": {
|
||||
"url": "https://example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
**overrides,
|
||||
}
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_derived_id_churns_when_connection_fields_change(self):
|
||||
"""The behavior the pin exists to escape: editing the url mints a brand-new id."""
|
||||
manager = MCPServerManager()
|
||||
|
||||
await manager.load_servers_from_config(self._config())
|
||||
before = next(iter(manager.config_mcp_servers))
|
||||
|
||||
manager.config_mcp_servers.clear()
|
||||
await manager.load_servers_from_config(self._config(url="https://prod.example.com/mcp"))
|
||||
after = next(iter(manager.config_mcp_servers))
|
||||
|
||||
assert before != after
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinned_id_survives_url_transport_auth_and_alias_edits(self):
|
||||
manager = MCPServerManager()
|
||||
|
||||
await manager.load_servers_from_config(self._config(server_id="docs-prod-1"))
|
||||
assert list(manager.config_mcp_servers) == ["docs-prod-1"]
|
||||
assert manager.config_mcp_servers["docs-prod-1"].server_id == "docs-prod-1"
|
||||
|
||||
manager.config_mcp_servers.clear()
|
||||
await manager.load_servers_from_config(
|
||||
self._config(
|
||||
server_id="docs-prod-1",
|
||||
url="https://prod.example.com/mcp",
|
||||
transport=MCPTransport.sse,
|
||||
auth_type=MCPAuth.bearer_token,
|
||||
alias="docs",
|
||||
)
|
||||
)
|
||||
|
||||
assert list(manager.config_mcp_servers) == ["docs-prod-1"]
|
||||
assert manager.config_mcp_servers["docs-prod-1"].url == "https://prod.example.com/mcp"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_absent_server_id_keeps_the_derived_hash(self):
|
||||
manager = MCPServerManager()
|
||||
|
||||
await manager.load_servers_from_config(self._config())
|
||||
|
||||
derived = manager._generate_stable_server_id(
|
||||
server_name="docs_server",
|
||||
url="https://example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=None,
|
||||
alias=None,
|
||||
)
|
||||
assert list(manager.config_mcp_servers) == [derived]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("bad_value", ["", " ", 123, True, ["docs-prod-1"]])
|
||||
async def test_blank_or_non_string_server_id_is_rejected(self, bad_value: Any):
|
||||
manager = MCPServerManager()
|
||||
|
||||
with pytest.raises(ValueError, match="server_id must be a non-empty string"):
|
||||
await manager.load_servers_from_config(self._config(server_id=bad_value))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_two_servers_pinning_the_same_id_are_rejected(self):
|
||||
manager = MCPServerManager()
|
||||
config: Dict[str, Any] = {
|
||||
"docs_server": {"url": "https://a.example.com/mcp", "server_id": "shared-id"},
|
||||
"wiki_server": {"url": "https://b.example.com/mcp", "server_id": "shared-id"},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="already used by MCP server 'docs_server'"):
|
||||
await manager.load_servers_from_config(config)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinned_id_colliding_with_a_derived_id_is_rejected(self):
|
||||
"""A pin that lands on another entry's derived hash collides just as hard."""
|
||||
manager = MCPServerManager()
|
||||
derived = manager._generate_stable_server_id(
|
||||
server_name="docs_server",
|
||||
url="https://a.example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=None,
|
||||
alias=None,
|
||||
)
|
||||
config: Dict[str, Any] = {
|
||||
"docs_server": {"url": "https://a.example.com/mcp", "transport": MCPTransport.http},
|
||||
"wiki_server": {"url": "https://b.example.com/mcp", "server_id": derived},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="already used by MCP server 'docs_server'"):
|
||||
await manager.load_servers_from_config(config)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinned_id_colliding_with_a_db_backed_server_is_rejected(self):
|
||||
"""get_registry() is ``config | registry``, so the db row would hide the config server.
|
||||
|
||||
The registry is seeded by hand because on a real startup the config loads before the
|
||||
database does, so this check only fires on a later reload. The startup ordering is covered
|
||||
by ``test_db_row_arriving_on_a_pinned_config_id_warns``; the warning there is not redundant.
|
||||
"""
|
||||
manager = MCPServerManager()
|
||||
manager.registry["db-uuid-1"] = MCPServer(
|
||||
server_id="db-uuid-1",
|
||||
name="db_server",
|
||||
transport=MCPTransport.http,
|
||||
url="https://db.example.com/mcp",
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="belongs to a database-backed MCP server"):
|
||||
await manager.load_servers_from_config(self._config(server_id="db-uuid-1"))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_derived_id_matching_a_db_backed_server_is_not_rejected(self):
|
||||
"""Only a pinned id is an authoring error; a hash collision must not fail startup."""
|
||||
manager = MCPServerManager()
|
||||
derived = manager._generate_stable_server_id(
|
||||
server_name="docs_server",
|
||||
url="https://example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=None,
|
||||
alias=None,
|
||||
)
|
||||
manager.registry[derived] = MCPServer(
|
||||
server_id=derived,
|
||||
name="db_server",
|
||||
transport=MCPTransport.http,
|
||||
url="https://db.example.com/mcp",
|
||||
)
|
||||
|
||||
await manager.load_servers_from_config(self._config())
|
||||
|
||||
assert derived in manager.config_mcp_servers
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinned_id_is_stripped_of_surrounding_whitespace(self):
|
||||
manager = MCPServerManager()
|
||||
|
||||
await manager.load_servers_from_config(self._config(server_id=" docs-prod-1 "))
|
||||
|
||||
assert list(manager.config_mcp_servers) == ["docs-prod-1"]
|
||||
|
||||
@staticmethod
|
||||
async def _reload_with_db_server(manager: MCPServerManager, server_id: str, db_name: str = "db_server") -> None:
|
||||
row = LiteLLM_MCPServerTable(
|
||||
server_id=server_id,
|
||||
server_name=db_name,
|
||||
alias=db_name,
|
||||
url="https://db.example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
)
|
||||
raw_row = MagicMock()
|
||||
raw_row.model_dump.return_value = row.model_dump()
|
||||
repository = MagicMock()
|
||||
repository.table.find_many = AsyncMock(return_value=[raw_row])
|
||||
built = MCPServer(
|
||||
server_id=server_id,
|
||||
name=db_name,
|
||||
server_name=db_name,
|
||||
url="https://db.example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
)
|
||||
with (
|
||||
patch( # test-quality-ok: the db reload path has no seam but its own repository
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository",
|
||||
return_value=repository,
|
||||
),
|
||||
patch( # test-quality-ok: same, the prisma client is fetched inside the reload
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
|
||||
return_value=MagicMock(),
|
||||
),
|
||||
patch.object(manager, "build_mcp_server_from_table", new=AsyncMock(return_value=built)),
|
||||
):
|
||||
await manager.reload_servers_from_database()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_db_row_arriving_on_a_pinned_config_id_warns(self, caplog):
|
||||
"""The db row loads after config on startup, so the config server is hidden then, not at load."""
|
||||
manager = MCPServerManager()
|
||||
await manager.load_servers_from_config(self._config(server_id="docs-prod-1"))
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
await self._reload_with_db_server(manager, "docs-prod-1")
|
||||
|
||||
assert any("docs-prod-1" in m and "database entry takes precedence" in m for m in caplog.messages)
|
||||
assert manager.get_registry()["docs-prod-1"].url == "https://db.example.com/mcp"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_db_row_with_a_distinct_id_does_not_warn(self, caplog):
|
||||
manager = MCPServerManager()
|
||||
await manager.load_servers_from_config(self._config(server_id="docs-prod-1"))
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
await self._reload_with_db_server(manager, "db-uuid-1")
|
||||
|
||||
assert all("database entry takes precedence" not in m for m in caplog.messages)
|
||||
assert set(manager.get_registry()) == {"docs-prod-1", "db-uuid-1"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinned_id_matching_another_entrys_server_name_is_rejected(self):
|
||||
"""expand_permission_list resolves against registry keys first, so this steals the grants."""
|
||||
manager = MCPServerManager()
|
||||
|
||||
with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"):
|
||||
await manager.load_servers_from_config(
|
||||
{
|
||||
"wiki_server": {"url": "https://wiki.example.com/mcp", "transport": MCPTransport.http},
|
||||
"docs_server": {
|
||||
"server_id": "wiki_server",
|
||||
"url": "https://example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinned_id_matching_another_entrys_alias_is_rejected(self):
|
||||
manager = MCPServerManager()
|
||||
|
||||
with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"):
|
||||
await manager.load_servers_from_config(
|
||||
{
|
||||
"wiki_server": {
|
||||
"alias": "wiki",
|
||||
"url": "https://wiki.example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
"docs_server": {
|
||||
"server_id": "wiki",
|
||||
"url": "https://example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinning_a_servers_own_name_is_allowed(self):
|
||||
"""The most natural pin an operator writes; it resolves to the same server either way."""
|
||||
manager = MCPServerManager()
|
||||
|
||||
await manager.load_servers_from_config(self._config(server_id="docs_server"))
|
||||
|
||||
assert list(manager.config_mcp_servers) == ["docs_server"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinning_a_servers_own_alias_is_allowed(self):
|
||||
manager = MCPServerManager()
|
||||
|
||||
await manager.load_servers_from_config(self._config(alias="docs", server_id="docs"))
|
||||
|
||||
assert list(manager.config_mcp_servers) == ["docs"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("aliasing_entry_first", [True, False])
|
||||
async def test_pinning_own_name_that_is_another_entrys_alias_is_rejected(self, aliasing_entry_first: bool):
|
||||
"""A grant naming 'docs_server' reaches both servers unpinned; the pin would narrow it to one."""
|
||||
manager = MCPServerManager()
|
||||
wiki = (
|
||||
"wiki_server",
|
||||
{"alias": "docs_server", "url": "https://wiki.example.com/mcp", "transport": MCPTransport.http},
|
||||
)
|
||||
docs = (
|
||||
"docs_server",
|
||||
{"server_id": "docs_server", "url": "https://example.com/mcp", "transport": MCPTransport.http},
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"):
|
||||
await manager.load_servers_from_config(dict((wiki, docs) if aliasing_entry_first else (docs, wiki)))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinning_own_name_that_is_another_entrys_mapped_alias_is_rejected(self):
|
||||
manager = MCPServerManager()
|
||||
|
||||
with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"):
|
||||
await manager.load_servers_from_config(
|
||||
{
|
||||
"wiki_server": {"url": "https://wiki.example.com/mcp", "transport": MCPTransport.http},
|
||||
"docs_server": {
|
||||
"server_id": "docs_server",
|
||||
"url": "https://example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
},
|
||||
mcp_aliases={"docs_server": "wiki_server"},
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinning_own_alias_shared_with_a_later_entry_is_rejected(self):
|
||||
"""Nothing rejects duplicate aliases, so the first entry's pin would answer the second's grants."""
|
||||
manager = MCPServerManager()
|
||||
|
||||
with pytest.raises(ValueError, match="server_name or alias of MCP server 'docs_server'"):
|
||||
await manager.load_servers_from_config(
|
||||
{
|
||||
"wiki_server": {
|
||||
"alias": "shared",
|
||||
"server_id": "shared",
|
||||
"url": "https://wiki.example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
"docs_server": {
|
||||
"alias": "shared",
|
||||
"url": "https://example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_own_name_pin_resolves_grants_like_the_unpinned_name(self):
|
||||
"""The negative control: a sole-owner self-pin must keep loading and answer the same grants."""
|
||||
manager = MCPServerManager()
|
||||
|
||||
await manager.load_servers_from_config(
|
||||
{
|
||||
"wiki_server": {"alias": "wiki", "url": "https://wiki.example.com/mcp", "transport": MCPTransport.http},
|
||||
"docs_server": {
|
||||
"server_id": "docs_server",
|
||||
"url": "https://example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
}
|
||||
)
|
||||
wiki_id = next(sid for sid, server in manager.config_mcp_servers.items() if server.alias == "wiki")
|
||||
|
||||
assert manager.expand_permission_list(["docs_server"]) == ["docs_server"]
|
||||
assert manager.expand_permission_list(["wiki"]) == [wiki_id]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_derived_id_is_not_checked_against_names(self):
|
||||
"""Unpinned configs must keep loading; only a pinned id can be an authoring error."""
|
||||
manager = MCPServerManager()
|
||||
|
||||
await manager.load_servers_from_config(
|
||||
{
|
||||
"wiki_server": {"url": "https://wiki.example.com/mcp", "transport": MCPTransport.http},
|
||||
"docs_server": {"url": "https://example.com/mcp", "transport": MCPTransport.http},
|
||||
}
|
||||
)
|
||||
|
||||
assert len(manager.config_mcp_servers) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shadow_warning_is_not_repeated_on_every_reload(self, caplog):
|
||||
"""reload_servers_from_database runs on the config-reload timer; one warning, not one a tick."""
|
||||
manager = MCPServerManager()
|
||||
await manager.load_servers_from_config(self._config(server_id="docs-prod-1"))
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
await self._reload_with_db_server(manager, "docs-prod-1")
|
||||
first_round = [m for m in caplog.messages if "database entry takes precedence" in m]
|
||||
await self._reload_with_db_server(manager, "docs-prod-1")
|
||||
second_round = [m for m in caplog.messages if "database entry takes precedence" in m]
|
||||
|
||||
assert len(first_round) == 1
|
||||
assert second_round == first_round
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shadow_warning_fires_again_when_the_shadowed_set_changes(self, caplog):
|
||||
manager = MCPServerManager()
|
||||
await manager.load_servers_from_config(self._config(server_id="docs-prod-1"))
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
await self._reload_with_db_server(manager, "docs-prod-1")
|
||||
await self._reload_with_db_server(manager, "db-uuid-1")
|
||||
await self._reload_with_db_server(manager, "docs-prod-1")
|
||||
|
||||
assert len([m for m in caplog.messages if "database entry takes precedence" in m]) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinned_id_matching_a_mapped_alias_is_rejected(self):
|
||||
"""An alias can also arrive from litellm_settings.mcp_aliases; it is reserved just the same."""
|
||||
manager = MCPServerManager()
|
||||
|
||||
with pytest.raises(ValueError, match="server_name or alias of MCP server 'wiki_server'"):
|
||||
await manager.load_servers_from_config(
|
||||
{
|
||||
"wiki_server": {"url": "https://wiki.example.com/mcp", "transport": MCPTransport.http},
|
||||
"docs_server": {
|
||||
"server_id": "wiki",
|
||||
"url": "https://example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
},
|
||||
{"wiki": "wiki_server"},
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinning_a_servers_own_mapped_alias_is_allowed(self):
|
||||
manager = MCPServerManager()
|
||||
|
||||
await manager.load_servers_from_config(
|
||||
self._config(server_id="docs"),
|
||||
{"docs": "docs_server"},
|
||||
)
|
||||
|
||||
assert list(manager.config_mcp_servers) == ["docs"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mapped_alias_for_an_unknown_server_reserves_nothing(self):
|
||||
"""A dangling mcp_aliases entry is never applied, so it must not fail an unrelated pin."""
|
||||
manager = MCPServerManager()
|
||||
|
||||
await manager.load_servers_from_config(
|
||||
self._config(server_id="wiki"),
|
||||
{"wiki": "a_server_that_does_not_exist"},
|
||||
)
|
||||
|
||||
assert list(manager.config_mcp_servers) == ["wiki"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_id_that_is_a_db_server_name_warns(self, caplog):
|
||||
"""The mirror of the shadow case: here the config entry captures the db server's grants."""
|
||||
manager = MCPServerManager()
|
||||
await manager.load_servers_from_config(self._config(server_id="db_server"))
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
await self._reload_with_db_server(manager, "db-uuid-1")
|
||||
|
||||
assert any("db_server" in m and "name or alias of a database-backed" in m for m in caplog.messages)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_capture_warning_is_not_repeated_on_every_reload(self, caplog):
|
||||
manager = MCPServerManager()
|
||||
await manager.load_servers_from_config(self._config(server_id="db_server"))
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
await self._reload_with_db_server(manager, "db-uuid-1")
|
||||
await self._reload_with_db_server(manager, "db-uuid-1")
|
||||
|
||||
assert len([m for m in caplog.messages if "name or alias of a database-backed" in m]) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_id_unrelated_to_db_names_does_not_warn(self, caplog):
|
||||
manager = MCPServerManager()
|
||||
await manager.load_servers_from_config(self._config(server_id="docs-prod-1"))
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
await self._reload_with_db_server(manager, "db-uuid-1")
|
||||
|
||||
assert all("name or alias of a database-backed" not in m for m in caplog.messages)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mapped_alias_for_a_server_with_its_own_alias_reserves_nothing(self):
|
||||
"""load_servers_from_config ignores the mapping when the entry sets alias, so it is free."""
|
||||
manager = MCPServerManager()
|
||||
|
||||
await manager.load_servers_from_config(
|
||||
{
|
||||
"wiki_server": {
|
||||
"alias": "wiki_prod",
|
||||
"url": "https://wiki.example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
"docs_server": {
|
||||
"server_id": "wiki",
|
||||
"url": "https://example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
},
|
||||
{"wiki": "wiki_server"},
|
||||
)
|
||||
|
||||
assert "wiki" in manager.config_mcp_servers
|
||||
assert len(manager.config_mcp_servers) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_the_first_mapped_alias_for_a_server_is_reserved(self):
|
||||
"""Only the first mapping is applied, so pinning the second one must still load."""
|
||||
manager = MCPServerManager()
|
||||
|
||||
await manager.load_servers_from_config(
|
||||
{
|
||||
"wiki_server": {"url": "https://wiki.example.com/mcp", "transport": MCPTransport.http},
|
||||
"docs_server": {
|
||||
"server_id": "wiki_two",
|
||||
"url": "https://example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
},
|
||||
{"wiki_one": "wiki_server", "wiki_two": "wiki_server"},
|
||||
)
|
||||
|
||||
assert "wiki_two" in manager.config_mcp_servers
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_name_is_reported_before_any_entry_body_is_read(self):
|
||||
"""The identifier index walks every entry up front, so a bad name must still fail on the name."""
|
||||
with pytest.raises(Exception, match="Server name cannot contain"):
|
||||
await MCPServerManager().load_servers_from_config({"my-server": None})
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_shadowing_db_server_reports_only_the_shadow_warning(self, caplog):
|
||||
"""The db row wins the id outright, so the capture message would contradict the shadow one."""
|
||||
manager = MCPServerManager()
|
||||
await manager.load_servers_from_config(self._config(server_id="db_server"))
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
await self._reload_with_db_server(manager, "db_server")
|
||||
|
||||
assert any("database entry takes precedence" in m for m in caplog.messages)
|
||||
assert all("name or alias of a database-backed" not in m for m in caplog.messages)
|
||||
assert manager.get_registry()["db_server"].url == "https://db.example.com/mcp"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_explicitly_blank_alias_still_blocks_the_mapping(self):
|
||||
"""The loader only consults mcp_aliases when the key is absent, so a blank alias frees it."""
|
||||
manager = MCPServerManager()
|
||||
|
||||
await manager.load_servers_from_config(
|
||||
{
|
||||
"wiki_server": {
|
||||
"alias": "",
|
||||
"url": "https://wiki.example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
"docs_server": {
|
||||
"server_id": "wiki",
|
||||
"url": "https://example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
},
|
||||
{"wiki": "wiki_server"},
|
||||
)
|
||||
|
||||
assert "wiki" in manager.config_mcp_servers
|
||||
assert manager.config_mcp_servers["wiki"].url == "https://example.com/mcp"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_row_that_shadows_one_id_still_reports_capturing_another(self, caplog):
|
||||
"""Skipping is per identifier, not per row, so the second collision is not lost."""
|
||||
manager = MCPServerManager()
|
||||
await manager.load_servers_from_config(
|
||||
{
|
||||
"docs_server": {
|
||||
"server_id": "shadow_x",
|
||||
"url": "https://example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
"wiki_server": {
|
||||
"server_id": "capture_y",
|
||||
"url": "https://wiki.example.com/mcp",
|
||||
"transport": MCPTransport.http,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
await self._reload_with_db_server(manager, "shadow_x", db_name="capture_y")
|
||||
|
||||
assert any("shadow_x" in m and "database entry takes precedence" in m for m in caplog.messages)
|
||||
assert any("capture_y" in m and "name or alias of a database-backed" in m for m in caplog.messages)
|
||||
|
||||
|
||||
class TestLitellmAdmissionKeyIsNeverTheSubjectToken:
|
||||
"""The bearer that admitted the request as a LiteLLM key must not be sent to the IdP as the
|
||||
RFC 8693 subject_token (or ID-JAG assertion). Only ``x-litellm-api-key`` disambiguates: with it
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue