diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 82b820d8cd9..9f70ada1497 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 9a1fffd5a67..b9e96f136e0 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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 diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index cb680dc8b86..2018d6f4f98 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -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: 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..e69f17222a1 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_logging_obj.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_logging_obj.py new file mode 100644 index 00000000000..274bfe59870 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_logging_obj.py @@ -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 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 d60fff66c44..e79764a06b5 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,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): """