Fix SSE server connection handling for MCP client
### Summary
This PR improves the MCP (Model Context Protocol) client's SSE (Server-Sent Events) server connection handling by replacing the generic string parameter with a proper `SseServerParameters` class.
### Changes
- **Breaking Change**: Changed `server_params` type from `Union[StdioServerParameters, str]` to `Union[StdioServerParameters, SseServerParameters]`
- Added import for `SseServerParameters` from `mcp.client.session_group`
- Updated SSE client connection to use structured parameters instead of a simple URL string
- Fixed error message to correctly reflect the expected parameter types
- Improved logging by changing info-level log to debug-level for consistency
### Details
#### Before
The SSE client connection only accepted a URL string:
```python
async with self._client(self._server_params) as (read, write):
```
#### After
Now properly unpacks SSE server parameters:
```python
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):
```
### Benefits
- **Type Safety**: Stronger type checking with dedicated `SseServerParameters` class
- **Extended Configuration**: Support for custom headers (authentication), timeouts, and SSE-specific settings
- **Better Error Messages**: Clear type error messages when incorrect parameters are provided
- **Improved Debugging**: Debug logging of SSE server parameters for troubleshooting
### Migration Guide
Users need to update their SSE server initialization:
```python
# Before
client = MCPClient("https://example.com/sse")
# After
from mcp.client.session_group import SseServerParameters
client = MCPClient(SseServerParameters(
url="https://example.com/sse",
headers={"Authorization": "Bearer token"},
timeout=30,
sse_read_timeout=60
))
```
### Testing
- [ ] Tested with StdioServerParameters (unchanged behavior)
- [ ] Tested with SseServerParameters with various configurations
- [ ] Verified error handling for invalid parameter types
---
This is a necessary change to support production-ready SSE connections with proper authentication and timeout handling.
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user