mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
feat(mcp): add provider registration configuration and management
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
86ff6f994c
commit
20cf6ba70b
19 changed files with 1106 additions and 154 deletions
|
|
@ -37315,6 +37315,14 @@
|
|||
"title": "MCPEnvVarScope",
|
||||
"type": "string"
|
||||
},
|
||||
"MCPExecutionMode": {
|
||||
"enum": [
|
||||
"litellm",
|
||||
"provider"
|
||||
],
|
||||
"title": "MCPExecutionMode",
|
||||
"type": "string"
|
||||
},
|
||||
"NewMCPServerRequest": {
|
||||
"properties": {
|
||||
"alias": {
|
||||
|
|
@ -37506,10 +37514,8 @@
|
|||
"title": "Env Vars"
|
||||
},
|
||||
"execution_mode": {
|
||||
"const": "litellm",
|
||||
"default": "litellm",
|
||||
"title": "Execution Mode",
|
||||
"type": "string"
|
||||
"$ref": "#/components/schemas/MCPExecutionMode",
|
||||
"default": "litellm"
|
||||
},
|
||||
"extra_headers": {
|
||||
"anyOf": [
|
||||
|
|
@ -37797,14 +37803,21 @@
|
|||
"title": "Tool Name To Display Name"
|
||||
},
|
||||
"transport": {
|
||||
"default": "sse",
|
||||
"enum": [
|
||||
"sse",
|
||||
"http",
|
||||
"stdio"
|
||||
"anyOf": [
|
||||
{
|
||||
"enum": [
|
||||
"sse",
|
||||
"http",
|
||||
"stdio"
|
||||
],
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Transport",
|
||||
"type": "string"
|
||||
"default": "sse",
|
||||
"title": "Transport"
|
||||
},
|
||||
"url": {
|
||||
"anyOf": [
|
||||
|
|
@ -40359,10 +40372,8 @@
|
|||
"title": "Env Vars"
|
||||
},
|
||||
"execution_mode": {
|
||||
"const": "litellm",
|
||||
"default": "litellm",
|
||||
"title": "Execution Mode",
|
||||
"type": "string"
|
||||
"$ref": "#/components/schemas/MCPExecutionMode",
|
||||
"default": "litellm"
|
||||
},
|
||||
"extra_headers": {
|
||||
"anyOf": [
|
||||
|
|
@ -40650,14 +40661,21 @@
|
|||
"title": "Tool Name To Display Name"
|
||||
},
|
||||
"transport": {
|
||||
"default": "sse",
|
||||
"enum": [
|
||||
"sse",
|
||||
"http",
|
||||
"stdio"
|
||||
"anyOf": [
|
||||
{
|
||||
"enum": [
|
||||
"sse",
|
||||
"http",
|
||||
"stdio"
|
||||
],
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Transport",
|
||||
"type": "string"
|
||||
"default": "sse",
|
||||
"title": "Transport"
|
||||
},
|
||||
"url": {
|
||||
"anyOf": [
|
||||
|
|
@ -44427,10 +44445,8 @@
|
|||
"title": "Env Vars"
|
||||
},
|
||||
"execution_mode": {
|
||||
"const": "litellm",
|
||||
"default": "litellm",
|
||||
"title": "Execution Mode",
|
||||
"type": "string"
|
||||
"$ref": "#/components/schemas/MCPExecutionMode",
|
||||
"default": "litellm"
|
||||
},
|
||||
"extra_headers": {
|
||||
"anyOf": [
|
||||
|
|
@ -44718,14 +44734,21 @@
|
|||
"title": "Tool Name To Display Name"
|
||||
},
|
||||
"transport": {
|
||||
"default": "sse",
|
||||
"enum": [
|
||||
"sse",
|
||||
"http",
|
||||
"stdio"
|
||||
"anyOf": [
|
||||
{
|
||||
"enum": [
|
||||
"sse",
|
||||
"http",
|
||||
"stdio"
|
||||
],
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Transport",
|
||||
"type": "string"
|
||||
"default": "sse",
|
||||
"title": "Transport"
|
||||
},
|
||||
"url": {
|
||||
"anyOf": [
|
||||
|
|
@ -45241,14 +45264,21 @@
|
|||
"title": "Tool Name To Display Name"
|
||||
},
|
||||
"transport": {
|
||||
"default": "sse",
|
||||
"enum": [
|
||||
"sse",
|
||||
"http",
|
||||
"stdio"
|
||||
"anyOf": [
|
||||
{
|
||||
"enum": [
|
||||
"sse",
|
||||
"http",
|
||||
"stdio"
|
||||
],
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Transport",
|
||||
"type": "string"
|
||||
"default": "sse",
|
||||
"title": "Transport"
|
||||
},
|
||||
"url": {
|
||||
"anyOf": [
|
||||
|
|
@ -47706,6 +47736,14 @@
|
|||
"title": "MCPEnvVarScope",
|
||||
"type": "string"
|
||||
},
|
||||
"MCPExecutionMode": {
|
||||
"enum": [
|
||||
"litellm",
|
||||
"provider"
|
||||
],
|
||||
"title": "MCPExecutionMode",
|
||||
"type": "string"
|
||||
},
|
||||
"NewMCPServerRequest": {
|
||||
"properties": {
|
||||
"alias": {
|
||||
|
|
@ -47897,10 +47935,8 @@
|
|||
"title": "Env Vars"
|
||||
},
|
||||
"execution_mode": {
|
||||
"const": "litellm",
|
||||
"default": "litellm",
|
||||
"title": "Execution Mode",
|
||||
"type": "string"
|
||||
"$ref": "#/components/schemas/MCPExecutionMode",
|
||||
"default": "litellm"
|
||||
},
|
||||
"extra_headers": {
|
||||
"anyOf": [
|
||||
|
|
@ -48188,14 +48224,21 @@
|
|||
"title": "Tool Name To Display Name"
|
||||
},
|
||||
"transport": {
|
||||
"default": "sse",
|
||||
"enum": [
|
||||
"sse",
|
||||
"http",
|
||||
"stdio"
|
||||
"anyOf": [
|
||||
{
|
||||
"enum": [
|
||||
"sse",
|
||||
"http",
|
||||
"stdio"
|
||||
],
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Transport",
|
||||
"type": "string"
|
||||
"default": "sse",
|
||||
"title": "Transport"
|
||||
},
|
||||
"url": {
|
||||
"anyOf": [
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import enum
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType as MappingProxyType # noqa: PLC0414 # public re-export
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple, TypeAlias
|
||||
|
|
@ -15,7 +16,9 @@ from pydantic import (
|
|||
Field,
|
||||
Json,
|
||||
JsonValue,
|
||||
ModelWrapValidatorHandler,
|
||||
PositiveInt,
|
||||
PrivateAttr,
|
||||
TypeAdapter,
|
||||
field_validator,
|
||||
model_validator,
|
||||
|
|
@ -65,13 +68,13 @@ from litellm.types.mcp import (
|
|||
MCPCredentials,
|
||||
MCPDisabledCapabilities,
|
||||
MCPExecutionMode,
|
||||
MCPGatewayExecutionMode,
|
||||
MCPTransport,
|
||||
MCPTransportType,
|
||||
MCPUpstreamProtocol,
|
||||
validate_mcp_protocol_transport,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPInfo
|
||||
from litellm.types.mcp_server.provider_registration import active_registration_extras, normalize_provider_registration
|
||||
from litellm.types.proxy.agent_identity import ManagedAgentContext
|
||||
from litellm.types.proxy.auth.special_headers import (
|
||||
SpecialHeaders as SpecialHeaders, # noqa: PLC0414 # public re-export
|
||||
|
|
@ -1388,12 +1391,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
|
||||
|
|
@ -1454,6 +1457,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")
|
||||
|
|
@ -1465,7 +1472,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")
|
||||
|
|
@ -1497,11 +1512,13 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
|
|||
|
||||
|
||||
class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
|
||||
_provider_extra_fields: tuple[str, ...] = PrivateAttr(default=())
|
||||
|
||||
server_id: str
|
||||
server_name: str | None = None
|
||||
alias: str | None = None
|
||||
description: str | None = None
|
||||
transport: MCPTransportType = MCPTransport.sse
|
||||
transport: MCPTransportType | None = MCPTransport.sse
|
||||
auth_type: MCPAuthType | None = None
|
||||
credentials: MCPCredentials | None = None
|
||||
url: str | None = None
|
||||
|
|
@ -1546,6 +1563,28 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
|
|||
max_concurrent_requests: int | None = None
|
||||
rpm: int | None = Field(default=None, ge=0)
|
||||
|
||||
@model_validator(mode="wrap")
|
||||
@classmethod
|
||||
def retain_extra_field_names(
|
||||
cls,
|
||||
values: object,
|
||||
handler: ModelWrapValidatorHandler["UpdateMCPServerRequest"],
|
||||
) -> "UpdateMCPServerRequest":
|
||||
validated: Final = handler(values)
|
||||
if isinstance(values, Mapping):
|
||||
validated._provider_extra_fields = active_registration_extras(
|
||||
TypeAdapter(dict[str, object]).validate_python(values),
|
||||
cls.model_fields.keys(),
|
||||
)
|
||||
return validated
|
||||
|
||||
def provider_registration_patch(self, fields_set: AbstractSet[str] | None = None) -> Mapping[str, object]:
|
||||
if self._provider_extra_fields:
|
||||
raise ValueError(
|
||||
f"Provider-executed MCP registrations do not accept {', '.join(self._provider_extra_fields)}"
|
||||
)
|
||||
return self.model_dump(include=set(fields_set) if fields_set is not None else None, exclude_unset=True)
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def reject_execution_mode_changes(cls, values: object) -> object:
|
||||
|
|
@ -1555,7 +1594,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(
|
||||
|
|
|
|||
|
|
@ -50,6 +50,7 @@ from pydantic import TypeAdapter
|
|||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.proxy.mcp_server.catalog import catalog_operation, public_catalog_operation
|
||||
from litellm.types.mcp import MCPExecutionMode
|
||||
|
||||
try:
|
||||
from prisma.errors import RecordNotFoundError, UniqueViolationError
|
||||
|
|
@ -257,6 +258,7 @@ if MCP_AVAILABLE:
|
|||
)
|
||||
from litellm.proxy.mcp_server.server_resolution import (
|
||||
MCPServerTargetCatalog,
|
||||
ResolvedMCPServer,
|
||||
authorize_mcp_server,
|
||||
resolve_mcp_server,
|
||||
)
|
||||
|
|
@ -272,6 +274,7 @@ if MCP_AVAILABLE:
|
|||
MCPCredentials,
|
||||
MCPGatewaySessionsResponse,
|
||||
MCPGatewaySessionsTerminateResponse,
|
||||
MCPTransportType,
|
||||
MCPUpstreamProtocol,
|
||||
normalize_upstream_header_name,
|
||||
validate_mcp_protocol_transport,
|
||||
|
|
@ -1342,13 +1345,34 @@ if MCP_AVAILABLE:
|
|||
aligned with the cards actually rendered: an admin in view_all mode sees
|
||||
every server even when their key carries no per-server MCP grant.
|
||||
"""
|
||||
if _get_user_mcp_management_mode() == "view_all" and not _is_restricted_virtual_key_request(user_api_key_dict):
|
||||
return await global_mcp_server_manager.get_all_mcp_servers_unfiltered()
|
||||
unrestricted: Final = not _is_restricted_virtual_key_request(user_api_key_dict)
|
||||
view_all: Final = unrestricted and _get_user_mcp_management_mode() == "view_all"
|
||||
view_providers: Final = unrestricted and user_api_key_has_admin_view(user_api_key_dict)
|
||||
inventory: Final = (
|
||||
await global_mcp_server_manager.get_all_mcp_servers_unfiltered() if view_all or view_providers else []
|
||||
)
|
||||
visible_inventory: Final = [
|
||||
server
|
||||
for server in inventory
|
||||
if view_providers
|
||||
or getattr(server, "execution_mode", MCPExecutionMode.litellm) != MCPExecutionMode.provider
|
||||
]
|
||||
if view_all:
|
||||
return visible_inventory
|
||||
|
||||
aggregated: Final[dict[str, LiteLLM_MCPServerTable]] = {}
|
||||
for auth_context in await build_effective_auth_contexts(user_api_key_dict):
|
||||
for server in await global_mcp_server_manager.get_all_allowed_mcp_servers(user_api_key_auth=auth_context):
|
||||
aggregated.setdefault(server.server_id, server)
|
||||
provider_servers: Final = [
|
||||
server
|
||||
for server in visible_inventory
|
||||
if getattr(server, "execution_mode", MCPExecutionMode.litellm) == MCPExecutionMode.provider
|
||||
]
|
||||
allowed_by_context: Final = [
|
||||
await global_mcp_server_manager.get_all_allowed_mcp_servers(user_api_key_auth=auth_context)
|
||||
for auth_context in await build_effective_auth_contexts(user_api_key_dict)
|
||||
]
|
||||
aggregated: Final = {
|
||||
server.server_id: server
|
||||
for server in (*chain.from_iterable(reversed(allowed_by_context)), *provider_servers)
|
||||
}
|
||||
return list(aggregated.values())
|
||||
|
||||
@catalog_operation(lambda: global_mcp_server_manager)
|
||||
|
|
@ -1949,20 +1973,40 @@ if MCP_AVAILABLE:
|
|||
--header 'Authorization: Bearer your_api_key_here'
|
||||
```
|
||||
"""
|
||||
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
|
||||
|
||||
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
|
||||
|
||||
client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
|
||||
is_admin_view: Final = user_api_key_has_admin_view(user_api_key_dict)
|
||||
is_restricted_virtual_key: Final = _is_restricted_virtual_key_request(user_api_key_dict)
|
||||
resolved: Final = await resolve_mcp_server(
|
||||
server_id,
|
||||
manager=global_mcp_server_manager,
|
||||
db_lookup=lambda sid: get_mcp_server(prisma_client, sid),
|
||||
id_client_ip=client_ip,
|
||||
name_client_ip=client_ip,
|
||||
match_name=True,
|
||||
snapshot: Final = global_mcp_server_manager.catalog.current()
|
||||
config_match: Final = (
|
||||
next(
|
||||
(
|
||||
server
|
||||
for server in snapshot.management_servers.values()
|
||||
if server.is_config
|
||||
and server.execution_mode == MCPExecutionMode.provider
|
||||
and server_id in (server.server_id, server.server_name, server.alias)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if snapshot is not None
|
||||
else None
|
||||
)
|
||||
resolved: Final = (
|
||||
ResolvedMCPServer(table=config_match, runtime=None, source="registry")
|
||||
if config_match is not None
|
||||
else await resolve_mcp_server(
|
||||
server_id,
|
||||
manager=global_mcp_server_manager,
|
||||
db_lookup=lambda sid: get_mcp_server(
|
||||
get_prisma_client_or_throw("Database not connected. Connect a database to your proxy"),
|
||||
sid,
|
||||
),
|
||||
id_client_ip=client_ip,
|
||||
name_client_ip=client_ip,
|
||||
match_name=True,
|
||||
)
|
||||
)
|
||||
authorized: Final = await authorize_mcp_server(
|
||||
resolved,
|
||||
|
|
@ -1992,18 +2036,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 +2134,12 @@ if MCP_AVAILABLE:
|
|||
payload,
|
||||
touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
|
||||
)
|
||||
except UniqueViolationError as exc:
|
||||
raise HTTPException(
|
||||
status_code=409, detail="An MCP registration already uses this identity or provider endpoint"
|
||||
) from exc
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error creating mcp server: %s", e)
|
||||
raise HTTPException(
|
||||
|
|
@ -3190,7 +3243,7 @@ if MCP_AVAILABLE:
|
|||
TypeAdapter[MCPUpstreamProtocol](MCPUpstreamProtocol).validate_python(
|
||||
(effective_info or {}).get("protocol_version", "auto")
|
||||
),
|
||||
effective_transport,
|
||||
TypeAdapter[MCPTransportType](MCPTransportType).validate_python(effective_transport),
|
||||
)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
|
|
@ -3302,12 +3355,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)
|
||||
|
|
|
|||
|
|
@ -23,8 +23,10 @@ from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
|
|||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.models.mcp_server import LiteLLM_MCPServerTable
|
||||
from litellm.proxy.mcp_server.result_conversion import age_freshness, aggregate_freshness
|
||||
from litellm.types.mcp import MCPExecutionMode
|
||||
from litellm.types.mcp_server.provider_registration import canonical_provider_mcp_url
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mcp.types import ListToolsResult, PaginatedRequestParams, PaginatedResult
|
||||
|
|
@ -60,6 +62,7 @@ class CatalogSnapshot:
|
|||
identity: str
|
||||
tools: Mapping[str, MCPTool]
|
||||
routing: MutableMapping[str, str]
|
||||
management_servers: Mapping[str, LiteLLM_MCPServerTable]
|
||||
|
||||
|
||||
def _configuration_identity(server: MCPServer) -> str:
|
||||
|
|
@ -90,7 +93,41 @@ def _check_oauth_revision(selected: MCPServer, candidate: MCPServer | None) -> N
|
|||
raise HTTPException(status_code=503, detail="OAuth metadata discovery changed repeatedly; retry shortly")
|
||||
|
||||
|
||||
def _snapshot(manager: MCPServerManager, database_identity: str) -> CatalogSnapshot:
|
||||
def _provider_registration_conflicts(first: LiteLLM_MCPServerTable, second: LiteLLM_MCPServerTable) -> bool:
|
||||
if MCPExecutionMode.provider not in (first.execution_mode, second.execution_mode):
|
||||
return False
|
||||
identifiers: Final = frozenset(
|
||||
value.lower() for value in (first.server_id, first.server_name, first.alias) if value
|
||||
)
|
||||
other_identifiers: Final = frozenset(
|
||||
value.lower() for value in (second.server_id, second.server_name, second.alias) if value
|
||||
)
|
||||
same_provider_url: Final = (
|
||||
first.execution_mode == second.execution_mode == MCPExecutionMode.provider
|
||||
and first.url is not None
|
||||
and second.url is not None
|
||||
and canonical_provider_mcp_url(first.url) == canonical_provider_mcp_url(second.url)
|
||||
)
|
||||
return bool(identifiers & other_identifiers) or same_provider_url
|
||||
|
||||
|
||||
def validate_provider_registration_conflicts(
|
||||
config_servers: Sequence[LiteLLM_MCPServerTable],
|
||||
database_servers: Sequence[LiteLLM_MCPServerTable] = (),
|
||||
) -> None:
|
||||
for index, config_server in enumerate(config_servers):
|
||||
for other in (*config_servers[index + 1 :], *database_servers):
|
||||
if _provider_registration_conflicts(config_server, other):
|
||||
raise ValueError(
|
||||
f"MCP config registration '{config_server.server_name}' conflicts with registration '{other.server_name}'"
|
||||
)
|
||||
|
||||
|
||||
def _snapshot(
|
||||
manager: MCPServerManager,
|
||||
database_identity: str,
|
||||
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
|
||||
|
|
@ -98,6 +135,7 @@ def _snapshot(manager: MCPServerManager, database_identity: str) -> CatalogSnaps
|
|||
serialized: Final = json.dumps(
|
||||
(
|
||||
database_identity,
|
||||
tuple(server.model_dump(mode="json") for server in manager.catalog.config_registrations.values()),
|
||||
tuple(sorted(manager.registry)),
|
||||
tuple((key, _configuration_identity(value)) for key, value in sorted(manager.config_mcp_servers.items())),
|
||||
),
|
||||
|
|
@ -108,6 +146,13 @@ def _snapshot(manager: MCPServerManager, database_identity: str) -> CatalogSnaps
|
|||
hashlib.sha256(serialized.encode()).hexdigest(),
|
||||
MappingProxyType(dict(global_mcp_tool_registry.published_tools)),
|
||||
dict(manager.published_tool_routes),
|
||||
MappingProxyType(
|
||||
{
|
||||
**{key: value.model_copy(deep=True) for key, value in manager.catalog.config_registrations.items()},
|
||||
**{key: manager.build_mcp_server_table(value) for key, value in servers.items()},
|
||||
**{key: value.model_copy(deep=True) for key, value in registration_only_servers.items()},
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -116,6 +161,8 @@ class CatalogSnapshots:
|
|||
self.manager = manager
|
||||
self._refresh_lock = asyncio.Lock()
|
||||
self._database_identity = ""
|
||||
self._registration_only_servers: Mapping[str, LiteLLM_MCPServerTable] = MappingProxyType({})
|
||||
self.config_registrations: Mapping[str, LiteLLM_MCPServerTable] = MappingProxyType({})
|
||||
self._arrival_ticket = 0
|
||||
self._completed_ticket = 0
|
||||
self._shared_snapshot: CatalogSnapshot | None = None
|
||||
|
|
@ -129,6 +176,24 @@ class CatalogSnapshots:
|
|||
"mcp_catalog_routing", default=None
|
||||
)
|
||||
|
||||
def validate_config_inventory(self, servers: Sequence[LiteLLM_MCPServerTable]) -> None:
|
||||
database_servers: Final = (
|
||||
*(self.manager.build_mcp_server_table(server) for server in self.manager.registry.values()),
|
||||
*self._registration_only_servers.values(),
|
||||
)
|
||||
validate_provider_registration_conflicts(servers, database_servers)
|
||||
|
||||
def publish_config_registrations(self, servers: Mapping[str, LiteLLM_MCPServerTable]) -> None:
|
||||
self.config_registrations = MappingProxyType(dict(servers))
|
||||
self._shared_snapshot = None
|
||||
|
||||
def validate_database_registration(self, server: LiteLLM_MCPServerTable) -> None:
|
||||
config_servers: Final = (
|
||||
*self.config_registrations.values(),
|
||||
*(self.manager.build_mcp_server_table(value) for value in self.manager.config_mcp_servers.values()),
|
||||
)
|
||||
validate_provider_registration_conflicts(config_servers, (server,))
|
||||
|
||||
def current(self) -> CatalogSnapshot | None:
|
||||
scoped: Final = self._operation.get()
|
||||
return scoped[0] if scoped is not None and not scoped[1].is_set() else None
|
||||
|
|
@ -150,11 +215,11 @@ class CatalogSnapshots:
|
|||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
return _snapshot(self.manager, self._database_identity)
|
||||
return _snapshot(self.manager, self._database_identity, self._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 +241,10 @@ class CatalogSnapshots:
|
|||
async with self.operation() as snapshot:
|
||||
return snapshot.servers
|
||||
|
||||
async def management_inventory(self) -> Mapping[str, LiteLLM_MCPServerTable]:
|
||||
async with self.operation() as snapshot:
|
||||
return snapshot.management_servers
|
||||
|
||||
def assert_current(self, server: MCPServer) -> None:
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
|
@ -208,6 +277,9 @@ class CatalogSnapshots:
|
|||
shared,
|
||||
servers=MappingProxyType({key: value.model_copy(deep=True) for key, value in shared.servers.items()}),
|
||||
routing=routing,
|
||||
management_servers=MappingProxyType(
|
||||
{key: value.model_copy(deep=True) for key, value in shared.management_servers.items()}
|
||||
),
|
||||
)
|
||||
closed: Final = asyncio.Event()
|
||||
token: Final = self._operation.set((snapshot, closed))
|
||||
|
|
@ -288,7 +360,7 @@ class CatalogSnapshots:
|
|||
self._completed_ticket = covered
|
||||
raise
|
||||
else:
|
||||
self._shared_snapshot = _snapshot(self.manager, self._database_identity)
|
||||
self._shared_snapshot = _snapshot(self.manager, self._database_identity, self._registration_only_servers)
|
||||
self._applied_revision = revision
|
||||
self._completed_ticket = covered
|
||||
finally:
|
||||
|
|
@ -490,6 +562,24 @@ class CatalogSnapshots:
|
|||
).hexdigest()
|
||||
verbose_logger.info("Found %s MCP servers in database", len(raw_rows))
|
||||
|
||||
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
|
||||
)
|
||||
}
|
||||
)
|
||||
config_servers: Final = (
|
||||
*self.config_registrations.values(),
|
||||
*(self.manager.build_mcp_server_table(value) for value in self.manager.config_mcp_servers.values()),
|
||||
)
|
||||
if self.config_registrations or registration_only_servers:
|
||||
validate_provider_registration_conflicts(
|
||||
config_servers, tuple(LiteLLM_MCPServerTable.model_validate(row.model_dump()) for row in raw_rows)
|
||||
)
|
||||
previous_registry: Final = self.manager.registry
|
||||
new_registry: Final = await self._stage_servers(raw_rows, reuse_unchanged=reuse_unchanged)
|
||||
|
||||
|
|
@ -569,6 +659,7 @@ class CatalogSnapshots:
|
|||
", ".join(sorted(capturing_config_server_ids)),
|
||||
)
|
||||
self._warned_capturing_config_server_ids = capturing_config_server_ids
|
||||
self._registration_only_servers = registration_only_servers
|
||||
|
||||
|
||||
def catalog_operation(
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import binascii
|
|||
import hashlib
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast
|
||||
|
|
@ -57,7 +58,13 @@ from litellm.repositories.verification_token_repository import (
|
|||
VerificationTokenRepository,
|
||||
)
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
from litellm.types.mcp import MCPCredentials, MCPTransportType, MCPUpstreamProtocol, validate_mcp_protocol_transport
|
||||
from litellm.types.mcp import (
|
||||
MCPCredentials,
|
||||
MCPExecutionMode,
|
||||
MCPTransportType,
|
||||
MCPUpstreamProtocol,
|
||||
validate_mcp_protocol_transport,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPInfo, PinnedMCPTool
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -452,6 +459,8 @@ def _prepare_mcp_server_data(
|
|||
data_dict[json_map_field] = {}
|
||||
else:
|
||||
data_dict = data.model_dump(exclude_none=True)
|
||||
if isinstance(data, NewMCPServerRequest) and data.execution_mode == MCPExecutionMode.provider:
|
||||
data_dict["transport"] = data.transport
|
||||
# Ensure alias is always present in the dict (even if None)
|
||||
if "alias" not in data_dict:
|
||||
data_dict["alias"] = getattr(data, "alias", None)
|
||||
|
|
@ -1076,16 +1085,46 @@ async def delete_mcp_server(
|
|||
return deleted_server # pyright: ignore[reportReturnType] # prisma row, not domain LiteLLM_MCPServerTable
|
||||
|
||||
|
||||
def _validate_mcp_config_registration_write(values: Mapping[str, object]) -> None:
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager
|
||||
|
||||
catalog: Final = global_mcp_server_manager.catalog
|
||||
if not catalog.config_registrations and values.get("execution_mode") != MCPExecutionMode.provider:
|
||||
return
|
||||
identity: Final = {
|
||||
key: values[key]
|
||||
for key in (
|
||||
"server_id",
|
||||
"server_name",
|
||||
"alias",
|
||||
"url",
|
||||
"transport",
|
||||
"execution_mode",
|
||||
)
|
||||
if key in values
|
||||
}
|
||||
try:
|
||||
catalog.validate_database_registration(LiteLLM_MCPServerTable.model_validate(identity))
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
|
||||
|
||||
async def create_mcp_server(
|
||||
prisma_client: PrismaClient, data: NewMCPServerRequest, touched_by: str
|
||||
) -> LiteLLM_MCPServerTable:
|
||||
"""
|
||||
Create a new mcp server record in the db
|
||||
"""
|
||||
if data.execution_mode == MCPExecutionMode.provider:
|
||||
provider_record: Final = await create_mcp_server_if_identifier_free(prisma_client, data, touched_by)
|
||||
if isinstance(provider_record, McpIdentifierConflict):
|
||||
raise HTTPException(status_code=409, detail="An MCP registration already uses this identity")
|
||||
return provider_record
|
||||
if data.server_id is None:
|
||||
data.server_id = str(uuid.uuid4())
|
||||
|
||||
# Use helper to prepare data with proper JSON serialization
|
||||
_validate_mcp_config_registration_write(data.model_dump())
|
||||
data_dict: Final = _prepare_mcp_server_data(data)
|
||||
|
||||
# Add audit fields
|
||||
|
|
@ -1110,6 +1149,7 @@ async def create_mcp_server_if_identifier_free(
|
|||
if data.server_id is None:
|
||||
data.server_id = str(uuid.uuid4())
|
||||
|
||||
_validate_mcp_config_registration_write(data.model_dump())
|
||||
data_dict: Final = _prepare_mcp_server_data(data)
|
||||
data_dict["created_by"] = touched_by
|
||||
data_dict["updated_by"] = touched_by
|
||||
|
|
@ -1246,35 +1286,75 @@ 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: AbstractSet[str] | None = None,
|
||||
) -> NewMCPServerRequest:
|
||||
try:
|
||||
patch: Final = data.provider_registration_patch(fields_set)
|
||||
registration: Final = LiteLLM_MCPServerTable.model_validate(stored.model_dump())
|
||||
merged: Final = {**registration.model_dump(include=set(NewMCPServerRequest.model_fields)), **patch}
|
||||
return NewMCPServerRequest.model_validate(merged)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
|
||||
|
||||
def _prepare_provider_registration_update(
|
||||
stored: "prisma_db_models.LiteLLM_MCPServerTable",
|
||||
data: UpdateMCPServerRequest,
|
||||
touched_by: str,
|
||||
fields_set: AbstractSet[str] | None,
|
||||
) -> Mapping[str, object]:
|
||||
from litellm.types.mcp_server.provider_registration import ProviderMCPRegistration
|
||||
|
||||
validated: Final = _validate_provider_registration_update(stored, data, fields_set)
|
||||
_validate_mcp_config_registration_write(validated.model_dump())
|
||||
supplied_fields: Final = fields_set if fields_set is not None else data.model_fields_set
|
||||
normalized_patch: Final = UpdateMCPServerRequest.model_validate(
|
||||
validated.model_dump(
|
||||
include=set(supplied_fields & ProviderMCPRegistration.model_fields.keys()) | {"server_id", "mcp_info"}
|
||||
)
|
||||
)
|
||||
serialized: Final = _prepare_mcp_server_data(normalized_patch, exclude_unset=True, fields_set=set(supplied_fields))
|
||||
return {
|
||||
**serialized,
|
||||
"updated_by": touched_by,
|
||||
}
|
||||
|
||||
|
||||
async def _update_mcp_server_row(
|
||||
prisma_client: PrismaClient,
|
||||
*,
|
||||
server_id: str,
|
||||
data_dict: Mapping[str, object],
|
||||
prepare_update: Callable[["prisma_db_models.LiteLLM_MCPServerTable"], Mapping[str, object]] | None = None,
|
||||
expected_updated_at: datetime | None = None,
|
||||
) -> "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 +1368,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 +1376,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 +1395,24 @@ 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())
|
||||
)
|
||||
|
||||
if registration is not None:
|
||||
_validate_mcp_config_registration_write({**registration.model_dump(), **data.model_dump(exclude_unset=True)})
|
||||
|
||||
# Use helper to prepare data with proper JSON serialization.
|
||||
# exclude_unset=True makes this a true partial update: fields the caller did
|
||||
# not provide are not written, so they keep their existing DB value instead
|
||||
|
|
@ -1322,15 +1420,17 @@ async def update_mcp_server(
|
|||
prepared_data: Final = _prepare_mcp_server_data(data, exclude_unset=True, fields_set=fields_set)
|
||||
|
||||
# Pre-fetch existing record once if we need it for auth_type, url, or credential logic
|
||||
existing = None
|
||||
has_credentials: Final = "credentials" in prepared_data and prepared_data["credentials"] is not None
|
||||
# An explicit token-exchange column write (set or clear) also migrates the
|
||||
# legacy blob copies below, so the existing row is needed for those updates.
|
||||
explicit_te_write: Final = bool(_TOKEN_EXCHANGE_COLUMN_FIELDS & prepared_data.keys())
|
||||
url_provided: Final = "url" in prepared_data and prepared_data["url"] is not None
|
||||
issuer_provided: Final = "issuer" in prepared_data
|
||||
if data.auth_type or has_credentials or explicit_te_write or url_provided or issuer_provided:
|
||||
existing = await _db_find_mcp_server_row(prisma_client, data.server_id)
|
||||
existing: Final = (
|
||||
registration
|
||||
if data.auth_type or has_credentials or explicit_te_write or url_provided or issuer_provided
|
||||
else None
|
||||
)
|
||||
|
||||
auth_type_changed: Final = bool(
|
||||
data.auth_type
|
||||
|
|
@ -1442,9 +1542,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
|
||||
|
||||
|
|
|
|||
|
|
@ -89,6 +89,7 @@ from litellm.proxy._types import (
|
|||
MCPEnvVar,
|
||||
MCPTransport,
|
||||
MCPTransportType,
|
||||
NewMCPServerRequest,
|
||||
SpecialMCPServerNames,
|
||||
UserAPIKeyAuth,
|
||||
is_per_server_oauth_discovery_eligible,
|
||||
|
|
@ -408,6 +409,7 @@ class MCPServerConfig(TypedDict, total=False):
|
|||
|
||||
server_id: ReadOnly[str]
|
||||
execution_mode: ReadOnly[MCPExecutionMode]
|
||||
tool_name_to_description: ReadOnly[Mapping[str, str]]
|
||||
protocol_version: ReadOnly[MCPUpstreamProtocol]
|
||||
alias: str
|
||||
description: str
|
||||
|
|
@ -2558,13 +2560,83 @@ class MCPServerManager:
|
|||
store; a rowless server that is NOT config-declared is a throwaway temp/session server whose
|
||||
client must not be persisted. This never overrides the row-existence check: a server that has
|
||||
a LiteLLM_MCPServerTable row is always resolved to that row first."""
|
||||
return server_id in self.config_mcp_servers
|
||||
return server_id in self.config_mcp_servers or server_id in self.catalog.config_registrations
|
||||
|
||||
def _config_registration(
|
||||
self,
|
||||
name: str,
|
||||
values: MCPServerConfig,
|
||||
aliases: Mapping[str, str] | None,
|
||||
) -> LiteLLM_MCPServerTable:
|
||||
alias: Final = _assigned_alias(name, values, aliases)
|
||||
pinned_id: Final = _pinned_config_server_id(values.get("server_id"), name)
|
||||
if values.get("execution_mode", MCPExecutionMode.litellm) == MCPExecutionMode.provider:
|
||||
request: Final = NewMCPServerRequest.model_validate(
|
||||
{
|
||||
**{field: value for field, value in values.items() if field != "access_groups"},
|
||||
"server_name": name,
|
||||
"alias": alias,
|
||||
"mcp_access_groups": values.get("access_groups", []),
|
||||
}
|
||||
)
|
||||
stable_id: Final = hashlib.sha256(f"provider|{name}|{request.url}|{alias or ''}".encode()).hexdigest()[:32]
|
||||
return LiteLLM_MCPServerTable.model_validate(
|
||||
{
|
||||
**request.model_dump(exclude_none=True),
|
||||
"server_id": pinned_id or stable_id,
|
||||
"transport": request.transport,
|
||||
"is_config": True,
|
||||
}
|
||||
)
|
||||
validate_mcp_gateway_execution(values.get("execution_mode", MCPExecutionMode.litellm))
|
||||
return LiteLLM_MCPServerTable(
|
||||
server_id=pinned_id
|
||||
or self._generate_stable_server_id(
|
||||
name,
|
||||
values.get("url") or "",
|
||||
values.get("transport", MCPTransport.http),
|
||||
values.get("auth_type"),
|
||||
alias,
|
||||
),
|
||||
server_name=name,
|
||||
alias=alias,
|
||||
url=values.get("url"),
|
||||
transport=values.get("transport", MCPTransport.http),
|
||||
is_config=True,
|
||||
)
|
||||
|
||||
async def load_servers_from_config(
|
||||
self,
|
||||
mcp_servers_config: dict[str, MCPServerConfig],
|
||||
mcp_aliases: dict[str, str] | None = None,
|
||||
):
|
||||
) -> None:
|
||||
_validate_config_server_names(mcp_servers_config)
|
||||
registrations: Final = tuple(
|
||||
self._config_registration(name, config, mcp_aliases) for name, config in mcp_servers_config.items()
|
||||
)
|
||||
self.catalog.validate_config_inventory(registrations)
|
||||
providers: Final = {
|
||||
server.server_id: server for server in registrations if server.execution_mode == MCPExecutionMode.provider
|
||||
}
|
||||
provider_names: Final = frozenset(server.server_name for server in providers.values())
|
||||
for server in tuple(self.config_mcp_servers.values()):
|
||||
if server.server_name in provider_names or server.server_id in providers:
|
||||
self._cleanup_server_tool_routing_artifacts(server)
|
||||
self.invalidate_oauth_discovery_state(server.server_id)
|
||||
self.config_mcp_servers.pop(server.server_id, None)
|
||||
gateway_config: Final = {
|
||||
name: config
|
||||
for name, config in mcp_servers_config.items()
|
||||
if config.get("execution_mode", MCPExecutionMode.litellm) != MCPExecutionMode.provider
|
||||
}
|
||||
await self._load_gateway_servers_from_config(gateway_config, mcp_aliases)
|
||||
self.catalog.publish_config_registrations(providers)
|
||||
|
||||
async def _load_gateway_servers_from_config(
|
||||
self,
|
||||
mcp_servers_config: Mapping[str, MCPServerConfig],
|
||||
mcp_aliases: Mapping[str, str] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Load the MCP Servers from the config
|
||||
|
||||
|
|
@ -7869,16 +7941,8 @@ class MCPServerManager:
|
|||
_build_mcp_server_table = build_mcp_server_table
|
||||
|
||||
async def get_all_mcp_servers_unfiltered(self) -> list[LiteLLM_MCPServerTable]:
|
||||
"""Return all MCP servers from registry without applying access controls."""
|
||||
|
||||
registry: Final = self.get_registry()
|
||||
if not registry:
|
||||
return []
|
||||
|
||||
servers: Final[list[LiteLLM_MCPServerTable]] = []
|
||||
for server in registry.values():
|
||||
servers.append(self.build_mcp_server_table(server))
|
||||
return servers
|
||||
"""Return the cached management inventory without applying access controls."""
|
||||
return list((await self.catalog.management_inventory()).values())
|
||||
|
||||
async def get_all_mcp_servers_with_health_unfiltered(
|
||||
self, server_ids: list[str] | None = None
|
||||
|
|
|
|||
|
|
@ -1671,6 +1671,9 @@ if MCP_AVAILABLE:
|
|||
Returns:
|
||||
The dict returned by *operation*, or an error dict on failure.
|
||||
"""
|
||||
if request.transport is None:
|
||||
return {"error": "Provider-executed MCP registrations do not support gateway connections"}
|
||||
|
||||
try:
|
||||
client_id, client_secret, scopes = _extract_credentials(request)
|
||||
|
||||
|
|
|
|||
|
|
@ -7612,17 +7612,18 @@ class ProxyConfig:
|
|||
|
||||
global_agent_registry.load_agents_from_config(agent_config)
|
||||
|
||||
from litellm.proxy.mcp_server.utils import is_mcp_available
|
||||
|
||||
mcp_servers_config: Final = config.get("mcp_servers", None)
|
||||
if mcp_servers_config:
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
if mcp_servers_config or is_mcp_available():
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager
|
||||
|
||||
# Get mcp_aliases from litellm_settings if available
|
||||
litellm_settings: Final = config.get("litellm_settings", {})
|
||||
mcp_aliases: Final = litellm_settings.get("mcp_aliases", None)
|
||||
|
||||
await global_mcp_server_manager.load_servers_from_config(mcp_servers_config, mcp_aliases)
|
||||
if mcp_servers_config:
|
||||
litellm_settings: Final = config.get("litellm_settings", {})
|
||||
mcp_aliases: Final = litellm_settings.get("mcp_aliases", None)
|
||||
await global_mcp_server_manager.load_servers_from_config(mcp_servers_config, mcp_aliases)
|
||||
elif global_mcp_server_manager.catalog.config_registrations:
|
||||
global_mcp_server_manager.catalog.publish_config_registrations({})
|
||||
|
||||
## VECTOR STORES
|
||||
vector_store_registry_config: Final = config.get("vector_store_registry", None)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,6 @@ from litellm.types.mcp import (
|
|||
MCPAuth,
|
||||
MCPAuthType,
|
||||
MCPExecutionMode,
|
||||
MCPGatewayExecutionMode,
|
||||
MCPTokenEndpointAuthMethod,
|
||||
MCPTransportType,
|
||||
MCPUpstreamProtocol,
|
||||
|
|
@ -98,7 +97,7 @@ def parse_pinned_tools(value: object) -> dict[str, PinnedMCPTool] | None:
|
|||
class MCPServer(LiteLLMBaseModel):
|
||||
server_id: str
|
||||
name: str
|
||||
execution_mode: MCPGatewayExecutionMode = MCPExecutionMode.litellm
|
||||
execution_mode: Literal[MCPExecutionMode.litellm] = MCPExecutionMode.litellm
|
||||
alias: str | None = None
|
||||
server_name: str | None = None
|
||||
url: str | None = None
|
||||
|
|
|
|||
129
litellm/types/mcp_server/provider_registration.py
Normal file
129
litellm/types/mcp_server/provider_registration.py
Normal file
|
|
@ -0,0 +1,129 @@
|
|||
import json
|
||||
from collections.abc import Collection, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import Final, Literal
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
from pydantic import AnyHttpUrl, BaseModel, ConfigDict, Field, TypeAdapter, field_validator, model_validator
|
||||
from typing_extensions import Self
|
||||
|
||||
from litellm.types.mcp import MCPExecutionMode
|
||||
|
||||
_STR: Final = TypeAdapter(str)
|
||||
_MAP: Final = TypeAdapter(Mapping[str, object])
|
||||
_URL: Final = TypeAdapter(AnyHttpUrl)
|
||||
|
||||
|
||||
def _is_neutral_registration_value(value: object) -> bool:
|
||||
if isinstance(value, str):
|
||||
return value.strip() in ("", "null", "{}", "[]")
|
||||
return value is None or value is False or value in ({}, [])
|
||||
|
||||
|
||||
def active_registration_extras(values: Mapping[str, object], known_fields: Collection[str]) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
field
|
||||
for field, value in values.items()
|
||||
if field not in known_fields and not _is_neutral_registration_value(value)
|
||||
)
|
||||
|
||||
|
||||
class ProviderMCPRegistration(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
execution_mode: Literal[MCPExecutionMode.provider] = MCPExecutionMode.provider
|
||||
server_id: str | None = None
|
||||
server_name: str | None = None
|
||||
alias: str | None = None
|
||||
description: str | None = None
|
||||
url: str
|
||||
allowed_tools: Sequence[str]
|
||||
mcp_info: Mapping[str, object] = Field(default_factory=dict, validate_default=True)
|
||||
mcp_access_groups: Sequence[str] = ()
|
||||
tool_name_to_display_name: Mapping[str, str] | None = None
|
||||
tool_name_to_description: Mapping[str, str] | None = None
|
||||
instructions: str | None = None
|
||||
source_url: str | None = None
|
||||
allow_all_keys: bool = False
|
||||
available_on_public_internet: Literal[True] = True
|
||||
protocol_version: Literal["auto"] = "auto"
|
||||
approval_status: str | None = None
|
||||
submitted_by: str | None = None
|
||||
submitted_at: datetime | None = None
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def discard_neutral_extras(cls, values: object) -> object:
|
||||
if not isinstance(values, Mapping):
|
||||
return values
|
||||
fields: Final = _MAP.validate_python(values)
|
||||
return {
|
||||
field: value
|
||||
for field, value in fields.items()
|
||||
if field in cls.model_fields or not _is_neutral_registration_value(value)
|
||||
}
|
||||
|
||||
@field_validator("url", mode="before")
|
||||
@classmethod
|
||||
def normalize_url(cls, value: object) -> str:
|
||||
return canonical_provider_mcp_url(value)
|
||||
|
||||
@field_validator("allowed_tools")
|
||||
@classmethod
|
||||
def normalize_allowed_tools(cls, tools: Sequence[str]) -> tuple[str, ...]:
|
||||
if any(not tool.strip() or any(char in tool for char in "*?[]") for tool in tools):
|
||||
raise ValueError("allowed_tools must contain nonblank exact tool names without wildcards")
|
||||
return tuple(dict.fromkeys(tools))
|
||||
|
||||
@field_validator("mcp_info", mode="before")
|
||||
@classmethod
|
||||
def normalize_metadata(cls, value: object) -> Mapping[str, object]:
|
||||
info: Final = _MAP.validate_python(_decoded_json(value) or {})
|
||||
if "tool_allowlist_enforced" in info and info["tool_allowlist_enforced"] is not True:
|
||||
raise ValueError("Provider registrations require mcp_info.tool_allowlist_enforced=true")
|
||||
if info.get("protocol_version", "auto") != "auto":
|
||||
raise ValueError("Provider registrations do not configure the upstream MCP protocol")
|
||||
return {**info, "tool_allowlist_enforced": True}
|
||||
|
||||
@field_validator("tool_name_to_description", "tool_name_to_display_name", mode="before")
|
||||
@classmethod
|
||||
def decode_stored_metadata(cls, value: object) -> object:
|
||||
return _decoded_json(value)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_tool_descriptions(self) -> Self:
|
||||
if not (self.tool_name_to_description or {}).keys() <= set(self.allowed_tools):
|
||||
raise ValueError("Tool descriptions must reference allowed_tools")
|
||||
return self
|
||||
|
||||
|
||||
def canonical_provider_mcp_url(value: object) -> str:
|
||||
if value is None:
|
||||
raise ValueError("Provider MCP url is required")
|
||||
url: Final = _STR.validate_python(value)
|
||||
if any(char.isspace() or ord(char) < 32 for char in url) or "\\" in url:
|
||||
raise ValueError("Provider MCP url must not contain whitespace, control characters, or backslashes")
|
||||
_URL.validate_python(url)
|
||||
parsed: Final = urlsplit(url)
|
||||
if parsed.scheme.lower() != "https" or not parsed.hostname or parsed.username is not None or "#" in url:
|
||||
raise ValueError("Provider MCP url must use HTTPS without userinfo or a fragment")
|
||||
host: Final = parsed.hostname.lower().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 normalize_provider_registration(values: Mapping[str, object]) -> Mapping[str, object]:
|
||||
registration: Final = ProviderMCPRegistration.model_validate(values)
|
||||
return {
|
||||
**registration.model_dump(mode="json", exclude_unset=True),
|
||||
"mcp_info": registration.mcp_info,
|
||||
"transport": None,
|
||||
"auth_type": None,
|
||||
"credentials": None,
|
||||
}
|
||||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -851,3 +851,204 @@ async def test_provider_mode_is_rejected_before_upstream_auth_or_connection(boun
|
|||
)
|
||||
with pytest.raises(ValidationError):
|
||||
await attempt
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_management_and_runtime_inventory_share_refresh_and_request_snapshot(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.models.mcp_server import LiteLLM_MCPServerTable
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
|
||||
|
||||
provider: Final = LiteLLM_MCPServerTable(
|
||||
server_id="provider", execution_mode="provider", transport=None,
|
||||
url="https://provider.example/mcp", allowed_tools=["search"],
|
||||
mcp_info={"tool_allowlist_enforced": True},
|
||||
)
|
||||
updated: Final = provider.model_copy(update={"allowed_tools": []})
|
||||
read_rows: Final = AsyncMock(side_effect=([provider], [updated], RuntimeError("offline"), []))
|
||||
read_revision: Final = AsyncMock(return_value=SimpleNamespace(reload_revision=1))
|
||||
tables: Final = SimpleNamespace(
|
||||
litellm_mcpservertable=SimpleNamespace(find_many=read_rows),
|
||||
litellm_config=SimpleNamespace(find_unique=read_revision),
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables))
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
manager: Final = MCPServerManager()
|
||||
async with manager.catalog.operation():
|
||||
assert (await manager.get_all_mcp_servers_unfiltered())[0].allowed_tools == ["search"]
|
||||
assert await manager.catalog.list() == {}
|
||||
read_revision.return_value = SimpleNamespace(reload_revision=2)
|
||||
assert (await manager.get_all_mcp_servers_unfiltered())[0].allowed_tools == ["search"]
|
||||
assert (await manager.get_all_mcp_servers_unfiltered())[0].allowed_tools == []
|
||||
assert await manager.catalog.list() == {}
|
||||
assert read_rows.await_count == 2
|
||||
read_revision.return_value = SimpleNamespace(reload_revision=3)
|
||||
with pytest.raises(HTTPException) as unavailable:
|
||||
await manager.get_all_mcp_servers_unfiltered()
|
||||
assert unavailable.value.status_code == 503
|
||||
assert await manager.get_all_mcp_servers_unfiltered() == []
|
||||
assert await manager.catalog.list() == {}
|
||||
assert read_rows.await_count == 4
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("pinned_id", [None, "configured-docs"])
|
||||
async def test_provider_config_inventory_reloads_without_entering_gateway(
|
||||
monkeypatch: pytest.MonkeyPatch, pinned_id: str | None,
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
manager: Final = MCPServerManager()
|
||||
config: Final = {
|
||||
"docs": {
|
||||
"execution_mode": "provider", "url": "https://DOCS.example:443/mcp", "allowed_tools": ["search"],
|
||||
"tool_name_to_description": {"search": "Search documentation"},
|
||||
**({"server_id": pinned_id} if pinned_id else {}),
|
||||
},
|
||||
}
|
||||
await manager.load_servers_from_config(config)
|
||||
first: Final = (await manager.get_all_mcp_servers_unfiltered())[0]
|
||||
assert first.is_config and first.execution_mode == "provider"
|
||||
assert first.url == "https://docs.example/mcp"
|
||||
assert first.allowed_tools == ["search"]
|
||||
assert first.tool_name_to_description == {"search": "Search documentation"}
|
||||
assert first.transport is None and first.credentials is None
|
||||
assert first.mcp_info["tool_allowlist_enforced"] is True
|
||||
assert not manager.get_registry()
|
||||
async with manager.catalog.operation():
|
||||
await manager.load_servers_from_config({"docs": {
|
||||
**config["docs"], "url": first.url, "allowed_tools": [], "tool_name_to_description": {},
|
||||
}})
|
||||
assert (await manager.get_all_mcp_servers_unfiltered())[0].allowed_tools == ["search"]
|
||||
updated: Final = (await manager.get_all_mcp_servers_unfiltered())[0]
|
||||
assert updated.server_id == first.server_id
|
||||
assert updated.allowed_tools == []
|
||||
await manager.load_servers_from_config({})
|
||||
assert await manager.get_all_mcp_servers_unfiltered() == []
|
||||
assert not manager.get_registry()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("field,value", [
|
||||
("authentication_token", "synthetic"), ("client_secret", "synthetic"),
|
||||
("aws_access_key_id", "synthetic"), ("client_private_key", "synthetic"),
|
||||
("allow_sampling", True), ("protocol_version", "2025-03-26"),
|
||||
("gateway_custom_option", {"enabled": True}),
|
||||
])
|
||||
async def test_provider_config_rejects_active_config_only_connection_fields(
|
||||
monkeypatch: pytest.MonkeyPatch, field: str, value: object,
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
manager: Final = MCPServerManager()
|
||||
config: Final = {"execution_mode": "provider", "url": "https://docs.example/mcp", "allowed_tools": ["search"]}
|
||||
await manager.load_servers_from_config({"docs": config})
|
||||
previous: Final = await manager.get_all_mcp_servers_unfiltered()
|
||||
with pytest.raises(ValueError, match=field):
|
||||
await manager.load_servers_from_config({"docs": {**config, field: value}})
|
||||
assert await manager.get_all_mcp_servers_unfiltered() == previous
|
||||
assert not manager.get_registry()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("other", [
|
||||
{"execution_mode": "provider", "url": "https://DOCS.example:443/mcp", "allowed_tools": []},
|
||||
{"execution_mode": "provider", "url": "https://other.example/mcp", "allowed_tools": [], "alias": "docs"},
|
||||
{"url": "https://gateway.example/mcp", "transport": "http", "server_id": "docs"},
|
||||
])
|
||||
async def test_provider_config_rejects_ambiguous_config_identities(
|
||||
monkeypatch: pytest.MonkeyPatch, other: dict[str, object],
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
manager: Final = MCPServerManager()
|
||||
with pytest.raises(ValueError, match="conflicts"):
|
||||
await manager.load_servers_from_config({
|
||||
"docs": {"execution_mode": "provider", "url": "https://docs.example/mcp", "allowed_tools": []},
|
||||
"other": other,
|
||||
})
|
||||
assert not manager.get_registry()
|
||||
assert await manager.get_all_mcp_servers_unfiltered() == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_config_replaces_gateway_entry_without_retaining_execution(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
manager: Final = MCPServerManager()
|
||||
await manager.load_servers_from_config({"docs": {"transport": "http", "url": "https://docs.example/mcp"}})
|
||||
assert manager.get_mcp_server_by_name("docs") is not None
|
||||
await manager.load_servers_from_config({"docs": {
|
||||
"execution_mode": "provider", "url": "https://docs.example/mcp", "allowed_tools": [],
|
||||
}})
|
||||
assert manager.get_mcp_server_by_name("docs") is None
|
||||
assert not manager.get_registry()
|
||||
assert (await manager.get_all_mcp_servers_unfiltered())[0].execution_mode == "provider"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("database_first", [False, True])
|
||||
async def test_provider_config_rejects_database_url_conflict_in_either_loading_order(
|
||||
monkeypatch: pytest.MonkeyPatch, database_first: bool,
|
||||
) -> None:
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.models.mcp_server import LiteLLM_MCPServerTable
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
|
||||
|
||||
row: Final = LiteLLM_MCPServerTable(
|
||||
server_id="database-docs", server_name="database-docs", execution_mode="provider", transport=None,
|
||||
url="https://docs.example/mcp", allowed_tools=["search"],
|
||||
)
|
||||
tables: Final = SimpleNamespace(
|
||||
litellm_mcpservertable=SimpleNamespace(find_many=AsyncMock(return_value=[row])),
|
||||
litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=SimpleNamespace(reload_revision=1))),
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables))
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
manager: Final = MCPServerManager()
|
||||
config: Final = {"docs": {"execution_mode": "provider", "url": row.url, "allowed_tools": []}}
|
||||
if database_first:
|
||||
await manager.get_all_mcp_servers_unfiltered()
|
||||
with pytest.raises(ValueError, match="conflicts"):
|
||||
await manager.load_servers_from_config(config)
|
||||
else:
|
||||
await manager.load_servers_from_config(config)
|
||||
with pytest.raises(HTTPException) as unavailable:
|
||||
await manager.get_all_mcp_servers_unfiltered()
|
||||
assert unavailable.value.status_code == 503
|
||||
assert not manager.get_registry()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_config_changes_invalidate_cached_inventory_without_database_revision_change(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
|
||||
|
||||
tables: Final = SimpleNamespace(
|
||||
litellm_mcpservertable=SimpleNamespace(find_many=AsyncMock(return_value=[])),
|
||||
litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=SimpleNamespace(reload_revision=1))),
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables))
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
manager: Final = MCPServerManager()
|
||||
initial: Final = {"execution_mode": "provider", "url": "https://docs.example/mcp", "allowed_tools": ["search"]}
|
||||
await manager.load_servers_from_config({"docs": initial})
|
||||
assert (await manager.get_all_mcp_servers_unfiltered())[0].allowed_tools == ["search"]
|
||||
await manager.load_servers_from_config({"docs": {**initial, "allowed_tools": []}})
|
||||
assert (await manager.get_all_mcp_servers_unfiltered())[0].allowed_tools == []
|
||||
await manager.load_servers_from_config({})
|
||||
assert await manager.get_all_mcp_servers_unfiltered() == []
|
||||
|
|
|
|||
|
|
@ -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"}},
|
||||
{"gateway_custom_option": {"enabled": True}},
|
||||
],
|
||||
)
|
||||
def test_provider_update_validates_the_merged_stored_registration(patch: dict[str, object]) -> None:
|
||||
from litellm.proxy.mcp_server.db import _validate_provider_registration_update
|
||||
|
||||
stored = models.LiteLLM_MCPServerTable.model_construct(
|
||||
server_id="provider", execution_mode="provider", url="https://docs.example/mcp", transport=None,
|
||||
allowed_tools=["search"], mcp_info={"tool_allowlist_enforced": True},
|
||||
)
|
||||
with pytest.raises(HTTPException) as error:
|
||||
_validate_provider_registration_update(stored, UpdateMCPServerRequest(server_id="provider", **patch))
|
||||
assert error.value.status_code == 400
|
||||
|
||||
|
||||
def test_provider_partial_update_preserves_tools_and_empty_list_stays_enforced() -> None:
|
||||
from litellm.proxy.mcp_server.db import _validate_provider_registration_update
|
||||
|
||||
stored = models.LiteLLM_MCPServerTable.model_construct(
|
||||
server_id="provider", execution_mode="provider", url="https://docs.example/mcp", transport=None,
|
||||
allowed_tools=["search"], mcp_info={"tool_allowlist_enforced": True, "description": "Docs"},
|
||||
)
|
||||
renamed = _validate_provider_registration_update(stored, UpdateMCPServerRequest(server_id="provider", description="New"))
|
||||
assert renamed.allowed_tools == ["search"]
|
||||
assert renamed.mcp_info == {"tool_allowlist_enforced": True, "description": "Docs"}
|
||||
cleared = _validate_provider_registration_update(stored, UpdateMCPServerRequest(server_id="provider", allowed_tools=[]))
|
||||
assert cleared.allowed_tools == []
|
||||
assert cleared.mcp_info["tool_allowlist_enforced"] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["description", "instructions", "source_url", "tool_name_to_description"])
|
||||
def test_provider_update_clears_explicit_null_metadata_without_writing_omitted_fields(field: str) -> None:
|
||||
from litellm.proxy.mcp_server.db import _prepare_provider_registration_update
|
||||
|
||||
stored = models.LiteLLM_MCPServerTable.model_construct(
|
||||
server_id="provider", execution_mode="provider", url="https://docs.example/mcp", transport=None,
|
||||
allowed_tools=["search"], mcp_info={"tool_allowlist_enforced": True},
|
||||
description="Docs", instructions="Search", source_url="https://docs.example",
|
||||
tool_name_to_description={"search": "Search"},
|
||||
)
|
||||
patch = UpdateMCPServerRequest(server_id="provider", **{field: None})
|
||||
serialized = _prepare_provider_registration_update(stored, patch, "admin", None)
|
||||
assert serialized[field] == ("{}" if field == "tool_name_to_description" else None)
|
||||
assert "allowed_tools" not in serialized
|
||||
assert "url" not in serialized
|
||||
assert json.loads(serialized["mcp_info"])["tool_allowlist_enforced"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_accepts_neutral_connection_settings_without_credentials_or_gateway_transport():
|
||||
neutral = {
|
||||
"transport": None, "auth_type": None, "credentials": {}, "static_headers": {},
|
||||
"extra_headers": [], "env_vars": [], "is_byok": False, "command": None,
|
||||
}
|
||||
created = NewMCPServerRequest(
|
||||
server_name="provider", execution_mode="provider", url="https://docs.example/mcp",
|
||||
allowed_tools=["search"], **neutral,
|
||||
)
|
||||
assert created.transport is None
|
||||
assert created.credentials is None
|
||||
prisma = _mock_prisma()
|
||||
stored = models.LiteLLM_MCPServerTable.model_construct(
|
||||
server_id="provider", execution_mode="provider", url=created.url, transport=None,
|
||||
allowed_tools=["search"], mcp_info={"tool_allowlist_enforced": True},
|
||||
)
|
||||
prisma.db.litellm_mcpservertable.find_unique.return_value = stored
|
||||
prisma.db.litellm_mcpservertable.update.return_value = stored
|
||||
updated = await update_mcp_server(prisma, UpdateMCPServerRequest(server_id="provider", **neutral), "admin")
|
||||
assert updated.allowed_tools == ["search"]
|
||||
written = prisma.db.litellm_mcpservertable.update.call_args.kwargs["data"]
|
||||
assert not neutral.keys() & written.keys()
|
||||
assert "allowed_tools" not in written
|
||||
assert json.loads(written["mcp_info"])["tool_allowlist_enforced"] is True
|
||||
|
|
|
|||
|
|
@ -5106,3 +5106,16 @@ async def test_rest_guardrail_bearer_excludes_admission_key(
|
|||
assert (authorization.removeprefix("Bearer ") or None) == expected
|
||||
finally:
|
||||
global_mcp_tool_registry.unregister_tools_with_prefix("guarded-")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_registration_does_not_open_gateway_connection() -> None:
|
||||
request: Final = NewMCPServerRequest(
|
||||
execution_mode="provider", url="https://mcp.example.com/mcp", allowed_tools=["read"]
|
||||
)
|
||||
|
||||
async def unexpected_operation(_client: object) -> dict[str, object]:
|
||||
pytest.fail("Provider registrations must not open a gateway client")
|
||||
|
||||
result: Final = await rest_endpoints._execute_with_mcp_client(request, unexpected_operation)
|
||||
assert result == {"error": "Provider-executed MCP registrations do not support gateway connections"}
|
||||
|
|
|
|||
|
|
@ -6593,3 +6593,34 @@ async def test_proxy_config_validates_advertised_mcp_versions_at_load(tmp_path,
|
|||
return
|
||||
with pytest.raises(ValidationError):
|
||||
await ProxyConfig().load_config(router=None, config_file_path=str(config))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_config_management_and_removal_without_database(monkeypatch):
|
||||
from starlette.requests import Request
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import fetch_all_mcp_servers, fetch_mcp_server
|
||||
from litellm.proxy.mcp_server import mcp_server_manager
|
||||
from litellm.proxy.management_endpoints import mcp_management_endpoints
|
||||
|
||||
manager = mcp_server_manager.MCPServerManager()
|
||||
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
|
||||
monkeypatch.setattr(mcp_management_endpoints, "global_mcp_server_manager", manager)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
config = proxy_server.ProxyConfig()
|
||||
await config._init_non_llm_configs({"mcp_servers": {"config_docs": {
|
||||
"server_id": "config-docs", "execution_mode": "provider", "url": "https://docs.example/mcp",
|
||||
"allowed_tools": ["search"],
|
||||
}}})
|
||||
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
inventory = await fetch_all_mcp_servers(admin, team_id=None, user_id=None)
|
||||
assert len(inventory) == 1 and inventory[0].is_config
|
||||
request = Request({"type": "http", "headers": [], "client": ("127.0.0.1", 12345)})
|
||||
detail = await fetch_mcp_server(request, "config_docs", admin, include_reachability=True)
|
||||
assert detail.server_id == "config-docs" and detail.allowed_tools == ["search"]
|
||||
assert detail.last_health_check is None
|
||||
await config._init_non_llm_configs({})
|
||||
assert await fetch_all_mcp_servers(admin, team_id=None, user_id=None) == []
|
||||
assert not manager.get_registry()
|
||||
|
|
|
|||
|
|
@ -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])
|
||||
|
|
|
|||
96
tests/unit/types/mcp_server/test_provider_registration.py
Normal file
96
tests/unit/types/mcp_server/test_provider_registration.py
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
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"},
|
||||
{"gateway_custom_option": {"enabled": True}},
|
||||
{"timeout": 0},
|
||||
{"mcp_info": {"tool_allowlist_enforced": False}},
|
||||
{"mcp_info": {"tool_allowlist_enforced": None}},
|
||||
{"tool_name_to_description": {"not_allowed": "Not in the list"}},
|
||||
],
|
||||
)
|
||||
def test_provider_registration_rejects_unbounded_tools_and_connection_settings(patch: dict[str, object]) -> None:
|
||||
with pytest.raises(ValueError, match=r"allowed_tools|valid list|Provider|Tool descriptions"):
|
||||
normalize_provider_registration({"url": "https://docs.example/mcp", "allowed_tools": ["search"], **patch})
|
||||
|
||||
|
||||
def test_empty_tool_list_is_enforced_without_transport_or_stored_credentials() -> None:
|
||||
request: Final = NewMCPServerRequest(
|
||||
execution_mode="provider", url="https://docs.example/mcp", allowed_tools=[], mcp_info={"description": "Docs"}
|
||||
)
|
||||
prepared: Final = _prepare_mcp_server_data(request)
|
||||
assert prepared["allowed_tools"] == []
|
||||
assert request.mcp_info == {"description": "Docs", "tool_allowlist_enforced": True}
|
||||
assert prepared["transport"] is None
|
||||
assert not prepared.get("credentials")
|
||||
assert prepared.get("auth_type") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("neutral", [None, False, "", {}, []])
|
||||
def test_provider_registration_ignores_neutral_unknown_settings(neutral: object) -> None:
|
||||
request: Final = NewMCPServerRequest(
|
||||
execution_mode="provider", url="https://docs.example/mcp", allowed_tools=["search"],
|
||||
gateway_custom_option=neutral,
|
||||
)
|
||||
assert request.allowed_tools == ["search"]
|
||||
assert request.transport is None
|
||||
assert "gateway_custom_option" not in request.model_dump()
|
||||
|
||||
|
||||
def test_provider_registration_preserves_supported_management_fields() -> None:
|
||||
request: Final = NewMCPServerRequest(
|
||||
execution_mode="provider", server_name="docs", alias="documentation",
|
||||
url="https://docs.example/mcp", allowed_tools=["search"], description="Documentation",
|
||||
instructions="Search docs", source_url="https://docs.example", mcp_access_groups=["research"],
|
||||
allow_all_keys=True, tool_name_to_description={"search": "Search docs"},
|
||||
tool_name_to_display_name={"search": "Search"}, mcp_info={"custom_metadata": "retained"},
|
||||
)
|
||||
assert request.server_name == "docs" and request.alias == "documentation"
|
||||
assert request.description == "Documentation" and request.instructions == "Search docs"
|
||||
assert request.source_url == "https://docs.example" and request.mcp_access_groups == ["research"]
|
||||
assert request.allow_all_keys
|
||||
assert request.tool_name_to_description == {"search": "Search docs"}
|
||||
assert request.tool_name_to_display_name == {"search": "Search"}
|
||||
assert request.mcp_info == {"custom_metadata": "retained", "tool_allowlist_enforced": True}
|
||||
14
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
14
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -39540,12 +39540,8 @@ export interface components {
|
|||
};
|
||||
/** Env Vars */
|
||||
env_vars?: components["schemas"]["MCPEnvVar"][] | null;
|
||||
/**
|
||||
* Execution Mode
|
||||
* @default litellm
|
||||
* @constant
|
||||
*/
|
||||
execution_mode: "litellm";
|
||||
/** @default litellm */
|
||||
execution_mode: components["schemas"]["MCPExecutionMode"];
|
||||
/** Extra Headers */
|
||||
extra_headers?: string[] | null;
|
||||
/** Instructions */
|
||||
|
|
@ -39624,9 +39620,8 @@ export interface components {
|
|||
/**
|
||||
* Transport
|
||||
* @default sse
|
||||
* @enum {string}
|
||||
*/
|
||||
transport: "sse" | "http" | "stdio";
|
||||
transport: ("sse" | "http" | "stdio") | null;
|
||||
/** Url */
|
||||
url?: string | null;
|
||||
};
|
||||
|
|
@ -49300,9 +49295,8 @@ export interface components {
|
|||
/**
|
||||
* Transport
|
||||
* @default sse
|
||||
* @enum {string}
|
||||
*/
|
||||
transport: "sse" | "http" | "stdio";
|
||||
transport: ("sse" | "http" | "stdio") | null;
|
||||
/** Url */
|
||||
url?: string | null;
|
||||
};
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue