From c720cfc7c7cb6e16e319ddff8238fd7f54bc5105 Mon Sep 17 00:00:00 2001
From: Yousif Astarabadi <6870090+yousifa@users.noreply.github.com>
Date: Tue, 24 Jun 2025 19:15:32 -0700
Subject: [PATCH] updated streamablehttp to use StreamableHttpParameters type

---
 src/pipecat/services/mcp_service.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/pipecat/services/mcp_service.py b/src/pipecat/services/mcp_service.py
index 3e1b04681..2f61343dd 100644
--- a/src/pipecat/services/mcp_service.py
+++ b/src/pipecat/services/mcp_service.py
@@ -9,7 +9,7 @@ from pipecat.utils.base_object import BaseObject
 
 try:
     from mcp import ClientSession, StdioServerParameters
-    from mcp.client.session_group import SseServerParameters
+    from mcp.client.session_group import SseServerParameters, StreamableHttpParameters
     from mcp.client.sse import sse_client
     from mcp.client.stdio import stdio_client
     from mcp.client.streamable_http import streamablehttp_client
@@ -22,13 +22,12 @@ except ModuleNotFoundError as e:
 class MCPClient(BaseObject):
     def __init__(
         self,
-        server_params: Union[StdioServerParameters, SseServerParameters],
+        server_params: Union[StdioServerParameters, SseServerParameters, StreamableHttpParameters],
         **kwargs,
     ):
         super().__init__(**kwargs)
         self._server_params = server_params
         self._session = ClientSession
-        self._additional_headers = additional_headers or {}
         
         if isinstance(server_params, StdioServerParameters):
             self._client = stdio_client
@@ -36,7 +35,7 @@ class MCPClient(BaseObject):
         elif isinstance(server_params, SseServerParameters):
             self._client = sse_client
             self._register_tools = self._sse_register_tools
-        elif isinstance(server_params, str) and streamable_http:
+        elif isinstance(server_params, StreamableHttpParameters):
             self._client = streamablehttp_client
             self._register_tools = self._streamable_http_register_tools
         else:
@@ -182,7 +181,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, headers=self._additional_headers) as (read_stream, write_stream, _):
+                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_stream, write_stream, _):
                     async with self._session(read_stream, write_stream) as session:
                         await session.initialize()
                         await self._call_tool(session, function_name, arguments, result_callback)
@@ -194,7 +198,12 @@ class MCPClient(BaseObject):
 
         logger.debug("Starting registration of mcp.run tools using streamable HTTP")
 
-        async with self._client(self._server_params, headers=self._additional_headers) as (read_stream, write_stream, _):
+        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_stream, write_stream, _):
             async with self._session(read_stream, write_stream) as session:
                 await session.initialize()
                 tools_schema = await self._list_tools(session, mcp_tool_wrapper, llm)