scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-24 05:45:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72

fix(proxy): enforce MCP trust filter before membership checks

Browse source 
Apply trust scoring inside execute_mcp_tool so SSE and REST paths
cannot bypass the filter, wire prompt/resource routes, sanitize URLs
sent to Dominion Observatory, and keep permission vs trust 403 messages distinct.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
marcusacosta 2026-06-24 11:41:20 -07:00
parent 41bf08cc2b
commit b83e47102f
4 changed files with 226 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

59
litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py
View file

@ -12,6 +12,7 @@ import asyncio
import hashlib
from dataclasses import dataclass
from typing import TYPE_CHECKING, Annotated, Callable, Literal, Protocol, Sequence, cast
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
import httpx
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
@ -29,6 +30,25 @@ from litellm.types.llms.custom_http import httpxSpecialProvider
if TYPE_CHECKING:
from litellm.types.mcp_server.mcp_server_manager import MCPServer
TRUST_SCORE_THRESHOLD_DENIED = (
"Server does not meet the configured trust score threshold."
)
_SENSITIVE_QUERY_PARAM_NAMES = frozenset(
{
"access_token",
"api_key",
"apikey",
"auth",
"credential",
"credentials",
"key",
"password",
"secret",
"token",
}
)
class MCPTrustScoringConfig(BaseModel):
model_config = ConfigDict(frozen=True)
@ -147,6 +167,42 @@ def _normalize_server_url(server_url: str) -> str:
return server_url.rstrip("/")
def sanitize_url_for_trust_lookup(server_url: str) -> str:
parsed = urlparse(server_url)
hostname = parsed.hostname
if not hostname:
return _normalize_server_url(server_url)
port_suffix = f":{parsed.port}" if parsed.port else ""
netloc = f"{hostname}{port_suffix}"
filtered_query = tuple(
(key, value)
for key, value in parse_qsl(parsed.query, keep_blank_values=True)
if key.lower() not in _SENSITIVE_QUERY_PARAM_NAMES
)
sanitized = parsed._replace(
netloc=netloc,
query=urlencode(filtered_query),
fragment="",
)
return urlunparse(sanitized)
def assert_requested_server_passes_trust_filter(
*,
filtered_servers: Sequence[MCPServer],
server_id: str,
) -> None:
from fastapi import HTTPException
trusted_server_ids = {server.server_id for server in filtered_servers}
if server_id not in trusted_server_ids:
raise HTTPException(
status_code=403,
detail=TRUST_SCORE_THRESHOLD_DENIED,
)
def _cache_key_for_url(server_url: str) -> str:
digest = hashlib.sha256(
_normalize_server_url(server_url).encode("utf-8")
@ -303,7 +359,8 @@ class MCPTrustScoringClient:
)
async def get_trust_score(self, server_url: str) -> DominionObservatoryTrustResult:
normalized_url = _normalize_server_url(server_url)
sanitized_url = sanitize_url_for_trust_lookup(server_url)
normalized_url = _normalize_server_url(sanitized_url)
cached = await self._read_cached_result( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed
normalized_url
)

14
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -380,6 +380,7 @@ if MCP_AVAILABLE:
from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import (
apply_trust_filter_to_allowed_mcp_servers,
assert_requested_server_passes_trust_filter,
)
allowed_mcp_servers = list(
@ -387,6 +388,10 @@ if MCP_AVAILABLE:
allowed_mcp_servers
)
)
assert_requested_server_passes_trust_filter(
filtered_servers=allowed_mcp_servers,
server_id=canonical_server_id,
)
return allowed_mcp_servers, canonical_server_id
@ -476,13 +481,20 @@ if MCP_AVAILABLE:
from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import (
apply_trust_filter_to_allowed_mcp_servers,
assert_requested_server_passes_trust_filter,
)
return list(
filtered_servers = list(
await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed
allowed_mcp_servers
)
)
assert_requested_server_passes_trust_filter(
filtered_servers=filtered_servers,
server_id=server_id,
)
return filtered_servers
async def _list_tools_for_single_server(
server_id: str,

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

@ -1917,6 +1917,16 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
)
from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import (
apply_trust_filter_to_allowed_mcp_servers,
)
allowed_mcp_servers = list(
await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed
allowed_mcp_servers
)
)
# Get prompts from each allowed server
all_prompts = []
for server in allowed_mcp_servers:
@ -1976,6 +1986,16 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
)
from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import (
apply_trust_filter_to_allowed_mcp_servers,
)
allowed_mcp_servers = list(
await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed
allowed_mcp_servers
)
)
all_resources: List[Resource] = []
for server in allowed_mcp_servers:
if server is None:
@ -2032,6 +2052,16 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
)
from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import (
apply_trust_filter_to_allowed_mcp_servers,
)
allowed_mcp_servers = list(
await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed
allowed_mcp_servers
)
)
all_resource_templates: List[ResourceTemplate] = []
for server in allowed_mcp_servers:
if server is None:
@ -2496,6 +2526,17 @@ if MCP_AVAILABLE:
Returns:
CallToolResult: Tool execution result
"""
from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import (
apply_trust_filter_to_allowed_mcp_servers,
assert_requested_server_passes_trust_filter,
)
allowed_mcp_servers = list(
await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed
allowed_mcp_servers
)
)
# Track resolved MCP server for both permission checks and dispatch
mcp_server: Optional[MCPServer] = None
requested_server_id: Optional[str] = kwargs.get("requested_server_id")
@ -2602,6 +2643,12 @@ if MCP_AVAILABLE:
if mcp_server is None:
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
if mcp_server is not None:
assert_requested_server_passes_trust_filter(
filtered_servers=allowed_mcp_servers,
server_id=mcp_server.server_id,
)
if mcp_server:
standard_logging_mcp_tool_call["mcp_server_cost_info"] = (
mcp_server.mcp_info or {}
@ -2893,6 +2940,23 @@ if MCP_AVAILABLE:
detail="User not allowed to get this prompt.",
)
from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import (
apply_trust_filter_to_allowed_mcp_servers,
TRUST_SCORE_THRESHOLD_DENIED,
)
allowed_mcp_servers = list(
await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed
allowed_mcp_servers
)
)
if not allowed_mcp_servers:
raise HTTPException(
status_code=403,
detail=TRUST_SCORE_THRESHOLD_DENIED,
)
# Extract server name from prefixed prompt name
original_prompt_name, server_name = split_server_prefix_from_name(name)
@ -2943,6 +3007,23 @@ if MCP_AVAILABLE:
detail="User not allowed to read this resource.",
)
from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import (
apply_trust_filter_to_allowed_mcp_servers,
TRUST_SCORE_THRESHOLD_DENIED,
)
allowed_mcp_servers = list(
await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed
allowed_mcp_servers
)
)
if not allowed_mcp_servers:
raise HTTPException(
status_code=403,
detail=TRUST_SCORE_THRESHOLD_DENIED,
)
if len(allowed_mcp_servers) != 1:
raise HTTPException(
status_code=400,

80
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py
View file

@ -8,11 +8,13 @@ import pytest
from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import (
MCPTrustScoringClient,
MCPTrustScoringConfig,
assert_requested_server_passes_trust_filter,
filter_mcp_servers_by_trust,
initialize_mcp_trust_scoring_from_config,
normalize_trust_score,
parse_trust_response,
rank_mcp_servers_by_trust,
sanitize_url_for_trust_lookup,
)
from litellm.proxy._types import MCPTransport
from litellm.types.mcp import MCPTransportType
@ -119,6 +121,60 @@ def test_normalize_trust_score_maps_do_percent_scale_to_fraction() -> None:
assert abs(low_trust - 0.01) < 1e-9
def test_sanitize_url_for_trust_lookup_strips_userinfo_and_sensitive_query_params() -> (
None
):
sanitized = sanitize_url_for_trust_lookup(
"https://secret-token@mcp.example.com:8443/mcp?api_key=leak&safe=1#frag"
)
assert sanitized == "https://mcp.example.com:8443/mcp?safe=1"
def test_assert_requested_server_passes_trust_filter_rejects_untrusted_server() -> None:
from fastapi import HTTPException
trusted = _server(server_id="srv-trusted")
with pytest.raises(HTTPException) as exc_info:
assert_requested_server_passes_trust_filter(
filtered_servers=[trusted],
server_id="srv-untrusted",
)
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_get_trust_score_forwards_sanitized_url_to_do() -> None:
http_client = _FakeHttpClient(
[
_trust_response(
payload={
"found": True,
"server_url": "https://mcp.example.com/mcp",
"trust_score": 90,
}
)
]
)
config = MCPTrustScoringConfig(enabled=True)
client = _isolated_client(
config,
http_client,
)
await client.get_trust_score(
"https://secret@mcp.example.com/mcp?api_key=leak&safe=1"
)
assert http_client.calls == [
(
config.api_url,
{"url": "https://mcp.example.com/mcp?safe=1"},
)
]
def test_parse_trust_response_found_server() -> None:
result = parse_trust_response(
server_url="https://alpha.example.com/mcp",
@ -195,7 +251,9 @@ async def test_filter_servers_by_trust_excludes_low_score() -> None:
filtered = await client.filter_servers_by_trust(
[
_server(server_id="srv-1", name="alpha", url="https://alpha.example.com/mcp"),
_server(
server_id="srv-1", name="alpha", url="https://alpha.example.com/mcp"
),
_server(server_id="srv-2", name="beta", url="https://beta.example.com/mcp"),
]
)
@ -205,7 +263,9 @@ async def test_filter_servers_by_trust_excludes_low_score() -> None:
@pytest.mark.asyncio
async def test_filter_servers_fail_open_on_lookup_error() -> None:
http_client = _FakeHttpClient([_trust_response(status_code=503, payload={"error": "busy"})])
http_client = _FakeHttpClient(
[_trust_response(status_code=503, payload={"error": "busy"})]
)
client = _isolated_client(
MCPTrustScoringConfig(enabled=True, fail_open=True),
http_client,
@ -220,7 +280,9 @@ async def test_filter_servers_fail_open_on_lookup_error() -> None:
@pytest.mark.asyncio
async def test_filter_servers_fail_closed_on_lookup_error() -> None:
http_client = _FakeHttpClient([_trust_response(status_code=503, payload={"error": "busy"})])
http_client = _FakeHttpClient(
[_trust_response(status_code=503, payload={"error": "busy"})]
)
client = _isolated_client(
MCPTrustScoringConfig(enabled=True, fail_open=False),
http_client,
@ -275,7 +337,9 @@ async def test_rank_servers_by_trust_orders_highest_first() -> None:
ranked = await client.rank_servers_by_trust(
[
_server(server_id="srv-2", name="beta", url="https://beta.example.com/mcp"),
_server(server_id="srv-1", name="alpha", url="https://alpha.example.com/mcp"),
_server(
server_id="srv-1", name="alpha", url="https://alpha.example.com/mcp"
),
]
)
@ -314,7 +378,9 @@ async def test_module_helpers_use_explicit_client() -> None:
@pytest.mark.asyncio
async def test_manager_list_tools_applies_trust_filter(monkeypatch: pytest.MonkeyPatch) -> None:
async def test_manager_list_tools_applies_trust_filter(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
@ -347,7 +413,9 @@ async def test_manager_list_tools_applies_trust_filter(monkeypatch: pytest.Monke
)
manager = MCPServerManager()
alpha = _server(server_id="srv-1", name="alpha", url="https://alpha.example.com/mcp")
alpha = _server(
server_id="srv-1", name="alpha", url="https://alpha.example.com/mcp"
)
beta = _server(server_id="srv-2", name="beta", url="https://beta.example.com/mcp")
manager.registry = {"srv-1": alpha, "srv-2": beta}