mirror of
https://github.com/mcp-use/mcp-use.git
synced 2026-09-14 18:06:01 +08:00
Pietro Zullo
18f51eddb8
* move: session * move: client * move:config * move: exceptions * move: connectors -> client.connectors * move: task_managers -> client.task_managers * move: observability -> agents.observability * move: middleware -> client.middleware * move: managers -> agents.managers * move: auth -> client.auth * move: adapters -> agents.adapters * test: bc test * test: bc test * chore: remove relative imports * fix tests * chore: update docs * chore: deprecation warnings * chore: fix dangerous non absolute import * fix: remove legacy imports * telemetry
69 lines
2.1 KiB
Python
69 lines
2.1 KiB
Python
"""
|
|
WebSocket connection management for MCP implementations.
|
|
|
|
This module provides a connection manager for WebSocket-based MCP connections.
|
|
"""
|
|
|
|
from typing import Any
|
|
|
|
from mcp.client.websocket import websocket_client
|
|
|
|
from mcp_use.client.task_managers.base import ConnectionManager
|
|
from mcp_use.logging import logger
|
|
|
|
|
|
class WebSocketConnectionManager(ConnectionManager[tuple[Any, Any]]):
|
|
"""Connection manager for WebSocket-based MCP connections.
|
|
|
|
This class handles the lifecycle of WebSocket connections, ensuring proper
|
|
connection establishment and cleanup.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
url: str,
|
|
headers: dict[str, str] | None = None,
|
|
):
|
|
"""Initialize a new WebSocket connection manager.
|
|
|
|
Args:
|
|
url: The WebSocket URL to connect to
|
|
headers: Optional HTTP headers
|
|
"""
|
|
super().__init__()
|
|
self.url = url
|
|
self.headers = headers or {}
|
|
|
|
async def _establish_connection(self) -> tuple[Any, Any]:
|
|
"""Establish a WebSocket connection.
|
|
|
|
Returns:
|
|
The established WebSocket connection
|
|
|
|
Raises:
|
|
Exception: If connection cannot be established
|
|
"""
|
|
logger.debug(f"Connecting to WebSocket: {self.url}")
|
|
# Create the context manager
|
|
# Note: The current MCP websocket_client implementation doesn't support headers
|
|
# If headers need to be passed, this would need to be updated when MCP supports it
|
|
self._ws_ctx = websocket_client(self.url)
|
|
|
|
# Enter the context manager
|
|
read_stream, write_stream = await self._ws_ctx.__aenter__()
|
|
|
|
# Return the streams
|
|
return (read_stream, write_stream)
|
|
|
|
async def _close_connection(self) -> None:
|
|
"""Close the WebSocket connection."""
|
|
if self._ws_ctx:
|
|
# Exit the context manager
|
|
try:
|
|
logger.debug("Closing WebSocket connection")
|
|
await self._ws_ctx.__aexit__(None, None, None)
|
|
except Exception as e:
|
|
logger.warning(f"Error closing WebSocket connection: {e}")
|
|
finally:
|
|
self._ws_ctx = None
|