AgentHub/server/lib/antigravity-usage.js

/**
 * @file Antigravity CLI sessions and token usage, read from `agy`'s own
 * per-conversation SQLite files (`~/.gemini/antigravity-cli/conversations/
 * <conversation id>.db`). This is how an Antigravity session reaches
 * `sessions` (as `tool = 'antigravity'`) and its tokens reach `token_usage`;
 * AgentHub's Antigravity hooks, when the user installs them, add its events and
 * a lane's live status but no tokens.
 *
 * What a conversation file holds (agy 1.2.x, verified against the binary's own
 * proto descriptors — nothing here is a documented format):
 * - its file name is the conversation id the hooks report;
 * - `trajectory_metadata_blob` — protobuf, whose first `file://` URI is the
 *   workspace the conversation ran in;
 * - one `gen_metadata` row per model call — protobuf, field 1 a message whose
 *   field 19 is the model name and field 4 `ModelUsageStats`: 2 input (NOT
 *   including cache reads), 3 output (INCLUDING thinking), 4 cache write,
 *   5 cache read.
 * The session's name and start come from its transcript
 * (`brain/<id>/.system_generated/logs/transcript_full.jsonl`, first
 * `USER_INPUT`). A file is re-read whole when its mtime (or its non-empty
 * `-wal`'s) changes — they are small — and the mtime per file lives in
 * `antigravity_conversations`, so a session the user deleted stays deleted
 * until agy writes to it again. Opening read-only still lets SQLite leave an
 * EMPTY `-wal` and a `-shm` beside a cleanly closed file, which is why an empty
 * `-wal` is not activity. A file that is not agy's is recorded too (with no
 * session), and tried again only once it changes or the server restarts.
 *
 * The status is the session's last activity — its file, its transcript and its
 * newest hook event: `agy` fires no hook when it exits, nor when a run ends on
 * an API error or a permission headless mode refused, so a session gone quiet
 * settles to completed (its agents too) and sets its lane idle, when the hooks
 * left it running and no dashboard pane owns its status. Quiet is 10 minutes
 * for a session the hooks report on (a long tool call writes nothing until it
 * ends; after a Stop agy may be waiting for the user, and it fires no hook on
 * exit), else 2 minutes of file activity. An error Stop's `error` is kept.
 * Only the CLI's directory is read, never the Antigravity IDE's.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const fs = require("node:fs");
const path = require("node:path");
// The driver db.js chose: better-sqlite3 has no build in a plugin install
// (installed without scripts), and only db.js falls back to node:sqlite.
const { db, stmts, Database } = require("../db");
const { claimToolSessions } = require("./dashboard-runs");
const { getAgyCliDir } = require("./antigravity-home");
const pricingSync = require("./pricing-sync");

const ACTIVE_MS = 120_000; // the window codex-usage.js calls a session active
const HOOK_QUIET_MS = 600_000; // with hooks: a tool call may run long, a prompt may wait
const FUTURE_SLACK_MS = 1000;
const FILE_NAME = /^([0-9a-f-]{36})\.db$/;
const NOT_AGYS = /SQLITE_NOTADB|SQLITE_CORRUPT|no such table/;

/**
 * A protobuf message's fields as `[number, wireType, value]` — a varint as a
 * Number (exact below 2^53, which every count is), a length-delimited field as
 * a Buffer. Throws on a malformed buffer.
 */
function pbFields(buf) {
  // node:sqlite hands a BLOB over as a plain Uint8Array, whose toString()
  // ignores its encoding — model names and file URIs would decode as "1,2,3".
  if (!Buffer.isBuffer(buf)) buf = Buffer.from(buf.buffer, buf.byteOffset, buf.byteLength);
  const out = [];
  let i = 0;
  const varint = () => {
    let r = 0;
    for (let s = 0; ; s += 7) {
      if (i >= buf.length || s > 63) throw new Error("bad varint");
      const b = buf[i++];
      r += (b & 0x7f) * 2 ** s;
      if (b < 0x80) return r;
    }
  };
  while (i < buf.length) {
    const key = varint();
    const field = Math.floor(key / 8);
    const wire = key % 8;
    if (field === 0) throw new Error("field 0");
    if (wire === 0) out.push([field, 0, varint()]);
    else if (wire === 1) i += 8;
    else if (wire === 5) i += 4;
    else if (wire === 2) {
      const n = varint();
      if (i + n > buf.length) throw new Error("truncated");
      out.push([field, 2, buf.subarray(i, i + n)]);
      i += n;
    } else throw new Error(`wire type ${wire}`);
  }
  if (i > buf.length) throw new Error("truncated");
  return out;
}

const field = (fields, n, wire) => fields.find((f) => f[0] === n && f[1] === wire)?.[2];

/** One `gen_metadata` row → `{model, input, output, cacheRead, cacheWrite}`, or null. */
function generationUsage(buf) {
  try {
    const gen = field(pbFields(buf), 1, 2);
    if (!gen) return null;
    const g = pbFields(gen);
    const stats = field(g, 4, 2);
    if (!stats) return null;
    const s = pbFields(stats);
    const n = (k) => field(s, k, 0) || 0;
    if (field(s, 2, 0) === undefined && field(s, 3, 0) === undefined) return null;
    const name = field(g, 19, 2)?.toString("utf8");
    return {
      model: name && /^[\w.:/-]{1,100}$/.test(name) ? name : "unknown",
      input: n(2),
      output: n(3),
      cacheWrite: n(4),
      cacheRead: n(5),
    };
  } catch {
    return null;
  }
}

/** The first `file://` URI anywhere in a protobuf blob, as a path; null otherwise. */
function firstFileUri(buf, depth = 0) {
  let fields;
  try {
    fields = pbFields(buf);
  } catch {
    return null;
  }
  for (const [, wire, v] of fields) {
    if (wire !== 2) continue;
    if (v.subarray(0, 7).toString("latin1") === "file://") {
      try {
        return decodeURIComponent(new URL(v.toString("utf8")).pathname) || null;
      } catch {
        continue;
      }
    }
    if (depth < 6) {
      const found = firstFileUri(v, depth + 1);
      if (found) return found;
    }
  }
  return null;
}

/** This account's agy dir, then every member's (S3: their own home). */
const agyDirs = () => [
  getAgyCliDir(),
  ...require("./run-as")
    .memberHomes()
    .map((h) => path.join(h, ".gemini", "antigravity-cli")),
];

// A conversation id is a UUID, so its transcript is wherever it exists.
const transcriptFiles = (id) =>
  agyDirs().flatMap((cli) => {
    const dir = path.join(cli, "brain", id, ".system_generated", "logs");
    return ["transcript_full.jsonl", "transcript.jsonl"].map((name) => path.join(dir, name));
  });

/** The transcript's last write, 0 when it has none. */
function transcriptMtime(id) {
  for (const file of transcriptFiles(id)) {
    try {
      return Math.floor(fs.statSync(file).mtimeMs);
    } catch {
      /* try the next name */
    }
  }
  return 0;
}

/** Name (first prompt line, 60 chars) and start from the transcript's first user step. */
function transcriptFacts(id) {
  for (const file of transcriptFiles(id)) {
    let text;
    try {
      const fd = require("./run-as").openSafe(file);
      try {
        const buf = Buffer.alloc(Math.min(fs.fstatSync(fd).size, 1024 * 1024));
        fs.readSync(fd, buf, 0, buf.length, 0);
        text = buf.toString("utf8");
      } finally {
        fs.closeSync(fd);
      }
    } catch {
      continue;
    }
    const facts = { startedAt: null, firstPrompt: null };
    for (const line of text.split("\n")) {
      let o;
      try {
        o = JSON.parse(line);
      } catch {
        continue;
      }
      facts.startedAt ||= typeof o.created_at === "string" ? o.created_at : null;
      if (o.type === "USER_INPUT" && typeof o.content === "string") {
        const req = /<USER_REQUEST>\s*([\s\S]*?)\s*<\/USER_REQUEST>/.exec(o.content);
        const prompt = (req ? req[1] : o.content).trim();
        if (prompt) {
          facts.firstPrompt = prompt.split("\n")[0].slice(0, 60);
          break;
        }
      }
    }
    return facts;
  }
  return { startedAt: null, firstPrompt: null };
}

/** Everything AgentHub reads from one conversation file. */
function readConversation(file) {
  // S3: SQLite opens by NAME, and its unchecked `-journal`/`-wal` beside it
  // (a member's FIFO there blocks this thread forever). A member's file and
  // WAL are read through the guard into a private copy, opened writable so
  // SQLite can replay the WAL. The host's own is opened in place.
  const { ownerForRead, readFileSafe } = require("./run-as");
  let tmp = null;
  try {
    let open = file;
    if (require("./linux-accounts").s3Enabled() && ownerForRead(file)) {
      tmp = fs.mkdtempSync(path.join(require("node:os").tmpdir(), "agenthub-agy-"));
      open = path.join(tmp, "conversation.db");
      fs.writeFileSync(open, readFileSafe(file, null, { max: 256 << 20 })); // ponytail: whole file in memory, capped
      try {
        fs.writeFileSync(`${open}-wal`, readFileSafe(`${file}-wal`, null, { max: 256 << 20 }));
      } catch {
        /* no WAL (or not theirs): what is in the main file */
      }
    }
    const conn = new Database(open, { readonly: !tmp, fileMustExist: true });
    try {
      const blob = conn.prepare("SELECT data FROM trajectory_metadata_blob LIMIT 1").get();
      const usage = {};
      let model = null;
      for (const { data } of conn.prepare("SELECT data FROM gen_metadata ORDER BY idx").iterate()) {
        const u = data && generationUsage(data);
        if (!u) continue;
        model = u.model;
        const t = (usage[u.model] ||= { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 });
        for (const k of Object.keys(t)) t[k] += u[k];
      }
      return {
        cwd: blob?.data ? firstFileUri(blob.data) : null,
        usage,
        model,
      };
    } finally {
      conn.close();
    }
  } finally {
    if (tmp) fs.rmSync(tmp, { recursive: true, force: true });
  }
}

/** The conversation files, each with the newer mtime of the file and its non-empty `-wal`. */
function conversationFiles() {
  return agyDirs().flatMap((cli) => conversationFilesIn(path.join(cli, "conversations")));
}

function conversationFilesIn(dir) {
  let names;
  try {
    names = fs.readdirSync(dir);
  } catch {
    return [];
  }
  const out = [];
  for (const name of names) {
    const m = FILE_NAME.exec(name);
    if (!m) continue;
    const file = path.join(dir, name);
    try {
      let mtime = fs.statSync(file).mtimeMs;
      try {
        // Our own read-only open leaves an empty one behind; it holds nothing.
        const wal = fs.statSync(`${file}-wal`);
        if (wal.size > 0) mtime = Math.max(mtime, wal.mtimeMs);
      } catch {
        /* no WAL right now */
      }
      out.push({ path: file, id: m[1], mtime: Math.floor(mtime) });
    } catch {
      /* gone between readdir and stat */
    }
  }
  return out;
}

const getSeen = db.prepare("SELECT * FROM antigravity_conversations WHERE path = ?");
const upsertSeen = db.prepare(
  `INSERT INTO antigravity_conversations (path, session_id, mtime_ms, updated_at)
   VALUES (?, ?, ?, strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
   ON CONFLICT(path) DO UPDATE SET session_id = excluded.session_id,
     mtime_ms = excluded.mtime_ms, updated_at = excluded.updated_at`
);
const insertSession = db.prepare(
  `INSERT INTO sessions (id, name, status, cwd, model, started_at, ended_at, updated_at, metadata, tool)
   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'antigravity')`
);
const updateSession = db.prepare(
  `UPDATE sessions SET name = ?, status = ?, model = COALESCE(?, model), ended_at = ?,
     updated_at = ?, cwd = COALESCE(cwd, ?), started_at = COALESCE(?, started_at)
   WHERE id = ? AND tool = 'antigravity'`
);
const settle = db.prepare(
  "UPDATE sessions SET status = 'completed', ended_at = ? WHERE id = ? AND tool = 'antigravity'"
);
const clearUsage = db.prepare("DELETE FROM token_usage WHERE session_id = ?");
const clearAwaiting = db.prepare(
  "UPDATE sessions SET awaiting_input_since = NULL, awaiting_reason = NULL WHERE id = ? AND tool = 'antigravity'"
);
const stuckLanes = db.prepare(
  `SELECT id FROM lanes WHERE tool = 'antigravity' AND session_id = ? AND status = 'running'
     AND (run_id IS NULL OR run_id = '')`
);

/** The lanes a quiet session left `running` (no Stop came), set idle. */
function releaseLanes(sessionId, broadcastLane) {
  for (const { id } of stuckLanes.all(sessionId)) {
    try {
      require("./lanes").updateLane(id, { status: "idle" });
      broadcastLane(id);
    } catch {
      /* a lane is never worth failing the pass for */
    }
  }
}

/**
 * The session's last activity (file, transcript, newest hook) and its quiet
 * window: 10 minutes for a session the hooks report on — mid-turn a long tool
 * call writes nothing, and after a Stop agy may be waiting for the user (it
 * fires no hook on exit, so an open agy and a closed one look the same) — else
 * 2 minutes of file activity.
 */
function activity(id, fileMtime) {
  const ev = stmts.lastEvent.get(id);
  return {
    last: Math.max(fileMtime, transcriptMtime(id), ev ? Date.parse(ev.created_at) || 0 : 0),
    window: ev ? HOOK_QUIET_MS : ACTIVE_MS,
    stopped: ev?.event_type === "Stop",
  };
}

// A write "in the future" (a clock that runs ahead) is not proof of activity.
const fresh = (a, now) => a.last - now < FUTURE_SLACK_MS && now - a.last < a.window;

/** A quiet session's agents, completed like a SessionEnd would. */
function completeAgents(sessionId, at, broadcast) {
  for (const agent of stmts.listAgentsBySession.all(sessionId)) {
    if (agent.status === "completed" || agent.status === "error") continue;
    stmts.updateAgent.run(null, "completed", null, null, at, null, agent.id);
    broadcast("agent_updated", stmts.getAgent.get(agent.id));
  }
  stmts.clearSessionAgentsAwaitingInput.run(sessionId);
  clearAwaiting.run(sessionId); // not stmts' own: that one stamps updated_at = now
}

/** Write one conversation's session row and usage. @returns the broadcast type, or null. */
function writeSession(f, conv, now) {
  const id = f.id;
  const existing = stmts.getSession.get(id);
  if (existing && existing.tool !== "antigravity") return null;
  // S3: a member's conversation names only their own folder and session.
  if (!require("./run-as").sameOwner(f.path, conv.cwd, existing?.cwd)) return null;
  const facts = transcriptFacts(id);
  const a = activity(id, f.mtime);
  const lastWrite = new Date(a.last).toISOString();
  const active = fresh(a, now);
  let status = active ? "active" : "completed";
  let endedAt = active ? null : lastWrite;
  // A run the hooks saw fail stays an error until the next turn.
  if (existing?.status === "error" && a.stopped) ({ status, ended_at: endedAt } = existing);
  const name =
    facts.firstPrompt ||
    (existing && !/^Antigravity [0-9a-f]{8}$/.test(existing.name) ? existing.name : null) ||
    `Antigravity ${id.slice(0, 8)}`;
  db.transaction(() => {
    if (existing)
      updateSession.run(
        name,
        status,
        conv.model,
        endedAt,
        lastWrite,
        conv.cwd,
        facts.startedAt,
        id
      );
    else
      insertSession.run(
        id,
        name,
        status,
        conv.cwd,
        conv.model,
        facts.startedAt || lastWrite,
        endedAt,
        lastWrite,
        JSON.stringify({ source: "antigravity-conversation" })
      );
    claimToolSessions(id); // its workspace, before the broadcast
    // Every usage row of an Antigravity session is this module's: rewritten whole.
    clearUsage.run(id);
    for (const [m, t] of Object.entries(conv.usage)) {
      stmts.setTokenUsage.run(
        id,
        m,
        "standard",
        "global",
        "standard",
        t.input,
        t.output,
        t.cacheRead,
        t.cacheWrite,
        0,
        0,
        0,
        0
      );
    }
  })();
  return existing ? "session_updated" : "session_created";
}

let running = null;

/**
 * One pass over every conversation file: a changed one is re-read, an
 * unchanged one only has an `active` session settled once it went quiet.
 * @returns {Promise<{files: number, sessions: number}>}
 */
function ingestAntigravityUsage({
  broadcast = () => {},
  broadcastLane = () => {},
  now = Date.now(),
} = {}) {
  if (running) return running;
  running = (async () => {
    let read = 0;
    let written = 0;
    for (const f of conversationFiles()) {
      try {
        const seen = getSeen.get(f.path);
        if (seen && seen.mtime_ms === f.mtime) {
          const a = activity(f.id, f.mtime);
          if (fresh(a, now)) continue;
          const s = stmts.getSession.get(f.id);
          if (s && s.tool === "antigravity" && s.status === "active") {
            const at = new Date(a.last).toISOString();
            settle.run(at, s.id);
            completeAgents(s.id, at, broadcast);
            broadcast("session_updated", stmts.getSession.get(s.id));
            written += 1;
          }
          releaseLanes(f.id, broadcastLane);
          continue;
        }
        let conv;
        try {
          conv = readConversation(f.path);
        } catch (err) {
          // Not agy's (or not yet): tried again once it changes, or after a
          // restart (a newer AgentHub may read it). Anything else — a busy
          // file, our own bug — is retried on the next pass.
          if (NOT_AGYS.test(`${err.code} ${err.message}`)) upsertSeen.run(f.path, "", f.mtime);
          throw err;
        }
        read += 1;
        const type = writeSession(f, conv, now);
        upsertSeen.run(f.path, f.id, f.mtime);
        if (type) {
          written += 1;
          const s = stmts.getSession.get(f.id);
          if (s?.status === "completed") {
            completeAgents(f.id, s.ended_at, broadcast);
            releaseLanes(f.id, broadcastLane);
          }
          broadcast(type, stmts.getSession.get(f.id));
        }
      } catch (err) {
        if (getSeen.get(f.path)?.mtime_ms !== f.mtime) {
          console.warn(`[antigravity-usage] skipped ${f.path}: ${err.message}`);
        }
      }
      await new Promise((r) => setImmediate(r));
    }
    return { files: read, sessions: written };
  })().finally(() => {
    running = null;
  });
  return running;
}

/**
 * First pass at once (in the background), then every
 * DASHBOARD_ANTIGRAVITY_POLL_MS (default 15 s; 0 disables).
 */
function startAntigravityUsagePoll(broadcast, broadcastLane) {
  const POLL_MS = process.env.DASHBOARD_ANTIGRAVITY_POLL_MS
    ? Number(process.env.DASHBOARD_ANTIGRAVITY_POLL_MS)
    : 15_000;
  if (!Number.isFinite(POLL_MS) || POLL_MS <= 0) return;
  // Files an older build could not read get another try.
  db.prepare("DELETE FROM antigravity_conversations WHERE session_id = ''").run();
  const tick = () => {
    ingestAntigravityUsage({ broadcast, broadcastLane })
      .then((r) => {
        if (r.sessions > 0 && pricingSync.syncEnabled()) void pricingSync.syncPricing(db);
      })
      .catch(() => {});
  };
  tick();
  const timer = setInterval(tick, POLL_MS);
  if (timer.unref) timer.unref();
}

module.exports = {
  pbFields,
  generationUsage,
  firstFileUri,
  readConversation,
  ingestAntigravityUsage,
  startAntigravityUsagePoll,
};