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);
    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.throws(() => ptyAttach.validateRunId("agenthub-lane-42-4"), /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], ["set-option", "-t", "agenthub-lane-1", "mouse", "on"]);
  });

  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();
    ptyAttach.attach(ws, "agenthub-lane-1", { cols: 80, rows: 24 });

    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("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);
  });
});