diff --git a/tests/e2e/coverage_registry/guardrail.yaml b/tests/e2e/coverage_registry/guardrail.yaml index 68722fbbb96..d54c12ba6dc 100644 --- a/tests/e2e/coverage_registry/guardrail.yaml +++ b/tests/e2e/coverage_registry/guardrail.yaml @@ -31,3 +31,4 @@ - {id: guardrail.tool_policy.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/tool_policy/tool_policy_guardrail.py", rationale: "Tool-use policy enforcement"} - {id: guardrail.mcp_security.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [mcp_operations], source: "guardrail_hooks/mcp_security", rationale: "MCP protocol security"} - {id: guardrail.llm_as_a_judge.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/llm_as_a_judge", rationale: "LLM-based judgment guardrail"} +- {id: guardrail.litellm_content_filter.pre_mcp_call.blocks, module: guardrail, tier: P1, hook_point: pre_mcp_call, assertions: [blocks], exercised_on: [mcp_operations], source: "guardrail_hooks/litellm_content_filter/content_filter.py:_scan_mcp_tool_call_arguments", rationale: "A general content-filter guardrail configured mode=pre_mcp_call blocks a banned keyword in an MCP tool call's arguments before it reaches the upstream MCP server; a clean argument passes"} diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index f68fdf63b3f..b0aa4c68e3a 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -85,6 +85,37 @@ class McpToolsListResponse(BaseModel): return None +class BlockedWordSpec(BaseModel): + keyword: str + action: str = "BLOCK" + + +class ContentFilterMcpParams(BaseModel): + """litellm_content_filter params scoped to the MCP tool-call hook. mode is + pre_mcp_call because a pre_call config silently no-ops on the tools/call path + (the event type is rewritten to pre_mcp_call for call_mcp_tool), and default_on + is required there because per-key/request guardrail selection is dropped from + the synthetic MCP request the hook sees.""" + + guardrail: str = "litellm_content_filter" + mode: str = "pre_mcp_call" + default_on: bool = True + blocked_words: list[BlockedWordSpec] + + +class GuardrailSpecBody(BaseModel): + guardrail_name: str + litellm_params: ContentFilterMcpParams + + +class GuardrailCreateBody(BaseModel): + guardrail: GuardrailSpecBody + + +class GuardrailCreateResponse(BaseModel): + guardrail_id: str + + class McpCallToolBody(BaseModel): name: str arguments: dict[str, McpToolArg] @@ -186,6 +217,35 @@ class McpClient: response_type=McpToolsListResponse, ) + def register_mcp_content_filter(self, *, name: str, blocked_keyword: str) -> str: + """Register a default-on content-filter guardrail that runs on the MCP + tool-call hook (pre_mcp_call) and blocks a single keyword. The keyword is + unique per test, so default_on only ever intercepts this test's own + banned tool call on the shared proxy.""" + return unwrap( + self.proxy.transport.post( + "/guardrails", + headers=self.proxy.transport.master, + json=GuardrailCreateBody( + guardrail=GuardrailSpecBody( + guardrail_name=name, + litellm_params=ContentFilterMcpParams( + blocked_words=[BlockedWordSpec(keyword=blocked_keyword)], + ), + ) + ), + response_type=GuardrailCreateResponse, + ) + ).guardrail_id + + def delete_guardrail(self, guardrail_id: str) -> None: + _ = self.proxy.transport.delete( + f"/guardrails/{guardrail_id}", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + def call_tool( self, key: str, diff --git a/tests/e2e/mcp/test_mcp_guardrail_e2e.py b/tests/e2e/mcp/test_mcp_guardrail_e2e.py new file mode 100644 index 00000000000..d6f1dbb76f6 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_guardrail_e2e.py @@ -0,0 +1,94 @@ +"""Live e2e: a guardrail on the MCP tool-call path blocks banned content in the +tool arguments before the call reaches the upstream MCP server. + +A general litellm_content_filter guardrail is configured with mode=pre_mcp_call +(the event type the proxy rewrites pre_call to for a call_mcp_tool) and default_on +(per-key/request guardrail selection is dropped from the synthetic MCP request the +hook sees, so default_on is how it attaches to tools/call). The banned keyword is +unique per run, so default_on only ever intercepts this test's own banned call. + +Against the real Datadog MCP server, calling search_datadog_logs with the banned +keyword in the query is blocked with HTTP 400 attributed to the pre_mcp_call hook, +and the tool never runs; the same guardrail lets a clean query through to Datadog. +This is the enforced half (the block) plus the pass-through half in one spec. +""" + +from __future__ import annotations + +import pytest + +from datadog_mcp import SEARCH_LOGS_TOOL, assert_dd_mcp_creds, register_datadog_mcp +from e2e_config import DD_SEARCH_FROM, unique_marker +from e2e_http import Result, Success, UnknownApiError, unwrap +from lifecycle import ResourceManager +from mcp_client import McpCallToolResponse, McpClient, McpToolArguments + +pytestmark = pytest.mark.e2e + + +class TestMcpToolCallGuardrail: + @pytest.mark.covers( + "guardrail.litellm_content_filter.pre_mcp_call.blocks", + exercised_on=["mcp_operations"], + ) + def test_content_filter_blocks_banned_keyword_in_tool_args( + self, client: McpClient, resources: ResourceManager + ) -> None: + assert_dd_mcp_creds() + marker = unique_marker() + banned_keyword = f"e2eblocked{marker}" + + guardrail_id = client.register_mcp_content_filter( + name=f"e2e-mcp-cf-{marker}", blocked_keyword=banned_keyword + ) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + + server_id = register_datadog_mcp(client, resources) + key = client.generate_key(user_id=f"e2e-mcp-guard-{marker}", mcp_servers=[server_id]) + resources.defer(lambda: client.proxy.delete_key(key)) + + tools = unwrap(client.list_tools(key)) + tool_name = tools.tool_name_containing(server_id, SEARCH_LOGS_TOOL) + assert tool_name is not None, ( + f"granted key never saw {SEARCH_LOGS_TOOL} on server {server_id}; " + f"tools={tools.tool_names_for_server(server_id)}" + ) + + def search(query: str) -> Result[McpCallToolResponse]: + arguments: McpToolArguments = { + "query": query, + "from": DD_SEARCH_FROM, + "to": "now", + "max_tokens": 500, + "telemetry": {"intent": "e2e mcp guardrail check"}, + } + return client.call_tool(key, server_id=server_id, name=tool_name, arguments=arguments) + + blocked = search(f"tell me about {banned_keyword}") + match blocked: + case UnknownApiError(status_code=status, body=body): + assert status == 400, ( + f"a banned keyword in the tool arguments must block the MCP tool call with " + f"400, got {status}: {body[:300]}" + ) + assert banned_keyword in body or "content blocked" in body.lower(), ( + f"the block must name the content-filter reason, got: {body[:300]}" + ) + assert "pre_mcp_call" in body, ( + f"the block must be attributed to the MCP tool-call hook (pre_mcp_call), got: {body[:300]}" + ) + case _: + pytest.fail( + f"content_filter did not block a banned keyword in an MCP tool call; got {blocked}" + ) + + allowed = search(f"e2e-clean-{marker}") + match allowed: + case Success(data=result): + assert result.is_error is not True, ( + f"a clean MCP tool call must reach the server and not error, got: {result}" + ) + case _: + pytest.fail( + f"a clean MCP tool call must pass the guardrail and reach the server; got {allowed}" + )