AgentHub/server/__tests__/pty-attach.test.js

/**
 * @file pty-attach.test.js
 * @description Unit tests for the PTY-attach helper's runId validation and
 * data/control framing, using a fake node-pty implementation (no real tmux
 * or PTY spawned).
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */
const { describe, it, beforeEach, afterEach } = require("node:test");
const assert = require("node:assert/strict");
const { EventEmitter } = require("node:events");
const ptyAttach = require("../lib/pty-attach");
const tmux = require("../lib/tmux");

function makeFakePty() {
  const emitter = new EventEmitter();
  const writes = [];
  const pty = {
    onData: (fn) => emitter.on("data", fn),
    onExit: (fn) => emitter.on("exit", fn),
    write: (d) => writes.push(d),
    resize: (cols, rows) => writes.push({ resize: [cols, rows] }),
    kill: () => emitter.emit("exit", { exitCode: 0 }),
    __emitter: emitter,
    __writes: writes,
  };
  return pty;
}

function makeFakeWs() {
  const emitter = new EventEmitter();
  const sent = [];
  return {
    on: (evt, fn) => emitter.on(evt, fn),
    send: (data, opts) => sent.push({ data, binary: !!(opts && opts.binary) }),
    __emitter: emitter,
    __sent: sent,
  };
}

describe("pty-attach", () => {
  let fakePty;
  beforeEach(() => {
    fakePty = makeFakePty();
    ptyAttach.__setSpawnImpl(() => fakePty);
    ptyAttach.__resetGroups();
    tmux.__setExecImpl(() => "");
  });
  afterEach(() => {
    tmux.__reset();
  });

  it("rejects a runId that doesn't match agenthub-lane-<digits>", () => {
    assert.throws(() => ptyAttach.validateRunId("../../etc/passwd"), /EBADRUNID/);
    assert.throws(() => ptyAttach.validateRunId("some-other-session"), /EBADRUNID/);
  });

  it("accepts a well-formed runId", () => {
    assert.doesNotThrow(() => ptyAttach.validateRunId("agenthub-lane-42"));
    assert.doesNotThrow(() => ptyAttach.validateRunId("agenthub-lane-42-3"));
    assert.doesNotThrow(() => ptyAttach.validateRunId("agenthub-lane-42-9"));
    assert.throws(() => ptyAttach.validateRunId("agenthub-lane-42-10"), /EBADRUNID/);
    assert.throws(() => ptyAttach.validateRunId("agenthub-lane-42-0"), /EBADRUNID/);
  });

  it("turns tmux mouse mode on for the session it attaches to", () => {
    const calls = [];
    tmux.__setExecImpl((args) => {
      calls.push(args);
      return "";
    });
    ptyAttach.attach(makeFakeWs(), "agenthub-lane-1", { cols: 80, rows: 24 });
    assert.deepEqual(calls[0].slice(0, 5), [
      "set-option",
      "-t",
      "=agenthub-lane-1:",
      "mouse",
      "on",
    ]);
  });

  it("attaches with the RUNNING tmux server's own binary, not the one on PATH", () => {
    // A tmux upgraded while its server runs leaves a newer client on PATH that
    // the old server answers with "open terminal failed: not a terminal".
    tmux.__setExecImpl((args) => (args[0] === "display-message" ? `${process.pid}\n` : ""));
    let file;
    ptyAttach.__setSpawnImpl((f) => {
      file = f;
      return fakePty;
    });
    ptyAttach.attach(makeFakeWs(), "agenthub-lane-1", { cols: 80, rows: 24 });
    assert.equal(file, process.platform === "linux" ? `/proc/${process.pid}/exe` : "tmux");
  });

  it("falls back to tmux on PATH when the server's pid cannot be read", () => {
    tmux.__setExecImpl((args) => {
      if (args[0] === "display-message") throw new Error("no server running");
      return "";
    });
    let file;
    ptyAttach.__setSpawnImpl((f) => {
      file = f;
      return fakePty;
    });
    ptyAttach.attach(makeFakeWs(), "agenthub-lane-1", { cols: 80, rows: 24 });
    assert.equal(file, "tmux");
  });

  it("still attaches even if the mouse-mode set-option call fails", () => {
    tmux.__setExecImpl(() => {
      throw new Error("no such session: agenthub-lane-1");
    });
    assert.doesNotThrow(() =>
      ptyAttach.attach(makeFakeWs(), "agenthub-lane-1", { cols: 80, rows: 24 })
    );
  });

  it("wires PTY data to binary WS frames and WS binary frames to PTY writes", () => {
    const ws = makeFakeWs();
    // Linux: a Windows attach also says `wheel` first (its own test below).
    ptyAttach.attach(ws, "agenthub-lane-1", { cols: 80, rows: 24, platform: "linux" });

    fakePty.__emitter.emit("data", "hello from claude");
    assert.equal(ws.__sent.length, 1);
    assert.equal(ws.__sent[0].data, "hello from claude");
    assert.equal(ws.__sent[0].binary, true);

    ws.__emitter.emit("message", Buffer.from("typed text"), { binary: true });
    assert.deepEqual(fakePty.__writes[0], "typed text");
  });

  it("routes a JSON text frame with type resize to pty.resize", () => {
    const ws = makeFakeWs();
    ptyAttach.attach(ws, "agenthub-lane-1", { cols: 80, rows: 24 });
    ws.__emitter.emit(
      "message",
      Buffer.from(JSON.stringify({ type: "resize", cols: 100, rows: 40 })),
      {
        binary: false,
      }
    );
    assert.deepEqual(fakePty.__writes[0], { resize: [100, 40] });
  });

  it("forwards a plain text WS frame (keystrokes) to pty.write", () => {
    // xterm.js's onData hands the browser a plain string, and
    // WebSocket.send(string) always emits a TEXT frame — so every keystroke
    // arrives here as non-binary. This must reach the pty, not be dropped as
    // an unparseable control message.
    const ws = makeFakeWs();
    ptyAttach.attach(ws, "agenthub-lane-1", { cols: 80, rows: 24 });
    ws.__emitter.emit("message", Buffer.from("ls -la\r"), { binary: false });
    assert.deepEqual(fakePty.__writes[0], "ls -la\r");
  });

  it("sends an exit control message and closes on PTY exit", () => {
    const ws = makeFakeWs();
    let closed = false;
    ws.close = () => {
      closed = true;
    };
    ptyAttach.attach(ws, "agenthub-lane-1", { cols: 80, rows: 24 });
    fakePty.__emitter.emit("exit", { exitCode: 0 });
    const last = ws.__sent[ws.__sent.length - 1];
    assert.equal(last.binary, false);
    assert.deepEqual(JSON.parse(last.data), { type: "exit", code: 0 });
    assert.equal(closed, true);
  });

  it("holds a second viewer at the first one's size until it takes the lead", () => {
    // Two browsers of different sizes on one pane: tmux's `window-size latest`
    // resized the window under whichever was not typing. One drives; the other
    // renders the driver's grid.
    const ptys = [];
    ptyAttach.__setSpawnImpl((_f, _a, opts) => {
      const p = makeFakePty();
      p.opts = opts;
      ptys.push(p);
      return p;
    });
    const resize = (ws, cols, rows, drive) =>
      ws.__emitter.emit(
        "message",
        Buffer.from(JSON.stringify({ type: "resize", cols, rows, drive })),
        {
          binary: false,
        }
      );
    const a = makeFakeWs();
    const b = makeFakeWs();
    ptyAttach.attach(a, "agenthub-lane-1", { cols: 200, rows: 50 });
    ptyAttach.attach(b, "agenthub-lane-1", { cols: 100, rows: 30 });
    assert.deepEqual([ptys[1].opts.cols, ptys[1].opts.rows], [200, 50]);
    assert.deepEqual(JSON.parse(b.__sent[0].data), { type: "size", cols: 200, rows: 50 });

    resize(b, 100, 30); // its own fit, not a takeover: nothing moves
    assert.equal(ptys[1].__writes.length, 0);

    resize(b, 100, 30, true); // a keystroke in b
    assert.deepEqual(ptys[1].__writes.at(-1), { resize: [100, 30] });
    assert.deepEqual(ptys[0].__writes.at(-1), { resize: [100, 30] });
    assert.deepEqual(JSON.parse(a.__sent.at(-1).data), { type: "size", cols: 100, rows: 30 });

    b.__emitter.emit("close"); // the driver leaves: the rest size themselves
    assert.deepEqual(JSON.parse(a.__sent.at(-1).data), { type: "lead" });
  });

  it("kills the PTY attach process when the WS connection closes", () => {
    const ws = makeFakeWs();
    ptyAttach.attach(ws, "agenthub-lane-1", { cols: 80, rows: 24 });
    let killed = false;
    fakePty.kill = () => {
      killed = true;
    };
    ws.__emitter.emit("close");
    assert.equal(killed, true);
  });
});

describe("attachBinary on Windows", () => {
  it("names tmux with its extension — node-pty there never adds `.exe`", () => {
    // A bare `tmux` made node-pty answer "File not found" on every attach, so
    // the browser's terminal opened, closed and reconnected every 1.5 s.
    const saved = process.env.PATH;
    process.env.PATH = "";
    try {
      assert.equal(ptyAttach.attachBinary("win32"), "tmux.exe");
    } finally {
      process.env.PATH = saved;
    }
  });

  it("attaches through node-pty's bundled conpty.dll on Windows only", () => {
    // The inbox ConPTY re-renders psmux's output and leaves stray cells
    // (torn left edge, bars mid-screen) wherever its glyph widths disagree
    // with xterm.js's.
    for (const [platform, dll] of [
      ["win32", true],
      ["linux", undefined],
    ]) {
      let opts;
      ptyAttach.__setSpawnImpl((_f, _a, o) => {
        opts = o;
        return makeFakePty();
      });
      ptyAttach.attach(makeFakeWs(), "agenthub-lane-1", { cols: 80, rows: 24, platform });
      assert.equal(opts.useConptyDll, dll);
      ptyAttach.__resetGroups();
    }
  });

  it("holds a Windows pane's input that follows a paste until psmux has delivered it (psmux#766)", async () => {
    const paste = "\x1b[200~hello\x1b[201~";
    for (const [platform, held] of [
      ["win32", true],
      ["linux", false],
    ]) {
      const fakePty = makeFakePty();
      ptyAttach.__setSpawnImpl(() => fakePty);
      tmux.__setExecImpl(() => "");
      const ws = makeFakeWs();
      ptyAttach.attach(ws, "agenthub-lane-1", { cols: 80, rows: 24, platform });
      ws.__emitter.emit("message", Buffer.from(paste), false);
      ws.__emitter.emit("message", Buffer.from("\r"), false);
      assert.deepEqual(fakePty.__writes, held ? [paste] : [paste, "\r"], platform);
      if (held) {
        await new Promise((r) => setTimeout(r, 800));
        assert.deepEqual(fakePty.__writes, [paste, "\r"]);
      }
      ptyAttach.__resetGroups();
    }
    tmux.__reset();
  });

  it("tells a Windows attach to send the wheel itself (psmux never tracks the mouse)", () => {
    for (const [platform, told] of [
      ["win32", true],
      ["linux", false],
    ]) {
      const ws = makeFakeWs();
      ptyAttach.attach(ws, "agenthub-lane-1", { cols: 80, rows: 24, platform });
      const said = ws.__sent.filter((m) => !m.binary).map((m) => JSON.parse(m.data));
      assert.equal(
        said.some((m) => m.type === "wheel"),
        told,
        platform
      );
      ptyAttach.__resetGroups();
    }
  });
});