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:
ezun-kim
2025-05-24 22:35:57 +09:00
parent f86d002ceb
commit 3da711ba8b

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)