AgentHub/server/websocket.js

/**
 * @file WebSocket functionalities for real-time communication with clients, including connection management, heartbeat for detecting dead connections, and broadcasting messages to all connected clients.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const { WebSocketServer } = require("ws");
const { isHostAllowed, isWebSocketAuthorized } = require("./lib/security");

let wss = null;

function initWebSocket(server) {
  // `noServer: true` + a manual, path-checked `server.on("upgrade", ...)`
  // rather than the `{server, path}` shorthand: that shorthand's own
  // internal upgrade listener calls `handleUpgrade` for EVERY upgrade on the
  // shared http.Server (path filtering happens inside `handleUpgrade`,
  // which `abortHandshake`s with 400 on a mismatch) — so it was answering,
  // and killing, `/ws-pty/*` upgrades before the PTY server's own listener
  // (registered below by `initPtyWebSocket`) ever got a chance to run.
  wss = new WebSocketServer({
    noServer: true,
    maxPayload: 64 * 1024,
    verifyClient(info, done) {
      if (!isHostAllowed(info.req.headers.host)) return done(false, 403, "host not allowed");
      if (!isWebSocketAuthorized(info.req)) return done(false, 401, "unauthorized");
      return done(true);
    },
  });

  server.on("upgrade", (req, socket, head) => {
    const url = new URL(req.url, "http://localhost");
    if (url.pathname !== "/ws") return; // not ours — `/ws-pty/*` handles its own path.
    // closeWebSocket() nulls `wss` but cannot unregister this listener, so an
    // upgrade arriving mid-shutdown would throw an unhandled TypeError and take
    // the process down instead of letting it exit gracefully. Clients retry, so
    // dropping the socket is the correct answer here.
    if (!wss) {
      socket.destroy();
      return;
    }
    wss.handleUpgrade(req, socket, head, (ws) => {
      wss.emit("connection", ws, req);
    });
  });

  wss.on("connection", (ws) => {
    ws.isAlive = true;
    ws.on("pong", () => {
      ws.isAlive = true;
    });
    ws.on("error", (err) => {
      // Log but don't crash — client disconnects are normal
      if (err.code !== "ECONNRESET") {
        console.warn("[WS] client error:", err.code || err.message);
      }
    });
  });

  // Heartbeat every 30s to detect dead connections
  const interval = setInterval(() => {
    if (!wss) {
      clearInterval(interval);
      return;
    }
    wss.clients.forEach((ws) => {
      if (!ws.isAlive) {
        ws.terminate();
        return;
      }
      ws.isAlive = false;
      ws.ping();
    });
  }, 30000);
  interval.unref();

  wss.on("close", () => {
    clearInterval(interval);
  });

  return wss;
}

let ptyWss = null;
const PTY_PATH_RE = /^\/ws-pty\/(agenthub-lane-\d+)$/;

/**
 * The attaching terminal's own size, off the attach URL's query. tmux resizes
 * the window to whichever client attached most recently, so a hardcoded 80×24
 * here is not a harmless default: the pane's program redraws its whole UI at
 * 80 columns and only gets the real width back when the client's first
 * `resize` message lands, leaving 80-column wreckage in the scrollback that
 * shows up as torn left-margin fragments when the user scrolls back
 * (issue #33). Bounded, because this number reaches a `TIOCSWINSZ` ioctl.
 */
function sizeFromQuery(params) {
  const read = (name, fallback) => {
    const parsed = Number.parseInt(params.get(name) ?? "", 10);
    return Number.isInteger(parsed) && parsed >= 2 && parsed <= 1000 ? parsed : fallback;
  };
  return { cols: read("cols", 80), rows: read("rows", 24) };
}

/**
 * Second WebSocket server, dedicated to the terminal-run PTY transport
 * (`/ws-pty/:runId`). Kept separate from the `/ws` JSON-broadcast path so
 * raw binary PTY frames never have to coexist with the typed
 * `{type, data, timestamp}` envelope the rest of the app relies on. Reuses
 * the exact same auth guard as `/ws`.
 *
 * `runId` is a path SEGMENT, not a fixed string, so this can't use the `ws`
 * library's `{server, path}` shorthand (that option only matches an exact
 * string). Instead this server is created with `noServer: true` and the
 * upgrade is handled manually — the same `server.on("upgrade", ...)` pattern
 * `ws` itself uses internally, just filtered to `/ws-pty/*` first so `/ws`'s
 * own upgrade handling (already registered by `initWebSocket`) is untouched.
 */
function initPtyWebSocket(server) {
  const { attach } = require("./lib/pty-attach");
  ptyWss = new WebSocketServer({ noServer: true, maxPayload: 1024 * 1024 });

  ptyWss.on("connection", (ws, runId, size) => {
    try {
      attach(ws, runId, size || { cols: 80, rows: 24 });
    } catch (err) {
      try {
        ws.close(1008, err.message);
      } catch {
        /* ignore */
      }
    }
  });

  server.on("upgrade", (req, socket, head) => {
    const url = new URL(req.url, "http://localhost");
    const match = url.pathname.match(PTY_PATH_RE);
    if (!match) return; // not ours — the `/ws` WebSocketServer (registered
    // by initWebSocket, also attached to this same http.Server) handles its
    // own path independently and ignores upgrades it doesn't match too.
    if (!isHostAllowed(req.headers.host)) {
      socket.destroy();
      return;
    }
    if (!isWebSocketAuthorized(req)) {
      socket.destroy();
      return;
    }
    // Same shutdown race as `/ws` above — and far easier to hit here, since
    // TerminalView re-attaches every 1.5s while a run console is open.
    if (!ptyWss) {
      socket.destroy();
      return;
    }
    ptyWss.handleUpgrade(req, socket, head, (ws) => {
      ptyWss.emit("connection", ws, match[1], sizeFromQuery(url.searchParams));
    });
  });

  return ptyWss;
}

function broadcast(type, data) {
  if (!wss) return;
  const message = JSON.stringify({ type, data, timestamp: new Date().toISOString() });
  wss.clients.forEach((client) => {
    if (client.readyState === 1) {
      try {
        client.send(message);
      } catch {
        // Client closed between readyState check and send — safe to ignore
      }
    }
  });
}

function getConnectionCount() {
  if (!wss) return 0;
  let count = 0;
  wss.clients.forEach((client) => {
    if (client.readyState === 1) count++;
  });
  return count;
}

/**
 * Tear down the WebSocket server for a graceful shutdown. Open WS clients keep
 * their underlying TCP sockets alive, which prevents http.Server#close() from
 * ever completing — under `node --watch` that turns every restart into a
 * multi-second "waiting for graceful termination" stall. Terminating the
 * clients first lets the HTTP server drain and close promptly.
 */
function closeWebSocket() {
  if (wss) {
    wss.clients.forEach((client) => {
      try {
        client.terminate();
      } catch {
        /* already gone */
      }
    });
    try {
      wss.close();
    } catch {
      /* ignore */
    }
    wss = null;
  }
  if (ptyWss) {
    ptyWss.clients.forEach((client) => {
      try {
        client.terminate();
      } catch {
        /* already gone */
      }
    });
    try {
      ptyWss.close();
    } catch {
      /* ignore */
    }
    ptyWss = null;
  }
}

module.exports = {
  initWebSocket,
  initPtyWebSocket,
  broadcast,
  getConnectionCount,
  closeWebSocket,
  sizeFromQuery,
};