mcp-use/dist/src/task_managers/websocket.js

A utility library for integrating Model Context Protocol (MCP) with LangChain, Zod, and related tools. Provides helpers for schema conversion, event streaming, and SDK usage.

import WS from 'ws';
import { logger } from '../logging.js';
import { ConnectionManager } from './base.js';
export class WebSocketConnectionManager extends ConnectionManager {
    url;
    headers;
    _ws = null;
    /**
     * @param url     The WebSocket URL to connect to.
     * @param headers Optional headers to include in the connection handshake.
     */
    constructor(url, headers = {}) {
        super();
        this.url = url;
        this.headers = headers;
    }
    /** Establish a WebSocket connection and wait until it is open. */
    async establishConnection() {
        logger.debug(`Connecting to WebSocket: ${this.url}`);
        return new Promise((resolve, reject) => {
            const ws = new WS(this.url, { headers: this.headers });
            this._ws = ws;
            const onOpen = () => {
                // eslint-disable-next-line ts/no-use-before-define
                cleanup();
                logger.debug('WebSocket connected successfully');
                resolve(ws);
            };
            const onError = (err) => {
                // eslint-disable-next-line ts/no-use-before-define
                cleanup();
                logger.error(`Failed to connect to WebSocket: ${err}`);
                reject(err);
            };
            const cleanup = () => {
                ws.off('open', onOpen);
                ws.off('error', onError);
            };
            // Register listeners (browser vs Node API differences handled)
            ws.on('open', onOpen);
            ws.on('error', onError);
        });
    }
    /** Cleanly close the WebSocket connection. */
    async closeConnection(connection) {
        logger.debug('Closing WebSocket connection');
        return new Promise((resolve) => {
            const onClose = () => {
                connection.off('close', onClose);
                this._ws = null;
                resolve();
            };
            if (connection.readyState === WS.CLOSED) {
                onClose();
                return;
            }
            connection.on('close', onClose);
            try {
                connection.close();
            }
            catch (e) {
                logger.warn(`Error closing WebSocket connection: ${e}`);
                onClose();
            }
        });
    }
}