mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(guardrails): record MCP tool guardrail evaluations and blocks in usage monitor
MCP tool calls run their guardrails against a throwaway LLM-shaped dict
built by `ProxyLogging._convert_mcp_to_llm_format`, not against the dict
the tool call is logged from. `@log_guardrail_information` therefore
appended `standard_logging_guardrail_information` to that throwaway
dict's metadata bucket, where `get_standard_logging_object_payload`
never saw it, so the Guardrails Monitor reported zero evaluations and
zero blocks for all MCP traffic.
Thread the request's `litellm_logging_obj` into `pre_call_tool_check`
and `_create_during_hook_task` and bridge the guardrail records onto it:
- Seed `data["litellm_logging_obj"]`, which unified guardrails read and
pass into `apply_guardrail`.
- Call `_sync_guardrail_info_to_logging_obj` in a `finally`, which is
what native guardrails need and what makes the block path work: a
blocked call raises straight out of `pre_call_tool_check`, so the
record has to be attached before the exception leaves the frame.
Only the guardrail evaluation records are copied. The synthetic
request's messages and tool arguments are deliberately left behind --
they can carry end-user data and nothing in the monitor needs them.
In `call_mcp_tool`, flush the failure handlers before
`post_call_failure_hook` so the `status="failure"` standard logging
object exists when `_ProxyDBLogger.async_post_call_failure_hook` writes
the spend-log row the monitor's "Total Blocked" counts. Both handlers
gate on `should_run_logging("sync_failure")` / `("async_failure")` and
then mark it, so the `@client` wrapper's own post-raise logging is a
no-op and nothing is double-counted -- the same pattern
`_fire_mcp_tool_call_logging` already uses for `isError=True`.
Threaded through every MCP tool entry point: the managed-server path,
the local-OpenAPI registry path, the legacy registry fallback, and the
Responses API's `_execute_tool_calls`.
This commit is contained in:
parent
2bc9bb4a9d
commit
9858d021ee
6 changed files with 560 additions and 10 deletions
|
|
@ -13,7 +13,7 @@ import json
|
|||
import os
|
||||
import re
|
||||
import time
|
||||
from collections.abc import AsyncIterator, Callable, Sequence
|
||||
from collections.abc import AsyncIterator, Callable, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast
|
||||
from urllib.parse import ParseResult, urlparse
|
||||
|
|
@ -46,6 +46,9 @@ from litellm.constants import (
|
|||
)
|
||||
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
|
||||
from litellm.experimental_mcp_client.client import MCPClient, MCPSigV4Auth
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
_sync_guardrail_info_to_logging_obj, # pyright: ignore[reportPrivateUsage] - the same bridge @log_guardrail_information uses; reimplementing it here would fork the metadata-key logic
|
||||
)
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
|
||||
|
|
@ -162,6 +165,7 @@ if TYPE_CHECKING:
|
|||
from mcp.types import CreateMessageRequestParams
|
||||
|
||||
from litellm.caching.caching import InMemoryCache
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.mcp_server.mcp_toolset import MCPToolset
|
||||
|
||||
try:
|
||||
|
|
@ -1233,6 +1237,35 @@ def _create_elicitation_callback():
|
|||
return _elicitation_callback
|
||||
|
||||
|
||||
def _record_mcp_guardrail_evaluations(
|
||||
synthetic_llm_data: dict[str, Any], # mutable-ok: `_sync_guardrail_info_to_logging_obj` takes a concrete dict
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None",
|
||||
) -> None:
|
||||
"""Bridge guardrail decision records off an MCP synthetic request onto the request's logger.
|
||||
|
||||
MCP guardrails run against a throwaway LLM-shaped dict from
|
||||
``ProxyLogging._convert_mcp_to_llm_format``, so ``@log_guardrail_information``
|
||||
files ``standard_logging_guardrail_information`` in that dict's metadata bucket,
|
||||
which ``get_standard_logging_object_payload`` never reads. Native (non-unified)
|
||||
guardrails receive no ``logging_obj`` kwarg, so the decorator cannot bridge on
|
||||
their behalf; this calls the same helper it would have.
|
||||
|
||||
Only the decision records move. The synthetic request's messages and tool
|
||||
arguments stay behind: they can carry end-user data, and the monitor needs none
|
||||
of it.
|
||||
"""
|
||||
if litellm_logging_obj is None:
|
||||
return
|
||||
|
||||
try:
|
||||
_sync_guardrail_info_to_logging_obj(synthetic_llm_data, litellm_logging_obj)
|
||||
except Exception as e: # noqa: BLE001 # callers run this from a `finally` on the block path
|
||||
# The breadth is the point. Narrowing to the knowable AttributeError/TypeError
|
||||
# would let an unexpected type escape that ``finally`` and replace the guardrail's
|
||||
# block with a bookkeeping error.
|
||||
verbose_logger.warning("Failed to record MCP guardrail evaluation for logging: %s", e)
|
||||
|
||||
|
||||
class MCPServerManager:
|
||||
_STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$")
|
||||
|
||||
|
|
@ -4543,6 +4576,7 @@ class MCPServerManager:
|
|||
proxy_logging_obj: ProxyLogging | None,
|
||||
server: MCPServer,
|
||||
raw_headers: dict[str, str] | None = None,
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Run pre-call checks and guardrail hooks for an MCP tool call.
|
||||
|
|
@ -4552,6 +4586,10 @@ class MCPServerManager:
|
|||
present. An absent logger must never be able to turn an authorization
|
||||
decision into a no-op.
|
||||
|
||||
``litellm_logging_obj`` is the request's logger, and it is what lands a
|
||||
``pre_mcp_call`` evaluation (or a block) on the spend-log row the Guardrails
|
||||
Monitor counts. It stays optional so callers that do no logging are unchanged.
|
||||
|
||||
Returns a dict that may contain:
|
||||
- "arguments": hook-modified tool arguments (only if changed)
|
||||
- "extra_headers": headers injected by pre_mcp_call guardrail hooks
|
||||
|
|
@ -4610,8 +4648,13 @@ class MCPServerManager:
|
|||
# Create MCP request object for processing
|
||||
mcp_request_obj: Final = proxy_logging_obj._create_mcp_request_object_from_kwargs(pre_hook_kwargs)
|
||||
|
||||
# Convert to LLM format for existing guardrail compatibility
|
||||
# Convert to LLM format for existing guardrail compatibility.
|
||||
# Unified guardrails read the seeded logger off the request dict and pass it
|
||||
# into ``apply_guardrail``, so ``@log_guardrail_information`` bridges their
|
||||
# evaluations itself; the ``finally`` below covers native guardrails, which
|
||||
# never receive it. Same seeding the pass-through routes do.
|
||||
synthetic_llm_data: Final = proxy_logging_obj._convert_mcp_to_llm_format(mcp_request_obj, pre_hook_kwargs)
|
||||
synthetic_llm_data["litellm_logging_obj"] = litellm_logging_obj
|
||||
|
||||
try:
|
||||
# Use standard pre_call_hook
|
||||
|
|
@ -4636,6 +4679,12 @@ class MCPServerManager:
|
|||
# Re-raise guardrail exceptions to properly fail the MCP call
|
||||
verbose_logger.error("Guardrail blocked MCP tool call pre call: %s", e)
|
||||
raise e
|
||||
finally:
|
||||
# ``finally`` rather than after the ``try``: a block raises straight out of
|
||||
# here, and the failure spend-log row that "Total Blocked" counts is built
|
||||
# from this logger further up the stack, so the record has to be attached
|
||||
# before the exception leaves this frame.
|
||||
_record_mcp_guardrail_evaluations(synthetic_llm_data, litellm_logging_obj)
|
||||
|
||||
return hook_result
|
||||
|
||||
|
|
@ -4647,8 +4696,14 @@ class MCPServerManager:
|
|||
user_api_key_auth: UserAPIKeyAuth | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
start_time: datetime.datetime,
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
):
|
||||
"""Create and return a during hook task for MCP tool calls."""
|
||||
"""Create and return a during hook task for MCP tool calls.
|
||||
|
||||
``litellm_logging_obj`` is the request's logger; see ``pre_call_tool_check``.
|
||||
The task is awaited before the tool call's success logging runs, so a
|
||||
``during_mcp_call`` evaluation recorded on it is serialized with that call.
|
||||
"""
|
||||
from litellm.types.llms.base import HiddenParams
|
||||
from litellm.types.mcp import MCPDuringCallRequestObject
|
||||
|
||||
|
|
@ -4667,15 +4722,23 @@ class MCPServerManager:
|
|||
"user_api_key_auth": user_api_key_auth,
|
||||
}
|
||||
|
||||
# Seeded for the same reason as in ``pre_call_tool_check``.
|
||||
synthetic_llm_data: Final = proxy_logging_obj._convert_mcp_to_llm_format(request_obj, during_hook_kwargs)
|
||||
synthetic_llm_data["litellm_logging_obj"] = litellm_logging_obj
|
||||
|
||||
return asyncio.create_task(
|
||||
proxy_logging_obj.during_call_hook(
|
||||
user_api_key_dict=user_api_key_auth,
|
||||
data=synthetic_llm_data,
|
||||
call_type=CallTypes.call_mcp_tool.value,
|
||||
)
|
||||
)
|
||||
# Wrapped so the bridge runs inside the task: the caller only holds the task and
|
||||
# gathers it later, so there is no other point that still sees a block here.
|
||||
async def _run_during_call_hook() -> Mapping[str, Any] | None:
|
||||
try:
|
||||
return await proxy_logging_obj.during_call_hook(
|
||||
user_api_key_dict=user_api_key_auth,
|
||||
data=synthetic_llm_data,
|
||||
call_type=CallTypes.call_mcp_tool.value,
|
||||
)
|
||||
finally:
|
||||
_record_mcp_guardrail_evaluations(synthetic_llm_data, litellm_logging_obj)
|
||||
|
||||
return asyncio.create_task(_run_during_call_hook())
|
||||
|
||||
def _get_call_semaphore(self, mcp_server: MCPServer) -> asyncio.Semaphore | None:
|
||||
limit: Final = mcp_server.max_concurrent_requests
|
||||
|
|
@ -5204,6 +5267,7 @@ class MCPServerManager:
|
|||
oauth2_headers: dict[str, str] | None = None,
|
||||
raw_headers: dict[str, str] | None = None,
|
||||
host_progress_callback: Callable | None = None,
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> CallToolResult:
|
||||
"""
|
||||
Call a tool with the given name and arguments
|
||||
|
|
@ -5216,6 +5280,9 @@ class MCPServerManager:
|
|||
mcp_auth_header: MCP auth header (deprecated)
|
||||
mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value}
|
||||
proxy_logging_obj: Optional ProxyLogging object for hook integration
|
||||
litellm_logging_obj: Optional request logger the guardrail hooks record
|
||||
their evaluations onto, so MCP guardrail activity reaches the
|
||||
Guardrails Monitor. See ``pre_call_tool_check``
|
||||
|
||||
|
||||
Returns:
|
||||
|
|
@ -5246,6 +5313,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"]
|
||||
|
|
@ -5260,6 +5328,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)
|
||||
|
||||
|
|
|
|||
|
|
@ -2824,6 +2824,7 @@ if MCP_AVAILABLE:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
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.
|
||||
|
|
@ -2962,6 +2963,7 @@ if MCP_AVAILABLE:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
server=prefix_server,
|
||||
raw_headers=raw_headers,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
)
|
||||
if "arguments" in hook_result:
|
||||
arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args
|
||||
|
|
@ -3149,6 +3151,20 @@ if MCP_AVAILABLE:
|
|||
traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG)
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
# Ordering is load-bearing. ``_ProxyDBLogger.async_post_call_failure_hook``,
|
||||
# reached below, writes the failure spend-log row from this logger's
|
||||
# ``standard_logging_object``, which only exists once the failure handlers
|
||||
# have run. Flush them first or the row lands with
|
||||
# ``guardrail_information=None`` and a guardrail block is never counted.
|
||||
#
|
||||
# Not double-logged: both handlers gate on ``should_run_logging`` and then
|
||||
# mark it, so the ``@client`` wrapper's own post-raise logging no-ops on this
|
||||
# logger, same as ``_fire_mcp_tool_call_logging`` does for ``isError=True``.
|
||||
if litellm_logging_obj is not None:
|
||||
end_time: Final = datetime.now() # noqa: DTZ005 # naive to match `start_time`, which it is subtracted from
|
||||
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,
|
||||
|
|
@ -3326,6 +3342,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
|
||||
|
|
|
|||
|
|
@ -798,6 +798,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 proxy_logging_obj:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,126 @@
|
|||
"""Tests for guardrail-block recording 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 that the Guardrails Monitor's
|
||||
"Total Blocked" counts is written by ``_ProxyDBLogger.async_post_call_failure_hook``
|
||||
(reached via ``proxy_logging_obj.post_call_failure_hook``), which reads
|
||||
``standard_logging_object`` off the request's logging obj -- and that only exists
|
||||
once ``failure_handler`` / ``async_failure_handler`` have 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. These tests
|
||||
pin that ordering.
|
||||
|
||||
``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. Note that this means they do not exercise the wrapper's dedup path; that
|
||||
dedup rests on ``should_run_logging("sync_failure")`` / ``("async_failure")``,
|
||||
which has its own coverage in the logging tests.
|
||||
|
||||
``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, so a fake module is injected into ``sys.modules``
|
||||
to satisfy that lazy import without loading it.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import sys
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._experimental.mcp_server import server
|
||||
|
||||
|
||||
class _RecordingLoggingObj:
|
||||
"""Stands in for ``LiteLLMLoggingObj``, recording the failure flush the fix
|
||||
makes so the test can assert it happens before ``post_call_failure_hook``."""
|
||||
|
||||
def __init__(self, order: list) -> None:
|
||||
self._order = order
|
||||
self.failure_calls = 0
|
||||
self.async_failure_calls = 0
|
||||
|
||||
def failure_handler(self, *_args, **_kwargs) -> None:
|
||||
self.failure_calls += 1
|
||||
self._order.append("failure_handler")
|
||||
|
||||
async def async_failure_handler(self, *_args, **_kwargs) -> None:
|
||||
self.async_failure_calls += 1
|
||||
self._order.append("async_failure_handler")
|
||||
|
||||
|
||||
async def _call_block(logging_obj, order: list, *, user_api_key_auth=mock.sentinel.auth):
|
||||
"""Drive ``call_mcp_tool`` into its except path via ``arguments=None``, which
|
||||
raises ``HTTPException(400)`` before any server-manager call, and return once it
|
||||
re-raises."""
|
||||
|
||||
async def _record_post_call_failure_hook(**_kwargs) -> None:
|
||||
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 # pyright: ignore[reportAttributeAccessIssue]
|
||||
|
||||
with mock.patch.dict(sys.modules, {"litellm.proxy.proxy_server": fake_proxy_server}):
|
||||
with contextlib.suppress(HTTPException):
|
||||
await server.call_mcp_tool.__wrapped__(
|
||||
name="t",
|
||||
arguments=None,
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_block_flushes_failure_before_post_call_failure_hook():
|
||||
order: list = []
|
||||
await _call_block(_RecordingLoggingObj(order), order)
|
||||
|
||||
assert order == ["failure_handler", "async_failure_handler", "post_call_failure_hook"], order
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_block_flushes_each_handler_exactly_once():
|
||||
"""Each handler runs once, so the block yields exactly one counted row rather
|
||||
than double-counting on the shared logging obj."""
|
||||
order: list = []
|
||||
obj = _RecordingLoggingObj(order)
|
||||
await _call_block(obj, order)
|
||||
|
||||
assert (obj.failure_calls, obj.async_failure_calls) == (1, 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_block_flushes_failure_for_anonymous_calls():
|
||||
"""With no ``user_api_key_auth`` the failure handlers still run, so OTel and the
|
||||
other failure sinks see the block.
|
||||
|
||||
``post_call_failure_hook`` stays gated on auth, matching the pre-existing
|
||||
contract: SpendLogs rows are attributable billing/audit records and the
|
||||
downstream DB logger dereferences authenticated key, budget, and route data.
|
||||
Counting anonymous MCP blocks needs a counter that does not live in SpendLogs,
|
||||
which is a separate design change, not part of this fix.
|
||||
"""
|
||||
order: list = []
|
||||
obj = _RecordingLoggingObj(order)
|
||||
await _call_block(obj, order, user_api_key_auth=None)
|
||||
|
||||
assert order == ["failure_handler", "async_failure_handler"], order
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_absent_logging_obj_still_calls_hook_and_skips_flush():
|
||||
"""Without a logging obj the flush is skipped (no crash) but
|
||||
``post_call_failure_hook`` still fires. Byte-equivalent to stock behavior for
|
||||
that branch; its value is as a mutation-killer for the ``is not None`` guard."""
|
||||
order: list = []
|
||||
await _call_block(None, order)
|
||||
|
||||
assert order == ["post_call_failure_hook"], order
|
||||
|
|
@ -0,0 +1,301 @@
|
|||
"""Tests for MCP guardrail evaluations reaching the Guardrails Monitor.
|
||||
|
||||
MCP tool calls run their guardrails against a throwaway LLM-shaped dict built by
|
||||
``ProxyLogging._convert_mcp_to_llm_format``, not against the dict the tool call is
|
||||
logged from. ``@log_guardrail_information`` therefore appends
|
||||
``standard_logging_guardrail_information`` to that throwaway dict's metadata
|
||||
bucket, where ``get_standard_logging_object_payload`` never sees it, so the
|
||||
Guardrails Monitor reported zero evaluations and zero blocks for MCP traffic.
|
||||
|
||||
``pre_call_tool_check`` and ``_create_during_hook_task`` now take the request's
|
||||
``litellm_logging_obj`` and bridge those records onto it. These tests pin both the
|
||||
seeding (which unified guardrails consume off ``data["litellm_logging_obj"]``) and
|
||||
the bridge (which native guardrails depend on), including on the block path.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
from typing import Any
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.proxy._experimental.mcp_server import mcp_server_manager as MOD
|
||||
|
||||
|
||||
class _FakeLoggingObj:
|
||||
"""Minimal stand-in for ``LiteLLMLoggingObj``.
|
||||
|
||||
``_sync_guardrail_info_to_logging_obj`` reads exactly these two attributes,
|
||||
and the spend-log payload is built from ``litellm_params["metadata"]``, so a
|
||||
real ``Logging`` instance would add setup cost without adding coverage.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.litellm_params: dict[str, Any] = {"metadata": {}}
|
||||
self.model_call_details: dict[str, Any] = {"litellm_params": self.litellm_params}
|
||||
|
||||
@property
|
||||
def recorded_guardrails(self) -> list:
|
||||
return self.litellm_params["metadata"].get("standard_logging_guardrail_information", [])
|
||||
|
||||
|
||||
def _bare_manager() -> MOD.MCPServerManager:
|
||||
"""An ``MCPServerManager`` without running ``__init__``.
|
||||
|
||||
The authorization/validation helpers on the path are stubbed out so the test
|
||||
reaches the guardrail hooks; they have their own coverage elsewhere.
|
||||
"""
|
||||
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(*_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
mgr.check_tool_permission_for_key_team = _ok
|
||||
return mgr
|
||||
|
||||
|
||||
def _fake_proxy_logging(capture: dict, *, guardrail_effect=None):
|
||||
"""A ``proxy_logging_obj`` double whose hooks capture the data they receive.
|
||||
|
||||
``guardrail_effect`` stands in for a guardrail: it is handed the synthetic
|
||||
request dict so it can append a guardrail record (and optionally raise, the
|
||||
way a blocking guardrail does).
|
||||
"""
|
||||
plo = mock.MagicMock()
|
||||
plo._create_mcp_request_object_from_kwargs.return_value = mock.MagicMock()
|
||||
# Mirror the real conversion's metadata bucket so a test can prove it survives.
|
||||
plo._convert_mcp_to_llm_format.side_effect = lambda *_a, **_k: {
|
||||
"metadata": {"headers": {"x-forwarded-for": "1.2.3.4"}}
|
||||
}
|
||||
|
||||
async def _hook(*, user_api_key_dict, data, call_type) -> None:
|
||||
del user_api_key_dict # captured shape is what matters, not the auth double
|
||||
capture["data"] = data
|
||||
capture["call_type"] = call_type
|
||||
if guardrail_effect is not None:
|
||||
guardrail_effect(data)
|
||||
|
||||
plo.pre_call_hook.side_effect = _hook
|
||||
plo.during_call_hook.side_effect = _hook
|
||||
return plo
|
||||
|
||||
|
||||
def _record_guardrail(status: str = "success"):
|
||||
"""Write a guardrail record the way ``@log_guardrail_information`` does."""
|
||||
|
||||
def _effect(data: dict) -> None:
|
||||
data.setdefault("metadata", {}).setdefault("standard_logging_guardrail_information", []).append(
|
||||
{"guardrail_name": "test-guardrail", "guardrail_status": status}
|
||||
)
|
||||
|
||||
return _effect
|
||||
|
||||
|
||||
def _blocking_guardrail():
|
||||
record = _record_guardrail(status="guardrail_intervened")
|
||||
|
||||
def _effect(data: dict) -> None:
|
||||
record(data)
|
||||
raise GuardrailRaisedException(guardrail_name="test-guardrail", message="blocked")
|
||||
|
||||
return _effect
|
||||
|
||||
|
||||
async def _run_pre_call(mgr, plo, logging_obj) -> dict:
|
||||
return await 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={},
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_seeds_request_logging_obj_for_unified_guardrails():
|
||||
"""Unified guardrails read ``data["litellm_logging_obj"]`` and pass it into
|
||||
``apply_guardrail``, whose ``@log_guardrail_information`` wrapper bridges the
|
||||
evaluation onto that logger itself. Drop the seed and that path records
|
||||
nothing."""
|
||||
capture: dict = {}
|
||||
logging_obj = _FakeLoggingObj()
|
||||
await _run_pre_call(_bare_manager(), _fake_proxy_logging(capture), logging_obj)
|
||||
|
||||
assert capture["data"]["litellm_logging_obj"] is logging_obj
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_keeps_synthetic_request_headers_metadata():
|
||||
"""The seed must not clobber the metadata bucket ``_convert_mcp_to_llm_format``
|
||||
builds: guardrails such as ``MCPJWTSigner`` read ``metadata["headers"]`` off
|
||||
it."""
|
||||
capture: dict = {}
|
||||
await _run_pre_call(_bare_manager(), _fake_proxy_logging(capture), _FakeLoggingObj())
|
||||
|
||||
assert capture["data"]["metadata"]["headers"] == {"x-forwarded-for": "1.2.3.4"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_bridges_allowed_evaluation_onto_request_logger():
|
||||
"""An allowed ``pre_mcp_call`` evaluation must land on the request logger, which
|
||||
is what the monitor's "Total Evaluations" counts."""
|
||||
capture: dict = {}
|
||||
logging_obj = _FakeLoggingObj()
|
||||
plo = _fake_proxy_logging(capture, guardrail_effect=_record_guardrail())
|
||||
|
||||
await _run_pre_call(_bare_manager(), plo, logging_obj)
|
||||
|
||||
assert logging_obj.recorded_guardrails == [{"guardrail_name": "test-guardrail", "guardrail_status": "success"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_bridges_blocked_evaluation_before_reraising():
|
||||
"""A block raises straight out of ``pre_call_tool_check``, and the failure
|
||||
spend-log row that "Total Blocked" counts is built from this logger further up
|
||||
the stack. So the record has to be attached before the exception leaves the
|
||||
frame -- hence the bridge lives in a ``finally``."""
|
||||
capture: dict = {}
|
||||
logging_obj = _FakeLoggingObj()
|
||||
plo = _fake_proxy_logging(capture, guardrail_effect=_blocking_guardrail())
|
||||
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
await _run_pre_call(_bare_manager(), plo, logging_obj)
|
||||
|
||||
assert logging_obj.recorded_guardrails == [
|
||||
{"guardrail_name": "test-guardrail", "guardrail_status": "guardrail_intervened"}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_without_logging_obj_is_unchanged():
|
||||
"""Callers that thread no logger are unaffected: the seed is an explicit
|
||||
``None`` (which every consumer reads via ``.get``) and nothing is bridged.
|
||||
Guards against the bridge assuming a logger exists."""
|
||||
capture: dict = {}
|
||||
plo = _fake_proxy_logging(capture, guardrail_effect=_record_guardrail())
|
||||
mgr = _bare_manager()
|
||||
|
||||
result = await 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 result == {}
|
||||
assert capture["data"]["litellm_logging_obj"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_during_hook_seeds_and_bridges_onto_request_logger():
|
||||
"""``during_mcp_call`` evaluations need the same treatment. The task is awaited
|
||||
before the tool call's success logging runs, so the record is serialized with
|
||||
that call."""
|
||||
capture: dict = {}
|
||||
logging_obj = _FakeLoggingObj()
|
||||
plo = _fake_proxy_logging(capture, guardrail_effect=_record_guardrail())
|
||||
|
||||
await _bare_manager()._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=logging_obj,
|
||||
)
|
||||
|
||||
assert capture["data"]["litellm_logging_obj"] is logging_obj
|
||||
assert logging_obj.recorded_guardrails == [{"guardrail_name": "test-guardrail", "guardrail_status": "success"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_during_hook_bridges_even_when_hook_raises():
|
||||
"""A during-call guardrail block must still be recorded before the task's
|
||||
exception propagates to the ``asyncio.gather`` in ``call_tool``."""
|
||||
capture: dict = {}
|
||||
logging_obj = _FakeLoggingObj()
|
||||
plo = _fake_proxy_logging(capture, guardrail_effect=_blocking_guardrail())
|
||||
|
||||
task = _bare_manager()._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=logging_obj,
|
||||
)
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
await task
|
||||
|
||||
assert logging_obj.recorded_guardrails == [
|
||||
{"guardrail_name": "test-guardrail", "guardrail_status": "guardrail_intervened"}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_failure_does_not_mask_a_guardrail_block():
|
||||
"""Recording is best-effort bookkeeping. If the bridge itself raises, the guardrail's
|
||||
block must still be what the caller sees, not a bookkeeping error.
|
||||
|
||||
The bridge is forced to fail by making the logger's ``model_call_details`` raise, and
|
||||
the swallow is asserted (not just the surviving exception type) so the test cannot go
|
||||
vacuous if a refactor stops the bridge from touching that attribute.
|
||||
"""
|
||||
capture: dict = {}
|
||||
plo = _fake_proxy_logging(capture, guardrail_effect=_blocking_guardrail())
|
||||
|
||||
broken_logging_obj = mock.MagicMock()
|
||||
type(broken_logging_obj).model_call_details = mock.PropertyMock(side_effect=RuntimeError("boom"))
|
||||
|
||||
with mock.patch.object(MOD.verbose_logger, "warning") as warn:
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
await _run_pre_call(_bare_manager(), plo, broken_logging_obj)
|
||||
|
||||
assert warn.call_count == 1, "the bridge did not actually fail, so this test proves nothing"
|
||||
assert "boom" in str(warn.call_args)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_tool_threads_logging_obj_into_both_hooks():
|
||||
"""``call_tool`` is the single entry point every MCP dispatch route funnels
|
||||
through, so it must hand the logger to both guardrail hook sites."""
|
||||
mgr = _bare_manager()
|
||||
logging_obj = _FakeLoggingObj()
|
||||
seen: dict = {}
|
||||
|
||||
async def _fake_pre_call_tool_check(**kwargs):
|
||||
seen["pre_call"] = kwargs.get("litellm_logging_obj")
|
||||
return {}
|
||||
|
||||
def _fake_during_hook_task(**kwargs):
|
||||
seen["during_call"] = kwargs.get("litellm_logging_obj")
|
||||
return asyncio.get_running_loop().create_future()
|
||||
|
||||
mgr.pre_call_tool_check = _fake_pre_call_tool_check
|
||||
mgr._create_during_hook_task = _fake_during_hook_task
|
||||
mgr._resolve_mcp_server_for_tool_call = lambda server_name, name: mock.MagicMock(spec_path=None)
|
||||
mgr._resolve_oauth2_headers_for_tool_call = mock.AsyncMock(return_value=None)
|
||||
mgr._call_regular_mcp_tool = mock.AsyncMock(return_value=mock.MagicMock())
|
||||
|
||||
with mock.patch.object(MOD, "_resolve_byok_mcp_auth_header", mock.AsyncMock(return_value=None)):
|
||||
await mgr.call_tool(
|
||||
server_name="s",
|
||||
name="t",
|
||||
arguments={},
|
||||
proxy_logging_obj=mock.MagicMock(),
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
assert seen == {"pre_call": logging_obj, "during_call": logging_obj}
|
||||
|
|
@ -450,6 +450,42 @@ 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 no guardrail evaluation is bridged onto
|
||||
the request logger, 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):
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue