diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_add_mcp_execution_mode/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_add_mcp_execution_mode/migration.sql new file mode 100644 index 00000000000..c812c53237b --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_add_mcp_execution_mode/migration.sql @@ -0,0 +1,16 @@ +SET LOCAL lock_timeout = '10s'; + +DO $$ +BEGIN + CREATE TYPE "MCPExecutionMode" AS ENUM ('litellm', 'provider'); +EXCEPTION + WHEN duplicate_object THEN NULL; +END $$; + +ALTER TABLE "LiteLLM_MCPServerTable" +ADD COLUMN IF NOT EXISTS "execution_mode" "MCPExecutionMode" NOT NULL DEFAULT 'litellm', +ALTER COLUMN "transport" DROP NOT NULL; + +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MCPServerTable_provider_url_key" +ON "LiteLLM_MCPServerTable" ("url") +WHERE "execution_mode" = 'provider'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index c007b047087..58d8fb3623c 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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") diff --git a/litellm/models/mcp_server.py b/litellm/models/mcp_server.py index cdbda7b1b70..ef00c4743bf 100644 --- a/litellm/models/mcp_server.py +++ b/litellm/models/mcp_server.py @@ -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: diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index ab85b6f8842..42bcf3eb61c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -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": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9d84770b288..88b06ee1a97 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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): diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f6b836bc4be..a99cc47ecaf 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -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 diff --git a/litellm/proxy/mcp_server/catalog.py b/litellm/proxy/mcp_server/catalog.py index d2ff322ba5c..f8899fb948d 100644 --- a/litellm/proxy/mcp_server/catalog.py +++ b/litellm/proxy/mcp_server/catalog.py @@ -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 ( diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index 013336e1d5d..25c32952f7d 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -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) diff --git a/litellm/proxy/mcp_server/upstream.py b/litellm/proxy/mcp_server/upstream.py index 3dd245947db..d07f4955a3a 100644 --- a/litellm/proxy/mcp_server/upstream.py +++ b/litellm/proxy/mcp_server/upstream.py @@ -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 = ( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index c007b047087..58d8fb3623c 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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") diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 3fa9adef913..3556198723e 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -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" diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 1d3d5911eb1..6080bbcabec 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -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 diff --git a/schema.prisma b/schema.prisma index c007b047087..58d8fb3623c 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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") diff --git a/tests/unit/models/test_models.py b/tests/unit/models/test_models.py index e2c0edf6d9d..8bb44cb620f 100644 --- a/tests/unit/models/test_models.py +++ b/tests/unit/models/test_models.py @@ -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): diff --git a/tests/unit/proxy/mcp_server/test_catalog.py b/tests/unit/proxy/mcp_server/test_catalog.py index 45c648b81d2..384a412c4f1 100644 --- a/tests/unit/proxy/mcp_server/test_catalog.py +++ b/tests/unit/proxy/mcp_server/test_catalog.py @@ -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 diff --git a/tests/unit/proxy/mcp_server/test_db_credentials.py b/tests/unit/proxy/mcp_server/test_db_credentials.py index db2572ef954..5be76def8d8 100644 --- a/tests/unit/proxy/mcp_server/test_db_credentials.py +++ b/tests/unit/proxy/mcp_server/test_db_credentials.py @@ -486,6 +486,7 @@ def map_algorithm(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatc def _prisma_map_row(data: dict[str, object], quoted: bool = False) -> PrismaMCPServer: return PrismaMCPServer.model_validate({ + "execution_mode": "litellm", "transport": "http", "mcp_access_groups": [], "allowed_tools": [], "extra_headers": [], "args": [], "allow_all_keys": False, "available_on_public_internet": True, "delegate_auth_to_upstream": False, "oauth_passthrough": False, "per_server_oauth_discovery": False, "is_byok": False, "byok_description": [], diff --git a/tests/unit/proxy/mcp_server/test_mcp_env_vars.py b/tests/unit/proxy/mcp_server/test_mcp_env_vars.py index 57ac6c0be2d..87bd388124b 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_env_vars.py +++ b/tests/unit/proxy/mcp_server/test_mcp_env_vars.py @@ -1551,6 +1551,7 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri from prisma.models import LiteLLM_MCPServerTable return LiteLLM_MCPServerTable.model_validate({ + "execution_mode": "litellm", "server_id": "srv-returned", "transport": "http", "mcp_access_groups": [], "allowed_tools": [], "extra_headers": [], "args": [], "allow_all_keys": False, "available_on_public_internet": True, "delegate_auth_to_upstream": False, "oauth_passthrough": False, "per_server_oauth_discovery": False, diff --git a/tests/unit/proxy/mcp_server/test_mcp_server.py b/tests/unit/proxy/mcp_server/test_mcp_server.py index 456b1994e2d..55959d04646 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server.py @@ -1560,6 +1560,7 @@ async def test_add_update_server_with_alias(): # Create a mock LiteLLM_MCPServerTable with alias mock_mcp_server = MagicMock() + mock_mcp_server.execution_mode = "litellm" mock_mcp_server.server_id = "test-server-123" mock_mcp_server.alias = "my_alias" mock_mcp_server.server_name = "Test Server" @@ -1624,6 +1625,7 @@ async def test_add_update_server_without_alias(): # Create a mock LiteLLM_MCPServerTable without alias mock_mcp_server = MagicMock() + mock_mcp_server.execution_mode = "litellm" mock_mcp_server.server_id = "test-server-123" mock_mcp_server.alias = None mock_mcp_server.server_name = "Test Server" @@ -1688,6 +1690,7 @@ async def test_add_update_server_fallback_to_server_id(): # Create a mock LiteLLM_MCPServerTable without alias or server_name mock_mcp_server = MagicMock() + mock_mcp_server.execution_mode = "litellm" mock_mcp_server.server_id = "test-server-123" mock_mcp_server.alias = None mock_mcp_server.server_name = None diff --git a/tests/unit/proxy/mcp_server/test_mcp_sigv4_auth.py b/tests/unit/proxy/mcp_server/test_mcp_sigv4_auth.py index 5350e8a9f41..28f96c3bdd2 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_sigv4_auth.py +++ b/tests/unit/proxy/mcp_server/test_mcp_sigv4_auth.py @@ -829,6 +829,7 @@ class TestSigV4BuildFromTable: table_record = MagicMock() table_record.server_id = "test-sigv4-server" + table_record.execution_mode = "litellm" table_record.server_name = "sigv4_server" table_record.alias = None table_record.description = None @@ -899,6 +900,7 @@ class TestSigV4BuildFromTable: table_record = MagicMock() table_record.server_id = "test-bearer-server" + table_record.execution_mode = "litellm" table_record.server_name = "bearer_server" table_record.alias = None table_record.description = None diff --git a/tests/unit/proxy/test__types.py b/tests/unit/proxy/test__types.py index 40e66f1da65..7d57f246b1e 100644 --- a/tests/unit/proxy/test__types.py +++ b/tests/unit/proxy/test__types.py @@ -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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 15d440c0ad4..111ada46919 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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 */