AgentHub/server/__tests__/cursor-pool.test.js

/**
 * @file cursor-pool.test.js
 * @description The Cursor account pool (server/lib/cursor-pool.js), fake keys
 * only: key validation and file modes, the balancer (fewest live Cursor panes,
 * a pin wins), the key reaching a `cursor-agent` pane only through tmux `-e
 * CURSOR_API_KEY` and never its argv, the account recorded on the run, pins
 * for `cursor` and the account_pins rebuild that drops the old tool CHECK, and
 * an empty pool leaving every spawn as it was.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const { describe, it, before, after, beforeEach } = require("node:test");
const assert = require("node:assert/strict");
const path = require("node:path");
const fs = require("node:fs");
const os = require("node:os");

const TMP = fs.realpathSync(fs.mkdtempSync(path.join(os.tmpdir(), "cursor-pool-test-")));
process.env.DASHBOARD_DB_PATH = path.join(TMP, "dashboard.db");
process.env.CLAUDE_HOME = path.join(TMP, "claude");

const { db } = require("../db");
const pool = require("../lib/cursor-pool");
const pins = require("../lib/account-pins");
const tmux = require("../lib/tmux");
const pty = require("../lib/pty-run");

const KEY = (c) => `key_${c.repeat(40)}`;
let tmuxArgv;

const envOf = (argv, name) =>
  argv
    .slice(0, argv.indexOf("--"))
    .find((x) => String(x).startsWith(`${name}=`))
    ?.slice(name.length + 1);

function run(id, account, status = "running") {
  db.prepare(
    `INSERT INTO dashboard_runs (id, mode, cwd, status, started_at, account, tool)
     VALUES (?, 'interactive', '/tmp', ?, ?, ?, 'cursor')`
  ).run(id, status, new Date().toISOString(), account);
}

before(() => {
  tmux.__setExecImpl((args) => {
    if (args[0] === "has-session") throw Object.assign(new Error("gone"), { status: 1 });
    if (args[0] === "new-session") tmuxArgv = args;
    return "";
  });
});

after(() => {
  tmux.__setExecImpl(null);
  fs.rmSync(TMP, { recursive: true, force: true });
});

beforeEach(() => {
  fs.rmSync(pool.poolDir(), { recursive: true, force: true });
  db.prepare("DELETE FROM dashboard_runs").run();
  tmuxArgv = null;
});

describe("cursor pool: accounts", () => {
  it("stores a key 0600 in a 0700 dir and never lists it", () => {
    const added = pool.addAccount("team", { apiKey: KEY("a") });
    assert.deepEqual(Object.keys(added).sort(), ["addedAt", "kind", "name"]);
    assert.equal(fs.statSync(pool.poolDir()).mode & 0o777, 0o700);
    assert.equal(fs.statSync(path.join(pool.poolDir(), "team.json")).mode & 0o777, 0o600);
    const listed = pool.poolStatus();
    assert.equal(listed.length, 1);
    assert.ok(!JSON.stringify(listed).includes(KEY("a")));
    assert.equal(listed[0].quotaPercent, null);
  });

  it("refuses a bad name, a bad key and a taken name", () => {
    assert.throws(() => pool.addAccount("../x", { apiKey: KEY("a") }), { code: "EBADNAME" });
    assert.throws(() => pool.addAccount("x", { apiKey: "short" }), { code: "EBADTOKEN" });
    assert.throws(() => pool.addAccount("x", { apiKey: `${KEY("a")}\nX=1` }), {
      code: "EBADTOKEN",
    });
    pool.addAccount("x", { apiKey: KEY("a") });
    assert.throws(() => pool.addAccount("x", { apiKey: KEY("b") }), { code: "EEXISTS" });
    pool.removeAccount("x");
    assert.throws(() => pool.removeAccount("x"), { code: "ENOACCOUNT" });
  });

  it("picks the account with the fewest live panes", () => {
    pool.addAccount("a", { apiKey: KEY("a") });
    pool.addAccount("b", { apiKey: KEY("b") });
    run("r1", "a");
    run("r2", "b", "completed");
    assert.equal(pool.pickAccount().name, "b");
  });
});

describe("cursor pool: spawn", () => {
  it("an empty pool leaves a Cursor spawn exactly as before", () => {
    pty.spawnRun({ laneId: 7, cwd: "/tmp", tool: "cursor" });
    assert.equal(envOf(tmuxArgv, "CURSOR_API_KEY"), undefined);
    const row = db.prepare("SELECT account FROM dashboard_runs WHERE id = 'agenthub-lane-7'").get();
    assert.equal(row.account, null);
  });

  it("the key reaches the pane through tmux -e only, and the run names its account", () => {
    pool.addAccount("a", { apiKey: KEY("a") });
    pty.spawnRun({ laneId: 7, cwd: "/tmp", tool: "cursor", initialPrompt: "hi" });
    assert.equal(envOf(tmuxArgv, "CURSOR_API_KEY"), KEY("a"));
    const after = tmuxArgv.slice(tmuxArgv.indexOf("--"));
    assert.ok(!after.some((x) => String(x).includes(KEY("a"))), "never in the pane's argv");
    const row = db.prepare("SELECT account FROM dashboard_runs WHERE id = 'agenthub-lane-7'").get();
    assert.equal(row.account, "a");

    // Other tools never get it.
    tmuxArgv = null;
    pty.spawnRun({ laneId: 8, cwd: "/tmp", tool: "claude" });
    assert.ok(!tmuxArgv.some((x) => String(x).includes("CURSOR_API_KEY")));
  });
});

describe("cursor pool: pins", () => {
  it("account_pins takes a cursor pin, and removing the account drops it", () => {
    const sql = db
      .prepare("SELECT sql FROM sqlite_master WHERE type='table' AND name='account_pins'")
      .get().sql;
    assert.ok(!sql.includes("CHECK(tool IN"), "no tool CHECK left");
    assert.ok(pins.TOOLS.includes("cursor"));
    const uid = db
      .prepare("INSERT INTO users (username, password_hash) VALUES ('pinned', 'x')")
      .run().lastInsertRowid;
    pool.addAccount("a", { apiKey: KEY("a") });
    db.prepare("INSERT INTO account_pins (user_id, tool, account) VALUES (?, 'cursor', 'a')").run(
      uid
    );
    assert.equal(pins.listPins("cursor").length, 1);
    pool.removeAccount("a");
    assert.equal(pins.listPins("cursor").length, 0);
  });
});