AgentHub/server/lib/omp-share.js

/**
 * @file omp-share.js
 * @description Server mode with S3: the HOST account's Oh My Pi (omp) login,
 * shared with every member's omp pane without any member ever holding it.
 * An admin's switch (Settings → AI tools → Oh My Pi login sharing,
 * `GET/PUT /api/ai-tools/omp-sharing`), off by default.
 *
 * While it is on, this server keeps two omp daemons of its own running, as the
 * host, on loopback only:
 *  - `omp auth-broker serve` over the host's own `~/.omp` credential store —
 *    the vault. Its bearer token (`~/.omp/auth-broker.token`, 0600 in the
 *    host's home) reaches nothing but the gateway's env.
 *  - `omp auth-gateway serve --clients=<file>` — an API gateway that answers
 *    `/v1/messages`, `/v1/chat/completions`, `/v1/responses`, `/v1/models`
 *    with the broker's credentials put in upstream. Its own unscoped token
 *    stays in the host's home too.
 * A member's omp pane reaches the gateway as a SCOPED client: it is handed one
 * key (`AGENTHUB_OMP_KEY`, through tmux -e / its launch file) that the
 * clients file had already listed — by SHA-256 only — before the pane started
 * (the gateway re-reads that file only every 10 s; see `takeToken`), and its
 * `~/.omp/agent/models.yml` gets one `agenthub` provider whose `apiKey` is
 * that env var's NAME. A scoped client can call models but never read a
 * credential or the gateway's usage, and its key stops working within the
 * next sync + the gateway's re-read once its pane is gone (`syncClients`) — so
 * an agent can spend through the login while its pane lives, never carry it
 * away. omp's gateway serves only its built-in providers' models (a custom
 * `models.yml` provider of the host, such as a 9Router, is not shared).
 *
 * Nothing here returns or logs a token. Local mode and a server without S3
 * never start anything.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const crypto = require("node:crypto");
const fs = require("node:fs");
const os = require("node:os");
const path = require("node:path");
const net = require("node:net");
const { spawn } = require("node:child_process");

const PROVIDER = "agenthub";
const TOKEN_ENV = "AGENTHUB_OMP_KEY";
const SYNC_MS = 15_000;

const port = (name, fallback) => {
  const n = parseInt(process.env[name] || "", 10);
  return Number.isInteger(n) && n > 0 && n < 65536 ? n : fallback;
};
const brokerPort = () => port("AGENTHUB_OMP_BROKER_PORT", 4826);
const gatewayPort = () => port("AGENTHUB_OMP_GATEWAY_PORT", 4827);

const dir = () => path.join(require("./claude-home").getDataDir(), "omp-share");
const onFlag = () => path.join(dir(), "enabled");
const clientsFile = () => path.join(dir(), "clients.json");
const pidsFile = () => path.join(dir(), "pids.json");
const tokensFile = () => path.join(dir(), "tokens.json");
const logFile = (name) => path.join(dir(), `${name}.log`);
/** The host's own omp config dir, where the broker keeps its token. */
const ompDir = () => process.env.AGENTHUB_OMP_DIR || path.join(os.homedir(), ".omp");

function available() {
  return require("./linux-accounts").s3Enabled();
}

const enabled = () => fs.existsSync(onFlag());

function ensureDir() {
  fs.mkdirSync(dir(), { recursive: true, mode: 0o700 });
}

// ── tokens and the clients file ───────────────────────────────────────────
//
// The gateway re-reads its clients file only every 10 s, and an omp that
// starts while its key is still refused hangs for good (its model discovery
// fails once and is never retried). So a pane never gets a key minted for
// it: it takes one of STANDBY keys that have been listed for a while, and a
// fresh one takes that key's place in the queue. `tokens.json` (0600, the
// host's data dir) keeps the queue and which run holds which key, so both
// survive a restart; a key whose pane is gone leaves the file at the next sync.

const STANDBY = 4;
/** A run's key stays listed this long before its pane must show up in tmux. */
const GRACE_MS = 60_000;

const sha = (t) => crypto.createHash("sha256").update(t).digest("hex");
const newToken = () => `ahomp_${crypto.randomBytes(24).toString("base64url")}`;

function readTokens() {
  try {
    const j = JSON.parse(fs.readFileSync(tokensFile(), "utf8"));
    return {
      standby: Array.isArray(j.standby) ? j.standby.filter((t) => typeof t === "string") : [],
      used: j.used && typeof j.used === "object" ? j.used : {},
    };
  } catch {
    return { standby: [], used: {} };
  }
}

function writeAtomic(file, text) {
  ensureDir();
  const tmp = `${file}.${process.pid}.tmp`;
  fs.writeFileSync(tmp, text, { mode: 0o600 });
  fs.renameSync(tmp, file);
}

function refill(t) {
  while (t.standby.length < STANDBY) t.standby.push(newToken());
}

/** The gateway's view: every queued key and every key a run holds — hashes only. */
function clientsFor(t) {
  const out = {};
  for (const s of t.standby) out[`standby-${sha(s).slice(0, 16)}`] = { tokenSha256: sha(s) };
  for (const [id, u] of Object.entries(t.used)) {
    if (u && typeof u.hash === "string") out[id] = { tokenSha256: u.hash };
  }
  return out;
}

let lastClients = null;

function save(t) {
  writeAtomic(tokensFile(), `${JSON.stringify(t)}\n`);
  const text = `${JSON.stringify(clientsFor(t), null, 2)}\n`;
  if (text === lastClients && fs.existsSync(clientsFile())) return;
  writeAtomic(clientsFile(), text);
  lastClients = text;
}

/** A key for run `runId`, already listed with the gateway; its place in the queue is refilled. */
function takeToken(runId) {
  const t = readTokens();
  refill(t);
  const token = t.standby.shift();
  t.used[runId] = { hash: sha(token), at: Date.now() };
  refill(t);
  save(t);
  return token;
}

/**
 * Drop the keys of runs whose pane is gone (after GRACE_MS), keep the queue
 * full, and rewrite both files. `live`: the tmux session names, or null when
 * tmux could not say — then nothing is dropped.
 */
function pruneTokens(t, live, now = Date.now()) {
  if (live) {
    for (const [id, u] of Object.entries(t.used)) {
      if (!live.has(id) && now - (Number(u?.at) || 0) > GRACE_MS) delete t.used[id];
    }
  }
  refill(t);
  return t;
}

/** Rewrite the clients file from the live omp panes (a gone pane's key drops out). */
function syncClients() {
  if (!enabled() || !available()) return;
  try {
    save(pruneTokens(readTokens(), require("./tmux").liveSessions()));
  } catch (err) {
    console.warn("omp share: clients sync failed:", err.message);
  }
}

/** Sharing turned off: no key stays listed, queued or held. */
function revokeAll() {
  lastClients = null;
  if (fs.existsSync(tokensFile()) || fs.existsSync(clientsFile())) {
    save({ standby: [], used: {} });
  }
}

// ── a member's models.yml ─────────────────────────────────────────────────

function providerEntry() {
  return {
    baseUrl: `http://127.0.0.1:${gatewayPort()}`,
    apiKey: TOKEN_ENV,
    api: "anthropic-messages",
    authHeader: true,
    discovery: { type: "proxy" },
  };
}

/**
 * The text of `models.yml` with our provider set, or null when it is already
 * there as is. Throws on a file that is not a YAML mapping — it is the
 * member's, and is never overwritten with something it did not say.
 */
function mergedModels(text) {
  const yaml = require("js-yaml");
  const doc = text && text.trim() ? yaml.load(text) : {};
  if (!doc || typeof doc !== "object" || Array.isArray(doc)) {
    throw new Error("models.yml is not a YAML mapping");
  }
  if (
    doc.providers != null &&
    (typeof doc.providers !== "object" || Array.isArray(doc.providers))
  ) {
    throw new Error("models.yml `providers` is not a mapping");
  }
  const want = providerEntry();
  const have = doc.providers?.[PROVIDER];
  if (have && JSON.stringify(have) === JSON.stringify(want)) return null;
  doc.providers = { ...(doc.providers || {}), [PROVIDER]: want };
  return yaml.dump(doc, { lineWidth: -1, noRefs: true });
}

/** Point a member's omp at the gateway: their `~/.omp/agent/models.yml`, written AS them. */
function writeMemberModels(acct) {
  const runAs = require("./run-as");
  const agentDir = path.join(acct.home, ".omp", "agent");
  const file = path.join(agentDir, "models.yml");
  let text = "";
  try {
    text = runAs.readFileSafe(file, "utf8");
  } catch (err) {
    if (err.code !== "ENOENT") throw err;
    runAs.mkdirAs(acct, agentDir);
  }
  const next = mergedModels(text);
  if (next !== null) runAs.writeAs(acct, file, "replace", next);
}

/** The host's own default omp model (`modelRoles.default` in `~/.omp/agent/config.yml`), or null. */
function hostDefaultModel() {
  try {
    const doc = require("js-yaml").load(
      fs.readFileSync(path.join(ompDir(), "agent", "config.yml"), "utf8")
    );
    const m = doc?.modelRoles?.default;
    return typeof m === "string" && m.includes("/") ? m : null;
  } catch {
    return null;
  }
}

/**
 * The `--model` a shared pane runs on: the requested model, or the host's
 * default, as the `agenthub` provider's id (`agenthub/<provider>/<model>`).
 * Without it omp falls back to the provider's FIRST listed model — a retired
 * `claude-3-5-sonnet` on the gateway, which answers 404 and retries forever —
 * or to a built-in provider the member has no login for.
 */
function shareModel(requested, hostDefault = hostDefaultModel()) {
  const want = requested || hostDefault;
  if (!want) return requested || null;
  if (want.startsWith(`${PROVIDER}/`)) return want;
  if (want.includes("/")) return `${PROVIDER}/${want}`;
  const provider = hostDefault ? hostDefault.split("/")[0] : "anthropic";
  return `${PROVIDER}/${provider}/${want}`;
}

/**
 * Called by `spawnRun` for an omp pane: `{env, model(), place()}` — env to add,
 * the `--model` to run on (`shareModel`) and the home to prepare — or
 * null when sharing is off (the pane then runs on the member's own login).
 * Only a MEMBER's pane — the host's omp already holds the login itself.
 */
function forSpawn(member, runId) {
  if (!member || !enabled() || !available()) return null;
  const token = takeToken(runId);
  return {
    // The key itself, never in argv: tmux -e, or a member's launch file.
    env: { OMP_SKIP_SETUP: "1", [TOKEN_ENV]: token },
    model: (requested) => shareModel(requested),
    place() {
      try {
        writeMemberModels(member);
      } catch (err) {
        console.warn(`omp share: ${member.login}'s models.yml left as is:`, err.message);
      }
    },
  };
}

// ── the two daemons ───────────────────────────────────────────────────────

const procs = { broker: null, gateway: null };
const status = { broker: null, gateway: null }; // { pid, startedAt } | { error }
let timers = [];
let stopping = false;

function ompBinary() {
  return require("./ai-tools").detect("omp").path;
}

function readPids() {
  try {
    return JSON.parse(fs.readFileSync(pidsFile(), "utf8")) || {};
  } catch {
    return {};
  }
}

function writePids() {
  ensureDir();
  const pids = {};
  for (const [k, p] of Object.entries(procs)) if (p?.pid) pids[k] = p.pid;
  fs.writeFileSync(pidsFile(), JSON.stringify(pids), { mode: 0o600 });
}

/** A daemon an earlier server left behind (KillMode=process keeps children): ours by its argv only. */
function killStale() {
  for (const [name, pid] of Object.entries(readPids())) {
    if (!Number.isInteger(pid)) continue;
    let argv = "";
    try {
      argv = fs.readFileSync(`/proc/${pid}/cmdline`, "utf8");
    } catch {
      continue;
    }
    if (argv.includes(`auth-${name}\0serve`)) {
      try {
        process.kill(pid, "SIGTERM");
      } catch {
        /* gone already */
      }
    }
  }
}

function portOpen(p) {
  return new Promise((resolve) => {
    const s = net.connect({ host: "127.0.0.1", port: p });
    s.once("connect", () => (s.destroy(), resolve(true)));
    s.once("error", () => resolve(false));
  });
}

async function waitPort(p, ms = 15_000) {
  for (const end = Date.now() + ms; Date.now() < end;) {
    if (await portOpen(p)) return true;
    await new Promise((r) => setTimeout(r, 250));
  }
  return false;
}

/** The env a daemon gets: the host's basics only, never AgentHub's own secrets. */
function daemonEnv(extra = {}) {
  const keep = ["HOME", "USER", "LOGNAME", "PATH", "LANG", "LC_ALL", "TMPDIR"];
  const env = {};
  for (const k of keep) if (process.env[k]) env[k] = process.env[k];
  return { ...env, OMP_SKIP_SETUP: "1", ...extra };
}

function launch(name, bin, args, env) {
  ensureDir();
  const out = fs.openSync(logFile(name), "a", 0o600);
  const child = spawn(bin, args, { env, stdio: ["ignore", out, out], windowsHide: true });
  fs.closeSync(out);
  procs[name] = child;
  status[name] = { pid: child.pid, startedAt: new Date().toISOString() };
  child.on("error", (err) => (status[name] = { error: err.message }));
  child.on("exit", (code, signal) => {
    // One this module stopped (or replaced) itself is no crash.
    if (procs[name] !== child) return;
    procs[name] = null;
    if (stopping || !enabled()) return;
    status[name] = { error: `exited (${signal || code})` };
    // The pair restarts together: the gateway holds the broker's token.
    restartSoon();
  });
  writePids();
  return child;
}

let restartTimer = null;
let backoff = 5_000;
function restartSoon() {
  if (restartTimer) return;
  restartTimer = setTimeout(() => {
    restartTimer = null;
    stopDaemons();
    start().catch(() => {});
    backoff = Math.min(backoff * 2, 5 * 60_000);
  }, backoff);
  if (restartTimer.unref) restartTimer.unref();
}

async function startDaemons() {
  const bin = ompBinary();
  if (!bin) {
    status.broker = status.gateway = { error: "omp is not installed" };
    return;
  }
  killStale();
  // The gateway refuses to start without the clients file it names, and
  // must list the queued keys from its first read.
  save(pruneTokens(readTokens(), null));
  const bport = brokerPort();
  const gport = gatewayPort();
  for (const p of [bport, gport]) {
    if (await waitPortClosed(p)) continue;
    status.broker = status.gateway = { error: `port ${p} is taken` };
    return;
  }
  launch("broker", bin, ["auth-broker", "serve", `--bind=127.0.0.1:${bport}`], daemonEnv());
  if (!(await waitPort(bport))) {
    status.broker = { error: "the broker did not start listening" };
    return restartSoon();
  }
  let token;
  try {
    token = fs.readFileSync(path.join(ompDir(), "auth-broker.token"), "utf8").trim();
  } catch {
    status.broker = { error: "the broker wrote no token" };
    return restartSoon();
  }
  launch(
    "gateway",
    bin,
    ["auth-gateway", "serve", `--bind=127.0.0.1:${gport}`, `--clients=${clientsFile()}`],
    daemonEnv({
      OMP_AUTH_BROKER_URL: `http://127.0.0.1:${bport}`,
      OMP_AUTH_BROKER_TOKEN: token,
    })
  );
  if (await waitPort(gport)) backoff = 5_000;
  else status.gateway = { error: "the gateway did not start listening" };
}

/** A port a previous daemon is still letting go of gets a few seconds. */
async function waitPortClosed(p, ms = 5_000) {
  for (const end = Date.now() + ms; Date.now() < end;) {
    if (!(await portOpen(p))) return true;
    await new Promise((r) => setTimeout(r, 250));
  }
  return false;
}

function stopDaemons() {
  for (const [name, child] of Object.entries(procs)) {
    if (!child) continue;
    procs[name] = null;
    try {
      child.kill("SIGTERM");
    } catch {
      /* gone already */
    }
  }
  try {
    fs.rmSync(pidsFile(), { force: true });
  } catch {
    /* nothing recorded */
  }
}

/** Start the daemons and the clients sync — when sharing is on, under S3 only. */
async function start() {
  stopping = false;
  if (!available() || !enabled()) return;
  if (!timers.length) {
    const t = setInterval(syncClients, SYNC_MS);
    if (t.unref) t.unref();
    timers.push(t);
  }
  syncClients();
  if (procs.broker || procs.gateway) return;
  await startDaemons();
}

/** Stop everything this module started (the server's shutdown, or the switch turned off). */
function stop() {
  stopping = true;
  for (const t of timers) clearInterval(t);
  timers = [];
  if (restartTimer) clearTimeout(restartTimer);
  restartTimer = null;
  stopDaemons();
  status.broker = status.gateway = null;
}

async function setEnabled(on) {
  ensureDir();
  if (on) {
    fs.writeFileSync(onFlag(), "", { mode: 0o600 });
    await start();
  } else {
    fs.rmSync(onFlag(), { force: true });
    stop();
    // No pane keeps a way in once sharing is off.
    revokeAll();
  }
  return state();
}

/** What the settings card shows — never a token. */
function state() {
  const live = (name) => (procs[name] ? status[name] : null);
  return {
    available: available(),
    enabled: enabled(),
    installed: !!ompBinary(),
    broker: live("broker")
      ? { running: true }
      : { running: false, error: status.broker?.error || null },
    gateway: live("gateway")
      ? { running: true, port: gatewayPort() }
      : { running: false, error: status.gateway?.error || null },
  };
}

module.exports = {
  available,
  enabled,
  setEnabled,
  start,
  stop,
  state,
  forSpawn,
  syncClients,
  // exported for the tests
  mergedModels,
  providerEntry,
  shareModel,
  takeToken,
  pruneTokens,
  clientsFor,
  readTokens,
  PROVIDER,
  TOKEN_ENV,
};