AgentHub/server/lib/pty-attach.js

/**
 * @file pty-attach.js
 * @description Bridges one WebSocket connection to an `@lydell/node-pty`-backed
 * `tmux attach-session` process, run with the live tmux server's own binary
 * (`attachBinary`). Binary WS frames carry raw PTY bytes in
 * both directions; text WS frames carry small JSON control messages
 * (`resize`, and an outbound `exit` sent once when the pane process/tmux
 * session ends). Multiple browser tabs each get their own PTY attach
 * process, and ONE of them sizes the window: the DRIVER — the first to attach,
 * then whichever sent `{type:"resize", drive:true}` last (the client does on a
 * click or a keystroke). Every other attach is a follower: its PTY is held at
 * the driver's size and told so (`{type:"size", cols, rows}`) so its xterm
 * renders the same grid, scaled to fit, instead of each browser resizing tmux
 * under the others. A driver that leaves sends `{type:"lead"}` to the rest.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const fs = require("fs");
const tmux = require("./tmux");

const RUN_ID_RE = /^agenthub-lane-\d+(?:-[1-9])?$/;

/**
 * Reject anything that isn't exactly `agenthub-lane-<digits>[-<1-9>]` before it can ever
 * reach a tmux/PTY command — the trust boundary for this WS path, since a
 * validated runId is the only thing standing between an authenticated WS
 * client and naming an arbitrary session on the host.
 */
function validateRunId(runId) {
  if (typeof runId !== "string" || !RUN_ID_RE.test(runId)) {
    const err = new Error(`EBADRUNID: invalid runId: ${runId}`);
    err.code = "EBADRUNID";
    throw err;
  }
}

// Test seam — real implementation set in Step 4.
let spawnImpl = null;
function __setSpawnImpl(fn) {
  spawnImpl = fn;
}

/**
 * The binary to run `tmux attach` with: the RUNNING server's own executable
 * (`/proc/<pid>/exe`, which still works after a package upgrade replaced the
 * file), else `tmux` on PATH. A client must match its server — tmux installed
 * in a newer version while the server runs makes PATH's client fail every
 * attach with "open terminal failed: not a terminal", though every other tmux
 * command still works. Linux only; elsewhere PATH's tmux as before — on
 * Windows by its full file name: node-pty there looks the file up on PATH as
 * given, never adding `.exe`, so a bare `tmux` failed every attach with
 * "File not found" (the browser saw a socket that opened and closed, again
 * every 1.5 s).
 */
function attachBinary(platform = process.platform) {
  if (platform === "win32") return require("./ai-tools").which("tmux") || "tmux.exe";
  const pid = tmux.serverPid();
  if (!pid) return "tmux";
  const exe = `/proc/${pid}/exe`;
  try {
    fs.accessSync(exe, fs.constants.X_OK);
    return exe;
  } catch {
    return "tmux";
  }
}

// runId -> { driver, viewers: Set<{ ws, pty, cols, rows }> }
const groups = new Map();

const tell = (v, msg) => {
  try {
    v.ws.send(JSON.stringify(msg), { binary: false });
  } catch {
    /* client gone — its close handler cleans up */
  }
};

/** `v` sizes the window now: every other attach is held at its size and told. */
function drive(group, v) {
  group.driver = v;
  v.pty.resize(v.cols, v.rows);
  for (const o of group.viewers) {
    if (o === v) continue;
    o.pty.resize(v.cols, v.rows);
    tell(o, { type: "size", cols: v.cols, rows: v.rows });
  }
}

/**
 * Attach `ws` to the tmux session `runId`. Spawns one PTY-backed
 * `tmux attach-session -t <runId>` per call.
 */
function attach(ws, runId, { cols, rows }) {
  validateRunId(runId);
  // Retroactively covers sessions started before mouse mode was wired into
  // newSession (or before a server restart picked that change up) — see
  // tmux.js's enableMouse.
  try {
    tmux.enableMouse(runId);
  } catch {
    /* session gone or tmux unavailable — attach below will surface that */
  }
  let group = groups.get(runId);
  if (!group) groups.set(runId, (group = { driver: null, viewers: new Set() }));
  const lead = group.driver;
  // A developer's pane is on that account's own tmux server, attached AS that
  // account (S3: never a host client on a socket the account controls).
  // node-pty forks natively, so drop-caps' child_process patch never sees it:
  // the client must not keep the unit's CAP_DAC_READ_SEARCH (S3).
  const [file, args] = require("./drop-caps").ptyArgv(
    ...(tmux.memberAttach(runId) || [attachBinary(), ["attach-session", "-t", tmux.session(runId)]])
  );
  // A follower attaches at the driver's size: tmux sizes the window to the
  // newest client, so its own size would resize the driver's pane first.
  const pty = spawnImpl(file, args, {
    name: "xterm-256color",
    cols: lead ? lead.cols : cols || 80,
    rows: lead ? lead.rows : rows || 24,
  });
  const me = { ws, pty, cols: cols || 80, rows: rows || 24 };
  group.viewers.add(me);
  if (lead) tell(me, { type: "size", cols: lead.cols, rows: lead.rows });
  else group.driver = me;

  // Backpressure: a slow viewer (a remote browser on a thin link) must not
  // queue a chatty pane's output in this process without bound. Past 1 MB
  // unsent the pty is paused — tmux holds the pane and redraws — and resumed
  // once the socket drains below 256 KB.
  let paused = false;
  const drained = () => {
    if (paused && ws.bufferedAmount < 256 * 1024) {
      paused = false;
      pty.resume?.();
    }
  };
  pty.onData((data) => {
    try {
      ws.send(data, { binary: true }, drained);
      if (!paused && ws.bufferedAmount > 1024 * 1024) {
        paused = true;
        pty.pause?.();
      }
    } catch {
      /* client gone between data event and send — safe to ignore */
    }
  });

  pty.onExit(({ exitCode }) => {
    try {
      ws.send(JSON.stringify({ type: "exit", code: exitCode }), { binary: false });
    } catch {
      /* ignore */
    }
    try {
      ws.close();
    } catch {
      /* ignore */
    }
  });

  ws.on("message", (data, isBinary) => {
    // node-pty attach processes for `tmux attach` are already tolerant of
    // resize mid-stream; the `ws` lib passes isBinary either as the second
    // callback arg (newer) or via `data.binary` on some transports — this
    // helper's own tests exercise the `{binary}` option shape used above.
    const binary = typeof isBinary === "boolean" ? isBinary : !!(isBinary && isBinary.binary);
    const text = data.toString("utf8");
    if (binary) {
      pty.write(text);
      return;
    }
    // The browser's WebSocket API sends a JS string as a text frame, and
    // xterm.js's onData callback hands over plain strings — so every
    // keystroke arrives here as text, not binary. Only a JSON control frame
    // (matched by this same prefix check the client uses for output) is
    // NOT keystroke input; everything else must reach the pty or typing
    // does nothing.
    if (text.startsWith('{"type"')) {
      let msg;
      try {
        msg = JSON.parse(text);
      } catch {
        pty.write(text);
        return;
      }
      if (msg && msg.type === "resize" && Number.isFinite(msg.cols) && Number.isFinite(msg.rows)) {
        me.cols = msg.cols;
        me.rows = msg.rows;
        // A follower's own size is remembered, not applied: its PTY stays at
        // the driver's until it takes the lead.
        if (msg.drive || group.driver === me || !group.driver) drive(group, me);
        return;
      }
    }
    pty.write(text);
  });

  ws.on("close", () => {
    group.viewers.delete(me);
    if (!group.viewers.size) groups.delete(runId);
    else if (group.driver === me) {
      group.driver = null;
      for (const o of group.viewers) tell(o, { type: "lead" });
    }
    try {
      pty.kill();
    } catch {
      /* already gone */
    }
  });

  return pty;
}

// Real spawn implementation — lazy-required so unit tests never load the
// native PTY addon unless they explicitly opt in. Uses @lydell/node-pty (a
// drop-in-API-compatible fork of node-pty) rather than node-pty itself:
// node-pty ships prebuilt binaries for darwin/win32 only, so on Linux it
// needs a native build via its install script — but the plugin install path
// runs `npm install --ignore-scripts` deliberately (see plugin-bootstrap.js)
// to avoid requiring a build toolchain on the user's machine. @lydell/node-pty
// instead ships the platform binary as a regular optionalDependency
// (@lydell/node-pty-linux-x64 etc.), so a plain --ignore-scripts install still
// resolves a working native binding with no compiler needed.
__setSpawnImpl((...args) => require("@lydell/node-pty").spawn(...args));

const __resetGroups = () => groups.clear();

module.exports = { attach, attachBinary, validateRunId, __setSpawnImpl, __resetGroups };