scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

feat(mcp): store execution mode and isolate provider registrations

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mrinal 2026-11-03 16:18:13 +00:00 • committed by Mrinal Chanshetty
parent 44b85d0c5c
commit 436181f313
17 changed files with 293 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_add_mcp_execution_mode/migration.sql Normal file
View file

@ -0,0 +1,9 @@
CREATE TYPE "MCPExecutionMode" AS ENUM ('litellm', 'provider');
ALTER TABLE "LiteLLM_MCPServerTable"
ADD COLUMN "execution_mode" "MCPExecutionMode" NOT NULL DEFAULT 'litellm',
ALTER COLUMN "transport" DROP NOT NULL;
CREATE UNIQUE INDEX "LiteLLM_MCPServerTable_provider_url_key"
ON "LiteLLM_MCPServerTable" ("url")
WHERE "execution_mode" = 'provider';

8
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -358,6 +358,11 @@ model LiteLLM_ObjectPermissionTable {
agents_table LiteLLM_AgentsTable[]
}
enum MCPExecutionMode {
litellm
provider
}
// Holds the MCP server configuration
model LiteLLM_MCPServerTable {
server_id String @id @default(uuid())
@ -367,7 +372,8 @@ model LiteLLM_MCPServerTable {
instructions String?
url String?
spec_path String?
transport String @default("sse")
execution_mode MCPExecutionMode @default(litellm)
transport String? @default("sse")
auth_type String?
credentials Json? @default("{}")
created_at DateTime? @default(now()) @map("created_at")

16
litellm/models/mcp_server.py
View file

@ -11,10 +11,11 @@ from datetime import datetime
from types import MappingProxyType
from typing import Literal
from pydantic import Field, ValidationInfo, field_validator
from pydantic import Field, ValidationInfo, field_validator, model_validator
from typing_extensions import Self
from litellm.types.llms.base import LiteLLMPydanticObjectBase
from litellm.types.mcp import MCPAuthType, MCPCredentials, MCPTransportType
from litellm.types.mcp import MCPAuthType, MCPCredentials, MCPExecutionMode, MCPTransportType
from litellm.types.mcp_server.mcp_server_manager import MCPInfo, PinnedMCPTool, parse_pinned_tools
@ -56,7 +57,8 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
description: str | None = None
url: str | None = None
spec_path: str | None = None
transport: MCPTransportType
execution_mode: MCPExecutionMode = MCPExecutionMode.litellm
transport: MCPTransportType | None
auth_type: MCPAuthType | None = None
credentials: MCPCredentials | None = None
instructions: str | None = None
@ -121,6 +123,14 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
reviewed_at: datetime | None = None
review_notes: str | None = None
@model_validator(mode="after")
def validate_execution_transport(self) -> Self:
if self.execution_mode == MCPExecutionMode.provider and self.transport is not None:
raise ValueError("Provider-executed MCP registrations require transport=null")
if self.execution_mode == MCPExecutionMode.litellm and self.transport is None:
raise ValueError("LiteLLM-executed MCP registrations require a transport")
return self
@field_validator("pinned_tools", mode="before")
@classmethod
def decode_stored_pinned_tools(cls, value: object) -> dict[str, PinnedMCPTool] | None:

86
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -37505,6 +37505,12 @@
],
"title": "Env Vars"
},
"execution_mode": {
"const": "litellm",
"default": "litellm",
"title": "Execution Mode",
"type": "string"
},
"extra_headers": {
"anyOf": [
{
@ -39333,6 +39339,10 @@
],
"title": "Env Vars"
},
"execution_mode": {
"$ref": "#/components/schemas/MCPExecutionMode",
"default": "litellm"
},
"extra_headers": {
"items": {
"type": "string"
@ -39715,13 +39725,20 @@
"title": "Tool Name To Display Name"
},
"transport": {
"enum": [
"sse",
"http",
"stdio"
"anyOf": [
{
"enum": [
"sse",
"http",
"stdio"
],
"type": "string"
},
{
"type": "null"
}
],
"title": "Transport",
"type": "string"
"title": "Transport"
},
"updated_at": {
"anyOf": [
@ -40143,6 +40160,14 @@
"title": "MCPEnvVarScope",
"type": "string"
},
"MCPExecutionMode": {
"enum": [
"litellm",
"provider"
],
"title": "MCPExecutionMode",
"type": "string"
},
"NewMCPServerRequest": {
"properties": {
"alias": {
@ -40333,6 +40358,12 @@
],
"title": "Env Vars"
},
"execution_mode": {
"const": "litellm",
"default": "litellm",
"title": "Execution Mode",
"type": "string"
},
"extra_headers": {
"anyOf": [
{
@ -42532,6 +42563,10 @@
],
"title": "Env Vars"
},
"execution_mode": {
"$ref": "#/components/schemas/MCPExecutionMode",
"default": "litellm"
},
"extra_headers": {
"items": {
"type": "string"
@ -42914,13 +42949,20 @@
"title": "Tool Name To Display Name"
},
"transport": {
"enum": [
"sse",
"http",
"stdio"
"anyOf": [
{
"enum": [
"sse",
"http",
"stdio"
],
"type": "string"
},
{
"type": "null"
}
],
"title": "Transport",
"type": "string"
"title": "Transport"
},
"updated_at": {
"anyOf": [
@ -43559,6 +43601,14 @@
"title": "MCPEnvVarScope",
"type": "string"
},
"MCPExecutionMode": {
"enum": [
"litellm",
"provider"
],
"title": "MCPExecutionMode",
"type": "string"
},
"MCPGatewaySession": {
"description": "One live stateful Streamable HTTP session held by this proxy worker.",
"properties": {
@ -44376,6 +44426,12 @@
],
"title": "Env Vars"
},
"execution_mode": {
"const": "litellm",
"default": "litellm",
"title": "Execution Mode",
"type": "string"
},
"extra_headers": {
"anyOf": [
{
@ -47840,6 +47896,12 @@
],
"title": "Env Vars"
},
"execution_mode": {
"const": "litellm",
"default": "litellm",
"title": "Execution Mode",
"type": "string"
},
"extra_headers": {
"anyOf": [
{

10
litellm/proxy/_types.py
View file

@ -64,6 +64,8 @@ from litellm.types.mcp import (
MCPAuthType,
MCPCredentials,
MCPDisabledCapabilities,
MCPExecutionMode,
MCPGatewayExecutionMode,
MCPTransport,
MCPTransportType,
MCPUpstreamProtocol,
@ -1386,6 +1388,7 @@ def _validate_mcp_transport_fields(values: object) -> None:
class NewMCPServerRequest(LiteLLMPydanticObjectBase):
execution_mode: MCPGatewayExecutionMode = MCPExecutionMode.litellm
server_id: str | None = None
server_name: str | None = None
alias: str | None = None
@ -1543,6 +1546,13 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
max_concurrent_requests: int | None = None
rpm: int | None = Field(default=None, ge=0)
@model_validator(mode="before")
@classmethod
def reject_execution_mode_changes(cls, values: object) -> object:
if isinstance(values, dict) and "execution_mode" in values:
raise ValueError("execution_mode is immutable; omit it when updating an MCP registration")
return values
@model_validator(mode="after")
def validate_protocol_transport(self) -> "UpdateMCPServerRequest":
if not {"transport", "mcp_info"}.issubset(self.model_fields_set):

2
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -3181,7 +3181,7 @@ if MCP_AVAILABLE:
raise HTTPException(
status_code=503, detail="Cannot validate MCP configuration while stored state is unavailable"
)
if stored is None:
if stored is None or stored.transport is None:
return
effective_transport: Final = payload.transport if "transport" in fields_set else stored.transport
effective_info: Final = payload.mcp_info if "mcp_info" in fields_set else stored.mcp_info

3
litellm/proxy/mcp_server/catalog.py
View file

@ -24,6 +24,7 @@ from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.mcp_server.result_conversion import age_freshness, aggregate_freshness
from litellm.types.mcp import MCPExecutionMode
if TYPE_CHECKING:
from mcp.types import ListToolsResult, PaginatedRequestParams, PaginatedResult
@ -415,6 +416,8 @@ class CatalogSnapshots:
for row in rows:
try:
server = LiteLLM_MCPServerTable.model_validate(row.model_dump())
if server.execution_mode == MCPExecutionMode.provider:
continue
existing_server = previous_registry.get(server.server_id)
if (

16
litellm/proxy/mcp_server/mcp_server_manager.py
View file

@ -230,8 +230,10 @@ from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.mcp import (
DEFAULT_SUBJECT_TOKEN_TYPE,
MCPAuth,
MCPExecutionMode,
MCPTokenEndpointAuthMethod,
MCPUpstreamProtocol,
validate_mcp_gateway_execution,
without_header,
)
from litellm.types.mcp_server.mcp_server_manager import (
@ -405,6 +407,7 @@ class MCPServerConfig(TypedDict, total=False):
whatever the admin wrote, and each read applies its own default."""
server_id: ReadOnly[str]
execution_mode: ReadOnly[MCPExecutionMode]
protocol_version: ReadOnly[MCPUpstreamProtocol]
alias: str
description: str
@ -2583,6 +2586,7 @@ class MCPServerManager:
for server_name, raw_server_config in mcp_servers_config.items():
server_config: MCPServerConfig = raw_server_config
validate_mcp_gateway_execution(server_config.get("execution_mode", MCPExecutionMode.litellm))
_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()
@ -3219,6 +3223,7 @@ class MCPServerManager:
env_vars_are_encrypted: bool | None = None,
register_oauth_discovery: bool = True,
) -> MCPServer:
validate_mcp_gateway_execution(mcp_server.execution_mode)
_mcp_info: Final[MCPInfo] = mcp_server.mcp_info or {}
env_dict: Final = _deserialize_json_dict(getattr(mcp_server, "env", None))
static_headers_dict: Final = _deserialize_json_dict(getattr(mcp_server, "static_headers", None))
@ -3485,6 +3490,10 @@ class MCPServerManager:
self.initialize_tool_name_to_mcp_server_name_mapping()
async def add_server(self, mcp_server: LiteLLM_MCPServerTable):
if mcp_server.execution_mode == MCPExecutionMode.provider:
if mcp_server.server_id in self.registry:
self.remove_server(mcp_server)
return
# The runtime registry is the allowlist for tool calls and health
# probes (which spawn the underlying transport, including stdio
# subprocesses). Match the eligibility set used by the bulk DB
@ -3518,7 +3527,11 @@ class MCPServerManager:
# If a previously-active server has been moved out of the active
# state, evict any stale registry entry so subsequent tool calls and
# health probes can't reach it.
if mcp_server.approval_status not in (None, "active", "approved"):
if mcp_server.execution_mode == MCPExecutionMode.provider or mcp_server.approval_status not in (
None,
"active",
"approved",
):
evicted = self.registry.pop(mcp_server.server_id, None)
if evicted is None and mcp_server.server_name:
evicted = self.registry.pop(mcp_server.server_name, None)
@ -4388,6 +4401,7 @@ class MCPServerManager:
protocol_version_override: MCPUpstreamProtocol | None = None,
for_tool_listing: bool = False,
) -> MCPClient:
validate_mcp_gateway_execution(server.execution_mode)
self.catalog.assert_current(server)
record_auth_resolution(server.server_id, AuthResolution.unresolved)
resolved_server: Final = await self.ensure_oauth_metadata_discovered(server)

2
litellm/proxy/mcp_server/upstream.py
View file

@ -40,6 +40,7 @@ from litellm.types.mcp import (
MCPTransportType,
MCPUpstreamProtocol,
has_header,
validate_mcp_gateway_execution,
without_header,
)
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -210,6 +211,7 @@ async def prepare_upstream_client(
upstream_transport: httpx2.AsyncBaseTransport | None = None,
for_tool_listing: bool = False,
) -> MCPClient:
validate_mcp_gateway_execution(server.execution_mode)
transport: Final[MCPTransportType] = server.transport or MCPTransport.sse
server_spec: Final = None if transport == MCPTransport.stdio else to_server_spec_fail_closed(server)
spec: Final = (

8
litellm/proxy/schema.prisma
View file

@ -358,6 +358,11 @@ model LiteLLM_ObjectPermissionTable {
agents_table LiteLLM_AgentsTable[]
}
enum MCPExecutionMode {
litellm
provider
}
// Holds the MCP server configuration
model LiteLLM_MCPServerTable {
server_id String @id @default(uuid())
@ -367,7 +372,8 @@ model LiteLLM_MCPServerTable {
instructions String?
url String?
spec_path String?
transport String @default("sse")
execution_mode MCPExecutionMode @default(litellm)
transport String? @default("sse")
auth_type String?
credentials Json? @default("{}")
created_at DateTime? @default(now()) @map("created_at")

15
litellm/types/mcp.py
View file

@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, TypeAlias
from urllib.parse import urlsplit
import httpx
from pydantic import ConfigDict, Field
from pydantic import ConfigDict, Field, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm.types.llms.base import HiddenParams, LiteLLMBaseModel
@ -24,6 +24,19 @@ else:
MCPTextContent = Any
class MCPExecutionMode(str, enum.Enum):
litellm = "litellm"
provider = "provider"
MCPGatewayExecutionMode: TypeAlias = Literal[MCPExecutionMode.litellm]
_GATEWAY_EXECUTION_MODE: Final = TypeAdapter[MCPGatewayExecutionMode](MCPGatewayExecutionMode)
def validate_mcp_gateway_execution(value: object) -> MCPGatewayExecutionMode:
return _GATEWAY_EXECUTION_MODE.validate_python(value)
class MCPTransport(str, enum.Enum):
sse = "sse"
http = "http"

3
litellm/types/mcp_server/mcp_server_manager.py
View file

@ -11,6 +11,8 @@ from litellm.types.mcp import (
DEFAULT_SUBJECT_TOKEN_TYPE,
MCPAuth,
MCPAuthType,
MCPExecutionMode,
MCPGatewayExecutionMode,
MCPTokenEndpointAuthMethod,
MCPTransportType,
MCPUpstreamProtocol,
@ -96,6 +98,7 @@ def parse_pinned_tools(value: object) -> dict[str, PinnedMCPTool] | None:
class MCPServer(LiteLLMBaseModel):
server_id: str
name: str
execution_mode: MCPGatewayExecutionMode = MCPExecutionMode.litellm
alias: str | None = None
server_name: str | None = None
url: str | None = None

8
schema.prisma
View file

@ -358,6 +358,11 @@ model LiteLLM_ObjectPermissionTable {
agents_table LiteLLM_AgentsTable[]
}
enum MCPExecutionMode {
litellm
provider
}
// Holds the MCP server configuration
model LiteLLM_MCPServerTable {
server_id String @id @default(uuid())
@ -367,7 +372,8 @@ model LiteLLM_MCPServerTable {
instructions String?
url String?
spec_path String?
transport String @default("sse")
execution_mode MCPExecutionMode @default(litellm)
transport String? @default("sse")
auth_type String?
credentials Json? @default("{}")
created_at DateTime? @default(now()) @map("created_at")

20
tests/unit/models/test_models.py
View file

@ -546,6 +546,26 @@ class TestMCPServerTable:
with pytest.raises(ValidationError):
LiteLLM_MCPServerTable(server_id="s1")
def test_provider_registration_round_trips_without_connection_fields(self) -> None:
record: Final = LiteLLM_MCPServerTable(
server_id="provider-docs",
execution_mode="provider",
transport=None,
credentials=None,
url="https://docs.example/mcp",
allowed_tools=["search_docs"],
)
restored: Final = LiteLLM_MCPServerTable.model_validate_json(record.model_dump_json())
assert restored.execution_mode == "provider"
assert restored.transport is None
assert restored.credentials is None
assert restored.allowed_tools == ["search_docs"]
@pytest.mark.parametrize("mode,transport", [("provider", "http"), ("litellm", None)])
def test_execution_mode_requires_the_matching_transport_shape(self, mode: str, transport: str | None) -> None:
with pytest.raises(ValidationError, match="transport"):
LiteLLM_MCPServerTable(server_id="s1", execution_mode=mode, transport=transport)
class TestSpendLogs:
def test_spend_logs_creation(self):

79
tests/unit/proxy/mcp_server/test_catalog.py
View file

@ -772,3 +772,82 @@ def test_partial_optional_catalog_never_advertises_freshness() -> None:
assert result.ttl_ms == 0
assert result.cache_scope == "private"
@pytest.mark.asyncio
async def test_provider_records_do_not_enter_or_survive_in_the_gateway_catalog() -> None:
from datetime import datetime, timezone
from litellm.models.mcp_server import LiteLLM_MCPServerTable
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
from litellm.types.mcp_server.mcp_server_manager import MCPServer
manager: Final = MCPServerManager()
updated_at: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
record: Final = LiteLLM_MCPServerTable(
server_id="docs", execution_mode="provider", transport=None, url="https://docs.example/mcp", updated_at=updated_at
)
await manager.add_server(record)
assert manager.get_mcp_server_by_id("docs") is None
legacy: Final = MCPServer(server_id="docs", name="docs", transport="http", url=record.url, updated_at=updated_at)
manager.registry[legacy.server_id] = legacy
assert await manager.catalog._stage_servers((record,), reuse_unchanged=True) == {}
await manager.update_server(record)
assert manager.get_mcp_server_by_id("docs") is None
manager.registry[legacy.server_id] = legacy
await manager.add_server(record)
assert manager.get_mcp_server_by_id("docs") is None
@pytest.mark.asyncio
async def test_provider_table_cannot_be_converted_to_a_gateway_client() -> None:
from pydantic import ValidationError
from litellm.models.mcp_server import LiteLLM_MCPServerTable
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
manager: Final = MCPServerManager()
record: Final = LiteLLM_MCPServerTable(
server_id="docs", execution_mode="provider", transport=None, url="https://docs.example/mcp"
)
with pytest.raises(ValidationError):
await manager.build_mcp_server_from_table(record)
@pytest.mark.asyncio
async def test_config_cannot_silently_load_provider_mode_as_a_gateway_server() -> None:
from pydantic import ValidationError
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
from litellm.types.mcp import MCPAuth, MCPExecutionMode
manager: Final = MCPServerManager()
with pytest.raises(ValidationError):
await manager.load_servers_from_config(
{"docs": {"execution_mode": MCPExecutionMode.provider, "url": "https://docs.example/mcp", "auth_type": MCPAuth.oauth2}}
)
assert manager.get_registry() == {}
@pytest.mark.asyncio
@pytest.mark.parametrize("boundary", ["manager", "upstream"])
async def test_provider_mode_is_rejected_before_upstream_auth_or_connection(boundary: str) -> None:
from pydantic import ValidationError
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager
from litellm.proxy.mcp_server.upstream import prepare_upstream_client
from litellm.types.mcp import MCPExecutionMode
from litellm.types.mcp_server.mcp_server_manager import MCPServer
manager: Final = MCPServerManager()
server: Final = MCPServer(server_id="docs", name="docs", transport="http", auth_type="oauth2").model_copy(
update={"execution_mode": MCPExecutionMode.provider, "url": "https://docs.example/mcp"}
)
attempt: Final = (
manager.create_mcp_client(server)
if boundary == "manager"
else prepare_upstream_client(server, provider=manager._cred_provider, root_path="", protocol_version="auto")
)
with pytest.raises(ValidationError):
await attempt

14
tests/unit/proxy/test__types.py
View file

@ -50,6 +50,20 @@ 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")
@pytest.mark.parametrize("mode", ["litellm", "provider", None])
def test_update_requests_cannot_replace_a_registrations_execution_mode(mode: str | None) -> None:
with pytest.raises(ValidationError, match="execution_mode is immutable"):
UpdateMCPServerRequest(server_id="docs", execution_mode=mode)
def _server_only_marker_value(marker: str, value: str) -> str | tuple[str, ...]:
return (value,) if marker == "mcp_session_resource_server_ids" else value

20
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -35259,6 +35259,8 @@ export interface components {
};
/** Env Vars */
env_vars?: components["schemas"]["MCPEnvVar"][] | null;
/** @default litellm */
execution_mode: components["schemas"]["MCPExecutionMode"];
/** Extra Headers */
extra_headers?: string[];
/** Has User Credential */
@ -35358,11 +35360,8 @@ export interface components {
tool_name_to_display_name?: {
[key: string]: string;
} | null;
/**
* Transport
* @enum {string}
*/
transport: "sse" | "http" | "stdio";
/** Transport */
transport: ("sse" | "http" | "stdio") | null;
/** Updated At */
updated_at?: string | null;
/** Updated By */
@ -38300,6 +38299,11 @@ export interface components {
* @enum {string}
*/
MCPEnvVarScope: "global" | "user";
/**
* MCPExecutionMode
* @enum {string}
*/
MCPExecutionMode: "litellm" | "provider";
/**
* MCPGatewaySession
* @description One live stateful Streamable HTTP session held by this proxy worker.
@ -39536,6 +39540,12 @@ export interface components {
};
/** Env Vars */
env_vars?: components["schemas"]["MCPEnvVar"][] | null;
/**
* Execution Mode
* @default litellm
* @constant
*/
execution_mode: "litellm";
/** Extra Headers */
extra_headers?: string[] | null;
/** Instructions */