AgentHub/server/lib/account-pins.js

/**
 * @file account-pins.js
 * @description Account pins: an admin's choice that one server-mode user's
 * new sessions of a tool run on one named pool account instead of the one the
 * balancer would pick (claude-pool.js / codex-pool.js / cursor-pool.js `pick`) — the Users
 * page's "switch a user's account". One pin per user and tool
 * (`account_pins`); the user is the workspace's OWNER
 * (`lanes.owner_id`), never whoever clicked Start. A pin is a preference the
 * pool still bounds: an account a pane may not use (expired, or a login the
 * pane's home cannot take) falls back to the balancer. A Claude login is
 * written into the user's home right away (`applyPin`), over the pool login
 * their running panes use, so those move to it without a restart.
 *
 * `setAccountUsers` is the Users page's Edit of one account: who runs on it.
 *
 * A user may also CHOOSE an account of their own (`account_choices`, Settings
 * → AI tools → Pool account, `/api/me/pool-accounts`): `preferredFor` — the
 * balancers' question — answers the admin's pin first, then that choice. A
 * choice is only ever among the accounts the user's group may use
 * (pool-groups.js `allowed`), checked when it is made and again at every
 * spawn, so an account later taken from the group falls back to the balancer.
 * A pinned user cannot choose (`ELOCKED`): the admin's pin is the decision.
 *
 * Server mode only: in local mode there are no users and nothing is pinned.
 * Removing the account drops its pins (`dropAccount`); deleting the user drops
 * theirs (ON DELETE CASCADE, and `auth.deleteUser`).
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const TOOLS = ["claude", "codex", "cursor"];
const POOLS = { claude: "./claude-pool", codex: "./codex-pool", cursor: "./cursor-pool" };

const err = (code, message) => Object.assign(new Error(message), { code });
const db = () => require("../db").db;
const poolOf = (tool) => require(POOLS[tool]);

function validTool(tool) {
  if (!TOOLS.includes(tool)) throw err("EBADTOOL", `tool: ${TOOLS.join(", ")}`);
  return tool;
}

/** The account name `ownerId`'s new `tool` sessions are pinned to, or null. */
function pinFor(ownerId, tool) {
  if (!require("./auth").isServerMode() || !Number.isInteger(ownerId)) return null;
  try {
    return (
      db()
        .prepare("SELECT account FROM account_pins WHERE user_id = ? AND tool = ?")
        .get(ownerId, tool)?.account ?? null
    );
  } catch {
    return null; // a spawn never fails over a pin
  }
}

/** The account `userId` chose for `tool` themselves, or null. */
function choiceFor(userId, tool) {
  if (!require("./auth").isServerMode() || !Number.isInteger(userId)) return null;
  try {
    return (
      db()
        .prepare("SELECT account FROM account_choices WHERE user_id = ? AND tool = ?")
        .get(userId, tool)?.account ?? null
    );
  } catch {
    return null; // a spawn never fails over a choice
  }
}

/**
 * The account the balancer should prefer for `ownerId`'s new `tool` pane: the
 * admin's pin, else the owner's own choice while their group still allows it.
 */
function preferredFor(ownerId, tool) {
  const pin = pinFor(ownerId, tool);
  if (pin) return pin;
  const choice = choiceFor(ownerId, tool);
  if (!choice) return null;
  return require("./pool-groups").allowed(ownerId, tool)(choice) ? choice : null;
}

/**
 * `userId`'s own choice of `account` for `tool`, or back to the balancer
 * (`account` null). Refused while an admin pins them (`ELOCKED`), and for an
 * account not in the pool or not their group's (`ENOACCOUNT`).
 */
function setChoice(userId, tool, account) {
  validTool(tool);
  if (!Number.isInteger(userId)) throw err("ENOUSER", `no user ${userId}`);
  if (pinFor(userId, tool)) {
    throw err("ELOCKED", `an admin has pinned your ${tool} sessions to one account`);
  }
  if (account == null || account === "") {
    db().prepare("DELETE FROM account_choices WHERE user_id = ? AND tool = ?").run(userId, tool);
    return null;
  }
  const name = String(account);
  if (!poolOf(tool).hasAccount(name) || !require("./pool-groups").allowed(userId, tool)(name)) {
    throw err("ENOACCOUNT", `no ${tool} account ${name} you may use`);
  }
  db()
    .prepare(
      `INSERT INTO account_choices (user_id, tool, account) VALUES (?, ?, ?)
       ON CONFLICT(user_id, tool) DO UPDATE SET account = excluded.account,
         chosen_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')`
    )
    .run(userId, tool, name);
  return name;
}

/**
 * What Settings → AI tools → Pool account shows `userId`: per tool with at
 * least one account they may use, those accounts (never a token), the admin's
 * pin and their own choice.
 */
function myAccounts(userId) {
  const groups = require("./pool-groups");
  const out = [];
  for (const tool of TOOLS) {
    let status = [];
    try {
      status = poolOf(tool).poolStatus();
    } catch {
      status = [];
    }
    const ok = groups.allowed(userId, tool);
    const pin = pinFor(userId, tool);
    const accounts = status
      .filter((a) => ok(a.name))
      .map((a) => ({
        name: a.name,
        kind: a.kind,
        email: a.email ?? null,
        expired: Boolean(a.expired),
        quotaPercent: a.quotaPercent ?? null,
        full: Boolean(a.full),
      }));
    if (accounts.length === 0) continue;
    const choice = choiceFor(userId, tool);
    out.push({ tool, pin, choice: choice && ok(choice) ? choice : null, accounts });
  }
  return out;
}

/** Every pin of `tool`, with its user's name: `[{userId, username, account, pinnedAt}]`. */
function listPins(tool) {
  return db()
    .prepare(
      `SELECT p.user_id AS userId, u.username, p.account, p.pinned_at AS pinnedAt
         FROM account_pins p JOIN users u ON u.id = p.user_id
        WHERE p.tool = ? ORDER BY u.username`
    )
    .all(validTool(tool));
}

/**
 * Pin `userId`'s new `tool` sessions to `account`, or unpin them (`account`
 * null). Returns `{userId, username, account}`.
 */
function setPin(userId, tool, account) {
  validTool(tool);
  const user = Number.isInteger(userId)
    ? db().prepare("SELECT id, username FROM users WHERE id = ?").get(userId)
    : null;
  if (!user) throw err("ENOUSER", `no user ${userId}`);
  if (account == null || account === "") {
    db().prepare("DELETE FROM account_pins WHERE user_id = ? AND tool = ?").run(user.id, tool);
    return { userId: user.id, username: user.username, account: null };
  }
  const name = String(account);
  if (!poolOf(tool).hasAccount(name)) {
    throw err("ENOACCOUNT", `no ${tool} account ${name} in the pool`);
  }
  db()
    .prepare(
      `INSERT INTO account_pins (user_id, tool, account) VALUES (?, ?, ?)
       ON CONFLICT(user_id, tool) DO UPDATE SET account = excluded.account,
         pinned_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')`
    )
    .run(user.id, tool, name);
  return { userId: user.id, username: user.username, account: name };
}

/**
 * The Users page's Edit of one account: exactly `userIds` run on `account`
 * afterwards. A user pinned to another account of the tool moves here; one
 * pinned here but left out goes back to the balancer. One transaction.
 * Returns `{added, removed}` (each `[{userId, username}]`).
 */
function setAccountUsers(tool, account, userIds) {
  validTool(tool);
  const name = String(account ?? "");
  if (!poolOf(tool).hasAccount(name)) {
    throw err("ENOACCOUNT", `no ${tool} account ${name} in the pool`);
  }
  if (!Array.isArray(userIds) || !userIds.every(Number.isInteger)) {
    throw err("EBADUSER", "userIds: an array of user ids");
  }
  const want = [...new Set(userIds)];
  const users = want.map((id) => {
    const u = db().prepare("SELECT id, username FROM users WHERE id = ?").get(id);
    if (!u) throw err("ENOUSER", `no user ${id}`);
    return u;
  });
  const before = listPins(tool);
  const here = new Set(before.filter((p) => p.account === name).map((p) => p.userId));
  const added = users.filter((u) => !here.has(u.id));
  const removed = before.filter((p) => p.account === name && !want.includes(p.userId));
  db().transaction(() => {
    for (const u of added) setPin(u.id, tool, name);
    for (const p of removed) setPin(p.userId, tool, null);
  })();
  return {
    added: added.map((u) => ({ userId: u.id, username: u.username })),
    removed: removed.map((p) => ({ userId: p.userId, username: p.username })),
  };
}

/**
 * Make a pin — or the user's own choice — take effect now where it can: a Claude login is written into
 * the user's home at once (claude-pool `switchHome`), so their running panes
 * move to it without a restart. Anything else goes in at the next spawn.
 * Answers what happened (see `switchHome`).
 */
async function applyPin(userId, tool) {
  const pin = preferredFor(userId, tool);
  if (!pin || tool !== "claude") return "next-session";
  const user = db().prepare("SELECT username FROM users WHERE id = ?").get(userId);
  const member = user ? require("./run-as").accountForUser(user.username) : null;
  return poolOf(tool).switchHome(member, pin);
}

/** A removed account takes its pins, choices — and its place in every group — with it. */
function dropAccount(tool, account) {
  try {
    db().prepare("DELETE FROM account_pins WHERE tool = ? AND account = ?").run(tool, account);
    db().prepare("DELETE FROM account_choices WHERE tool = ? AND account = ?").run(tool, account);
  } catch {
    /* no database (a bare CLI test): nothing pinned */
  }
  require("./pool-groups").dropAccount(tool, account);
}

module.exports = {
  TOOLS,
  poolOf,
  pinFor,
  choiceFor,
  preferredFor,
  setChoice,
  myAccounts,
  listPins,
  setPin,
  setAccountUsers,
  applyPin,
  dropAccount,
};