diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 76a28344856..e32ae85f9ad 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -6,6 +6,7 @@ LiteLLM MCP Server Routes import asyncio import contextlib from datetime import datetime +import traceback from typing import Any, AsyncIterator, Dict, List, Optional, Tuple, Union, cast from fastapi import FastAPI, HTTPException @@ -13,6 +14,7 @@ from pydantic import AnyUrl, ConfigDict from starlette.types import Receive, Scope, Send from litellm._logging import verbose_logger +from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, @@ -25,7 +27,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer -from litellm.types.utils import StandardLoggingMCPToolCall +from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall from litellm.utils import client # Check if MCP is available @@ -1320,33 +1322,6 @@ if MCP_AVAILABLE: content=cast(Any, local_content), isError=False ) - ######################################################### - # Post MCP Tool Call Hook - # Allow modifying the MCP tool call response before it is returned to the user - ######################################################### - 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, - ) - # Set call_type to call_mcp_tool so cost calculator recognizes it - from litellm.types.utils import CallTypes - - litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value - # Trigger success logging to build standard_logging_object and call callbacks - # async_success_handler will: - # 1. Call _success_handler_helper_fn which recognizes call_mcp_tool - # 2. Call _process_hidden_params_and_response_cost which: - # - Calculates cost via _response_cost_calculator -> MCPCostCalculator - # - Builds standard_logging_object - # 3. Call async_log_success_event on all callbacks - await litellm_logging_obj.async_success_handler( - result=response, start_time=start_time, end_time=end_time - ) return response @client @@ -1365,49 +1340,84 @@ if MCP_AVAILABLE: Call a specific tool with the provided arguments (handles prefixed tool names). """ start_time = datetime.now() - if arguments is None: - raise HTTPException( - status_code=400, detail="Request arguments are required" + litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get( + "litellm_logging_obj", None + ) + + try: + if arguments is None: + raise HTTPException( + status_code=400, detail="Request arguments are required" + ) + + ## CHECK IF USER IS ALLOWED TO CALL THIS TOOL + allowed_mcp_server_ids = ( + await global_mcp_server_manager.get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + ) ) - ## CHECK IF USER IS ALLOWED TO CALL THIS TOOL - allowed_mcp_server_ids = ( - await global_mcp_server_manager.get_allowed_mcp_servers( + allowed_mcp_servers: List[MCPServer] = [] + for allowed_mcp_server_id in allowed_mcp_server_ids: + allowed_server = global_mcp_server_manager.get_mcp_server_by_id( + allowed_mcp_server_id + ) + if allowed_server is not None: + allowed_mcp_servers.append(allowed_server) + + allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names( + mcp_servers=mcp_servers, + allowed_mcp_servers=allowed_mcp_servers, + ) + if not allowed_mcp_servers: + raise HTTPException( + status_code=403, + detail="User not allowed to call this tool.", + ) + + # Delegate to execute_mcp_tool for execution + response = await execute_mcp_tool( + name=name, + arguments=arguments, + allowed_mcp_servers=allowed_mcp_servers, + start_time=start_time, user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + **kwargs, ) - ) - - allowed_mcp_servers: List[MCPServer] = [] - for allowed_mcp_server_id in allowed_mcp_server_ids: - allowed_server = global_mcp_server_manager.get_mcp_server_by_id( - allowed_mcp_server_id + except Exception as e: + traceback_str = traceback.format_exc( + limit=MAXIMUM_TRACEBACK_LINES_TO_LOG ) - if allowed_server is not None: - allowed_mcp_servers.append(allowed_server) + from litellm.proxy.proxy_server import proxy_logging_obj + + if proxy_logging_obj and user_api_key_auth: + await proxy_logging_obj.post_call_failure_hook( + request_data=kwargs, + original_exception=e, + user_api_key_dict=user_api_key_auth, + route="/mcp/call_tool", + traceback_str=traceback_str, + ) + raise - allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names( - mcp_servers=mcp_servers, - allowed_mcp_servers=allowed_mcp_servers, - ) - if not allowed_mcp_servers: - raise HTTPException( - status_code=403, - detail="User not allowed to call this tool.", + 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, ) - - # Delegate to execute_mcp_tool for execution - return await execute_mcp_tool( - name=name, - arguments=arguments, - allowed_mcp_servers=allowed_mcp_servers, - start_time=start_time, - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - **kwargs, - ) + 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( name: str,