diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a1adda2bc95..0ad372064cd 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index f237529b319..17457c3362f 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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 diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 56818717c09..0321034dffe 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -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: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py new file mode 100644 index 00000000000..64d926bc5e3 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py new file mode 100644 index 00000000000..24e6d2de10d --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py @@ -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} diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index 4981caa10c3..418c716b1a1 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -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): """