From f4623a1325d4017ec1e47ca27799f2b0bec0be8c Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 6 Jul 2026 17:09:54 -0700 Subject: [PATCH] fix(model_armor): scan MCP tool calls for pre_mcp_call / during_mcp_call modes (#32296) ModelArmorGuardrail.async_pre_call_hook and async_moderation_hook hardcoded their inner should_run_guardrail event type to pre_call / during_call. The central dispatcher already remaps call_mcp_tool -> pre_mcp_call/during_mcp_call and passes the outer gate, but Model Armor's redundant inner gate then rejected MCP calls for a guardrail configured with mode pre_mcp_call/during_mcp_call, so tool-call content was silently skipped. Remap call_mcp_tool -> pre_mcp_call/during_mcp_call in both hooks, matching the existing behavior of the noma and cisco guardrails. Adds regression tests covering both hooks (scan runs on MCP calls, still skipped for chat traffic). Generated with AI Co-Authored-By: Claude Code Co-authored-by: eugene-yao-zocdoc --- .../model_armor/model_armor.py | 5 + .../guardrail_hooks/test_model_armor.py | 124 ++++++++++++++++++ 2 files changed, 129 insertions(+) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index bebd9b28745..8a4e79b31ab 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -39,6 +39,7 @@ from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import ( from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ( + CallTypes, CallTypesLiteral, Choices, GuardrailStatus, @@ -477,6 +478,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): ) event_type = GuardrailEventHooks.pre_call + if call_type == CallTypes.call_mcp_tool.value: + event_type = GuardrailEventHooks.pre_mcp_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data @@ -574,6 +577,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): ) event_type = GuardrailEventHooks.during_call + if call_type == CallTypes.call_mcp_tool.value: + event_type = GuardrailEventHooks.during_mcp_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 64df9ee7ab5..19c200bdaf0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -2940,3 +2940,127 @@ def test_accumulated_responses_are_redactable_as_a_list(): assert "secret-one" not in blob assert "secret-two" not in blob assert blob.count("[REDACTED]") == 2 + + +def _mcp_synthetic_data(tool_name: str = "send_email", arguments: dict = None): + """Mirror ProxyLogging._convert_mcp_to_llm_format: an MCP tool call rendered as a + synthetic user message so the existing prompt-scanning path can inspect it.""" + if arguments is None: + arguments = {"to": "user@example.com", "body": "some content"} + return { + "model": "mcp-tool-call", + "messages": [ + { + "role": "user", + "content": f"Tool: {tool_name}\nArguments: {arguments}", + } + ], + "metadata": {"guardrails": ["model-armor-test"]}, + "mcp_tool_name": tool_name, + "mcp_arguments": arguments, + } + + +@pytest.mark.asyncio +async def test_pre_call_hook_scans_mcp_tool_call_when_configured_for_pre_mcp_call(): + """A guardrail configured with mode `pre_mcp_call` must scan MCP tool calls. + + Regression: async_pre_call_hook hardcoded its event-type gate to `pre_call`, so a + `pre_mcp_call` guardrail's own inner should_run_guardrail check returned False for an + MCP call (call_type=call_mcp_tool) and the scan was skipped entirely -- letting + sensitive content in tool arguments through unscanned. The gate must remap + call_mcp_tool -> pre_mcp_call. + """ + guardrail = _make_guardrail(event_hook="pre_mcp_call", mask_request_content=True) + data = _mcp_synthetic_data() + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ) as mock_post: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=data, + call_type="call_mcp_tool", + ) + + mock_post.assert_called() + + +@pytest.mark.asyncio +async def test_pre_call_hook_skips_chat_traffic_when_configured_for_pre_mcp_call(): + """A `pre_mcp_call` guardrail must NOT scan ordinary chat completions -- the remap is + scoped to MCP calls, so a `completion` call_type still fails the gate and is skipped.""" + guardrail = _make_guardrail(event_hook="pre_mcp_call", mask_request_content=True) + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello there"}], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ) as mock_post: + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=data, + call_type="completion", + ) + + assert result == data + mock_post.assert_not_called() + + +@pytest.mark.asyncio +async def test_moderation_hook_scans_mcp_tool_call_when_configured_for_during_mcp_call(): + """A guardrail configured with mode `during_mcp_call` must scan MCP tool calls. + + Regression: async_moderation_hook hardcoded its event-type gate to `during_call`, so a + `during_mcp_call` guardrail skipped MCP calls (call_type=call_mcp_tool). The gate must + remap call_mcp_tool -> during_mcp_call. + """ + guardrail = _make_guardrail(event_hook="during_mcp_call", mask_request_content=True) + data = _mcp_synthetic_data() + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ) as mock_post: + await guardrail.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(), + call_type="call_mcp_tool", + ) + + mock_post.assert_called() + + +@pytest.mark.asyncio +async def test_moderation_hook_skips_chat_traffic_when_configured_for_during_mcp_call(): + """A `during_mcp_call` guardrail must NOT scan ordinary chat completions.""" + guardrail = _make_guardrail(event_hook="during_mcp_call", mask_request_content=True) + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello there"}], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ) as mock_post: + result = await guardrail.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + ) + + assert result == data + mock_post.assert_not_called()