From 3035305c76bfaa3326e9a4fdb1230156db139c7e Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 10 Oct 2025 07:15:32 +0900 Subject: [PATCH 1/2] fix: Ensure MCP client stays open during tool call --- .../mcp_server/mcp_server_manager.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 3c27a65f06a..2efd27a1aa9 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1098,13 +1098,14 @@ class MCPServerManager: extra_headers=extra_headers, ) - async with client: - # Use the original tool name (without prefix) for the actual call - call_tool_params = MCPCallToolRequestParams( - name=original_tool_name, - arguments=arguments, - ) - tasks.append(asyncio.create_task(client.call_tool(call_tool_params))) + call_tool_params = MCPCallToolRequestParams( + name=original_tool_name, + arguments=arguments, + ) + async def _call_tool_via_client(client, params): + async with client: + return await client.call_tool(params) + tasks.append(asyncio.create_task(_call_tool_via_client(client, call_tool_params))) try: mcp_responses = await asyncio.gather(*tasks) From 2b1c061d35384b6639f2f2b715c18619a5f85a1c Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 10 Oct 2025 07:19:03 +0900 Subject: [PATCH 2/2] fix: lint error --- .../mcp_server/mcp_server_manager.py | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 2efd27a1aa9..53dfd9f42c4 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -354,12 +354,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( @@ -879,7 +879,6 @@ class MCPServerManager: proxy_logging_obj: ProxyLogging, server: MCPServer, ): - ## check if the tool is allowed or banned for the given server if not self.check_allowed_or_banned_tools(name, server): raise HTTPException( @@ -959,7 +958,7 @@ class MCPServerManager: verbose_logger.error(f"Guardrail blocked MCP tool call pre call: {str(e)}") raise e - async def call_tool( + async def call_tool( # noqa: PLR0915 self, name: str, arguments: Dict[str, Any], @@ -1102,10 +1101,14 @@ class MCPServerManager: name=original_tool_name, arguments=arguments, ) + async def _call_tool_via_client(client, params): async with client: return await client.call_tool(params) - tasks.append(asyncio.create_task(_call_tool_via_client(client, call_tool_params))) + + tasks.append( + asyncio.create_task(_call_tool_via_client(client, call_tool_params)) + ) try: mcp_responses = await asyncio.gather(*tasks)