From b2146713e5d576ddd62a9ef516c8cb2fefdee919 Mon Sep 17 00:00:00 2001 From: "Jugal D. Bhatt" <55304795+jugaldb@users.noreply.github.com> Date: Tue, 5 Aug 2025 16:26:36 -0700 Subject: [PATCH] refactor mcp guardrails (#13238) --- litellm/integrations/custom_logger.py | 54 ---- .../mcp_server/mcp_server_manager.py | 77 +++-- litellm/proxy/utils.py | 266 ++++-------------- 3 files changed, 102 insertions(+), 295 deletions(-) diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index ded5ccca76..b5c7101dde 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -34,8 +34,6 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp import ( - MCPDuringCallRequestObject, - MCPDuringCallResponseObject, MCPPostCallResponseObject, MCPPreCallRequestObject, MCPPreCallResponseObject, @@ -412,59 +410,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ######################################################### # MCP TOOL CALL HOOKS ######################################################### - async def async_pre_mcp_tool_call_hook( - self, - kwargs, - request_obj: MCPPreCallRequestObject, - start_time, - end_time - ) -> Optional[MCPPreCallResponseObject]: - """ - This hook gets called before the MCP tool call is made. - Useful for: - - Validating tool calls before execution - - Modifying arguments before they are sent to the MCP server - - Implementing access control and rate limiting - - Adding custom metadata or tracking information - - Args: - kwargs: The logging kwargs containing model call details - request_obj: MCPPreCallRequestObject containing tool name, arguments, and metadata - start_time: Start time of the request - end_time: End time of the request - - Returns: - MCPPreCallResponseObject with validation results and any modifications - """ - return None - - async def async_during_mcp_tool_call_hook( - self, - kwargs, - request_obj: MCPDuringCallRequestObject, - start_time, - end_time - ) -> Optional[MCPDuringCallResponseObject]: - """ - This hook gets called during the MCP tool call execution. - - Useful for: - - Concurrent monitoring and validation during tool execution - - Implementing timeouts and cancellation logic - - Real-time cost tracking and billing - - Performance monitoring and metrics collection - - Args: - kwargs: The logging kwargs containing model call details - request_obj: MCPDuringCallRequestObject containing tool execution context - start_time: Start time of the request - end_time: End time of the request - - Returns: - MCPDuringCallResponseObject with execution control decisions - """ - return None async def async_post_mcp_tool_call_hook( self, kwargs, response_obj: MCPPostCallResponseObject, start_time, end_time diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index e1ea9918e6..77444219be 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -644,6 +644,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" ######################################################### if proxy_logging_obj: pre_hook_kwargs = { @@ -651,24 +652,32 @@ class MCPServerManager: "arguments": arguments, "server_name": server_name_from_prefix, "user_api_key_auth": user_api_key_auth, + "user_api_key_user_id": getattr(user_api_key_auth, 'user_id', None) if user_api_key_auth else None, + "user_api_key_team_id": getattr(user_api_key_auth, 'team_id', None) if user_api_key_auth else None, + "user_api_key_end_user_id": getattr(user_api_key_auth, 'end_user_id', None) if user_api_key_auth else None, + "user_api_key_hash": getattr(user_api_key_auth, 'api_key_hash', None) if user_api_key_auth else None, } + + # Create MCP request object for processing + mcp_request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs(pre_hook_kwargs) + + # Convert to LLM format for existing guardrail compatibility + synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(mcp_request_obj, pre_hook_kwargs) + try: - pre_hook_result = await proxy_logging_obj.async_pre_mcp_tool_call_hook( - kwargs=pre_hook_kwargs, - request_obj=None, # Will be created in the hook - start_time=start_time, - end_time=start_time, + # Use standard pre_call_hook with call_type="mcp_call" + 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 ) - - if pre_hook_result: - # Apply any argument modifications - if pre_hook_result.get("modified_arguments"): - arguments = pre_hook_result["modified_arguments"] - except ( - BlockedPiiEntityError, - GuardrailRaisedException, - HTTPException, - ) as e: + if modified_data: + # Convert response back to MCP format and apply modifications + modified_kwargs = proxy_logging_obj._convert_mcp_hook_response_to_kwargs(modified_data, pre_hook_kwargs) + if modified_kwargs.get("arguments") != arguments: + arguments = modified_kwargs["arguments"] + + except (BlockedPiiEntityError, GuardrailRaisedException, HTTPException) as e: # Re-raise guardrail exceptions to properly fail the MCP call verbose_logger.error( f"Guardrail blocked MCP tool call pre call: {str(e)}" @@ -699,22 +708,34 @@ class MCPServerManager: name=original_tool_name, arguments=arguments, ) - - # Initialize during_hook_task as None - during_hook_task = None tasks = [] - # Start during hook if proxy_logging_obj is available if proxy_logging_obj: + # Create synthetic LLM data for during hook processing + from litellm.types.mcp import MCPDuringCallRequestObject + from litellm.types.llms.base import HiddenParams + + request_obj = MCPDuringCallRequestObject( + tool_name=name, + arguments=arguments, + server_name=server_name_from_prefix, + start_time=start_time.timestamp() if start_time else None, + hidden_params=HiddenParams(), + ) + + during_hook_kwargs = { + "name": name, + "arguments": arguments, + "server_name": server_name_from_prefix, + "user_api_key_auth": user_api_key_auth, + } + + synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(request_obj, during_hook_kwargs) + during_hook_task = asyncio.create_task( - proxy_logging_obj.async_during_mcp_tool_call_hook( - kwargs={ - "name": name, - "arguments": arguments, - "server_name": server_name_from_prefix, - }, - request_obj=None, # Will be created in the hook - start_time=start_time, - end_time=start_time, + proxy_logging_obj.during_call_hook( + user_api_key_dict=user_api_key_auth, + data=synthetic_llm_data, + call_type="mcp_call" #type: ignore ) ) tasks.append(during_hook_task) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1415660827..0e6b242bdb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -55,11 +55,7 @@ from litellm import ( from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceLogging, ServiceTypes from litellm.caching.caching import DualCache, RedisCache -from litellm.exceptions import ( - BlockedPiiEntityError, - GuardrailRaisedException, - RejectedRequestError, -) +from litellm.exceptions import RejectedRequestError from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting @@ -452,108 +448,6 @@ class ProxyLogging: litellm_parent_otel_span=None, ) - async def async_pre_mcp_tool_call_hook( - self, - kwargs: dict, - request_obj: Any, - start_time: datetime, - end_time: datetime, - ) -> Optional[Any]: - """ - Pre MCP Tool Call Hook - - Use this to validate and modify MCP tool calls before execution. - Reuses existing LLM guardrail logic by converting MCP calls to message format. - """ - from litellm.types.llms.base import HiddenParams - from litellm.types.mcp import MCPPreCallRequestObject - - callbacks = self.get_combined_callback_list( - dynamic_success_callbacks=getattr(self, "dynamic_success_callbacks", None), - global_callbacks=litellm.success_callback, - ) - - # Create the request object if it's not already one - if not isinstance(request_obj, MCPPreCallRequestObject): - # Convert UserAPIKeyAuth object to dict if needed - user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict( - kwargs.get("user_api_key_auth") - ) - - request_obj = MCPPreCallRequestObject( - tool_name=kwargs.get("name", ""), - arguments=kwargs.get("arguments", {}), - server_name=kwargs.get("server_name"), - user_api_key_auth=user_api_key_auth_dict, - hidden_params=HiddenParams(), - ) - - for callback in callbacks: - try: - _callback: Optional[CustomLogger] = None - if isinstance(callback, str): - from typing import cast - - from litellm import _custom_logger_compatible_callbacks_literal - - _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( - cast(_custom_logger_compatible_callbacks_literal, callback) - ) - else: - _callback = callback # type: ignore - - if _callback is not None and isinstance(_callback, CustomGuardrail): - from litellm.types.guardrails import GuardrailEventHooks - - # Check if guardrail should be run for pre_call hook (reusing existing logic) - if ( - _callback.should_run_guardrail( - data=kwargs, event_type=GuardrailEventHooks.pre_mcp_call - ) - is not True - ): - continue - - # Convert MCP tool call to LLM message format for existing guardrail logic - synthetic_llm_data = self._convert_mcp_to_llm_format( - request_obj, kwargs - ) - # Reuse existing LLM guardrail logic - user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict( - kwargs.get("user_api_key_auth") - ) - - result = await _callback.async_pre_call_hook( - user_api_key_dict=user_api_key_auth_dict, # type: ignore - cache=self.call_details["user_api_key_cache"], - data=synthetic_llm_data, - call_type="mcp_call", - ) - - # Convert result back to MCP response format if blocked/modified - if result is not None: - mcp_response = self._convert_llm_result_to_mcp_response( - result, request_obj - ) - if mcp_response is not None: - return self._parse_pre_mcp_call_hook_response( - response=mcp_response, original_request=request_obj - ) - - except ( - BlockedPiiEntityError, - GuardrailRaisedException, - HTTPException, - ) as e: - # Re-raise guardrail exceptions so they can be properly handled - raise e - except Exception as e: - verbose_proxy_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( - str(e) - ) - ) - return None def _convert_user_api_key_auth_to_dict(self, user_api_key_auth_obj): """ @@ -567,7 +461,7 @@ class ProxyLogging: elif hasattr(user_api_key_auth_obj, "__dict__"): # If it's a regular object, convert to dict return user_api_key_auth_obj.__dict__ - return user_api_key_auth_obj + return {} def _convert_mcp_to_llm_format(self, request_obj, kwargs: dict) -> dict: """ @@ -765,8 +659,6 @@ class ProxyLogging: """ Convert LLM guardrail result back to MCP during call response format. """ - from litellm.types.mcp import MCPDuringCallResponseObject - # If result is an exception, it means the guardrail wants to stop execution if isinstance(llm_result, Exception): return MCPDuringCallResponseObject( @@ -836,112 +728,39 @@ class ProxyLogging: } return result - async def async_during_mcp_tool_call_hook( - self, - kwargs: dict, - request_obj: Any, - start_time: datetime, - end_time: datetime, - ) -> Optional[Any]: + def _create_mcp_request_object_from_kwargs(self, kwargs: dict) -> "MCPPreCallRequestObject": """ - During MCP Tool Call Hook - - Use this for concurrent monitoring and validation during tool execution. - Reuses existing LLM guardrail logic by converting MCP calls to message format. + Helper function to create MCPPreCallRequestObject from kwargs for standard pre_call_hook. """ from litellm.types.llms.base import HiddenParams - from litellm.types.mcp import MCPDuringCallRequestObject + from litellm.types.mcp import MCPPreCallRequestObject - callbacks = self.get_combined_callback_list( - dynamic_success_callbacks=getattr(self, "dynamic_success_callbacks", None), - global_callbacks=litellm.success_callback, + user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(kwargs.get("user_api_key_auth")) + + return MCPPreCallRequestObject( + tool_name=kwargs.get("name", ""), + arguments=kwargs.get("arguments", {}), + server_name=kwargs.get("server_name"), + user_api_key_auth=user_api_key_auth_dict, + hidden_params=HiddenParams(), ) - # Create the request object if it's not already one - if not isinstance(request_obj, MCPDuringCallRequestObject): - request_obj = MCPDuringCallRequestObject( - tool_name=kwargs.get("name", ""), - arguments=kwargs.get("arguments", {}), - server_name=kwargs.get("server_name"), - start_time=start_time.timestamp() if start_time else None, - hidden_params=HiddenParams(), - ) - - for callback in callbacks: - try: - _callback: Optional[CustomLogger] = None - if isinstance(callback, str): - from typing import cast - - from litellm import _custom_logger_compatible_callbacks_literal - - _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( - cast(_custom_logger_compatible_callbacks_literal, callback) - ) - else: - _callback = callback # type: ignore - - if _callback is not None and isinstance(_callback, CustomGuardrail): - from litellm.types.guardrails import GuardrailEventHooks - - # Check if guardrail should be run for during_call hook (reusing existing logic) - if ( - _callback.should_run_guardrail( - data=kwargs, event_type=GuardrailEventHooks.during_mcp_call - ) - is not True - ): - continue - # Convert MCP tool call to LLM message format for existing guardrail logic - synthetic_llm_data = self._convert_mcp_to_llm_format( - request_obj, kwargs - ) - - # Reuse existing LLM guardrail logic for during call - user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict( - kwargs.get("user_api_key_auth") - ) - - result = await _callback.async_moderation_hook( - data=synthetic_llm_data, - user_api_key_dict=user_api_key_auth_dict, # type: ignore - call_type="mcp_call", - ) - # Convert result back to MCP response format if blocked/modified - if result is not None: - mcp_response = self._convert_llm_result_to_mcp_during_response( - result, request_obj - ) - if mcp_response is not None: - return self._parse_during_mcp_call_hook_response( - response=mcp_response - ) - - except Exception as e: - raise e - verbose_proxy_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( - str(e) - ) - ) - return None - - def _parse_during_mcp_call_hook_response( - self, response: MCPDuringCallResponseObject - ) -> Dict[str, Any]: + def _convert_mcp_hook_response_to_kwargs(self, response_data: Optional[dict], original_kwargs: dict) -> dict: """ - Parse the response from the during_mcp_tool_call_hook - - 1. Check if execution should continue - 2. Handle any error messages - 3. Apply any hidden parameter updates + Helper function to convert pre_call_hook response back to kwargs for MCP usage. """ - result = { - "should_continue": response.should_continue, - "error_message": response.error_message, - "hidden_params": response.hidden_params, - } - return result + if not response_data: + return original_kwargs + + # Apply any argument modifications from the hook response + modified_kwargs = original_kwargs.copy() + + # If the response contains modified arguments, apply them + if response_data.get("modified_arguments"): + modified_kwargs["arguments"] = response_data["modified_arguments"] + + return modified_kwargs + async def process_pre_call_hook_response(self, response, data, call_type): if isinstance(response, Exception): @@ -975,6 +794,7 @@ class ProxyLogging: "audio_transcription", "pass_through_endpoint", "rerank", + "mcp_call", ], ) -> None: pass @@ -993,6 +813,7 @@ class ProxyLogging: "audio_transcription", "pass_through_endpoint", "rerank", + "mcp_call", ], ) -> dict: pass @@ -1010,6 +831,7 @@ class ProxyLogging: "audio_transcription", "pass_through_endpoint", "rerank", + "mcp_call", ], ) -> Optional[dict]: """ @@ -1081,10 +903,14 @@ class ProxyLogging: _callback = callback # type: ignore if _callback is not None and isinstance(_callback, CustomGuardrail): from litellm.types.guardrails import GuardrailEventHooks - + + event_type = GuardrailEventHooks.pre_call + if call_type == "mcp_call": + event_type = GuardrailEventHooks.pre_mcp_call + if ( _callback.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call + data=data, event_type=event_type ) is not True ): @@ -1108,6 +934,9 @@ class ProxyLogging: and _callback.__class__.async_pre_call_hook != CustomLogger.async_pre_call_hook ): + if call_type == "mcp_call" and user_api_key_dict is None: + continue + response = await _callback.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=self.call_details["user_api_key_cache"], @@ -1126,7 +955,7 @@ class ProxyLogging: async def during_call_hook( self, data: dict, - user_api_key_dict: UserAPIKeyAuth, + user_api_key_dict: Optional[UserAPIKeyAuth], call_type: Literal[ "completion", "responses", @@ -1134,6 +963,7 @@ class ProxyLogging: "image_generation", "moderation", "audio_transcription", + "mcp_call", ], ): """ @@ -1156,16 +986,26 @@ class ProxyLogging: # Main - V2 Guardrails implementation from litellm.types.guardrails import GuardrailEventHooks + event_type = GuardrailEventHooks.during_call + if call_type == "mcp_call": + event_type = GuardrailEventHooks.during_mcp_call + if ( callback.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.during_call + data=data, event_type=event_type ) is not True ): continue + # Convert user_api_key_dict to proper format for async_moderation_hook + if call_type == "mcp_call": + user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(user_api_key_dict) + else: + user_api_key_auth_dict = user_api_key_dict + await callback.async_moderation_hook( data=data, - user_api_key_dict=user_api_key_dict, + user_api_key_dict=user_api_key_auth_dict, # type: ignore call_type=call_type, ) except Exception as e: