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");
const { isWebSocketLoggedIn, isServerMode, userFromRequest } = require("./lib/auth");
const { broadcastAudience, audienceIncludes, canSeeRun } = require("./lib/scope");

let wss = null;
// How often an open terminal's login is re-checked in server mode.
const PTY_RECHECK_MS = Number(process.env.AGENTHUB_PTY_RECHECK_MS) || 10000;

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");
      if (!isWebSocketLoggedIn(info.req)) return done(false, 401, "log in first");
      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, req) => {
    for (const fn of connectionListeners) {
      try {
        fn(ws);
      } catch {
        /* a listener never breaks a connection */
      }
    }
    // Server mode: who this socket belongs to, for broadcast filtering, and
    // the cookie to re-check it by (a logout or a disabled user ends it).
    if (isServerMode()) {
      ws.user = userFromRequest(req);
      ws.cookie = req.headers.cookie || "";
    }
    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;
      }
      if (ws.cookie !== undefined) {
        ws.user = userFromRequest({ headers: { cookie: ws.cookie } });
        if (!ws.user) {
          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+(?:-[1-9])?)$/;

/**
 * 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, cookie) => {
    // Server mode: who attached, re-checked on the beat below.
    if (cookie !== undefined) {
      ws.cookie = cookie;
      ws.runId = runId;
    }
    try {
      attach(ws, runId, size || { cols: 80, rows: 24 });
    } catch (err) {
      // The browser only sees a socket that closed and retries every 1.5 s;
      // this line is the one place that says why.
      console.warn(`[WS] attach ${runId} failed:`, err.message);
      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) || !isWebSocketLoggedIn(req)) {
      socket.destroy();
      return;
    }
    // Server mode: a terminal is a write, so a developer or an analyst
    // attaches only to a pane of their own workspace.
    if (isServerMode() && !canSeeRun({ user: userFromRequest(req), method: "POST" }, match[1])) {
      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),
        isServerMode() ? req.headers.cookie || "" : undefined
      );
    });
  });

  // Server mode: a terminal stays open only while its user may still use it —
  // logged in, not disabled, still owning the pane's workspace.
  const recheck = setInterval(() => {
    if (!ptyWss) {
      clearInterval(recheck);
      return;
    }
    ptyWss.clients.forEach((ws) => {
      if (ws.cookie === undefined) return;
      const user = userFromRequest({ headers: { cookie: ws.cookie } });
      if (!user || !canSeeRun({ user, method: "POST" }, ws.runId)) {
        try {
          ws.close(1008, "log in again");
        } catch {
          /* already gone */
        }
        ws.terminate();
      }
    });
  }, PTY_RECHECK_MS);
  recheck.unref();

  return ptyWss;
}

function broadcast(type, data) {
  if (!wss) return;
  const message = JSON.stringify({ type, data, timestamp: new Date().toISOString() });
  // Server mode: worked out once, and only when a member is connected.
  let audience;
  const reaches = (client) => {
    if (!client.user || client.user.role !== "developer")
      return !isServerMode() || Boolean(client.user);
    if (audience === undefined) {
      try {
        audience = broadcastAudience(type, data);
      } catch {
        audience = { staffOnly: true };
      }
    }
    return audienceIncludes(audience, client.user);
  };
  let stripped;
  wss.clients.forEach((client) => {
    if (client.readyState === 1 && reaches(client)) {
      try {
        const member = client.user && client.user.role === "developer";
        if (member && audience && audience.strip) {
          stripped ??= JSON.stringify({ type, data: {}, timestamp: new Date().toISOString() });
          client.send(stripped);
          return;
        }
        client.send(message);
      } catch {
        // Client closed between readyState check and send — safe to ignore
      }
    }
  });
}

/**
 * Subscribe to new `/ws` browser connections (the update scheduler checks once
 * when a browser returns after its ticks were skipped). Returns an unsubscribe.
 */
const connectionListeners = new Set();
function onConnection(fn) {
  connectionListeners.add(fn);
  return () => connectionListeners.delete(fn);
}

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,
  onConnection,
  closeWebSocket,
  sizeFromQuery,
};