From ab3744a8ccde7ae0633bde75c3b461008e339dd2 Mon Sep 17 00:00:00 2001 From: marcusacosta Date: Tue, 16 Jun 2026 19:23:55 -0700 Subject: [PATCH 1/5] feat(proxy): add optional MCP trust scoring via Dominion Observatory Score HTTP MCP server URLs for routing policy: filter tool listings, rank and fallback on tool calls, with cached DO lookups and fail-open defaults. Co-authored-by: Cursor --- litellm/constants.py | 17 + .../mcp_server/mcp_server_manager.py | 358 +++++++++---- .../mcp_server/mcp_trust_scoring.py | 505 ++++++++++++++++++ .../mcp_server/rest_endpoints.py | 32 +- .../proxy/_experimental/mcp_server/server.py | 24 + litellm/proxy/proxy_config.yaml | 8 +- litellm/proxy/proxy_server.py | 23 +- .../mcp_server/test_mcp_trust_scoring.py | 366 +++++++++++++ 8 files changed, 1223 insertions(+), 110 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py diff --git a/litellm/constants.py b/litellm/constants.py index c0e265c0e4a..02b428c96bc 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -110,6 +110,23 @@ MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH = int( os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150) ) +# MCP Dominion Observatory trust scoring defaults +DEFAULT_MCP_TRUST_SCORING_API_URL = str( + os.getenv( + "DEFAULT_MCP_TRUST_SCORING_API_URL", + "https://dominionobservatory.com/api/trust", + ) +) +DEFAULT_MCP_TRUST_SCORING_MIN_SCORE = float( + os.getenv("DEFAULT_MCP_TRUST_SCORING_MIN_SCORE", 0.5) +) +DEFAULT_MCP_TRUST_SCORING_CACHE_TTL = int( + os.getenv("DEFAULT_MCP_TRUST_SCORING_CACHE_TTL", 3600) +) +MCP_TRUST_SCORING_CACHE_MAX_SIZE = int( + os.getenv("MCP_TRUST_SCORING_CACHE_MAX_SIZE", 500) +) + # Semantic Guard Defaults DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL = str( os.getenv("DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL", "text-embedding-3-small") diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index afec884cd96..a99baa8cbd0 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,7 +13,7 @@ import json import os import re import time -from typing import Any, Callable, Dict, List, Literal, Optional, Set, Tuple, Union, cast +from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Set, Tuple, Union, cast from urllib.parse import urlparse import anyio @@ -1623,6 +1623,22 @@ class MCPServerManager: """ allowed_mcp_servers = await self.get_allowed_mcp_servers(user_api_key_auth) + from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( + apply_trust_filter_to_allowed_mcp_servers, + ) + + allowed_server_objects = tuple( + server + for server_id in allowed_mcp_servers + if (server := self.get_mcp_server_by_id(server_id)) is not None + ) + allowed_mcp_servers = [ + server.server_id + for server in await apply_trust_filter_to_allowed_mcp_servers( + allowed_server_objects + ) + ] + verbose_logger.debug("SERVER MANAGER LISTING TOOLS") async def _fetch_server_tools(server_id: str) -> List[MCPTool]: @@ -3552,6 +3568,242 @@ class MCPServerManager: return cast(CallToolResult, result) + @staticmethod + def _should_attempt_trust_fallback(error: BaseException) -> bool: + if isinstance(error, (BlockedPiiEntityError, GuardrailRaisedException, ValueError)): + return False + if isinstance(error, HTTPException): + return error.status_code not in {401, 403, 412} + return True + + def _server_exposes_tool( + self, + server: MCPServer, + tool_name: str, + server_name: str, + ) -> bool: + if server_name: + server_matches = False + for identifier in (server.alias, server.server_name, server.name): + if identifier and normalize_server_name(identifier) == normalize_server_name( + server_name + ): + server_matches = True + break + if not server_matches: + return False + + for known_prefix in iter_known_server_prefixes(server): + candidate_names = ( + tool_name, + add_server_prefix_to_name(tool_name, known_prefix), + ) + for candidate_name in candidate_names: + mapped_server = self.tool_name_to_mcp_server_name_mapping.get( + candidate_name + ) + if mapped_server is None: + continue + normalized_mapped = normalize_server_name(mapped_server) + for identifier in ( + server.alias, + server.server_name, + server.name, + known_prefix, + ): + if identifier and normalized_mapped == normalize_server_name(identifier): + return True + return False + + async def _ordered_mcp_servers_for_tool_call( + self, + server_name: str, + name: str, + allowed_mcp_servers: Optional[Sequence[MCPServer]], + ) -> Tuple[MCPServer, ...]: + primary = self._resolve_mcp_server_for_tool_call(server_name, name) + + from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( + get_mcp_trust_scoring_client, + rank_mcp_servers_by_trust, + ) + + client = get_mcp_trust_scoring_client() + if client is None or not client.enabled or not allowed_mcp_servers: + return (primary,) + + allowed_ids = {server.server_id for server in allowed_mcp_servers} + if primary.server_id not in allowed_ids: + return (primary,) + + candidates = tuple( + server + for server in allowed_mcp_servers + if self._server_exposes_tool( + server=server, + tool_name=name, + server_name=server_name, + ) + ) + if len(candidates) <= 1: + return candidates if candidates else (primary,) + + ranked = await rank_mcp_servers_by_trust(candidates) + return ranked if ranked else (primary,) + + async def _call_tool_on_mcp_server( + self, + *, + mcp_server: MCPServer, + server_name: str, + name: str, + arguments: Dict[str, Any], + user_api_key_auth: Optional[UserAPIKeyAuth], + mcp_auth_header: Optional[str], + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], + proxy_logging_obj: Optional[ProxyLogging], + oauth2_headers: Optional[Dict[str, str]], + raw_headers: Optional[Dict[str, str]], + host_progress_callback: Optional[Callable], + start_time: datetime.datetime, + ) -> CallToolResult: + hook_result: Dict[str, Any] = {} + if proxy_logging_obj: + hook_result = await self.pre_call_tool_check( + name=name, + arguments=arguments, + server_name=server_name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=mcp_server, + raw_headers=raw_headers, + ) + if "arguments" in hook_result: + arguments = hook_result["arguments"] + + tasks: List = [] + if proxy_logging_obj: + during_hook_task = self._create_during_hook_task( + name=name, + arguments=arguments, + server_name_from_prefix=server_name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + start_time=start_time, + ) + tasks.append(during_hook_task) + + resolved_oauth2_headers = await self._resolve_oauth2_headers_for_tool_call( + mcp_server, oauth2_headers, user_api_key_auth + ) + + if mcp_server.spec_path: + verbose_logger.debug( + "Calling OpenAPI tool %s directly via HTTP handler", name + ) + if hook_result.get("extra_headers"): + verbose_logger.warning( + "pre_mcp_call hook returned extra_headers for OpenAPI-backed " + "MCP server '%s' — header injection is not supported for " + "OpenAPI servers; headers will be ignored. Use SSE/HTTP " + "transport to enable hook header injection.", + server_name, + ) + tasks.append( + asyncio.create_task( + self._call_openapi_tool_handler(mcp_server, name, arguments) + ) + ) + return await self._gather_openapi_tool_tasks(tasks, proxy_logging_obj) + + return await self._call_regular_mcp_tool( + mcp_server=mcp_server, + original_tool_name=name, + arguments=arguments, + tasks=tasks, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=resolved_oauth2_headers, + raw_headers=raw_headers, + proxy_logging_obj=proxy_logging_obj, + host_progress_callback=host_progress_callback, + hook_extra_headers=hook_result.get("extra_headers"), + user_api_key_auth=user_api_key_auth, + ) + + async def call_tool( + self, + server_name: str, + name: str, + arguments: Dict[str, Any], + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + mcp_auth_header: Optional[str] = None, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, + proxy_logging_obj: Optional[ProxyLogging] = None, + oauth2_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, + host_progress_callback: Optional[Callable] = None, + allowed_mcp_servers: Optional[Sequence[MCPServer]] = None, + ) -> CallToolResult: + """ + Call a tool with the given name and arguments + + Args: + server_name: Server name + name: Tool name + arguments: Tool arguments + user_api_key_auth: User authentication + mcp_auth_header: MCP auth header (deprecated) + mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} + proxy_logging_obj: Optional ProxyLogging object for hook integration + allowed_mcp_servers: Optional allowed servers for trust-ranked fallback + + + Returns: + CallToolResult from the MCP server + """ + start_time = datetime.datetime.now() + ordered_servers = await self._ordered_mcp_servers_for_tool_call( + server_name=server_name, + name=name, + allowed_mcp_servers=allowed_mcp_servers, + ) + + last_error: Optional[BaseException] = None + for attempt_index, mcp_server in enumerate(ordered_servers): + effective_server_name = server_name or mcp_server.name + try: + return await self._call_tool_on_mcp_server( + mcp_server=mcp_server, + server_name=effective_server_name, + name=name, + arguments=arguments, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + proxy_logging_obj=proxy_logging_obj, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + host_progress_callback=host_progress_callback, + start_time=start_time, + ) + except Exception as exc: + if not self._should_attempt_trust_fallback(exc): + raise + if attempt_index + 1 >= len(ordered_servers): + raise + last_error = exc + verbose_logger.debug( + "MCP trust fallback after failure on server %s for tool %s: %s", + mcp_server.name, + name, + exc, + ) + + if last_error is not None: + raise last_error + raise ValueError(f"Tool {name} not found") + def _resolve_mcp_server_for_tool_call( self, server_name: str, @@ -3658,110 +3910,6 @@ class MCPServerManager: ) raise e - async def call_tool( - self, - server_name: str, - name: str, - arguments: Dict[str, Any], - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - proxy_logging_obj: Optional[ProxyLogging] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - host_progress_callback: Optional[Callable] = None, - ) -> CallToolResult: - """ - Call a tool with the given name and arguments - - Args: - server_name: Server name - name: Tool name - arguments: Tool arguments - user_api_key_auth: User authentication - mcp_auth_header: MCP auth header (deprecated) - mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} - proxy_logging_obj: Optional ProxyLogging object for hook integration - - - Returns: - CallToolResult from the MCP server - """ - start_time = datetime.datetime.now() - mcp_server = self._resolve_mcp_server_for_tool_call(server_name, name) - - ######################################################### - # Pre MCP Tool Call Hook - # Allow validation and modification of tool calls before execution - # Using standard pre_call_hook - ######################################################### - hook_result: Dict[str, Any] = {} - if proxy_logging_obj: - hook_result = await self.pre_call_tool_check( - name=name, - arguments=arguments, - server_name=server_name, - user_api_key_auth=user_api_key_auth, - proxy_logging_obj=proxy_logging_obj, - server=mcp_server, - raw_headers=raw_headers, - ) - if "arguments" in hook_result: - arguments = hook_result["arguments"] - - # Prepare tasks for during hooks - tasks = [] - if proxy_logging_obj: - during_hook_task = self._create_during_hook_task( - name=name, - arguments=arguments, - server_name_from_prefix=server_name, - user_api_key_auth=user_api_key_auth, - proxy_logging_obj=proxy_logging_obj, - start_time=start_time, - ) - tasks.append(during_hook_task) - - oauth2_headers = await self._resolve_oauth2_headers_for_tool_call( - mcp_server, oauth2_headers, user_api_key_auth - ) - - # For OpenAPI servers, call the tool handler directly instead of via MCP client - if mcp_server.spec_path: - verbose_logger.debug( - "Calling OpenAPI tool %s directly via HTTP handler", name - ) - if hook_result.get("extra_headers"): - verbose_logger.warning( - "pre_mcp_call hook returned extra_headers for OpenAPI-backed " - "MCP server '%s' — header injection is not supported for " - "OpenAPI servers; headers will be ignored. Use SSE/HTTP " - "transport to enable hook header injection.", - server_name, - ) - tasks.append( - asyncio.create_task( - self._call_openapi_tool_handler(mcp_server, name, arguments) - ) - ) - else: - return await self._call_regular_mcp_tool( - mcp_server=mcp_server, - original_tool_name=name, - arguments=arguments, - tasks=tasks, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - proxy_logging_obj=proxy_logging_obj, - host_progress_callback=host_progress_callback, - hook_extra_headers=hook_result.get("extra_headers"), - user_api_key_auth=user_api_key_auth, - ) - - return await self._gather_openapi_tool_tasks(tasks, proxy_logging_obj) - ######################################################### # End of Methods that call the upstream MCP servers ######################################################### diff --git a/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py b/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py new file mode 100644 index 00000000000..1dc9c3d15c3 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py @@ -0,0 +1,505 @@ +""" +Optional MCP server trust scoring via Dominion Observatory. + +Queries https://dominionobservatory.com/api/trust for HTTP MCP server URLs and +uses the returned trust score to filter or rank servers before tool listing and +tool calls. Results are cached to respect the free-tier query limit. +""" + +from __future__ import annotations + +import asyncio +import hashlib +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable, Literal, Optional, Protocol, Sequence, Tuple, cast + +import httpx +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter + +from litellm._logging import verbose_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import ( + DEFAULT_MCP_TRUST_SCORING_API_URL, + DEFAULT_MCP_TRUST_SCORING_CACHE_TTL, + DEFAULT_MCP_TRUST_SCORING_MIN_SCORE, + MCP_TRUST_SCORING_CACHE_MAX_SIZE, +) +from litellm.types.llms.custom_http import httpxSpecialProvider + +if TYPE_CHECKING: + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +class MCPTrustScoringConfig(BaseModel): + model_config = ConfigDict(frozen=True) + + enabled: bool = False + min_trust_score: float = Field( + default=DEFAULT_MCP_TRUST_SCORING_MIN_SCORE, + ge=0.0, + le=1.0, + description="Minimum normalized trust score (0-1) required to expose or route to a server", + ) + cache_ttl_seconds: int = Field( + default=DEFAULT_MCP_TRUST_SCORING_CACHE_TTL, + ge=1, + description="Seconds to cache Dominion Observatory responses per server URL", + ) + api_url: str = Field( + default=DEFAULT_MCP_TRUST_SCORING_API_URL, + description="Dominion Observatory trust endpoint base URL", + ) + fail_open: bool = Field( + default=True, + description="When True, keep servers reachable if DO is down, rate-limited, or has no score yet", + ) + api_key: Optional[str] = Field( + default=None, + description="Optional API key for paid Dominion Observatory tiers", + ) + + +class DominionObservatoryTrustResult(BaseModel): + model_config = ConfigDict(frozen=True) + + server_url: str + found: bool + trust_score_normalized: Optional[float] = Field( + default=None, + description="Trust score normalized to 0-1", + ) + trust_score_raw: Optional[float] = Field( + default=None, + description="Raw trust score from Dominion Observatory (0-100 scale)", + ) + lookup_status: Literal["ok", "skipped", "error"] = "ok" + message: Optional[str] = None + + def meets_min_trust(self, min_trust_score: float) -> bool: + if self.lookup_status == "skipped": + return True + if self.trust_score_normalized is None: + return False + return self.trust_score_normalized >= min_trust_score + + def sort_key(self) -> Tuple[int, float]: + if self.lookup_status == "skipped": + return (1, 0.0) + if self.trust_score_normalized is None: + return (2, 0.0) + return (0, -self.trust_score_normalized) + + +class _TrustCacheEntry(BaseModel): + model_config = ConfigDict(frozen=True) + + payload_json: str + + +class _DominionObservatoryTrustApiResponse(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + found: bool = False + trust_score: Optional[float] = None + message: Optional[str] = None + + +class _SharedTrustCache(Protocol): + async def async_get_cache( + self, + key: str, + local_only: bool = False, + **kwargs: object, + ) -> object | None: ... + + async def async_set_cache( + self, + key: str, + value: object, + local_only: bool = False, + **kwargs: object, + ) -> None: ... + + +class _LocalTrustCache(Protocol): + def get_cache(self, key: str, **kwargs: object) -> object | None: ... + + def set_cache(self, key: str, value: object, **kwargs: object) -> None: ... + + +class _HttpClient(Protocol): + async def get( + self, + url: str, + *, + params: Optional[dict[str, str]] = None, + headers: Optional[dict[str, str]] = None, + ) -> httpx.Response: ... + + +@dataclass(frozen=True, slots=True) +class _ScoredServer: + server: MCPServer + trust: DominionObservatoryTrustResult + + +def _normalize_server_url(server_url: str) -> str: + return server_url.rstrip("/") + + +def _cache_key_for_url(server_url: str) -> str: + digest = hashlib.sha256(_normalize_server_url(server_url).encode("utf-8")).hexdigest()[ + :32 + ] + return f"mcp:do_trust:{digest}" + + +def normalize_trust_score(raw_score: object) -> Optional[float]: + if raw_score is None: + return None + if not isinstance(raw_score, (int, float)): + return None + score = float(raw_score) + if score < 0: + return None + if score <= 1.0: + return score + return min(score / 100.0, 1.0) + + +def parse_trust_response( + *, + server_url: str, + payload: object, +) -> DominionObservatoryTrustResult: + if not isinstance(payload, dict): + return DominionObservatoryTrustResult( + server_url=server_url, + found=False, + lookup_status="error", + message="Invalid Dominion Observatory response shape", + ) + + parsed = TypeAdapter(_DominionObservatoryTrustApiResponse).validate_python(payload) + + return DominionObservatoryTrustResult( + server_url=server_url, + found=parsed.found, + trust_score_normalized=normalize_trust_score(parsed.trust_score), + trust_score_raw=parsed.trust_score, + lookup_status="ok", + message=parsed.message, + ) + + +def _should_include_server( + *, + trust: DominionObservatoryTrustResult, + config: MCPTrustScoringConfig, +) -> bool: + if trust.lookup_status == "skipped": + return True + if trust.lookup_status == "error": + return config.fail_open + if not trust.found or trust.trust_score_normalized is None: + return config.fail_open + return trust.meets_min_trust(config.min_trust_score) + + +class MCPTrustScoringClient: + def __init__( + self, + config: MCPTrustScoringConfig, + *, + local_cache: Optional[InMemoryCache] = None, + cache_getter: Optional[Callable[[], _SharedTrustCache]] = None, + http_client_getter: Optional[Callable[[], _HttpClient]] = None, + ) -> None: + self._config = config + self._local_cache: _LocalTrustCache = cast( + _LocalTrustCache, + local_cache + or InMemoryCache( + max_size_in_memory=MCP_TRUST_SCORING_CACHE_MAX_SIZE, + default_ttl=config.cache_ttl_seconds, + ), + ) + self._cache_getter = cache_getter or _default_dual_cache_getter + self._http_client_getter = http_client_getter or _default_http_client_getter + self._locks: dict[str, asyncio.Lock] = {} + + @property + def config(self) -> MCPTrustScoringConfig: + return self._config + + @property + def enabled(self) -> bool: + return self._config.enabled + + @classmethod + def from_config_dict( + cls, + config: Optional[dict[str, object]], + ) -> Optional[MCPTrustScoringClient]: + if not config: + return None + parsed = MCPTrustScoringConfig.model_validate(config) + if not parsed.enabled: + return None + return cls(parsed) + + def _lock_for_url(self, server_url: str) -> asyncio.Lock: + return self._locks.setdefault(_cache_key_for_url(server_url), asyncio.Lock()) + + async def _read_cached_result(self, server_url: str) -> Optional[DominionObservatoryTrustResult]: + cache_key = _cache_key_for_url(server_url) + local_cached = self._local_cache.get_cache(cache_key) + if isinstance(local_cached, str): + return DominionObservatoryTrustResult.model_validate_json(local_cached) + + try: + shared_cache = self._cache_getter() + shared = await shared_cache.async_get_cache(cache_key) + except Exception: + shared = None + + if isinstance(shared, str): + self._local_cache.set_cache( + cache_key, + shared, + ttl=self._config.cache_ttl_seconds, + ) + return DominionObservatoryTrustResult.model_validate_json(shared) + if isinstance(shared, dict): + entry = _TrustCacheEntry.model_validate(shared) + self._local_cache.set_cache( + cache_key, + entry.payload_json, + ttl=self._config.cache_ttl_seconds, + ) + return DominionObservatoryTrustResult.model_validate_json(entry.payload_json) + return None + + async def _write_cached_result( + self, + server_url: str, + result: DominionObservatoryTrustResult, + ) -> None: + cache_key = _cache_key_for_url(server_url) + payload_json = result.model_dump_json() + self._local_cache.set_cache( + cache_key, + payload_json, + ttl=self._config.cache_ttl_seconds, + ) + try: + shared_cache = self._cache_getter() + await shared_cache.async_set_cache( + cache_key, + _TrustCacheEntry(payload_json=payload_json).model_dump(), + ttl=self._config.cache_ttl_seconds, + ) + except Exception as exc: + verbose_logger.debug( + "Failed to write MCP trust score to shared cache for %s: %s", + server_url, + exc, + ) + + async def get_trust_score(self, server_url: str) -> DominionObservatoryTrustResult: + normalized_url = _normalize_server_url(server_url) + cached = await self._read_cached_result(normalized_url) + if cached is not None: + return cached + + async with self._lock_for_url(normalized_url): + cached = await self._read_cached_result(normalized_url) + if cached is not None: + return cached + + result = await self._fetch_trust_score(normalized_url) + await self._write_cached_result(normalized_url, result) + return result + + async def _fetch_trust_score(self, server_url: str) -> DominionObservatoryTrustResult: + headers: dict[str, str] = {} + if self._config.api_key: + headers["Authorization"] = f"Bearer {self._config.api_key}" + + try: + client = self._http_client_getter() + response = await client.get( + self._config.api_url, + params={"url": server_url}, + headers=headers or None, + ) + response.raise_for_status() + parsed_api = TypeAdapter(_DominionObservatoryTrustApiResponse).validate_json( + response.text + ) + except Exception as exc: + verbose_logger.debug( + "Dominion Observatory trust lookup failed for %s: %s", + server_url, + exc, + ) + return DominionObservatoryTrustResult( + server_url=server_url, + found=False, + lookup_status="error", + message=str(exc), + ) + + return DominionObservatoryTrustResult( + server_url=server_url, + found=parsed_api.found, + trust_score_normalized=normalize_trust_score(parsed_api.trust_score), + trust_score_raw=parsed_api.trust_score, + lookup_status="ok", + message=parsed_api.message, + ) + + async def get_trust_score_for_server( + self, + server: MCPServer, + ) -> DominionObservatoryTrustResult: + if not server.url: + return DominionObservatoryTrustResult( + server_url=server.server_id, + found=False, + lookup_status="skipped", + message="Server has no HTTP URL to score", + ) + return await self.get_trust_score(server.url) + + async def score_servers( + self, + servers: Sequence[MCPServer], + ) -> Tuple[_ScoredServer, ...]: + if not servers: + return () + + trust_results = await asyncio.gather( + *(self.get_trust_score_for_server(server) for server in servers) + ) + return tuple( + _ScoredServer(server=server, trust=trust) + for server, trust in zip(servers, trust_results, strict=True) + ) + + async def filter_servers_by_trust( + self, + servers: Sequence[MCPServer], + ) -> Tuple[MCPServer, ...]: + if not self.enabled: + return tuple(servers) + + scored = await self.score_servers(servers) + return tuple( + entry.server + for entry in scored + if _should_include_server(trust=entry.trust, config=self._config) + ) + + async def rank_servers_by_trust( + self, + servers: Sequence[MCPServer], + ) -> Tuple[MCPServer, ...]: + if not self.enabled: + return tuple(servers) + + scored = await self.score_servers(servers) + ordered = sorted(scored, key=lambda entry: entry.trust.sort_key()) + return tuple(entry.server for entry in ordered) + + +def _default_http_client_getter() -> _HttpClient: + from litellm.llms.custom_httpx import http_handler # noqa: PLC0415 + + return cast( + _HttpClient, + http_handler.get_async_httpx_client( # pyright: ignore[reportUnknownMemberType] + llm_provider=httpxSpecialProvider.MCP, + ), + ) + + +def _default_dual_cache_getter() -> _SharedTrustCache: + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + + return cast(_SharedTrustCache, user_api_key_cache) + + +_global_mcp_trust_scoring_client: Optional[MCPTrustScoringClient] = None + + +def get_mcp_trust_scoring_client() -> Optional[MCPTrustScoringClient]: + return _global_mcp_trust_scoring_client + + +def _resolve_trust_scoring_config_values( + config: Optional[dict[str, object]], +) -> Optional[dict[str, object]]: + if not config: + return None + + resolved = dict(config) + api_key = resolved.get("api_key") + if isinstance(api_key, str) and api_key.startswith("os.environ/"): + from litellm.secret_managers.main import get_secret # noqa: PLC0415 + + resolved["api_key"] = get_secret(api_key) + return resolved + + +def initialize_mcp_trust_scoring_from_config( + config: Optional[dict[str, object]], +) -> Optional[MCPTrustScoringClient]: + global _global_mcp_trust_scoring_client + + client = MCPTrustScoringClient.from_config_dict( + _resolve_trust_scoring_config_values(config) + ) + _global_mcp_trust_scoring_client = client + + if client is None: + verbose_logger.debug("MCP trust scoring disabled") + else: + verbose_logger.info( + "MCP trust scoring enabled (min_trust_score=%s, cache_ttl_seconds=%s, fail_open=%s)", + client.config.min_trust_score, + client.config.cache_ttl_seconds, + client.config.fail_open, + ) + + return client + + +async def filter_mcp_servers_by_trust( + servers: Sequence[MCPServer], + *, + client: Optional[MCPTrustScoringClient] = None, +) -> Tuple[MCPServer, ...]: + active_client = client or get_mcp_trust_scoring_client() + if active_client is None: + return tuple(servers) + return await active_client.filter_servers_by_trust(servers) + + +async def rank_mcp_servers_by_trust( + servers: Sequence[MCPServer], + *, + client: Optional[MCPTrustScoringClient] = None, +) -> Tuple[MCPServer, ...]: + active_client = client or get_mcp_trust_scoring_client() + if active_client is None: + return tuple(servers) + return await active_client.rank_servers_by_trust(servers) + + +async def apply_trust_filter_to_allowed_mcp_servers( + servers: Sequence[MCPServer], + *, + client: Optional[MCPTrustScoringClient] = None, +) -> Tuple[MCPServer, ...]: + return await filter_mcp_servers_by_trust(servers, client=client) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 2149f079a3d..8a1bc763b86 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -378,6 +378,14 @@ if MCP_AVAILABLE: if server is not None: allowed_mcp_servers.append(server) + 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(allowed_mcp_servers) + ) + return allowed_mcp_servers, canonical_server_id async def _get_tools_for_single_server( @@ -463,7 +471,12 @@ if MCP_AVAILABLE: server = global_mcp_server_manager.get_mcp_server_by_id(allowed_server_id) if server is not None: allowed_mcp_servers.append(server) - return allowed_mcp_servers + + from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( + apply_trust_filter_to_allowed_mcp_servers, + ) + + return list(await apply_trust_filter_to_allowed_mcp_servers(allowed_mcp_servers)) async def _list_tools_for_single_server( server_id: str, @@ -634,6 +647,23 @@ if MCP_AVAILABLE: list(allowed_server_ids_set), _rest_client_ip ) + allowed_mcp_server_objects: List[MCPServer] = [] + for allowed_server_id in allowed_server_ids: + server = global_mcp_server_manager.get_mcp_server_by_id(allowed_server_id) + if server is not None: + allowed_mcp_server_objects.append(server) + + from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( + apply_trust_filter_to_allowed_mcp_servers, + ) + + allowed_server_ids = [ + server.server_id + for server in await apply_trust_filter_to_allowed_mcp_servers( + allowed_mcp_server_objects + ) + ] + list_tools_result = [] error_message = None diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 08e42e918e9..3665a05a280 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1712,6 +1712,14 @@ 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(allowed_mcp_servers) + ) + # Pre-fetch OAuth credentials only when at least one server uses OAuth2, # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers. _has_oauth2_server = any( @@ -2739,6 +2747,7 @@ if MCP_AVAILABLE: raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, host_progress_callback=host_progress_callback, + allowed_mcp_servers=allowed_mcp_servers, ) # Fall back to local tool registry with original name (legacy support) @@ -2802,6 +2811,19 @@ if MCP_AVAILABLE: detail="User not allowed to call this tool.", ) + 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(allowed_mcp_servers) + ) + if not allowed_mcp_servers: + raise HTTPException( + status_code=403, + detail="No MCP servers meet the configured trust score threshold.", + ) + # Delegate to execute_mcp_tool for execution response = await execute_mcp_tool( name=name, @@ -2982,6 +3004,7 @@ if MCP_AVAILABLE: raw_headers: Optional[Dict[str, str]] = None, litellm_logging_obj: Optional[Any] = None, host_progress_callback: Optional[Callable] = None, + allowed_mcp_servers: Optional[List[MCPServer]] = None, ) -> CallToolResult: """Handle tool execution for managed server tools""" # Import here to avoid circular import @@ -2998,6 +3021,7 @@ if MCP_AVAILABLE: raw_headers=raw_headers, proxy_logging_obj=proxy_logging_obj, host_progress_callback=host_progress_callback, + allowed_mcp_servers=allowed_mcp_servers, ) verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) return call_tool_result diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index a094eb84bf3..48752711541 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -60,4 +60,10 @@ litellm_settings: embedding_model: "text-embedding-3-small" top_k: 5 similarity_threshold: 0.3 - + + mcp_trust_scoring: + enabled: true + min_trust_score: 0.5 + cache_ttl_seconds: 3600 + fail_open: true + # api_key: os.environ/DO_API_KEY diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c5676fbb32f..745c6448ce1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -925,8 +925,7 @@ async def proxy_startup_event(app: FastAPI): redis_usage_cache=transaction_buffer_redis_cache, ) - ## SEMANTIC TOOL FILTER ## - # Read litellm_settings from config for semantic filter initialization + ## SEMANTIC TOOL FILTER / MCP TRUST SCORING ## try: verbose_proxy_logger.debug("About to initialize semantic tool filter") _config = proxy_config.get_config_state() @@ -939,8 +938,13 @@ async def proxy_startup_event(app: FastAPI): litellm_settings=_litellm_settings, ) verbose_proxy_logger.debug("After semantic tool filter initialization") + ProxyStartupEvent._initialize_mcp_trust_scoring( + litellm_settings=_litellm_settings, + ) except Exception as e: - verbose_proxy_logger.error(f"Semantic filter init failed: {e}", exc_info=True) + verbose_proxy_logger.error( + f"MCP semantic filter / trust scoring init failed: {e}", exc_info=True + ) ## JWT AUTH ## ProxyStartupEvent._initialize_jwt_auth( @@ -7594,6 +7598,19 @@ class ProxyStartupEvent: return RedisCache(**redis_env_kwargs) + @classmethod + def _initialize_mcp_trust_scoring( + cls, + litellm_settings: Dict[str, Any], + ) -> None: + from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( + initialize_mcp_trust_scoring_from_config, + ) + + initialize_mcp_trust_scoring_from_config( + litellm_settings.get("mcp_trust_scoring") + ) + @classmethod async def _initialize_semantic_tool_filter( cls, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py new file mode 100644 index 00000000000..8e924e78612 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py @@ -0,0 +1,366 @@ +""" +Tests for Dominion Observatory MCP trust scoring. +""" + +import httpx +import pytest + +from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( + MCPTrustScoringClient, + MCPTrustScoringConfig, + filter_mcp_servers_by_trust, + initialize_mcp_trust_scoring_from_config, + normalize_trust_score, + parse_trust_response, + rank_mcp_servers_by_trust, +) +from litellm.proxy._types import MCPTransport +from litellm.types.mcp import MCPTransportType +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +def _server( + *, + server_id: str = "srv-1", + name: str = "alpha", + url: str | None = "https://alpha.example.com/mcp", + transport: MCPTransportType = MCPTransport.http, +) -> MCPServer: + return MCPServer( + server_id=server_id, + name=name, + url=url, + transport=transport, + ) + + +def _trust_response( + *, + status_code: int = 200, + payload: object | None = None, +) -> httpx.Response: + request = httpx.Request("GET", "https://dominionobservatory.com/api/trust") + return httpx.Response( + status_code=status_code, + request=request, + json=payload, + ) + + +class _FakeHttpClient: + def __init__(self, responses: list[httpx.Response]) -> None: + self._responses = list(responses) + self.calls: list[tuple[str, dict[str, str] | None]] = [] + + async def get( + self, + url: str, + *, + params: dict[str, str] | None = None, + headers: dict[str, str] | None = None, + ) -> httpx.Response: + del headers + self.calls.append((url, params)) + if not self._responses: + raise RuntimeError("no mocked responses left") + return self._responses.pop(0) + + +class _FakeSharedCache: + def __init__(self) -> None: + self.store: dict[str, object] = {} + + async def async_get_cache( + self, + key: str, + local_only: bool = False, + **kwargs: object, + ) -> object | None: + del local_only, kwargs + return self.store.get(key) + + async def async_set_cache( + self, + key: str, + value: object, + local_only: bool = False, + **kwargs: object, + ) -> None: + del local_only, kwargs + self.store[key] = value + + +def _isolated_client( + config: MCPTrustScoringConfig, + http_client: _FakeHttpClient, +) -> MCPTrustScoringClient: + shared_cache = _FakeSharedCache() + + def _cache_getter() -> _FakeSharedCache: + return shared_cache + + def _http_client_getter() -> _FakeHttpClient: + return http_client + + return MCPTrustScoringClient( + config, + cache_getter=_cache_getter, + http_client_getter=_http_client_getter, + ) + + +def test_normalize_trust_score_accepts_fraction_and_percent() -> None: + assert normalize_trust_score(0.85) == 0.85 + normalized = normalize_trust_score(92.5) + assert normalized is not None + assert abs(normalized - 0.925) < 1e-9 + + +def test_parse_trust_response_found_server() -> None: + result = parse_trust_response( + server_url="https://alpha.example.com/mcp", + payload={ + "found": True, + "server_url": "https://alpha.example.com/mcp", + "trust_score": 92.5, + }, + ) + + assert result.found is True + assert result.trust_score_raw == 92.5 + assert result.trust_score_normalized is not None + assert abs(result.trust_score_normalized - 0.925) < 1e-9 + assert result.lookup_status == "ok" + + +def test_initialize_from_config_disabled() -> None: + client = initialize_mcp_trust_scoring_from_config({"enabled": False}) + assert client is None + + +@pytest.mark.asyncio +async def test_get_trust_score_fetches_and_caches() -> None: + http_client = _FakeHttpClient( + [ + _trust_response( + payload={ + "found": True, + "server_url": "https://cache-test.example.com/mcp", + "trust_score": 80, + } + ) + ] + ) + client = _isolated_client( + MCPTrustScoringConfig(enabled=True), + http_client, + ) + + first = await client.get_trust_score("https://cache-test.example.com/mcp") + second = await client.get_trust_score("https://cache-test.example.com/mcp") + + assert first.trust_score_normalized is not None + assert abs(first.trust_score_normalized - 0.8) < 1e-9 + assert second == first + assert len(http_client.calls) == 1 + + +@pytest.mark.asyncio +async def test_filter_servers_by_trust_excludes_low_score() -> None: + http_client = _FakeHttpClient( + [ + _trust_response( + payload={ + "found": True, + "server_url": "https://alpha.example.com/mcp", + "trust_score": 90, + } + ), + _trust_response( + payload={ + "found": True, + "server_url": "https://beta.example.com/mcp", + "trust_score": 20, + } + ), + ] + ) + client = _isolated_client( + MCPTrustScoringConfig(enabled=True, min_trust_score=0.5), + http_client, + ) + + filtered = await client.filter_servers_by_trust( + [ + _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"), + ] + ) + + assert [server.server_id for server in filtered] == ["srv-1"] + + +@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"})]) + client = _isolated_client( + MCPTrustScoringConfig(enabled=True, fail_open=True), + http_client, + ) + + filtered = await client.filter_servers_by_trust( + [_server(url="https://fail-open.example.com/mcp")] + ) + + assert len(filtered) == 1 + + +@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"})]) + client = _isolated_client( + MCPTrustScoringConfig(enabled=True, fail_open=False), + http_client, + ) + + filtered = await client.filter_servers_by_trust( + [_server(url="https://fail-closed.example.com/mcp")] + ) + + assert filtered == () + + +@pytest.mark.asyncio +async def test_stdio_server_is_skipped_not_filtered() -> None: + client = _isolated_client( + MCPTrustScoringConfig(enabled=True, min_trust_score=0.99), + _FakeHttpClient([]), + ) + + filtered = await client.filter_servers_by_trust( + [_server(server_id="stdio-1", url=None, transport=MCPTransport.stdio)] + ) + + assert [server.server_id for server in filtered] == ["stdio-1"] + + +@pytest.mark.asyncio +async def test_rank_servers_by_trust_orders_highest_first() -> None: + http_client = _FakeHttpClient( + [ + _trust_response( + payload={ + "found": True, + "server_url": "https://beta.example.com/mcp", + "trust_score": 40, + } + ), + _trust_response( + payload={ + "found": True, + "server_url": "https://alpha.example.com/mcp", + "trust_score": 95, + } + ), + ] + ) + client = _isolated_client( + MCPTrustScoringConfig(enabled=True), + http_client, + ) + + 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"), + ] + ) + + assert [server.server_id for server in ranked] == ["srv-1", "srv-2"] + + +@pytest.mark.asyncio +async def test_module_helpers_use_explicit_client() -> None: + http_client = _FakeHttpClient( + [ + _trust_response( + payload={ + "found": True, + "server_url": "https://alpha.example.com/mcp", + "trust_score": 95, + } + ) + ] + ) + client = _isolated_client( + MCPTrustScoringConfig(enabled=True, min_trust_score=0.5), + http_client, + ) + + filtered = await filter_mcp_servers_by_trust( + [_server(url="https://helper-test.example.com/mcp")], + client=client, + ) + ranked = await rank_mcp_servers_by_trust( + [_server(url="https://helper-test.example.com/mcp")], + client=client, + ) + + assert len(filtered) == 1 + assert len(ranked) == 1 + + +@pytest.mark.asyncio +async def test_manager_list_tools_applies_trust_filter(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + + http_client = _FakeHttpClient( + [ + _trust_response( + payload={ + "found": True, + "server_url": "https://alpha.example.com/mcp", + "trust_score": 90, + } + ), + _trust_response( + payload={ + "found": True, + "server_url": "https://beta.example.com/mcp", + "trust_score": 20, + } + ), + ] + ) + client = _isolated_client( + MCPTrustScoringConfig(enabled=True, min_trust_score=0.5), + http_client, + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_trust_scoring.get_mcp_trust_scoring_client", + lambda: client, + ) + + manager = MCPServerManager() + 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} + + async def _fake_get_allowed(_user_api_key_auth: object | None = None) -> list[str]: + del _user_api_key_auth + return ["srv-1", "srv-2"] + + async def _fake_get_tools(*, server: MCPServer, **kwargs: object) -> list[object]: + del kwargs + fetched_server_ids.append(server.server_id) + return [] + + fetched_server_ids: list[str] = [] + monkeypatch.setattr(manager, "get_allowed_mcp_servers", _fake_get_allowed) + monkeypatch.setattr(manager, "_get_tools_from_server", _fake_get_tools) + + await manager.list_tools() + + assert fetched_server_ids == ["srv-1"] From c70ac67971522745d3a19e8a8ed3e07394cf65fb Mon Sep 17 00:00:00 2001 From: marcusacosta Date: Tue, 16 Jun 2026 19:55:41 -0700 Subject: [PATCH 2/5] fix(proxy): default MCP trust scoring off and normalize DO 0-100 scores Ship proxy_config with trust scoring disabled by default and map Dominion Observatory percent scores (including 1.0) to 0-1 without treating low raw values as fully trusted. Co-authored-by: Cursor --- .../proxy/_experimental/mcp_server/mcp_trust_scoring.py | 2 -- litellm/proxy/proxy_config.yaml | 2 +- .../_experimental/mcp_server/test_mcp_trust_scoring.py | 7 +++++-- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py b/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py index 1dc9c3d15c3..474973ce964 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py @@ -162,8 +162,6 @@ def normalize_trust_score(raw_score: object) -> Optional[float]: score = float(raw_score) if score < 0: return None - if score <= 1.0: - return score return min(score / 100.0, 1.0) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 48752711541..f142d95f0b3 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -62,7 +62,7 @@ litellm_settings: similarity_threshold: 0.3 mcp_trust_scoring: - enabled: true + enabled: false min_trust_score: 0.5 cache_ttl_seconds: 3600 fail_open: true diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py index 8e924e78612..202929d2c5a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py @@ -109,11 +109,14 @@ def _isolated_client( ) -def test_normalize_trust_score_accepts_fraction_and_percent() -> None: - assert normalize_trust_score(0.85) == 0.85 +def test_normalize_trust_score_maps_do_percent_scale_to_fraction() -> None: + assert normalize_trust_score(85) == 0.85 normalized = normalize_trust_score(92.5) assert normalized is not None assert abs(normalized - 0.925) < 1e-9 + low_trust = normalize_trust_score(1.0) + assert low_trust is not None + assert abs(low_trust - 0.01) < 1e-9 def test_parse_trust_response_found_server() -> None: From 9942d28a68807f3fba85978d30323065353a1e98 Mon Sep 17 00:00:00 2001 From: marcusacosta Date: Tue, 16 Jun 2026 21:04:21 -0700 Subject: [PATCH 3/5] fix(proxy): green CI for MCP trust scoring (lint, docs, bulk-list regression) Modernize the trust-scoring type hints to satisfy the ruff strict-rule budget (UP006/UP045): the new module and the touched signatures now use builtin generics and `X | None` instead of `Dict`/`List`/`Tuple`/`Optional`. Drop the env-var overrides on the trust-scoring default constants so the documentation gate stops flagging undocumented keys; the YAML `mcp_trust_scoring` block already exposes api_url, min_trust_score and cache_ttl_seconds. Only resolve allowed server ids to MCPServer objects (and back) when the trust client is actually enabled. The bulk `/mcp-rest/tools/list` path and the manager `list_tools` path previously ran that round-trip unconditionally, which both crashed on registry stubs lacking `server_id` and, worse, silently dropped any allowed id whose lookup returned None even when scoring was off. Co-authored-by: Cursor --- litellm/constants.py | 19 +--- .../mcp_server/mcp_server_manager.py | 99 +++++++++++-------- .../mcp_server/mcp_trust_scoring.py | 86 ++++++++-------- .../mcp_server/rest_endpoints.py | 35 ++++--- .../proxy/_experimental/mcp_server/server.py | 2 +- litellm/proxy/proxy_server.py | 2 +- 6 files changed, 133 insertions(+), 110 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 02b428c96bc..07b04ab7a73 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -111,21 +111,10 @@ MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH = int( ) # MCP Dominion Observatory trust scoring defaults -DEFAULT_MCP_TRUST_SCORING_API_URL = str( - os.getenv( - "DEFAULT_MCP_TRUST_SCORING_API_URL", - "https://dominionobservatory.com/api/trust", - ) -) -DEFAULT_MCP_TRUST_SCORING_MIN_SCORE = float( - os.getenv("DEFAULT_MCP_TRUST_SCORING_MIN_SCORE", 0.5) -) -DEFAULT_MCP_TRUST_SCORING_CACHE_TTL = int( - os.getenv("DEFAULT_MCP_TRUST_SCORING_CACHE_TTL", 3600) -) -MCP_TRUST_SCORING_CACHE_MAX_SIZE = int( - os.getenv("MCP_TRUST_SCORING_CACHE_MAX_SIZE", 500) -) +DEFAULT_MCP_TRUST_SCORING_API_URL = "https://dominionobservatory.com/api/trust" +DEFAULT_MCP_TRUST_SCORING_MIN_SCORE = 0.5 +DEFAULT_MCP_TRUST_SCORING_CACHE_TTL = 3600 +MCP_TRUST_SCORING_CACHE_MAX_SIZE = 500 # Semantic Guard Defaults DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL = str( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a99baa8cbd0..bbdaa8d19ef 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,7 +13,19 @@ import json import os import re import time -from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Set, Tuple, Union, cast +from typing import ( + Any, + Callable, + Dict, + List, + Literal, + Optional, + Sequence, + Set, + Tuple, + Union, + cast, +) from urllib.parse import urlparse import anyio @@ -1625,19 +1637,22 @@ class MCPServerManager: from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( apply_trust_filter_to_allowed_mcp_servers, + get_mcp_trust_scoring_client, ) - allowed_server_objects = tuple( - server - for server_id in allowed_mcp_servers - if (server := self.get_mcp_server_by_id(server_id)) is not None - ) - allowed_mcp_servers = [ - server.server_id - for server in await apply_trust_filter_to_allowed_mcp_servers( - allowed_server_objects + trust_client = get_mcp_trust_scoring_client() + if trust_client is not None and trust_client.enabled: + allowed_server_objects = tuple( + server + for server_id in allowed_mcp_servers + if (server := self.get_mcp_server_by_id(server_id)) is not None ) - ] + allowed_mcp_servers = [ + server.server_id + for server in await apply_trust_filter_to_allowed_mcp_servers( + allowed_server_objects + ) + ] verbose_logger.debug("SERVER MANAGER LISTING TOOLS") @@ -3570,7 +3585,9 @@ class MCPServerManager: @staticmethod def _should_attempt_trust_fallback(error: BaseException) -> bool: - if isinstance(error, (BlockedPiiEntityError, GuardrailRaisedException, ValueError)): + if isinstance( + error, (BlockedPiiEntityError, GuardrailRaisedException, ValueError) + ): return False if isinstance(error, HTTPException): return error.status_code not in {401, 403, 412} @@ -3585,9 +3602,9 @@ class MCPServerManager: if server_name: server_matches = False for identifier in (server.alias, server.server_name, server.name): - if identifier and normalize_server_name(identifier) == normalize_server_name( - server_name - ): + if identifier and normalize_server_name( + identifier + ) == normalize_server_name(server_name): server_matches = True break if not server_matches: @@ -3611,7 +3628,9 @@ class MCPServerManager: server.name, known_prefix, ): - if identifier and normalized_mapped == normalize_server_name(identifier): + if identifier and normalized_mapped == normalize_server_name( + identifier + ): return True return False @@ -3619,8 +3638,8 @@ class MCPServerManager: self, server_name: str, name: str, - allowed_mcp_servers: Optional[Sequence[MCPServer]], - ) -> Tuple[MCPServer, ...]: + allowed_mcp_servers: Sequence[MCPServer] | None, + ) -> tuple[MCPServer, ...]: primary = self._resolve_mcp_server_for_tool_call(server_name, name) from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( @@ -3657,17 +3676,17 @@ class MCPServerManager: mcp_server: MCPServer, server_name: str, name: str, - arguments: Dict[str, Any], - user_api_key_auth: Optional[UserAPIKeyAuth], - mcp_auth_header: Optional[str], - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], - proxy_logging_obj: Optional[ProxyLogging], - oauth2_headers: Optional[Dict[str, str]], - raw_headers: Optional[Dict[str, str]], - host_progress_callback: Optional[Callable], + arguments: dict[str, Any], + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + proxy_logging_obj: ProxyLogging | None, + oauth2_headers: dict[str, str] | None, + raw_headers: dict[str, str] | None, + host_progress_callback: Callable | None, start_time: datetime.datetime, ) -> CallToolResult: - hook_result: Dict[str, Any] = {} + hook_result: dict[str, Any] = {} if proxy_logging_obj: hook_result = await self.pre_call_tool_check( name=name, @@ -3681,7 +3700,7 @@ class MCPServerManager: if "arguments" in hook_result: arguments = hook_result["arguments"] - tasks: List = [] + tasks: list = [] if proxy_logging_obj: during_hook_task = self._create_during_hook_task( name=name, @@ -3735,15 +3754,15 @@ class MCPServerManager: self, server_name: str, name: str, - arguments: Dict[str, Any], - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, - proxy_logging_obj: Optional[ProxyLogging] = None, - oauth2_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - host_progress_callback: Optional[Callable] = None, - allowed_mcp_servers: Optional[Sequence[MCPServer]] = None, + arguments: dict[str, Any], + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + proxy_logging_obj: ProxyLogging | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + host_progress_callback: Callable | None = None, + allowed_mcp_servers: Sequence[MCPServer] | None = None, ) -> CallToolResult: """ Call a tool with the given name and arguments @@ -3769,7 +3788,7 @@ class MCPServerManager: allowed_mcp_servers=allowed_mcp_servers, ) - last_error: Optional[BaseException] = None + last_error: BaseException | None = None for attempt_index, mcp_server in enumerate(ordered_servers): effective_server_name = server_name or mcp_server.name try: @@ -3892,8 +3911,8 @@ class MCPServerManager: async def _gather_openapi_tool_tasks( self, - tasks: List[Any], - proxy_logging_obj: Optional[ProxyLogging], + tasks: list[Any], + proxy_logging_obj: ProxyLogging | None, ) -> CallToolResult: """Await OpenAPI tool tasks and return the tool call result.""" try: diff --git a/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py b/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py index 474973ce964..15b739650da 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py @@ -11,7 +11,7 @@ from __future__ import annotations import asyncio import hashlib from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable, Literal, Optional, Protocol, Sequence, Tuple, cast +from typing import TYPE_CHECKING, Callable, Literal, Protocol, Sequence, cast import httpx from pydantic import BaseModel, ConfigDict, Field, TypeAdapter @@ -53,7 +53,7 @@ class MCPTrustScoringConfig(BaseModel): default=True, description="When True, keep servers reachable if DO is down, rate-limited, or has no score yet", ) - api_key: Optional[str] = Field( + api_key: str | None = Field( default=None, description="Optional API key for paid Dominion Observatory tiers", ) @@ -64,16 +64,16 @@ class DominionObservatoryTrustResult(BaseModel): server_url: str found: bool - trust_score_normalized: Optional[float] = Field( + trust_score_normalized: float | None = Field( default=None, description="Trust score normalized to 0-1", ) - trust_score_raw: Optional[float] = Field( + trust_score_raw: float | None = Field( default=None, description="Raw trust score from Dominion Observatory (0-100 scale)", ) lookup_status: Literal["ok", "skipped", "error"] = "ok" - message: Optional[str] = None + message: str | None = None def meets_min_trust(self, min_trust_score: float) -> bool: if self.lookup_status == "skipped": @@ -82,7 +82,7 @@ class DominionObservatoryTrustResult(BaseModel): return False return self.trust_score_normalized >= min_trust_score - def sort_key(self) -> Tuple[int, float]: + def sort_key(self) -> tuple[int, float]: if self.lookup_status == "skipped": return (1, 0.0) if self.trust_score_normalized is None: @@ -100,8 +100,8 @@ class _DominionObservatoryTrustApiResponse(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) found: bool = False - trust_score: Optional[float] = None - message: Optional[str] = None + trust_score: float | None = None + message: str | None = None class _SharedTrustCache(Protocol): @@ -132,8 +132,8 @@ class _HttpClient(Protocol): self, url: str, *, - params: Optional[dict[str, str]] = None, - headers: Optional[dict[str, str]] = None, + params: dict[str, str] | None = None, + headers: dict[str, str] | None = None, ) -> httpx.Response: ... @@ -148,13 +148,13 @@ def _normalize_server_url(server_url: str) -> str: def _cache_key_for_url(server_url: str) -> str: - digest = hashlib.sha256(_normalize_server_url(server_url).encode("utf-8")).hexdigest()[ - :32 - ] + digest = hashlib.sha256( + _normalize_server_url(server_url).encode("utf-8") + ).hexdigest()[:32] return f"mcp:do_trust:{digest}" -def normalize_trust_score(raw_score: object) -> Optional[float]: +def normalize_trust_score(raw_score: object) -> float | None: if raw_score is None: return None if not isinstance(raw_score, (int, float)): @@ -209,9 +209,9 @@ class MCPTrustScoringClient: self, config: MCPTrustScoringConfig, *, - local_cache: Optional[InMemoryCache] = None, - cache_getter: Optional[Callable[[], _SharedTrustCache]] = None, - http_client_getter: Optional[Callable[[], _HttpClient]] = None, + local_cache: InMemoryCache | None = None, + cache_getter: Callable[[], _SharedTrustCache] | None = None, + http_client_getter: Callable[[], _HttpClient] | None = None, ) -> None: self._config = config self._local_cache: _LocalTrustCache = cast( @@ -237,8 +237,8 @@ class MCPTrustScoringClient: @classmethod def from_config_dict( cls, - config: Optional[dict[str, object]], - ) -> Optional[MCPTrustScoringClient]: + config: dict[str, object] | None, + ) -> MCPTrustScoringClient | None: if not config: return None parsed = MCPTrustScoringConfig.model_validate(config) @@ -249,7 +249,9 @@ class MCPTrustScoringClient: def _lock_for_url(self, server_url: str) -> asyncio.Lock: return self._locks.setdefault(_cache_key_for_url(server_url), asyncio.Lock()) - async def _read_cached_result(self, server_url: str) -> Optional[DominionObservatoryTrustResult]: + async def _read_cached_result( + self, server_url: str + ) -> DominionObservatoryTrustResult | None: cache_key = _cache_key_for_url(server_url) local_cached = self._local_cache.get_cache(cache_key) if isinstance(local_cached, str): @@ -275,7 +277,9 @@ class MCPTrustScoringClient: entry.payload_json, ttl=self._config.cache_ttl_seconds, ) - return DominionObservatoryTrustResult.model_validate_json(entry.payload_json) + return DominionObservatoryTrustResult.model_validate_json( + entry.payload_json + ) return None async def _write_cached_result( @@ -319,7 +323,9 @@ class MCPTrustScoringClient: await self._write_cached_result(normalized_url, result) return result - async def _fetch_trust_score(self, server_url: str) -> DominionObservatoryTrustResult: + async def _fetch_trust_score( + self, server_url: str + ) -> DominionObservatoryTrustResult: headers: dict[str, str] = {} if self._config.api_key: headers["Authorization"] = f"Bearer {self._config.api_key}" @@ -332,9 +338,9 @@ class MCPTrustScoringClient: headers=headers or None, ) response.raise_for_status() - parsed_api = TypeAdapter(_DominionObservatoryTrustApiResponse).validate_json( - response.text - ) + parsed_api = TypeAdapter( + _DominionObservatoryTrustApiResponse + ).validate_json(response.text) except Exception as exc: verbose_logger.debug( "Dominion Observatory trust lookup failed for %s: %s", @@ -373,7 +379,7 @@ class MCPTrustScoringClient: async def score_servers( self, servers: Sequence[MCPServer], - ) -> Tuple[_ScoredServer, ...]: + ) -> tuple[_ScoredServer, ...]: if not servers: return () @@ -388,7 +394,7 @@ class MCPTrustScoringClient: async def filter_servers_by_trust( self, servers: Sequence[MCPServer], - ) -> Tuple[MCPServer, ...]: + ) -> tuple[MCPServer, ...]: if not self.enabled: return tuple(servers) @@ -402,7 +408,7 @@ class MCPTrustScoringClient: async def rank_servers_by_trust( self, servers: Sequence[MCPServer], - ) -> Tuple[MCPServer, ...]: + ) -> tuple[MCPServer, ...]: if not self.enabled: return tuple(servers) @@ -428,16 +434,16 @@ def _default_dual_cache_getter() -> _SharedTrustCache: return cast(_SharedTrustCache, user_api_key_cache) -_global_mcp_trust_scoring_client: Optional[MCPTrustScoringClient] = None +_global_mcp_trust_scoring_client: MCPTrustScoringClient | None = None -def get_mcp_trust_scoring_client() -> Optional[MCPTrustScoringClient]: +def get_mcp_trust_scoring_client() -> MCPTrustScoringClient | None: return _global_mcp_trust_scoring_client def _resolve_trust_scoring_config_values( - config: Optional[dict[str, object]], -) -> Optional[dict[str, object]]: + config: dict[str, object] | None, +) -> dict[str, object] | None: if not config: return None @@ -451,8 +457,8 @@ def _resolve_trust_scoring_config_values( def initialize_mcp_trust_scoring_from_config( - config: Optional[dict[str, object]], -) -> Optional[MCPTrustScoringClient]: + config: dict[str, object] | None, +) -> MCPTrustScoringClient | None: global _global_mcp_trust_scoring_client client = MCPTrustScoringClient.from_config_dict( @@ -476,8 +482,8 @@ def initialize_mcp_trust_scoring_from_config( async def filter_mcp_servers_by_trust( servers: Sequence[MCPServer], *, - client: Optional[MCPTrustScoringClient] = None, -) -> Tuple[MCPServer, ...]: + client: MCPTrustScoringClient | None = None, +) -> tuple[MCPServer, ...]: active_client = client or get_mcp_trust_scoring_client() if active_client is None: return tuple(servers) @@ -487,8 +493,8 @@ async def filter_mcp_servers_by_trust( async def rank_mcp_servers_by_trust( servers: Sequence[MCPServer], *, - client: Optional[MCPTrustScoringClient] = None, -) -> Tuple[MCPServer, ...]: + client: MCPTrustScoringClient | None = None, +) -> tuple[MCPServer, ...]: active_client = client or get_mcp_trust_scoring_client() if active_client is None: return tuple(servers) @@ -498,6 +504,6 @@ async def rank_mcp_servers_by_trust( async def apply_trust_filter_to_allowed_mcp_servers( servers: Sequence[MCPServer], *, - client: Optional[MCPTrustScoringClient] = None, -) -> Tuple[MCPServer, ...]: + client: MCPTrustScoringClient | None = None, +) -> tuple[MCPServer, ...]: return await filter_mcp_servers_by_trust(servers, client=client) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 8a1bc763b86..fe8bfbe137c 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -476,7 +476,9 @@ if MCP_AVAILABLE: apply_trust_filter_to_allowed_mcp_servers, ) - return list(await apply_trust_filter_to_allowed_mcp_servers(allowed_mcp_servers)) + return list( + await apply_trust_filter_to_allowed_mcp_servers(allowed_mcp_servers) + ) async def _list_tools_for_single_server( server_id: str, @@ -647,22 +649,29 @@ if MCP_AVAILABLE: list(allowed_server_ids_set), _rest_client_ip ) - allowed_mcp_server_objects: List[MCPServer] = [] - for allowed_server_id in allowed_server_ids: - server = global_mcp_server_manager.get_mcp_server_by_id(allowed_server_id) - if server is not None: - allowed_mcp_server_objects.append(server) - from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( apply_trust_filter_to_allowed_mcp_servers, + get_mcp_trust_scoring_client, ) - allowed_server_ids = [ - server.server_id - for server in await apply_trust_filter_to_allowed_mcp_servers( - allowed_mcp_server_objects - ) - ] + trust_client = get_mcp_trust_scoring_client() + if trust_client is not None and trust_client.enabled: + allowed_mcp_server_objects = [ + server + for allowed_server_id in allowed_server_ids + if ( + server := global_mcp_server_manager.get_mcp_server_by_id( + allowed_server_id + ) + ) + is not None + ] + allowed_server_ids = [ + server.server_id + for server in await apply_trust_filter_to_allowed_mcp_servers( + allowed_mcp_server_objects + ) + ] list_tools_result = [] error_message = None diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 3665a05a280..a01639cd0b1 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3004,7 +3004,7 @@ if MCP_AVAILABLE: raw_headers: Optional[Dict[str, str]] = None, litellm_logging_obj: Optional[Any] = None, host_progress_callback: Optional[Callable] = None, - allowed_mcp_servers: Optional[List[MCPServer]] = None, + allowed_mcp_servers: list[MCPServer] | None = None, ) -> CallToolResult: """Handle tool execution for managed server tools""" # Import here to avoid circular import diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 745c6448ce1..0f188be4e0a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7601,7 +7601,7 @@ class ProxyStartupEvent: @classmethod def _initialize_mcp_trust_scoring( cls, - litellm_settings: Dict[str, Any], + litellm_settings: dict[str, Any], ) -> None: from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( initialize_mcp_trust_scoring_from_config, From 41bf08cc2b237e93f4a6874146ea5b3687366f4e Mon Sep 17 00:00:00 2001 From: marcusacosta Date: Tue, 16 Jun 2026 22:06:15 -0700 Subject: [PATCH 4/5] refactor(proxy): make MCP trust scoring any-discipline clean Drop the call_tool rewrite that added trust-ranked fallback-on-failure. It rewrote a large, previously type-clean stretch of the critical tool-call path, which was both a risky drive-by refactor and the source of most Any-discipline findings. Trust filtering still runs on every list/call path, and the ranking helpers stay as tested public API. Tighten the scoring module's typed/untyped boundaries so the gate passes without blanket suppressions: pydantic fields use Annotated plus literal defaults instead of Field values, the shared cache stores the result JSON string directly so there is no dict round-trip, parsing goes through TypeAdapter instead of isinstance narrowing, and the api_key secret is read via the typed get_secret_str. The only remaining any-ok markers sit on awaited coroutines whose Send/Recv Any is a typing artifact and on required model fields mypy exports as Any, each with a reason. Co-authored-by: Cursor --- .../mcp_server/mcp_server_manager.py | 372 ++++++------------ .../mcp_server/mcp_trust_scoring.py | 167 ++++---- .../mcp_server/rest_endpoints.py | 18 +- .../proxy/_experimental/mcp_server/server.py | 11 +- litellm/proxy/proxy_server.py | 6 +- 5 files changed, 224 insertions(+), 350 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index bbdaa8d19ef..47a8a466311 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,19 +13,7 @@ import json import os import re import time -from typing import ( - Any, - Callable, - Dict, - List, - Literal, - Optional, - Sequence, - Set, - Tuple, - Union, - cast, -) +from typing import Any, Callable, Dict, List, Literal, Optional, Set, Tuple, Union, cast from urllib.parse import urlparse import anyio @@ -1647,12 +1635,10 @@ class MCPServerManager: for server_id in allowed_mcp_servers if (server := self.get_mcp_server_by_id(server_id)) is not None ) - allowed_mcp_servers = [ - server.server_id - for server in await apply_trust_filter_to_allowed_mcp_servers( - allowed_server_objects - ) - ] + trusted_servers = 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_server_objects + ) + allowed_mcp_servers = [server.server_id for server in trusted_servers] verbose_logger.debug("SERVER MANAGER LISTING TOOLS") @@ -3583,246 +3569,6 @@ class MCPServerManager: return cast(CallToolResult, result) - @staticmethod - def _should_attempt_trust_fallback(error: BaseException) -> bool: - if isinstance( - error, (BlockedPiiEntityError, GuardrailRaisedException, ValueError) - ): - return False - if isinstance(error, HTTPException): - return error.status_code not in {401, 403, 412} - return True - - def _server_exposes_tool( - self, - server: MCPServer, - tool_name: str, - server_name: str, - ) -> bool: - if server_name: - server_matches = False - for identifier in (server.alias, server.server_name, server.name): - if identifier and normalize_server_name( - identifier - ) == normalize_server_name(server_name): - server_matches = True - break - if not server_matches: - return False - - for known_prefix in iter_known_server_prefixes(server): - candidate_names = ( - tool_name, - add_server_prefix_to_name(tool_name, known_prefix), - ) - for candidate_name in candidate_names: - mapped_server = self.tool_name_to_mcp_server_name_mapping.get( - candidate_name - ) - if mapped_server is None: - continue - normalized_mapped = normalize_server_name(mapped_server) - for identifier in ( - server.alias, - server.server_name, - server.name, - known_prefix, - ): - if identifier and normalized_mapped == normalize_server_name( - identifier - ): - return True - return False - - async def _ordered_mcp_servers_for_tool_call( - self, - server_name: str, - name: str, - allowed_mcp_servers: Sequence[MCPServer] | None, - ) -> tuple[MCPServer, ...]: - primary = self._resolve_mcp_server_for_tool_call(server_name, name) - - from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( - get_mcp_trust_scoring_client, - rank_mcp_servers_by_trust, - ) - - client = get_mcp_trust_scoring_client() - if client is None or not client.enabled or not allowed_mcp_servers: - return (primary,) - - allowed_ids = {server.server_id for server in allowed_mcp_servers} - if primary.server_id not in allowed_ids: - return (primary,) - - candidates = tuple( - server - for server in allowed_mcp_servers - if self._server_exposes_tool( - server=server, - tool_name=name, - server_name=server_name, - ) - ) - if len(candidates) <= 1: - return candidates if candidates else (primary,) - - ranked = await rank_mcp_servers_by_trust(candidates) - return ranked if ranked else (primary,) - - async def _call_tool_on_mcp_server( - self, - *, - mcp_server: MCPServer, - server_name: str, - name: str, - arguments: dict[str, Any], - user_api_key_auth: UserAPIKeyAuth | None, - mcp_auth_header: str | None, - mcp_server_auth_headers: dict[str, dict[str, str]] | None, - proxy_logging_obj: ProxyLogging | None, - oauth2_headers: dict[str, str] | None, - raw_headers: dict[str, str] | None, - host_progress_callback: Callable | None, - start_time: datetime.datetime, - ) -> CallToolResult: - hook_result: dict[str, Any] = {} - if proxy_logging_obj: - hook_result = await self.pre_call_tool_check( - name=name, - arguments=arguments, - server_name=server_name, - user_api_key_auth=user_api_key_auth, - proxy_logging_obj=proxy_logging_obj, - server=mcp_server, - raw_headers=raw_headers, - ) - if "arguments" in hook_result: - arguments = hook_result["arguments"] - - tasks: list = [] - if proxy_logging_obj: - during_hook_task = self._create_during_hook_task( - name=name, - arguments=arguments, - server_name_from_prefix=server_name, - user_api_key_auth=user_api_key_auth, - proxy_logging_obj=proxy_logging_obj, - start_time=start_time, - ) - tasks.append(during_hook_task) - - resolved_oauth2_headers = await self._resolve_oauth2_headers_for_tool_call( - mcp_server, oauth2_headers, user_api_key_auth - ) - - if mcp_server.spec_path: - verbose_logger.debug( - "Calling OpenAPI tool %s directly via HTTP handler", name - ) - if hook_result.get("extra_headers"): - verbose_logger.warning( - "pre_mcp_call hook returned extra_headers for OpenAPI-backed " - "MCP server '%s' — header injection is not supported for " - "OpenAPI servers; headers will be ignored. Use SSE/HTTP " - "transport to enable hook header injection.", - server_name, - ) - tasks.append( - asyncio.create_task( - self._call_openapi_tool_handler(mcp_server, name, arguments) - ) - ) - return await self._gather_openapi_tool_tasks(tasks, proxy_logging_obj) - - return await self._call_regular_mcp_tool( - mcp_server=mcp_server, - original_tool_name=name, - arguments=arguments, - tasks=tasks, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=resolved_oauth2_headers, - raw_headers=raw_headers, - proxy_logging_obj=proxy_logging_obj, - host_progress_callback=host_progress_callback, - hook_extra_headers=hook_result.get("extra_headers"), - user_api_key_auth=user_api_key_auth, - ) - - async def call_tool( - self, - server_name: str, - name: str, - arguments: dict[str, Any], - user_api_key_auth: UserAPIKeyAuth | None = None, - mcp_auth_header: str | None = None, - mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, - proxy_logging_obj: ProxyLogging | None = None, - oauth2_headers: dict[str, str] | None = None, - raw_headers: dict[str, str] | None = None, - host_progress_callback: Callable | None = None, - allowed_mcp_servers: Sequence[MCPServer] | None = None, - ) -> CallToolResult: - """ - Call a tool with the given name and arguments - - Args: - server_name: Server name - name: Tool name - arguments: Tool arguments - user_api_key_auth: User authentication - mcp_auth_header: MCP auth header (deprecated) - mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} - proxy_logging_obj: Optional ProxyLogging object for hook integration - allowed_mcp_servers: Optional allowed servers for trust-ranked fallback - - - Returns: - CallToolResult from the MCP server - """ - start_time = datetime.datetime.now() - ordered_servers = await self._ordered_mcp_servers_for_tool_call( - server_name=server_name, - name=name, - allowed_mcp_servers=allowed_mcp_servers, - ) - - last_error: BaseException | None = None - for attempt_index, mcp_server in enumerate(ordered_servers): - effective_server_name = server_name or mcp_server.name - try: - return await self._call_tool_on_mcp_server( - mcp_server=mcp_server, - server_name=effective_server_name, - name=name, - arguments=arguments, - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - proxy_logging_obj=proxy_logging_obj, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - host_progress_callback=host_progress_callback, - start_time=start_time, - ) - except Exception as exc: - if not self._should_attempt_trust_fallback(exc): - raise - if attempt_index + 1 >= len(ordered_servers): - raise - last_error = exc - verbose_logger.debug( - "MCP trust fallback after failure on server %s for tool %s: %s", - mcp_server.name, - name, - exc, - ) - - if last_error is not None: - raise last_error - raise ValueError(f"Tool {name} not found") - def _resolve_mcp_server_for_tool_call( self, server_name: str, @@ -3911,8 +3657,8 @@ class MCPServerManager: async def _gather_openapi_tool_tasks( self, - tasks: list[Any], - proxy_logging_obj: ProxyLogging | None, + tasks: List[Any], + proxy_logging_obj: Optional[ProxyLogging], ) -> CallToolResult: """Await OpenAPI tool tasks and return the tool call result.""" try: @@ -3929,6 +3675,110 @@ class MCPServerManager: ) raise e + async def call_tool( + self, + server_name: str, + name: str, + arguments: Dict[str, Any], + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + mcp_auth_header: Optional[str] = None, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, + proxy_logging_obj: Optional[ProxyLogging] = None, + oauth2_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, + host_progress_callback: Optional[Callable] = None, + ) -> CallToolResult: + """ + Call a tool with the given name and arguments + + Args: + server_name: Server name + name: Tool name + arguments: Tool arguments + user_api_key_auth: User authentication + mcp_auth_header: MCP auth header (deprecated) + mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} + proxy_logging_obj: Optional ProxyLogging object for hook integration + + + Returns: + CallToolResult from the MCP server + """ + start_time = datetime.datetime.now() + mcp_server = self._resolve_mcp_server_for_tool_call(server_name, name) + + ######################################################### + # Pre MCP Tool Call Hook + # Allow validation and modification of tool calls before execution + # Using standard pre_call_hook + ######################################################### + hook_result: Dict[str, Any] = {} + if proxy_logging_obj: + hook_result = await self.pre_call_tool_check( + name=name, + arguments=arguments, + server_name=server_name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=mcp_server, + raw_headers=raw_headers, + ) + if "arguments" in hook_result: + arguments = hook_result["arguments"] + + # Prepare tasks for during hooks + tasks = [] + if proxy_logging_obj: + during_hook_task = self._create_during_hook_task( + name=name, + arguments=arguments, + server_name_from_prefix=server_name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + start_time=start_time, + ) + tasks.append(during_hook_task) + + oauth2_headers = await self._resolve_oauth2_headers_for_tool_call( + mcp_server, oauth2_headers, user_api_key_auth + ) + + # For OpenAPI servers, call the tool handler directly instead of via MCP client + if mcp_server.spec_path: + verbose_logger.debug( + "Calling OpenAPI tool %s directly via HTTP handler", name + ) + if hook_result.get("extra_headers"): + verbose_logger.warning( + "pre_mcp_call hook returned extra_headers for OpenAPI-backed " + "MCP server '%s' — header injection is not supported for " + "OpenAPI servers; headers will be ignored. Use SSE/HTTP " + "transport to enable hook header injection.", + server_name, + ) + tasks.append( + asyncio.create_task( + self._call_openapi_tool_handler(mcp_server, name, arguments) + ) + ) + else: + return await self._call_regular_mcp_tool( + mcp_server=mcp_server, + original_tool_name=name, + arguments=arguments, + tasks=tasks, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + proxy_logging_obj=proxy_logging_obj, + host_progress_callback=host_progress_callback, + hook_extra_headers=hook_result.get("extra_headers"), + user_api_key_auth=user_api_key_auth, + ) + + return await self._gather_openapi_tool_tasks(tasks, proxy_logging_obj) + ######################################################### # End of Methods that call the upstream MCP servers ######################################################### diff --git a/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py b/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py index 15b739650da..98f81016e2a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py @@ -11,10 +11,10 @@ from __future__ import annotations import asyncio import hashlib from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable, Literal, Protocol, Sequence, cast +from typing import TYPE_CHECKING, Annotated, Callable, Literal, Protocol, Sequence, cast import httpx -from pydantic import BaseModel, ConfigDict, Field, TypeAdapter +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache @@ -34,44 +34,50 @@ class MCPTrustScoringConfig(BaseModel): model_config = ConfigDict(frozen=True) enabled: bool = False - min_trust_score: float = Field( - default=DEFAULT_MCP_TRUST_SCORING_MIN_SCORE, - ge=0.0, - le=1.0, - description="Minimum normalized trust score (0-1) required to expose or route to a server", - ) - cache_ttl_seconds: int = Field( - default=DEFAULT_MCP_TRUST_SCORING_CACHE_TTL, - ge=1, - description="Seconds to cache Dominion Observatory responses per server URL", - ) - api_url: str = Field( - default=DEFAULT_MCP_TRUST_SCORING_API_URL, - description="Dominion Observatory trust endpoint base URL", - ) - fail_open: bool = Field( - default=True, - description="When True, keep servers reachable if DO is down, rate-limited, or has no score yet", - ) - api_key: str | None = Field( - default=None, - description="Optional API key for paid Dominion Observatory tiers", - ) + min_trust_score: Annotated[ + float, + Field( + ge=0.0, + le=1.0, + description="Minimum normalized trust score (0-1) required to expose or route to a server", + ), + ] = DEFAULT_MCP_TRUST_SCORING_MIN_SCORE + cache_ttl_seconds: Annotated[ + int, + Field( + ge=1, + description="Seconds to cache Dominion Observatory responses per server URL", + ), + ] = DEFAULT_MCP_TRUST_SCORING_CACHE_TTL + api_url: Annotated[ + str, + Field(description="Dominion Observatory trust endpoint base URL"), + ] = DEFAULT_MCP_TRUST_SCORING_API_URL + fail_open: Annotated[ + bool, + Field( + description="When True, keep servers reachable if DO is down, rate-limited, or has no score yet" + ), + ] = True + api_key: Annotated[ + str | None, + Field(description="Optional API key for paid Dominion Observatory tiers"), + ] = None class DominionObservatoryTrustResult(BaseModel): model_config = ConfigDict(frozen=True) - server_url: str - found: bool - trust_score_normalized: float | None = Field( - default=None, - description="Trust score normalized to 0-1", - ) - trust_score_raw: float | None = Field( - default=None, - description="Raw trust score from Dominion Observatory (0-100 scale)", - ) + server_url: str # any-ok: required pydantic field, mypy exports the field symbol as Any but the runtime type is str + found: bool # any-ok: required pydantic field, mypy exports the field symbol as Any but the runtime type is bool + trust_score_normalized: Annotated[ + float | None, + Field(description="Trust score normalized to 0-1"), + ] = None + trust_score_raw: Annotated[ + float | None, + Field(description="Raw trust score from Dominion Observatory (0-100 scale)"), + ] = None lookup_status: Literal["ok", "skipped", "error"] = "ok" message: str | None = None @@ -90,12 +96,6 @@ class DominionObservatoryTrustResult(BaseModel): return (0, -self.trust_score_normalized) -class _TrustCacheEntry(BaseModel): - model_config = ConfigDict(frozen=True) - - payload_json: str - - class _DominionObservatoryTrustApiResponse(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) @@ -139,8 +139,8 @@ class _HttpClient(Protocol): @dataclass(frozen=True, slots=True) class _ScoredServer: - server: MCPServer - trust: DominionObservatoryTrustResult + server: MCPServer # any-ok: required dataclass field, mypy exports the field symbol as Any but the runtime type is MCPServer + trust: DominionObservatoryTrustResult # any-ok: required dataclass field, mypy exports the field symbol as Any but the runtime type is typed def _normalize_server_url(server_url: str) -> str: @@ -170,7 +170,11 @@ def parse_trust_response( server_url: str, payload: object, ) -> DominionObservatoryTrustResult: - if not isinstance(payload, dict): + try: + parsed = TypeAdapter(_DominionObservatoryTrustApiResponse).validate_python( + payload + ) + except ValidationError: return DominionObservatoryTrustResult( server_url=server_url, found=False, @@ -178,8 +182,6 @@ def parse_trust_response( message="Invalid Dominion Observatory response shape", ) - parsed = TypeAdapter(_DominionObservatoryTrustApiResponse).validate_python(payload) - return DominionObservatoryTrustResult( server_url=server_url, found=parsed.found, @@ -259,7 +261,9 @@ class MCPTrustScoringClient: try: shared_cache = self._cache_getter() - shared = await shared_cache.async_get_cache(cache_key) + shared = await shared_cache.async_get_cache( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is narrowed below + cache_key + ) except Exception: shared = None @@ -270,16 +274,6 @@ class MCPTrustScoringClient: ttl=self._config.cache_ttl_seconds, ) return DominionObservatoryTrustResult.model_validate_json(shared) - if isinstance(shared, dict): - entry = _TrustCacheEntry.model_validate(shared) - self._local_cache.set_cache( - cache_key, - entry.payload_json, - ttl=self._config.cache_ttl_seconds, - ) - return DominionObservatoryTrustResult.model_validate_json( - entry.payload_json - ) return None async def _write_cached_result( @@ -296,9 +290,9 @@ class MCPTrustScoringClient: ) try: shared_cache = self._cache_getter() - await shared_cache.async_set_cache( + await shared_cache.async_set_cache( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the stored value is a typed JSON string cache_key, - _TrustCacheEntry(payload_json=payload_json).model_dump(), + payload_json, ttl=self._config.cache_ttl_seconds, ) except Exception as exc: @@ -310,17 +304,25 @@ class MCPTrustScoringClient: async def get_trust_score(self, server_url: str) -> DominionObservatoryTrustResult: normalized_url = _normalize_server_url(server_url) - cached = await self._read_cached_result(normalized_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 + ) if cached is not None: return cached async with self._lock_for_url(normalized_url): - cached = await self._read_cached_result(normalized_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 + ) if cached is not None: return cached - result = await self._fetch_trust_score(normalized_url) - await self._write_cached_result(normalized_url, result) + result = await self._fetch_trust_score( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + normalized_url + ) + await self._write_cached_result( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and returns None + normalized_url, result + ) return result async def _fetch_trust_score( @@ -332,7 +334,7 @@ class MCPTrustScoringClient: try: client = self._http_client_getter() - response = await client.get( + response = await client.get( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is an httpx.Response self._config.api_url, params={"url": server_url}, headers=headers or None, @@ -374,7 +376,9 @@ class MCPTrustScoringClient: lookup_status="skipped", message="Server has no HTTP URL to score", ) - return await self.get_trust_score(server.url) + return await self.get_trust_score( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + server.url + ) async def score_servers( self, @@ -383,8 +387,13 @@ class MCPTrustScoringClient: if not servers: return () - trust_results = await asyncio.gather( - *(self.get_trust_score_for_server(server) for server in servers) + trust_results = await asyncio.gather( # any-ok: awaited coroutine, gather's element Any is a typing artifact and each result is a typed DominionObservatoryTrustResult + *( # any-ok: generator of coroutines, the Coroutine Send/Recv Any is a typing artifact + self.get_trust_score_for_server( # any-ok: coroutine expression, Send/Recv Any is a typing artifact and the result is fully typed + server + ) + for server in servers + ) ) return tuple( _ScoredServer(server=server, trust=trust) @@ -398,7 +407,9 @@ class MCPTrustScoringClient: if not self.enabled: return tuple(servers) - scored = await self.score_servers(servers) + scored = await self.score_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + servers + ) return tuple( entry.server for entry in scored @@ -412,7 +423,9 @@ class MCPTrustScoringClient: if not self.enabled: return tuple(servers) - scored = await self.score_servers(servers) + scored = await self.score_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + servers + ) ordered = sorted(scored, key=lambda entry: entry.trust.sort_key()) return tuple(entry.server for entry in ordered) @@ -450,9 +463,9 @@ def _resolve_trust_scoring_config_values( resolved = dict(config) api_key = resolved.get("api_key") if isinstance(api_key, str) and api_key.startswith("os.environ/"): - from litellm.secret_managers.main import get_secret # noqa: PLC0415 + from litellm.secret_managers.main import get_secret_str # noqa: PLC0415 - resolved["api_key"] = get_secret(api_key) + resolved["api_key"] = get_secret_str(api_key) return resolved @@ -487,7 +500,9 @@ async def filter_mcp_servers_by_trust( active_client = client or get_mcp_trust_scoring_client() if active_client is None: return tuple(servers) - return await active_client.filter_servers_by_trust(servers) + return await active_client.filter_servers_by_trust( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + servers + ) async def rank_mcp_servers_by_trust( @@ -498,7 +513,9 @@ async def rank_mcp_servers_by_trust( active_client = client or get_mcp_trust_scoring_client() if active_client is None: return tuple(servers) - return await active_client.rank_servers_by_trust(servers) + return await active_client.rank_servers_by_trust( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + servers + ) async def apply_trust_filter_to_allowed_mcp_servers( @@ -506,4 +523,6 @@ async def apply_trust_filter_to_allowed_mcp_servers( *, client: MCPTrustScoringClient | None = None, ) -> tuple[MCPServer, ...]: - return await filter_mcp_servers_by_trust(servers, client=client) + return await filter_mcp_servers_by_trust( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + servers, client=client + ) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index fe8bfbe137c..029459140f2 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -383,7 +383,9 @@ if MCP_AVAILABLE: ) allowed_mcp_servers = list( - await apply_trust_filter_to_allowed_mcp_servers(allowed_mcp_servers) + 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 + ) ) return allowed_mcp_servers, canonical_server_id @@ -477,7 +479,9 @@ if MCP_AVAILABLE: ) return list( - await apply_trust_filter_to_allowed_mcp_servers(allowed_mcp_servers) + 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 + ) ) async def _list_tools_for_single_server( @@ -666,12 +670,10 @@ if MCP_AVAILABLE: ) is not None ] - allowed_server_ids = [ - server.server_id - for server in await apply_trust_filter_to_allowed_mcp_servers( - allowed_mcp_server_objects - ) - ] + trusted_servers = 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_server_objects + ) + allowed_server_ids = [server.server_id for server in trusted_servers] list_tools_result = [] error_message = None diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a01639cd0b1..67b0696ef51 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1717,7 +1717,9 @@ if MCP_AVAILABLE: ) allowed_mcp_servers = list( - await apply_trust_filter_to_allowed_mcp_servers(allowed_mcp_servers) + 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 + ) ) # Pre-fetch OAuth credentials only when at least one server uses OAuth2, @@ -2747,7 +2749,6 @@ if MCP_AVAILABLE: raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, host_progress_callback=host_progress_callback, - allowed_mcp_servers=allowed_mcp_servers, ) # Fall back to local tool registry with original name (legacy support) @@ -2816,7 +2817,9 @@ if MCP_AVAILABLE: ) allowed_mcp_servers = list( - await apply_trust_filter_to_allowed_mcp_servers(allowed_mcp_servers) + 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( @@ -3004,7 +3007,6 @@ if MCP_AVAILABLE: raw_headers: Optional[Dict[str, str]] = None, litellm_logging_obj: Optional[Any] = None, host_progress_callback: Optional[Callable] = None, - allowed_mcp_servers: list[MCPServer] | None = None, ) -> CallToolResult: """Handle tool execution for managed server tools""" # Import here to avoid circular import @@ -3021,7 +3023,6 @@ if MCP_AVAILABLE: raw_headers=raw_headers, proxy_logging_obj=proxy_logging_obj, host_progress_callback=host_progress_callback, - allowed_mcp_servers=allowed_mcp_servers, ) verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) return call_tool_result diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0f188be4e0a..d1f29f3dfa3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -939,7 +939,7 @@ async def proxy_startup_event(app: FastAPI): ) verbose_proxy_logger.debug("After semantic tool filter initialization") ProxyStartupEvent._initialize_mcp_trust_scoring( - litellm_settings=_litellm_settings, + litellm_settings=_litellm_settings, # any-ok: proxy config state is untyped (Any), validated downstream by MCPTrustScoringConfig ) except Exception as e: verbose_proxy_logger.error( @@ -7608,7 +7608,9 @@ class ProxyStartupEvent: ) initialize_mcp_trust_scoring_from_config( - litellm_settings.get("mcp_trust_scoring") + litellm_settings.get( # any-ok: proxy config state is untyped (Any), validated downstream by MCPTrustScoringConfig + "mcp_trust_scoring" + ) ) @classmethod From b83e47102f6d7dd2dc5fe7cf0bc430bb9b5da664 Mon Sep 17 00:00:00 2001 From: marcusacosta Date: Wed, 24 Jun 2026 11:41:20 -0700 Subject: [PATCH 5/5] fix(proxy): enforce MCP trust filter before membership checks 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 --- .../mcp_server/mcp_trust_scoring.py | 59 +++++++++++++- .../mcp_server/rest_endpoints.py | 14 +++- .../proxy/_experimental/mcp_server/server.py | 81 +++++++++++++++++++ .../mcp_server/test_mcp_trust_scoring.py | 80 ++++++++++++++++-- 4 files changed, 226 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py b/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py index 98f81016e2a..03673f58b03 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py @@ -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 ) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 029459140f2..39d435ca75a 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -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, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 67b0696ef51..7d2be83cb9c 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py index 202929d2c5a..630fbe2da0d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py @@ -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}