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 (server mode: each user only what they may see, and an admin/analyst's session frames marked `mine` when the session is their own).
 * @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, exactAudience, audienceIncludes, canSeeRun } = require("./lib/scope");

let wss = null;
// Session-scoped broadcasts an admin/analyst gets with `mine` (see broadcast).
const MINE_TYPES = new Set([
  "session_created",
  "session_updated",
  "agent_created",
  "agent_updated",
  "new_event",
]);
/**
 * Upgrades refused, by reason (`host`, `token`, `login`, `scope`), since start
 * — and a refusal's one warning line per reason a minute. A browser only sees
 * a socket that closed and retries; this is where the server says why.
 */
const refusals = { host: 0, token: 0, login: 0, scope: 0 };
const lastRefusalLog = new Map();
const REFUSAL_LOG_MS = 60_000;
const REFUSAL_STATUS = {
  host: [403, "Forbidden"],
  token: [401, "Unauthorized"],
  login: [401, "Unauthorized"],
  scope: [403, "Forbidden"],
};
const REFUSAL_WHY = {
  host: "Host header not allowed (DASHBOARD_ALLOWED_HOSTS)",
  token: "missing or invalid dashboard token (DASHBOARD_TOKEN)",
  login: "not logged in",
  scope: "not this user's pane",
};
function noteRefusal(reason, pathname) {
  refusals[reason] = (refusals[reason] || 0) + 1;
  const now = Date.now();
  if (now - (lastRefusalLog.get(reason) || 0) >= REFUSAL_LOG_MS) {
    lastRefusalLog.set(reason, now);
    console.warn(`[WS] refused ${pathname}: ${REFUSAL_WHY[reason]} (${refusals[reason]} so far)`);
  }
}
function refuseUpgrade(socket, reason, pathname) {
  noteRefusal(reason, pathname);
  const [code, text] = REFUSAL_STATUS[reason] || [400, "Bad Request"];
  // Like ws's own abortHandshake: answer, then drop once it is written.
  try {
    socket.once("finish", () => socket.destroy());
    socket.end(`HTTP/1.1 ${code} ${text}\r\nConnection: close\r\nContent-Length: 0\r\n\r\n`);
  } catch {
    socket.destroy();
  }
}
/** A copy of the refusal counters (for the tests and diagnostics). */
const getUpgradeRefusals = () => ({ ...refusals });

// 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) {
      const refuse = (reason, code, msg) => {
        noteRefusal(reason, "/ws");
        return done(false, code, msg);
      };
      if (!isHostAllowed(info.req.headers.host)) return refuse("host", 403, "host not allowed");
      if (!isWebSocketAuthorized(info.req)) return refuse("token", 401, "unauthorized");
      if (!isWebSocketLoggedIn(info.req)) return refuse("login", 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.
 * No size at all (a terminal on a hidden tab has none to give) stays
 * `undefined`: the attach then takes the window's own size and resizes nothing.
 */
function sizeFromQuery(params) {
  const read = (name) => {
    const parsed = Number.parseInt(params.get(name) ?? "", 10);
    return Number.isInteger(parsed) && parsed >= 2 && parsed <= 1000 ? parsed : undefined;
  };
  const cols = read("cols");
  const rows = read("rows");
  return cols && rows ? { cols, rows } : { cols: undefined, rows: undefined };
}

/**
 * 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;
      // A browser that vanished without a close (a sleeping laptop) answers
      // no ping: the beat below ends it, and its presence with it.
      ws.isAlive = true;
      ws.on("pong", () => {
        ws.isAlive = true;
      });
    }
    // Server mode: who has this pane open (pane-presence.js), and the first
    // keystroke of a collaborator-era attach written to the audit log.
    if (cookie !== undefined) {
      const presence = require("./lib/pane-presence");
      const user = userFromRequest({ headers: { cookie } });
      if (user) {
        presence.join(runId, user);
        ws.on("close", () => presence.leave(runId, user));
        let typed = false;
        ws.on("message", (data, isBinary) => {
          if (typed) return;
          const binary = typeof isBinary === "boolean" ? isBinary : false;
          const text = data.toString("utf8");
          if (!binary && text.startsWith('{"type"')) return; // a control frame
          // A mouse report, a focus change or the terminal answering tmux.
          if (!require("./lib/pty-attach").isTyping(text)) return;
          typed = true;
          presence.auditInput(runId, user);
        });
      }
    }
    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)) return refuseUpgrade(socket, "host", url.pathname);
    if (!isWebSocketAuthorized(req)) return refuseUpgrade(socket, "token", url.pathname);
    if (!isWebSocketLoggedIn(req)) return refuseUpgrade(socket, "login", url.pathname);
    // 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])) {
      return refuseUpgrade(socket, "scope", url.pathname);
    }
    // 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) || ws.isAlive === false) {
        try {
          ws.close(1008, "log in again");
        } catch {
          /* already gone */
        }
        ws.terminate();
        return;
      }
      ws.isAlive = false;
      try {
        ws.ping();
      } catch {
        /* closing */
      }
    });
  }, PTY_RECHECK_MS);
  recheck.unref();

  return ptyWss;
}

function broadcast(type, data) {
  if (!wss) return;
  const timestamp = new Date().toISOString();
  const message = JSON.stringify({ type, data, timestamp });
  // Server mode: worked out once, and only when a member is connected.
  let audience;
  // A message for exactly these users (a member's own config change): no
  // one else, whatever their role. Pure, so asked once for every broadcast.
  let only = null;
  try {
    only = exactAudience(type, data);
  } catch {
    only = new Set();
  }
  const reaches = (client) => {
    if (only) return Boolean(client.user && only.has(client.user.id));
    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;
  // Server mode: an admin/analyst gets every user's session news, so a
  // session-scoped frame tells them whether it is their OWN (`mine`) — the
  // pet counts only those. Worked out once per broadcast, on first need.
  const owned = MINE_TYPES.has(type) && isServerMode();
  let mineMsg;
  let otherMsg;
  const forStaff = (client) => {
    if (audience === undefined) {
      try {
        audience = broadcastAudience(type, data);
      } catch {
        audience = { staffOnly: true };
      }
    }
    const mine = Boolean(audience && audience.owners && audience.owners.has(client.user.id));
    if (mine) return (mineMsg ??= JSON.stringify({ type, data, timestamp, mine: true }));
    return (otherMsg ??= JSON.stringify({ type, data, timestamp, mine: false }));
  };
  wss.clients.forEach((client) => {
    if (client.readyState === 1 && reaches(client)) {
      try {
        const member = client.user && client.user.role === "developer";
        if (!only && member && audience && audience.strip) {
          stripped ??= JSON.stringify({ type, data: {}, timestamp: new Date().toISOString() });
          client.send(stripped);
          return;
        }
        if (owned && !only && client.user && !member) {
          client.send(forStaff(client));
          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;
  }
}

/** Tests only: broadcast to these fake clients (`{readyState, user, send}`). */
function _setClientsForTest(clients) {
  wss = clients ? { clients } : null;
}

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