feat(mcp): scan and mask MCP tool results via post_mcp_call guardrails

litellm/integrations/custom_logger.py

@@ -519,7 +519,13 @@ async def async_post_mcp_tool_call_hook(
         """
         This log gets called after the MCP tool call is made.
 
-        Useful if you want to modiy the standard logging payload after the MCP tool call is made.
+        Useful if you want to modify the standard logging payload after the MCP tool call is made.
+
+        To change what the caller sends back to the MCP client, mutate ``response_obj``
+        in place: every call site discards the returned object, because the
+        dispatcher unwraps it to ``mcp_tool_call_response`` (a raw content list, not
+        a ``CallToolResult``) which the tool-call paths cannot forward. Guardrails
+        that mask or reject tool output should use ``post_mcp_call`` instead.
         """
         return None
 

litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py

@@ -13,11 +13,22 @@
 
 from typing import TYPE_CHECKING, Any, Dict, Optional
 
+from fastapi import HTTPException
 from mcp.types import Tool as MCPTool
 
 from litellm._logging import verbose_proxy_logger
 from litellm.experimental_mcp_client.tools import transform_mcp_tool_to_openai_tool
 from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
+from litellm.proxy._experimental.mcp_server.utils import (
+    json_string_leaves,
+    json_unrewritable_labels,
+    mcp_content_item_text,
+    mcp_tool_result_content_list,
+    mcp_tool_result_structured_content,
+    set_mcp_tool_result_structured_content,
+    with_json_string_leaves,
+    with_mcp_content_item_text,
+)
 from litellm.types.llms.openai import (
     ChatCompletionToolParam,
     ChatCompletionToolParamFunctionChunk,
@@ -92,7 +103,93 @@ async def process_output_response(
         user_api_key_dict: Optional[Any] = None,
         request_data: Optional[dict] = None,
     ) -> Any:
-        verbose_proxy_logger.debug(
-            "MCP Guardrail: Output processing not implemented for MCP tools",
+        """Scan the text content of an MCP tool result and write masked text back.
+
+        The content list is rewritten in place (only the entries the guardrail
+        actually changed) rather than returned as a new result: the same object is
+        already referenced by the logging payload captured before this hook runs,
+        so a copy would leave the unmasked text in the spend log / span. A
+        guardrail that rejects the result raises, and the exception propagates to
+        the caller.
+
+        ``structuredContent`` is scanned and masked too, in the same
+        ``apply_guardrail`` call: it is serialized to the client alongside
+        ``content``, so a value living only there would otherwise reach the
+        client unscanned.
+        """
+        content = mcp_tool_result_content_list(response)
+        text_blocks = (
+            tuple(
+                (index, text) for index, item in enumerate(content) if (text := mcp_content_item_text(item)) is not None
+            )
+            if content is not None
+            else ()
+        )
+
+        structured = mcp_tool_result_structured_content(response)
+        structured_leaves = json_string_leaves(structured) if structured is not None else ()
+        structured_labels = json_unrewritable_labels(structured) if structured is not None else ()
+        if structured_leaves is None or structured_labels is None:
+            raise HTTPException(
+                status_code=400,
+                detail={
+                    "error": (
+                        "Content blocked: MCP tool result structuredContent is nested too deeply to be scanned "
+                        "by the configured guardrail"
+                    )
+                },
+            )
+
+        if not text_blocks and not structured_leaves and not structured_labels:
+            verbose_proxy_logger.debug("MCP Guardrail: tool result has no scannable text, nothing to do")
+            return response
+
+        originals = (
+            tuple(text for _, text in text_blocks) + tuple(text for _, text in structured_leaves) + structured_labels
         )
+        guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
+            inputs=GenericGuardrailAPIInputs(texts=list(originals)),
+            request_data=request_data if request_data is not None else {},
+            input_type="response",
+            logging_obj=litellm_logging_obj,
+        )
+        masked_texts = guardrailed_inputs.get("texts") if guardrailed_inputs else None
+        if masked_texts is None:
+            return response
+        if len(masked_texts) != len(originals):
+            verbose_proxy_logger.warning(
+                "MCP Guardrail: guardrail returned %d texts for %d tool result texts; leaving the result unmasked",
+                len(masked_texts),
+                len(originals),
+            )
+            return response
+
+        split = len(text_blocks)
+        if content is not None:
+            for (index, original), masked in zip(text_blocks, masked_texts[:split]):
+                if masked != original:
+                    content[index] = with_mcp_content_item_text(content[index], masked)
+
+        label_start = split + len(structured_leaves)
+        if any(masked != original for original, masked in zip(structured_labels, masked_texts[label_start:])):
+            raise HTTPException(
+                status_code=400,
+                detail={
+                    "error": (
+                        "Content blocked: MCP tool result matched a masking rule on a non-rewritable field "
+                        "(a structuredContent key or numeric value), which cannot be redacted without changing "
+                        "the payload contract"
+                    )
+                },
+            )
+
+        structured_replacements = {
+            path: masked
+            for (path, original), masked in zip(structured_leaves, masked_texts[split:label_start])
+            if masked != original
+        }
+        if structured_replacements:
+            set_mcp_tool_result_structured_content(
+                response, with_json_string_leaves(structured, structured_replacements)
+            )
         return response

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -12,6 +12,11 @@
 from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
 
 from litellm._logging import verbose_logger
+from litellm.exceptions import (
+    BlockedPiiEntityError,
+    GuardrailRaisedException,
+    ModifyResponseException,
+)
 from litellm.proxy._experimental.mcp_server.exceptions import (
     MCPServerListError,
     MCPUpstreamAuthError,
@@ -33,6 +38,8 @@
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 
 if TYPE_CHECKING:
+    from mcp.types import CallToolResult
+
     from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload
 from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
 from litellm.types.mcp import MCPAuth
@@ -51,6 +58,13 @@
     tags=["mcp"],
 )
 
+_MCP_GUARDRAIL_REJECTIONS = (
+    BlockedPiiEntityError,
+    GuardrailRaisedException,
+    ModifyResponseException,
+    HTTPException,
+)
+
 
 def _connection_error_message(exc: BaseException) -> str:
     if isinstance(exc, httpx.LocalProtocolError):
@@ -99,9 +113,17 @@ async def _safe_fire_mcp_tool_call_logging(
         end_time: datetime,
         user_api_key_auth: UserAPIKeyAuth | None = None,
         request_data: Mapping[str, object] | None = None,
-    ) -> None:
+    ) -> "CallToolResult":
+        """Fire post-call logging, returning the tool result to send to the client.
+
+        ``post_mcp_call`` guardrails already ran on ``execute_mcp_tool``'s return
+        path, so the result arriving here is the guardrailed one. A guardrail
+        rejection raised by a native ``async_post_mcp_tool_call_hook`` is still
+        re-raised rather than swallowed as a logging failure, which would return
+        the unguarded result.
+        """
         if logging_obj is None:
-            return
+            return result
         logging_results = await asyncio.gather(
             _fire_mcp_tool_call_logging(
                 logging_obj,
@@ -113,11 +135,13 @@ async def _safe_fire_mcp_tool_call_logging(
             ),
             return_exceptions=True,
         )
-        logging_error = logging_results[0]
-        if isinstance(logging_error, asyncio.CancelledError):
-            raise logging_error
-        if isinstance(logging_error, BaseException):
-            verbose_logger.warning("MCP tool call logging failed (continuing): %s", logging_error)
+        outcome = logging_results[0]
+        if isinstance(outcome, (asyncio.CancelledError, *_MCP_GUARDRAIL_REJECTIONS)):
+            raise outcome
+        if isinstance(outcome, BaseException):
+            verbose_logger.warning("MCP tool call logging failed (continuing): %s", outcome)
+            return result
+        return outcome
 
     def _relay_upstream_auth_http_exception(e: MCPUpstreamAuthError, request: Request) -> HTTPException:
         """Convert a client-forwarded pass-through upstream 401 into an HTTPException that preserves the
@@ -196,15 +220,14 @@ async def _handle_virtual_mcp_tool(
             raw_headers=virtual_raw_headers,
             litellm_logging_obj=virtual_logging_obj,
         )
-        await _safe_fire_mcp_tool_call_logging(
+        return await _safe_fire_mcp_tool_call_logging(
             virtual_logging_obj,
             result,
             _tool_start_time,
             datetime.now(),
             user_api_key_auth=user_api_key_dict,
             request_data=data,
         )
-        return result
 
     def _get_server_auth_header(
         server,
@@ -998,15 +1021,14 @@ async def call_tool_rest_api(
                 litellm_logging_obj=data.get("litellm_logging_obj"),
                 requested_server_id=canonical_server_id,
             )
-            await _safe_fire_mcp_tool_call_logging(
+            return await _safe_fire_mcp_tool_call_logging(
                 logging_obj,
                 result,
                 _tool_start_time,
                 datetime.now(),
                 user_api_key_auth=user_api_key_dict,
                 request_data=data,
             )
-            return result
         except MCPMissingUserEnvVarsError as e:
             verbose_logger.info(
                 "MCP tool call missing per-user env vars: server_id=%s missing=%s",

litellm/proxy/_experimental/mcp_server/server.py

@@ -2910,7 +2910,38 @@ async def execute_mcp_tool(
             local_content = await _handle_local_mcp_tool(original_tool_name, arguments)
             response = CallToolResult(content=cast(Any, local_content), isError=False)
 
-        return response
+        return await _run_post_mcp_call_guardrails(
+            result=response,
+            litellm_logging_obj=litellm_logging_obj,
+            user_api_key_auth=user_api_key_auth,
+            request_data=kwargs,
+        )
+
+    async def _run_post_mcp_call_guardrails(
+        result: CallToolResult,
+        litellm_logging_obj: LiteLLMLoggingObj | None,
+        user_api_key_auth: UserAPIKeyAuth | None,
+        request_data: Mapping[str, object],
+    ) -> CallToolResult:
+        """Run ``post_mcp_call`` guardrails over an executed tool result.
+
+        Lives on ``execute_mcp_tool``'s return path rather than inside
+        ``_fire_mcp_tool_call_logging`` so enforcement never depends on logging
+        being configured, and so every dispatch route gets it: the MCP protocol
+        handler, the REST endpoint, and tool search all funnel through here.
+        A guardrail that rejects the result raises, matching ``pre_mcp_call``.
+        """
+        from litellm.proxy.proxy_server import proxy_logging_obj
+
+        if proxy_logging_obj is None:
+            return result
+        return await proxy_logging_obj.post_mcp_call_hook(
+            response=result,
+            request_data=(
+                litellm_logging_obj.model_call_details if litellm_logging_obj is not None else dict(request_data)
+            ),
+            user_api_key_dict=user_api_key_auth,
+        )
 
     _MCP_CREDENTIAL_REQUEST_FIELDS = frozenset(
         {
@@ -2929,8 +2960,14 @@ async def _fire_mcp_tool_call_logging(
         end_time: datetime,
         user_api_key_auth: UserAPIKeyAuth | None = None,
         request_data: Mapping[str, object] | None = None,
-    ) -> None:
-        """Fire post-call logging for an executed MCP tool call.
+    ) -> CallToolResult:
+        """Fire post-call logging for an executed MCP tool call, returning the result to send.
+
+        The returned result is what the caller must forward to the client: a
+        ``post_mcp_call`` guardrail may rewrite the tool output (e.g. mask
+        sensitive values) or reject it, in which case its exception propagates.
+        Guardrails run before the success/failure logging so the masked text, not
+        the raw one, is what gets logged.
 
         A result with ``isError=True`` is logged as a failure (``status="failure"``
         payload, so OTel marks the span ERROR) while the HTTP wire behavior stays
@@ -2946,6 +2983,8 @@ async def _fire_mcp_tool_call_logging(
         stripped before the dict is handed to ``post_call_failure_hook``
         callbacks.
         """
+        from litellm.proxy.proxy_server import proxy_logging_obj
+
         logging_obj.post_call(original_response=result)
         await logging_obj.async_post_mcp_tool_call_hook(
             kwargs=logging_obj.model_call_details,
@@ -2957,7 +2996,7 @@ async def _fire_mcp_tool_call_logging(
         error_message = extract_mcp_tool_result_error_message(result)
         if error_message is None:
             await logging_obj.async_success_handler(result=result, start_time=start_time, end_time=end_time)
-            return
+            return result
 
         logging_obj.has_run_logging(event_type="sync_success")
         logging_obj.has_run_logging(event_type="async_success")
@@ -2966,8 +3005,7 @@ async def _fire_mcp_tool_call_logging(
         await logging_obj.async_failure_handler(tool_error, "", start_time, end_time)
 
         if user_api_key_auth is None:
-            return
-        from litellm.proxy.proxy_server import proxy_logging_obj
+            return result
 
         if proxy_logging_obj:
             sanitized_request_data = {
@@ -2979,6 +3017,7 @@ async def _fire_mcp_tool_call_logging(
                 user_api_key_dict=user_api_key_auth,
                 route="/mcp/call_tool",
             )
+        return result
 
     @client
     async def call_mcp_tool(
@@ -3062,7 +3101,7 @@ async def call_mcp_tool(
             raise
 
         if litellm_logging_obj:
-            await _fire_mcp_tool_call_logging(
+            response = await _fire_mcp_tool_call_logging(
                 logging_obj=litellm_logging_obj,
                 result=response,
                 start_time=start_time,