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

/**
 * @file Tests server/lib/antigravity-usage.js — Antigravity CLI sessions and
 * tokens from agy's per-conversation SQLite files: the protobuf reader, a
 * session created as `tool = 'antigravity'` with its name, cwd, model and
 * per-model tokens (input excluding cache reads, output including thinking),
 * an unchanged file never re-read, a quiet one settled to completed (and the
 * lane a run with no Stop left running set idle, unless a pane owns it), a
 * session the user deleted staying deleted, a non-antigravity row with the
 * same id never touched, the empty `-wal` of our own read-only open not being
 * activity, a hook-reported turn kept running through a long quiet tool call,
 * an error Stop's status kept, and a file that is not agy's read only once. The Gemini dir and the DB are temp dirs; the fixture
 * files are built in the shape agy 1.2.13 writes.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const os = require("node:os");
const fs = require("node:fs");
const path = require("node:path");

const ROOT = fs.mkdtempSync(path.join(os.tmpdir(), "agenthub-agy-usage-"));
process.env.DASHBOARD_DB_PATH = path.join(ROOT, "dashboard.db");
process.env.AGENTHUB_GEMINI_DIR = path.join(ROOT, "gemini");

const { describe, it, beforeEach, after } = require("node:test");
const assert = require("node:assert/strict");
const Database = require("better-sqlite3");
const { db } = require("../db");
const agy = require("../lib/antigravity-usage");
const lanes = require("../lib/lanes");

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

const CLI = path.join(process.env.AGENTHUB_GEMINI_DIR, "antigravity-cli");
const CONV = path.join(CLI, "conversations");

// --- a minimal protobuf writer, enough for the fixtures ---
function varint(n) {
  const out = [];
  while (n >= 0x80) {
    out.push((n % 128) | 0x80);
    n = Math.floor(n / 128);
  }
  out.push(n);
  return Buffer.from(out);
}
const pb = (num, v) =>
  typeof v === "number"
    ? Buffer.concat([varint(num * 8), varint(v)])
    : Buffer.concat([varint(num * 8 + 2), varint(v.length), Buffer.from(v)]);
const msg = (...parts) => Buffer.concat(parts);

function generation({ model, input, output, cacheRead = 0, cacheWrite = 0 }) {
  const stats = msg(
    pb(1, 1318),
    pb(2, input),
    pb(3, output),
    ...(cacheWrite ? [pb(4, cacheWrite)] : []),
    pb(5, cacheRead),
    pb(7, "bot-1"),
    pb(9, Math.floor(output / 2))
  );
  return msg(pb(2, ""), pb(4, "req"), pb(1, msg(pb(3, 1318), pb(4, stats), pb(19, model))));
}

function conversation(
  id,
  { cwd = "/work/app", gens = [], prompt = "Fix the login bug\nplease", wal = false } = {}
) {
  fs.mkdirSync(CONV, { recursive: true });
  const file = path.join(CONV, `${id}.db`);
  fs.rmSync(file, { force: true });
  const c = new Database(file);
  c.exec(`CREATE TABLE trajectory_meta (trajectory_id text, cascade_id text, trajectory_type integer, source integer, PRIMARY KEY (trajectory_id));
    CREATE TABLE gen_metadata (idx integer, data blob, size integer NOT NULL DEFAULT 0, PRIMARY KEY (idx));
    CREATE TABLE trajectory_metadata_blob (id text DEFAULT "main", data blob, PRIMARY KEY (id));`);
  c.prepare("INSERT INTO trajectory_meta VALUES ('t1', ?, 4, 17)").run(id);
  c.prepare("INSERT INTO trajectory_metadata_blob VALUES ('main', ?)").run(
    msg(pb(1, msg(pb(2, `file://${cwd.replace(/ /g, "%20")}`))))
  );
  gens.forEach((g, i) =>
    c.prepare("INSERT INTO gen_metadata VALUES (?, ?, 0)").run(i, generation(g))
  );
  // agy's files are WAL; a clean close deletes the -wal, as agy's exit does.
  if (wal) c.pragma("journal_mode = WAL");
  c.close();
  const logs = path.join(CLI, "brain", id, ".system_generated", "logs");
  fs.mkdirSync(logs, { recursive: true });
  fs.writeFileSync(
    path.join(logs, "transcript_full.jsonl"),
    JSON.stringify({
      step_index: 0,
      source: "USER_EXPLICIT",
      type: "USER_INPUT",
      status: "DONE",
      created_at: "2026-09-29T04:38:44Z",
      content: `<USER_REQUEST>\n${prompt}\n</USER_REQUEST>\n<ADDITIONAL_METADATA>x</ADDITIONAL_METADATA>`,
    }) + "\n"
  );
  return file;
}

function age(id, ms) {
  const t = new Date(Date.now() - ms);
  fs.utimesSync(path.join(CONV, `${id}.db`), t, t);
  fs.utimesSync(
    path.join(CLI, "brain", id, ".system_generated", "logs", "transcript_full.jsonl"),
    t,
    t
  );
}

const session = (id) => db.prepare("SELECT * FROM sessions WHERE id = ?").get(id);
const usage = (id) =>
  db
    .prepare(
      "SELECT model, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens FROM token_usage WHERE session_id = ? ORDER BY model"
    )
    .all(id);

const ID = "57f3e409-9608-422b-8040-503f1bc7b5b8";

beforeEach(() => {
  fs.rmSync(process.env.AGENTHUB_GEMINI_DIR, { recursive: true, force: true });
  db.exec(
    "DELETE FROM events; DELETE FROM agents; DELETE FROM token_usage; DELETE FROM sessions; DELETE FROM antigravity_conversations;"
  );
});

const hookEvent = (id, type, msAgo) =>
  db
    .prepare("INSERT INTO events (session_id, event_type, created_at) VALUES (?, ?, ?)")
    .run(id, type, new Date(Date.now() - msAgo).toISOString());

describe("antigravity usage", () => {
  it("reads a generation's usage and model, and refuses a malformed buffer", () => {
    const u = agy.generationUsage(
      generation({ model: "gemini-3.8-flash", input: 12271, output: 153, cacheRead: 40 })
    );
    assert.deepEqual(u, {
      model: "gemini-3.8-flash",
      input: 12271,
      output: 153,
      cacheWrite: 0,
      cacheRead: 40,
    });
    assert.equal(agy.generationUsage(Buffer.from([0x0a, 0x7f, 1])), null);
    assert.equal(agy.generationUsage(msg(pb(4, "x"))), null);
    assert.equal(agy.firstFileUri(msg(pb(1, msg(pb(2, "file:///a%20b/c"))))), "/a b/c");
  });

  it("creates the session with its name, cwd, model and per-model tokens", async () => {
    conversation(ID, {
      cwd: "/work/my app",
      gens: [
        { model: "gemini-3.8-flash", input: 100, output: 10, cacheRead: 1000 },
        { model: "gemini-3.8-flash", input: 50, output: 5, cacheWrite: 7 },
        { model: "gemini-3.1-pro", input: 1, output: 2 },
      ],
    });
    const r = await agy.ingestAntigravityUsage();
    assert.deepEqual(r, { files: 1, sessions: 1 });
    const s = session(ID);
    assert.equal(s.tool, "antigravity");
    assert.equal(s.name, "Fix the login bug");
    assert.equal(s.cwd, "/work/my app");
    assert.equal(s.model, "gemini-3.1-pro");
    assert.equal(s.started_at, "2026-09-29T04:38:44Z");
    assert.equal(s.status, "active");
    assert.equal(s.transcript_path, null);
    assert.deepEqual(usage(ID), [
      {
        model: "gemini-3.1-pro",
        input_tokens: 1,
        output_tokens: 2,
        cache_read_tokens: 0,
        cache_write_tokens: 0,
      },
      {
        model: "gemini-3.8-flash",
        input_tokens: 150,
        output_tokens: 15,
        cache_read_tokens: 1000,
        cache_write_tokens: 7,
      },
    ]);
    assert.deepEqual(await agy.ingestAntigravityUsage(), { files: 0, sessions: 0 });
  });

  it("settles a session to completed once its file has gone quiet", async () => {
    conversation(ID, { gens: [{ model: "gemini-3.8-flash", input: 1, output: 1 }] });
    await agy.ingestAntigravityUsage();
    assert.equal(session(ID).status, "active");
    // Unchanged since the last read, just older: no re-read, only the status.
    const seen = db.prepare("SELECT mtime_ms FROM antigravity_conversations").get().mtime_ms;
    const r = await agy.ingestAntigravityUsage({ now: seen + 5 * 60_000 });
    assert.deepEqual(r, { files: 0, sessions: 1 });
    assert.equal(session(ID).status, "completed");
    // Read fresh when already old: completed straight away, ended at its last write.
    db.exec("DELETE FROM sessions; DELETE FROM antigravity_conversations;");
    age(ID, 10 * 60_000);
    await agy.ingestAntigravityUsage();
    assert.equal(session(ID).status, "completed");
    assert.ok(session(ID).ended_at);
  });

  it("sets idle the lane a run with no Stop left running, once the session is quiet", async () => {
    const cwd = fs.mkdtempSync(path.join(ROOT, "lane-"));
    const lane = lanes.createLane({ title: "a", cwd, tool: "antigravity" });
    lanes.updateLane(lane.id, { status: "running", session_id: ID });
    const other = "11111111-2222-3333-4444-555555555555";
    const cwd2 = fs.mkdtempSync(path.join(ROOT, "lane-"));
    const paned = lanes.createLane({ title: "b", cwd: cwd2, tool: "antigravity" });
    lanes.updateLane(paned.id, { status: "running", session_id: other, run_id: "agenthub-lane-9" });
    conversation(ID, { gens: [{ model: "gemini-3.8-flash", input: 1, output: 1 }] });
    conversation(other, { gens: [{ model: "gemini-3.8-flash", input: 1, output: 1 }] });
    const released = [];
    const broadcastLane = (id) => released.push(id);
    await agy.ingestAntigravityUsage({ broadcastLane });
    assert.equal(lanes.getLane(lane.id).status, "running"); // still being written
    const seen = db.prepare("SELECT max(mtime_ms) m FROM antigravity_conversations").get().m;
    await agy.ingestAntigravityUsage({ broadcastLane, now: seen + 5 * 60_000 });
    assert.equal(lanes.getLane(lane.id).status, "idle");
    assert.equal(lanes.getLane(paned.id).status, "running");
    assert.deepEqual(released, [lane.id]);
  });

  it("keeps a deleted session deleted until agy writes to it again", async () => {
    conversation(ID, { gens: [{ model: "gemini-3.8-flash", input: 1, output: 1 }] });
    age(ID, 10 * 60_000);
    await agy.ingestAntigravityUsage();
    db.prepare("DELETE FROM sessions WHERE id = ?").run(ID);
    await agy.ingestAntigravityUsage();
    assert.equal(session(ID), undefined);
    conversation(ID, { gens: [{ model: "gemini-3.8-flash", input: 2, output: 1 }] });
    await agy.ingestAntigravityUsage();
    assert.equal(session(ID).tool, "antigravity");
  });

  it("never touches a session of another tool with the same id", async () => {
    db.prepare(
      "INSERT INTO sessions (id, name, status, tool, started_at, updated_at) VALUES (?, 'mine', 'active', 'claude', '2026-09-29T00:00:00Z', '2026-09-29T00:00:00Z')"
    ).run(ID);
    conversation(ID, { gens: [{ model: "gemini-3.8-flash", input: 1, output: 1 }] });
    await agy.ingestAntigravityUsage();
    assert.equal(session(ID).name, "mine");
    assert.deepEqual(usage(ID), []);
  });

  it("keeps a hook-made session's cwd and a name the user saw, over a placeholder", async () => {
    db.prepare(
      "INSERT INTO sessions (id, name, status, tool, cwd, started_at, updated_at) VALUES (?, ?, 'active', 'antigravity', '/from/hook', '2026-09-29T00:00:00Z', '2026-09-29T00:00:00Z')"
    ).run(ID, `Antigravity ${ID.slice(0, 8)}`);
    conversation(ID, {
      cwd: "/from/db",
      gens: [{ model: "gemini-3.8-flash", input: 1, output: 1 }],
    });
    await agy.ingestAntigravityUsage();
    assert.equal(session(ID).cwd, "/from/hook");
    assert.equal(session(ID).name, "Fix the login bug");
  });

  it("skips a file that is not a conversation database, and carries on", async () => {
    fs.mkdirSync(CONV, { recursive: true });
    fs.writeFileSync(path.join(CONV, "00000000-0000-0000-0000-000000000000.db"), "not sqlite");
    conversation(ID, { gens: [{ model: "gemini-3.8-flash", input: 1, output: 1 }] });
    const r = await agy.ingestAntigravityUsage();
    assert.equal(r.sessions, 1);
    assert.ok(session(ID));
    // Recorded, so it is not re-read (nor warned about) every pass.
    assert.deepEqual(await agy.ingestAntigravityUsage(), { files: 0, sessions: 0 });
  });

  it("does not take the empty -wal its own read-only open leaves for activity", async () => {
    conversation(ID, { wal: true, gens: [{ model: "gemini-3.8-flash", input: 1, output: 1 }] });
    age(ID, 10 * 60_000);
    await agy.ingestAntigravityUsage();
    const first = session(ID);
    assert.equal(first.status, "completed");
    assert.deepEqual(await agy.ingestAntigravityUsage(), { files: 0, sessions: 0 });
    assert.equal(session(ID).status, "completed");
    assert.equal(session(ID).updated_at, first.updated_at);
  });

  it("keeps a turn the hooks report running through a long quiet tool call", async () => {
    const cwd = fs.mkdtempSync(path.join(ROOT, "lane-"));
    const lane = lanes.createLane({ title: "h", cwd, tool: "antigravity" });
    conversation(ID, { gens: [{ model: "gemini-3.8-flash", input: 1, output: 1 }] });
    age(ID, 10 * 60_000);
    await agy.ingestAntigravityUsage();
    // A hook three minutes ago: agy is mid-turn, its file untouched for ten.
    db.prepare("UPDATE sessions SET status = 'active', ended_at = NULL WHERE id = ?").run(ID);
    db.prepare(
      "INSERT INTO agents (id, session_id, name, type, status, started_at) VALUES (?, ?, 'Main', 'main', 'working', '2026-09-29T00:00:00Z')"
    ).run(`${ID}-main`, ID);
    lanes.updateLane(lane.id, { status: "running", session_id: ID });
    hookEvent(ID, "PostToolUse", 3 * 60_000);
    await agy.ingestAntigravityUsage();
    assert.equal(session(ID).status, "active");
    assert.equal(lanes.getLane(lane.id).status, "running");
    // Ten minutes after that hook, it is over.
    await agy.ingestAntigravityUsage({ now: Date.now() + 8 * 60_000 });
    assert.equal(session(ID).status, "completed");
    assert.equal(lanes.getLane(lane.id).status, "idle");
    assert.equal(
      db.prepare("SELECT status FROM agents WHERE id = ?").get(`${ID}-main`).status,
      "completed"
    );
  });

  it("keeps a session waiting after a Stop active for ten quiet minutes", async () => {
    conversation(ID, { gens: [{ model: "gemini-3.8-flash", input: 1, output: 1 }] });
    age(ID, 10 * 60_000);
    await agy.ingestAntigravityUsage();
    db.prepare("UPDATE sessions SET status = 'active', ended_at = NULL WHERE id = ?").run(ID);
    hookEvent(ID, "Stop", 3 * 60_000);
    await agy.ingestAntigravityUsage();
    assert.equal(session(ID).status, "active");
    await agy.ingestAntigravityUsage({ now: Date.now() + 8 * 60_000 });
    assert.equal(session(ID).status, "completed");
  });

  it("gives a file it could not read another try after a restart", async () => {
    fs.mkdirSync(CONV, { recursive: true });
    const bad = path.join(CONV, "00000000-0000-0000-0000-000000000000.db");
    fs.writeFileSync(bad, "not sqlite");
    await agy.ingestAntigravityUsage();
    const seen = () => db.prepare("SELECT session_id FROM antigravity_conversations").all();
    assert.deepEqual(seen(), [{ session_id: "" }]);
    const saved = process.env.DASHBOARD_ANTIGRAVITY_POLL_MS;
    process.env.DASHBOARD_ANTIGRAVITY_POLL_MS = "3600000";
    try {
      agy.startAntigravityUsagePoll(
        () => {},
        () => {}
      );
      // Cleared before the first pass reads it again.
      assert.ok(seen().length <= 1);
      await agy.ingestAntigravityUsage();
    } finally {
      if (saved === undefined) delete process.env.DASHBOARD_ANTIGRAVITY_POLL_MS;
      else process.env.DASHBOARD_ANTIGRAVITY_POLL_MS = saved;
    }
  });

  it("keeps the error a failed run's Stop gave the session", async () => {
    conversation(ID, { gens: [{ model: "gemini-3.8-flash", input: 1, output: 1 }] });
    await agy.ingestAntigravityUsage();
    db.prepare("UPDATE sessions SET status = 'error' WHERE id = ?").run(ID);
    hookEvent(ID, "Stop", 1000);
    conversation(ID, { gens: [{ model: "gemini-3.8-flash", input: 2, output: 1 }] });
    await agy.ingestAntigravityUsage();
    assert.equal(session(ID).status, "error");
    assert.equal(usage(ID)[0].input_tokens, 2);
  });
});