diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 8f524690be1..7105f1ae6fb 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -181,6 +181,7 @@ async def aresponses_api_with_mcp( ) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform( user_api_key_auth=user_api_key_auth, mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, + litellm_trace_id=kwargs.get("litellm_trace_id"), ) openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai( original_mcp_tools diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 297ccf4355e..805a1958552 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -6,10 +6,10 @@ from typing import ( Dict, Iterable, List, + Literal, Optional, Tuple, Union, - Literal, ) from litellm._logging import verbose_logger @@ -29,6 +29,7 @@ from litellm.utils import Rules, function_setup if TYPE_CHECKING: from mcp.types import Tool as MCPTool + from litellm.proxy.utils import ProxyLogging else: MCPTool = Any @@ -97,6 +98,7 @@ class LiteLLM_Proxy_MCP_Handler: async def _get_mcp_tools_from_manager( user_api_key_auth: Any, mcp_tools_with_litellm_proxy: Optional[Iterable[ToolParam]], + litellm_trace_id: Optional[str] = None, ) -> tuple[List[MCPTool], List[str]]: """ Get available tools from the MCP server manager. @@ -109,13 +111,13 @@ class LiteLLM_Proxy_MCP_Handler: List of MCP tools List names of allowed MCP servers """ - from litellm.proxy._experimental.mcp_server.server import ( - _get_tools_from_mcp_servers, - _get_allowed_mcp_servers_from_mcp_server_names, - ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + from litellm.proxy._experimental.mcp_server.server import ( + _get_allowed_mcp_servers_from_mcp_server_names, + _get_tools_from_mcp_servers, + ) mcp_servers: List[str] = [] if mcp_tools_with_litellm_proxy: @@ -136,6 +138,7 @@ class LiteLLM_Proxy_MCP_Handler: mcp_server_auth_headers=None, log_list_tools_to_spendlogs=True, list_tools_log_source="responses", + litellm_trace_id=litellm_trace_id, ) allowed_mcp_server_ids = ( await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) @@ -239,7 +242,9 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod async def _process_mcp_tools_to_openai_format( - user_api_key_auth: Any, mcp_tools_with_litellm_proxy: List[ToolParam] + user_api_key_auth: Any, + mcp_tools_with_litellm_proxy: List[ToolParam], + litellm_trace_id: Optional[str] = None, ) -> tuple[List[Any], dict[str, str]]: """ Centralized method to process MCP tools through the complete pipeline. @@ -247,6 +252,7 @@ class LiteLLM_Proxy_MCP_Handler: Args: user_api_key_auth: User authentication info for access control mcp_tools_with_litellm_proxy: ToolParam objects with server_url starting with "litellm_proxy" + litellm_trace_id: Optional trace ID for linking list_mcp_tools spend logs to parent request Returns: List of tools in OpenAI format ready to be sent to the LLM @@ -258,6 +264,7 @@ class LiteLLM_Proxy_MCP_Handler: ) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform( user_api_key_auth, mcp_tools_with_litellm_proxy, + litellm_trace_id=litellm_trace_id, ) openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai( @@ -268,7 +275,9 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod async def _process_mcp_tools_without_openai_transform( - user_api_key_auth: Any, mcp_tools_with_litellm_proxy: List[ToolParam] + user_api_key_auth: Any, + mcp_tools_with_litellm_proxy: List[ToolParam], + litellm_trace_id: Optional[str] = None, ) -> tuple[List[Any], dict[str, str]]: """ Process MCP tools through filtering and deduplication pipeline without OpenAI transformation. @@ -291,6 +300,7 @@ class LiteLLM_Proxy_MCP_Handler: ) = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( user_api_key_auth=user_api_key_auth, mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, + litellm_trace_id=litellm_trace_id, ) # Step 2: Filter tools based on allowed_tools parameter @@ -495,14 +505,13 @@ class LiteLLM_Proxy_MCP_Handler: """Execute tool calls and return results.""" from fastapi import HTTPException + from litellm._uuid import uuid from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) from litellm.proxy.proxy_server import proxy_logging_obj - from litellm._uuid import uuid - tool_results = [] tool_call_id: Optional[str] = None rules_obj = Rules() @@ -1025,7 +1034,6 @@ class LiteLLM_Proxy_MCP_Handler: List of MCP tool execution events for streaming """ from litellm._uuid import uuid - from litellm.responses.mcp.mcp_streaming_iterator import create_mcp_call_events tool_execution_events: List[Any] = [] @@ -1108,8 +1116,8 @@ class LiteLLM_Proxy_MCP_Handler: """Add custom output elements to the final response for MCP tool execution.""" # Import the required classes for creating output items import json - from litellm._uuid import uuid + from litellm._uuid import uuid from litellm.types.responses.main import GenericResponseOutputItem, OutputText # Create output element for initial MCP tools