scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Feat/mcp preserve tool metadata calltoolresult (#17561)

Browse source 
* feat(mcp): preserve tool metadata and full CallToolResult in MCP gateway

This PR fixes two issues that prevented ChatGPT from rendering MCP UI widgets
when proxied through LiteLLM:

1. Preserve Tool Metadata in tools/list
   - Modified _create_prefixed_tools() to mutate tools in place instead of
     reconstructing them, preserving all fields including metadata/_meta
   - This ensures ChatGPT can see 'openai/outputTemplate' URIs in tools/list
     and will call resources/read to fetch widgets

2. Preserve Full CallToolResult (structuredContent + metadata)
   - Changed call_mcp_tool() and _handle_managed_mcp_tool() to return full
     CallToolResult objects instead of just content
   - Updated error handlers to return CallToolResult with isError flag
   - Wrapped local tool results in CallToolResult objects
   - This preserves structuredContent and metadata fields needed for widget rendering

Files changed:
- litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
- litellm/proxy/_experimental/mcp_server/server.py

Fixes issues where ChatGPT could not render MCP UI widgets when using
LiteLLM as an MCP gateway.

* feat(mcp): Preserve tool metadata and return full CallToolResult for ChatGPT UI widgets

- Preserve metadata and _meta fields when creating prefixed tools
- Return full CallToolResult instead of just content list
- Ensures ChatGPT can discover and render UI widgets via openai/outputTemplate
- Fixes metadata stripping that prevented widget rendering in ChatGPT

Changes:
- mcp_server_manager.py: Mutate tools in place to preserve all fields including metadata
- server.py: Return CallToolResult with structuredContent and metadata preserved
- Added test to verify metadata preservation

* fix: guard cost calculator when BaseModel lacks _hidden_params

---------

Co-authored-by: Afroz Ahmad <aahmad@Afrozs-MacBook-Pro.local>
Co-authored-by: Afroz Ahmad <aahmad@KNDMCPTMZH3.sephoraus.com>
This commit is contained in:
YutaSaito 2025-12-06 09:15:22 +09:00 • committed by GitHub
parent 342723eb12
commit b5133c4c7d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 177 additions and 60 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

36
litellm/cost_calculator.py
View file

@ -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
)

13
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -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(

91
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -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.

10
tests/mcp_tests/test_mcp_logging.py
View file

@ -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

87
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_metadata_preservation.py Normal file
View file

@ -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__])