From 03ec3f2cf23e0de57d7fc1055f78e0dec2356359 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 4 Sep 2026 00:44:16 +0000 Subject: [PATCH] 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> --- .../guardrail_translation/handler.py | 28 +++++++------- .../test_mcp_guardrail_handler.py | 37 ++++++++++++++----- type-discipline-budget.json | 2 +- 3 files changed, 42 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py index 385ff3ced56..c0235077ecd 100644 --- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py +++ b/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: 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 c4543b10d3b..26a4ba7e1c4 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 @@ -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=[""]) 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): diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ab1a793e09d..ecdb157c368 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16468 + "limit": 16466 }, "LIT011": { "limit": 5514