AgentHub/server/__tests__/budgets.test.js

/**
 * @file Tests for spend budgets and dollar alerts: period bounds in local time,
 * the `budget` and `cost_threshold` rule types (validation, firing once per
 * threshold per period), GET /api/budgets/status and the read-only rule
 * preview POST /api/alerts/rules/preview.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */
"use strict";

const { describe, it, before, after, beforeEach } = 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 http = require("node:http");

const ROOT = fs.realpathSync(fs.mkdtempSync(path.join(os.tmpdir(), "budgets-")));
process.env.DASHBOARD_DB_PATH = path.join(ROOT, "t.db");
process.env.CLAUDE_HOME = path.join(ROOT, "claude");
process.env.DASHBOARD_DATA_DIR = path.join(ROOT, "data");
process.env.DASHBOARD_REMOTE_SYNC_MS = "0";
process.env.DASHBOARD_LIVENESS_PROBE = "0";

const { createApp } = require("../index");
const { db, stmts } = require("../db");
const { periodBounds, budgetSpend } = require("../lib/budgets");
const alerts = require("../lib/alerts");

const M = 1_000_000;
let server;
let base;

function call(method, url, body) {
  return new Promise((resolve, reject) => {
    const req = http.request(
      `${base}${url}`,
      { method, headers: body ? { "content-type": "application/json" } : {} },
      (res) => {
        let text = "";
        res.on("data", (c) => (text += c));
        res.on("end", () => resolve({ status: res.statusCode, body: JSON.parse(text) }));
      }
    );
    req.on("error", reject);
    if (body) req.write(JSON.stringify(body));
    req.end();
  });
}

const iso = (ms) => new Date(ms).toISOString();
function session(id, tool, startedAt, updatedAt = startedAt) {
  db.prepare(
    "INSERT INTO sessions (id, name, status, started_at, updated_at, tool) VALUES (?, ?, 'completed', ?, ?, ?)"
  ).run(id, id, startedAt, updatedAt, tool);
}
const spend = (id, input) =>
  stmts.setTokenUsage.run(
    id,
    "budget-model",
    "standard",
    "global",
    "standard",
    input,
    0,
    0,
    0,
    0,
    0,
    0,
    0
  );

before(async () => {
  server = http.createServer(createApp());
  await new Promise((r) => server.listen(0, "127.0.0.1", r));
  base = `http://127.0.0.1:${server.address().port}`;
  db.prepare(
    "INSERT OR REPLACE INTO model_pricing (model_pattern, display_name, input_per_mtok, output_per_mtok, cache_read_per_mtok, cache_write_per_mtok) VALUES ('budget-model%', 'B', 1, 0, 0, 0)"
  ).run();
});
after(() => {
  server.close();
  fs.rmSync(ROOT, { recursive: true, force: true });
});
beforeEach(() => {
  db.exec(
    "DELETE FROM alert_events; DELETE FROM alert_rules; DELETE FROM token_usage; DELETE FROM sessions;"
  );
  alerts.invalidateRuleCache();
});

describe("periodBounds", () => {
  // 2026-10-05 20:00 UTC = 2026-10-06 03:00 in UTC+7 (tz_offset -420), a Tuesday.
  const now = Date.parse("2026-10-05T20:00:00Z");
  it("starts a day at local midnight", () => {
    const b = periodBounds("day", now, -420);
    assert.equal(b.start, "2026-10-06");
    assert.equal(b.end, "2026-10-06");
    assert.equal(b.startUtc, "2026-10-05T17:00:00.000Z");
    assert.equal(b.days, 1);
    assert.ok(Math.abs(b.elapsed - 3 / 24) < 1e-9);
  });
  it("starts a week on Monday and a month on the 1st", () => {
    assert.equal(periodBounds("week", now, -420).start, "2026-10-05");
    assert.equal(periodBounds("week", now, -420).end, "2026-10-11");
    const m = periodBounds("month", now, -420);
    assert.equal(m.start, "2026-10-01");
    assert.equal(m.end, "2026-10-31");
    assert.equal(m.days, 31);
  });
  it("walks back whole periods, which count as fully elapsed", () => {
    const m = periodBounds("month", now, 0, 1);
    assert.equal(m.start, "2026-09-01");
    assert.equal(m.end, "2026-09-30");
    assert.equal(m.elapsed, 1);
  });
});

describe("budget rules", () => {
  it("validates the config", async () => {
    const bad = await call("POST", "/api/alerts/rules", {
      name: "b",
      rule_type: "budget",
      config: { period: "year", limit_usd: 5 },
    });
    assert.equal(bad.status, 400);
    const ok = await call("POST", "/api/alerts/rules", {
      name: "b",
      rule_type: "budget",
      config: { period: "day", limit_usd: 5, tool: "codex", tz_offset: -420 },
    });
    assert.equal(ok.status, 201);
    assert.deepEqual(ok.body.rule.config, {
      period: "day",
      limit_usd: 5,
      thresholds: [80, 100],
      tool: "codex",
      tz_offset: -420,
    });
  });

  it("measures the period's spend in scope and projects it", () => {
    const now = Date.now();
    session("today", "codex", iso(now - 60_000));
    spend("today", 2 * M); // $2
    session("other-tool", "claude", iso(now - 60_000));
    spend("other-tool", 9 * M);
    session("yesterday", "codex", iso(now - 3 * 86_400_000));
    spend("yesterday", 7 * M);
    const s = budgetSpend(
      { period: "day", limit_usd: 4, thresholds: [50], tool: "codex" },
      { now }
    );
    assert.equal(s.spend, 2);
    assert.equal(s.pct, 50);
    assert.ok(s.projected >= 2);
  });

  it("fires each crossed threshold once per period, and the status lists it", async () => {
    const now = Date.now();
    session("s1", "claude", iso(now - 60_000));
    spend("s1", 9 * M); // $9 of $10
    const created = await call("POST", "/api/alerts/rules", {
      name: "Daily",
      rule_type: "budget",
      config: { period: "day", limit_usd: 10, thresholds: [50, 80, 100] },
    });
    alerts.sweepTimeRules();
    alerts.sweepTimeRules();
    const fired = db.prepare("SELECT details FROM alert_events ORDER BY id").all();
    assert.deepEqual(
      fired.map((f) => JSON.parse(f.details).threshold),
      [50, 80]
    );
    const status = await call("GET", "/api/budgets/status");
    const b = status.body.budgets.find((x) => x.id === created.body.rule.id);
    assert.equal(b.spend, 9);
    assert.equal(b.pct, 90);
    assert.deepEqual(b.crossed, [50, 80]);
  });
});

describe("cost_threshold", () => {
  it("fires when a session's cost reaches the threshold", async () => {
    session("pricey", "claude", iso(Date.now()));
    spend("pricey", 3 * M);
    await call("POST", "/api/alerts/rules", {
      name: "Expensive",
      rule_type: "cost_threshold",
      config: { usd: 2.5 },
    });
    alerts.evaluateEvent({ session_id: "pricey", event_type: "Stop" });
    const a = db.prepare("SELECT message, details FROM alert_events").get();
    assert.match(a.message, /\$3\.00/);
    assert.equal(JSON.parse(a.details).cost_usd, 3);
  });
});

describe("rule preview", () => {
  it("replays an event pattern with its count window and cooldown, firing nothing", async () => {
    session("p1", "claude", iso(Date.now() - 3_600_000));
    const t0 = Date.now() - 30 * 60_000;
    for (const [i, m] of [0, 1, 2, 20].entries()) {
      db.prepare(
        "INSERT INTO events (session_id, event_type, tool_name, summary, created_at) VALUES ('p1', 'PostToolUse', 'Bash', ?, ?)"
      ).run(`e${i}`, iso(t0 + m * 60_000));
    }
    const r = await call("POST", "/api/alerts/rules/preview", {
      rule_type: "event_pattern",
      config: { tool_name: "Bash", count: 2, window_minutes: 5 },
      cooldown_seconds: 300,
    });
    assert.equal(r.status, 200);
    // e1 (2 in window) fires; e2 is inside the cooldown; e20 is alone in its window.
    assert.equal(r.body.fires, 1);
    assert.equal(r.body.sessions, 1);
    assert.equal(db.prepare("SELECT COUNT(*) c FROM alert_events").get().c, 0);
  });
  it("lists past budget periods and refuses what cannot be replayed", async () => {
    const b = await call("POST", "/api/alerts/rules/preview", {
      rule_type: "budget",
      config: { period: "week", limit_usd: 1 },
    });
    assert.equal(b.body.periods.length, 6);
    const s = await call("POST", "/api/alerts/rules/preview", {
      rule_type: "status_duration",
      config: { status: "working", minutes: 5 },
    });
    assert.equal(s.body.supported, false);
    const bad = await call("POST", "/api/alerts/rules/preview", {
      rule_type: "budget",
      config: { period: "week", limit_usd: 1 },
      days: 0,
    });
    assert.equal(bad.status, 400);
  });
});