tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-11 11:04:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #18630 from BerriAI/litellm_fix_mcp_guardrail

Browse Source 
fix: MCP handling in unified guardrail
This commit is contained in:
YutaSaito
2026-01-05 12:25:14 +09:00
committed by GitHub
parent 399579f8ea 0509fc79da
commit fcabc059ca
8 changed files with 337 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/__init__.py
+16
View File
@@ -45,6 +45,7 @@ def get_cost_for_web_search_request(
return 0.0
elif custom_llm_provider == "xai":
from .xai.cost_calculator import cost_per_web_search_request
return cost_per_web_search_request(usage=usage, model_info=model_info)
else:
return None
@@ -110,6 +111,21 @@ def discover_guardrail_translation_mappings() -> (
verbose_logger.error(f"Error processing {module_path}: {e}")
continue
try:
from litellm.proxy._experimental.mcp_server.guardrail_translation import (
guardrail_translation_mappings as mcp_guardrail_translation_mappings,
)
discovered_mappings.update(mcp_guardrail_translation_mappings)
verbose_logger.debug(
"Loaded MCP guardrail translation mappings: %s",
list(mcp_guardrail_translation_mappings.keys()),
)
except ImportError:
verbose_logger.debug(
"MCP guardrail translation mappings not available; skipping"
)
verbose_logger.debug(
f"Discovered {len(discovered_mappings)} guardrail translation mappings: {list(discovered_mappings.keys())}"
)
litellm/proxy/_experimental/mcp_server/guardrail_translation/__init__.py
+16
View File
@@ -0,0 +1,16 @@
"""Guardrail translation mapping for MCP tool calls."""
from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import (
MCPGuardrailTranslationHandler,
)
from litellm.types.utils import CallTypes
# This mapping lives alongside the MCP server implementation because MCP
# integrations are managed by the proxy subsystem, not litellm.llms providers.
# Unified guardrails import this module explicitly to register the handler.
guardrail_translation_mappings = {
CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler,
}
__all__ = ["guardrail_translation_mappings", "MCPGuardrailTranslationHandler"]
litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py
+89
View File
@@ -0,0 +1,89 @@
"""
MCP Guardrail Handler for Unified Guardrails.
This handler works with the synthetic "messages" payload generated by
`ProxyLogging._convert_mcp_to_llm_format`, which always produces a single user
message whose `content` string encodes the MCP tool name and arguments. The
handler simply feeds that text through the configured guardrail and writes the
result back onto the message.
"""
from typing import TYPE_CHECKING, Any, Dict, Optional
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from mcp.types import CallToolResult
class MCPGuardrailTranslationHandler(BaseTranslation):
"""Guardrail translation handler for MCP tool calls."""
async def process_input_messages(
self,
data: Dict[str, Any],
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Optional[Any] = None,
) -> Dict[str, Any]:
messages = data.get("messages")
if not isinstance(messages, list) or not messages:
verbose_proxy_logger.debug("MCP Guardrail: No messages to process")
return data
first_message = messages[0]
content: Optional[str] = None
if isinstance(first_message, dict):
content = first_message.get("content")
else:
content = getattr(first_message, "content", None)
if not isinstance(content, str):
verbose_proxy_logger.debug(
"MCP Guardrail: Message content missing or not a string",
)
return data
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=GenericGuardrailAPIInputs(texts=[content]),
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = (
guardrailed_inputs.get("texts", []) if guardrailed_inputs else []
)
if guardrailed_texts:
new_content = guardrailed_texts[0]
if isinstance(first_message, dict):
first_message["content"] = new_content
else:
setattr(first_message, "content", new_content)
verbose_proxy_logger.debug(
"MCP Guardrail: Updated content for tool %s",
data.get("mcp_tool_name"),
)
else:
verbose_proxy_logger.debug(
"MCP Guardrail: Guardrail returned no text updates for tool %s",
data.get("mcp_tool_name"),
)
return data
async def process_output_response(
self,
response: "CallToolResult",
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Optional[Any] = None,
user_api_key_dict: Optional[Any] = None,
) -> Any:
# Not implemented: MCP guardrail translation never calls this path today.
verbose_proxy_logger.debug(
"MCP Guardrail: Output processing not implemented for MCP tools",
)
return response
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+5 -4
View File
@@ -30,6 +30,7 @@ from pydantic import AnyUrl
import litellm
from litellm._logging import verbose_logger
from litellm.types.utils import CallTypes
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.experimental_mcp_client.client import MCPClient
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
@@ -1676,11 +1677,11 @@ class MCPServerManager:
)
try:
# Use standard pre_call_hook with call_type="mcp_call"
# Use standard pre_call_hook
modified_data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_auth, # type: ignore
data=synthetic_llm_data,
call_type="mcp_call", # type: ignore
call_type=CallTypes.call_mcp_tool.value,
)
if modified_data:
# Convert response back to MCP format and apply modifications
@@ -1737,7 +1738,7 @@ class MCPServerManager:
proxy_logging_obj.during_call_hook(
user_api_key_dict=user_api_key_auth,
data=synthetic_llm_data,
call_type="mcp_call", # type: ignore
call_type=CallTypes.call_mcp_tool.value,
)
)
@@ -1893,7 +1894,7 @@ class MCPServerManager:
#########################################################
# Pre MCP Tool Call Hook
# Allow validation and modification of tool calls before execution
# Using standard pre_call_hook with call_type="mcp_call"
# Using standard pre_call_hook
#########################################################
if proxy_logging_obj:
await self.pre_call_tool_check(
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
+12 -5
View File
@@ -28,7 +28,6 @@ class UnifiedLLMGuardrails(CustomLogger):
self,
**kwargs,
):
# store kwargs as optional_params
self.optional_params = kwargs
@@ -63,6 +62,9 @@ class UnifiedLLMGuardrails(CustomLogger):
return data
event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call
if call_type == CallTypes.call_mcp_tool.value:
event_type = GuardrailEventHooks.pre_mcp_call
if (
guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type)
is not True
@@ -114,6 +116,9 @@ class UnifiedLLMGuardrails(CustomLogger):
return data
event_type: GuardrailEventHooks = GuardrailEventHooks.during_call
if call_type == CallTypes.call_mcp_tool.value:
event_type = GuardrailEventHooks.during_mcp_call
if (
guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type)
is not True
@@ -128,7 +133,10 @@ class UnifiedLLMGuardrails(CustomLogger):
endpoint_guardrail_translation_mappings = (
load_guardrail_translation_mappings()
)
if call_type is not None and CallTypes(call_type) not in endpoint_guardrail_translation_mappings:
if (
call_type is not None
and CallTypes(call_type) not in endpoint_guardrail_translation_mappings
):
return data
endpoint_translation = endpoint_guardrail_translation_mappings[
@@ -180,8 +188,8 @@ class UnifiedLLMGuardrails(CustomLogger):
call_type: Optional[CallTypesLiteral] = None
if user_api_key_dict.request_route is not None:
call_types = get_call_types_for_route(user_api_key_dict.request_route)
if call_types is not None and len(call_types) > 0: # type: ignore
call_type = call_types[0] # type: ignore
if call_types is not None and len(call_types) > 0: # type: ignore
call_type = call_types[0] # type: ignore
if call_type is None:
call_type = _infer_call_type(call_type=None, completion_response=response) # type: ignore
@@ -330,7 +338,6 @@ class UnifiedLLMGuardrails(CustomLogger):
# Process chunk based on sampling rate
if chunk_counter % sampling_rate == 0:
verbose_proxy_logger.debug(
"Processing streaming chunk %s (sampling_rate=%s) with guardrail %s",
chunk_counter,
litellm/proxy/utils.py
+37 -35
View File
@@ -151,25 +151,25 @@ def _get_email_logger_class():
"""
Determine which email logger class to use based on environment variables.
Priority: SendGrid > Resend > SMTP > BaseEmailLogger (fallback)
Returns:
The email logger class to use, or None if BaseEmailLogger is not available
"""
if BaseEmailLogger is None:
return None
# Check for SendGrid API key
if SendGridEmailLogger is not None and os.getenv("SENDGRID_API_KEY"):
return SendGridEmailLogger
# Check for Resend API key
if ResendEmailLogger is not None and os.getenv("RESEND_API_KEY"):
return ResendEmailLogger
# Check for SMTP configuration
if SMTPEmailLogger is not None and os.getenv("SMTP_HOST"):
return SMTPEmailLogger
# Fallback to BaseEmailLogger (though it won't actually send emails)
return BaseEmailLogger
@@ -452,7 +452,6 @@ class ProxyLogging:
litellm.logging_callback_manager.add_litellm_callback(self.service_logging_obj) # type: ignore
for callback in litellm.callbacks:
if isinstance(callback, str):
callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( # type: ignore
cast(_custom_logger_compatible_callbacks_literal, callback),
internal_usage_cache=self.internal_usage_cache.dual_cache,
@@ -965,7 +964,7 @@ class ProxyLogging:
# Determine the event type based on call type
event_type = GuardrailEventHooks.pre_call
if call_type == "mcp_call":
if call_type == CallTypes.call_mcp_tool.value:
event_type = GuardrailEventHooks.pre_mcp_call
# Check if the guardrail should run for this request
@@ -1038,7 +1037,6 @@ class ProxyLogging:
data.pop("prompt_id", None)
if custom_logger and prompt_spec is not None:
(
model,
messages,
@@ -1261,7 +1259,7 @@ class ProxyLogging:
from litellm.types.guardrails import GuardrailEventHooks
event_type = GuardrailEventHooks.during_call
if call_type == "mcp_call":
if call_type == CallTypes.call_mcp_tool.value:
event_type = GuardrailEventHooks.during_mcp_call
if (
@@ -1270,7 +1268,7 @@ class ProxyLogging:
):
continue
# Convert user_api_key_dict to proper format for async_moderation_hook
if call_type == "mcp_call":
if call_type == CallTypes.call_mcp_tool.value:
user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(
user_api_key_dict
)
@@ -1288,7 +1286,6 @@ class ProxyLogging:
call_type=call_type,
)
else:
guardrail_task = callback.async_moderation_hook(
data=data,
user_api_key_dict=user_api_key_auth_dict, # type: ignore
@@ -1337,7 +1334,7 @@ class ProxyLogging:
if self.alerting is None:
# do nothing if alerting is not switched on
return
if "slack" in self.alerting:
await self.slack_alerting_instance.budget_alerts(
type=type,
@@ -1548,7 +1545,10 @@ class ProxyLogging:
traceback_str=traceback_str,
)
# If callback returned an HTTPException, use it (first one wins)
if isinstance(hook_result, HTTPException) and transformed_exception is None:
if (
isinstance(hook_result, HTTPException)
and transformed_exception is None
):
transformed_exception = hook_result
except HTTPException as e:
# If callback raised an HTTPException, use it (first one wins)
@@ -1849,7 +1849,6 @@ class ProxyLogging:
current_response = response
for callback in litellm.callbacks:
_callback: Optional[CustomLogger] = None
if isinstance(callback, str):
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
@@ -3568,11 +3567,13 @@ class ProxyUpdateSpend:
)
# Atomically read and remove logs to process (protected by lock)
async with prisma_client._spend_log_transactions_lock:
logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL]
logs_to_process = prisma_client.spend_log_transactions[
:MAX_LOGS_PER_INTERVAL
]
# Remove the logs we're about to process
prisma_client.spend_log_transactions = (
prisma_client.spend_log_transactions[len(logs_to_process):]
)
prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[
len(logs_to_process) :
]
start_time = time.time()
try:
for i in range(n_retry_times + 1):
@@ -3675,9 +3676,7 @@ async def update_spend( # noqa: PLR0915
# Check queue size with lock protection
async with prisma_client._spend_log_transactions_lock:
queue_size = len(prisma_client.spend_log_transactions)
verbose_proxy_logger.debug(
"Spend Logs transactions: {}".format(queue_size)
)
verbose_proxy_logger.debug("Spend Logs transactions: {}".format(queue_size))
# Process spend log transactions when called directly.
# This keeps backwards compatibility with the old behavior.
@@ -3699,19 +3698,19 @@ async def update_spend_logs_job(
):
"""
Job to process spend_log_transactions queue.
This job is triggered based on queue size rather than time.
Processes spend log transactions when the queue reaches a threshold.
"""
n_retry_times = 3
# Check queue size with lock protection
async with prisma_client._spend_log_transactions_lock:
queue_size = len(prisma_client.spend_log_transactions)
if queue_size == 0:
return
await ProxyUpdateSpend.update_spend_logs(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
@@ -3728,7 +3727,7 @@ async def _monitor_spend_logs_queue(
"""
Background task that monitors the spend_log_transactions queue size
and triggers processing when the threshold is reached.
Args:
prisma_client: Prisma client instance
db_writer_client: Optional HTTP handler for external spend logs endpoint
@@ -3738,23 +3737,23 @@ async def _monitor_spend_logs_queue(
SPEND_LOG_QUEUE_POLL_INTERVAL,
SPEND_LOG_QUEUE_SIZE_THRESHOLD,
)
threshold = SPEND_LOG_QUEUE_SIZE_THRESHOLD
base_interval = SPEND_LOG_QUEUE_POLL_INTERVAL
max_backoff = 30.0 # Maximum backoff interval in seconds
backoff_multiplier = 1.5 # Exponential backoff multiplier
current_interval = base_interval
verbose_proxy_logger.info(
f"Starting spend logs queue monitor (threshold: {threshold}, poll_interval: {base_interval}s)"
)
while True:
try:
# Check queue size with lock protection
async with prisma_client._spend_log_transactions_lock:
queue_size = len(prisma_client.spend_log_transactions)
if queue_size > 0:
if queue_size >= threshold:
verbose_proxy_logger.debug(
@@ -3767,8 +3766,10 @@ async def _monitor_spend_logs_queue(
f"Spend logs queue size ({queue_size}) below threshold ({threshold}), processing with backoff"
)
# Exponential backoff when below threshold but still processing
current_interval = min(current_interval * backoff_multiplier, max_backoff)
current_interval = min(
current_interval * backoff_multiplier, max_backoff
)
await update_spend_logs_job(
prisma_client=prisma_client,
db_writer_client=db_writer_client,
@@ -3776,8 +3777,10 @@ async def _monitor_spend_logs_queue(
)
else:
# Exponential backoff when no logs to process
current_interval = min(current_interval * backoff_multiplier, max_backoff)
current_interval = min(
current_interval * backoff_multiplier, max_backoff
)
await asyncio.sleep(current_interval)
except Exception as e:
verbose_proxy_logger.error(
@@ -3788,7 +3791,6 @@ async def _monitor_spend_logs_queue(
await asyncio.sleep(current_interval)
def _raise_failed_update_spend_exception(
e: Exception, start_time: float, proxy_logging_obj: ProxyLogging
):
tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py
+78
View File
@@ -0,0 +1,78 @@
"""Tests for the MCP guardrail translation handler."""
import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import (
MCPGuardrailTranslationHandler,
)
class MockGuardrail(CustomGuardrail):
"""Simple guardrail mock that records invocations."""
def __init__(self, return_texts=None):
super().__init__(guardrail_name="mock-mcp-guardrail")
self.return_texts = return_texts
self.call_count = 0
self.last_inputs = None
async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
self.call_count += 1
self.last_inputs = inputs
if self.return_texts is not None:
return {"texts": self.return_texts}
texts = inputs.get("texts", [])
return {"texts": [f"{text} [SAFE]" for text in texts]}
@pytest.mark.asyncio
async def test_process_input_messages_updates_content():
"""Handler should update the synthetic message content when guardrail modifies text."""
handler = MCPGuardrailTranslationHandler()
guardrail = MockGuardrail()
original_content = "Tool: weather\nArguments: {'city': 'tokyo'}"
data = {
"messages": [{"role": "user", "content": original_content}],
"mcp_tool_name": "weather",
}
result = await handler.process_input_messages(data, guardrail)
assert result["messages"][0]["content"].endswith("[SAFE]")
assert guardrail.last_inputs == {"texts": [original_content]}
assert guardrail.call_count == 1
@pytest.mark.asyncio
async def test_process_input_messages_skips_when_no_messages():
"""Handler should skip guardrail invocation if messages array is missing or empty."""
handler = MCPGuardrailTranslationHandler()
guardrail = MockGuardrail()
data = {"mcp_tool_name": "noop"}
result = await handler.process_input_messages(data, guardrail)
assert result == data
assert guardrail.call_count == 0
@pytest.mark.asyncio
async def test_process_input_messages_handles_empty_guardrail_result():
"""Handler should leave content untouched when guardrail returns no text updates."""
handler = MCPGuardrailTranslationHandler()
guardrail = MockGuardrail(return_texts=[])
original_content = "Tool: calendar\nArguments: {'date': '2024-12-25'}"
data = {
"messages": [{"role": "user", "content": original_content}],
"mcp_tool_name": "calendar",
}
result = await handler.process_input_messages(data, guardrail)
assert result["messages"][0]["content"] == original_content
assert guardrail.call_count == 1
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
+84
View File
@@ -0,0 +1,84 @@
"""Tests for unified guardrail."""
import pytest
from litellm.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import (
MCPGuardrailTranslationHandler,
)
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import unified_guardrail as unified_module
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import CallTypes
class RecordingGuardrail(CustomGuardrail):
"""Records the event types it is asked to run for."""
def __init__(self):
super().__init__(guardrail_name="recording-guardrail")
self.event_history = []
def should_run_guardrail(self, data, event_type): # type: ignore[override]
self.event_history.append(event_type)
return True
async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
return {"texts": inputs.get("texts", [])}
@pytest.fixture(autouse=True)
def _inject_mcp_handler_mapping():
"""Inject MCP handler mapping so the unified guardrail can run inside tests."""
unified_module.endpoint_guardrail_translation_mappings = {
CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler,
}
yield
unified_module.endpoint_guardrail_translation_mappings = None
@pytest.mark.asyncio
async def test_pre_call_hook_uses_mcp_event_type():
"""pre_call hook should swap to GuardrailEventHooks.pre_mcp_call for MCP calls."""
handler = UnifiedLLMGuardrails()
guardrail = RecordingGuardrail()
cache = DualCache()
data = {
"guardrail_to_apply": guardrail,
"messages": [{"role": "user", "content": "Tool: test\nArguments: {}"}],
"model": "mcp-tool-call",
}
await handler.async_pre_call_hook(
user_api_key_dict=None,
cache=cache,
data=data,
call_type=CallTypes.call_mcp_tool.value,
)
assert guardrail.event_history == [GuardrailEventHooks.pre_mcp_call]
@pytest.mark.asyncio
async def test_moderation_hook_uses_mcp_event_type():
"""moderation hook should request GuardrailEventHooks.during_mcp_call for MCP calls."""
handler = UnifiedLLMGuardrails()
guardrail = RecordingGuardrail()
data = {
"guardrail_to_apply": guardrail,
"messages": [{"role": "user", "content": "Tool: test\nArguments: {}"}],
"model": "mcp-tool-call",
}
await handler.async_moderation_hook(
data=data,
user_api_key_dict=None,
call_type=CallTypes.call_mcp_tool.value,
)
assert guardrail.event_history == [GuardrailEventHooks.during_mcp_call]