AgentHub/server/__tests__/server-mode-aggregates.test.js

/**
 * @file server-mode-aggregates.test.js
 * @description Server mode's row ownership on the aggregate endpoints: GET
 * /api/stats, /api/analytics, /api/query (pages and CSV export),
 * /api/pricing/cost, /api/budgets/status and /api/ai-tools/:id/usage count,
 * total and price only a member's own sessions, while an admin, an analyst and
 * local mode keep seeing every session.
 * @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 path = require("node:path");
const fs = require("node:fs");
const os = require("node:os");

const TMP = fs.mkdtempSync(path.join(os.tmpdir(), "server-mode-aggregates-"));
process.env.DASHBOARD_DB_PATH = path.join(TMP, "dashboard.db");
process.env.DASHBOARD_LIVENESS_PROBE = "0";
process.env.AGENTHUB_PRICING_SYNC = "0";

const { createApp, startServer } = require("../index");
const { db, stmts } = require("../db");
const auth = require("../lib/auth");

let server;
let BASE;
const cookies = {};
const lanes = {};

async function get(p, cookie) {
  const res = await fetch(BASE + p, { headers: cookie ? { cookie } : {} });
  const text = await res.text();
  assert.equal(res.status, 200, `${p}: ${text}`);
  return res.headers.get("content-type").includes("json") ? JSON.parse(text) : text;
}

async function login(username, password) {
  const res = await fetch(BASE + "/api/auth/login", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ username, password }),
  });
  assert.equal(res.status, 200);
  return res.headers.get("set-cookie").split(";")[0];
}

// Each session spends a distinct number of MTok at $1 each, so every total
// says exactly which sessions it counted.
const SESSIONS = [
  // id, tool, stamped lane, MTok
  ["agg-a", "claude", "mia", 1],
  ["agg-b", "codex", "max", 2],
  ["agg-host", "claude", null, 4],
  ["agg-run", "claude", null, 8], // unstamped, held by a run of mia's workspace
];

before(async () => {
  server = await startServer(createApp(), 0);
  BASE = `http://127.0.0.1:${server.address().port}`;
  const mia = await auth.createUser("mia", "member-pass-1", "member");
  const max = await auth.createUser("max", "member-pass-2", "member");
  await auth.createUser("boss", "admin-pass-1", "admin");
  await auth.createUser("ana", "analyst-pass-1", "analyst");

  const lane = db.prepare("INSERT INTO lanes (cwd, tool, owner_id) VALUES (?, 'claude', ?)");
  lanes.mia = Number(lane.run("/agg/mia", mia.id).lastInsertRowid);
  lanes.max = Number(lane.run("/agg/max", max.id).lastInsertRowid);
  db.prepare(
    "INSERT INTO model_pricing (model_pattern, display_name, input_per_mtok, output_per_mtok, cache_read_per_mtok, cache_write_per_mtok) VALUES ('agg-model%', 'Agg', 1, 0, 0, 0)"
  ).run();
  for (const [id, tool, owner, mtok] of SESSIONS) {
    db.prepare(
      "INSERT INTO sessions (id, name, status, cwd, tool, lane_id, updated_at) VALUES (?, ?, 'active', '/agg', ?, ?, strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))"
    ).run(id, id, tool, owner ? lanes[owner] : null);
    db.prepare(
      "INSERT INTO token_usage (session_id, model, input_tokens) VALUES (?, 'agg-model', ?)"
    ).run(id, mtok * 1e6);
    db.prepare("INSERT INTO agents (id, session_id, name, type) VALUES (?, ?, 'Main', 'main')").run(
      `${id}-main`,
      id
    );
    db.prepare(
      "INSERT INTO events (session_id, event_type, tool_name, summary) VALUES (?, 'PreToolUse', 'Bash', 'x')"
    ).run(id);
  }
  db.prepare(
    "INSERT INTO dashboard_runs (id, session_id, mode, cwd, status, lane_id) VALUES (?, 'agg-run', 'interactive', '/agg', 'exited', ?)"
  ).run(`agenthub-lane-${lanes.mia}@old`, lanes.mia);
  stmts.insertAlertRule.run(
    "agg-budget",
    "Org budget",
    "budget",
    JSON.stringify({ period: "month", limit_usd: 100, thresholds: [50] }),
    1,
    0
  );

  process.env.AGENTHUB_MODE = "server";
  cookies.mia = await login("mia", "member-pass-1");
  cookies.max = await login("max", "member-pass-2");
  cookies.boss = await login("boss", "admin-pass-1");
  cookies.ana = await login("ana", "analyst-pass-1");
});

after(() => {
  delete process.env.AGENTHUB_MODE;
  server?.close();
  fs.rmSync(TMP, { recursive: true, force: true });
});

/** Every aggregate one viewer gets, reduced to the numbers that tell rows apart. */
async function numbers(cookie) {
  const stats = await get("/api/stats", cookie);
  const analytics = await get("/api/analytics", cookie);
  const ranged = await get("/api/analytics?from=2000-01-01&to=2999-12-31", cookie);
  const sessions = await get("/api/query?entity=sessions", cookie);
  const events = await get("/api/query?entity=events", cookie);
  const agents = await get("/api/query?entity=agents", cookie);
  const csv = await get("/api/query?entity=sessions&format=csv", cookie);
  const cost = await get("/api/pricing/cost", cookie);
  const claudeCost = await get("/api/pricing/cost?tool=claude", cookie);
  const budgets = await get("/api/budgets/status", cookie);
  const usage = await get("/api/ai-tools/claude/usage", cookie);
  return {
    stats: [stats.total_sessions, stats.total_agents, stats.total_events],
    analytics: [analytics.overview.total_sessions, analytics.tokens.total_input / 1e6],
    analyticsCost: analytics.total_cost,
    rangedCost: ranged.total_cost,
    query: [sessions.total, events.total, agents.total],
    queryIds: sessions.rows.map((r) => r.id).sort(),
    csvRows: csv.trim().split("\r\n").length - 1,
    cost: cost.total_cost,
    claudeCost: claudeCost.total_cost,
    budget: budgets.budgets.find((b) => b.id === "agg-budget").spend,
    usage: [usage.sessions, usage.cost],
  };
}

const EVERYONE = {
  stats: [4, 4, 4],
  analytics: [4, 15],
  analyticsCost: 15,
  rangedCost: 15,
  query: [4, 4, 4],
  queryIds: ["agg-a", "agg-b", "agg-host", "agg-run"],
  csvRows: 4,
  cost: 15,
  claudeCost: 13,
  budget: 15,
  usage: [3, 13],
};

describe("aggregates in server mode", () => {
  it("count a member's own sessions only, unstamped ones their runs hold included", async () => {
    assert.deepEqual(await numbers(cookies.mia), {
      stats: [2, 2, 2],
      analytics: [2, 9],
      analyticsCost: 9,
      rangedCost: 9,
      query: [2, 2, 2],
      queryIds: ["agg-a", "agg-run"],
      csvRows: 2,
      cost: 9,
      claudeCost: 9,
      budget: 9,
      usage: [2, 9],
    });
  });

  it("give another member theirs", async () => {
    assert.deepEqual(await numbers(cookies.max), {
      stats: [1, 1, 1],
      analytics: [1, 2],
      analyticsCost: 2,
      rangedCost: 2,
      query: [1, 1, 1],
      queryIds: ["agg-b"],
      csvRows: 1,
      cost: 2,
      claudeCost: 0,
      budget: 2,
      usage: [0, 0],
    });
  });

  it("show an admin and an analyst everything", async () => {
    assert.deepEqual(await numbers(cookies.boss), EVERYONE);
    assert.deepEqual(await numbers(cookies.ana), EVERYONE);
  });

  it("refuse a member who names another member's workspace", async () => {
    const res = await fetch(`${BASE}/api/pricing/cost?lane=${lanes.max}`, {
      headers: { cookie: cookies.mia },
    });
    assert.equal(res.status, 404);
  });
});

describe("aggregates in local mode", () => {
  it("are unchanged: every session, no login", async () => {
    delete process.env.AGENTHUB_MODE;
    try {
      assert.deepEqual(await numbers(), EVERYONE);
    } finally {
      process.env.AGENTHUB_MODE = "server";
    }
  });
});