mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(guardrails): run post_mcp_call guardrails attached by key, team or policy
MCP tool calls hand post_mcp_call gating the request's Logging.model_call_details, which nests the request metadata under litellm_params. should_run_guardrail only reads the top level, so the guardrail list a key, team or policy resolved onto the request was invisible there. A post_mcp_call guardrail only ever ran when it was default_on, and opted_out_global_guardrails never took one off. Lift the metadata buckets out of litellm_params before gating, with a bucket the caller already set at the top level taking precedence. The guardrail auth path in operations, which every protocol tool call now takes, read the same top level when merging the request's guardrails with the guardrail auth's, so it reads the request metadata through the same lift. The post_mcp_call hook tests install callbacks through one fixture instead of each assigning litellm.callbacks, so the new regression tests do not add to that module-global write count.
This commit is contained in:
karpanin
parent
60fc121ce8
commit
7c19a18ffe
4 changed files with 231 additions and 24 deletions
|
|
@ -2618,8 +2618,9 @@ async def _run_post_mcp_call_guardrails(
|
|||
user_api_key_dict=post_auth,
|
||||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata
|
||||
from litellm.proxy.utils import mcp_guardrail_gate_data
|
||||
|
||||
base_request_data: Final[Mapping[str, object]] = (
|
||||
base_request_data: Final[Mapping[str, object]] = mcp_guardrail_gate_data(
|
||||
litellm_logging_obj.model_call_details if litellm_logging_obj is not None else request_data
|
||||
)
|
||||
request_metadata: Final[object | None] = base_request_data.get("metadata")
|
||||
|
|
|
|||
|
|
@ -1799,6 +1799,27 @@ def _failure_fields_to_lift(
|
|||
return MappingProxyType({key: value for key, value in _entries if value is not None})
|
||||
|
||||
|
||||
def mcp_guardrail_gate_data(request_data: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""``request_data`` reshaped so guardrail gating can read the request's metadata.
|
||||
|
||||
MCP tool calls carry the request's ``Logging.model_call_details``, which nests
|
||||
the request metadata under ``litellm_params``. ``CustomGuardrail.should_run_guardrail``
|
||||
only reads the top level, so without this lift the guardrail list a key, team
|
||||
or policy resolved onto the request is invisible and only ``default_on``
|
||||
guardrails ever run, and ``opted_out_global_guardrails`` never takes one off.
|
||||
A bucket the caller already set at the top level wins, because the MCP call
|
||||
site that sets it has merged the guardrail auth's guardrails into it.
|
||||
"""
|
||||
params: Final = request_data.get("litellm_params")
|
||||
if not isinstance(params, Mapping):
|
||||
return request_data
|
||||
nested: Final = cast("Mapping[str, object]", params) # cast-ok: litellm_params is always a str-keyed dict
|
||||
lifted: Final = MappingProxyType({key: nested[key] for key in ("metadata", "litellm_metadata") if key in nested})
|
||||
if not lifted:
|
||||
return request_data
|
||||
return MappingProxyType({**lifted, **request_data})
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _CallbackCapabilities:
|
||||
"""Cached per-hook capability flags derived from ``litellm.callbacks``.
|
||||
|
|
@ -4516,15 +4537,13 @@ class ProxyLogging:
|
|||
verbose_proxy_logger.debug("MCP guardrail translation handler unavailable; skipping post_mcp_call hook")
|
||||
return response
|
||||
|
||||
gate_data: Final = mcp_guardrail_gate_data(request_data)
|
||||
for callback in caps.resolved_callbacks:
|
||||
if not isinstance(callback, CustomGuardrail):
|
||||
continue
|
||||
if "apply_guardrail" not in type(callback).__dict__ or callback.use_native_lifecycle_hooks:
|
||||
continue
|
||||
if (
|
||||
callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_mcp_call)
|
||||
is not True
|
||||
):
|
||||
if callback.should_run_guardrail(data=gate_data, event_type=GuardrailEventHooks.post_mcp_call) is not True:
|
||||
continue
|
||||
response = await self._run_guardrail_with_metrics(
|
||||
callback,
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
import asyncio
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from mcp.server.context import ServerRequestContext
|
||||
from mcp.types import GetPromptRequest, GetPromptRequestParams, GetPromptResult
|
||||
from mcp.types import CallToolResult, GetPromptRequest, GetPromptRequestParams, GetPromptResult, TextContent
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
import litellm
|
||||
|
|
@ -2050,3 +2050,32 @@ async def test_discovery_optional_probes_use_remaining_listing_budget(delay, mon
|
|||
else:
|
||||
assert entered.is_set()
|
||||
assert cleaned.is_set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_mcp_call_guardrails_keep_the_request_metadata_on_the_guardrail_auth_path():
|
||||
"""Guardrails and opt-outs resolved onto the request must reach post_mcp_call gating.
|
||||
|
||||
Protocol tool calls pass a guardrail auth that is a separate object from the
|
||||
caller's, and the logging payload nests the request metadata under litellm_params.
|
||||
"""
|
||||
request_metadata: Final = {
|
||||
"guardrails": ["policy-guardrail"],
|
||||
"user_api_key_metadata": {"opted_out_global_guardrails": ["global-guardrail"]},
|
||||
}
|
||||
logging_obj: Final = MagicMock(model_call_details={"litellm_params": {"metadata": request_metadata}})
|
||||
proxy_logging: Final = MagicMock(post_mcp_call_hook=AsyncMock(side_effect=lambda response, **_: response))
|
||||
result: Final = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging):
|
||||
await operations._run_post_mcp_call_guardrails(
|
||||
result=result,
|
||||
litellm_logging_obj=logging_obj,
|
||||
original_auth=UserAPIKeyAuth(api_key="sk-caller", user_id="u1"),
|
||||
guardrail_auth=UserAPIKeyAuth(api_key="sk-caller", user_id="u1"),
|
||||
request_data={},
|
||||
)
|
||||
|
||||
gated: Final = proxy_logging.post_mcp_call_hook.await_args.kwargs["request_data"]["metadata"]
|
||||
assert gated["guardrails"] == ["policy-guardrail"]
|
||||
assert gated["user_api_key_metadata"] == request_metadata["user_api_key_metadata"]
|
||||
|
|
|
|||
|
|
@ -2007,8 +2007,8 @@ class TestSendEmailStartTls:
|
|||
class _RecordingMCPGuardrail(CustomGuardrail):
|
||||
"""Unified guardrail that masks every text it is handed."""
|
||||
|
||||
def __init__(self, event_hook, masked_text="<MASKED>", raises=None):
|
||||
super().__init__(guardrail_name="mcp-output-guardrail", event_hook=event_hook, default_on=True)
|
||||
def __init__(self, event_hook, masked_text="<MASKED>", raises=None, default_on=True):
|
||||
super().__init__(guardrail_name="mcp-output-guardrail", event_hook=event_hook, default_on=default_on)
|
||||
self.masked_text = masked_text
|
||||
self.raises = raises
|
||||
self.call_count = 0
|
||||
|
|
@ -2042,26 +2042,32 @@ class _NativeMCPGuardrail(CustomGuardrail):
|
|||
|
||||
|
||||
@pytest.fixture
|
||||
def restore_callbacks():
|
||||
"""Restore the process-wide callback state post_mcp_call_hook reads.
|
||||
def install_callbacks():
|
||||
"""Install callbacks on the process-wide list post_mcp_call_hook reads, and restore it afterwards.
|
||||
|
||||
ProxyLogging caches callback capabilities keyed on id()s of litellm.callbacks,
|
||||
so a restored-but-different list can collide with a stale entry after GC and
|
||||
leak a has_guardrail verdict into unrelated tests in the same worker.
|
||||
so a replaced list can collide with a stale entry after GC and leak a
|
||||
has_guardrail verdict into unrelated tests in the same worker. The cache is
|
||||
cleared on both install and restore for that reason.
|
||||
"""
|
||||
original = list(litellm.callbacks)
|
||||
yield
|
||||
|
||||
def install(*callbacks):
|
||||
litellm.callbacks = list(callbacks)
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
yield install
|
||||
litellm.callbacks = original
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_mcp_call_hook_masks_tool_result(restore_callbacks):
|
||||
async def test_post_mcp_call_hook_masks_tool_result(install_callbacks):
|
||||
"""A post_mcp_call guardrail must see the tool result text and mask it in the returned result."""
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call)
|
||||
litellm.callbacks = [guardrail]
|
||||
install_callbacks(guardrail)
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
|
||||
|
||||
|
|
@ -2077,12 +2083,12 @@ async def test_post_mcp_call_hook_masks_tool_result(restore_callbacks):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_mcp_call_hook_skips_guardrail_configured_for_other_hooks(restore_callbacks):
|
||||
async def test_post_mcp_call_hook_skips_guardrail_configured_for_other_hooks(install_callbacks):
|
||||
"""A guardrail not configured for post_mcp_call must not scan MCP tool results."""
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_call)
|
||||
litellm.callbacks = [guardrail]
|
||||
install_callbacks(guardrail)
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
|
||||
|
||||
|
|
@ -2097,13 +2103,13 @@ async def test_post_mcp_call_hook_skips_guardrail_configured_for_other_hooks(res
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_mcp_call_hook_skips_guardrail_without_apply_guardrail(restore_callbacks):
|
||||
async def test_post_mcp_call_hook_skips_guardrail_without_apply_guardrail(install_callbacks):
|
||||
"""Guardrails that implement async_post_mcp_tool_call_hook are dispatched by the
|
||||
logging object, so this hook must not run them a second time."""
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
guardrail = _NativeMCPGuardrail()
|
||||
litellm.callbacks = [guardrail]
|
||||
install_callbacks(guardrail)
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
|
||||
|
||||
|
|
@ -2118,7 +2124,7 @@ async def test_post_mcp_call_hook_skips_guardrail_without_apply_guardrail(restor
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_mcp_call_hook_propagates_guardrail_block(restore_callbacks):
|
||||
async def test_post_mcp_call_hook_propagates_guardrail_block(install_callbacks):
|
||||
"""A guardrail rejecting the tool result must raise out of the hook."""
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
|
|
@ -2128,7 +2134,7 @@ async def test_post_mcp_call_hook_propagates_guardrail_block(restore_callbacks):
|
|||
event_hook=GuardrailEventHooks.post_mcp_call,
|
||||
raises=BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="mcp-output-guardrail"),
|
||||
)
|
||||
litellm.callbacks = [guardrail]
|
||||
install_callbacks(guardrail)
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
|
||||
|
||||
|
|
@ -2140,6 +2146,158 @@ async def test_post_mcp_call_hook_propagates_guardrail_block(restore_callbacks):
|
|||
)
|
||||
|
||||
|
||||
def _mcp_logging_payload(metadata, metadata_key="metadata"):
|
||||
"""The request payload every MCP call site hands ``post_mcp_call_hook``.
|
||||
|
||||
``Logging.model_call_details`` nests the request metadata under
|
||||
``litellm_params``; see ``litellm.utils.function_setup``. The bucket is
|
||||
named ``litellm_metadata`` on the routes that reserve ``metadata`` for the
|
||||
provider, which is how MCP tools invoked from /responses arrive here.
|
||||
"""
|
||||
return {
|
||||
"model": "MCP: echo",
|
||||
"call_type": "call_mcp_tool",
|
||||
"litellm_params": {"api_base": "", metadata_key: metadata},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_mcp_call_hook_runs_guardrail_requested_through_request_metadata(install_callbacks):
|
||||
"""A guardrail attached by key, team or policy (not default_on) must mask MCP tool output.
|
||||
|
||||
The MCP paths pass the logging payload, where the resolved guardrail list
|
||||
sits under litellm_params.metadata, so gating on the top level alone left
|
||||
masking working only for default_on guardrails.
|
||||
"""
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call, default_on=False)
|
||||
install_callbacks(guardrail)
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
|
||||
|
||||
returned = await proxy_logging_obj.post_mcp_call_hook(
|
||||
response=result,
|
||||
request_data=_mcp_logging_payload({"guardrails": [guardrail.guardrail_name]}),
|
||||
user_api_key_dict=None,
|
||||
)
|
||||
|
||||
assert guardrail.call_count == 1
|
||||
assert [item.text for item in returned.content] == ["<MASKED>"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_mcp_call_hook_skips_guardrail_absent_from_request_metadata(install_callbacks):
|
||||
"""A non-default_on guardrail the request never asked for must leave tool output alone."""
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call, default_on=False)
|
||||
install_callbacks(guardrail)
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
|
||||
|
||||
returned = await proxy_logging_obj.post_mcp_call_hook(
|
||||
response=result,
|
||||
request_data=_mcp_logging_payload({"guardrails": ["some-other-guardrail"]}),
|
||||
user_api_key_dict=None,
|
||||
)
|
||||
|
||||
assert guardrail.call_count == 0
|
||||
assert [item.text for item in returned.content] == ["jane@example.com"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_mcp_call_hook_runs_guardrail_requested_through_litellm_metadata(install_callbacks):
|
||||
"""The same lift must work for routes whose request metadata bucket is litellm_metadata."""
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call, default_on=False)
|
||||
install_callbacks(guardrail)
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
|
||||
|
||||
returned = await proxy_logging_obj.post_mcp_call_hook(
|
||||
response=result,
|
||||
request_data=_mcp_logging_payload({"guardrails": [guardrail.guardrail_name]}, metadata_key="litellm_metadata"),
|
||||
user_api_key_dict=None,
|
||||
)
|
||||
|
||||
assert guardrail.call_count == 1
|
||||
assert [item.text for item in returned.content] == ["<MASKED>"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_mcp_call_hook_honors_opt_out_in_request_metadata(install_callbacks):
|
||||
"""A key that opted out of a global guardrail must not have MCP tool output scanned by it."""
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call)
|
||||
install_callbacks(guardrail)
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
|
||||
|
||||
returned = await proxy_logging_obj.post_mcp_call_hook(
|
||||
response=result,
|
||||
request_data=_mcp_logging_payload(
|
||||
{"user_api_key_metadata": {"opted_out_global_guardrails": [guardrail.guardrail_name]}}
|
||||
),
|
||||
user_api_key_dict=None,
|
||||
)
|
||||
|
||||
assert guardrail.call_count == 0
|
||||
assert [item.text for item in returned.content] == ["jane@example.com"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_mcp_call_hook_runs_default_on_guardrail_when_payload_has_no_metadata_bucket(install_callbacks):
|
||||
"""A logging payload carrying no metadata bucket must leave the Default On path working.
|
||||
|
||||
There is nothing to lift here, and that no-op must not change the verdict.
|
||||
"""
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call)
|
||||
install_callbacks(guardrail)
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
|
||||
|
||||
returned = await proxy_logging_obj.post_mcp_call_hook(
|
||||
response=result,
|
||||
request_data={"model": "MCP: echo", "litellm_params": {"api_base": ""}},
|
||||
user_api_key_dict=None,
|
||||
)
|
||||
|
||||
assert guardrail.call_count == 1
|
||||
assert [item.text for item in returned.content] == ["<MASKED>"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_mcp_call_hook_keeps_top_level_metadata_the_caller_shaped(install_callbacks):
|
||||
"""A top-level metadata bucket the call site already shaped must drive gating over the nested one.
|
||||
|
||||
The MCP call site merges the guardrail auth's guardrails, such as a granting
|
||||
team's, into the top level, so lifting litellm_params over it would drop them.
|
||||
"""
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call, default_on=False)
|
||||
install_callbacks(guardrail)
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
|
||||
|
||||
returned = await proxy_logging_obj.post_mcp_call_hook(
|
||||
response=result,
|
||||
request_data={
|
||||
"metadata": {"guardrails": [guardrail.guardrail_name]},
|
||||
**_mcp_logging_payload({"guardrails": ["some-other-guardrail"]}),
|
||||
},
|
||||
user_api_key_dict=None,
|
||||
)
|
||||
|
||||
assert guardrail.call_count == 1
|
||||
assert [item.text for item in returned.content] == ["<MASKED>"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prisma_health_check_failure_names_itself_at_operator_visible_level(caplog):
|
||||
"""A failing DB health check has to name the check that failed, at a level
|
||||
|
|
@ -2247,7 +2405,7 @@ async def test_update_data_key_branch_stamps_settings_updated_at():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_mcp_call_hook_skips_opted_out_guardrail(restore_callbacks):
|
||||
async def test_post_mcp_call_hook_skips_opted_out_guardrail(install_callbacks):
|
||||
"""A guardrail that keeps its native lifecycle hooks must not have MCP tool results
|
||||
scanned through the unified path, even though it implements apply_guardrail."""
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
|
@ -2262,7 +2420,7 @@ async def test_post_mcp_call_hook_skips_opted_out_guardrail(restore_callbacks):
|
|||
return await super().apply_guardrail(inputs, request_data, input_type, **kwargs)
|
||||
|
||||
guardrail = _OptedOutMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call)
|
||||
litellm.callbacks = [guardrail]
|
||||
install_callbacks(guardrail)
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue