Merge pull request #1887 from ezun-kim/feat/mcp-sse-params

Fix SSE server connection handling for MCP client
This commit is contained in:
Vanessa Pyne
2025-06-24 12:58:05 -05:00
committed by GitHub

View File

@@ -8,8 +8,8 @@ from pipecat.adapters.schemas.tools_schema import ToolsSchema
from pipecat.utils.base_object import BaseObject
try:
from mcp import ClientSession, StdioServerParameters, types
from mcp.client.session import ClientSession
from mcp import ClientSession, StdioServerParameters
from mcp.client.session_group import SseServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
except ModuleNotFoundError as e:
@@ -21,7 +21,7 @@ except ModuleNotFoundError as e:
class MCPClient(BaseObject):
def __init__(
self,
server_params: Union[StdioServerParameters, str],
server_params: Union[StdioServerParameters, SseServerParameters],
**kwargs,
):
super().__init__(**kwargs)
@@ -30,12 +30,12 @@ class MCPClient(BaseObject):
if isinstance(server_params, StdioServerParameters):
self._client = stdio_client
self._register_tools = self._stdio_register_tools
elif isinstance(server_params, str):
elif isinstance(server_params, SseServerParameters):
self._client = sse_client
self._register_tools = self._sse_register_tools
else:
raise TypeError(
f"{self} invalid argument type: `server_params` must be either StdioServerParameters or an SSE server url string."
f"{self} invalid argument type: `server_params` must be either StdioServerParameters or SseServerParameters."
)
async def register_tools(self, llm) -> ToolsSchema:
@@ -90,7 +90,12 @@ class MCPClient(BaseObject):
logger.debug(f"Executing tool '{function_name}' with call ID: {tool_call_id}")
logger.trace(f"Tool arguments: {json.dumps(arguments, indent=2)}")
try:
async with self._client(self._server_params) as (read, write):
async with self._client(
url=self._server_params.url,
headers=self._server_params.headers,
timeout=self._server_params.timeout,
sse_read_timeout=self._server_params.sse_read_timeout
) as (read, write):
async with self._session(read, write) as session:
await session.initialize()
await self._call_tool(session, function_name, arguments, result_callback)
@@ -98,12 +103,16 @@ class MCPClient(BaseObject):
error_msg = f"Error calling mcp tool {function_name}: {str(e)}"
logger.error(error_msg)
logger.exception("Full exception details:")
await result_callback(error_msg)
logger.debug("Starting registration of mcp.run tools")
tool_schemas: List[FunctionSchema] = []
async with self._client(self._server_params) as (read, write):
await result_callback(error_msg)\
logger.debug(f"SSE server parameters: {self._server_params}")
async with self._client(
url=self._server_params.url,
headers=self._server_params.headers,
timeout=self._server_params.timeout,
sse_read_timeout=self._server_params.sse_read_timeout
) as (read, write):
async with self._session(read, write) as session:
await session.initialize()
tools_schema = await self._list_tools(session, mcp_tool_wrapper, llm)