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 `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.
 *
 * 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"];

const err = (code, message) => Object.assign(new Error(message), { code });
const db = () => require("../db").db;
const poolOf = (tool) => require(tool === "codex" ? "./codex-pool" : "./claude-pool");

function validTool(tool) {
  if (!TOOLS.includes(tool)) throw err("EBADTOOL", "tool: claude or codex");
  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
  }
}

/** 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 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 = pinFor(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 with it. */
function dropAccount(tool, account) {
  try {
    db().prepare("DELETE FROM account_pins WHERE tool = ? AND account = ?").run(tool, account);
  } catch {
    /* no database (a bare CLI test): nothing pinned */
  }
}

module.exports = { TOOLS, pinFor, listPins, setPin, setAccountUsers, applyPin, dropAccount };