fix(mcp): fail closed when guardrail rewrites cannot be mapped to MCP arguments

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

@@ -66,17 +66,15 @@ def _argument_replacements(
 
     Only leaves the guardrail actually rewrote are returned, so a guardrail that
     detects nothing leaves the outbound tool call byte-identical. A guardrail that
-    returns the wrong number of texts is paired against nothing, because a
-    positional write-back would scramble the arguments rather than mask them.
+    returns the wrong number of texts fails closed, because a positional write-back
+    would scramble the arguments rather than mask them.
     """
-    usable: Final = masked_texts if masked_texts is not None and len(masked_texts) == len(argument_leaves) else ()
     if masked_texts is not None and len(masked_texts) != len(argument_leaves):
-        verbose_proxy_logger.warning(
-            "MCP Guardrail: guardrail returned %d texts for %d tool call argument strings; leaving arguments unmasked",
-            len(masked_texts),
-            len(argument_leaves),
+        raise _blocked(
+            f"guardrail returned {len(masked_texts)} texts for {len(argument_leaves)} MCP tool call argument strings, "
+            "so the redaction cannot be mapped back to the arguments"
         )
-    return {path: masked for (path, original), masked in zip(argument_leaves, usable) if masked != original}
+    return {path: masked for (path, original), masked in zip(argument_leaves, masked_texts or ()) if masked != original}
 
 
 def _conflicting_rewrite_paths(
@@ -96,10 +94,12 @@ def _conflicting_rewrite_paths(
     including a payload reshaped so the leaves no longer line up, because the
     write-back is positional and would land a redaction on the wrong value.
     """
+    if tuple(path for path, _ in scanned_leaves) != tuple(path for path, _ in current_leaves):
+        return tuple(replacements)
     return tuple(
         path
-        for (path, scanned), (current_path, current) in zip(scanned_leaves, current_leaves)
-        if path in replacements and (current_path != path or current not in (scanned, replacements[path]))
+        for (path, scanned), (_, current) in zip(scanned_leaves, current_leaves)
+        if path in replacements and current not in (scanned, replacements[path])
     )
 
 
@@ -242,12 +242,10 @@ class MCPGuardrailTranslationHandler(BaseTranslation):
         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),
+            raise _blocked(
+                f"guardrail returned {len(masked_texts)} texts for {len(originals)} MCP tool result texts, "
+                "so the redaction cannot be mapped back to the result"
             )
-            return response
 
         split: Final = len(text_blocks)
         if content is not None:

tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py

@@ -221,18 +221,19 @@ async def test_clean_arguments_are_not_overridden():
 
 
 @pytest.mark.asyncio
-async def test_guardrail_returning_wrong_text_count_leaves_arguments_alone():
-    """Write-back is positional, so a length mismatch must not scramble arguments."""
+async def test_guardrail_returning_wrong_text_count_blocks_the_call():
+    """Write-back is positional, so a length mismatch must block the call."""
     handler = MCPGuardrailTranslationHandler()
     guardrail = ArgumentMaskingGuardrail(texts_override=["only", "two", "texts"])
 
     arguments = {"query": "contact jane.doe@example.com about the invoice"}
     data = {"mcp_tool_name": "search", "mcp_arguments": arguments}
 
-    result = await handler.process_input_messages(data, guardrail)
+    with pytest.raises(HTTPException) as exc_info:
+        await handler.process_input_messages(data, guardrail)
 
-    assert "modified_arguments" not in result
-    assert result["mcp_arguments"] == arguments
+    assert exc_info.value.status_code == 400
+    assert "modified_arguments" not in data
 
 
 @pytest.mark.asyncio
@@ -325,6 +326,23 @@ async def test_arguments_reshaped_under_the_guardrail_fail_closed():
     assert exc_info.value.status_code == 400
 
 
+@pytest.mark.asyncio
+async def test_arguments_shortened_under_the_guardrail_fail_closed():
+    handler = MCPGuardrailTranslationHandler()
+    guardrail = ReshapingGuardrail({"padding": "jane.doe@example.com"})
+
+    data = {
+        "mcp_tool_name": "search",
+        "mcp_arguments": {"padding": "x", "secret": "jane.doe@example.com"},
+    }
+
+    with pytest.raises(HTTPException) as exc_info:
+        await handler.process_input_messages(data, guardrail)
+
+    assert exc_info.value.status_code == 400
+    assert "jane.doe@example.com" not in str(data.get("modified_arguments"))
+
+
 @pytest.mark.asyncio
 async def test_a_renamed_argument_key_blocks_rather_than_dropping_the_mask():
     """The leak this closes: same text, new path, so the write-back would find nothing.
@@ -586,8 +604,8 @@ async def test_process_output_response_handles_result_without_content():
 
 
 @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."""
+async def test_process_output_response_blocks_on_text_count_mismatch():
+    """A guardrail returning the wrong number of texts must block the result."""
     handler = MCPGuardrailTranslationHandler()
     guardrail = MaskingGuardrail(masked_texts=["<EMAIL_ADDRESS>"])
     result = CallToolResult(
@@ -598,9 +616,10 @@ async def test_process_output_response_leaves_result_unmasked_on_text_count_mism
         isError=False,
     )
 
-    returned = await handler.process_output_response(response=result, guardrail_to_apply=guardrail)
+    with pytest.raises(HTTPException) as exc_info:
+        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"]
+    assert exc_info.value.status_code == 400
 
 
 class SubstitutingGuardrail(CustomGuardrail):

type-discipline-budget.json

@@ -27,7 +27,7 @@
     "limit": 0
   },
   "LIT010": {
-    "limit": 16468
+    "limit": 16466
   },
   "LIT011": {
     "limit": 5514