Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py: 16%
89 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2MCP Guardrail Handler for Unified Guardrails.
4Converts an MCP call_tool (name + arguments) into the OpenAI-compatible shape
5apply_guardrail expects: the tool as a single-entry ``tools`` definition, and
6every string leaf of the call arguments as ``texts`` so text guardrails can
7detect and mask sensitive values in the payload. Works with the synthetic
8request from ProxyLogging._convert_mcp_to_llm_format.
10Note: For MCP tool definitions (schema) -> OpenAI tools=[], see
11litellm.experimental_mcp_client.tools.transform_mcp_tool_to_openai_tool
12when you have a full MCP Tool from list_tools. Here we only have the call
13payload (name + arguments) so we just build the tool definition.
14"""
16from collections.abc import Mapping, Sequence
17from typing import TYPE_CHECKING, Any, Final
19from fastapi import HTTPException
20from mcp.types import Tool as MCPTool
22from litellm._logging import verbose_proxy_logger
23from litellm.experimental_mcp_client.tools import transform_mcp_tool_to_openai_tool
24from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
25from litellm.proxy._experimental.mcp_server.utils import (
26 MAX_STRUCTURED_CONTENT_SCAN_DEPTH,
27 JSONLeafPath,
28 json_string_leaves,
29 json_unrewritable_labels,
30 mcp_content_item_text,
31 mcp_tool_result_content_list,
32 mcp_tool_result_structured_content,
33 set_mcp_tool_result_structured_content,
34 with_json_string_leaves,
35 with_mcp_content_item_text,
36)
37from litellm.types.llms.openai import (
38 ChatCompletionToolParam,
39 ChatCompletionToolParamFunctionChunk,
40)
41from litellm.types.utils import GenericGuardrailAPIInputs
43if TYPE_CHECKING: 43 ↛ 44line 43 didn't jump to line 44 because the condition on line 43 was never true
44 from mcp.types import CallToolResult
46 from litellm.integrations.custom_guardrail import CustomGuardrail
47 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
50def _blocked(reason: str) -> HTTPException:
51 return HTTPException(status_code=400, detail={"error": f"Content blocked: {reason}"})
54def _too_deeply_nested() -> HTTPException:
55 return _blocked(
56 f"MCP tool call arguments exceed the maximum nesting depth of {MAX_STRUCTURED_CONTENT_SCAN_DEPTH} "
57 "and cannot be scanned by the configured guardrail"
58 )
61def _argument_replacements(
62 argument_leaves: tuple[tuple[JSONLeafPath, str], ...],
63 masked_texts: Sequence[str] | None,
64) -> Mapping[JSONLeafPath, str]:
65 """Positionally pair the guardrail's returned texts with the leaves they came from.
67 Only leaves the guardrail actually rewrote are returned, so a guardrail that
68 detects nothing leaves the outbound tool call byte-identical. A guardrail that
69 returns the wrong number of texts fails closed, because a positional write-back
70 would scramble the arguments rather than mask them.
71 """
72 if masked_texts is not None and len(masked_texts) != len(argument_leaves):
73 raise _blocked(
74 f"guardrail returned {len(masked_texts)} texts for {len(argument_leaves)} MCP tool call argument strings, "
75 "so the redaction cannot be mapped back to the arguments"
76 )
77 return {path: masked for (path, original), masked in zip(argument_leaves, masked_texts or ()) if masked != original}
80def _conflicting_rewrite_paths(
81 scanned_leaves: tuple[tuple[JSONLeafPath, str], ...],
82 current_leaves: tuple[tuple[JSONLeafPath, str], ...],
83 replacements: Mapping[JSONLeafPath, str],
84) -> tuple[JSONLeafPath, ...]:
85 """Paths another guardrail already rewrote differently from what this one wants.
87 Guardrails opted into ``run_in_parallel`` all scan the same payload snapshot, so
88 each one returns a full replacement string derived from the *original* leaf. Two
89 of them rewriting one leaf to different values cannot be merged: writing either
90 result discards the other guardrail's redaction. A leaf still holding the text
91 this guardrail was handed, or already holding this guardrail's own replacement,
92 is safe to write; the latter is how a guardrail that masks the arguments itself
93 as well as through ``texts`` gets there first. Anything else fails closed,
94 including a payload reshaped so the leaves no longer line up, because the
95 write-back is positional and would land a redaction on the wrong value.
96 """
97 if tuple(path for path, _ in scanned_leaves) != tuple(path for path, _ in current_leaves):
98 return tuple(replacements)
99 return tuple(
100 path
101 for (path, scanned), (_, current) in zip(scanned_leaves, current_leaves)
102 if path in replacements and current not in (scanned, replacements[path])
103 )
106def _conflicting_rewrite(paths: tuple[JSONLeafPath, ...]) -> HTTPException:
107 return _blocked(
108 "two guardrails running concurrently rewrote the same MCP tool call "
109 f"argument{'s' if len(paths) > 1 else ''} "
110 f"({', '.join('.'.join(str(part) for part in path) for path in paths)}); "
111 "their redactions cannot be merged. Remove run_in_parallel from one of them so they "
112 "run in sequence."
113 )
116class MCPGuardrailTranslationHandler(BaseTranslation):
117 """Guardrail translation handler for MCP tool calls (passes a single tool_call to guardrail)."""
119 async def process_input_messages(
120 self,
121 data: dict[str, Any],
122 guardrail_to_apply: "CustomGuardrail",
123 litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
124 ) -> dict[str, Any]:
125 mcp_tool_name: Final = data.get("mcp_tool_name") or data.get("name")
126 mcp_arguments: Final[object] = data.get("mcp_arguments") or data.get("arguments")
127 mcp_tool_description: Final = data.get("mcp_tool_description") or data.get("description")
129 if not mcp_tool_name:
130 verbose_proxy_logger.debug("MCP Guardrail: mcp_tool_name missing")
131 return data
133 # Convert MCP input via transform_mcp_tool_to_openai_tool, then map to litellm
134 # ChatCompletionToolParam (openai SDK type has incompatible strict/cache_control).
135 mcp_tool: Final = MCPTool(
136 name=mcp_tool_name,
137 description=mcp_tool_description or "",
138 input_schema={}, # mutable-ok: call payload has no schema; guardrail gets args from request_data
139 )
140 openai_tool: Final = transform_mcp_tool_to_openai_tool(mcp_tool)
141 fn: Final = openai_tool["function"]
142 tool_def: Final[ChatCompletionToolParam] = {
143 "type": "function",
144 "function": ChatCompletionToolParamFunctionChunk(
145 name=fn["name"],
146 description=fn.get("description") or "",
147 parameters=fn.get("parameters")
148 or {
149 "type": "object",
150 "properties": {},
151 "additionalProperties": False,
152 },
153 strict=fn.get("strict", False) or False, # Default to False if None
154 ),
155 }
156 argument_leaves: Final = json_string_leaves(mcp_arguments)
157 if argument_leaves is None:
158 raise _too_deeply_nested()
159 inputs: Final[GenericGuardrailAPIInputs] = GenericGuardrailAPIInputs(
160 tools=[tool_def],
161 texts=[text for _, text in argument_leaves],
162 )
164 guarded: Final = await guardrail_to_apply.apply_guardrail(
165 inputs=inputs,
166 request_data=data,
167 input_type="request",
168 logging_obj=litellm_logging_obj,
169 )
170 replacements: Final = _argument_replacements(
171 argument_leaves=argument_leaves,
172 masked_texts=guarded.get("texts") if guarded else None,
173 )
174 if not replacements:
175 return data
177 current_arguments: Final[object] = data.get("mcp_arguments") or data.get("arguments")
178 current_leaves: Final = json_string_leaves(current_arguments)
179 if current_leaves is None:
180 raise _too_deeply_nested()
181 conflicting: Final = _conflicting_rewrite_paths(argument_leaves, current_leaves, replacements)
182 if conflicting:
183 raise _conflicting_rewrite(conflicting)
184 masked_arguments: Final = with_json_string_leaves(current_arguments, replacements)
185 data["mcp_arguments"] = masked_arguments # rebind-ok: preserve the mask for the outbound MCP call
186 data["modified_arguments"] = masked_arguments # rebind-ok: expose the applied mask to the caller
187 return data
189 async def process_output_response(
190 self,
191 response: "CallToolResult",
192 guardrail_to_apply: "CustomGuardrail",
193 litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
194 user_api_key_dict: Any | None = None,
195 request_data: dict | None = None,
196 ) -> Any:
197 """Scan the text content of an MCP tool result and write masked text back.
199 The content list is rewritten in place (only the entries the guardrail
200 actually changed) rather than returned as a new result: the same object is
201 already referenced by the logging payload captured before this hook runs,
202 so a copy would leave the unmasked text in the spend log / span. A
203 guardrail that rejects the result raises, and the exception propagates to
204 the caller.
206 ``structuredContent`` is scanned and masked too, in the same
207 ``apply_guardrail`` call: it is serialized to the client alongside
208 ``content``, so a value living only there would otherwise reach the
209 client unscanned.
210 """
211 content: Final = mcp_tool_result_content_list(response)
212 text_blocks: Final = (
213 tuple(
214 (index, text) for index, item in enumerate(content) if (text := mcp_content_item_text(item)) is not None
215 )
216 if content is not None
217 else ()
218 )
220 structured: Final = mcp_tool_result_structured_content(response)
221 structured_leaves: Final = json_string_leaves(structured) if structured is not None else ()
222 structured_labels: Final = json_unrewritable_labels(structured) if structured is not None else ()
223 if structured_leaves is None or structured_labels is None:
224 raise _blocked(
225 "MCP tool result structuredContent is nested too deeply to be scanned by the configured guardrail"
226 )
228 if not text_blocks and not structured_leaves and not structured_labels:
229 verbose_proxy_logger.debug("MCP Guardrail: tool result has no scannable text, nothing to do")
230 return response
232 originals: Final = (
233 tuple(text for _, text in text_blocks) + tuple(text for _, text in structured_leaves) + structured_labels
234 )
235 guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
236 inputs=GenericGuardrailAPIInputs(texts=list(originals)),
237 request_data=request_data if request_data is not None else {},
238 input_type="response",
239 logging_obj=litellm_logging_obj,
240 )
241 masked_texts: Final = guardrailed_inputs.get("texts") if guardrailed_inputs else None
242 if masked_texts is None:
243 return response
244 if len(masked_texts) != len(originals):
245 raise _blocked(
246 f"guardrail returned {len(masked_texts)} texts for {len(originals)} MCP tool result texts, "
247 "so the redaction cannot be mapped back to the result"
248 )
250 split: Final = len(text_blocks)
251 if content is not None:
252 for (index, original), masked in zip(text_blocks, masked_texts[:split]):
253 if masked != original:
254 content[index] = with_mcp_content_item_text(content[index], masked)
256 label_start: Final = split + len(structured_leaves)
257 if any(masked != original for original, masked in zip(structured_labels, masked_texts[label_start:])):
258 raise _blocked(
259 "MCP tool result matched a masking rule on a non-rewritable field "
260 "(a structuredContent key or numeric value), which cannot be redacted without changing "
261 "the payload contract"
262 )
264 structured_replacements: Final = {
265 path: masked
266 for (path, original), masked in zip(structured_leaves, masked_texts[split:label_start])
267 if masked != original
268 }
269 if structured_replacements:
270 set_mcp_tool_result_structured_content(
271 response, with_json_string_leaves(structured, structured_replacements)
272 )
273 return response