diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 57eab6d29a3..4407134fa92 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -860,9 +860,9 @@ def completion_cost( # noqa: PLR0915 or isinstance(completion_response, dict) ): # tts returns a custom class if isinstance(completion_response, dict): - usage_obj: Optional[Union[dict, Usage]] = ( - completion_response.get("usage", {}) - ) + usage_obj: Optional[ + Union[dict, Usage] + ] = completion_response.get("usage", {}) else: usage_obj = getattr(completion_response, "usage", {}) if isinstance(usage_obj, BaseModel) and not _is_known_usage_objects( @@ -1066,13 +1066,14 @@ def completion_cost( # noqa: PLR0915 # If model is like "tavily-search", construct "tavily/search" for cost lookup search_model = f"{custom_llm_provider}/search" - prompt_cost, completion_cost_result = ( - search_provider_cost_per_query( - model=search_model, - custom_llm_provider=custom_llm_provider, - number_of_queries=number_of_queries, - optional_params=optional_params, - ) + ( + prompt_cost, + completion_cost_result, + ) = search_provider_cost_per_query( + model=search_model, + custom_llm_provider=custom_llm_provider, + number_of_queries=number_of_queries, + optional_params=optional_params, ) # Return the total cost (prompt_cost + completion_cost, but for search it's just prompt_cost) @@ -1080,11 +1081,13 @@ def completion_cost( # noqa: PLR0915 # Apply discount original_cost = _final_cost - _final_cost, discount_percent, discount_amount = ( - _apply_cost_discount( - base_cost=_final_cost, - custom_llm_provider=custom_llm_provider, - ) + ( + _final_cost, + discount_percent, + discount_amount, + ) = _apply_cost_discount( + base_cost=_final_cost, + custom_llm_provider=custom_llm_provider, ) # Store cost breakdown in logging object if available @@ -1329,9 +1332,8 @@ def response_cost_calculator( response_cost = 0.0 else: if isinstance(response_object, BaseModel): - response_object._hidden_params["optional_params"] = optional_params - if hasattr(response_object, "_hidden_params"): + response_object._hidden_params["optional_params"] = optional_params provider_response_cost = get_response_cost_from_hidden_params( response_object._hidden_params ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d42fd80cbd7..92d863230c3 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1276,15 +1276,14 @@ class MCPServerManager: name_to_use = prefixed_name if add_prefix else tool.name - tool_obj = MCPTool( - name=name_to_use, - description=tool.description, - inputSchema=tool.inputSchema, - ) - prefixed_tools.append(tool_obj) + # Preserve all tool fields including metadata/_meta by mutating the original tool + # Similar to how _create_prefixed_prompts works + original_name = tool.name + tool.name = name_to_use + prefixed_tools.append(tool) # Update tool to server mapping for resolution (support both forms) - self.tool_name_to_mcp_server_name_mapping[tool.name] = prefix + self.tool_name_to_mcp_server_name_mapping[original_name] = prefix self.tool_name_to_mcp_server_name_mapping[prefixed_name] = prefix verbose_logger.info( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 412a6de0059..edf53e99573 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -6,7 +6,7 @@ LiteLLM MCP Server Routes import asyncio import contextlib from datetime import datetime -from typing import Any, AsyncIterator, Dict, List, Optional, Tuple, Union +from typing import Any, AsyncIterator, Dict, List, Optional, Tuple, Union, cast from fastapi import FastAPI, HTTPException from pydantic import AnyUrl, ConfigDict @@ -72,7 +72,13 @@ if MCP_AVAILABLE: auth_context_var, ) from mcp.server.streamable_http_manager import StreamableHTTPSessionManager - from mcp.types import EmbeddedResource, ImageContent, Prompt, TextContent + from mcp.types import ( + CallToolResult, + EmbeddedResource, + ImageContent, + Prompt, + TextContent, + ) from mcp.types import Tool as MCPTool from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( @@ -234,7 +240,7 @@ if MCP_AVAILABLE: @server.call_tool() async def mcp_server_tool_call( name: str, arguments: Dict[str, Any] | None - ) -> List[Union[TextContent, ImageContent, EmbeddedResource]]: + ) -> CallToolResult: """ Call a specific tool with the provided arguments @@ -300,26 +306,37 @@ if MCP_AVAILABLE: ) except BlockedPiiEntityError as e: verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {str(e)}") - # Return error as text content for MCP protocol - return [ - TextContent( - text=f"Error: Blocked PII entity detected - {str(e)}", type="text" - ) - ] + return CallToolResult( + content=[ + TextContent( + text=f"Error: Blocked PII entity detected - {str(e)}", + type="text", + ) + ], + isError=True, + ) except GuardrailRaisedException as e: verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {str(e)}") - # Return error as text content for MCP protocol - return [ - TextContent(text=f"Error: Guardrail violation - {str(e)}", type="text") - ] + return CallToolResult( + content=[ + TextContent( + text=f"Error: Guardrail violation - {str(e)}", type="text" + ) + ], + isError=True, + ) except HTTPException as e: verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}") - # Return error as text content for MCP protocol - return [TextContent(text=f"Error: {str(e.detail)}", type="text")] + return CallToolResult( + content=[TextContent(text=f"Error: {str(e.detail)}", type="text")], + isError=True, + ) except Exception as e: verbose_logger.exception(f"MCP mcp_server_tool_call - error: {e}") - # Return error as text content for MCP protocol - return [TextContent(text=f"Error: {str(e)}", type="text")] + return CallToolResult( + content=[TextContent(text=f"Error: {str(e)}", type="text")], + isError=True, + ) return response @@ -1173,7 +1190,7 @@ if MCP_AVAILABLE: oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, **kwargs: Any, - ) -> List[Union[TextContent, ImageContent, EmbeddedResource]]: + ) -> CallToolResult: """ Call a specific tool with the provided arguments (handles prefixed tool names) """ @@ -1237,9 +1254,9 @@ if MCP_AVAILABLE: "litellm_logging_obj", None ) if litellm_logging_obj: - litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( - standard_logging_mcp_tool_call - ) + litellm_logging_obj.model_call_details[ + "mcp_tool_call_metadata" + ] = standard_logging_mcp_tool_call litellm_logging_obj.model = f"MCP: {name}" # Check if tool exists in local registry first (for OpenAPI-based tools) # These tools are registered with their prefixed names @@ -1247,7 +1264,8 @@ if MCP_AVAILABLE: local_tool = global_mcp_tool_registry.get_tool(name) if local_tool: verbose_logger.debug(f"Executing local registry tool: {name}") - response = await _handle_local_mcp_tool(name, arguments) + local_content = await _handle_local_mcp_tool(name, arguments) + response = CallToolResult(content=cast(Any, local_content), isError=False) # Try managed MCP server tool (pass the full prefixed name) # Primary and recommended way to use external MCP servers @@ -1279,7 +1297,12 @@ if MCP_AVAILABLE: # Deprecated: Local MCP Server Tool ######################################################### else: - response = await _handle_local_mcp_tool(original_tool_name, arguments) + local_content = await _handle_local_mcp_tool( + original_tool_name, arguments + ) + response = CallToolResult( + content=cast(Any, local_content), isError=False + ) ######################################################### # Post MCP Tool Call Hook @@ -1432,7 +1455,7 @@ if MCP_AVAILABLE: oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, litellm_logging_obj: Optional[Any] = None, - ) -> List[Union[TextContent, ImageContent, EmbeddedResource]]: + ) -> CallToolResult: """Handle tool execution for managed server tools""" # Import here to avoid circular import from litellm.proxy.proxy_server import proxy_logging_obj @@ -1449,7 +1472,7 @@ if MCP_AVAILABLE: proxy_logging_obj=proxy_logging_obj, ) verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) - return call_tool_result.content # type: ignore[return-value] + return call_tool_result async def _handle_local_mcp_tool( name: str, arguments: Dict[str, Any] @@ -1741,14 +1764,16 @@ if MCP_AVAILABLE: ) auth_context_var.set(auth_user) - def get_auth_context() -> Tuple[ - Optional[UserAPIKeyAuth], - Optional[str], - Optional[List[str]], - Optional[Dict[str, Dict[str, str]]], - Optional[Dict[str, str]], - Optional[Dict[str, str]], - ]: + def get_auth_context() -> ( + Tuple[ + Optional[UserAPIKeyAuth], + Optional[str], + Optional[List[str]], + Optional[Dict[str, Dict[str, str]]], + Optional[Dict[str, str]], + Optional[Dict[str, str]], + ] + ): """ Get the UserAPIKeyAuth from the auth context variable. diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py index 44acd60ad49..d27626dd447 100644 --- a/tests/mcp_tests/test_mcp_logging.py +++ b/tests/mcp_tests/test_mcp_logging.py @@ -109,7 +109,11 @@ async def test_mcp_cost_tracking(): # Add assertions assert response is not None - response_list = list(response) # Convert iterable to list + # Handle CallToolResult - access .content for the list of content items + if isinstance(response, CallToolResult): + response_list = response.content + else: + response_list = list(response) # Convert iterable to list for backward compatibility assert len(response_list) == 1 assert isinstance(response_list[0], TextContent) assert response_list[0].text == "Test response" @@ -238,8 +242,8 @@ async def test_mcp_cost_tracking_per_tool(): assert response1 is not None assert response2 is not None - response_list_1 = list(response1) - response_list_2 = list(response2) + response_list_1 = list(response1.content) + response_list_2 = list(response2.content) assert len(response_list_1) == 1 assert len(response_list_2) == 1 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py new file mode 100644 index 00000000000..6311b6d74b2 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py @@ -0,0 +1,87 @@ +""" +Tests for MCP metadata preservation. + +This module tests that tool metadata is preserved when creating prefixed tools, +which is critical for ChatGPT UI widget rendering. +""" + +import sys + +import pytest + +# Add the parent directory to the path so we can import litellm +sys.path.insert(0, "../../../../../") + +from mcp.types import Tool as MCPTool + +from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager +from litellm.proxy._types import MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +class TestMCPMetadataPreservation: + """Test that metadata is preserved when creating prefixed tools""" + + def test_create_prefixed_tools_preserves_metadata(self): + """Test that _create_prefixed_tools preserves metadata and _meta fields""" + manager = MCPServerManager() + + # Create a mock server + mock_server = MCPServer( + server_id="test-server-1", + name="test_server", + alias="test", + server_name="Test Server", + url="https://test-server.com/mcp", + transport=MCPTransport.http, + ) + + # Create a tool with metadata + tool_with_metadata = MCPTool( + name="hello_widget", + description="Display a greeting widget", + inputSchema={"type": "object", "properties": {}}, + ) + # Add metadata using setattr since MCPTool might not have it in the constructor + tool_with_metadata.metadata = { + "openai/outputTemplate": "ui://widget/hello.html", + "openai/widgetDescription": "A greeting widget", + } + tool_with_metadata._meta = { + "openai/toolInvocation/invoking": "Preparing greeting...", + } + + # Create prefixed tools + prefixed_tools = manager._create_prefixed_tools( + [tool_with_metadata], mock_server, add_prefix=True + ) + + # Verify + assert len(prefixed_tools) == 1 + prefixed_tool = prefixed_tools[0] + + # Check that name is prefixed + assert prefixed_tool.name == "test-hello_widget" + + # Check that metadata is preserved + assert hasattr(prefixed_tool, "metadata") + assert prefixed_tool.metadata == { + "openai/outputTemplate": "ui://widget/hello.html", + "openai/widgetDescription": "A greeting widget", + } + + # Check that _meta is preserved + assert hasattr(prefixed_tool, "_meta") + assert prefixed_tool._meta == { + "openai/toolInvocation/invoking": "Preparing greeting...", + } + + # Check that other fields are preserved + assert prefixed_tool.description == "Display a greeting widget" + assert prefixed_tool.inputSchema == {"type": "object", "properties": {}} + + + +if __name__ == "__main__": + pytest.main([__file__]) +