From caf5f7f8aeb248b64fa95d246b7376a8428656ce Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Wed, 21 Jan 2026 14:51:56 +0900 Subject: [PATCH] test: add test --- .../mcp_server/test_mcp_server.py | 135 ++++++++++++++++ .../mcp/test_litellm_proxy_mcp_handler.py | 151 +++++++++++++++++- 2 files changed, 283 insertions(+), 3 deletions(-) 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 7aa176aaac..8c29cf6a59 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 @@ -1817,3 +1817,138 @@ class TestMCPServerManagerReload: mock_get_all.assert_awaited_once() mock_build.assert_awaited_once_with(db_row) assert manager.registry["server-1"] is rebuilt_server + + +@pytest.mark.asyncio +async def test_call_mcp_tool_logs_failure_via_post_call_failure_hook(): + """ + Regression test for 6267f168...: + Ensure proxy-side `call_mcp_tool` logs failures via `proxy_logging_obj.post_call_failure_hook`. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + call_mcp_tool, + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport, UserAPIKeyAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP server not available") + + mock_server = MCPServer( + server_id="server-123", + name="test_server", + alias="test_server", + server_name="test_server", + url="https://test-server.com/mcp", + transport=MCPTransport.http, + mcp_info={"server_name": "test_server"}, + ) + + proxy_logging_mock = MagicMock() + proxy_logging_mock.post_call_failure_hook = AsyncMock() + + user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + + with patch.object( + global_mcp_server_manager, + "get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[mock_server.server_id], + ), patch.object( + global_mcp_server_manager, + "get_mcp_server_by_id", + return_value=mock_server, + ), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers_from_mcp_server_names", + new_callable=AsyncMock, + return_value=[mock_server], + ), patch( + "litellm.proxy._experimental.mcp_server.server.execute_mcp_tool", + new_callable=AsyncMock, + side_effect=Exception("boom"), + ), patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + proxy_logging_mock, + ): + with pytest.raises(Exception): + await call_mcp_tool( + name="test_server-any_tool", + arguments={"x": 1}, + user_api_key_auth=user_auth, + litellm_call_id="cid", + ) + + proxy_logging_mock.post_call_failure_hook.assert_awaited_once() + assert ( + proxy_logging_mock.post_call_failure_hook.await_args.kwargs.get("route") + == "/mcp/call_tool" + ) + + +@pytest.mark.asyncio +async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enabled(): + """ + Regression test for 872e5b98...: + Ensure list-tools logging path calls `async_success_handler` when enabled. + """ + try: + from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers + from litellm.proxy._types import UserAPIKeyAuth + except ImportError: + pytest.skip("MCP server not available") + + user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + + server_a = MagicMock(name="server_a_obj") + server_a.name = "server_a" + server_a.alias = "server_a" + server_a.server_name = "server_a" + server_a.server_id = "a" + server_a.auth_type = None + server_a.extra_headers = None + + tool_1 = MagicMock() + tool_1.name = "server_a-tool_1" + + dummy_logging_obj = MagicMock() + dummy_logging_obj.model_call_details = {"metadata": {"spend_logs_metadata": {}}} + dummy_logging_obj.async_success_handler = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server_a]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers", + return_value=(None, None), + ), patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + ) as mock_manager, patch( + "litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools", + side_effect=lambda tools, _server: tools, + ), patch( + "litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions", + new=AsyncMock(side_effect=lambda tools, **_: tools), + ), patch( + "litellm.proxy._experimental.mcp_server.server.function_setup", + return_value=(dummy_logging_obj, None), + ): + mock_manager._get_tools_from_server = AsyncMock(return_value=[tool_1]) + + tools = await _get_tools_from_mcp_servers( + user_api_key_auth=user_auth, + mcp_auth_header=None, + mcp_servers=["server_a"], + mcp_server_auth_headers=None, + log_list_tools_to_spendlogs=True, + list_tools_log_source="mcp_protocol", + ) + + assert tools == [tool_1] + dummy_logging_obj.async_success_handler.assert_awaited_once() + assert dummy_logging_obj.async_success_handler.await_args.kwargs["result"] == [tool_1] + + spend_meta = dummy_logging_obj.model_call_details["metadata"]["spend_logs_metadata"] + assert spend_meta["tool_count_total"] == 1 + assert spend_meta["allowed_server_count"] == 1 + assert spend_meta["per_server_tool_counts"]["server_a"] == 1 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 b632e72f56..15fdc7bd0c 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,12 +1,15 @@ import sys import types -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException +import importlib from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) +from typing import Any, cast from litellm.types.utils import ModelResponse from litellm.types.responses.main import OutputFunctionToolCall @@ -22,7 +25,9 @@ def _setup_mcp_call_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) fake_manager = types.SimpleNamespace( - call_tool=AsyncMock(return_value=_DummyMCPResult()) + call_tool=AsyncMock(return_value=_DummyMCPResult()), + # Newer logging path calls this to enrich spend logs metadata + _get_mcp_server_from_tool_name=MagicMock(return_value=None), ) monkeypatch.setattr( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", @@ -31,6 +36,15 @@ def _setup_mcp_call_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: return fake_manager.call_tool +def _setup_proxy_logging(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + """Patch proxy_logging_obj so failure hook can be asserted.""" + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock() + proxy_module = types.SimpleNamespace(proxy_logging_obj=proxy_logging_obj) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) + return proxy_logging_obj.post_call_failure_hook + + def test_deduplicate_mcp_tools_single_allowed_server(): tools = [{"name": "search"}, {"name": "search"}] # duplicate on purpose @@ -184,7 +198,7 @@ def test_create_follow_up_input_handles_response_function_tool_call(): ) follow_up = LiteLLM_Proxy_MCP_Handler._create_follow_up_input( - response=response, + response=cast(Any, response), tool_results=[], original_input=None, ) @@ -216,6 +230,8 @@ async def test_execute_tool_calls_strips_server_prefix(monkeypatch): 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["name"] == "read_wiki_structure" @@ -236,6 +252,8 @@ async def test_execute_tool_calls_keeps_tool_name_without_prefix(monkeypatch): 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["name"] == tool_name @@ -256,4 +274,131 @@ async def test_execute_tool_calls_keeps_tool_name_when_equal_to_server(monkeypat 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["name"] == tool_name + + +@pytest.mark.asyncio +async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkeypatch): + """ + Regression test for ae4d92ad...: + Ensure responses-side MCP tool execution logs failures via proxy_logging_obj.post_call_failure_hook. + """ + post_call_failure_hook = _setup_proxy_logging(monkeypatch) + + fake_manager = types.SimpleNamespace( + call_tool=AsyncMock( + side_effect=HTTPException(status_code=500, detail="boom") + ) + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + + tool_name = "deepwiki-read_wiki_structure" + tool_calls = [ + {"id": "call-err", "function": {"name": tool_name, "arguments": "{}"}} + ] + + user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user") + + results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=tool_calls, + user_api_key_auth=user_auth, + litellm_call_id="cid", + litellm_trace_id="tid", + ) + + assert len(results) == 1 + assert results[0]["tool_call_id"] == "call-err" + assert results[0]["name"] == tool_name + + post_call_failure_hook.assert_awaited_once() + assert post_call_failure_hook.await_args is not None + assert ( + post_call_failure_hook.await_args.kwargs.get("route") + == "/responses/mcp/call_tool" + ) + + +@pytest.mark.asyncio +async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_function_setup( + monkeypatch, +): + """ + Regression test for ae4d92ad...: + Ensure litellm_call_id / litellm_trace_id are forwarded into function_setup kwargs. + """ + _setup_proxy_logging(monkeypatch) + call_tool_mock = _setup_mcp_call_environment(monkeypatch) + + captured = {} + + def fake_function_setup(*_args, **kwargs): + captured.update(kwargs) + return None, None + + # NOTE: Don't patch via dotted string path here because `litellm.responses` + # is a function attribute on the `litellm` package (shadowing the submodule), + # which breaks monkeypatch's importpath resolution. + handler_module = importlib.import_module( + "litellm.responses.mcp.litellm_proxy_mcp_handler" + ) + monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) + + 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, + litellm_call_id="cid", + litellm_trace_id="tid", + ) + + # Ensure the tool call was attempted (sanity) + assert call_tool_mock.await_count == 1 + + assert captured.get("litellm_call_id") == "cid" + assert captured.get("litellm_trace_id") == "tid" + + +@pytest.mark.asyncio +async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch): + """ + Regression test for 872e5b98...: + Ensure responses-side tool discovery enables list-tools SpendLogs logging flags. + """ + mock_get_tools = AsyncMock(return_value=[]) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.server._get_tools_from_mcp_servers", + mock_get_tools, + ) + + # Patch manager methods used by _get_mcp_tools_from_manager to avoid needing full UserAPIKeyAuth fields. + fake_manager = types.SimpleNamespace( + get_allowed_mcp_servers=AsyncMock(return_value=[]), + get_mcp_servers_from_ids=MagicMock(return_value=[]), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + + user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user") + tools, _server_names = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( + user_api_key_auth=user_auth, + mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}], + ) + + assert tools == [] + assert mock_get_tools.await_count == 1 + assert mock_get_tools.await_args is not None + assert mock_get_tools.await_args.kwargs["log_list_tools_to_spendlogs"] is True + assert mock_get_tools.await_args.kwargs["list_tools_log_source"] == "responses"