mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
test(e2e): cover a guardrail on the MCP tool-call path (content_filter pre_mcp_call)
Adds guardrail.litellm_content_filter.pre_mcp_call.blocks: against the real Datadog MCP server, a content_filter guardrail configured mode=pre_mcp_call blocks a banned keyword in an MCP tool call's arguments with HTTP 400 attributed to the pre_mcp_call hook, and lets a clean argument reach the upstream server. The guardrail attaches with default_on because per-key/request guardrail selection is dropped from the synthetic MCP request the hook sees; the banned keyword is unique per run so default_on only intercepts this test's own call. mode must be pre_mcp_call - a pre_call config silently no-ops on tools/call because the event type is rewritten for call_mcp_tool. Drives the tool directly via /mcp-rest/tools/call for a deterministic check of the same pre_mcp_call enforcement the OpenAI-SDK chat path hits when a model invokes an MCP tool.
This commit is contained in:
mubashir1osmani
parent
3b58a9eba0
commit
1d98326fc0
3 changed files with 155 additions and 0 deletions
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
94
tests/e2e/mcp/test_mcp_guardrail_e2e.py
Normal file
94
tests/e2e/mcp/test_mcp_guardrail_e2e.py
Normal file
|
|
@ -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}"
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue