From ab3744a8ccde7ae0633bde75c3b461008e339dd2 Mon Sep 17 00:00:00 2001 From: marcusacosta Date: Tue, 16 Jun 2026 19:23:55 -0700 Subject: [PATCH] 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"]