diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 5b5261fab6b..326abd5c6a3 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -574,11 +574,10 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac Useful if you want to modify the standard logging payload after the MCP tool call is made. - To change what the caller sends back to the MCP client, mutate ``response_obj`` - in place: every call site discards the returned object, because the - dispatcher unwraps it to ``mcp_tool_call_response`` (a raw content list, not - a ``CallToolResult``) which the tool-call paths cannot forward. Guardrails - that mask or reject tool output should use ``post_mcp_call`` instead. + Modify ``mcp_tool_call_response`` in place or return a replacement response + object to change what the caller sends back to the MCP client. Content rewrites + discard stale structured output and mark those results as tool errors. + Use ``post_mcp_call`` guardrails for schema-preserving structured redaction. """ return None diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 357ee51a44a..6e5a37a7226 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -221,7 +221,7 @@ from .initialize_dynamic_callback_params import ( from .specialty_caches.dynamic_logging_cache import DynamicLoggingCache if TYPE_CHECKING: - from mcp.types import EmbeddedResource, ImageContent, TextContent + from mcp.types import CallToolResult, EmbeddedResource, ImageContent, TextContent from litellm.integrations.otel.logger import OpenTelemetryV2 from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config @@ -1634,15 +1634,11 @@ class Logging(LiteLLMLoggingBaseClass): async def async_post_mcp_tool_call_hook( self, kwargs: dict, - response_obj: Any, + response_obj: "CallToolResult", start_time: datetime.datetime, end_time: datetime.datetime, - ): - """ - Post MCP Tool Call Hook - - Use this to modify the MCP tool call response before it is returned to the user. - """ + ) -> "CallToolResult": + """Apply ordered MCP content callbacks to the result returned to the caller.""" from litellm.types.llms.base import HiddenParams from litellm.types.mcp import MCPPostCallResponseObject @@ -1650,24 +1646,51 @@ class Logging(LiteLLMLoggingBaseClass): dynamic_success_callbacks=self.dynamic_success_callbacks, global_callbacks=litellm.success_callback, ) - post_mcp_tool_call_response_obj: Final[MCPPostCallResponseObject] = MCPPostCallResponseObject( - mcp_tool_call_response=response_obj, hidden_params=HiddenParams() - ) + hidden_params = HiddenParams() for callback in callbacks: try: if isinstance(callback, CustomLogger): - response: MCPPostCallResponseObject | None = await callback.async_post_mcp_tool_call_hook( - kwargs=kwargs, - response_obj=post_mcp_tool_call_response_obj, - start_time=start_time, - end_time=end_time, + original_content = copy.deepcopy(response_obj.content) + original_structured_content = copy.deepcopy(response_obj.structured_content) + callback_response = MCPPostCallResponseObject( + mcp_tool_call_response=copy.deepcopy(original_content), hidden_params=hidden_params ) - ###################################################################### - # if any of the callbacks modify the response, use the modified response - # current implementation returns the first modified response - ###################################################################### - if response is not None: - response_obj = self._parse_post_mcp_call_hook_response(response=response) + try: + response = await callback.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=callback_response, + start_time=start_time, + end_time=end_time, + ) + hook_content = ( + self._parse_post_mcp_call_hook_response(response=response) + if response is not None + else callback_response.mcp_tool_call_response + ) + if response is not None: + hidden_params = response.hidden_params + except Exception as e: + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e + ) + hook_content = None + structured_replacement_matches = ( + response_obj.structured_content != original_structured_content + and ( + hook_content is None + or hook_content == original_content + or response_obj.content == hook_content + ) + ) + if hook_content is not None and hook_content != original_content: + response_obj.content[:] = hook_content + if ( + response_obj.content != original_content + and response_obj.structured_content is not None + and not structured_replacement_matches + ): + response_obj.structured_content = None + response_obj.is_error = True except Exception as e: verbose_logger.exception("LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e) return response_obj diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index fcee3483e15..2bd6d186f24 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -2198,7 +2198,7 @@ async def _fire_mcp_tool_call_logging( from litellm.proxy.proxy_server import proxy_logging_obj logging_obj.post_call(original_response=result) - await logging_obj.async_post_mcp_tool_call_hook( + result = await logging_obj.async_post_mcp_tool_call_hook( kwargs=logging_obj.model_call_details, response_obj=result, start_time=start_time, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 88a2b92c680..319f10a9b22 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -872,7 +872,7 @@ class LiteLLM_Proxy_MCP_Handler: if litellm_logging_obj: try: litellm_logging_obj.post_call(original_response=result) - await litellm_logging_obj.async_post_mcp_tool_call_hook( + result = await litellm_logging_obj.async_post_mcp_tool_call_hook( kwargs=litellm_logging_obj.model_call_details, response_obj=result, start_time=start_time, diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index d38a11ec927..dfd250338ff 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -2962,7 +2962,7 @@ async def test_call_mcp_tool_uses_manager_permission_lookup(): mcp_info={"server_name": "test_server"}, ) - expected_response = [TextContent(type="text", text="ok")] + expected_response = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) with ( patch.object( @@ -3038,7 +3038,7 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission mcp_info={"server_name": "test_server"}, ) - expected_response = [TextContent(type="text", text="ok")] + expected_response = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) with ( patch.object( diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index b0fc725b50a..4541250e896 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -5,16 +5,15 @@ import json import logging import os import sys +import time from collections.abc import Callable, Iterator, Mapping from types import MappingProxyType from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch -import pytest - -import time - import httpx +import pytest +from mcp.types import AudioContent, CallToolResult, ImageContent, TextContent from openai._legacy_response import HttpxBinaryResponseContent import litellm @@ -51,6 +50,272 @@ def logging_obj(): ) +@pytest.mark.asyncio +async def test_async_post_mcp_tool_call_hook_preserves_and_returns_content(logging_obj): + from litellm.types.mcp import MCPPostCallResponseObject + + class RedactingLogger(CustomLogger): + async def async_post_mcp_tool_call_hook( + self, + kwargs: dict[str, object], + response_obj: MCPPostCallResponseObject, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> MCPPostCallResponseObject: + assert isinstance(response_obj.mcp_tool_call_response, list) + assert isinstance(response_obj.mcp_tool_call_response[0], TextContent) + response_obj.mcp_tool_call_response = [TextContent(type="text", text="[REDACTED]")] + return response_obj + + logging_obj.dynamic_success_callbacks = [RedactingLogger()] + result = CallToolResult(content=[TextContent(type="text", text="SECRET-1234")], isError=False) + + hooked_content = await logging_obj.async_post_mcp_tool_call_hook( + kwargs=logging_obj.model_call_details, + response_obj=result, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), + ) + + assert hooked_content.content == [TextContent(type="text", text="[REDACTED]")] + + +@pytest.mark.asyncio +async def test_async_post_mcp_tool_call_hook_chains_every_callback(logging_obj): + from litellm.types.mcp import MCPPostCallResponseObject + + class ReplacingLogger(CustomLogger): + def __init__(self, old: str, new: str) -> None: + super().__init__() + self.old: Final = old + self.new: Final = new + self.seen: list[str] = [] # mutable-ok: test records what each callback observed + + async def async_post_mcp_tool_call_hook( + self, + kwargs: dict[str, object], + response_obj: MCPPostCallResponseObject, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> MCPPostCallResponseObject: + first = response_obj.mcp_tool_call_response[0] + assert isinstance(first, TextContent) + self.seen.append(first.text) + return MCPPostCallResponseObject( + mcp_tool_call_response=[TextContent(type="text", text=first.text.replace(self.old, self.new))], + hidden_params=response_obj.hidden_params, + ) + + first_logger: Final = ReplacingLogger("SECRET", "[S]") + second_logger: Final = ReplacingLogger("1234", "[N]") + logging_obj.dynamic_success_callbacks = [first_logger, second_logger] + result = CallToolResult(content=[TextContent(type="text", text="SECRET-1234")], isError=False) + + hooked_content = await logging_obj.async_post_mcp_tool_call_hook( + kwargs=logging_obj.model_call_details, + response_obj=result, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), + ) + + assert first_logger.seen == ["SECRET-1234"] + assert second_logger.seen == ["[S]-1234"] + assert hooked_content.content == [TextContent(type="text", text="[S]-[N]")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["replace", "inplace", "empty", "inplace_none", "replace_none"]) +@pytest.mark.parametrize("structured", [False, True]) +async def test_mcp_content_rewrite_never_returns_stale_structured_data(logging_obj, mode, structured): + from litellm.types.llms.base import HiddenParams + from litellm.types.mcp import MCPPostCallResponseObject + + class Redactor(CustomLogger): + async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, end_time): + block = response_obj.mcp_tool_call_response[0] + assert isinstance(block, TextContent) + if mode in ("inplace", "inplace_none"): + block.text = "[REDACTED]" + return None if mode == "inplace_none" else response_obj + if mode == "replace_none": + response_obj.mcp_tool_call_response = [TextContent(type="text", text="[REDACTED]")] + return None + return MCPPostCallResponseObject( + mcp_tool_call_response=[] if mode == "empty" else [TextContent(type="text", text="[REDACTED]")], + hidden_params=HiddenParams(response_cost=0.25), + ) + + logging_obj.dynamic_success_callbacks = [Redactor()] + result = CallToolResult( + content=[TextContent(type="text", text="SECRET-1234")], + structured_content={"nested": {"secret": "SECRET-1234"}} if structured else None, + meta={"request": "trace-1"}, + ) + logging_obj.model_call_details["original_response"] = result + returned = await logging_obj.async_post_mcp_tool_call_hook( + kwargs=logging_obj.model_call_details, + response_obj=result, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), + ) + assert "SECRET-1234" not in result.model_dump_json(by_alias=True) + assert returned is result + assert result.content == ([] if mode == "empty" else [TextContent(type="text", text="[REDACTED]")]) + assert result.structured_content is None + assert result.is_error is structured + assert result.meta == {"request": "trace-1"} + assert logging_obj.model_call_details["original_response"] is result + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["none", "cost", "direct", "block", "exception"]) +async def test_mcp_callbacks_preserve_effective_result_and_cost(logging_obj, mode): + from litellm.types.llms.base import HiddenParams + from litellm.types.mcp import MCPPostCallResponseObject + + class Callback(CustomLogger): + async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, end_time): + if mode == "exception": + response_obj.mcp_tool_call_response[0].text = "discarded" + raise ValueError("non-blocking callback") + if mode in ("direct", "block"): + original = kwargs["original_response"] + original.content = [TextContent(type="text", text="safe")] + original.structured_content = {"result": "safe"} + original.is_error = mode == "block" + if mode == "none" or mode == "direct": + return None + return MCPPostCallResponseObject( + mcp_tool_call_response=response_obj.mcp_tool_call_response, + hidden_params=HiddenParams(response_cost=0.25), + ) + + logging_obj.dynamic_success_callbacks = [Callback()] + result = CallToolResult( + content=[TextContent(type="text", text="original")], structured_content={"result": "original"} + ) + logging_obj.model_call_details["original_response"] = result + returned = await logging_obj.async_post_mcp_tool_call_hook( + kwargs=logging_obj.model_call_details, + response_obj=result, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), + ) + expected = "safe" if mode in ("direct", "block") else "original" + assert result.content == [TextContent(type="text", text=expected)] + assert result.structured_content == {"result": expected} + assert result.is_error is (mode == "block") + assert returned is result + assert logging_obj.model_call_details.get("response_cost") == (0.25 if mode in ("cost", "block") else None) + + +@pytest.mark.asyncio +async def test_mcp_callback_cancellation_propagates_without_mutating_result(logging_obj): + class CancelledCallback(CustomLogger): + async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, end_time): + response_obj.mcp_tool_call_response[0].text = "partial" + raise asyncio.CancelledError + + logging_obj.dynamic_success_callbacks = [CancelledCallback()] + result = CallToolResult(content=[TextContent(type="text", text="original")]) + with pytest.raises(asyncio.CancelledError): + await logging_obj.async_post_mcp_tool_call_hook( + kwargs={}, + response_obj=result, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), + ) + assert result.content == [TextContent(type="text", text="original")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("callbacks", [[], ["prometheus"]]) +@pytest.mark.parametrize("is_error", [False, True]) +async def test_mcp_without_custom_callbacks_preserves_mixed_content(logging_obj, callbacks, is_error): + logging_obj.dynamic_success_callbacks = callbacks + result = CallToolResult( + content=[ + TextContent(type="text", text="ok"), + ImageContent(type="image", data="aW1n", mime_type="image/png"), + AudioContent(type="audio", data="c291bmQ=", mime_type="audio/wav"), + ], + structured_content={"result": "ok"}, + is_error=is_error, + ) + before = result.model_dump() + returned = await logging_obj.async_post_mcp_tool_call_hook( + kwargs={}, + response_obj=result, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), + ) + assert returned is result + assert returned.model_dump() == before + + +@pytest.mark.asyncio +@pytest.mark.parametrize("replace_structured", [False, True]) +@pytest.mark.parametrize("same_content", [False, True]) +async def test_mcp_native_structured_replacement_must_match_returned_content( + logging_obj, replace_structured, same_content +): + from litellm.types.mcp import MCPPostCallResponseObject + + class NativeReplacement(CustomLogger): + async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, end_time): + original = kwargs["original_response"] + original.content[0].text = "native-safe" + if replace_structured: + original.structured_content["result"] = "native-safe" + return MCPPostCallResponseObject( + mcp_tool_call_response=[TextContent(type="text", text="native-safe" if same_content else "final-safe")], + hidden_params=response_obj.hidden_params, + ) + + result = CallToolResult( + content=[TextContent(type="text", text="SECRET-1234")], + structured_content={"result": "SECRET-1234"}, + ) + logging_obj.dynamic_success_callbacks = [NativeReplacement()] + returned = await logging_obj.async_post_mcp_tool_call_hook( + kwargs={"original_response": result}, response_obj=result, + start_time=datetime.datetime.now(), end_time=datetime.datetime.now(), + ) + assert returned is result + assert result.content == [TextContent(type="text", text="native-safe" if same_content else "final-safe")] + assert result.structured_content == ({"result": "native-safe"} if replace_structured and same_content else None) + assert result.is_error is not (replace_structured and same_content) + assert "SECRET-1234" not in result.model_dump_json() + + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["none", "wrapper", "exception"]) +@pytest.mark.parametrize("structured", [False, True]) +async def test_mcp_direct_content_edit_invalidates_stale_structured_data(logging_obj, mode, structured): + class DirectRedactor(CustomLogger): + async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, end_time): + kwargs["original_response"].content[0].text = "[REDACTED]" + if mode == "exception": + raise ValueError("non-blocking callback after direct edit") + return response_obj if mode == "wrapper" else None + + result = CallToolResult( + content=[TextContent(type="text", text="SECRET-1234")], + structured_content={"result": "SECRET-1234"} if structured else None, + ) + logging_obj.dynamic_success_callbacks = [DirectRedactor()] + returned = await logging_obj.async_post_mcp_tool_call_hook( + kwargs={"original_response": result}, response_obj=result, + start_time=datetime.datetime.now(), end_time=datetime.datetime.now(), + ) + assert returned is result + assert result.content == [TextContent(type="text", text="[REDACTED]")] + assert result.structured_content is None + assert result.is_error is structured + assert "SECRET-1234" not in result.model_dump_json() + + def test_get_combined_callback_list_preserves_insertion_order(logging_obj): assert logging_obj.get_combined_callback_list( dynamic_success_callbacks=["prometheus", "langfuse", "datadog", "otel", "s3"], diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 6588ab4c87b..ea82c9b069d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -8755,7 +8755,7 @@ def _call_tool_result(is_error: bool, text: str) -> CallToolResult: def _mock_mcp_logging_obj() -> MagicMock: logging_obj = MagicMock() logging_obj.model_call_details = {} - logging_obj.async_post_mcp_tool_call_hook = AsyncMock() + logging_obj.async_post_mcp_tool_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response_obj"]) logging_obj.async_success_handler = AsyncMock() logging_obj.async_failure_handler = AsyncMock() return logging_obj @@ -8857,6 +8857,64 @@ async def test_fire_mcp_tool_call_logging_success_path_unchanged(): proxy_logging_mock.post_call_failure_hook.assert_not_awaited() +@pytest.mark.asyncio +async def test_fire_mcp_tool_call_logging_applies_hook_content(): + from litellm.integrations.custom_logger import CustomLogger + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy._experimental.mcp_server.server import ( + _fire_mcp_tool_call_logging, + ) + from litellm.types.mcp import MCPPostCallResponseObject + + class RedactingLogger(CustomLogger): + async def async_post_mcp_tool_call_hook( + self, + kwargs: dict[str, object], + response_obj: MCPPostCallResponseObject, + start_time: datetime, + end_time: datetime, + ) -> MCPPostCallResponseObject: + assert isinstance(response_obj.mcp_tool_call_response, list) + assert isinstance(response_obj.mcp_tool_call_response[0], TextContent) + response_obj.mcp_tool_call_response = [TextContent(type="text", text="[REDACTED]")] + return response_obj + + logging_obj = Logging( + model="MCP: weather/get_forecast", + messages=[{"role": "user", "content": "tool call"}], + stream=False, + call_type="call_mcp_tool", + start_time=datetime.now(), + litellm_call_id="test-mcp-hook-content", + function_id="test-fn", + dynamic_success_callbacks=[RedactingLogger()], + ) + proxy_logging_mock = _mock_mcp_proxy_logging() + result = CallToolResult( + content=[TextContent(type="text", text="SECRET-1234")], + structuredContent={"result": "SECRET-1234"}, + isError=False, + ) + + with patch( # test-quality-ok: [TQ008] inject proxy logging collaborator + "litellm.proxy.proxy_server.proxy_logging_obj", + proxy_logging_mock, + ): + hooked_result = await _fire_mcp_tool_call_logging( + logging_obj=logging_obj, + result=result, + start_time=datetime.now(), + end_time=datetime.now(), + user_api_key_auth=UserAPIKeyAuth(api_key="test-key", user_id="test-user"), + request_data={}, + ) + + assert isinstance(hooked_result.content[0], TextContent) + assert hooked_result.content[0].text == "[REDACTED]" + assert hooked_result.structured_content is None + assert hooked_result.is_error is True + + @pytest.mark.asyncio async def test_fire_mcp_tool_call_logging_iserror_without_auth_skips_failure_hook(): """Without a UserAPIKeyAuth the failure handlers still fire but the proxy @@ -8871,7 +8929,7 @@ async def test_fire_mcp_tool_call_logging_iserror_without_auth_skips_failure_hoo with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock): await _fire_mcp_tool_call_logging( logging_obj=logging_obj, - result={"isError": True, "content": [{"type": "text", "text": "denied"}]}, + result=_call_tool_result(True, "denied"), start_time=datetime.now(), end_time=datetime.now(), ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py index 826edab694d..2e3bf760e68 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py @@ -827,3 +827,38 @@ class TestCiscoAIDefenseJsonRpcSuccessEnvelope: assert unwrapped is verdict else: assert unwrapped == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["block", "redact"]) +async def test_cisco_native_hook_through_logging_preserves_sanitized_result(action): + from mcp.types import CallToolResult, TextContent + from litellm.litellm_core_utils.litellm_logging import Logging + + guardrail = _make_guardrail( + inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"] + ) + verdict = ( + _violation_response(url=MCP_URL) if action == "block" + else _redact_response(sanitized_text="[REDACTED]", url=MCP_URL) + ) + result = CallToolResult( + content=[TextContent(type="text", text="SECRET-1234")], + structured_content={"result": "SECRET-1234"}, + ) + logging_obj = Logging( + model="MCP: probe/search", messages=[], stream=False, call_type="call_mcp_tool", + start_time=datetime.now(), litellm_call_id="cisco-hook", function_id="cisco-hook", + dynamic_success_callbacks=[guardrail], + ) + logging_obj.model_call_details.update({"name": "search", "arguments": {}, "original_response": result}) + with _patch_inspection_post(guardrail, AsyncMock(return_value=verdict)): + returned = await logging_obj.async_post_mcp_tool_call_hook( + kwargs=logging_obj.model_call_details, response_obj=result, + start_time=datetime.now(), end_time=datetime.now(), + ) + assert returned is result + assert "SECRET-1234" not in returned.model_dump_json() + assert returned.is_error is (action == "block") + assert ("Blocked by Cisco AI Defense" if action == "block" else "[REDACTED]") in returned.content[0].text + assert returned.structured_content == {"result": returned.content[0].text} 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 83537c236a3..56e40206ebe 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 @@ -1,13 +1,15 @@ +import importlib import subprocess import sys import textwrap import types +from typing import Any, cast from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException +from mcp.types import CallToolResult, TextContent from openai.types.responses.tool_param import Mcp -import importlib from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing from litellm.responses import main as responses_main @@ -15,10 +17,9 @@ from litellm.responses.mcp import litellm_proxy_mcp_handler as mcp_handler_modul from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) -from typing import Any, cast from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.types.utils import ModelResponse from litellm.types.responses.main import OutputFunctionToolCall +from litellm.types.utils import ModelResponse class _DummyMCPResult: @@ -492,6 +493,157 @@ async def test_execute_tool_calls_threads_logging_obj_into_call_tool(monkeypatch assert call_tool_mock.await_args.kwargs["litellm_logging_obj"] is sentinel_logging_obj +@pytest.mark.asyncio +async def test_execute_tool_calls_applies_post_call_hook_content(monkeypatch): + proxy_module = types.SimpleNamespace(proxy_logging_obj=None) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) + + result = CallToolResult( + content=[TextContent(type="text", text="SECRET-1234")], + structuredContent={"result": "SECRET-1234"}, + isError=False, + ) + fake_manager = types.SimpleNamespace( + get_registry=MagicMock(return_value={}), + call_tool=AsyncMock(return_value=result), + _get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + + logging_obj = MagicMock() + logging_obj.model_call_details = {} + logging_obj.async_post_mcp_tool_call_hook = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="[REDACTED]")], is_error=True)) + 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: (logging_obj, None)) + + tool_name = "deepwiki-read_wiki_structure" + results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=[{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}], + user_api_key_auth=None, + ) + + assert results == [{"tool_call_id": "call-1", "result": "[REDACTED]", "name": tool_name}] + assert logging_obj.async_success_handler.await_args.kwargs["result"].content[0].text == "[REDACTED]" + assert logging_obj.async_success_handler.await_args.kwargs["result"].structured_content is None + + +@pytest.mark.asyncio +async def test_execute_tool_calls_returns_proxy_result_without_logging(monkeypatch): + result = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_mcp_call_hook = AsyncMock(side_effect=lambda response, **_: response) + monkeypatch.setitem( + sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(proxy_logging_obj=proxy_logging_obj) + ) + + fake_manager = types.SimpleNamespace( + get_registry=MagicMock(return_value={}), + call_tool=AsyncMock(return_value=result), + _get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") + monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (None, None)) + + tool_name = "deepwiki-read_wiki_structure" + results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=[{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}], + user_api_key_auth=None, + ) + + assert results == [{"tool_call_id": "call-1", "result": "ok", "name": tool_name}] + proxy_logging_obj.post_mcp_call_hook.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_execute_tool_calls_passes_logging_details_to_proxy_hook(monkeypatch): + result = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_mcp_call_hook = AsyncMock(side_effect=lambda response, **_: response) + monkeypatch.setitem( + sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(proxy_logging_obj=proxy_logging_obj) + ) + + fake_manager = types.SimpleNamespace( + get_registry=MagicMock(return_value={}), + call_tool=AsyncMock(return_value=result), + _get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + logging_obj = MagicMock() + logging_obj.model_call_details = {"request_id": "request-1"} + logging_obj.async_post_mcp_tool_call_hook = AsyncMock(return_value=result) + 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: (logging_obj, None)) + + tool_name = "deepwiki-read_wiki_structure" + results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=[{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}], + user_api_key_auth=None, + ) + + assert results == [{"tool_call_id": "call-1", "result": "ok", "name": tool_name}] + assert proxy_logging_obj.post_mcp_call_hook.await_args.kwargs["request_data"] == logging_obj.model_call_details + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_stage", ["post_call_hook", "success_handler"]) +async def test_execute_tool_calls_continues_when_post_call_logging_fails(monkeypatch, failure_stage: str): + proxy_module = types.SimpleNamespace(proxy_logging_obj=None) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) + + result = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) + fake_manager = types.SimpleNamespace( + get_registry=MagicMock(return_value={}), + call_tool=AsyncMock(return_value=result), + _get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + + logging_obj = MagicMock() + logging_obj.model_call_details = {} + logging_obj.post_call = MagicMock() + logging_obj.async_post_mcp_tool_call_hook = AsyncMock( + side_effect=RuntimeError("hook failed") if failure_stage == "post_call_hook" else None, + return_value=result, + ) + logging_obj.async_success_handler = AsyncMock( + side_effect=RuntimeError("success logging failed") if failure_stage == "success_handler" else None + ) + handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") + monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (logging_obj, None)) + + tool_name = "deepwiki-read_wiki_structure" + results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=[{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}], + user_api_key_auth=None, + ) + + assert results == [{"tool_call_id": "call-1", "result": "ok", "name": tool_name}] + + @pytest.mark.asyncio async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch): """