diff --git a/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py b/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py index 98f81016e2a..03673f58b03 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_trust_scoring.py @@ -12,6 +12,7 @@ import asyncio import hashlib from dataclasses import dataclass from typing import TYPE_CHECKING, Annotated, Callable, Literal, Protocol, Sequence, cast +from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import httpx from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError @@ -29,6 +30,25 @@ from litellm.types.llms.custom_http import httpxSpecialProvider if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer +TRUST_SCORE_THRESHOLD_DENIED = ( + "Server does not meet the configured trust score threshold." +) + +_SENSITIVE_QUERY_PARAM_NAMES = frozenset( + { + "access_token", + "api_key", + "apikey", + "auth", + "credential", + "credentials", + "key", + "password", + "secret", + "token", + } +) + class MCPTrustScoringConfig(BaseModel): model_config = ConfigDict(frozen=True) @@ -147,6 +167,42 @@ def _normalize_server_url(server_url: str) -> str: return server_url.rstrip("/") +def sanitize_url_for_trust_lookup(server_url: str) -> str: + parsed = urlparse(server_url) + hostname = parsed.hostname + if not hostname: + return _normalize_server_url(server_url) + + port_suffix = f":{parsed.port}" if parsed.port else "" + netloc = f"{hostname}{port_suffix}" + filtered_query = tuple( + (key, value) + for key, value in parse_qsl(parsed.query, keep_blank_values=True) + if key.lower() not in _SENSITIVE_QUERY_PARAM_NAMES + ) + sanitized = parsed._replace( + netloc=netloc, + query=urlencode(filtered_query), + fragment="", + ) + return urlunparse(sanitized) + + +def assert_requested_server_passes_trust_filter( + *, + filtered_servers: Sequence[MCPServer], + server_id: str, +) -> None: + from fastapi import HTTPException + + trusted_server_ids = {server.server_id for server in filtered_servers} + if server_id not in trusted_server_ids: + raise HTTPException( + status_code=403, + detail=TRUST_SCORE_THRESHOLD_DENIED, + ) + + def _cache_key_for_url(server_url: str) -> str: digest = hashlib.sha256( _normalize_server_url(server_url).encode("utf-8") @@ -303,7 +359,8 @@ class MCPTrustScoringClient: ) async def get_trust_score(self, server_url: str) -> DominionObservatoryTrustResult: - normalized_url = _normalize_server_url(server_url) + sanitized_url = sanitize_url_for_trust_lookup(server_url) + normalized_url = _normalize_server_url(sanitized_url) cached = await self._read_cached_result( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed normalized_url ) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 029459140f2..39d435ca75a 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -380,6 +380,7 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( apply_trust_filter_to_allowed_mcp_servers, + assert_requested_server_passes_trust_filter, ) allowed_mcp_servers = list( @@ -387,6 +388,10 @@ if MCP_AVAILABLE: allowed_mcp_servers ) ) + assert_requested_server_passes_trust_filter( + filtered_servers=allowed_mcp_servers, + server_id=canonical_server_id, + ) return allowed_mcp_servers, canonical_server_id @@ -476,13 +481,20 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( apply_trust_filter_to_allowed_mcp_servers, + assert_requested_server_passes_trust_filter, ) - return list( + filtered_servers = list( await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed allowed_mcp_servers ) ) + assert_requested_server_passes_trust_filter( + filtered_servers=filtered_servers, + server_id=server_id, + ) + + return filtered_servers async def _list_tools_for_single_server( server_id: str, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 67b0696ef51..7d2be83cb9c 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1917,6 +1917,16 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, ) + from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( + apply_trust_filter_to_allowed_mcp_servers, + ) + + allowed_mcp_servers = list( + await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + allowed_mcp_servers + ) + ) + # Get prompts from each allowed server all_prompts = [] for server in allowed_mcp_servers: @@ -1976,6 +1986,16 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, ) + from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( + apply_trust_filter_to_allowed_mcp_servers, + ) + + allowed_mcp_servers = list( + await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + allowed_mcp_servers + ) + ) + all_resources: List[Resource] = [] for server in allowed_mcp_servers: if server is None: @@ -2032,6 +2052,16 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, ) + from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( + apply_trust_filter_to_allowed_mcp_servers, + ) + + allowed_mcp_servers = list( + await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + allowed_mcp_servers + ) + ) + all_resource_templates: List[ResourceTemplate] = [] for server in allowed_mcp_servers: if server is None: @@ -2496,6 +2526,17 @@ if MCP_AVAILABLE: Returns: CallToolResult: Tool execution result """ + from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( + apply_trust_filter_to_allowed_mcp_servers, + assert_requested_server_passes_trust_filter, + ) + + allowed_mcp_servers = list( + await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + allowed_mcp_servers + ) + ) + # Track resolved MCP server for both permission checks and dispatch mcp_server: Optional[MCPServer] = None requested_server_id: Optional[str] = kwargs.get("requested_server_id") @@ -2602,6 +2643,12 @@ if MCP_AVAILABLE: if mcp_server is None: mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) + if mcp_server is not None: + assert_requested_server_passes_trust_filter( + filtered_servers=allowed_mcp_servers, + server_id=mcp_server.server_id, + ) + if mcp_server: standard_logging_mcp_tool_call["mcp_server_cost_info"] = ( mcp_server.mcp_info or {} @@ -2893,6 +2940,23 @@ if MCP_AVAILABLE: detail="User not allowed to get this prompt.", ) + from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( + apply_trust_filter_to_allowed_mcp_servers, + TRUST_SCORE_THRESHOLD_DENIED, + ) + + allowed_mcp_servers = list( + await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + allowed_mcp_servers + ) + ) + + if not allowed_mcp_servers: + raise HTTPException( + status_code=403, + detail=TRUST_SCORE_THRESHOLD_DENIED, + ) + # Extract server name from prefixed prompt name original_prompt_name, server_name = split_server_prefix_from_name(name) @@ -2943,6 +3007,23 @@ if MCP_AVAILABLE: detail="User not allowed to read this resource.", ) + from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( + apply_trust_filter_to_allowed_mcp_servers, + TRUST_SCORE_THRESHOLD_DENIED, + ) + + allowed_mcp_servers = list( + await apply_trust_filter_to_allowed_mcp_servers( # any-ok: awaited coroutine, Send/Recv Any is a typing artifact and the result is fully typed + allowed_mcp_servers + ) + ) + + if not allowed_mcp_servers: + raise HTTPException( + status_code=403, + detail=TRUST_SCORE_THRESHOLD_DENIED, + ) + if len(allowed_mcp_servers) != 1: raise HTTPException( status_code=400, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py index 202929d2c5a..630fbe2da0d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_trust_scoring.py @@ -8,11 +8,13 @@ import pytest from litellm.proxy._experimental.mcp_server.mcp_trust_scoring import ( MCPTrustScoringClient, MCPTrustScoringConfig, + assert_requested_server_passes_trust_filter, filter_mcp_servers_by_trust, initialize_mcp_trust_scoring_from_config, normalize_trust_score, parse_trust_response, rank_mcp_servers_by_trust, + sanitize_url_for_trust_lookup, ) from litellm.proxy._types import MCPTransport from litellm.types.mcp import MCPTransportType @@ -119,6 +121,60 @@ def test_normalize_trust_score_maps_do_percent_scale_to_fraction() -> None: assert abs(low_trust - 0.01) < 1e-9 +def test_sanitize_url_for_trust_lookup_strips_userinfo_and_sensitive_query_params() -> ( + None +): + sanitized = sanitize_url_for_trust_lookup( + "https://secret-token@mcp.example.com:8443/mcp?api_key=leak&safe=1#frag" + ) + + assert sanitized == "https://mcp.example.com:8443/mcp?safe=1" + + +def test_assert_requested_server_passes_trust_filter_rejects_untrusted_server() -> None: + from fastapi import HTTPException + + trusted = _server(server_id="srv-trusted") + with pytest.raises(HTTPException) as exc_info: + assert_requested_server_passes_trust_filter( + filtered_servers=[trusted], + server_id="srv-untrusted", + ) + + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_get_trust_score_forwards_sanitized_url_to_do() -> None: + http_client = _FakeHttpClient( + [ + _trust_response( + payload={ + "found": True, + "server_url": "https://mcp.example.com/mcp", + "trust_score": 90, + } + ) + ] + ) + config = MCPTrustScoringConfig(enabled=True) + client = _isolated_client( + config, + http_client, + ) + + await client.get_trust_score( + "https://secret@mcp.example.com/mcp?api_key=leak&safe=1" + ) + + assert http_client.calls == [ + ( + config.api_url, + {"url": "https://mcp.example.com/mcp?safe=1"}, + ) + ] + + def test_parse_trust_response_found_server() -> None: result = parse_trust_response( server_url="https://alpha.example.com/mcp", @@ -195,7 +251,9 @@ async def test_filter_servers_by_trust_excludes_low_score() -> None: filtered = await client.filter_servers_by_trust( [ - _server(server_id="srv-1", name="alpha", url="https://alpha.example.com/mcp"), + _server( + server_id="srv-1", name="alpha", url="https://alpha.example.com/mcp" + ), _server(server_id="srv-2", name="beta", url="https://beta.example.com/mcp"), ] ) @@ -205,7 +263,9 @@ async def test_filter_servers_by_trust_excludes_low_score() -> None: @pytest.mark.asyncio async def test_filter_servers_fail_open_on_lookup_error() -> None: - http_client = _FakeHttpClient([_trust_response(status_code=503, payload={"error": "busy"})]) + http_client = _FakeHttpClient( + [_trust_response(status_code=503, payload={"error": "busy"})] + ) client = _isolated_client( MCPTrustScoringConfig(enabled=True, fail_open=True), http_client, @@ -220,7 +280,9 @@ async def test_filter_servers_fail_open_on_lookup_error() -> None: @pytest.mark.asyncio async def test_filter_servers_fail_closed_on_lookup_error() -> None: - http_client = _FakeHttpClient([_trust_response(status_code=503, payload={"error": "busy"})]) + http_client = _FakeHttpClient( + [_trust_response(status_code=503, payload={"error": "busy"})] + ) client = _isolated_client( MCPTrustScoringConfig(enabled=True, fail_open=False), http_client, @@ -275,7 +337,9 @@ async def test_rank_servers_by_trust_orders_highest_first() -> None: ranked = await client.rank_servers_by_trust( [ _server(server_id="srv-2", name="beta", url="https://beta.example.com/mcp"), - _server(server_id="srv-1", name="alpha", url="https://alpha.example.com/mcp"), + _server( + server_id="srv-1", name="alpha", url="https://alpha.example.com/mcp" + ), ] ) @@ -314,7 +378,9 @@ async def test_module_helpers_use_explicit_client() -> None: @pytest.mark.asyncio -async def test_manager_list_tools_applies_trust_filter(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_manager_list_tools_applies_trust_filter( + monkeypatch: pytest.MonkeyPatch, +) -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, ) @@ -347,7 +413,9 @@ async def test_manager_list_tools_applies_trust_filter(monkeypatch: pytest.Monke ) manager = MCPServerManager() - alpha = _server(server_id="srv-1", name="alpha", url="https://alpha.example.com/mcp") + alpha = _server( + server_id="srv-1", name="alpha", url="https://alpha.example.com/mcp" + ) beta = _server(server_id="srv-2", name="beta", url="https://beta.example.com/mcp") manager.registry = {"srv-1": alpha, "srv-2": beta}