diff --git a/litellm/constants.py b/litellm/constants.py index a9edf135731..8f96296dc89 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -91,6 +91,12 @@ DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD = float( ) 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 = "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( 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 82b820d8cd9..00d580a9f41 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2662,6 +2662,23 @@ 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, + get_mcp_trust_scoring_client, + ) + + 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 + ) + 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") async def _fetch_server_tools(server_id: str) -> list[MCPTool]: 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..03673f58b03 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py @@ -0,0 +1,585 @@ +""" +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, 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 + +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 + +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) + + enabled: bool = False + 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 # 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 + + 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 _DominionObservatoryTrustApiResponse(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + found: bool = False + trust_score: float | None = None + message: str | None = 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: dict[str, str] | None = None, + headers: dict[str, str] | None = None, + ) -> httpx.Response: ... + + +@dataclass(frozen=True, slots=True) +class _ScoredServer: + 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: + 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") + ).hexdigest()[:32] + return f"mcp:do_trust:{digest}" + + +def normalize_trust_score(raw_score: object) -> float | None: + 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 + return min(score / 100.0, 1.0) + + +def parse_trust_response( + *, + server_url: str, + payload: object, +) -> DominionObservatoryTrustResult: + try: + parsed = TypeAdapter(_DominionObservatoryTrustApiResponse).validate_python( + payload + ) + except ValidationError: + return DominionObservatoryTrustResult( + server_url=server_url, + found=False, + lookup_status="error", + message="Invalid Dominion Observatory response shape", + ) + + 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: 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( + _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: dict[str, object] | None, + ) -> MCPTrustScoringClient | None: + 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 + ) -> DominionObservatoryTrustResult | None: + 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( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is narrowed below + 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) + 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( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the stored value is a typed JSON string + cache_key, + payload_json, + 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: + 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 + ) + if cached is not None: + return cached + + async with self._lock_for_url(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( # 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( + 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( # 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, + ) + 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( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + server.url + ) + + async def score_servers( + self, + servers: Sequence[MCPServer], + ) -> tuple[_ScoredServer, ...]: + if not servers: + return () + + 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) + 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( # 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 + 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( # 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) + + +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: MCPTrustScoringClient | None = None + + +def get_mcp_trust_scoring_client() -> MCPTrustScoringClient | None: + return _global_mcp_trust_scoring_client + + +def _resolve_trust_scoring_config_values( + config: dict[str, object] | None, +) -> dict[str, object] | None: + 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_str # noqa: PLC0415 + + resolved["api_key"] = get_secret_str(api_key) + return resolved + + +def initialize_mcp_trust_scoring_from_config( + config: dict[str, object] | None, +) -> MCPTrustScoringClient | None: + 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: MCPTrustScoringClient | None = 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( # 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( + servers: Sequence[MCPServer], + *, + client: MCPTrustScoringClient | None = 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( # 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( + servers: Sequence[MCPServer], + *, + client: MCPTrustScoringClient | None = None, +) -> tuple[MCPServer, ...]: + 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 26e4176e09b..6e55fba2835 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -477,6 +477,21 @@ 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, + 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 + ) + ) + assert_requested_server_passes_trust_filter( + filtered_servers=allowed_mcp_servers, + server_id=canonical_server_id, + ) + return allowed_mcp_servers, canonical_server_id async def _get_tools_for_single_server( @@ -550,7 +565,23 @@ 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, + assert_requested_server_passes_trust_filter, + ) + + 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, @@ -770,6 +801,28 @@ if MCP_AVAILABLE: list(allowed_server_ids_set), _rest_client_ip ) + from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( + apply_trust_filter_to_allowed_mcp_servers, + get_mcp_trust_scoring_client, + ) + + 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 + ] + 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 9a1fffd5a67..7aae23ff7bd 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1955,6 +1955,16 @@ if MCP_AVAILABLE: client_ip=client_ip, ) + 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 + ) + ) + # 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(getattr(s, "auth_type", None) == MCPAuth.oauth2 for s in allowed_mcp_servers) @@ -2154,6 +2164,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: @@ -2208,6 +2228,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: @@ -2259,6 +2289,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: @@ -2653,6 +2693,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") @@ -2746,6 +2797,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 {}).get( "mcp_server_cost_info" @@ -3030,6 +3087,21 @@ 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( # 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="No MCP servers meet the configured trust score threshold.", + ) + # Delegate to execute_mcp_tool for execution response = await execute_mcp_tool( name=name, @@ -3100,6 +3172,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) @@ -3150,6 +3239,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/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index a094eb84bf3..f142d95f0b3 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: false + 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 c44f5602c95..e693a102fd2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1029,8 +1029,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() @@ -1041,8 +1040,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, # any-ok: proxy config state is untyped (Any), validated downstream by MCPTrustScoringConfig + ) 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( @@ -7709,6 +7713,21 @@ class ProxyStartupEvent: return _build_redis_usage_cache_from_environment() + @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( # any-ok: proxy config state is untyped (Any), validated downstream by MCPTrustScoringConfig + "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..630fbe2da0d --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py @@ -0,0 +1,437 @@ +""" +Tests for Dominion Observatory MCP trust scoring. +""" + +import httpx +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 +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_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_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", + 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"]