refactor(mcp): extract upstream preparation and support modern clients

litellm/experimental_mcp_client/client.py

@@ -36,6 +36,7 @@
     METHOD_NOT_FOUND,
     REQUEST_TIMEOUT,
     ClientCapabilities,
+    DiscoverResult,
     ElicitationCapability,
     FormElicitationCapability,
     GetPromptRequestParams,
@@ -84,6 +85,7 @@
     MCPUpstreamProtocol,
     credential_redirect_hook,
     has_header,
+    validate_mcp_protocol_transport,
     without_header,
 )
 
@@ -401,7 +403,10 @@ def __init__(
         logging_callback: Callable | None = None,
         protocol_version: MCPUpstreamProtocol = "auto",
     ):
-        self.protocol_version: MCPUpstreamProtocol = TypeAdapter(MCPUpstreamProtocol).validate_python(protocol_version)
+        self.protocol_version: MCPUpstreamProtocol = TypeAdapter[MCPUpstreamProtocol](
+            MCPUpstreamProtocol
+        ).validate_python(protocol_version)
+        validate_mcp_protocol_transport(self.protocol_version, transport_type)
         self.server_url: str = server_url
         self.transport_type: MCPTransport = transport_type
         self.auth_type: MCPAuthType = auth_type
@@ -540,6 +545,17 @@ def _get_safe_stdio_env(self, provided_env: dict[str, str] | None) -> dict[str,
 
         return safe_env
 
+    async def _prepare_session(self, session: ClientSession) -> InitializeResult | DiscoverResult:
+        if self.protocol_version != "2026-07-28":
+            return await self._initialize_session(session)
+        discovery: Final = DiscoverResult.model_validate(await session.send_discover(self.protocol_version))
+        if self.protocol_version not in discovery.supported_versions:
+            raise MCPError(code=-32022, message="Upstream did not accept the configured MCP protocol version")
+        session.adopt(discovery)
+        if session.protocol_version != self.protocol_version:
+            raise MCPError(code=-32022, message="Upstream selected an unsupported MCP protocol version")
+        return discovery
+
     async def _initialize_session(self, session: ClientSession) -> InitializeResult:
         if self.protocol_version == "auto":
             automatic: Final = await session.initialize()
@@ -623,7 +639,7 @@ async def receive_message(
                     )
                     session: Final = await session_ctx.__aenter__()
                     try:
-                        init_result: Final = await self._initialize_session(session)
+                        init_result: Final = await self._prepare_session(session)
                         instructions: Final = getattr(init_result, "instructions", None)
                         self._last_initialize_instructions = (
                             instructions.strip() or None if isinstance(instructions, str) else None

litellm/proxy/_experimental/mcp_server/contracts.py

@@ -3,12 +3,28 @@
 from dataclasses import dataclass, field
 from datetime import datetime
 from types import MappingProxyType
-from typing import Final, Protocol
+from typing import TYPE_CHECKING, Final, Literal, Protocol
 
 from litellm.proxy._experimental.mcp_server.tool_outcome import WireCompat
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.types.mcp_server.mcp_server_manager import MCPServer
 
+if TYPE_CHECKING:
+    from litellm.proxy._experimental.mcp_server.server_resolution import ResolvedMCPServer
+
+
+class TargetCatalog(Protocol):
+    async def resolve(
+        self,
+        server_id: str,
+        caller: UserAPIKeyAuth,
+        *,
+        is_admin_view: bool,
+        not_found_detail: Mapping[str, str],
+        forbidden_detail: Mapping[str, str],
+        non_admin_missing: Literal["not_found", "forbidden"],
+    ) -> "ResolvedMCPServer": ...
+
 
 def copy_caller(auth: UserAPIKeyAuth | None) -> UserAPIKeyAuth | None:
     if auth is None:

litellm/proxy/_experimental/mcp_server/db.py

@@ -8,6 +8,7 @@
 from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast
 
 from fastapi import HTTPException
+from pydantic import TypeAdapter
 from typing_extensions import ReadOnly
 
 from litellm._logging import verbose_proxy_logger
@@ -52,8 +53,8 @@
     VerificationTokenRepository,
 )
 from litellm.types.llms.custom_http import httpxSpecialProvider
-from litellm.types.mcp import MCPCredentials
-from litellm.types.mcp_server.mcp_server_manager import PinnedMCPTool
+from litellm.types.mcp import MCPCredentials, MCPTransportType, MCPUpstreamProtocol, validate_mcp_protocol_transport
+from litellm.types.mcp_server.mcp_server_manager import MCPInfo, PinnedMCPTool
 
 if TYPE_CHECKING:
     from prisma import models as prisma_db_models
@@ -600,6 +601,7 @@ async def _mcp_server_write_if_identifier_free(
     alias: str | None,
     exclude_server_id: str | None,
     write: "Callable[[TableActions[prisma_db_models.LiteLLM_MCPServerTable]], Awaitable[prisma_db_models.LiteLLM_MCPServerTable | None]]",
+    lock_server_id: str | None = None,
 ) -> "prisma_db_models.LiteLLM_MCPServerTable | McpIdentifierConflict | None":
     """Run ``write`` only when no other live row owns ``server_name``/``alias``.
 
@@ -619,6 +621,10 @@ async def _mcp_server_write_if_identifier_free(
         )
         if conflict is not None:
             return conflict
+        if lock_server_id is not None:
+            await tx.execute_raw(
+                'SELECT server_id FROM "LiteLLM_MCPServerTable" WHERE server_id=$1 FOR UPDATE', lock_server_id
+            )
         return await write(tx.litellm_mcpservertable)
 
 
@@ -1100,43 +1106,72 @@ async def get_draft_mcp_server(
     return table
 
 
+def _validate_mcp_protocol_write(
+    stored: "prisma_db_models.LiteLLM_MCPServerTable", data_dict: Mapping[str, object]
+) -> None:
+    raw_info: Final = data_dict.get("mcp_info", stored.mcp_info)
+    adapter: Final = TypeAdapter[MCPInfo | None](MCPInfo | None)
+    try:
+        info: Final = (
+            adapter.validate_json(raw_info) if isinstance(raw_info, str) else adapter.validate_python(raw_info)
+        )
+        validate_mcp_protocol_transport(
+            TypeAdapter[MCPUpstreamProtocol](MCPUpstreamProtocol).validate_python(
+                (info or {}).get("protocol_version", "auto")
+            ),
+            TypeAdapter[MCPTransportType](MCPTransportType).validate_python(
+                data_dict.get("transport", stored.transport)
+            ),
+        )
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+
 async def _update_mcp_server_row(
     prisma_client: PrismaClient,
     *,
     server_id: str,
     data_dict: Mapping[str, object],
 ) -> "prisma_db_models.LiteLLM_MCPServerTable | McpIdentifierConflict | None":
     identifier_write: Final = any(field in data_dict for field in ("server_name", "alias"))
+    protocol_write: Final = bool({"transport", "mcp_info"}.intersection(data_dict))
 
     async def _update(
         table: "TableActions[prisma_db_models.LiteLLM_MCPServerTable]",
     ) -> "prisma_db_models.LiteLLM_MCPServerTable | None":
+        if protocol_write:
+            stored: Final = await table.find_unique(where={"server_id": server_id})
+            if stored is None:
+                return None
+            _validate_mcp_protocol_write(stored, data_dict)
         return await table.update(
             where={"server_id": server_id},
             data=data_dict,
         )
 
-    if not identifier_write:
+    if not identifier_write and not protocol_write:
         return await _update(_mcp_server_table_actions(prisma_client))
     if "alias" in data_dict and not data_dict["alias"] and "server_name" not in data_dict:
         # Clearing the alias drops the prefix to the stored server_name, which
         # may already belong to another row, so that name needs the check too.
         existing: Final = await _db_find_mcp_server_row(prisma_client, server_id)
         if existing is None:
-            return await _update(_mcp_server_table_actions(prisma_client))
+            return None
         return await _mcp_server_write_if_identifier_free(
             prisma_client,
             server_name=existing.server_name,
             alias=None,
             exclude_server_id=server_id,
             write=_update,
+            lock_server_id=server_id if protocol_write else None,
         )
     return await _mcp_server_write_if_identifier_free(
         prisma_client,
         server_name=_identifier_field(data_dict, "server_name"),
         alias=_identifier_field(data_dict, "alias"),
         exclude_server_id=server_id,
         write=_update,
+        lock_server_id=server_id if protocol_write else None,
     )