AgentHub/server/__tests__/subagent-usage.test.js

/**
 * @file Tests for counting every token a Claude session spent: the hook
 * writer adds the session's sub-agent transcripts (`lib/subagent-usage.js`)
 * to its main one — total and per hour, a sub-agent on the session's own
 * model included, and a sub-agent that grew alone is picked up — and the
 * importer raises a session it first saw with an empty transcript once the
 * transcript has grown (`importSession`).
 * @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("path");
const fs = require("fs");
const os = require("os");
const http = require("http");

const TMP = fs.mkdtempSync(path.join(os.tmpdir(), "agenthub-subusage-"));
process.env.DASHBOARD_DB_PATH = path.join(TMP, "dashboard.db");
process.env.CLAUDE_HOME = path.join(TMP, "home");
process.env.DASHBOARD_DATA_DIR = path.join(TMP, "data");

const { createApp, startServer } = require("../index");
const dbModule = require("../db");
const { db } = dbModule;
const importHistory = require("../../scripts/import-history");
const { subagentUsage } = require("../lib/subagent-usage");

let server;
let BASE;

function post(body) {
  return new Promise((resolve, reject) => {
    const url = new URL("/api/hooks/event", BASE);
    const payload = JSON.stringify(body);
    const r = http.request(
      {
        hostname: url.hostname,
        port: url.port,
        path: url.pathname,
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          "Content-Length": Buffer.byteLength(payload),
        },
      },
      (res) => {
        res.resume();
        res.on("end", () => resolve(res.statusCode));
      }
    );
    r.on("error", reject);
    r.write(payload);
    r.end();
  });
}

const line = (o) => JSON.stringify(o) + "\n";
const assistant = (id, inTok, ts, model = "claude-sonnet-5") => ({
  type: "assistant",
  timestamp: ts,
  message: {
    id,
    model,
    content: [{ type: "text", text: "ok" }],
    usage: { input_tokens: inTok, output_tokens: 1 },
  },
});

const total = (sid) =>
  db
    .prepare(
      "SELECT model, SUM(input_tokens + baseline_input) AS input FROM token_usage WHERE session_id = ? GROUP BY model ORDER BY model"
    )
    .all(sid)
    .map((r) => ({ ...r }));
const hours = (sid) =>
  db
    .prepare(
      "SELECT hour, model, SUM(input_tokens) AS input FROM token_usage_hourly WHERE session_id = ? GROUP BY hour, model ORDER BY hour, model"
    )
    .all(sid)
    .map((r) => ({ ...r }));

before(async () => {
  server = await startServer(createApp(), 0);
  BASE = `http://127.0.0.1:${server.address().port}`;
});

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

describe("sub-agent usage", () => {
  it("counts a session's sub-agents, own model included, in its total and its hours", async () => {
    const sid = "cdcdcdcd-1111-4111-8111-cdcdcdcdcdcd";
    const cwd = path.join(TMP, "work");
    fs.mkdirSync(cwd, { recursive: true });
    const tpath = path.join(TMP, `${sid}.jsonl`);
    fs.writeFileSync(tpath, line(assistant("m1", 100, "2026-10-22T01:00:00.000Z")));
    const subDir = path.join(TMP, sid, "subagents");
    fs.mkdirSync(subDir, { recursive: true });
    const sub1 = path.join(subDir, "agent-a1.jsonl");
    fs.writeFileSync(sub1, line(assistant("s1", 40, "2026-10-22T02:00:00.000Z")));
    fs.writeFileSync(
      path.join(subDir, "agent-a2.jsonl"),
      line(assistant("s2", 7, "2026-10-22T02:30:00.000Z", "claude-haiku-5-5"))
    );
    // A Workflow agent is booked by workflow-ingest under its own tier: not here.
    fs.mkdirSync(path.join(subDir, "workflows", "run1"), { recursive: true });
    fs.writeFileSync(
      path.join(subDir, "workflows", "run1", "agent-w.jsonl"),
      line(assistant("w1", 1000, "2026-10-22T02:00:00.000Z"))
    );
    const data = { session_id: sid, cwd, transcript_path: tpath };

    assert.equal(await post({ hook_type: "PostToolUse", data }), 200);
    assert.deepEqual(total(sid), [
      { model: "claude-haiku-5-5", input: 7 },
      { model: "claude-sonnet-5", input: 140 },
    ]);
    assert.deepEqual(hours(sid), [
      { hour: "2026-10-22T01", model: "claude-sonnet-5", input: 100 },
      { hour: "2026-10-22T02", model: "claude-haiku-5-5", input: 7 },
      { hour: "2026-10-22T02", model: "claude-sonnet-5", input: 40 },
    ]);

    // Only the sub-agent grew: the next hook still counts it, once.
    fs.appendFileSync(sub1, line(assistant("s3", 10, "2026-10-22T03:00:00.000Z")));
    assert.equal(await post({ hook_type: "PostToolUse", data: { ...data, tool_name: "x" } }), 200);
    assert.equal(await post({ hook_type: "PostToolUse", data: { ...data, tool_name: "y" } }), 200);
    assert.deepEqual(total(sid), [
      { model: "claude-haiku-5-5", input: 7 },
      { model: "claude-sonnet-5", input: 150 },
    ]);
    assert.equal(hours(sid).find((h) => h.hour === "2026-10-22T03").input, 10);

    // The importer's full re-read agrees and changes nothing (high-water).
    await importHistory.reconcileTokens(dbModule, { exact: false });
    assert.deepEqual(total(sid), [
      { model: "claude-haiku-5-5", input: 7 },
      { model: "claude-sonnet-5", input: 150 },
    ]);
  });

  it("is null for a session with no sub-agent transcript", () => {
    assert.equal(subagentUsage(path.join(TMP, "nope.jsonl")), null);
    assert.equal(subagentUsage(null), null);
  });

  it("raises an imported session first seen before its transcript had any usage", async () => {
    const sid = "efefefef-2222-4222-8222-efefefefefef";
    const proj = path.join(TMP, "projects", "-tmp-eval");
    fs.mkdirSync(proj, { recursive: true });
    const tpath = path.join(proj, `${sid}.jsonl`);
    const user = {
      type: "user",
      sessionId: sid,
      cwd: "/tmp/eval",
      timestamp: "2026-10-23T03:00:00.000Z",
      message: { role: "user", content: "fix the bug" },
    };
    fs.writeFileSync(tpath, line(user));
    const parse = async () => {
      const s = await importHistory.parseSessionFile(tpath);
      s.parsedSubagents = [];
      return s;
    };
    importHistory.importSession(dbModule, await parse());
    assert.deepEqual(total(sid), []);

    fs.appendFileSync(
      tpath,
      line({ ...assistant("e1", 500, "2026-10-23T03:01:00.000Z"), sessionId: sid })
    );
    const again = importHistory.importSession(dbModule, await parse());
    assert.equal(again.backfilled, true);
    assert.deepEqual(total(sid), [{ model: "claude-sonnet-5", input: 500 }]);
    // Nothing new: no backfill reported, no change.
    assert.equal(importHistory.importSession(dbModule, await parse()).backfilled, undefined);
    assert.deepEqual(total(sid), [{ model: "claude-sonnet-5", input: 500 }]);
  });
});