mcp-use/dist/src/connectors/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 { v4 as uuidv4 } from 'uuid';
import { logger } from '../logging.js';
import { WebSocketConnectionManager } from '../task_managers/websocket.js';
import { BaseConnector } from './base.js';
export class WebSocketConnector extends BaseConnector {
    url;
    headers;
    connectionManager = null;
    ws = null;
    receiverTask = null;
    pending = new Map();
    toolsCache = null;
    constructor(url, opts = {}) {
        super();
        this.url = url;
        this.headers = { ...(opts.headers ?? {}) };
        if (opts.authToken)
            this.headers.Authorization = `Bearer ${opts.authToken}`;
    }
    async connect() {
        if (this.connected) {
            logger.debug('Already connected to MCP implementation');
            return;
        }
        logger.debug(`Connecting via WebSocket: ${this.url}`);
        try {
            this.connectionManager = new WebSocketConnectionManager(this.url, this.headers);
            this.ws = await this.connectionManager.start();
            this.receiverTask = this.receiveLoop();
            this.connected = true;
            logger.debug('WebSocket connected successfully');
        }
        catch (e) {
            logger.error(`Failed to connect: ${e}`);
            await this.cleanupResources();
            throw e;
        }
    }
    async disconnect() {
        if (!this.connected) {
            logger.debug('Not connected to MCP implementation');
            return;
        }
        logger.debug('Disconnecting …');
        await this.cleanupResources();
        this.connected = false;
    }
    sendRequest(method, params = null) {
        if (!this.ws)
            throw new Error('WebSocket is not connected');
        const id = uuidv4();
        const payload = JSON.stringify({ id, method, params: params ?? {} });
        return new Promise((resolve, reject) => {
            this.pending.set(id, { resolve, reject });
            this.ws.send(payload, (err) => {
                if (err) {
                    this.pending.delete(id);
                    reject(err);
                }
            });
        });
    }
    async receiveLoop() {
        if (!this.ws)
            return;
        const socket = this.ws; // Node.ws or browser WS
        const onMessage = (msg) => {
            let data;
            try {
                data = JSON.parse(msg.data ?? msg);
            }
            catch (e) {
                logger.warn('Received non‑JSON frame', e);
                return;
            }
            const id = data.id;
            if (id && this.pending.has(id)) {
                const { resolve, reject } = this.pending.get(id);
                this.pending.delete(id);
                if ('result' in data)
                    resolve(data.result);
                else if ('error' in data)
                    reject(data.error);
            }
            else {
                logger.debug('Received unsolicited message', data);
            }
        };
        socket.addEventListener ? socket.addEventListener('message', onMessage) : socket.on('message', onMessage);
        // keep promise pending until close
        return new Promise((resolve) => {
            const onClose = () => {
                socket.removeEventListener ? socket.removeEventListener('message', onMessage) : socket.off('message', onMessage);
                this.rejectAll(new Error('WebSocket closed'));
                resolve();
            };
            socket.addEventListener ? socket.addEventListener('close', onClose) : socket.on('close', onClose);
        });
    }
    rejectAll(err) {
        for (const { reject } of this.pending.values())
            reject(err);
        this.pending.clear();
    }
    async initialize() {
        logger.debug('Initializing MCP session over WebSocket');
        const result = await this.sendRequest('initialize');
        const toolsList = await this.listTools();
        this.toolsCache = toolsList.map(t => t);
        logger.debug(`Initialized with ${this.toolsCache.length} tools`);
        return result;
    }
    async listTools() {
        const res = await this.sendRequest('tools/list');
        return res.tools ?? [];
    }
    async callTool(name, args) {
        return await this.sendRequest('tools/call', { name, arguments: args });
    }
    async listResources() {
        const resources = await this.sendRequest('resources/list');
        return { resources: Array.isArray(resources) ? resources : [] };
    }
    async readResource(uri) {
        const res = await this.sendRequest('resources/read', { uri });
        return { content: res.content, mimeType: res.mimeType };
    }
    async request(method, params = null) {
        return await this.sendRequest(method, params);
    }
    get tools() {
        if (!this.toolsCache)
            throw new Error('MCP client is not initialized');
        return this.toolsCache;
    }
    async cleanupResources() {
        // Stop receiver
        if (this.receiverTask)
            await this.receiverTask.catch(() => { });
        this.receiverTask = null;
        // Reject pending
        this.rejectAll(new Error('WebSocket disconnected'));
        // Stop connection manager → closes socket
        if (this.connectionManager) {
            await this.connectionManager.stop();
            this.connectionManager = null;
            this.ws = null;
        }
        this.toolsCache = null;
    }
}