AgentHub/server/__tests__/ai-tool-models.test.js

/**
 * @file Tests for server/lib/ai-tool-models.js and the model plumbing around
 * it: each tool's model list comes from that tool's own source (Codex's
 * models cache, Claude's aliases named after the newest priced version, a
 * catalog tool's vendor rows plus the models its sessions used), a model id is
 * validated before it reaches an argv, the catalog `model` field is checked,
 * and the per-tool default model persists.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const { describe, it, before, after } = require("node:test");
const assert = require("node:assert/strict");
const fs = require("node:fs");
const os = require("node:os");
const path = require("node:path");

const root = fs.mkdtempSync(path.join(os.tmpdir(), "ai-tool-models-"));
const builtin = path.join(root, "builtin");
const codexHome = path.join(root, "codex");
const data = path.join(root, "data");
for (const d of [builtin, codexHome, data]) fs.mkdirSync(d);

const saved = {
  data: process.env.DASHBOARD_DATA_DIR,
  codex: process.env.CODEX_HOME,
  db: process.env.DASHBOARD_DB_PATH,
};
process.env.DASHBOARD_DATA_DIR = data;
process.env.CODEX_HOME = codexHome;
process.env.DASHBOARD_DB_PATH = path.join(root, "test.db");

const { db } = require("../db");
const aiTools = require("../lib/ai-tools");
const { listModels } = require("../lib/ai-tool-models");

const write = (name, doc) => fs.writeFileSync(path.join(builtin, name), JSON.stringify(doc));

before(() => {
  write("claude.json", {
    id: "claude",
    label: "Claude Code",
    detect: ["claude"],
    launcher: "claude",
  });
  write("codex.json", { id: "codex", label: "Codex", detect: ["codex"], launcher: "codex" });
  write("gemini.json", {
    id: "gemini",
    label: "Gemini",
    detect: ["gemini"],
    prompt: { mode: "argv" },
    model: { flag: "--model", vendor: "google" },
  });
  write("plain.json", { id: "plain", label: "Plain", detect: ["plain"], prompt: { mode: "argv" } });
  write("badflag.json", {
    id: "badflag",
    label: "Bad",
    detect: ["badflag"],
    prompt: { mode: "argv" },
    model: { flag: "rm -rf" },
  });
  aiTools.__setCatalogDirs([builtin, path.join(root, "none")]);

  const price = db.prepare(
    "INSERT OR REPLACE INTO model_pricing (model_pattern, display_name, input_per_mtok) VALUES (?, ?, 1)"
  );
  db.prepare("DELETE FROM model_pricing").run();
  price.run("claude-sonnet-5-5%", "Claude Sonnet 5.5");
  price.run("claude-sonnet-5%", "Claude Sonnet 5");
  price.run("claude-opus-5-5%", "Claude Opus 5.5");
  price.run("gemini-3.8-flash%", "Gemini 3.8 Flash");
  price.run("gpt-6-sol%", "GPT-6 Sol");
  db.prepare(
    "INSERT INTO sessions (id, status, started_at, tool) VALUES ('g1', 'completed', ?, 'gemini')"
  ).run(new Date().toISOString());
  db.prepare("INSERT INTO token_usage (session_id, model) VALUES ('g1', 'gemini-3.1-pro')").run();
});

after(() => {
  aiTools.__setCatalogDirs(null);
  for (const [k, env] of [
    ["data", "DASHBOARD_DATA_DIR"],
    ["codex", "CODEX_HOME"],
    ["db", "DASHBOARD_DB_PATH"],
  ]) {
    if (saved[k] === undefined) delete process.env[env];
    else process.env[env] = saved[k];
  }
  fs.rmSync(root, { recursive: true, force: true });
});

describe("listModels", () => {
  it("reads Codex's own models cache, listed ones only, and nothing else from it", () => {
    fs.writeFileSync(
      path.join(codexHome, "models_cache.json"),
      JSON.stringify({
        identity: { email: "someone@example.com" },
        models: [
          { slug: "gpt-5.5", display_name: "GPT-5.5", visibility: "list", priority: 12 },
          { slug: "gpt-hidden", display_name: "Hidden", visibility: "hide", priority: 1 },
          {
            slug: "gpt-6-sol",
            display_name: "GPT-6-Sol",
            description: "Top",
            visibility: "list",
            priority: 2,
          },
        ],
      })
    );
    const out = listModels("codex", db);
    assert.equal(out.flag, true);
    assert.deepEqual(
      out.models.map((m) => m.id),
      ["gpt-6-sol", "gpt-5.5"]
    );
    assert.equal(out.models[0].label, "GPT-6-Sol");
    assert.equal(out.models[0].hint, "Top");
    assert.doesNotMatch(JSON.stringify(out), /someone@example\.com/);
  });

  it("answers an empty Codex list when the cache is missing or unreadable", () => {
    fs.writeFileSync(path.join(codexHome, "models_cache.json"), "{nope");
    assert.deepEqual(listModels("codex", db).models, []);
  });

  it("names Claude's aliases after the newest priced version and lists the versions", () => {
    const out = listModels("claude", db);
    const sonnet = out.models.find((m) => m.id === "sonnet");
    assert.match(sonnet.label, /5\.5/);
    const ids = out.models.map((m) => m.id);
    assert.ok(ids.includes("claude-sonnet-5-5"));
    assert.ok(ids.includes("claude-opus-5-5"));
    assert.ok(ids.indexOf("claude-sonnet-5-5") < ids.indexOf("claude-sonnet-5"));
  });

  it("gives a catalog tool its vendor's priced models plus those its sessions used", () => {
    const out = listModels("gemini", db);
    assert.equal(out.flag, true);
    const ids = out.models.map((m) => m.id);
    assert.ok(ids.includes("gemini-3.8-flash"));
    assert.ok(ids.includes("gemini-3.1-pro"));
    assert.ok(!ids.includes("gpt-6-sol"));
  });

  it("offers no picker to a tool whose catalog declares no model flag", () => {
    assert.deepEqual(listModels("plain", db), { flag: false, models: [] });
    assert.equal(aiTools.has("badflag"), false);
  });
});

describe("model ids and the default model", () => {
  it("accepts real model ids and refuses anything flag- or shell-shaped", () => {
    for (const ok of [
      "opus",
      "opus[1m]",
      "claude-sonnet-5-5",
      "openai/gpt-5.5",
      "gemini-3.1-pro",
    ]) {
      assert.equal(aiTools.validModel(ok), true, ok);
    }
    for (const bad of [
      "-p",
      "--dangerously-skip-permissions",
      "a b",
      "x;rm",
      "",
      "a".repeat(200),
    ]) {
      assert.equal(aiTools.validModel(bad), false, bad);
    }
  });

  it("persists a default model per tool next to enable/default, and clears it", () => {
    aiTools.setDefaultModel("gemini", "gemini-3.8-flash");
    aiTools.setEnabled("gemini", false);
    aiTools.setEnabled("gemini", true);
    assert.equal(
      aiTools.view().tools.find((t) => t.id === "gemini").defaultModel,
      "gemini-3.8-flash"
    );
    assert.equal(aiTools.view().tools.find((t) => t.id === "gemini").modelFlag, true);
    assert.throws(() => aiTools.setDefaultModel("gemini", "--x"), { code: "EBADMODEL" });
    aiTools.setDefaultModel("gemini", null);
    assert.equal(aiTools.view().tools.find((t) => t.id === "gemini").defaultModel, null);
  });
});