From 436181f313e012b87b42d022e7acf3edc7b0b4be Mon Sep 17 00:00:00 2001 From: mrinal Date: Tue, 3 Nov 2026 16:18:13 +0000 Subject: [PATCH 1/5] feat(mcp): store execution mode and isolate provider registrations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 9 ++ .../litellm_proxy_extras/schema.prisma | 8 +- litellm/models/mcp_server.py | 16 +++- litellm/proxy/_lazy_openapi_snapshot.json | 86 ++++++++++++++++--- litellm/proxy/_types.py | 10 +++ .../mcp_management_endpoints.py | 2 +- litellm/proxy/mcp_server/catalog.py | 3 + .../proxy/mcp_server/mcp_server_manager.py | 16 +++- litellm/proxy/mcp_server/upstream.py | 2 + litellm/proxy/schema.prisma | 8 +- litellm/types/mcp.py | 15 +++- .../types/mcp_server/mcp_server_manager.py | 3 + schema.prisma | 8 +- tests/unit/models/test_models.py | 20 +++++ tests/unit/proxy/mcp_server/test_catalog.py | 79 +++++++++++++++++ tests/unit/proxy/test__types.py | 14 +++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 20 +++-- 17 files changed, 293 insertions(+), 26 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_add_mcp_execution_mode/migration.sql 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..d2251cc1cf9 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_add_mcp_execution_mode/migration.sql @@ -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'; 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/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 */ From b6679a2a81970fe13144726218aba47e8643bd48 Mon Sep 17 00:00:00 2001 From: mrinal Date: Tue, 3 Nov 2026 16:28:38 +0000 Subject: [PATCH 2/5] fix(mcp): make execution-mode migration safe to retry Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) 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 index d2251cc1cf9..348e5d9cb63 100644 --- 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 @@ -1,9 +1,14 @@ -CREATE TYPE "MCPExecutionMode" AS ENUM ('litellm', 'provider'); +DO $$ +BEGIN + CREATE TYPE "MCPExecutionMode" AS ENUM ('litellm', 'provider'); +EXCEPTION + WHEN duplicate_object THEN NULL; +END $$; ALTER TABLE "LiteLLM_MCPServerTable" -ADD COLUMN "execution_mode" "MCPExecutionMode" NOT NULL DEFAULT 'litellm', +ADD COLUMN IF NOT EXISTS "execution_mode" "MCPExecutionMode" NOT NULL DEFAULT 'litellm', ALTER COLUMN "transport" DROP NOT NULL; -CREATE UNIQUE INDEX "LiteLLM_MCPServerTable_provider_url_key" +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MCPServerTable_provider_url_key" ON "LiteLLM_MCPServerTable" ("url") WHERE "execution_mode" = 'provider'; From d42ed1e5ce5f10b6bc16ed3c0b52872e1154bded Mon Sep 17 00:00:00 2001 From: mrinal Date: Tue, 3 Nov 2026 16:56:42 +0000 Subject: [PATCH 3/5] fix(mcp): align gateway fixtures and migration with CI contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../20261031000000_add_mcp_execution_mode/migration.sql | 2 ++ tests/unit/proxy/mcp_server/test_db_credentials.py | 1 + tests/unit/proxy/mcp_server/test_mcp_env_vars.py | 1 + tests/unit/proxy/mcp_server/test_mcp_server.py | 3 +++ tests/unit/proxy/mcp_server/test_mcp_sigv4_auth.py | 2 ++ 5 files changed, 9 insertions(+) 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 index 348e5d9cb63..c812c53237b 100644 --- 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 @@ -1,3 +1,5 @@ +SET LOCAL lock_timeout = '10s'; + DO $$ BEGIN CREATE TYPE "MCPExecutionMode" AS ENUM ('litellm', 'provider'); 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 From 780db6a4c8d9902af5e512d24563b5da8a5925ce Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Tue, 3 Nov 2026 16:57:59 -0800 Subject: [PATCH 4/5] fix(mcp): create the provider URL unique index after every database setup The partial unique index on provider MCP URLs exists only in migration SQL. prisma db push never creates it, and adopting a pushed database marks migrations applied without running their SQL, so those setups accepted duplicate provider URLs. Database setup now creates the index when it is missing, under the migration lock, and never fails startup. --- .../litellm_proxy_extras/utils.py | 63 +++++++++++++++ .../migrations/test_provider_url_index.py | 78 +++++++++++++++++++ 2 files changed, 141 insertions(+) create mode 100644 tests/integration/database/migrations/test_provider_url_index.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 8b19576d86e..c51fc45f89e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -67,6 +67,8 @@ _MIGRATION_TS_RE = re.compile(r"^(\d{14})_") _MIGRATION_DEADLOCK_MARKER = "deadlock detected" _MIGRATION_LOCK_TIMEOUT_MARKER = "canceling statement due to lock timeout" INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big") +MCP_SERVER_TABLE: Final = "LiteLLM_MCPServerTable" +PROVIDER_URL_INDEX: Final = "LiteLLM_MCPServerTable_provider_url_key" _TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$") _INVALID_LITELLM_INDEXES_SQL: Final = ( "SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) " @@ -940,6 +942,66 @@ class ProxyExtrasDBManager: return False return not remaining + @staticmethod + def ensure_provider_url_index(lock_timeout: str = "30s") -> bool: + """Create the partial unique index on provider MCP URLs when it is missing. + + The index exists only in migration SQL, because Prisma's schema cannot express a + partial index. `prisma db push` never creates it, and adopting a pushed database + marks migrations applied without running their SQL, so every database setup path + ends here. Never raises: returns True when the index exists, False when creation was + skipped or failed and will be retried on the next database setup run.""" + prisma_url: Final = os.getenv("DATABASE_URL") + if not prisma_url: + return False + try: + import psycopg + from psycopg import sql + except ImportError: + logger.warning( + "psycopg is not installed; skipping the provider MCP URL index check. " + "Install the litellm[extra_proxy] extra, which includes psycopg." + ) + return False + + schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public" + cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url) + index_name: Final = sql.Identifier(schema, PROVIDER_URL_INDEX) + try: + with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn: + present: Final = conn.execute("SELECT to_regclass(%s)", (index_name.as_string(conn),)).fetchone() + if present is not None and present[0] is not None: + return True + column: Final = conn.execute( + "SELECT 1 FROM information_schema.columns " + "WHERE table_schema = %s AND table_name = %s AND column_name = 'execution_mode'", + (schema, MCP_SERVER_TABLE), + ).fetchone() + if column is None: + return False + conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout))) + with held_migration_lock(conn) as held: + if not held: + logger.info( + "Another process holds the migration lock; leaving the provider MCP URL index " + "to the next database setup run" + ) + return False + conn.execute( + sql.SQL( + "CREATE UNIQUE INDEX IF NOT EXISTS {} ON {} (\"url\") WHERE \"execution_mode\" = 'provider'" + ).format(sql.Identifier(PROVIDER_URL_INDEX), sql.Identifier(schema, MCP_SERVER_TABLE)) + ) + except psycopg.Error as e: + logger.warning( + "Could not create the provider MCP URL unique index, will retry on the next database setup run. " + "Error: %s", + e, + ) + return False + logger.info("Created the provider MCP URL unique index that prisma db push and schema adoption skip") + return True + @staticmethod def _repair_under_migration_lock( conn: "psycopg.Connection[tuple[str, str, str]]", @@ -1422,6 +1484,7 @@ class ProxyExtrasDBManager: if not migrated: return False ProxyExtrasDBManager.repair_invalid_indexes() + ProxyExtrasDBManager.ensure_provider_url_index() ProxyExtrasDBManager.apply_replica_identity_full_if_requested() return True diff --git a/tests/integration/database/migrations/test_provider_url_index.py b/tests/integration/database/migrations/test_provider_url_index.py new file mode 100644 index 00000000000..94611463135 --- /dev/null +++ b/tests/integration/database/migrations/test_provider_url_index.py @@ -0,0 +1,78 @@ +import os +import uuid +from collections.abc import Iterator +from typing import Final + +import pytest +from litellm_proxy_extras.utils import MCP_SERVER_TABLE, PROVIDER_URL_INDEX, ProxyExtrasDBManager + +psycopg = pytest.importorskip("psycopg") + +pytestmark = pytest.mark.timeout(120) + +requires_db: Final = pytest.mark.skipif( + "DATABASE_URL" not in os.environ, + reason="requires a postgres database (DATABASE_URL)", +) + + +def _base_url() -> str: + return os.environ["DATABASE_URL"].split("?")[0] + + +def _scratch_schema(monkeypatch: pytest.MonkeyPatch, columns: str) -> Iterator[str]: + schema: Final = f"provider_url_{uuid.uuid4().hex[:8]}" + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE SCHEMA "{schema}"') + conn.execute(f'CREATE TABLE "{schema}"."{MCP_SERVER_TABLE}" ({columns})') + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}") + yield schema + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP SCHEMA "{schema}" CASCADE') + + +@pytest.fixture +def pushed_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + """What `prisma db push` or schema adoption leaves: the column, but not the SQL-only index.""" + yield from _scratch_schema(monkeypatch, 'server_id TEXT PRIMARY KEY, url TEXT, execution_mode TEXT NOT NULL') + + +@pytest.fixture +def older_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + yield from _scratch_schema(monkeypatch, "server_id TEXT PRIMARY KEY, url TEXT") + + +def _insert(schema: str, server_id: str, url: str, mode: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute( + f'INSERT INTO "{schema}"."{MCP_SERVER_TABLE}" (server_id, url, execution_mode) VALUES (%s, %s, %s)', + (server_id, url, mode), + ) + + +def _index_exists(schema: str) -> bool: + with psycopg.connect(_base_url(), autocommit=True) as conn: + row = conn.execute("SELECT to_regclass(%s)", (f'"{schema}"."{PROVIDER_URL_INDEX}"',)).fetchone() + return row is not None and row[0] is not None + + +@requires_db +def test_setup_creates_the_provider_url_index_a_pushed_database_lacks(pushed_schema: str) -> None: + assert _index_exists(pushed_schema) is False + + assert ProxyExtrasDBManager.ensure_provider_url_index() is True + assert ProxyExtrasDBManager.ensure_provider_url_index() is True + + assert _index_exists(pushed_schema) is True + _insert(pushed_schema, "gateway-a", "https://same.example/mcp", "litellm") + _insert(pushed_schema, "gateway-b", "https://same.example/mcp", "litellm") + _insert(pushed_schema, "provider-a", "https://same.example/mcp", "provider") + with pytest.raises(psycopg.errors.UniqueViolation): + _insert(pushed_schema, "provider-b", "https://same.example/mcp", "provider") + + +@requires_db +def test_setup_skips_the_provider_url_index_before_the_execution_mode_column_exists(older_schema: str) -> None: + assert ProxyExtrasDBManager.ensure_provider_url_index() is False + assert _index_exists(older_schema) is False From 86ff6f994c3ec346bb6442b35387ce6beaaf7161 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Tue, 3 Nov 2026 17:50:49 -0800 Subject: [PATCH 5/5] revert(mcp): drop startup creation of the provider URL index Duplicate provider URLs are now rejected on write by an app-level check under an advisory lock, and enforcement considers every registration of a URL, so nothing depends on the partial unique index. The index stays in the regular migration as an extra guard; startup no longer runs DDL outside prisma migrate. This reverts commit 54ec7f7e51. --- .../litellm_proxy_extras/utils.py | 63 --------------- .../migrations/test_provider_url_index.py | 78 ------------------- 2 files changed, 141 deletions(-) delete mode 100644 tests/integration/database/migrations/test_provider_url_index.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index c51fc45f89e..8b19576d86e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -67,8 +67,6 @@ _MIGRATION_TS_RE = re.compile(r"^(\d{14})_") _MIGRATION_DEADLOCK_MARKER = "deadlock detected" _MIGRATION_LOCK_TIMEOUT_MARKER = "canceling statement due to lock timeout" INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big") -MCP_SERVER_TABLE: Final = "LiteLLM_MCPServerTable" -PROVIDER_URL_INDEX: Final = "LiteLLM_MCPServerTable_provider_url_key" _TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$") _INVALID_LITELLM_INDEXES_SQL: Final = ( "SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) " @@ -942,66 +940,6 @@ class ProxyExtrasDBManager: return False return not remaining - @staticmethod - def ensure_provider_url_index(lock_timeout: str = "30s") -> bool: - """Create the partial unique index on provider MCP URLs when it is missing. - - The index exists only in migration SQL, because Prisma's schema cannot express a - partial index. `prisma db push` never creates it, and adopting a pushed database - marks migrations applied without running their SQL, so every database setup path - ends here. Never raises: returns True when the index exists, False when creation was - skipped or failed and will be retried on the next database setup run.""" - prisma_url: Final = os.getenv("DATABASE_URL") - if not prisma_url: - return False - try: - import psycopg - from psycopg import sql - except ImportError: - logger.warning( - "psycopg is not installed; skipping the provider MCP URL index check. " - "Install the litellm[extra_proxy] extra, which includes psycopg." - ) - return False - - schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public" - cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url) - index_name: Final = sql.Identifier(schema, PROVIDER_URL_INDEX) - try: - with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn: - present: Final = conn.execute("SELECT to_regclass(%s)", (index_name.as_string(conn),)).fetchone() - if present is not None and present[0] is not None: - return True - column: Final = conn.execute( - "SELECT 1 FROM information_schema.columns " - "WHERE table_schema = %s AND table_name = %s AND column_name = 'execution_mode'", - (schema, MCP_SERVER_TABLE), - ).fetchone() - if column is None: - return False - conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout))) - with held_migration_lock(conn) as held: - if not held: - logger.info( - "Another process holds the migration lock; leaving the provider MCP URL index " - "to the next database setup run" - ) - return False - conn.execute( - sql.SQL( - "CREATE UNIQUE INDEX IF NOT EXISTS {} ON {} (\"url\") WHERE \"execution_mode\" = 'provider'" - ).format(sql.Identifier(PROVIDER_URL_INDEX), sql.Identifier(schema, MCP_SERVER_TABLE)) - ) - except psycopg.Error as e: - logger.warning( - "Could not create the provider MCP URL unique index, will retry on the next database setup run. " - "Error: %s", - e, - ) - return False - logger.info("Created the provider MCP URL unique index that prisma db push and schema adoption skip") - return True - @staticmethod def _repair_under_migration_lock( conn: "psycopg.Connection[tuple[str, str, str]]", @@ -1484,7 +1422,6 @@ class ProxyExtrasDBManager: if not migrated: return False ProxyExtrasDBManager.repair_invalid_indexes() - ProxyExtrasDBManager.ensure_provider_url_index() ProxyExtrasDBManager.apply_replica_identity_full_if_requested() return True diff --git a/tests/integration/database/migrations/test_provider_url_index.py b/tests/integration/database/migrations/test_provider_url_index.py deleted file mode 100644 index 94611463135..00000000000 --- a/tests/integration/database/migrations/test_provider_url_index.py +++ /dev/null @@ -1,78 +0,0 @@ -import os -import uuid -from collections.abc import Iterator -from typing import Final - -import pytest -from litellm_proxy_extras.utils import MCP_SERVER_TABLE, PROVIDER_URL_INDEX, ProxyExtrasDBManager - -psycopg = pytest.importorskip("psycopg") - -pytestmark = pytest.mark.timeout(120) - -requires_db: Final = pytest.mark.skipif( - "DATABASE_URL" not in os.environ, - reason="requires a postgres database (DATABASE_URL)", -) - - -def _base_url() -> str: - return os.environ["DATABASE_URL"].split("?")[0] - - -def _scratch_schema(monkeypatch: pytest.MonkeyPatch, columns: str) -> Iterator[str]: - schema: Final = f"provider_url_{uuid.uuid4().hex[:8]}" - with psycopg.connect(_base_url(), autocommit=True) as conn: - conn.execute(f'CREATE SCHEMA "{schema}"') - conn.execute(f'CREATE TABLE "{schema}"."{MCP_SERVER_TABLE}" ({columns})') - monkeypatch.delenv("DIRECT_URL", raising=False) - monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}") - yield schema - with psycopg.connect(_base_url(), autocommit=True) as conn: - conn.execute(f'DROP SCHEMA "{schema}" CASCADE') - - -@pytest.fixture -def pushed_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: - """What `prisma db push` or schema adoption leaves: the column, but not the SQL-only index.""" - yield from _scratch_schema(monkeypatch, 'server_id TEXT PRIMARY KEY, url TEXT, execution_mode TEXT NOT NULL') - - -@pytest.fixture -def older_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: - yield from _scratch_schema(monkeypatch, "server_id TEXT PRIMARY KEY, url TEXT") - - -def _insert(schema: str, server_id: str, url: str, mode: str) -> None: - with psycopg.connect(_base_url(), autocommit=True) as conn: - conn.execute( - f'INSERT INTO "{schema}"."{MCP_SERVER_TABLE}" (server_id, url, execution_mode) VALUES (%s, %s, %s)', - (server_id, url, mode), - ) - - -def _index_exists(schema: str) -> bool: - with psycopg.connect(_base_url(), autocommit=True) as conn: - row = conn.execute("SELECT to_regclass(%s)", (f'"{schema}"."{PROVIDER_URL_INDEX}"',)).fetchone() - return row is not None and row[0] is not None - - -@requires_db -def test_setup_creates_the_provider_url_index_a_pushed_database_lacks(pushed_schema: str) -> None: - assert _index_exists(pushed_schema) is False - - assert ProxyExtrasDBManager.ensure_provider_url_index() is True - assert ProxyExtrasDBManager.ensure_provider_url_index() is True - - assert _index_exists(pushed_schema) is True - _insert(pushed_schema, "gateway-a", "https://same.example/mcp", "litellm") - _insert(pushed_schema, "gateway-b", "https://same.example/mcp", "litellm") - _insert(pushed_schema, "provider-a", "https://same.example/mcp", "provider") - with pytest.raises(psycopg.errors.UniqueViolation): - _insert(pushed_schema, "provider-b", "https://same.example/mcp", "provider") - - -@requires_db -def test_setup_skips_the_provider_url_index_before_the_execution_mode_column_exists(older_schema: str) -> None: - assert ProxyExtrasDBManager.ensure_provider_url_index() is False - assert _index_exists(older_schema) is False