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

refactor: drop MCPClient.connect and use run_with_session lifecycle (#16696)

Browse source 
Surface detailed connection errors by handling HTTP failures
This commit is contained in:
YutaSaito 2025-11-16 10:54:27 +09:00 • committed by GitHub
parent 6b9fbc5a36
commit 0b586d26fc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 106 additions and 280 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

219
litellm/experimental_mcp_client/client.py
View file

@ -5,7 +5,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
import asyncio
import base64
from datetime import timedelta
from typing import Callable, Dict, List, Optional, Union
from typing import Awaitable, Callable, Dict, List, Optional, TypeVar, Union
import httpx
from mcp import ClientSession, StdioServerParameters
@ -34,6 +34,9 @@ def to_basic_auth(auth_value: str) -> str:
return base64.b64encode(auth_value.encode("utf-8")).decode()
TSessionResult = TypeVar("TSessionResult")
class MCPClient:
"""
MCP Client supporting:
@ -58,12 +61,6 @@ class MCPClient:
self.auth_type: MCPAuthType = auth_type
self.timeout: float = timeout
self._mcp_auth_value: Optional[Union[str, Dict[str, str]]] = None
self._session: Optional[ClientSession] = None
self._context = None
self._transport_ctx = None
self._transport = None
self._session_ctx = None
self._task: Optional[asyncio.Task] = None
self.stdio_config: Optional[MCPStdioConfig] = stdio_config
self.extra_headers: Optional[Dict[str, str]] = extra_headers
self.ssl_verify: Optional[VerifyTypes] = ssl_verify
@ -71,33 +68,14 @@ class MCPClient:
if auth_value:
self.update_auth_value(auth_value)
async def __aenter__(self):
"""
Enable async context manager support.
Initializes the transport and session.
"""
try:
await self.connect()
return self
except Exception:
await self.disconnect()
raise
async def connect(self):
"""Initialize the transport and session."""
if self._session:
verbose_logger.debug(
f"MCP client already connected to {self.server_url or 'stdio'}"
)
return # Already connected
verbose_logger.info(
f"MCP client connecting to {self.server_url or 'stdio'} via {self.transport_type}"
)
async def run_with_session(
self, operation: Callable[[ClientSession], Awaitable[TSessionResult]]
) -> TSessionResult:
"""Open a session, run the provided coroutine, and clean up."""
transport_ctx = None
try:
if self.transport_type == MCPTransport.stdio:
# For stdio transport, use stdio_client with command-line parameters
if not self.stdio_config:
raise ValueError("stdio_config is required for stdio transport")
@ -106,117 +84,43 @@ class MCPClient:
args=self.stdio_config.get("args", []),
env=self.stdio_config.get("env", {}),
)
self._transport_ctx = stdio_client(server_params)
self._transport = await self._transport_ctx.__aenter__()
self._session_ctx = ClientSession(
self._transport[0], self._transport[1]
)
self._session = await self._session_ctx.__aenter__()
await self._session.initialize()
verbose_logger.info(
f"MCP client successfully connected via stdio: {self.stdio_config.get('command', '')}"
)
transport_ctx = stdio_client(server_params)
elif self.transport_type == MCPTransport.sse:
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
self._transport_ctx = sse_client(
transport_ctx = sse_client(
url=self.server_url,
timeout=self.timeout,
headers=headers,
httpx_client_factory=httpx_client_factory,
)
self._transport = await self._transport_ctx.__aenter__()
self._session_ctx = ClientSession(
self._transport[0], self._transport[1]
)
self._session = await self._session_ctx.__aenter__()
await self._session.initialize()
verbose_logger.info(
f"MCP client successfully connected via SSE to {self.server_url}"
)
else: # http
else:
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
verbose_logger.debug(
"litellm headers for streamablehttp_client: %s", headers
)
self._transport_ctx = streamablehttp_client(
transport_ctx = streamablehttp_client(
url=self.server_url,
timeout=timedelta(seconds=self.timeout),
headers=headers,
httpx_client_factory=httpx_client_factory,
)
self._transport = await self._transport_ctx.__aenter__()
self._session_ctx = ClientSession(
self._transport[0], self._transport[1]
)
self._session = await self._session_ctx.__aenter__()
await self._session.initialize()
verbose_logger.info(
f"MCP client successfully connected via HTTP to {self.server_url}"
)
except ValueError as e:
# Re-raise ValueError exceptions (like missing stdio_config)
verbose_logger.warning(f"MCP client connection failed: {str(e)}")
await self.disconnect()
if transport_ctx is None:
raise RuntimeError("Failed to create transport context")
async with transport_ctx as transport:
read_stream, write_stream = transport[0], transport[1]
session_ctx = ClientSession(read_stream, write_stream)
async with session_ctx as session:
await session.initialize()
return await operation(session)
except Exception:
verbose_logger.warning(
"MCP client run_with_session failed for %s", self.server_url or "stdio"
)
raise
except Exception as e:
verbose_logger.warning(f"MCP client connection failed: {str(e)}")
await self.disconnect()
# Don't raise other exceptions, let the calling code handle it gracefully
# This allows the server manager to continue with other servers
# Instead of raising, we'll let the calling code handle the failure
pass
async def __aexit__(self, exc_type, exc_val, exc_tb):
"""Cleanup when exiting context manager."""
await self.disconnect()
async def disconnect(self):
"""Clean up session and connections."""
verbose_logger.info(
f"MCP client disconnecting from {self.server_url or 'stdio'}"
)
if self._task and not self._task.done():
verbose_logger.debug("MCP client cancelling background task")
self._task.cancel()
try:
await self._task
except asyncio.CancelledError:
pass
if self._session:
try:
verbose_logger.debug("MCP client closing session")
await self._session_ctx.__aexit__(None, None, None) # type: ignore
except Exception as e:
verbose_logger.debug(
f"Error closing MCP session: {type(e).__name__}: {str(e)}"
)
pass
self._session = None
self._session_ctx = None
if self._transport_ctx:
try:
verbose_logger.debug("MCP client closing transport")
await self._transport_ctx.__aexit__(None, None, None)
except Exception as e:
verbose_logger.debug(
f"Error closing MCP transport: {type(e).__name__}: {str(e)}"
)
pass
self._transport_ctx = None
self._transport = None
if self._context:
try:
await self._context.__aexit__(None, None, None) # type: ignore
except Exception:
pass
self._context = None
def update_auth_value(self, mcp_auth_value: Union[str, Dict[str, str]]):
"""
@ -294,24 +198,11 @@ class MCPClient:
f"MCP client listing tools from {self.server_url or 'stdio'}"
)
if not self._session:
verbose_logger.debug("MCP client session not found, attempting to connect")
try:
await self.connect()
except Exception as e:
verbose_logger.error(
f"MCP client connection failed during list_tools: {type(e).__name__}: {str(e)}"
)
return []
if self._session is None:
verbose_logger.error(
"MCP client session is not initialized after connection attempt"
)
return []
async def _list_tools_operation(session: ClientSession):
return await session.list_tools()
try:
result = await self._session.list_tools()
result = await self.run_with_session(_list_tools_operation)
tool_count = len(result.tools)
tool_names = [tool.name for tool in result.tools]
verbose_logger.info(
@ -320,7 +211,6 @@ class MCPClient:
return result.tools
except asyncio.CancelledError:
verbose_logger.warning("MCP client list_tools was cancelled")
await self.disconnect()
raise
except Exception as e:
error_type = type(e).__name__
@ -339,7 +229,6 @@ class MCPClient:
"the MCP server may have crashed, disconnected, or timed out"
)
await self.disconnect()
# Return empty list instead of raising to allow graceful degradation
return []
@ -353,55 +242,21 @@ class MCPClient:
f"MCP client calling tool '{call_tool_request_params.name}' with arguments: {call_tool_request_params.arguments}"
)
if not self._session:
verbose_logger.warning(
"MCP client session not found, attempting to connect"
)
try:
await self.connect()
except Exception as e:
verbose_logger.error(
f"MCP client connection failed before tool call: {type(e).__name__}: {str(e)}"
)
return MCPCallToolResult(
content=[TextContent(type="text", text=f"{str(e)}")], isError=True
)
if self._session is None:
verbose_logger.error(
"MCP client session is not initialized after connection attempt"
)
return MCPCallToolResult(
content=[
TextContent(
type="text", text="MCP client session is not initialized"
)
],
isError=True,
)
# Check session and transport state before calling tool
verbose_logger.debug(
f"MCP client state before tool call - "
f"session: {'active' if self._session else 'none'}, "
f"transport: {'active' if self._transport else 'none'}, "
f"session_ctx: {'active' if self._session_ctx else 'none'}, "
f"transport_ctx: {'active' if self._transport_ctx else 'none'}"
)
try:
async def _call_tool_operation(session: ClientSession):
verbose_logger.debug("MCP client sending tool call to session")
tool_result = await self._session.call_tool(
return await session.call_tool(
name=call_tool_request_params.name,
arguments=call_tool_request_params.arguments,
)
try:
tool_result = await self.run_with_session(_call_tool_operation)
verbose_logger.info(
f"MCP client tool call '{call_tool_request_params.name}' completed successfully"
)
return tool_result
except asyncio.CancelledError:
verbose_logger.warning("MCP client tool call was cancelled")
await self.disconnect()
raise
except Exception as e:
import traceback
@ -424,11 +279,9 @@ class MCPClient:
if "BrokenResourceError" in error_type or "Broken" in error_type:
verbose_logger.error(
"MCP client detected broken connection/stream - "
"the MCP server may have crashed, disconnected, or timed out. "
"Session and transport will be disconnected."
"the MCP server may have crashed, disconnected, or timed out."
)
await self.disconnect()
# Return a default error result instead of raising
return MCPCallToolResult(
content=[

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

@ -385,12 +385,12 @@ class MCPServerManager:
)
# Update tool name to server name mapping (for both prefixed and base names)
self.tool_name_to_mcp_server_name_mapping[base_tool_name] = (
server_prefix
)
self.tool_name_to_mcp_server_name_mapping[prefixed_tool_name] = (
server_prefix
)
self.tool_name_to_mcp_server_name_mapping[
base_tool_name
] = server_prefix
self.tool_name_to_mcp_server_name_mapping[
prefixed_tool_name
] = server_prefix
registered_count += 1
verbose_logger.debug(
@ -714,12 +714,6 @@ class MCPServerManager:
f"Failed to get tools from server {server.name}: {str(e)}"
)
return []
finally:
if client:
try:
await client.disconnect()
except Exception:
pass
async def _descovery_metadata(
self,
@ -983,8 +977,6 @@ class MCPServerManager:
async def _list_tools_task():
try:
await client.connect()
tools = await client.list_tools()
verbose_logger.debug(f"Tools from {server_name}: {tools}")
return tools
@ -1439,14 +1431,12 @@ class MCPServerManager:
)
async def _call_tool_via_client(client, params):
async with client:
return await client.call_tool(params)
return await client.call_tool(params)
tasks.append(
asyncio.create_task(_call_tool_via_client(client, call_tool_params))
)
# IMPORTANT: Must await tasks INSIDE the context manager to keep connection alive
try:
mcp_responses = await asyncio.gather(*tasks)
except (

45
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -76,12 +76,12 @@ if MCP_AVAILABLE:
mcp_auth_header=server_auth_header,
add_prefix=False,
)
# Filter tools based on allowed_tools configuration
# Only filter if allowed_tools is explicitly configured (not None and not empty)
if server.allowed_tools is not None and len(server.allowed_tools) > 0:
tools = filter_tools_by_allowed_tools(tools, server)
return _create_tool_response_objects(tools, server.mcp_info)
########################################################
@ -212,7 +212,9 @@ if MCP_AVAILABLE:
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.proxy.proxy_server import add_litellm_data_to_request, proxy_config
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
try:
data = await request.json()
@ -222,21 +224,27 @@ if MCP_AVAILABLE:
user_api_key_dict=user_api_key_dict,
proxy_config=proxy_config,
)
# FIX: Extract MCP auth headers from request
# The UI sends bearer token in x-mcp-auth header and server-specific headers,
# but they weren't being extracted and passed to call_mcp_tool.
# This fix ensures auth headers are properly extracted from the HTTP request
# and passed through to the MCP server for authentication.
mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(request.headers)
mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(request.headers)
mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(
request.headers
)
mcp_server_auth_headers = (
MCPRequestHandler._get_mcp_server_auth_headers_from_headers(
request.headers
)
)
# Add extracted headers to data dict to pass to call_mcp_tool
if mcp_auth_header:
data["mcp_auth_header"] = mcp_auth_header
if mcp_server_auth_headers:
data["mcp_server_auth_headers"] = mcp_server_auth_headers
result = await call_mcp_tool(**data)
return result
except BlockedPiiEntityError as e:
@ -296,7 +304,6 @@ if MCP_AVAILABLE:
Returns:
Operation result or error response
"""
client = None
try:
client = global_mcp_server_manager._create_mcp_client(
server=MCPServer(
@ -315,13 +322,6 @@ if MCP_AVAILABLE:
except Exception as e:
verbose_logger.error(f"Error in MCP operation: {e}", exc_info=True)
return {"status": "error", "message": "An internal error has occurred."}
finally:
# Ensure client is properly disconnected before response is sent
if client is not None:
try:
await client.disconnect()
except Exception as e:
verbose_logger.warning(f"Error disconnecting MCP client: {e}")
@router.post("/test/connection")
async def test_connection(
@ -332,7 +332,10 @@ if MCP_AVAILABLE:
"""
async def _test_connection_operation(client):
await client.connect()
async def _noop(session):
return "ok"
await client.run_with_session(_noop)
return {"status": "ok"}
return await _execute_with_mcp_client(request, _test_connection_operation)
@ -347,7 +350,13 @@ if MCP_AVAILABLE:
"""
async def _list_tools_operation(client):
list_tools_result: List[MCPTool] = await client.list_tools()
async def _list_tools_session_operation(session):
return await session.list_tools()
list_tools_response = await client.run_with_session(
_list_tools_session_operation
)
list_tools_result: List[MCPTool] = list_tools_response.tools
model_dumped_tools: List[dict] = [
tool.model_dump() for tool in list_tools_result
]

11
tests/mcp_tests/test_mcp_client_unit.py
View file

@ -84,8 +84,8 @@ class TestMCPClientUnitTests:
@pytest.mark.asyncio
@patch("litellm.experimental_mcp_client.client.streamablehttp_client")
@patch("litellm.experimental_mcp_client.client.ClientSession")
async def test_connect(self, mock_session_class, mock_transport):
"""Test connecting to MCP server with authentication."""
async def test_run_with_session(self, mock_session_class, mock_transport):
"""Test run_with_session establishes session with auth headers."""
# Setup mocks
mock_transport_ctx = AsyncMock()
mock_transport.return_value = mock_transport_ctx
@ -102,7 +102,11 @@ class TestMCPClientUnitTests:
auth_type=MCPAuth.bearer_token,
auth_value="test_token",
)
await client.connect()
async def _operation(session):
return "ok"
await client.run_with_session(_operation)
# Verify transport was created with auth headers
call_args = mock_transport.call_args
@ -112,7 +116,6 @@ class TestMCPClientUnitTests:
# Verify session was initialized
mock_session_instance.initialize.assert_called_once()
assert client._session == mock_session_instance
@pytest.mark.asyncio
@patch("litellm.experimental_mcp_client.client.streamablehttp_client")

11
tests/mcp_tests/test_mcp_logging.py
View file

@ -54,8 +54,6 @@ async def test_mcp_cost_tracking():
inputSchema={"type": "object", "properties": {"test": {"type": "string"}}}
)
])
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
# Mock the MCPClient constructor
def mock_client_constructor(*args, **kwargs):
@ -118,7 +116,6 @@ async def test_mcp_cost_tracking():
assert response_list[0].text == "Test response"
# Verify client methods were called
mock_client.__aenter__.assert_called()
mock_client.call_tool.assert_called_once()
######
@ -153,9 +150,6 @@ async def test_mcp_cost_tracking_per_tool():
inputSchema={"type": "object", "properties": {"data": {"type": "string"}}}
)
])
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
mock_client.disconnect = AsyncMock(return_value=None)
# Mock the MCPClient constructor
def mock_client_constructor(*args, **kwargs):
@ -297,9 +291,6 @@ async def test_mcp_tool_call_hook():
inputSchema={"type": "object", "properties": {"test": {"type": "string"}}}
)
])
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
mock_client.disconnect = AsyncMock(return_value=None)
# Mock the MCPClient constructor
def mock_client_constructor(*args, **kwargs):
@ -351,4 +342,4 @@ async def test_mcp_tool_call_hook():
logged_standard_logging_payload = test_logger.standard_logging_payload
print("logged_standard_logging_payload", json.dumps(logged_standard_logging_payload, indent=4))
assert logged_standard_logging_payload is not None, "Standard logging payload should not be None"
assert logged_standard_logging_payload["response_cost"] == 1.42
assert logged_standard_logging_payload["response_cost"] == 1.42

5
tests/mcp_tests/test_mcp_server.py
View file

@ -68,8 +68,6 @@ async def test_mcp_server_manager_https_server():
mock_client = AsyncMock()
mock_client.list_tools = AsyncMock(return_value=mock_tools)
mock_client.call_tool = AsyncMock(return_value=mock_result)
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
# Mock the MCPClient constructor
def mock_client_constructor(*args, **kwargs):
@ -132,7 +130,6 @@ async def test_mcp_server_manager_https_server():
assert result.content[0].text == "Email sent successfully"
# Verify client methods were called
mock_client.__aenter__.assert_called()
mock_client.list_tools.assert_called()
mock_client.call_tool.assert_called_once()
@ -301,7 +298,6 @@ async def test_mcp_http_transport_call_tool_mock():
assert result.content[0].text == "Email sent successfully to test@example.com"
# Verify client methods were called
mock_client.__aenter__.assert_called()
mock_client.call_tool.assert_called_once()
@ -364,7 +360,6 @@ async def test_mcp_http_transport_call_tool_error_mock():
assert "Error: Invalid email address" in result.content[0].text
# Verify client methods were called
mock_client.__aenter__.assert_called()
mock_client.call_tool.assert_called_once()

25
tests/test_litellm/experimental_mcp_client/test_mcp_client.py
View file

@ -39,7 +39,10 @@ class TestMCPClient:
with pytest.raises(
ValueError, match="stdio_config is required for stdio transport"
):
await client.connect()
async def _noop(session):
return None
await client.run_with_session(_noop)
@pytest.mark.asyncio
@patch("litellm.experimental_mcp_client.client.stdio_client")
@ -65,7 +68,10 @@ class TestMCPClient:
client = MCPClient(transport_type=MCPTransport.stdio, stdio_config=stdio_config)
await client.connect()
async def _operation(session):
return "ok"
await client.run_with_session(_operation)
# Verify stdio_client was called with correct parameters
mock_stdio_client.assert_called_once()
@ -109,7 +115,10 @@ class TestMCPClient:
transport_type=MCPTransport.http,
)
await client.connect()
async def _operation(session):
return "ok"
await client.run_with_session(_operation)
# Verify streamablehttp_client was called
mock_streamablehttp_client.assert_called_once()
@ -157,7 +166,10 @@ class TestMCPClient:
ssl_verify=False,
)
await client.connect()
async def _operation(session):
return "ok"
await client.run_with_session(_operation)
# Verify sse_client was called
mock_sse_client.assert_called_once()
@ -208,7 +220,10 @@ class TestMCPClient:
ssl_verify=custom_ca_path,
)
await client.connect()
async def _operation(session):
return "ok"
await client.run_with_session(_operation)
# Verify streamablehttp_client was called
mock_streamablehttp_client.assert_called_once()

1
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -540,7 +540,6 @@ async def test_oauth2_headers_passed_to_mcp_client():
)
# Return a mock client that doesn't actually connect
mock_client = MagicMock()
mock_client.disconnect = AsyncMock()
return mock_client
# Mock _fetch_tools_with_timeout to avoid actual network calls

45
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -1347,10 +1347,10 @@ class TestMCPServerManager:
@pytest.mark.asyncio
async def test_call_tool_without_broken_pipe_error(self):
"""
Test that call_tool properly uses async context manager to avoid broken pipe errors.
This test ensures that tasks are awaited INSIDE the context manager, keeping the connection alive.
Test that call_tool awaits the client call even without a persistent context manager.
Ensures the gathered tasks still include the MCP client call result.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, MagicMock
from mcp.types import CallToolResult
@ -1369,42 +1369,17 @@ class TestMCPServerManager:
manager.tool_name_to_mcp_server_name_mapping["test_tool"] = "test-server"
manager.tool_name_to_mcp_server_name_mapping["test-server-test_tool"] = "test-server"
# Create mock client that tracks context manager usage
mock_client = MagicMock()
context_entered = False
context_exited = False
call_tool_called_inside_context = False
async def mock_aenter(self):
nonlocal context_entered
context_entered = True
return self
async def mock_aexit(self, exc_type, exc_val, exc_tb):
nonlocal context_exited
context_exited = True
# Verify that call_tool was called before context exit
assert (
call_tool_called_inside_context
), "call_tool must be awaited inside context manager"
return False
# Create mock client that tracks call_tool usage
mock_client = AsyncMock()
async def mock_call_tool(params):
nonlocal call_tool_called_inside_context
# Verify we're inside the context when this is called
assert context_entered, "call_tool called outside context manager"
assert not context_exited, "call_tool called after context exit"
call_tool_called_inside_context = True
# Return a mock CallToolResult
result = MagicMock(spec=CallToolResult)
result.content = [{"type": "text", "text": "Tool executed successfully"}]
result.isError = False
return result
mock_client.__aenter__ = mock_aenter
mock_client.__aexit__ = mock_aexit
mock_client.call_tool = mock_call_tool
mock_client.call_tool.side_effect = mock_call_tool
# Mock _create_mcp_client to return our mock client
manager._create_mcp_client = MagicMock(return_value=mock_client)
@ -1437,12 +1412,8 @@ class TestMCPServerManager:
assert result.isError is False
assert len(result.content) > 0
# Verify context manager was used properly
assert context_entered, "Context manager __aenter__ was not called"
assert context_exited, "Context manager __aexit__ was not called"
assert (
call_tool_called_inside_context
), "call_tool was not awaited inside context"
# Verify the MCP client call was awaited exactly once
assert mock_client.call_tool.await_count == 1
if __name__ == "__main__":