comfyui-mcp/dist/services/relay-client.js

Local-first, agent-native control plane for ComfyUI — MCP server + autonomous sidebar agent that drives your live graph in natural language on ANY LLM: Claude/ChatGPT/Gemini on your subscription (no API key), free local models via Ollama (fully offline),

// Relay client: the comfyui-mcp side of the comfyui-mcp-relay Worker/Durable
// Object (github.com/artokun/comfyui-mcp-relay, private). Dials OUT to the relay
// as the session's single "orchestrator" connection and wraps each
// relay-multiplexed panel-tab connection in a BridgeSocket shim so it can be fed
// straight into UiBridge.attachRelayConnection — from there it's indistinguishable
// from a directly-connected loopback socket.
//
// See that repo's README for the full wire protocol. Summary: the orchestrator's
// one physical socket carries a small {t,id,d} envelope (open/data/close per
// panel connection id); the panel leg itself carries raw, un-enveloped bridge
// frames, so comfyui-mcp-panel needs no changes to speak through the relay.
import { EventEmitter } from "node:events";
import WebSocket from "ws";
import { logger } from "../utils/logger.js";
const OPEN = 1;
const CLOSED = 3;
/** A relay-multiplexed panel connection, shaped to satisfy BridgeSocket. */
class RelaySocketShim extends EventEmitter {
    connId;
    sendUp;
    readyState = OPEN;
    constructor(connId, sendUp) {
        super();
        this.connId = connId;
        this.sendUp = sendUp;
    }
    send(data) {
        if (this.readyState !== OPEN)
            return;
        this.sendUp(this.connId, data);
    }
    close() {
        if (this.readyState === CLOSED)
            return;
        this.readyState = CLOSED;
        this.emit("close");
    }
    terminate() {
        this.close();
    }
    /** No-op: per-connection ping/pong isn't plumbed through the relay envelope in
     *  v1 (see the relay README's "known limitations"). Liveness of the whole
     *  relay connection is handled by RelayClient's own ping/pong below. */
    ping() { }
}
const PING_INTERVAL_MS = 25_000;
const RECONNECT_MAX_MS = 30_000;
export class RelayClient {
    opts;
    ws = null;
    shims = new Map();
    stopped = false;
    reconnectAttempt = 0;
    pingTimer = null;
    constructor(opts) {
        this.opts = opts;
    }
    start() {
        this.connect();
    }
    /** Resolves once the control connection is open (or rejects on the first
     *  connect attempt's error) — lets the caller confirm the relay is actually
     *  reachable before advertising its URL to the pod. */
    waitUntilOpen(timeoutMs = 15_000) {
        return new Promise((resolve, reject) => {
            const timer = setTimeout(() => reject(new Error("relay connect timed out")), timeoutMs);
            const sock = this.ws;
            if (sock?.readyState === WebSocket.OPEN) {
                clearTimeout(timer);
                resolve();
                return;
            }
            sock?.once("open", () => {
                clearTimeout(timer);
                resolve();
            });
            sock?.once("error", (err) => {
                clearTimeout(timer);
                reject(err);
            });
        });
    }
    controlUrl() {
        const u = new URL(`${this.opts.relayUrl.replace(/\/+$/, "")}/s/${this.opts.sessionId}`);
        u.searchParams.set("role", "orchestrator");
        u.searchParams.set("token", this.opts.token);
        if (this.opts.accessKey)
            u.searchParams.set("key", this.opts.accessKey);
        return u.toString();
    }
    connect() {
        if (this.stopped)
            return;
        const ws = new WebSocket(this.controlUrl());
        this.ws = ws;
        ws.on("open", () => {
            this.reconnectAttempt = 0;
            logger.info("[relay-client] connected to relay control channel");
            this.startPing();
        });
        ws.on("message", (buf) => {
            let msg;
            try {
                msg = JSON.parse(String(buf));
            }
            catch {
                return;
            }
            if (msg.t === "open" && msg.id) {
                const shim = new RelaySocketShim(msg.id, (id, data) => this.sendData(id, data));
                this.shims.set(msg.id, shim);
                this.opts.onAttach(shim);
            }
            else if (msg.t === "data" && msg.id) {
                this.shims.get(msg.id)?.emit("message", msg.d ?? "");
            }
            else if (msg.t === "close" && msg.id) {
                const shim = this.shims.get(msg.id);
                this.shims.delete(msg.id);
                shim?.close();
            }
            // "pong" needs no handling — receiving ANY message is itself proof of life.
        });
        ws.on("close", () => {
            this.stopPing();
            // Every relay-mediated panel connection dies with the control channel —
            // they'll reappear (with fresh connection ids) once the panel's own
            // reconnect logic redials and this client re-establishes the session.
            for (const shim of this.shims.values())
                shim.close();
            this.shims.clear();
            if (this.stopped)
                return;
            logger.warn("[relay-client] relay control channel closed — reconnecting");
            this.scheduleReconnect();
        });
        ws.on("error", (err) => {
            logger.warn(`[relay-client] relay error: ${err.message}`);
        });
    }
    sendData(connId, data) {
        if (this.ws?.readyState === WebSocket.OPEN) {
            this.ws.send(JSON.stringify({ t: "data", id: connId, d: data }));
        }
    }
    startPing() {
        this.stopPing();
        this.pingTimer = setInterval(() => {
            if (this.ws?.readyState === WebSocket.OPEN) {
                this.ws.send(JSON.stringify({ t: "ping" }));
            }
        }, PING_INTERVAL_MS);
        this.pingTimer.unref?.();
    }
    stopPing() {
        if (this.pingTimer) {
            clearInterval(this.pingTimer);
            this.pingTimer = null;
        }
    }
    scheduleReconnect() {
        const delay = Math.min(RECONNECT_MAX_MS, 1000 * 2 ** this.reconnectAttempt);
        this.reconnectAttempt += 1;
        setTimeout(() => this.connect(), delay);
    }
    stop() {
        this.stopped = true;
        this.stopPing();
        for (const shim of this.shims.values())
            shim.close();
        this.shims.clear();
        try {
            this.ws?.close();
        }
        catch {
            // already gone
        }
    }
}
//# sourceMappingURL=relay-client.js.map