fix(guardrails): record MCP tool guardrail evaluations and blocks in usage monitor

MCP tool calls reached the guardrail hooks without the request's litellm_logging_obj,
so pre_call_tool_check and _create_during_hook_task handed the guardrail-recording
bridge a None logging object; the bridge no-oped and the Guardrails Monitor reported
Total Evaluations 0 for MCP traffic. Separately, a pre-call guardrail block raised
before any failure handler ran, so the failure spend-log row written by
post_call_failure_hook -> get_logging_payload persisted with guardrail_information
None and the block was never counted, leaving Total Blocked at 0

The fix threads the logging object that MCPServerManager.call_tool already receives
through pre_call_tool_check and _create_during_hook_task and seeds it (and its shared
metadata dict) onto the synthetic guardrail request, so the hooks serialize their
evaluation into the same record the monitor counts. It also fires the failure handlers
before post_call_failure_hook so a blocked call's guardrail_information is populated
before the spend-log row is written; running them there flips the failure dedup flags,
so the @client wrapper's own post-raise failure logging short-circuits and does not
double-count

The raw MCP call_tool handler needs no change of its own: call_mcp_tool is
@client-decorated, so the wrapper builds the litellm_logging_obj and threads it into
kwargs already. The Responses API MCP path builds its own logging object, so it now
passes that into MCPServerManager.call_tool to cover that entry point too
This commit is contained in:
eugene-yao-zocdoc 2026-07-22 09:16:28 -04:00
parent bb6bb664b1
commit 5ccbedfb9e
6 changed files with 322 additions and 0 deletions

View file

@ -4433,6 +4433,7 @@ class MCPServerManager:
proxy_logging_obj: ProxyLogging,
server: MCPServer,
raw_headers: Optional[dict[str, str]] = None,
litellm_logging_obj: "Any | None" = None,
) -> dict[str, Any]:
"""
Run pre-call checks and guardrail hooks for an MCP tool call.
@ -4493,6 +4494,18 @@ class MCPServerManager:
# Convert to LLM format for existing guardrail compatibility
synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(mcp_request_obj, pre_hook_kwargs)
# Seed the real request logging object so the guardrail-recording bridge
# (custom_guardrail._sync_guardrail_info_to_logging_obj, run in the
# @log_guardrail_information finally) has a non-None logging_obj and can
# persist this pre_mcp_call evaluation into the spend log the Guardrails
# Monitor counts. Without it the synthetic dict carries no logging obj,
# the bridge no-ops, and Total Evaluations stays 0.
synthetic_llm_data["litellm_logging_obj"] = litellm_logging_obj
if litellm_logging_obj is not None:
logging_metadata = litellm_logging_obj.litellm_params.get("metadata")
if isinstance(logging_metadata, dict):
synthetic_llm_data["metadata"] = logging_metadata
hook_result: dict[str, Any] = {}
try:
# Use standard pre_call_hook
@ -4528,6 +4541,7 @@ class MCPServerManager:
user_api_key_auth: Optional[UserAPIKeyAuth],
proxy_logging_obj: ProxyLogging,
start_time: datetime.datetime,
litellm_logging_obj: "Any | None" = None,
):
"""Create and return a during hook task for MCP tool calls."""
from litellm.types.llms.base import HiddenParams
@ -4550,6 +4564,16 @@ class MCPServerManager:
synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(request_obj, during_hook_kwargs)
# Keep the synthetic guardrail request and the real request logger on the
# same metadata dict, as in pre_call_tool_check. The during hook is joined
# before success logging, so its completed evaluation is serialized with
# the tool call.
synthetic_llm_data["litellm_logging_obj"] = litellm_logging_obj
if litellm_logging_obj is not None:
logging_metadata = litellm_logging_obj.litellm_params.get("metadata")
if isinstance(logging_metadata, dict):
synthetic_llm_data["metadata"] = logging_metadata
return asyncio.create_task(
proxy_logging_obj.during_call_hook(
user_api_key_dict=user_api_key_auth,
@ -5085,6 +5109,7 @@ class MCPServerManager:
oauth2_headers: Optional[dict[str, str]] = None,
raw_headers: Optional[dict[str, str]] = None,
host_progress_callback: Optional[Callable] = None,
litellm_logging_obj: "Any | None" = None,
) -> CallToolResult:
"""
Call a tool with the given name and arguments
@ -5129,6 +5154,7 @@ class MCPServerManager:
proxy_logging_obj=proxy_logging_obj,
server=mcp_server,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
)
if "arguments" in hook_result:
arguments = hook_result["arguments"]
@ -5143,6 +5169,7 @@ class MCPServerManager:
user_api_key_auth=user_api_key_auth,
proxy_logging_obj=proxy_logging_obj,
start_time=start_time,
litellm_logging_obj=litellm_logging_obj,
)
tasks.append(during_hook_task)

View file

@ -2818,6 +2818,7 @@ if MCP_AVAILABLE:
proxy_logging_obj=proxy_logging_obj, # type: ignore[arg-type]
server=mcp_server,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
)
# `pre_call_tool_check` may return guardrail-modified
# arguments; honor them on the local path too.
@ -3055,6 +3056,20 @@ if MCP_AVAILABLE:
traceback_str = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG)
from litellm.proxy.proxy_server import proxy_logging_obj
# Build the status="failure" standard_logging_object BEFORE post_call_failure_hook.
# A pre-call guardrail block raises into this except, and the failure spend-log row is
# written by post_call_failure_hook -> get_logging_payload, which reads
# guardrail_information off the logging obj's standard_logging_object. That object is
# only populated once the failure handlers run, so without this the row persists with
# guardrail_information=None and the block is never counted on the Guardrails Monitor
# (totalBlocked stays 0). Running the failure handlers here also flips has_logged_sync_failure/
# has_logged_async_failure, so the @client wrapper's own post-raise failure logging
# short-circuits on should_run_logging and does not double-count this same logging obj.
end_time = datetime.now()
if litellm_logging_obj is not None:
litellm_logging_obj.failure_handler(e, traceback_str, start_time, end_time)
await litellm_logging_obj.async_failure_handler(e, traceback_str, start_time, end_time)
if proxy_logging_obj and user_api_key_auth:
await proxy_logging_obj.post_call_failure_hook(
request_data=kwargs,
@ -3232,6 +3247,7 @@ if MCP_AVAILABLE:
raw_headers=raw_headers,
proxy_logging_obj=proxy_logging_obj,
host_progress_callback=host_progress_callback,
litellm_logging_obj=litellm_logging_obj,
)
verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result)
return call_tool_result

View file

@ -793,6 +793,7 @@ class LiteLLM_Proxy_MCP_Handler:
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
proxy_logging_obj=proxy_logging_obj,
litellm_logging_obj=litellm_logging_obj,
)
if litellm_logging_obj:

View file

@ -0,0 +1,118 @@
"""Seam test for the guardrail-block recording fix in
``litellm.proxy._experimental.mcp_server.server.call_mcp_tool``.
A pre-call MCP guardrail block *raises* into ``call_mcp_tool``'s
``except Exception``. The failure spend-log row (the one the Guardrails Monitor
counts) is written by ``proxy_logging_obj.post_call_failure_hook`` ->
get_logging_payload, which reads ``guardrail_information`` off the logging obj's
``standard_logging_object``. That object is only populated once
``failure_handler`` / ``async_failure_handler`` run. So the failure handlers MUST
run *before* post_call_failure_hook, otherwise the row persists with
``guardrail_information=None`` and the block is never counted (totalBlocked stays
0). This test locks that ordering in.
``call_mcp_tool`` is wrapped by ``@client`` (litellm.utils.client), which uses
``functools.wraps`` and therefore exposes the raw undecorated coroutine as
``__wrapped__``. The tests drive ``__wrapped__`` directly so the except-block
ordering is observed in isolation, without the wrapper's own post-raise logging
firing. This mirrors the original config-repo seam test, which stubbed ``@client``
to an identity decorator for the same reason; like the original, it does not
exercise the wrapper's dedup path.
``proxy_logging_obj`` is imported lazily inside the except block via
``from litellm.proxy.proxy_server import proxy_logging_obj``; the real
proxy_server module is heavy (boto3 etc.), so a fake module is injected into
``sys.modules`` to satisfy that lazy import without loading it.
"""
import asyncio
import contextlib
import sys
import types
from unittest import mock
from fastapi import HTTPException
from litellm.proxy._experimental.mcp_server import server
class _RecordingLoggingObj:
"""Stands in for LiteLLMLoggingObj; records the failure-flush calls the fix
makes so the test can assert they happen before post_call_failure_hook."""
def __init__(self, order):
self._order = order
self.failure_calls = 0
self.async_failure_calls = 0
def has_run_logging(self, event_type):
self._order.append(("has_run_logging", event_type))
def failure_handler(self, *args, **kwargs):
self.failure_calls += 1
self._order.append("failure_handler")
async def async_failure_handler(self, *args, **kwargs):
self.async_failure_calls += 1
self._order.append("async_failure_handler")
def _call_block(logging_obj, order):
"""Drive call_mcp_tool into its except path via arguments=None (raises
HTTPException before any server-manager call) and return once it re-raises.
A fake ``litellm.proxy.proxy_server`` module is injected so the except
block's lazy ``from litellm.proxy.proxy_server import proxy_logging_obj``
resolves to a recording stand-in instead of importing the real (heavy)
module."""
async def _record_post_call_failure_hook(**kwargs):
order.append("post_call_failure_hook")
proxy_logging_obj = mock.MagicMock()
proxy_logging_obj.post_call_failure_hook.side_effect = _record_post_call_failure_hook
fake_proxy_server = types.ModuleType("litellm.proxy.proxy_server")
fake_proxy_server.proxy_logging_obj = proxy_logging_obj
with mock.patch.dict(sys.modules, {"litellm.proxy.proxy_server": fake_proxy_server}):
with contextlib.suppress(HTTPException):
asyncio.run(
server.call_mcp_tool.__wrapped__(
name="t",
arguments=None, # -> raise HTTPException(400) -> except Exception
user_api_key_auth=mock.MagicMock(), # truthy: gate the hook
litellm_logging_obj=logging_obj,
)
)
def test_block_flushes_failure_before_post_call_failure_hook():
order = []
obj = _RecordingLoggingObj(order)
_call_block(obj, order)
assert "post_call_failure_hook" in order, order
assert "failure_handler" in order, order
assert "async_failure_handler" in order, order
hook_at = order.index("post_call_failure_hook")
assert order.index("failure_handler") < hook_at, order
assert order.index("async_failure_handler") < hook_at, order
def test_block_flushes_each_handler_exactly_once():
order = []
obj = _RecordingLoggingObj(order)
_call_block(obj, order)
assert obj.failure_calls == 1, obj.failure_calls
assert obj.async_failure_calls == 1, obj.async_failure_calls
def test_absent_logging_obj_still_calls_hook_and_skips_flush():
# Guarded by `if litellm_logging_obj is not None`: without a logging obj the
# failure handlers are skipped (no crash) but post_call_failure_hook still
# fires. Byte-equivalent to stock behavior for that branch, so it passes on
# baseline too; its value is as a mutation-killer for that guard.
order = []
_call_block(None, order)
assert order == ["post_call_failure_hook"], order

View file

@ -0,0 +1,122 @@
"""Seam test for the MCP guardrail logging-obj threading in
``litellm.proxy._experimental.mcp_server.mcp_server_manager``.
It proves the fix's crux: ``pre_call_tool_check`` and ``_create_during_hook_task``
seed the real ``litellm_logging_obj`` (and its shared ``metadata`` dict) onto the
synthetic dict handed to the guardrail hooks (``data["litellm_logging_obj"]`` /
``data["metadata"]``). That is exactly what the custom_guardrail bridge reads to
record the evaluation into the spend log the monitor counts.
The real ``MCPServerManager`` is used via ``__new__`` (skipping ``__init__``); the
three pre-call guard/validation helpers on the path are stubbed, and a fake
``proxy_logging_obj`` is dependency-injected to capture the data each hook
receives. ``_create_during_hook_task`` does lazy in-function imports of
``litellm.types.*`` (HiddenParams, MCPDuringCallRequestObject) at call time;
those resolve against the installed litellm package.
"""
import asyncio
import datetime
from unittest import mock
from litellm.proxy._experimental.mcp_server import mcp_server_manager as MOD
SENTINEL = mock.MagicMock() # stands in for the real LiteLLMLoggingObj
SENTINEL.litellm_params = {"metadata": {}}
def _bare_manager():
"""An MCPServerManager instance without running __init__; stub the pre-call
guard/validation helpers this test's path touches so it reaches the seed."""
mgr = MOD.MCPServerManager.__new__(MOD.MCPServerManager)
mgr.check_allowed_or_banned_tools = lambda name, server: True
mgr.validate_allowed_params = lambda tool_name, arguments, server: None
async def _ok(*a, **k):
return None
mgr.check_tool_permission_for_key_team = _ok
return mgr
def _fake_proxy_logging(capture):
"""A proxy_logging_obj whose _convert_mcp_to_llm_format returns a fresh dict
and whose pre_call_hook/during_call_hook capture the data they receive."""
plo = mock.MagicMock()
plo._create_mcp_request_object_from_kwargs.return_value = mock.MagicMock()
plo._convert_mcp_to_llm_format.side_effect = lambda *a, **k: {}
async def _pre_call_hook(*, user_api_key_dict, data, call_type):
capture["pre"] = data
return None
async def _during_call_hook(*, user_api_key_dict, data, call_type):
capture["during"] = data
return None
plo.pre_call_hook.side_effect = _pre_call_hook
plo.during_call_hook.side_effect = _during_call_hook
return plo
def test_pre_call_seeds_real_logging_obj():
capture = {}
mgr = _bare_manager()
plo = _fake_proxy_logging(capture)
server = mock.MagicMock()
asyncio.run(
mgr.pre_call_tool_check(
name="t",
arguments={},
server_name="s",
user_api_key_auth=None,
proxy_logging_obj=plo,
server=server,
raw_headers={},
litellm_logging_obj=SENTINEL,
)
)
assert capture["pre"]["litellm_logging_obj"] is SENTINEL
assert capture["pre"]["metadata"] is SENTINEL.litellm_params["metadata"]
def test_during_hook_seeds_real_logging_obj():
capture = {}
mgr = _bare_manager()
plo = _fake_proxy_logging(capture)
async def _run():
task = mgr._create_during_hook_task(
name="t",
arguments={},
server_name_from_prefix="s",
user_api_key_auth=None,
proxy_logging_obj=plo,
start_time=datetime.datetime(2026, 7, 14),
litellm_logging_obj=SENTINEL,
)
await task
asyncio.run(_run())
assert capture["during"]["litellm_logging_obj"] is SENTINEL
assert capture["during"]["metadata"] is SENTINEL.litellm_params["metadata"]
def test_absent_logging_obj_defaults_to_none():
# When the caller does not thread a logging obj, the seed is None. Byte
# equivalent to stock behavior (key absent -> .get returns None).
capture = {}
mgr = _bare_manager()
plo = _fake_proxy_logging(capture)
asyncio.run(
mgr.pre_call_tool_check(
name="t",
arguments={},
server_name="s",
user_api_key_auth=None,
proxy_logging_obj=plo,
server=mock.MagicMock(),
raw_headers={},
)
)
assert capture["pre"]["litellm_logging_obj"] is None

View file

@ -450,6 +450,44 @@ async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_functio
assert captured.get("litellm_trace_id") == "tid"
@pytest.mark.asyncio
async def test_execute_tool_calls_threads_logging_obj_into_call_tool(monkeypatch):
"""The Responses-API MCP path must hand the request's litellm_logging_obj to
global_mcp_server_manager.call_tool, otherwise pre_call_tool_check /
_create_during_hook_task get None and the guardrail-recording bridge no-ops,
so MCP tool calls made through the Responses API report zero guardrail
evaluations in the monitor. Drop the litellm_logging_obj kwarg on the call_tool
invocation and this fails."""
_setup_proxy_logging(monkeypatch)
call_tool_mock = _setup_mcp_call_environment(monkeypatch)
sentinel_logging_obj = MagicMock()
sentinel_logging_obj.async_post_mcp_tool_call_hook = AsyncMock()
sentinel_logging_obj.async_success_handler = AsyncMock()
handler_module = importlib.import_module(
"litellm.responses.mcp.litellm_proxy_mcp_handler"
)
monkeypatch.setattr(
handler_module,
"function_setup",
lambda *_args, **_kwargs: (sentinel_logging_obj, None),
)
tool_name = "deepwiki-read_wiki_structure"
tool_calls = [{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}]
await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
tool_server_map={tool_name: "deepwiki"},
tool_calls=tool_calls,
user_api_key_auth=None,
)
assert call_tool_mock.await_count == 1
assert call_tool_mock.await_args is not None
assert call_tool_mock.await_args.kwargs["litellm_logging_obj"] is sentinel_logging_obj
@pytest.mark.asyncio
async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch):
"""