diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 2149f079a3d..794ea7d33a3 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,5 +1,6 @@ import asyncio import importlib +import traceback from datetime import datetime from typing import ( Any, @@ -757,7 +758,7 @@ if MCP_AVAILABLE: } @router.post("/tools/call", dependencies=[Depends(user_api_key_auth)]) - async def call_tool_rest_api( + async def call_tool_rest_api( # noqa: PLR0915 request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): @@ -854,20 +855,33 @@ if MCP_AVAILABLE: target_server, user_api_key_dict ) - # Call execute_mcp_tool directly (permission checks already done) - result = await execute_mcp_tool( - name=tool_name, - arguments=tool_arguments, - allowed_mcp_servers=allowed_mcp_servers, - start_time=datetime.now(), - user_api_key_auth=data.get("user_api_key_auth"), - mcp_auth_header=data.get("mcp_auth_header"), - mcp_server_auth_headers=data.get("mcp_server_auth_headers"), - oauth2_headers=user_oauth_extra_headers or data.get("oauth2_headers"), - raw_headers=data.get("raw_headers"), - litellm_logging_obj=data.get("litellm_logging_obj"), - requested_server_id=canonical_server_id, - ) + # Call execute_mcp_tool directly (permission checks already done). + # This endpoint bypasses the @client decorator that wraps + # call_mcp_tool, so it fires failure logging here the way @client does + # for the JSON-RPC path; execute_mcp_tool drives the success side. + start_time = datetime.now() + try: + result = await execute_mcp_tool( + name=tool_name, + arguments=tool_arguments, + allowed_mcp_servers=allowed_mcp_servers, + start_time=start_time, + user_api_key_auth=data.get("user_api_key_auth"), + mcp_auth_header=data.get("mcp_auth_header"), + mcp_server_auth_headers=data.get("mcp_server_auth_headers"), + oauth2_headers=user_oauth_extra_headers + or data.get("oauth2_headers"), + raw_headers=data.get("raw_headers"), + litellm_logging_obj=logging_obj, + requested_server_id=canonical_server_id, + ) + except Exception as e: + if logging_obj is not None: + logging_obj.call_type = CallTypes.call_mcp_tool.value + await logging_obj.async_failure_handler( + e, traceback.format_exc(), start_time, datetime.now() + ) + raise return result except MCPMissingUserEnvVarsError as e: verbose_logger.info( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 746fc4e7d3f..331f30c319b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2731,6 +2731,19 @@ if MCP_AVAILABLE: local_content = await _handle_local_mcp_tool(original_tool_name, arguments) response = CallToolResult(content=cast(Any, local_content), isError=False) + if litellm_logging_obj is not None: + litellm_logging_obj.post_call(original_response=response) + end_time = datetime.now() + await litellm_logging_obj.async_post_mcp_tool_call_hook( + kwargs=litellm_logging_obj.model_call_details, + response_obj=response, + start_time=start_time, + end_time=end_time, + ) + litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value + await litellm_logging_obj.async_success_handler( + result=response, start_time=start_time, end_time=end_time + ) return response @client @@ -2749,9 +2762,6 @@ if MCP_AVAILABLE: Call a specific tool with the provided arguments (handles prefixed tool names). """ start_time = datetime.now() - litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get( - "litellm_logging_obj", None - ) try: if arguments is None: @@ -2811,19 +2821,6 @@ if MCP_AVAILABLE: ) raise - if litellm_logging_obj: - litellm_logging_obj.post_call(original_response=response) - end_time = datetime.now() - await litellm_logging_obj.async_post_mcp_tool_call_hook( - kwargs=litellm_logging_obj.model_call_details, - response_obj=response, - start_time=start_time, - end_time=end_time, - ) - litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value - await litellm_logging_obj.async_success_handler( - result=response, start_time=start_time, end_time=end_time - ) return response async def mcp_get_prompt( 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 1c31f437363..ca70b6dd053 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 @@ -4301,6 +4301,71 @@ async def test_get_tools_from_mcp_servers_returns_tools_when_success_logging_fai dummy_logging_obj.async_success_handler.assert_awaited_once() +@pytest.mark.asyncio +async def test_execute_mcp_tool_drives_success_logging(): + """ + Regression: the REST/UI tool-call endpoint (`POST /mcp-rest/tools/call`) + invokes `execute_mcp_tool` directly instead of the `@client`-wrapped + `call_mcp_tool`, so `execute_mcp_tool` must drive the success logging itself. + Without it, `async_success_handler` never runs for dashboard-initiated tool + calls and no OTel MCP tool-call span (or spend log) is emitted. + """ + from mcp.types import CallToolResult, TextContent + + from litellm.proxy._experimental.mcp_server.server import ( + execute_mcp_tool, + global_mcp_tool_registry, + ) + from litellm.types.utils import CallTypes + + 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"}, + ) + + tool_result = CallToolResult( + content=[TextContent(type="text", text="ok")], isError=False + ) + + logging_obj = MagicMock() + logging_obj.model_call_details = {} + logging_obj.post_call = MagicMock() + logging_obj.async_post_mcp_tool_call_hook = AsyncMock() + logging_obj.async_success_handler = AsyncMock() + + user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server._handle_managed_mcp_tool", + new_callable=AsyncMock, + return_value=tool_result, + ) as mock_handle, + patch.object(global_mcp_tool_registry, "get_tool", return_value=None), + ): + result = await execute_mcp_tool( + name="test_server-some_tool", + arguments={"x": 1}, + allowed_mcp_servers=[server], + start_time=datetime.now(), + user_api_key_auth=user_auth, + litellm_logging_obj=logging_obj, + requested_server_id="server-123", + ) + + mock_handle.assert_awaited_once() + assert result is tool_result + logging_obj.async_post_mcp_tool_call_hook.assert_awaited_once() + logging_obj.async_success_handler.assert_awaited_once() + assert logging_obj.async_success_handler.await_args.kwargs["result"] is tool_result + assert logging_obj.call_type == CallTypes.call_mcp_tool.value + + def test_tool_name_matches_case_insensitive(): """Test that _tool_name_matches performs case-insensitive comparison. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 47c9396f121..cf37d1344c8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1,6 +1,6 @@ import json from typing import Any, Dict, Optional -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import httpx import pytest @@ -1069,6 +1069,101 @@ class TestCallToolRestAPI: assert captured["arguments"] == {"foo": "bar"} assert captured["allowed_mcp_servers"] == [stub_server] + async def test_fires_failure_logging_when_tool_call_raises(self, monkeypatch): + """ + Regression: the REST tool-call endpoint bypasses the @client decorator + that drives failure logging for the JSON-RPC path, so it must fire + async_failure_handler itself when execute_mcp_tool raises. Otherwise no + OTel MCP tool-call span is emitted for failed dashboard-initiated calls. + """ + from litellm.types.utils import CallTypes + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + class StubServer: + server_id = "server-1" + alias = "server-1" + server_name = "server-1" + name = "stub" + allowed_tools = None + mcp_info = {"server_name": "stub"} + available_on_public_internet = True + auth_type = None + + stub_server = StubServer() + + logging_obj = MagicMock() + logging_obj.async_failure_handler = AsyncMock() + + async def fake_common_processing(self, **kwargs): + return self.data, logging_obj + + tool_error = ValueError("upstream tool blew up") + + async def fake_execute_mcp_tool(**kwargs): + raise tool_error + + monkeypatch.setattr( + rest_endpoints, + "build_effective_auth_contexts", + fake_contexts, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_config", + {}, + raising=False, + ) + monkeypatch.setattr( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.common_processing_pre_call_logic", + fake_common_processing, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints, + "execute_mcp_tool", + fake_execute_mcp_tool, + raising=False, + ) + + request_payload = { + "server_id": "server-1", + "name": "demo-tool", + "arguments": {"foo": "bar"}, + } + request = _build_request( + path="/mcp-rest/tools/call", + method="POST", + json_body=request_payload, + ) + + with pytest.raises(HTTPException) as exc_info: + await rest_endpoints.call_tool_rest_api( + request, + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert exc_info.value.status_code == 500 + logging_obj.async_failure_handler.assert_awaited_once() + assert logging_obj.async_failure_handler.await_args.args[0] is tool_error + assert logging_obj.call_type == CallTypes.call_mcp_tool.value + class TestGetToolsForSingleServer: """Test _get_tools_for_single_server with object_permission filtering"""