scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-23 00:41:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

refactor(mcp): validate initialize body and general_settings with pydantic in the client allowlist

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-17 19:53:33 +00:00
parent 99eeb813c4
commit 21cafc8780
2 changed files with 29 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
litellm/proxy/_experimental/mcp_server/client_allowlist.py
View file

@ -4,11 +4,11 @@ Gateway-level allowlist of MCP client applications, matched against the
name is client-supplied, so this is a policy control and not a security boundary.
"""
import json
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final, Literal
from pydantic import TypeAdapter, ValidationError
from pydantic import BaseModel, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
@ -16,6 +16,19 @@ from litellm._logging import verbose_logger
MCP_ALLOWED_CLIENTS_SETTING: Final = "mcp_allowed_clients"
_ALLOWED_CLIENTS_ADAPTER: Final = TypeAdapter(list[str])
_GENERAL_SETTINGS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
class _ClientInfo(BaseModel):
name: str | None = None
class _InitializeParams(BaseModel):
clientInfo: _ClientInfo | None = None
class _InitializeRequest(BaseModel):
params: _InitializeParams | None = None
class MCPClientForbiddenBody(TypedDict):
@ -57,15 +70,21 @@ def parse_allowed_mcp_clients(raw_setting: object) -> frozenset[str] | None:
return frozenset()
def allowed_mcp_clients_from_general_settings(general_settings: object) -> frozenset[str] | None:
"""Reads the allowlist out of the proxy's untyped general_settings mapping."""
return parse_allowed_mcp_clients(
_GENERAL_SETTINGS_ADAPTER.validate_python(general_settings).get(MCP_ALLOWED_CLIENTS_SETTING)
)
def extract_mcp_client_name(body: bytes) -> str | None:
try:
data: Final = json.loads(body)
except (json.JSONDecodeError, UnicodeDecodeError):
request: Final = _InitializeRequest.model_validate_json(body)
except ValidationError:
return None
params: Final = data.get("params") if isinstance(data, dict) else None
client_info: Final = params.get("clientInfo") if isinstance(params, dict) else None
name: Final = client_info.get("name") if isinstance(client_info, dict) else None
return name if isinstance(name, str) and name else None
client_info: Final = request.params.clientInfo if request.params is not None else None
name: Final = client_info.name if client_info is not None else None
return name if name else None
def check_mcp_client_allowed(body: bytes, allowed_clients: frozenset[str] | None) -> MCPClientRejection | None:

4
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -465,8 +465,8 @@ if MCP_AVAILABLE:
)
from litellm.proxy._experimental.mcp_server.client_allowlist import (
MCP_ALLOWED_CLIENTS_SETTING,
allowed_mcp_clients_from_general_settings,
check_mcp_client_allowed,
parse_allowed_mcp_clients,
)
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
SERVER_OUTCOMES_META_KEY,
@ -3824,7 +3824,7 @@ if MCP_AVAILABLE:
def _load_allowed_mcp_clients() -> frozenset[str] | None:
from litellm.proxy.proxy_server import general_settings
return parse_allowed_mcp_clients(general_settings.get(MCP_ALLOWED_CLIENTS_SETTING))
return allowed_mcp_clients_from_general_settings(general_settings)
async def _reject_initialize_from_disallowed_client(
scope: Scope,