AgentHub/server/__tests__/omp-share.test.js

/**
 * @file omp-share.test.js
 * @description Oh My Pi login sharing (server/lib/omp-share.js): a member's
 * `models.yml` gets one `agenthub` provider pointing at the loopback gateway
 * with the key's env var NAME as its key (the rest of the file kept, a
 * non-mapping refused); a pane takes a key the clients file listed BEFORE it
 * started (the gateway re-reads it only every 10 s), and a gone pane's key
 * leaves it; a spawn is untouched while sharing is off, for a host pane, or
 * without S3; and the route refuses outside S3. No omp daemon is started.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const { describe, it, before, after, afterEach } = 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(), "omp-share-"));
process.env.DASHBOARD_DB_PATH = path.join(TMP, "dashboard.db");
process.env.DASHBOARD_LIVENESS_PROBE = "0";
process.env.LANES_ROOT = path.join(TMP, "lanes");
process.env.CLAUDE_HOME = path.join(TMP, ".claude");
process.env.DASHBOARD_DATA_DIR = path.join(TMP, "data");
process.env.AGENTHUB_OMP_GATEWAY_PORT = "4999";

const crypto = require("node:crypto");
const yaml = require("js-yaml");

const sha = (t) => crypto.createHash("sha256").update(t).digest("hex");
const { createApp, startServer } = require("../index");
const share = require("../lib/omp-share");
const linux = require("../lib/linux-accounts");
const runAs = require("../lib/run-as");
const tmux = require("../lib/tmux");

const realS3 = linux.s3Enabled;
const realLive = tmux.liveSessions;
const flag = path.join(TMP, "data", "omp-share", "enabled");
const clientsFile = path.join(TMP, "data", "omp-share", "clients.json");

function turnOn() {
  fs.mkdirSync(path.dirname(flag), { recursive: true });
  fs.writeFileSync(flag, "");
  linux.s3Enabled = () => true;
}

afterEach(() => {
  linux.s3Enabled = realS3;
  tmux.liveSessions = realLive;
  fs.rmSync(flag, { force: true });
});

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

describe("mergedModels", () => {
  it("adds the provider to an empty or missing file", () => {
    const doc = yaml.load(share.mergedModels(""));
    assert.deepEqual(doc, {
      providers: {
        agenthub: {
          baseUrl: "http://127.0.0.1:4999",
          apiKey: "AGENTHUB_OMP_KEY",
          api: "anthropic-messages",
          authHeader: true,
          discovery: { type: "proxy" },
        },
      },
    });
  });

  it("keeps the member's own providers and settings", () => {
    const text =
      "equivalence:\n  x: y\nproviders:\n  mine:\n    baseUrl: http://h/v1\n    apiKey: K\n";
    const doc = yaml.load(share.mergedModels(text));
    assert.equal(doc.providers.mine.baseUrl, "http://h/v1");
    assert.deepEqual(doc.equivalence, { x: "y" });
    assert.equal(doc.providers.agenthub.apiKey, "AGENTHUB_OMP_KEY");
  });

  it("writes nothing when the provider is already there as is", () => {
    const once = share.mergedModels("");
    assert.equal(share.mergedModels(once), null);
  });

  it("replaces a stale provider of ours", () => {
    const stale = "providers:\n  agenthub:\n    baseUrl: http://127.0.0.1:1\n";
    assert.equal(
      yaml.load(share.mergedModels(stale)).providers.agenthub.baseUrl,
      "http://127.0.0.1:4999"
    );
  });

  it("refuses a file that is not a mapping", () => {
    assert.throws(() => share.mergedModels("- a\n- b\n"), /mapping/);
    assert.throws(() => share.mergedModels("providers: [1]\n"), /mapping/);
  });
});

describe("shareModel", () => {
  const host = "anthropic/claude-sonnet-5-5";
  it("runs on the host's default when none was asked for", () => {
    assert.equal(share.shareModel(null, host), "agenthub/anthropic/claude-sonnet-5-5");
  });
  it("maps a requested model onto the gateway's provider", () => {
    assert.equal(share.shareModel("claude-opus-5-5", host), "agenthub/anthropic/claude-opus-5-5");
    assert.equal(share.shareModel("openai-codex/gpt-5.5", host), "agenthub/openai-codex/gpt-5.5");
    assert.equal(share.shareModel("agenthub/anthropic/x", host), "agenthub/anthropic/x");
  });
  it("leaves the choice to omp when there is nothing to go on", () => {
    assert.equal(share.shareModel(null, null), null);
  });
});

describe("forSpawn", () => {
  const member = { login: "ah-mia", home: path.join(TMP, "home-mia") };

  it("is null while sharing is off, for a host pane, or without S3", () => {
    assert.equal(share.forSpawn(member, "agenthub-lane-1"), null); // off
    fs.mkdirSync(path.dirname(flag), { recursive: true });
    fs.writeFileSync(flag, "");
    assert.equal(share.forSpawn(member, "agenthub-lane-1"), null); // no S3
    linux.s3Enabled = () => true;
    assert.equal(share.forSpawn(null, "agenthub-lane-1"), null); // the host's
  });

  it("hands out a key the gateway already listed, and writes models.yml AS the member", () => {
    turnOn();
    share.syncClients(); // the queue, listed before any pane asks
    const listedBefore = JSON.parse(fs.readFileSync(clientsFile, "utf8"));
    const writes = [];
    const real = { read: runAs.readFileSafe, mkdir: runAs.mkdirAs, write: runAs.writeAs };
    runAs.readFileSafe = () => {
      const e = new Error("missing");
      e.code = "ENOENT";
      throw e;
    };
    runAs.mkdirAs = (acct, dir) => writes.push(["mkdir", acct.login, dir]);
    runAs.writeAs = (acct, file, mode, content) => writes.push([mode, acct.login, file, content]);
    tmux.liveSessions = () => new Set();
    try {
      const s = share.forSpawn(member, "agenthub-lane-7");
      const key = s.env.AGENTHUB_OMP_KEY;
      assert.match(key, /^ahomp_/);
      assert.equal(s.env.OMP_SKIP_SETUP, "1");
      const hash = sha(key);
      assert.ok(
        Object.values(listedBefore).some((c) => c.tokenSha256 === hash),
        "the key was listed before the pane started"
      );
      const after = JSON.parse(fs.readFileSync(clientsFile, "utf8"));
      assert.deepEqual(after["agenthub-lane-7"], { tokenSha256: hash });
      // its place in the queue is filled again, and no plaintext reaches the clients file
      assert.equal(Object.keys(after).filter((k) => k.startsWith("standby-")).length, 4);
      assert.ok(!fs.readFileSync(clientsFile, "utf8").includes("ahomp_"));
      assert.equal((fs.statSync(clientsFile).mode & 0o777).toString(8), "600");
      s.place();
      const agentDir = path.join(member.home, ".omp", "agent");
      assert.deepEqual(writes[0], ["mkdir", "ah-mia", agentDir]);
      assert.equal(writes[1][0], "replace");
      assert.equal(writes[1][2], path.join(agentDir, "models.yml"));
      assert.equal(yaml.load(writes[1][3]).providers.agenthub.apiKey, "AGENTHUB_OMP_KEY");
    } finally {
      runAs.readFileSafe = real.read;
      runAs.mkdirAs = real.mkdir;
      runAs.writeAs = real.write;
    }
  });
});

describe("pruneTokens", () => {
  const now = 1_000_000;
  const tokens = () => ({
    standby: [],
    used: {
      "agenthub-lane-21": { hash: "b", at: now - 120_000 }, // live
      "agenthub-lane-22": { hash: "c", at: now - 120_000 }, // pane gone
      "agenthub-lane-23": { hash: "d", at: now - 5_000 }, // just started
    },
  });

  it("drops only a gone pane's key past the grace, and keeps the queue full", () => {
    const t = share.pruneTokens(tokens(), new Set(["agenthub-lane-21"]), now);
    assert.deepEqual(Object.keys(t.used).sort(), ["agenthub-lane-21", "agenthub-lane-23"]);
    assert.equal(t.standby.length, 4);
  });

  it("drops nothing when tmux could not be asked", () => {
    const t = share.pruneTokens(tokens(), null, now);
    assert.equal(Object.keys(t.used).length, 3);
  });
});

describe("/api/ai-tools/omp-sharing", () => {
  let server;
  let BASE;
  before(async () => {
    server = await startServer(createApp(), 0);
    BASE = `http://127.0.0.1:${server.address().port}`;
  });
  after(() => server.close());

  it("reports it unavailable and refuses the switch outside S3", async () => {
    const got = await (await fetch(`${BASE}/api/ai-tools/omp-sharing`)).json();
    assert.equal(got.available, false);
    assert.equal(got.enabled, false);
    const put = await fetch(`${BASE}/api/ai-tools/omp-sharing`, {
      method: "PUT",
      headers: { "content-type": "application/json", origin: BASE },
      body: JSON.stringify({ enabled: true }),
    });
    assert.equal(put.status, 409);
    assert.equal((await put.json()).error.code, "ENOS3");
    assert.equal(fs.existsSync(flag), false);
  });

  it("wants a boolean", async () => {
    const put = await fetch(`${BASE}/api/ai-tools/omp-sharing`, {
      method: "PUT",
      headers: { "content-type": "application/json", origin: BASE },
      body: JSON.stringify({ enabled: "yes" }),
    });
    assert.equal(put.status, 400);
  });
});