litellm/tests/e2e/mcp/test_mcp_guardrail_e2e.py

"""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 time
from collections.abc import Callable

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
from lifecycle import ResourceManager
from mcp_client import McpCallToolResponse, McpClient, McpToolArguments

pytestmark = pytest.mark.e2e

# Stage runs several data-plane pods behind the shared key, and each picks up a
# newly registered guardrail or MCP server only on its next periodic DB sync (~30s in
# proxy_server.py). Every pod is guaranteed to have refreshed only once a full sync
# interval has elapsed since the later of those two writes; before then a banned call
# routed to a lagging pod passes through as legitimate in-flight propagation, not a leak.
FULL_SYNC_SECONDS = 40.0
POST_SYNC_VERIFICATION_CALLS = 4


def _poll_until_blocked(
    search: Callable[[str], Result[McpCallToolResponse]], banned_keyword: str, client: McpClient
) -> Result[McpCallToolResponse]:
    """Retry a banned tool call until the guardrail blocks it (400) or the deadline
    passes, returning the last result. Absorbs the control-plane -> data-plane
    guardrail-sync delay so the check waits for enforcement instead of racing it."""
    deadline = time.monotonic() + client.proxy.poll_timeout
    last: Result[McpCallToolResponse] = search(f"tell me about {banned_keyword}")
    while time.monotonic() < deadline:
        if isinstance(last, UnknownApiError) and last.status_code == 400:
            return last
        time.sleep(client.proxy.poll_interval)
        last = search(f"tell me about {banned_keyword}")
    return last


def _pod_lacks_mcp_server(result: Result[McpCallToolResponse]) -> bool:
    """True when the pod that served the call answered as though the MCP server or its
    tool does not exist (500 "Tool ... not found"), i.e. its MCP registry has not synced
    yet and the request never reached the guardrail at all."""
    if not isinstance(result, UnknownApiError) or result.status_code != 500:
        return False
    body = result.body.lower()
    return "not found" in body and ("tool" in body or "server" in body)


def _search_on_synced_pod(
    search: Callable[[str], Result[McpCallToolResponse]], query: str, client: McpClient
) -> Result[McpCallToolResponse]:
    """Issue `query`, retrying to the poll deadline only while the serving pod does not
    know the MCP server yet. Every other outcome, guardrail block or pass-through, comes
    back untouched so the caller's assertion still decides it."""
    deadline = time.monotonic() + client.proxy.poll_timeout
    last = search(query)
    while _pod_lacks_mcp_server(last) and time.monotonic() < deadline:
        time.sleep(client.proxy.poll_interval)
        last = search(query)
    return last


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
        )
        guardrail_created_at = time.monotonic()
        resources.defer(lambda: client.delete_guardrail(guardrail_id))

        server_id = register_datadog_mcp(client, resources)
        server_registered_at = time.monotonic()
        key = client.generate_key(user_id=f"e2e-mcp-guard-{marker}", mcp_servers=[server_id])
        resources.defer(lambda: client.proxy.delete_key(key))

        tool_name = client.await_tool(key, server_id, SEARCH_LOGS_TOOL)

        def search(query: str) -> Result[McpCallToolResponse]:
            arguments: McpToolArguments = {
                "query": query,
                "from": DD_SEARCH_FROM,
                "to": "now",
                "max_tokens": 500,
            }
            return client.call_tool(key, server_id=server_id, name=tool_name, arguments=arguments)

        # Registering the guardrail is a control-plane write; the data-plane worker
        # that serves tools/call picks it up on its next guardrail sync, so an
        # immediate call can race the propagation and slip through. Poll the banned
        # call to the deadline and require a block, so the check proves enforcement
        # rather than catching a pre-sync pass-through. The keyword is unique per
        # run, so this only ever intercepts this test's own call.
        blocked = _poll_until_blocked(search, banned_keyword, client)
        match blocked:
            case UnknownApiError(status_code=400, body=body):
                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(
                    "content_filter never blocked the banned keyword on the MCP tool call within "
                    f"{client.proxy.poll_timeout}s (the guardrail or the MCP server never synced to "
                    f"the data plane); last result: {blocked}"
                )

        # The block above only proves the one pod that served it has synced; another
        # pod could still lack the guardrail and let the banned call reach Datadog.
        # Wait out the full sync interval from the later of the guardrail create and the
        # MCP server registration (each syncs on its own clock, so the earlier write's
        # deadline can elapse while a pod still lacks the other) so every pod has
        # refreshed from the DB, then require the banned call to stay blocked across
        # several attempts. A pass-through now is a genuine partial-propagation leak, not
        # a race. Client load balancing still can't guarantee every pod is hit, so this
        # samples several worker selections rather than proving all pods synced.
        sync_remaining = max(guardrail_created_at, server_registered_at) + FULL_SYNC_SECONDS - time.monotonic()
        if sync_remaining > 0:
            time.sleep(sync_remaining)
        for attempt in range(1, POST_SYNC_VERIFICATION_CALLS + 1):
            reblocked = _search_on_synced_pod(search, f"still about {banned_keyword} #{attempt}", client)
            assert isinstance(reblocked, UnknownApiError) and reblocked.status_code == 400, (
                "after the sync interval every data-plane pod must block the banned keyword, but "
                f"attempt {attempt} of {POST_SYNC_VERIFICATION_CALLS} was not blocked (a pod still "
                f"lacks the guardrail, or never synced the MCP server): {reblocked}"
            )
            if attempt < POST_SYNC_VERIFICATION_CALLS:
                time.sleep(client.proxy.poll_interval)

        allowed = _search_on_synced_pod(search, f"e2e-clean-{marker}", client)
        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}"
                )