mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 7c19a18ffe into 26ba02d3ce
This commit is contained in:
karpanin
GitHub
commit
0298158ff4
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")
|
||||
|
|
|
|||
|
|
@ -1825,6 +1825,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``.
|
||||
|
|
@ -4561,15 +4582,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"]
|
||||
|
|
|
|||
|
|
@ -2166,8 +2166,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
|
||||
|
|
@ -2201,26 +2201,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)
|
||||
|
||||
|
|
@ -2236,12 +2242,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)
|
||||
|
||||
|
|
@ -2256,13 +2262,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)
|
||||
|
||||
|
|
@ -2277,7 +2283,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
|
||||
|
||||
|
|
@ -2287,7 +2293,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)
|
||||
|
||||
|
|
@ -2299,6 +2305,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
|
||||
|
|
@ -2406,7 +2564,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
|
||||
|
|
@ -2421,7 +2579,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