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