AgentHub/server/lib/codex-usage.js

/**
 * @file Codex sessions and token usage, read from Codex's rollout journals
 * (`$CODEX_HOME/sessions/YYYY/MM/DD/rollout-<ts>-<thread>.jsonl`). This is how
 * a Codex session reaches `sessions` (as `tool = 'codex'`) and its tokens reach
 * `token_usage` — Codex's hooks, when the user installs them, add its events
 * and live status but no tokens.
 *
 * What a journal holds, and the two traps:
 * - line 1 `session_meta`: `id` = the thread, `session_id` = the ROOT session
 *   (a sub-agent thread's parent; absent on old CLIs, then `id`), `cwd`,
 *   `timestamp`, `forked_from_id` on a fork;
 * - `turn_context.model` is the model for the lines that follow;
 * - `token_count` events carry a cumulative `total_token_usage` and its delta
 *   `last_token_usage` — `input_tokens` INCLUDES `cached_input_tokens`,
 *   `output_tokens` includes reasoning. The same total is often sent twice.
 * - a FORK copies its parent's whole history into its own file, token events
 *   included, then carries on from the parent's total: every total vector the
 *   parent's file holds is skipped in the child, or the parent counts twice.
 * - a sub-agent thread has its own file; its usage is summed into the root.
 *
 * Progress per file lives in `codex_rollouts` (offset + parse state), so an
 * unchanged file is never re-read, a grown one is read from where it stopped,
 * and a session the user deleted stays deleted until its journal changes.
 * Totals are exact recomputes, written with `setTokenUsage`. A session is named
 * by Codex's own thread name (`session_index.jsonl`), else its first prompt. Rows are never
 * given a `transcript_path`, so the Claude transcript parsers skip them, and a
 * session row whose tool is not `codex` is never touched. A session Codex's
 * hooks report on (an event no older than its newest journal write, less two
 * minutes) keeps the status, end and last-active time the hooks gave it; this
 * module still writes its usage, model and name. Hooks that stopped reporting
 * (removed, no longer trusted) hand the status back to the journal.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const fs = require("node:fs");
const path = require("node:path");
const readline = require("node:readline");
const { db, stmts } = require("../db");
const { getCodexHome, codexHomes } = require("./codex-home");
const { allRolloutFiles, FUTURE_SLACK_MS } = require("./codex-ingest");
const pricingSync = require("./pricing-sync");

const ACTIVE_MS = 120_000; // same window codex-ingest.js calls a lane running
// Read at most this much of one journal per pass: a sync read of a 100 MB
// journal would block the event loop for about a second.
const MAX_CHUNK = 16 * 1024 * 1024;
const NEEDLES = [
  '"session_meta"',
  '"turn_context"',
  '"token_count"',
  '"user_message"', // a prompt, older CLIs
  '"UserMessage"', // a prompt, as an item_completed event (0.149+)
];
const FILE_NAME = /^rollout-\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}-(.+)\.jsonl$/;

function totalKey(t) {
  return [
    t.input_tokens,
    t.cached_input_tokens,
    t.cache_write_input_tokens,
    t.output_tokens,
    t.reasoning_output_tokens,
  ].join(",");
}

/** The text a user typed, from either prompt event shape; null otherwise. */
function promptOf(p) {
  let text = null;
  if (p.type === "user_message") text = p.message;
  else if (p.type === "item_completed" && p.item?.type === "UserMessage") {
    text = (p.item.content || [])
      .filter((c) => c && c.type === "text")
      .map((c) => c.text)
      .join("\n");
  }
  return typeof text === "string" && text.trim() ? text.trim() : null;
}

/**
 * Fold one chunk of journal text into `state` (mutated and returned). `skip`
 * holds total vectors that belong to a fork's parent.
 */
function parseRolloutChunk(text, state = {}, skip = null) {
  const s = state;
  s.usage ||= {};
  for (const line of text.split("\n")) {
    if (!line || !NEEDLES.some((n) => line.includes(n))) continue;
    let o;
    try {
      o = JSON.parse(line);
    } catch {
      continue;
    }
    const p = o.payload || {};
    if (o.type === "session_meta") {
      if (s.threadId) continue; // a fork repeats its ancestors' metas
      s.threadId = p.id || null;
      s.sessionId = p.session_id || p.id || null;
      s.cwd = p.cwd || null;
      s.startedAt = p.timestamp || o.timestamp || null;
      s.forkedFrom = p.forked_from_id || null;
    } else if (o.type === "turn_context") {
      if (p.model) s.model = p.model;
    } else if (o.type === "event_msg" && !s.firstPrompt && promptOf(p)) {
      // The session's label when Codex has no thread name for it — the same
      // first-line, 60-character rule hooks.js applies to a Claude prompt.
      s.firstPrompt = promptOf(p).split("\n")[0].slice(0, 60);
    } else if (o.type === "event_msg" && p.type === "token_count" && p.info?.total_token_usage) {
      const key = totalKey(p.info.total_token_usage);
      if (key === s.lastTotal) continue;
      s.lastTotal = key;
      // A fork's copied history comes first; its first total the parent never
      // wrote ends it, so later totals are never matched against the parent.
      if (s.forkedFrom && !s.pastFork) {
        if (skip?.has(key)) continue;
        if (skip) s.pastFork = true;
      }
      const l = p.info.last_token_usage || {};
      const u = (s.usage[s.model || "unknown"] ||= {
        input: 0,
        cacheRead: 0,
        cacheWrite: 0,
        output: 0,
      });
      const cached = l.cached_input_tokens || 0;
      u.input += Math.max(0, (l.input_tokens || 0) - cached);
      u.cacheRead += cached;
      u.cacheWrite += l.cache_write_input_tokens || 0;
      u.output += l.output_tokens || 0;
    }
  }
  return s;
}

/**
 * Bytes [offset, size), at most MAX_CHUNK, up to the last newline, as text +
 * bytes consumed + whether more is left. A window with no newline at all is
 * one line longer than MAX_CHUNK — never a token event — and is skipped whole.
 */
function readChunk(file, offset, size) {
  const buf = Buffer.alloc(Math.max(0, Math.min(size - offset, MAX_CHUNK)));
  const fd = fs.openSync(file, "r");
  try {
    fs.readSync(fd, buf, 0, buf.length, offset);
  } finally {
    fs.closeSync(fd);
  }
  let end = buf.lastIndexOf(0x0a) + 1; // 0 when no full line yet
  if (end === 0 && buf.length === MAX_CHUNK) end = MAX_CHUNK;
  return { text: buf.toString("utf8", 0, end), bytes: end, more: offset + end < size && end > 0 };
}

/** Every total vector in a parent journal: the history a fork copied. Streamed. */
async function parentTotals(file) {
  const keys = new Set();
  const lines = readline.createInterface({ input: fs.createReadStream(file), crlfDelay: Infinity });
  for await (const line of lines) {
    if (!line.includes('"token_count"')) continue;
    try {
      const t = JSON.parse(line).payload?.info?.total_token_usage;
      if (t) keys.add(totalKey(t));
    } catch {
      /* torn line */
    }
  }
  return keys;
}

function firstLineMeta(file) {
  try {
    const { text } = readChunk(file, 0, Math.min(fs.statSync(file).size, 1024 * 1024));
    const o = JSON.parse(text.split("\n")[0]);
    return o.type === "session_meta" ? o.payload || {} : {};
  } catch {
    return {};
  }
}

const threadCache = new Map(); // rollout path → {id, sessionId}; a first line never changes

/**
 * The thread a rollout journal belongs to — `{id, sessionId}` from its
 * session_meta, `sessionId` being the ROOT session — or null. Only a
 * `rollout-*.jsonl` under `$CODEX_HOME/sessions` (as given or resolved) is read: the path comes from a
 * hook payload, and a file elsewhere could name any session it likes.
 */
function threadOf(file) {
  if (typeof file !== "string" || !path.isAbsolute(file)) return null;
  const resolved = path.resolve(file);
  const under = codexHomes().some((h) => resolved.startsWith(path.join(h, "sessions") + path.sep));
  if (!under || !FILE_NAME.test(path.basename(resolved))) return null;
  if (threadCache.has(resolved)) return threadCache.get(resolved);
  const meta = firstLineMeta(resolved);
  const thread = meta.id ? { id: meta.id, sessionId: meta.session_id || meta.id } : null;
  if (!thread) return null; // not written yet: ask again next time
  if (threadCache.size >= 500) threadCache.clear();
  threadCache.set(resolved, thread);
  return thread;
}

/** `thread_name` per id from `$CODEX_HOME/session_index.jsonl` (last wins). */
function threadNames() {
  const names = new Map();
  let text = "";
  try {
    text = fs.readFileSync(path.join(getCodexHome(), "session_index.jsonl"), "utf8");
  } catch {
    return names;
  }
  for (const line of text.split("\n")) {
    try {
      const o = JSON.parse(line);
      if (o.id && o.thread_name) names.set(o.id, String(o.thread_name));
    } catch {
      /* blank or torn */
    }
  }
  return names;
}

const getRollout = db.prepare("SELECT * FROM codex_rollouts WHERE path = ?");
const upsertRollout = db.prepare(
  `INSERT INTO codex_rollouts (path, session_id, thread_id, size, mtime_ms, offset, state, updated_at)
   VALUES (?, ?, ?, ?, ?, ?, ?, strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
   ON CONFLICT(path) DO UPDATE SET session_id = excluded.session_id, thread_id = excluded.thread_id,
     size = excluded.size, mtime_ms = excluded.mtime_ms, offset = excluded.offset,
     state = excluded.state, updated_at = excluded.updated_at`
);
// One row per thread: the same journal met under a new path (CODEX_HOME moved,
// or reached through a symlink) replaces its old row instead of counting twice.
const dropOtherPaths = db.prepare("DELETE FROM codex_rollouts WHERE thread_id = ? AND path <> ?");
const rolloutsOf = db.prepare(
  "SELECT thread_id, mtime_ms, state FROM codex_rollouts WHERE session_id = ?"
);
const insertCodexSession = db.prepare(
  `INSERT INTO sessions (id, name, status, cwd, model, started_at, ended_at, updated_at, metadata, tool)
   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'codex')`
);
const updateCodexSession = db.prepare(
  `UPDATE sessions SET name = ?, status = ?, model = ?, ended_at = ?, updated_at = ?,
     cwd = COALESCE(?, cwd), started_at = COALESCE(?, started_at)
   WHERE id = ? AND tool = 'codex'`
);
const clearUsage = db.prepare("DELETE FROM token_usage WHERE session_id = ?");
// A session Codex's hooks report on (server/lib/codex-hooks.js) has events;
// its status is theirs, as a Claude session's is — while they still report:
// an event no older than the journal's newest write (less ACTIVE_MS). Hooks
// removed, or no longer trusted, hand the status back to the journal.
const lastEventAt = db.prepare("SELECT MAX(created_at) AS t FROM events WHERE session_id = ?");
const newestRollout = db.prepare(
  "SELECT MAX(mtime_ms) AS m FROM codex_rollouts WHERE session_id = ?"
);
function hooksCurrent(sessionId, newestMtime) {
  const t = Date.parse(lastEventAt.get(sessionId)?.t || "");
  return Number.isFinite(t) && t >= newestMtime - ACTIVE_MS;
}
// A journal written "in the future" (a clock that runs ahead) is not proof of
// activity — codex-ingest.js treats it the same way.
const fresh = (mtime, now) => mtime - now < FUTURE_SLACK_MS && now - mtime < ACTIVE_MS;

/** Re-derive one root session's row and usage from all its threads. */
function writeSession(sessionId, names, now) {
  const rows = rolloutsOf.all(sessionId).map((r) => ({ ...r, state: JSON.parse(r.state) }));
  if (rows.length === 0) return null;
  const existing = stmts.getSession.get(sessionId);
  if (existing && existing.tool !== "codex") return null;

  // The root thread's own file, when it has been read: a sub-agent file read
  // first must not fix the session's cwd, start or model.
  const ownRoot = rows.find((r) => r.thread_id === sessionId);
  const root = ownRoot || rows[0];
  const newest = rows.reduce((a, b) => (b.mtime_ms > a.mtime_ms ? b : a));
  const active = rows.some((r) => fresh(r.mtime_ms, now));
  const name =
    names.get(sessionId) ||
    root.state.firstPrompt ||
    existing?.name ||
    `Codex ${sessionId.slice(0, 8)}`;
  const model = (ownRoot && ownRoot.state.model) || newest.state.model || null;
  // A journal is appended to the end, so its mtime is when the session last
  // did anything — later than its last token event after an abort.
  let lastWrite = new Date(newest.mtime_ms).toISOString();
  let status = active ? "active" : "completed";
  let endedAt = active ? null : lastWrite;
  if (existing && hooksCurrent(sessionId, newest.mtime_ms)) {
    ({ status, ended_at: endedAt, updated_at: lastWrite } = existing);
  }

  const totals = {};
  for (const r of rows) {
    for (const [m, u] of Object.entries(r.state.usage || {})) {
      const t = (totals[m] ||= { input: 0, cacheRead: 0, cacheWrite: 0, output: 0 });
      for (const k of Object.keys(t)) t[k] += u[k] || 0;
    }
  }

  db.transaction(() => {
    // `updated_at` is the session's last activity (the list's "Last active"
    // and its sort), so it is the journal's last write, not ingestion time.
    if (existing)
      updateCodexSession.run(
        name,
        status,
        model,
        endedAt,
        lastWrite,
        ownRoot ? ownRoot.state.cwd : null,
        ownRoot ? ownRoot.state.startedAt : null,
        sessionId
      );
    else
      insertCodexSession.run(
        sessionId,
        name,
        status,
        root.state.cwd,
        model,
        root.state.startedAt || new Date(now).toISOString(),
        endedAt,
        lastWrite,
        JSON.stringify({ source: "codex-rollout" })
      );
    // Every usage row of a Codex session is this module's, so it is rewritten
    // whole: a model gone from a re-read journal leaves no stale row behind.
    clearUsage.run(sessionId);
    for (const [m, t] of Object.entries(totals)) {
      stmts.setTokenUsage.run(
        sessionId,
        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 journal. Yields to the event loop between files; a
 * second call while one runs joins it.
 * @returns {Promise<{files: number, sessions: number}>} files read, sessions written
 */
function ingestCodexUsage({ broadcast = () => {}, now = Date.now() } = {}) {
  if (running) return running;
  running = (async () => {
    const files = allRolloutFiles();
    const byThread = new Map();
    for (const f of files) {
      const m = FILE_NAME.exec(path.basename(f.path));
      if (m) byThread.set(m[1], f.path);
    }
    const touched = new Set();
    let read = 0;
    for (const f of files) {
      try {
        const mtime = Math.floor(f.mtime);
        const row = getRollout.get(f.path);
        let offset = row ? row.offset : 0;
        let state = row ? JSON.parse(row.state) : {};
        if (row && row.size === f.size && row.mtime_ms === mtime && !state.more) continue;
        if (f.size < offset) {
          offset = 0;
          state = {};
        }
        let skip = null;
        const parent = offset === 0 ? firstLineMeta(f.path).forked_from_id : state.forkedFrom;
        if (parent && !state.pastFork) {
          if (byThread.has(parent)) skip = await parentTotals(byThread.get(parent));
          else if (offset === 0)
            console.warn(`[codex-usage] fork parent ${parent} not found; ${f.path} may overcount`);
        }
        const { text, bytes, more } = readChunk(f.path, offset, f.size);
        state = parseRolloutChunk(text, state, skip);
        state.more = more;
        upsertRollout.run(
          f.path,
          state.sessionId || "",
          state.threadId || "",
          f.size,
          mtime,
          offset + bytes,
          JSON.stringify(state)
        );
        if (state.threadId) dropOtherPaths.run(state.threadId, f.path);
        read += 1;
        if (state.sessionId) touched.add(state.sessionId);
      } catch (err) {
        console.warn(`[codex-usage] skipped ${f.path}: ${err.message}`);
      }
      await new Promise((r) => setImmediate(r));
    }
    let written = 0;
    const names = touched.size ? threadNames() : new Map();
    for (const sid of touched) {
      try {
        const type = writeSession(sid, names, now);
        if (!type) continue;
        written += 1;
        broadcast(type, stmts.getSession.get(sid));
      } catch (err) {
        console.warn(`[codex-usage] session ${sid}: ${err.message}`);
      }
    }
    return { files: read, sessions: written };
  })().finally(() => {
    running = null;
  });
  return running;
}

/**
 * First pass at once (in the background), then every DASHBOARD_CODEX_POLL_MS
 * (default 15 s; 0 disables, like the codex lane poll it shares the knob with).
 * A pass also re-derives `active` → `completed`, so a session whose journal
 * went quiet is re-written once its files are older than ACTIVE_MS.
 */
function startCodexUsagePoll(broadcast) {
  const POLL_MS = process.env.DASHBOARD_CODEX_POLL_MS
    ? Number(process.env.DASHBOARD_CODEX_POLL_MS)
    : 15_000;
  if (!Number.isFinite(POLL_MS) || POLL_MS <= 0) return;
  const tick = () => {
    ingestCodexUsage({ broadcast })
      .then((r) => {
        settleActive(broadcast);
        // A pass that wrote sessions may have brought a model with no rate
        // row; price it now rather than at the sync's next 15-minute check.
        if (r.sessions > 0 && pricingSync.syncEnabled()) void pricingSync.syncPricing(db);
      })
      .catch(() => {});
  };
  tick();
  const timer = setInterval(tick, POLL_MS);
  if (timer.unref) timer.unref();
}

/** Codex sessions still `active` whose every journal has gone quiet. */
function settleActive(broadcast, now = Date.now()) {
  const quiet = db
    .prepare(
      `SELECT s.id FROM sessions s WHERE s.tool = 'codex' AND s.status = 'active'
         AND NOT EXISTS (SELECT 1 FROM codex_rollouts r WHERE r.session_id = s.id
                           AND r.mtime_ms > ? AND r.mtime_ms < ?)`
    )
    .all(now - ACTIVE_MS, now + FUTURE_SLACK_MS);
  if (quiet.length === 0) return;
  const names = threadNames();
  for (const { id } of quiet) {
    if (hooksCurrent(id, newestRollout.get(id)?.m || 0)) continue;
    if (writeSession(id, names, now)) broadcast("session_updated", stmts.getSession.get(id));
  }
}

module.exports = {
  parseRolloutChunk,
  ingestCodexUsage,
  settleActive,
  startCodexUsagePoll,
  threadOf,
};