AgentHub/server/__tests__/daily-cost.test.js

/**
 * @file The daily cost split of `GET /api/pricing/cost`: a session's tokens are
 * booked to the local days its `token_usage_hourly` rows fall on (the client's
 * `tz_offset` moving an hour across midnight), the rest of its total to its
 * start date, and the days add up to the total. Also `TranscriptCache`'s
 * `tokensByHour` — split by each message's timestamp, kept across an
 * incremental read — and the hook writer filling the table from it.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */
"use strict";

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 http = require("node:http");

const ROOT = fs.realpathSync(fs.mkdtempSync(path.join(os.tmpdir(), "daily-cost-")));
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 TranscriptCache = require("../lib/transcript-cache");

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();
  });
}

function session(id, startedAt) {
  db.prepare(
    "INSERT INTO sessions (id, name, status, started_at, updated_at) VALUES (?, ?, 'completed', ?, ?)"
  ).run(id, id, startedAt, startedAt);
}
const total = (id, input) =>
  stmts.setTokenUsage.run(
    id,
    "daily-model",
    "standard",
    "global",
    "standard",
    input,
    0,
    0,
    0,
    0,
    0,
    0,
    0
  );
const hour = (id, h, input) =>
  stmts.upsertHourlyUsage.run(
    id,
    h,
    "daily-model",
    "standard",
    "global",
    "standard",
    input,
    0,
    0,
    0,
    0,
    0,
    0,
    0
  );
const byDate = (r) => Object.fromEntries(r.body.daily_costs.map((d) => [d.date, d.cost]));
const line = (ts, id, input) =>
  JSON.stringify({
    type: "assistant",
    timestamp: ts,
    message: {
      id,
      model: "daily-model",
      role: "assistant",
      usage: { input_tokens: input, output_tokens: 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}`;
  // $1 per million input tokens, so a day's cost reads as millions of tokens.
  const put = await call("PUT", "/api/pricing", {
    model_pattern: "daily-model%",
    display_name: "Daily",
    input_per_mtok: 1,
    output_per_mtok: 0,
    cache_read_per_mtok: 0,
    cache_write_per_mtok: 0,
  });
  assert.equal(put.status, 200);
});

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

describe("GET /api/pricing/cost daily split", () => {
  it("books each hour to its local day and the rest to the start date", async () => {
    session("multi-day", "2026-10-02T03:00:00.000Z");
    total("multi-day", 7 * M); // 1 M recorded before hourly rows existed
    hour("multi-day", "2026-10-03T05", 2 * M);
    hour("multi-day", "2026-10-03T20", 1 * M); // 03:00 on the 4th at UTC+7
    hour("multi-day", "2026-10-05T01", 3 * M);

    const utc = byDate(await call("GET", "/api/pricing/cost?tz_offset=0"));
    assert.deepEqual(utc, { "2026-10-02": 1, "2026-10-03": 3, "2026-10-05": 3 });

    const ict = await call("GET", "/api/pricing/cost?tz_offset=-420");
    assert.deepEqual(byDate(ict), {
      "2026-10-02": 1,
      "2026-10-03": 2,
      "2026-10-04": 1,
      "2026-10-05": 3,
    });
    assert.equal(ict.body.total_cost, 7);
  });

  it("scales hours above the total down to it, so the days still add up", async () => {
    session("over", "2026-09-01T00:00:00.000Z");
    total("over", 2 * M);
    hour("over", "2026-09-02T00", 2 * M);
    hour("over", "2026-09-03T00", 2 * M);
    const r = await call("GET", "/api/pricing/cost?tz_offset=0");
    const d = byDate(r);
    assert.equal(d["2026-09-01"], undefined);
    assert.equal(d["2026-09-02"], 1);
    assert.equal(d["2026-09-03"], 1);
  });
});

describe("GET /api/pricing/cost?lane=", () => {
  it("prices only the sessions a workspace owns, including one its run holds unstamped", async () => {
    const lane = require("../lib/lanes").createLane({
      title: "w",
      cwd: fs.mkdtempSync(path.join(ROOT, "lane-")),
      tool: "claude",
    });
    session("owned", "2026-08-01T00:00:00.000Z");
    session("held", "2026-08-01T00:00:00.000Z");
    session("other", "2026-08-01T00:00:00.000Z");
    db.prepare("UPDATE sessions SET lane_id = ? WHERE id = 'owned'").run(lane.id);
    db.prepare(
      "INSERT INTO dashboard_runs (id, lane_id, session_id, cwd, status) VALUES ('run-held', ?, 'held', '/x', 'exited')"
    ).run(lane.id);
    total("owned", 2 * M);
    total("held", 1 * M);
    total("other", 5 * M);

    const r = await call("GET", `/api/pricing/cost?tz_offset=0&lane=${lane.id}`);
    assert.equal(r.status, 200);
    assert.equal(r.body.total_cost, 3);
    assert.equal((await call("GET", "/api/pricing/cost?lane=x")).status, 400);
  });
});

describe("TranscriptCache tokensByHour", () => {
  it("splits by message hour, counts a message once, and survives an incremental read", () => {
    const file = path.join(ROOT, "t.jsonl");
    fs.writeFileSync(
      file,
      [
        line("2026-10-02T10:15:00.000Z", "m1", 100),
        line("2026-10-02T10:15:01.000Z", "m1", 100), // a second block line of m1
        line("2026-10-02T10:40:00.000Z", "m2", 50),
      ].join("\n") + "\n"
    );
    const cache = new TranscriptCache();
    const hours = (r) =>
      Object.fromEntries(Object.values(r.tokensByHour).map((b) => [b.hour, b.input]));
    assert.deepEqual(hours(cache.extract(file)), { "2026-10-02T10": 150 });
    fs.appendFileSync(file, line("2026-10-04T23:59:00.000Z", "m3", 7) + "\n");
    assert.deepEqual(hours(cache.extract(file)), { "2026-10-02T10": 150, "2026-10-04T23": 7 });
  });

  it("is written by the hook writer", async () => {
    const file = path.join(ROOT, "hook.jsonl");
    fs.writeFileSync(
      file,
      [
        line("2026-10-02T10:00:00.000Z", "a", 1 * M),
        line("2026-10-03T10:00:00.000Z", "b", 2 * M),
      ].join("\n") + "\n"
    );
    const r = await call("POST", "/api/hooks/event", {
      hook_type: "Stop",
      data: { session_id: "hooked", transcript_path: file },
    });
    assert.equal(r.status, 200);
    const rows = db
      .prepare(
        "SELECT hour, input_tokens FROM token_usage_hourly WHERE session_id = 'hooked' ORDER BY hour"
      )
      .all();
    assert.deepEqual(
      rows.map((x) => [x.hour, x.input_tokens]),
      [
        ["2026-10-02T10", 1 * M],
        ["2026-10-03T10", 2 * M],
      ]
    );
  });
});