mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
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 <cursoragent@cursor.com>
This commit is contained in:
marcusacosta
parent
ce4111b800
commit
ab3744a8cc
8 changed files with 1223 additions and 110 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
#########################################################
|
||||
|
|
|
|||
505
litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py
Normal file
505
litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
Loading…
Add table
Reference in a new issue