[MCP Gateway] Litellm mcp pre and during guardrails (#13188)

* add guardrail support

* add guardrail support

* guardrails for MCP

* added changes

* add mcp guardrails

* added test

* add ui

* fix guardrail form

* working with cursor

* remvoe print

* fix mcp servertests

* fix mypy and remove console logs

* fix mypy and remove console logs

* fix mypy tests

Author: jugaldb

enterprise/enterprise_hooks/aporia_ai.py

@@ -173,6 +173,7 @@ async def async_moderation_hook(
             "moderation",
             "audio_transcription",
             "responses",
+            "mcp_call",
         ],
     ):
         from litellm.proxy.common_utils.callback_utils import (

enterprise/enterprise_hooks/google_text_moderation.py

@@ -95,6 +95,7 @@ async def async_moderation_hook(
             "moderation",
             "audio_transcription",
             "responses",
+            "mcp_call",
         ],
     ):
         """

enterprise/enterprise_hooks/openai_moderation.py

@@ -42,6 +42,7 @@ async def async_moderation_hook(
             "moderation",
             "audio_transcription",
             "responses",
+            "mcp_call",
         ],
     ):
         text = ""

enterprise/litellm_enterprise/enterprise_callbacks/llama_guard.py

@@ -105,6 +105,7 @@ async def async_moderation_hook(
             "moderation",
             "audio_transcription",
             "responses",
+            "mcp_call",
         ],
     ):
         """

enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py

@@ -127,6 +127,7 @@ async def async_moderation_hook(
             "moderation",
             "audio_transcription",
             "responses",
+            "mcp_call",
         ],
     ):
         """

enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py

@@ -147,6 +147,7 @@ async def async_pre_call_hook(
             "audio_transcription",
             "pass_through_endpoint",
             "rerank",
+            "mcp_call",
         ],
     ) -> Optional[Union[Exception, str, dict]]:
         """

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -290,6 +290,7 @@ async def async_pre_call_hook(
             "aretrieve_fine_tuning_job",
             "alist_fine_tuning_jobs",
             "acancel_fine_tuning_job",
+            "mcp_call",
         ],
     ) -> Union[Exception, str, Dict, None]:
         """

litellm/integrations/custom_guardrail.py

@@ -234,7 +234,6 @@ def should_run_guardrail(
         Returns True if the guardrail should be run on the event_type
         """
         requested_guardrails = self.get_guardrail_from_metadata(data)
-
         verbose_logger.debug(
             "inside should_run_guardrail for guardrail=%s event_type= %s guardrail_supported_event_hooks= %s requested_guardrails= %s self.default_on= %s",
             self.guardrail_name,
@@ -243,7 +242,6 @@ def should_run_guardrail(
             requested_guardrails,
             self.default_on,
         )
-
         if self.default_on is True:
             if self._event_hook_is_event_type(event_type):
                 if isinstance(self.event_hook, Mode):
@@ -287,7 +285,6 @@ def should_run_guardrail(
             )
             if result is not None:
                 return result
-
         return True
 
     def _event_hook_is_event_type(self, event_type: GuardrailEventHooks) -> bool:

litellm/integrations/custom_logger.py

@@ -281,6 +281,7 @@ async def async_pre_call_hook(
             "audio_transcription",
             "pass_through_endpoint",
             "rerank",
+            "mcp_call",
         ],
     ) -> Optional[
         Union[Exception, str, dict]
@@ -327,6 +328,7 @@ async def async_moderation_hook(
             "moderation",
             "audio_transcription",
             "responses",
+            "mcp_call",
         ],
     ) -> Any:
         pass

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -17,6 +17,8 @@
 from mcp.types import Tool as MCPTool
 
 from litellm._logging import verbose_logger
+from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
+from fastapi import HTTPException
 from litellm.experimental_mcp_client.client import MCPClient
 from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
     MCPRequestHandler,
@@ -592,22 +594,22 @@ async def call_tool(
                 "server_name": server_name_from_prefix,
                 "user_api_key_auth": user_api_key_auth,
             }
-            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,
-            )
-            
-            if pre_hook_result:
-                # Check if the call should proceed
-                if not pre_hook_result.get("should_proceed", True):
-                    error_message = pre_hook_result.get("error_message", "Tool call rejected by pre-hook")
-                    raise ValueError(error_message)
+            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,
+                )
 
-                # Apply any argument modifications
-                if pre_hook_result.get("modified_arguments"):
-                    arguments = pre_hook_result["modified_arguments"]
+                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:
+                # Re-raise guardrail exceptions to properly fail the MCP call
+                verbose_logger.error(f"Guardrail blocked MCP tool call pre call: {str(e)}")
+                raise e
 
         # Get server-specific auth header if available
         server_auth_header = None
@@ -627,6 +629,7 @@ async def call_tool(
         )
 
         async with client:
+
             # Use the original tool name (without prefix) for the actual call
             call_tool_params = MCPCallToolRequestParams(
                 name=original_tool_name,
@@ -635,40 +638,39 @@ async def call_tool(
 
             # Initialize during_hook_task as None
             during_hook_task = None
-            
+            tasks = []
             # Start during hook if proxy_logging_obj is available
             if proxy_logging_obj:
-                try:
-                    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,
-                        )
+                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,
                     )
-                except Exception as e:
-                    verbose_logger.warning(f"During hook error (non-blocking): {str(e)}")
-            
-            result = await client.call_tool(call_tool_params)
-            
-            #########################################################
-            # Check during hook result if it completed
-            #########################################################
-            if proxy_logging_obj and during_hook_task is not None:
-                try:
-                    during_hook_result = await during_hook_task
-                    if during_hook_result and not during_hook_result.get("should_continue", True):
-                        error_message = during_hook_result.get("error_message", "Tool call cancelled by during-hook")
-                        raise ValueError(error_message)
-                except Exception as e:
-                    verbose_logger.warning(f"During hook error (non-blocking): {str(e)}")
+                )
+                tasks.append(during_hook_task)
 
-            return result
+
+            tasks.append(asyncio.create_task(client.call_tool(call_tool_params)))
+            try:
+
+                mcp_responses = await asyncio.gather(*tasks)
+
+                # If proxy_logging_obj is None, the tool call result is at index 0
+                # If proxy_logging_obj is not None, the tool call result is at index 1 (after the during hook task)
+                result_index = 1 if proxy_logging_obj else 0
+                result = mcp_responses[result_index]
+                
+                return cast(CallToolResult, result)
+            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 during result check: {str(e)}")
+                    raise e
 
     #########################################################
     # End of Methods that call the upstream MCP servers