AgentHub/server/__tests__/pool-account-choice.test.js

/**
 * @file pool-account-choice.test.js
 * @description A user's own choice of pool account (account-pins.js
 * `setChoice`/`preferredFor`, `/api/me/pool-accounts`), fake accounts only:
 * the balancer runs the user's new panes on the chosen account; an admin's pin
 * wins and locks the choice (ELOCKED); only accounts the user's group allows
 * may be chosen (ENOACCOUNT), and one later taken from the group falls back
 * to the balancer; a removed account drops its choices; the routes list no
 * token, need a login cookie and refuse in local mode.
 * @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.mkdtempSync(path.join(os.tmpdir(), "pool-choice-test-"));
process.env.DASHBOARD_DB_PATH = path.join(TMP, "dashboard.db");
process.env.CLAUDE_HOME = path.join(TMP, ".claude");
process.env.DASHBOARD_LIVENESS_PROBE = "0";

const { createApp, startServer } = require("../index");
const { db } = require("../db");
const auth = require("../lib/auth");
const groups = require("../lib/pool-groups");
const cursorPool = require("../lib/cursor-pool");
const claudePool = require("../lib/claude-pool");
const pins = require("../lib/account-pins");

require("../lib/cc-usage").fetchUsage = async () => ({});

const KEY = (c) => `key_${c.repeat(40)}`;
const TOKEN = (c) => `sk-ant-oat01-${c.repeat(40)}`;
let server;
let BASE;
let MIA;
let LEO;

async function call(method, p, { cookie, body } = {}) {
  const res = await fetch(BASE + p, {
    method,
    headers: {
      "Content-Type": "application/json",
      Origin: BASE,
      ...(cookie ? { cookie } : {}),
    },
    body: body ? JSON.stringify(body) : undefined,
  });
  return { status: res.status, body: await res.json().catch(() => null) };
}

async function login(username, password) {
  const res = await fetch(BASE + "/api/auth/login", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ username, password }),
  });
  return res.headers.get("set-cookie").split(";")[0];
}

before(async () => {
  process.env.AGENTHUB_MODE = "server";
  server = await startServer(createApp(), 0);
  BASE = `http://127.0.0.1:${server.address().port}`;
  await auth.createUser("boss", "admin-pass-1", "admin");
  MIA = (await auth.createUser("mia", "member-pass-1", "developer")).id;
  LEO = (await auth.createUser("leo", "member-pass-2", "developer")).id;
});

after(() => {
  delete process.env.AGENTHUB_MODE;
  server?.close();
  fs.rmSync(TMP, { recursive: true, force: true });
});

beforeEach(() => {
  db.prepare("DELETE FROM pool_groups").run();
  db.prepare("DELETE FROM account_pins").run();
  db.prepare("DELETE FROM account_choices").run();
  for (const dir of [cursorPool.poolDir(), claudePool.poolDir()]) {
    fs.rmSync(dir, { recursive: true, force: true });
  }
  for (const c of ["a", "b", "c"]) cursorPool.addAccount(c, { apiKey: KEY(c) });
});

describe("pool account choice: the balancer", () => {
  it("runs the user's new panes on the account they chose, others unchanged", () => {
    assert.equal(cursorPool.pickAccount(MIA).name, "a");
    pins.setChoice(MIA, "cursor", "c");
    assert.equal(cursorPool.pickAccount(MIA).name, "c");
    assert.equal(cursorPool.pickAccount(LEO).name, "a");
    claudePool.addAccount("one", TOKEN("o"));
    claudePool.addAccount("two", TOKEN("t"));
    pins.setChoice(MIA, "claude", "two");
    assert.equal(claudePool.pickAccount(null, MIA).name, "two");
    pins.setChoice(MIA, "cursor", null);
    assert.equal(cursorPool.pickAccount(MIA).name, "a", "automatic again");
  });

  it("lets an admin's pin win, and refuses a choice while pinned", () => {
    pins.setChoice(MIA, "cursor", "b");
    pins.setPin(MIA, "cursor", "c");
    assert.equal(cursorPool.pickAccount(MIA).name, "c");
    assert.throws(() => pins.setChoice(MIA, "cursor", "a"), { code: "ELOCKED" });
    pins.setPin(MIA, "cursor", null);
    assert.equal(cursorPool.pickAccount(MIA).name, "b", "the choice comes back");
  });

  it("allows only the group's accounts, and falls back once one leaves it", () => {
    const g = groups.createGroup("A");
    groups.setGroupUsers(g.id, [MIA]);
    groups.setGroupAccounts(g.id, "cursor", ["b", "c"]);
    assert.throws(() => pins.setChoice(MIA, "cursor", "a"), { code: "ENOACCOUNT" });
    assert.throws(() => pins.setChoice(MIA, "cursor", "zz"), { code: "ENOACCOUNT" });
    pins.setChoice(MIA, "cursor", "c");
    assert.equal(cursorPool.pickAccount(MIA).name, "c");
    groups.setGroupAccounts(g.id, "cursor", ["b"]);
    assert.equal(pins.preferredFor(MIA, "cursor"), null);
    assert.equal(cursorPool.pickAccount(MIA).name, "b");
  });

  it("drops a removed account's choices", () => {
    pins.setChoice(MIA, "cursor", "b");
    cursorPool.removeAccount("b");
    assert.equal(pins.choiceFor(MIA, "cursor"), null);
  });
});

describe("pool account choice: /api/me/pool-accounts", () => {
  it("lists only the tools and accounts the user may use, never a token", async () => {
    const g = groups.createGroup("A");
    groups.setGroupUsers(g.id, [MIA]);
    groups.setGroupAccounts(g.id, "cursor", ["b"]);
    const cookie = await login("mia", "member-pass-1");
    const r = await call("GET", "/api/me/pool-accounts", { cookie });
    assert.equal(r.status, 200);
    assert.deepEqual(
      r.body.tools.map((t) => t.tool),
      ["cursor"]
    );
    assert.deepEqual(
      r.body.tools[0].accounts.map((a) => a.name),
      ["b"]
    );
    assert.ok(!JSON.stringify(r.body).includes(KEY("b")));
  });

  it("chooses, goes back to automatic, refuses while pinned, and is audited", async () => {
    const cookie = await login("leo", "member-pass-2");
    let r = await call("PUT", "/api/me/pool-accounts", {
      cookie,
      body: { tool: "cursor", account: "c" },
    });
    assert.equal(r.status, 200);
    assert.equal(r.body.choice, "c");
    assert.equal(r.body.switched, "next-session");
    assert.equal(cursorPool.pickAccount(LEO).name, "c");
    r = await call("PUT", "/api/me/pool-accounts", {
      cookie,
      body: { tool: "cursor", account: "nope" },
    });
    assert.equal(r.status, 404);
    assert.equal(r.body.error.code, "ENOACCOUNT");
    r = await call("PUT", "/api/me/pool-accounts", { cookie, body: { tool: "x", account: "c" } });
    assert.equal(r.status, 400);
    pins.setPin(LEO, "cursor", "a");
    r = await call("PUT", "/api/me/pool-accounts", {
      cookie,
      body: { tool: "cursor", account: "b" },
    });
    assert.equal(r.status, 409);
    assert.equal(r.body.error.code, "ELOCKED");
    pins.setPin(LEO, "cursor", null);
    r = await call("PUT", "/api/me/pool-accounts", {
      cookie,
      body: { tool: "cursor", account: null },
    });
    assert.equal(r.status, 200);
    assert.equal(r.body.choice, null);
    const audited = db
      .prepare("SELECT COUNT(*) AS n FROM audit_log WHERE path LIKE '%pool account%'")
      .get().n;
    assert.ok(audited >= 2);
  });

  it("needs a login", async () => {
    const r = await call("PUT", "/api/me/pool-accounts", {
      body: { tool: "cursor", account: "a" },
    });
    assert.equal(r.status, 401);
  });
});