diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 108928871b0c..8b831b55da3e 100644 --- a/litellm/integrations/custom_logger.py +++ b/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 diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py index b668833e638b..909925da00a8 100644 --- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py +++ b/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 diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index af3d966c95b1..9b51513f4ac2 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/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,7 +220,7 @@ 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, @@ -204,7 +228,6 @@ async def _handle_virtual_mcp_tool( user_api_key_auth=user_api_key_dict, request_data=data, ) - return result def _get_server_auth_header( server, @@ -998,7 +1021,7 @@ 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, @@ -1006,7 +1029,6 @@ async def call_tool_rest_api( 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", diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 14673cf12c13..06a3a5a61e4b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/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, diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 80a469b8c1a5..afd396adc4cc 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -4,6 +4,7 @@ import json import re +from collections.abc import MutableMapping, MutableSequence from typing import ( Any, Dict, @@ -434,6 +435,56 @@ def extract_mcp_tool_result_error_message(result: object) -> Optional[str]: return "MCP tool call returned isError=true" +def mcp_tool_result_content_list(result: object) -> MutableSequence[object] | None: # mutable-ok: see below + """The mutable content list of an MCP tool result, or ``None`` when it has none. + + Deliberately mutable: a guardrail masking the result rewrites entries in place, + because the logging payload captured before the guardrail runs references this + same list, so handing back a copy would leave the unmasked text in the spend log + and the OTel span. + + Accepts both ``mcp.types.CallToolResult`` objects and their dict + equivalents, duck-typed so the ``mcp`` package is not required. + """ + content: object = result.get("content") if isinstance(result, Mapping) else getattr(result, "content", None) + if isinstance(content, MutableSequence): + return content + return None + + +def mcp_content_item_text(item: object) -> str | None: + """The ``text`` of a rewritable MCP content item, or ``None``. + + Only mappings and Pydantic-style models report a text, because those are the + only shapes ``with_mcp_content_item_text`` can rewrite; a caller therefore + never reads text it would be unable to write back (e.g. masked by a + guardrail). Non-text content (images, embedded resources) has no ``text`` + and is reported as ``None``. + """ + text: object + if isinstance(item, Mapping): + text = item.get("text") + elif callable(getattr(item, "model_copy", None)): + text = getattr(item, "text", None) + else: + return None + return text if isinstance(text, str) else None + + +def with_mcp_content_item_text(item: object, text: str) -> object: + """A copy of an MCP content item carrying ``text`` instead of its own. + + Only meaningful for items ``mcp_content_item_text`` returned a text for; any + other item is returned unchanged. + """ + if isinstance(item, Mapping): + return {**item, "text": text} + model_copy = getattr(item, "model_copy", None) + if callable(model_copy): + return model_copy(update={"text": text}) + return item + + TOOL_DISPLAY_NAME_PATTERN = re.compile(r"^[a-zA-Z0-9_-]+$") @@ -618,3 +669,112 @@ def merge_mcp_headers( merged.update({str(k): str(v) for k, v in static_headers.items()}) return merged or None + + +# Local rather than litellm.constants: this module deliberately imports no litellm +# package, so pulling one in for a single integer would drag in litellm/__init__. +MAX_STRUCTURED_CONTENT_SCAN_DEPTH = 100 + + +JSONLeafPath = tuple[str | int, ...] + + +def _flatten_leaf_groups( + groups: Iterable[tuple[tuple[JSONLeafPath, str], ...] | None], +) -> tuple[tuple[JSONLeafPath, str], ...] | None: + """Concatenate child leaf groups, propagating the too-deep sentinel.""" + materialized = tuple(groups) + if any(group is None for group in materialized): + return None + return tuple(leaf for group in materialized if group is not None for leaf in group) + + +def json_string_leaves(value: object, path: JSONLeafPath = ()) -> tuple[tuple[JSONLeafPath, str], ...] | None: + """Depth-first, deterministically ordered string leaves of a JSON value. + + Returns ``None`` when the value is nested past ``MAX_STRUCTURED_CONTENT_SCAN_DEPTH``, + so the caller blocks rather than letting deeper values through unscanned; an + empty tuple means there was simply nothing to scan. A sentinel rather than an + exception because this module is reloaded by tests (see the note above the + environment-backed constants), which would give a custom exception class a new + identity and let it escape a caller's ``except``. + """ + if len(path) > MAX_STRUCTURED_CONTENT_SCAN_DEPTH: + return None + if isinstance(value, str): + return ((path, value),) + if isinstance(value, dict): + return _flatten_leaf_groups(json_string_leaves(item, (*path, key)) for key, item in value.items()) + if isinstance(value, list): + return _flatten_leaf_groups(json_string_leaves(item, (*path, index)) for index, item in enumerate(value)) + return () + + +def with_json_string_leaves( + value: object, + replacements: Mapping[JSONLeafPath, str], + path: JSONLeafPath = (), +) -> object: + """Rebuild a JSON value with the guardrail's rewritten string leaves.""" + if isinstance(value, str): + return replacements.get(path, value) + if isinstance(value, dict): + return {key: with_json_string_leaves(item, replacements, (*path, key)) for key, item in value.items()} + if isinstance(value, list): + return [with_json_string_leaves(item, replacements, (*path, index)) for index, item in enumerate(value)] + return value + + +def json_unrewritable_labels(value: object, path_depth: int = 0) -> tuple[str, ...] | None: + """Strings in a JSON value that carry meaning but cannot be rewritten. + + Dictionary keys and non-string scalars: masking either would change the + payload's contract rather than redact a value, so a caller scans these and + blocks on a match instead of rewriting, matching what the content filter + already does for MCP tool call arguments. ``None`` means the value is nested + past the scan depth, same contract as ``json_string_leaves``. + """ + if path_depth > MAX_STRUCTURED_CONTENT_SCAN_DEPTH: + return None + if isinstance(value, bool) or value is None or isinstance(value, str): + return () + if isinstance(value, (int, float)): + return (str(value),) + if isinstance(value, dict): + own = tuple(key for key in value if isinstance(key, str)) + nested = tuple(json_unrewritable_labels(item, path_depth + 1) for item in value.values()) + if any(group is None for group in nested): + return None + return own + tuple(label for group in nested if group is not None for label in group) + if isinstance(value, list): + nested = tuple(json_unrewritable_labels(item, path_depth + 1) for item in value) + if any(group is None for group in nested): + return None + return tuple(label for group in nested if group is not None for label in group) + return () + + +def mcp_tool_result_structured_content(result: object) -> object: + """The ``structuredContent`` of an MCP tool result, or ``None`` when it has none.""" + if isinstance(result, Mapping): + return result.get("structuredContent") + return getattr(result, "structuredContent", None) + + +def set_mcp_tool_result_structured_content(result: object, value: object) -> bool: + """Replace ``structuredContent`` in place; ``False`` when the shape does not carry it. + + In place for the same reason the content list is: the logging payload captured + before the guardrail ran references this object, so a copy would leave the + unmasked value in the spend log and the OTel span. + """ + if isinstance(result, MutableMapping): + result["structuredContent"] = value + return True + if not hasattr(result, "structuredContent"): + return False + try: + setattr(result, "structuredContent", value) # attribute name is fixed by the MCP result shape + return True + except (AttributeError, TypeError, ValueError): + return False diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index a0c822964a05..60e1947b752d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -75,6 +75,7 @@ def get_supported_event_hooks(cls) -> List[GuardrailEventHooks]: GuardrailEventHooks.post_call, GuardrailEventHooks.logging_only, GuardrailEventHooks.pre_mcp_call, + GuardrailEventHooks.post_mcp_call, ] # Class variables or attributes diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 924189fed4bb..1e80f15a7610 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -109,6 +109,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.llms import load_guardrail_translation_mappings from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( AlertType, @@ -172,6 +173,7 @@ from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams if TYPE_CHECKING: + from mcp.types import CallToolResult from opentelemetry.trace import Span as _Span from prisma.client import TransactionManager @@ -2470,6 +2472,59 @@ async def _run_one(callback: CustomGuardrail) -> None: if raised: raise raised[0] + async def post_mcp_call_hook( + self, + response: "CallToolResult", + request_data: Mapping[str, Any], + user_api_key_dict: UserAPIKeyAuth | None = None, + ) -> "CallToolResult": + """ + Run guardrails configured for ``post_mcp_call`` against an MCP tool result. + + The MCP counterpart of ``post_call_success_hook``: guardrails that + implement ``apply_guardrail`` see the tool result's text through the + unified guardrail seam (``MCPGuardrailTranslationHandler``), so a text + guardrail can mask sensitive values in the result without any MCP-specific + code of its own. Guardrails that instead implement + ``async_post_mcp_tool_call_hook`` are dispatched by + ``Logging.async_post_mcp_tool_call_hook`` and are not run here. + + A guardrail that rejects the result raises, and the exception propagates + (matching the inbound ``pre_mcp_call`` behavior) rather than being + swallowed into an unguarded result. + """ + caps = ProxyLogging._callback_capabilities() + if not caps.has_guardrail: + return response + + handler_cls = load_guardrail_translation_mappings().get(CallTypes.call_mcp_tool) + if handler_cls is None: + verbose_proxy_logger.debug("MCP guardrail translation handler unavailable; skipping post_mcp_call hook") + return response + + for callback in caps.resolved_callbacks: + if not isinstance(callback, CustomGuardrail): + continue + if "apply_guardrail" not in type(callback).__dict__: + continue + if ( + callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_mcp_call) + is not True + ): + continue + response = await self._run_guardrail_with_metrics( + callback, + handler_cls().process_output_response( + response=response, + guardrail_to_apply=callback, + litellm_logging_obj=request_data.get("litellm_logging_obj"), + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ), + "post_mcp_call", + ) + return response + async def post_call_response_headers_hook( self, data: dict, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index cb680dc8b865..9d30e40cd54d 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -795,20 +795,33 @@ async def _execute_tool_calls( proxy_logging_obj=proxy_logging_obj, ) + if proxy_logging_obj: + result = await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data=( + litellm_logging_obj.model_call_details + if litellm_logging_obj + else {"mcp_tool_name": tool_name} + ), + user_api_key_dict=user_api_key_auth, + ) + if litellm_logging_obj: try: litellm_logging_obj.post_call(original_response=result) - end_time = datetime.now() await litellm_logging_obj.async_post_mcp_tool_call_hook( kwargs=litellm_logging_obj.model_call_details, response_obj=result, start_time=start_time, - end_time=end_time, + end_time=datetime.now(), ) + except Exception: + verbose_logger.exception("Failed to run post-call logging for MCP tool call %s", tool_name) + try: await litellm_logging_obj.async_success_handler( result=result, start_time=start_time, - end_time=end_time, + end_time=datetime.now(), ) except Exception: verbose_logger.exception("Failed to log MCP tool call success for %s", tool_name) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index a324e71e2899..af419d8cb6ff 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -1043,6 +1043,7 @@ class GuardrailEventHooks(str, Enum): logging_only = "logging_only" pre_mcp_call = "pre_mcp_call" during_mcp_call = "during_mcp_call" + post_mcp_call = "post_mcp_call" realtime_input_transcription = "realtime_input_transcription" diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 0bc3cebdd5a1..244e17b46a1e 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -57,6 +57,9 @@ "_filter_mcp_argument_value", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by blocking the MCP call at the cap. "_redact_scanned_content", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by returning "[REDACTED]" at the cap. "_iter_fallback_targets", # max depth set (2 * ROUTER_MAX_FALLBACKS); fails closed by raising ValueError at the cap. + "json_string_leaves", # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); fails closed by raising at the cap so nothing goes unscanned. + "with_json_string_leaves", # transitively bounded: only runs on a tree json_string_leaves already walked under the cap. + "json_unrewritable_labels", # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); returns the None sentinel at the cap so the caller blocks. ] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py index 5dbad53948b6..2e286a237c4f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py @@ -1,11 +1,14 @@ """Tests for the MCP guardrail translation handler.""" import pytest +from mcp.types import CallToolResult, ImageContent, TextContent +from litellm.exceptions import BlockedPiiEntityError from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( MCPGuardrailTranslationHandler, ) +from litellm.types.utils import GenericGuardrailAPIInputs class MockGuardrail(CustomGuardrail): @@ -80,3 +83,304 @@ async def test_process_input_messages_handles_minimal_data(): tools = guardrail.last_inputs.get("tools", []) assert len(tools) == 1 assert tools[0]["function"]["name"] == "simple_tool" + + +class MaskingGuardrail(CustomGuardrail): + """Guardrail that rewrites every scanned text, recording what it saw.""" + + def __init__(self, masked_texts=None, raises=None): + super().__init__(guardrail_name="masking-mcp-guardrail") + self.masked_texts = masked_texts + self.raises = raises + self.call_count = 0 + self.last_inputs = None + self.last_input_type = None + self.last_request_data = None + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.call_count += 1 + self.last_inputs = inputs + self.last_input_type = input_type + self.last_request_data = request_data + if self.raises is not None: + raise self.raises + if self.masked_texts is None: + return inputs + return GenericGuardrailAPIInputs(texts=list(self.masked_texts)) + + +@pytest.mark.asyncio +async def test_process_output_response_masks_text_content(): + """Masked text returned by the guardrail must land in the tool result.""" + handler = MCPGuardrailTranslationHandler() + guardrail = MaskingGuardrail(masked_texts=["email ", "call "]) + result = CallToolResult( + content=[ + TextContent(type="text", text="email jane@example.com"), + TextContent(type="text", text="call 415-555-0132"), + ], + isError=False, + ) + + returned = await handler.process_output_response( + response=result, + guardrail_to_apply=guardrail, + request_data={"mcp_tool_name": "echo"}, + ) + + assert guardrail.call_count == 1 + assert guardrail.last_input_type == "response" + assert guardrail.last_inputs["texts"] == ["email jane@example.com", "call 415-555-0132"] + assert [item.text for item in returned.content] == ["email ", "call "] + assert [item.text for item in result.content] == ["email ", "call "] + + +@pytest.mark.asyncio +async def test_process_output_response_masks_dict_shaped_result(): + """A dict-shaped tool result (REST/JSON-RPC payload) must be masked too.""" + handler = MCPGuardrailTranslationHandler() + guardrail = MaskingGuardrail(masked_texts=[""]) + result = {"content": [{"type": "text", "text": "jane@example.com"}], "isError": False} + + returned = await handler.process_output_response(response=result, guardrail_to_apply=guardrail) + + assert returned["content"][0]["text"] == "" + assert returned["content"][0]["type"] == "text" + + +@pytest.mark.asyncio +async def test_process_output_response_propagates_block(): + """A guardrail rejecting the tool result must not be swallowed.""" + handler = MCPGuardrailTranslationHandler() + guardrail = MaskingGuardrail( + raises=BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="masking-mcp-guardrail") + ) + result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) + + with pytest.raises(BlockedPiiEntityError): + await handler.process_output_response(response=result, guardrail_to_apply=guardrail) + + +@pytest.mark.asyncio +async def test_process_output_response_skips_non_text_content(): + """A result carrying no text content must not be sent to the guardrail.""" + handler = MCPGuardrailTranslationHandler() + guardrail = MaskingGuardrail(masked_texts=["should not be used"]) + result = CallToolResult( + content=[ImageContent(type="image", data="aGk=", mimeType="image/png")], + isError=False, + ) + + returned = await handler.process_output_response(response=result, guardrail_to_apply=guardrail) + + assert guardrail.call_count == 0 + assert returned is result + + +@pytest.mark.asyncio +async def test_process_output_response_handles_result_without_content(): + """An unexpected result shape must be passed through, not crash the tool call.""" + handler = MCPGuardrailTranslationHandler() + guardrail = MaskingGuardrail(masked_texts=["should not be used"]) + + returned = await handler.process_output_response(response={"error": "boom"}, guardrail_to_apply=guardrail) + + assert guardrail.call_count == 0 + assert returned == {"error": "boom"} + + +@pytest.mark.asyncio +async def test_process_output_response_leaves_result_unmasked_on_text_count_mismatch(): + """A guardrail returning the wrong number of texts must not shuffle content.""" + handler = MCPGuardrailTranslationHandler() + guardrail = MaskingGuardrail(masked_texts=[""]) + result = CallToolResult( + content=[ + TextContent(type="text", text="jane@example.com"), + TextContent(type="text", text="415-555-0132"), + ], + isError=False, + ) + + returned = await handler.process_output_response(response=result, guardrail_to_apply=guardrail) + + assert [item.text for item in returned.content] == ["jane@example.com", "415-555-0132"] + + +class SubstitutingGuardrail(CustomGuardrail): + """Masks one substring wherever it appears, across however many texts it is given.""" + + def __init__(self, needle: str, replacement: str): + super().__init__(guardrail_name="substituting-mcp-guardrail") + self.needle = needle + self.replacement = replacement + self.seen_texts: list = [] + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.seen_texts = list(inputs.get("texts") or []) + return GenericGuardrailAPIInputs( + texts=[text.replace(self.needle, self.replacement) for text in self.seen_texts] + ) + + +@pytest.mark.asyncio +async def test_structured_content_is_masked_alongside_content(): + """structuredContent goes to the client too, so it must be masked, not just content.""" + handler = MCPGuardrailTranslationHandler() + guardrail = SubstitutingGuardrail("jane@example.com", "") + response = CallToolResult( + content=[TextContent(type="text", text="email jane@example.com")], + structuredContent={"contact": {"email": "jane@example.com"}, "balance": 42.0}, + isError=False, + ) + + returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert returned.content[0].text == "email " + assert returned.structuredContent == {"contact": {"email": ""}, "balance": 42.0} + + +@pytest.mark.asyncio +async def test_value_present_only_in_structured_content_is_masked(): + """The gap this closes: a sensitive value that never appears in the text content. + + Scanning only content would hand it to the guardrail never, so it would reach + the client unscanned behind a result that looks inspected. + """ + handler = MCPGuardrailTranslationHandler() + guardrail = SubstitutingGuardrail("jane@example.com", "") + response = CallToolResult( + content=[TextContent(type="text", text="lookup complete")], + structuredContent={"records": [{"email": "jane@example.com"}]}, + isError=False, + ) + + returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert "jane@example.com" in guardrail.seen_texts + assert returned.structuredContent == {"records": [{"email": ""}]} + assert returned.content[0].text == "lookup complete" + + +@pytest.mark.asyncio +async def test_structured_content_without_a_match_is_untouched(): + """Unrelated structured data keeps its values and its types.""" + handler = MCPGuardrailTranslationHandler() + guardrail = SubstitutingGuardrail("jane@example.com", "") + response = CallToolResult( + content=[TextContent(type="text", text="lookup complete")], + structuredContent={"record_id": "C-1001", "balance": 42.0, "active": True, "note": None}, + isError=False, + ) + + returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert returned.structuredContent == {"record_id": "C-1001", "balance": 42.0, "active": True, "note": None} + + +@pytest.mark.asyncio +async def test_structured_content_nested_too_deeply_is_blocked(): + """Too deep to walk must block rather than pass the deeper values unscanned.""" + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server.utils import MAX_STRUCTURED_CONTENT_SCAN_DEPTH + + handler = MCPGuardrailTranslationHandler() + guardrail = SubstitutingGuardrail("jane@example.com", "") + nested: dict = {"leaf": "jane@example.com"} + for _ in range(MAX_STRUCTURED_CONTENT_SCAN_DEPTH + 1): + nested = {"next": nested} + response = CallToolResult( + content=[TextContent(type="text", text="lookup complete")], + structuredContent=nested, + isError=False, + ) + + with pytest.raises(HTTPException) as exc_info: + await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert exc_info.value.status_code == 400 + + +def test_too_deep_json_returns_a_sentinel_rather_than_raising(): + """The too-deep signal must be a return value, not a custom exception. + + mcp_server/utils.py is reloaded by tests that override its environment-backed + constants, which gives any exception class defined there a fresh identity and + lets it escape a caller's except clause; under xdist that surfaced as a failure + in an unrelated shard. A sentinel has no identity to lose. Asserted directly on + the helper so this pins the contract without reloading the module and leaking + that reload into other tests. + """ + from litellm.proxy._experimental.mcp_server.utils import ( + MAX_STRUCTURED_CONTENT_SCAN_DEPTH, + json_string_leaves, + ) + + nested: dict = {"leaf": "jane@example.com"} + for _ in range(MAX_STRUCTURED_CONTENT_SCAN_DEPTH + 1): + nested = {"next": nested} + + assert json_string_leaves(nested) is None + assert json_string_leaves({"a": "b"}) == ((("a",), "b"),) + + +@pytest.mark.asyncio +async def test_sensitive_structured_content_key_is_blocked(): + """A dict key is client-visible but not rewritable, so a match must block. + + Maps keyed by an identifier are a common API shape, and renaming the key would + change the payload contract rather than redact a value; the content filter takes + the same position on MCP tool call arguments. + """ + from fastapi import HTTPException + + handler = MCPGuardrailTranslationHandler() + guardrail = SubstitutingGuardrail("jane@example.com", "") + response = CallToolResult( + content=[TextContent(type="text", text="lookup complete")], + structuredContent={"jane@example.com": {"balance": 42.0}}, + isError=False, + ) + + with pytest.raises(HTTPException) as exc_info: + await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert exc_info.value.status_code == 400 + assert "non-rewritable" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_sensitive_structured_content_numeric_value_is_blocked(): + """A numeric value cannot be masked in place either, so a match must block.""" + from fastapi import HTTPException + + handler = MCPGuardrailTranslationHandler() + guardrail = SubstitutingGuardrail("4155550199", "") + response = CallToolResult( + content=[TextContent(type="text", text="lookup complete")], + structuredContent={"phone": 4155550199}, + isError=False, + ) + + with pytest.raises(HTTPException) as exc_info: + await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_clean_structured_content_keys_do_not_block(): + """Ordinary keys and numbers must pass through untouched.""" + handler = MCPGuardrailTranslationHandler() + guardrail = SubstitutingGuardrail("jane@example.com", "") + response = CallToolResult( + content=[TextContent(type="text", text="email jane@example.com")], + structuredContent={"record_id": "C-1001", "balance": 42.0, "count": 3}, + isError=False, + ) + + returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail) + + assert returned.content[0].text == "email " + assert returned.structuredContent == {"record_id": "C-1001", "balance": 42.0, "count": 3} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 1753b0d92a8d..c0affdf46b35 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -7130,6 +7130,14 @@ def _mock_mcp_logging_obj() -> MagicMock: return logging_obj +def _mock_mcp_proxy_logging() -> MagicMock: + """ProxyLogging stand-in whose post_mcp_call_hook passes the result through.""" + proxy_logging_mock = MagicMock() + proxy_logging_mock.post_call_failure_hook = AsyncMock() + proxy_logging_mock.post_mcp_call_hook = AsyncMock(side_effect=lambda response, **_: response) + return proxy_logging_mock + + def test_extract_mcp_tool_result_error_message(): from litellm.proxy._experimental.mcp_server.utils import ( extract_mcp_tool_result_error_message, @@ -7160,8 +7168,7 @@ async def test_fire_mcp_tool_call_logging_iserror_logs_failure(): from litellm.proxy._experimental.mcp_server.exceptions import MCPToolResultError logging_obj = _mock_mcp_logging_obj() - proxy_logging_mock = MagicMock() - proxy_logging_mock.post_call_failure_hook = AsyncMock() + proxy_logging_mock = _mock_mcp_proxy_logging() user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock): @@ -7199,8 +7206,7 @@ async def test_fire_mcp_tool_call_logging_success_path_unchanged(): ) logging_obj = _mock_mcp_logging_obj() - proxy_logging_mock = MagicMock() - proxy_logging_mock.post_call_failure_hook = AsyncMock() + proxy_logging_mock = _mock_mcp_proxy_logging() result = _call_tool_result(False, "all good") with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock): @@ -7229,8 +7235,7 @@ async def test_fire_mcp_tool_call_logging_iserror_without_auth_skips_failure_hoo ) logging_obj = _mock_mcp_logging_obj() - proxy_logging_mock = MagicMock() - proxy_logging_mock.post_call_failure_hook = AsyncMock() + proxy_logging_mock = _mock_mcp_proxy_logging() with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock): await _fire_mcp_tool_call_logging( @@ -7256,8 +7261,7 @@ async def test_fire_mcp_tool_call_logging_strips_credentials_from_failure_hook() ) logging_obj = _mock_mcp_logging_obj() - proxy_logging_mock = MagicMock() - proxy_logging_mock.post_call_failure_hook = AsyncMock() + proxy_logging_mock = _mock_mcp_proxy_logging() user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") request_data = { "name": "explode", @@ -7528,8 +7532,7 @@ async def test_call_mcp_tool_skips_failure_hook_for_upstream_auth_error(): transport=MCPTransport.http, mcp_info={"server_name": "test_server"}, ) - proxy_logging_mock = MagicMock() - proxy_logging_mock.post_call_failure_hook = AsyncMock() + proxy_logging_mock = _mock_mcp_proxy_logging() user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") with ( @@ -7841,3 +7844,83 @@ async def test_delegate_challenges_only_when_bearer_absent(self): await self._run(delegate, None, has_stored_token=False) assert exc.value.status_code == 401 await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False) + + +@pytest.mark.asyncio +async def test_post_mcp_call_guardrails_return_the_rewritten_result(): + """The result a post_mcp_call guardrail rewrote must be what the caller sends back.""" + from litellm.proxy._experimental.mcp_server.server import ( + _run_post_mcp_call_guardrails, + ) + + logging_obj = _mock_mcp_logging_obj() + raw_result = _call_tool_result(False, "jane@example.com") + masked_result = _call_tool_result(False, "") + proxy_logging_mock = _mock_mcp_proxy_logging() + proxy_logging_mock.post_mcp_call_hook = AsyncMock(return_value=masked_result) + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock): + returned = await _run_post_mcp_call_guardrails( + result=raw_result, + litellm_logging_obj=logging_obj, + user_api_key_auth=UserAPIKeyAuth(api_key="test-key", user_id="test-user"), + request_data={}, + ) + + assert returned is masked_result + hook_kwargs = proxy_logging_mock.post_mcp_call_hook.await_args.kwargs + assert hook_kwargs["response"] is raw_result + assert hook_kwargs["request_data"] is logging_obj.model_call_details + + +@pytest.mark.asyncio +async def test_post_mcp_call_guardrails_run_without_a_logging_object(): + """Enforcement must not depend on logging being configured. + + A tool call dispatched without a litellm_logging_obj (tool search, and any + caller that omits it) would otherwise skip the guardrail entirely and return + the unscanned tool output to the client. + """ + from litellm.proxy._experimental.mcp_server.server import ( + _run_post_mcp_call_guardrails, + ) + + raw_result = _call_tool_result(False, "jane@example.com") + masked_result = _call_tool_result(False, "") + proxy_logging_mock = _mock_mcp_proxy_logging() + proxy_logging_mock.post_mcp_call_hook = AsyncMock(return_value=masked_result) + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock): + returned = await _run_post_mcp_call_guardrails( + result=raw_result, + litellm_logging_obj=None, + user_api_key_auth=UserAPIKeyAuth(api_key="test-key", user_id="test-user"), + request_data={"name": "fetch_record"}, + ) + + assert returned is masked_result + proxy_logging_mock.post_mcp_call_hook.assert_awaited_once() + assert proxy_logging_mock.post_mcp_call_hook.await_args.kwargs["request_data"] == {"name": "fetch_record"} + + +@pytest.mark.asyncio +async def test_post_mcp_call_guardrails_propagate_a_block(): + """A post_mcp_call guardrail rejection must propagate instead of returning the result.""" + from litellm.exceptions import BlockedPiiEntityError + from litellm.proxy._experimental.mcp_server.server import ( + _run_post_mcp_call_guardrails, + ) + + proxy_logging_mock = _mock_mcp_proxy_logging() + proxy_logging_mock.post_mcp_call_hook = AsyncMock( + side_effect=BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="presidio-mcp") + ) + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock): + with pytest.raises(BlockedPiiEntityError): + await _run_post_mcp_call_guardrails( + result=_call_tool_result(False, "jane@example.com"), + litellm_logging_obj=_mock_mcp_logging_obj(), + user_api_key_auth=None, + request_data={}, + ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 5c9612a055e6..cec62f79e335 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1710,6 +1710,89 @@ async def fake_execute_mcp_tool(**kwargs): assert captured["allowed_mcp_servers"] == [stub_server] fire_logging.assert_awaited_once() + async def test_returns_guardrail_rewritten_tool_result(self, monkeypatch): + """A post_mcp_call guardrail rewrite of the tool result must reach the REST caller, + not the raw result the upstream server returned.""" + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + class StubServer: + server_id = "server-1" + alias = "server-1" + server_name = "server-1" + name = "stub" + allowed_tools = None + mcp_info = {"server_name": "stub"} + available_on_public_internet = True + auth_type = None + + stub_server = StubServer() + + async def fake_add_litellm_data_to_request(**kwargs): + return kwargs.get("data", {}) + + async def fake_execute_mcp_tool(**kwargs): + return {"content": [{"type": "text", "text": "jane@example.com"}]} + + monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.add_litellm_data_to_request", + fake_add_litellm_data_to_request, + raising=False, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {}, raising=False) + monkeypatch.setattr(rest_endpoints, "execute_mcp_tool", fake_execute_mcp_tool, raising=False) + masked_result = {"content": [{"type": "text", "text": ""}]} + monkeypatch.setattr( + rest_endpoints, + "_fire_mcp_tool_call_logging", + AsyncMock(return_value=masked_result), + raising=False, + ) + + request = _build_request( + path="/mcp-rest/tools/call", + method="POST", + json_body={"server_id": "server-1", "name": "demo-tool", "arguments": {}}, + ) + + result = await rest_endpoints.call_tool_rest_api(request, user_api_key_dict=UserAPIKeyAuth()) + + assert result == masked_result + + async def test_success_logging_guardrail_rejection_propagates(self, monkeypatch): + """A guardrail rejecting the tool result must not be swallowed as a logging failure, + otherwise the unguarded result would still be returned to the caller.""" + from litellm.exceptions import BlockedPiiEntityError + + fire_logging = AsyncMock( + side_effect=BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="presidio-mcp") + ) + monkeypatch.setattr(rest_endpoints, "_fire_mcp_tool_call_logging", fire_logging, raising=False) + + with pytest.raises(BlockedPiiEntityError): + await rest_endpoints._safe_fire_mcp_tool_call_logging( + object(), {"result": "ok"}, datetime.now(), datetime.now() + ) + + fire_logging.assert_awaited_once() + @pytest.mark.parametrize("upstream_status", [401, 403]) async def test_call_tool_rest_relays_upstream_auth_failure(self, monkeypatch, upstream_status): """A pass-through call that hits an upstream 401/403 (surfaced by the manager as diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 4673807a135f..3421751d9621 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -7,8 +7,10 @@ from fastapi import HTTPException from litellm.caching.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import ProxyErrorTypes from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks sys.path.insert( 0, os.path.abspath("../../..") @@ -947,3 +949,139 @@ async def test_starttls_uses_verified_context(self, monkeypatch): assert isinstance(context, ssl.SSLContext) assert context.verify_mode == ssl.CERT_REQUIRED assert context.check_hostname is True + + +class _RecordingMCPGuardrail(CustomGuardrail): + """Unified guardrail that masks every text it is handed.""" + + def __init__(self, event_hook, masked_text="", raises=None): + super().__init__(guardrail_name="mcp-output-guardrail", event_hook=event_hook, default_on=True) + self.masked_text = masked_text + self.raises = raises + self.call_count = 0 + self.last_input_type = None + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.call_count += 1 + self.last_input_type = input_type + if self.raises is not None: + raise self.raises + return {"texts": [self.masked_text for _ in inputs.get("texts", [])]} + + +class _NativeMCPGuardrail(CustomGuardrail): + """Guardrail that only implements the MCP logging hook (cisco-style).""" + + def __init__(self): + super().__init__( + guardrail_name="native-mcp-guardrail", + event_hook=GuardrailEventHooks.post_mcp_call, + default_on=True, + ) + self.considered_count = 0 + + def should_run_guardrail(self, data, event_type): + self.considered_count += 1 + return super().should_run_guardrail(data=data, event_type=event_type) + + async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, end_time): + return None + + +@pytest.fixture +def restore_callbacks(): + """Restore the process-wide callback state post_mcp_call_hook reads. + + ProxyLogging caches callback capabilities keyed on id()s of litellm.callbacks, + so a restored-but-different list can collide with a stale entry after GC and + leak a has_guardrail verdict into unrelated tests in the same worker. + """ + original = list(litellm.callbacks) + yield + litellm.callbacks = original + ProxyLogging._callback_capabilities_cache.clear() + + +@pytest.mark.asyncio +async def test_post_mcp_call_hook_masks_tool_result(restore_callbacks): + """A post_mcp_call guardrail must see the tool result text and mask it in the returned result.""" + from mcp.types import CallToolResult, TextContent + + guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call) + litellm.callbacks = [guardrail] + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) + + returned = await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data={"mcp_tool_name": "echo"}, + user_api_key_dict=None, + ) + + assert guardrail.call_count == 1 + assert guardrail.last_input_type == "response" + assert [item.text for item in returned.content] == [""] + + +@pytest.mark.asyncio +async def test_post_mcp_call_hook_skips_guardrail_configured_for_other_hooks(restore_callbacks): + """A guardrail not configured for post_mcp_call must not scan MCP tool results.""" + from mcp.types import CallToolResult, TextContent + + guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_call) + litellm.callbacks = [guardrail] + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) + + returned = await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data={"mcp_tool_name": "echo"}, + user_api_key_dict=None, + ) + + assert guardrail.call_count == 0 + assert [item.text for item in returned.content] == ["jane@example.com"] + + +@pytest.mark.asyncio +async def test_post_mcp_call_hook_skips_guardrail_without_apply_guardrail(restore_callbacks): + """Guardrails that implement async_post_mcp_tool_call_hook are dispatched by the + logging object, so this hook must not run them a second time.""" + from mcp.types import CallToolResult, TextContent + + guardrail = _NativeMCPGuardrail() + litellm.callbacks = [guardrail] + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) + + returned = await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data={"mcp_tool_name": "echo"}, + user_api_key_dict=None, + ) + + assert guardrail.considered_count == 0 + assert [item.text for item in returned.content] == ["jane@example.com"] + + +@pytest.mark.asyncio +async def test_post_mcp_call_hook_propagates_guardrail_block(restore_callbacks): + """A guardrail rejecting the tool result must raise out of the hook.""" + from mcp.types import CallToolResult, TextContent + + from litellm.exceptions import BlockedPiiEntityError + + guardrail = _RecordingMCPGuardrail( + event_hook=GuardrailEventHooks.post_mcp_call, + raises=BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="mcp-output-guardrail"), + ) + litellm.callbacks = [guardrail] + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) + + with pytest.raises(BlockedPiiEntityError): + await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data={"mcp_tool_name": "echo"}, + user_api_key_dict=None, + ) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx index 17331014c577..251ce631beb0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx @@ -41,6 +41,7 @@ const modeDescriptions = { logging_only: "Logging Only - Only runs on logging callbacks without affecting the LLM call", pre_mcp_call: "Before MCP Tool Call - Runs before MCP tool execution and validates tool calls", during_mcp_call: "During MCP Tool Call - Runs in parallel with MCP tool execution for monitoring", + post_mcp_call: "After MCP Tool Call - Runs after MCP tool execution and checks the tool result", }; interface GuardrailPreset {