tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-17 02:23:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

[MCP Protocol header] fix issue with clients protocol header (#13112)

Browse Source 
* fix headers

* fix test

* fix ruff

* fix mypy
This commit is contained in:
Jugal D. Bhatt
2025-07-30 20:42:29 +05:30
committed by GitHub
parent e324f76859
commit 5fbe732ce1
4 changed files with 117 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/experimental_mcp_client/client.py
+10 -1
View File
@@ -183,9 +183,18 @@ class MCPClient:
elif self.auth_type == MCPAuth.api_key:
headers["X-API-Key"] = self._mcp_auth_value
headers["MCP-Protocol-Version"] = self.protocol_version.value
# Handle protocol version - it might be a string or enum
if hasattr(self.protocol_version, 'value'):
# It's an enum
protocol_version_str = self.protocol_version.value
else:
# It's a string
protocol_version_str = str(self.protocol_version)
headers["MCP-Protocol-Version"] = protocol_version_str
return headers
async def list_tools(self) -> List[MCPTool]:
"""List available tools from the server."""
if not self._session:
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+52 -25
View File
@@ -67,6 +67,34 @@ def _deserialize_env_dict(env_data: Any) -> Optional[Dict[str, str]]:
return env_data
def _convert_protocol_version_to_enum(protocol_version: Optional[str | MCPSpecVersionType]) -> MCPSpecVersionType:
"""
Convert string protocol version to MCPSpecVersion enum.
Args:
protocol_version: String protocol version, enum, or None
Returns:
MCPSpecVersionType: The enum value
"""
if not protocol_version:
return MCPSpecVersion.jun_2025 # type: ignore
# If it's already an MCPSpecVersion enum, return it
if isinstance(protocol_version, MCPSpecVersion):
return protocol_version # type: ignore
# If it's a string, try to match it to enum values
if isinstance(protocol_version, str):
for version in MCPSpecVersion:
if version.value == protocol_version:
return version # type: ignore
# If no match found, return default
verbose_logger.warning(f"Unknown protocol version '{protocol_version}', using default")
return MCPSpecVersion.jun_2025 # type: ignore
class MCPServerManager:
def __init__(self):
self.registry: Dict[str, MCPServer] = {}
@@ -231,7 +259,7 @@ class MCPServerManager:
server_name=getattr(mcp_server, 'server_name', None),
url=mcp_server.url,
transport=cast(MCPTransportType, mcp_server.transport),
spec_version=cast(MCPSpecVersionType, mcp_server.spec_version),
spec_version=_convert_protocol_version_to_enum(mcp_server.spec_version),
auth_type=cast(MCPAuthType, mcp_server.auth_type),
mcp_info=MCPInfo(
server_name=mcp_server.server_name or mcp_server.server_id,
@@ -354,6 +382,9 @@ class MCPServerManager:
"""
transport = server.transport or MCPTransport.sse
# Convert protocol version string to enum
protocol_version_enum = _convert_protocol_version_to_enum(protocol_version or server.spec_version)
# Handle stdio transport
if transport == MCPTransport.stdio:
# For stdio, we need to get the stdio config from the server
@@ -372,7 +403,7 @@ class MCPServerManager:
auth_value=mcp_auth_header or server.authentication_token,
timeout=60.0,
stdio_config=stdio_config,
protocol_version=cast(MCPSpecVersionType, protocol_version or server.spec_version),
protocol_version=protocol_version_enum,
)
else:
# For HTTP/SSE transports
@@ -383,7 +414,7 @@ class MCPServerManager:
auth_type=server.auth_type,
auth_value=mcp_auth_header or server.authentication_token,
timeout=60.0,
protocol_version=cast(MCPSpecVersionType, protocol_version or server.spec_version),
protocol_version=protocol_version_enum,
)
async def _get_tools_from_server(self, server: MCPServer, mcp_auth_header: Optional[str] = None, mcp_protocol_version: Optional[str] = None) -> List[MCPTool]:
@@ -400,6 +431,9 @@ class MCPServerManager:
verbose_logger.debug(f"Connecting to url: {server.url}")
verbose_logger.info(f"_get_tools_from_server for {server.name}...")
# Use protocol version from request if provided, otherwise use server's default
protocol_version = mcp_protocol_version if mcp_protocol_version else server.spec_version
client = None
try:
# Use protocol version from request if provided, otherwise use server's default
@@ -562,39 +596,32 @@ class MCPServerManager:
protocol_version=mcp_protocol_version,
)
#########################################################
# During MCP Tool Call Hook
# Allow concurrent monitoring and validation during execution
#########################################################
if litellm_logging_obj:
during_hook_kwargs = {
"name": name,
"arguments": arguments,
"server_name": server_name_from_prefix,
}
# Start the during hook in a separate task for concurrent execution
during_hook_task = asyncio.create_task(
litellm_logging_obj.async_during_mcp_tool_call_hook(
kwargs=during_hook_kwargs,
request_obj=None, # Will be created in the hook
start_time=start_time,
end_time=start_time,
)
)
async with client:
# Use the original tool name (without prefix) for the actual call
call_tool_params = MCPCallToolRequestParams(
name=original_tool_name,
arguments=arguments,
)
# Initialize during_hook_task as None
during_hook_task = None
# Start during hook if litellm_logging_obj is available
if litellm_logging_obj:
try:
during_hook_task = litellm_logging_obj.async_during_mcp_tool_call_hook(
kwargs=litellm_logging_obj.model_call_details,
start_time=start_time,
)
except Exception as e:
verbose_logger.warning(f"During hook error (non-blocking): {str(e)}")
result = await client.call_tool(call_tool_params)
#########################################################
# Check during hook result if it completed
#########################################################
if litellm_logging_obj and 'during_hook_task' in locals():
if litellm_logging_obj and during_hook_task is not None:
try:
during_hook_result = await during_hook_task
if during_hook_result and not during_hook_result.get("should_continue", True):
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
+3 -2
View File
@@ -23,6 +23,7 @@ router = APIRouter(
if MCP_AVAILABLE:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
_convert_protocol_version_to_enum,
)
from litellm.proxy._experimental.mcp_server.server import (
ListMCPToolsRestAPIResponseObject,
@@ -174,7 +175,7 @@ if MCP_AVAILABLE:
name=request.alias or request.server_name or "",
url=request.url,
transport=request.transport,
spec_version=request.spec_version,
spec_version=_convert_protocol_version_to_enum(request.spec_version),
auth_type=request.auth_type,
mcp_info=request.mcp_info,
),
@@ -203,7 +204,7 @@ if MCP_AVAILABLE:
name=request.alias or request.server_name or "",
url=request.url,
transport=request.transport,
spec_version=request.spec_version,
spec_version=_convert_protocol_version_to_enum(request.spec_version),
auth_type=request.auth_type,
mcp_info=request.mcp_info,
),
litellm/proxy/_experimental/mcp_server/server.py
+52 -39
View File
@@ -168,25 +168,30 @@ if MCP_AVAILABLE:
"""
List all available tools
"""
# Get user authentication from context variable
user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, mcp_protocol_version = get_auth_context()
verbose_logger.debug(
f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}"
)
verbose_logger.debug(
f"MCP list_tools - MCP servers from context: {mcp_servers}"
)
verbose_logger.debug(
f"MCP list_tools - MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}"
)
# Get mcp_servers from context variable
return await _list_mcp_tools(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
mcp_protocol_version=mcp_protocol_version,
)
try:
# Get user authentication from context variable
user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, mcp_protocol_version = get_auth_context()
verbose_logger.debug(
f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}"
)
verbose_logger.debug(
f"MCP list_tools - MCP servers from context: {mcp_servers}"
)
verbose_logger.debug(
f"MCP list_tools - MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}"
)
# Get mcp_servers from context variable
return await _list_mcp_tools(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
mcp_protocol_version=mcp_protocol_version,
)
except Exception as e:
verbose_logger.exception(f"Error in list_tools endpoint: {str(e)}")
# Return empty list instead of failing completely
return []
@server.call_tool()
async def mcp_server_tool_call(
@@ -362,29 +367,37 @@ if MCP_AVAILABLE:
"""
if not MCP_AVAILABLE:
return []
# Get tools from managed MCP servers
managed_tools = await _get_tools_from_mcp_servers(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
mcp_protocol_version=mcp_protocol_version,
)
# Get tools from managed MCP servers with error handling
managed_tools = []
try:
managed_tools = await _get_tools_from_mcp_servers(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
mcp_protocol_version=mcp_protocol_version,
)
except Exception as e:
verbose_logger.exception(f"Error getting tools from managed MCP servers: {str(e)}")
# Continue with empty managed tools list instead of failing completely
# Get tools from local registry
local_tools_raw = global_mcp_tool_registry.list_tools()
# Convert local tools to MCPTool format
local_tools = []
for tool in local_tools_raw:
# Convert from litellm.types.mcp_server.tool_registry.MCPTool to mcp.types.Tool
mcp_tool = MCPTool(
name=tool.name,
description=tool.description,
inputSchema=tool.input_schema
)
local_tools.append(mcp_tool)
try:
local_tools_raw = global_mcp_tool_registry.list_tools()
# Convert local tools to MCPTool format
for tool in local_tools_raw:
# Convert from litellm.types.mcp_server.tool_registry.MCPTool to mcp.types.Tool
mcp_tool = MCPTool(
name=tool.name,
description=tool.description,
inputSchema=tool.input_schema
)
local_tools.append(mcp_tool)
except Exception as e:
verbose_logger.exception(f"Error getting tools from local registry: {str(e)}")
# Continue with empty local tools list instead of failing completely
# Combine all tools
all_tools = managed_tools + local_tools