AgentHub/server/lib/omp-usage.js

/**
 * @file Oh My Pi (`omp`) sessions, subagents and token usage, read from the
 * session journals omp writes under `~/.omp/agent/sessions/<encoded cwd>/`.
 * This is how an omp session reaches `sessions` (as `tool = 'omp'`), its
 * tokens `token_usage` and its subagents `agents` — omp has no hooks.
 *
 * What a journal holds (omp 18.x, verified against files it wrote — not a
 * documented format):
 * - its file name is `<timestamp>_<session id>.jsonl`, and `omp --resume <id>`
 *   takes that id back;
 * - line 1 is a fixed-width `{"type":"title", title}` slot omp rewrites in
 *   place, line 2 the `{"type":"session", id, cwd, timestamp}` header;
 * - every assistant turn is a `{"type":"message"}` whose `message.usage` is
 *   `{input, output, cacheRead, cacheWrite}` (input NOT including cache reads)
 *   under `message.model`;
 * - a process that ends appends `{"type":"custom","customType":"session_exit"}`;
 * - a subagent the `task` tool spawns writes its own journal, in the same
 *   format, under the folder named like the session file
 *   (`<timestamp>_<session id>/<Agent name>.jsonl`, nested agents deeper).
 * Journals only grow (bar the title slot), so each is read from where the last
 * pass stopped; what it said so far is kept per file in `omp_sessions`, so a
 * session the user deleted stays deleted until omp writes to it again. A
 * session's tokens are the sum of its own journal and its subagents'.
 *
 * Status: completed once its process wrote `session_exit`, else active while
 * a journal changed in the last 2 minutes or a live dashboard pane holds the
 * session (an idle omp waiting for its user writes nothing), else completed.
 * S3: every member's own `~/.omp` is read too, only through `openSafeAsync`
 * (never a link, only a file that member owns), and only for a session whose
 * folder is that same member's.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const fs = require("node:fs");
const os = require("node:os");
const path = require("node:path");
const { db, stmts } = require("../db");
const { claimToolSessions } = require("./dashboard-runs");
const pricingSync = require("./pricing-sync");
const { readNew } = require("./jsonl-tail");

const ACTIVE_MS = 120_000; // the window codex-usage.js calls a session active
const FUTURE_SLACK_MS = 1000;
const SESSION_FILE = /^.+_([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})\.jsonl$/;
const SESSION_DIR = /^.+_([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})$/;
const HEAD_BYTES = 4096;
const MAX_DEPTH = 4; // nested subagents
const MAX_FILES = 20_000;

/** `~/.omp`, or `AGENTHUB_OMP_DIR`; read at call time. */
function getOmpDir() {
  return process.env.AGENTHUB_OMP_DIR || path.join(os.homedir(), ".omp");
}

/** A subagent's key: its path under the session folder, as one plain name. */
const agentKey = (rel) =>
  rel
    .replace(/\.jsonl$/, "")
    .split(path.sep)
    .join(".")
    .replace(/[^\w.-]/g, "_")
    .replace(/^\.+/, "")
    .slice(0, 120);

/** The agents-table id of a subagent journal. */
const agentRowId = (sessionId, key) => `omp-${sessionId}-${key}`;

/** Every journal: `{path, sessionId, agent ('' = the session's own), mtime, size}`. */
function journalFiles() {
  const members = require("./run-as")
    .memberHomes()
    .map((h) => path.join(h, ".omp"));
  const out = [];
  for (const dir of [getOmpDir(), ...members]) {
    journalFilesIn(path.join(dir, "agent", "sessions"), out);
  }
  return out;
}

function statFile(file) {
  try {
    const st = fs.lstatSync(file);
    return st.isFile() ? st : null;
  } catch {
    return null;
  }
}

function journalFilesIn(root, out) {
  let projects;
  try {
    projects = fs.readdirSync(root, { withFileTypes: true });
  } catch {
    return;
  }
  for (const p of projects) {
    if (!p.isDirectory()) continue;
    const dir = path.join(root, p.name);
    let entries;
    try {
      entries = fs.readdirSync(dir, { withFileTypes: true });
    } catch {
      continue;
    }
    // The session's own journal before its subagents': a subagent row needs
    // its session row.
    entries.sort((a, b) => Number(b.isFile()) - Number(a.isFile()));
    for (const e of entries) {
      if (out.length >= MAX_FILES) return;
      const full = path.join(dir, e.name);
      let m;
      if (e.isFile() && (m = SESSION_FILE.exec(e.name))) {
        const st = statFile(full);
        if (st) {
          out.push({
            path: full,
            sessionId: m[1],
            agent: "",
            mtime: Math.floor(st.mtimeMs),
            size: st.size,
          });
        }
      } else if (e.isDirectory() && (m = SESSION_DIR.exec(e.name))) {
        subagentFiles(full, full, m[1], 0, out);
      }
    }
  }
}

function subagentFiles(base, dir, sessionId, depth, out) {
  if (depth >= MAX_DEPTH) return;
  let entries;
  try {
    entries = fs.readdirSync(dir, { withFileTypes: true });
  } catch {
    return;
  }
  for (const e of entries) {
    if (out.length >= MAX_FILES) return;
    const full = path.join(dir, e.name);
    // Real folders and files only: a link is never followed.
    if (e.isDirectory()) subagentFiles(base, full, sessionId, depth + 1, out);
    else if (e.isFile() && e.name.endsWith(".jsonl")) {
      const st = statFile(full);
      const key = agentKey(path.relative(base, full));
      if (st && key) {
        out.push({
          path: full,
          sessionId,
          agent: key,
          mtime: Math.floor(st.mtimeMs),
          size: st.size,
        });
      }
    }
  }
}

/** A fresh per-journal state. */
const emptyState = () => ({
  v: 1,
  cwd: null,
  startedAt: null,
  lastAt: null,
  model: null,
  firstPrompt: null,
  exited: false,
  usage: {},
});

/** A message's text parts, joined. */
function textOf(content) {
  if (typeof content === "string") return content;
  if (!Array.isArray(content)) return "";
  return content
    .map((c) => (c && c.type === "text" && typeof c.text === "string" ? c.text : ""))
    .join("");
}

/** A model id without its provider (`anthropic/claude-sonnet-4-5` → `claude-sonnet-4-5`). */
const bareModel = (m) =>
  typeof m === "string" && m ? m.slice(m.lastIndexOf("/") + 1).slice(0, 128) : null;

const num = (v) => (Number.isFinite(v) && v > 0 ? Math.floor(v) : 0);

/** Fold one journal line into `state`. */
function applyEntry(state, e) {
  if (!e || typeof e !== "object") return;
  const at = typeof e.timestamp === "string" ? e.timestamp : null;
  switch (e.type) {
    case "session":
      if (typeof e.cwd === "string" && path.isAbsolute(e.cwd)) state.cwd = e.cwd;
      if (at && !state.startedAt) state.startedAt = at;
      break;
    case "model_change":
      state.model = bareModel(e.model) || state.model;
      state.exited = false;
      break;
    case "message": {
      const m = e.message || {};
      state.exited = false;
      if (at) state.lastAt = at;
      if (m.role === "user" && !state.firstPrompt) {
        const text = textOf(m.content).trim();
        if (text && !text.startsWith("<system-reminder>")) state.firstPrompt = text.slice(0, 500);
      }
      if (m.role === "assistant" && m.usage && typeof m.usage === "object") {
        const model = bareModel(m.model) || state.model || "unknown";
        state.model = model;
        const u = (state.usage[model] ||= { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 });
        u.input += num(m.usage.input);
        u.output += num(m.usage.output);
        u.cacheRead += num(m.usage.cacheRead);
        u.cacheWrite += num(m.usage.cacheWrite);
      }
      break;
    }
    case "custom":
      if (e.customType === "session_exit") state.exited = true;
      break;
    default:
      break;
  }
}

/** The journal's title slot (line 1), or null. */
async function readTitle(file) {
  let fh;
  try {
    fh = await require("./run-as").openSafeAsync(file);
    const buf = Buffer.alloc(HEAD_BYTES);
    const { bytesRead } = await fh.read(buf, 0, HEAD_BYTES, 0);
    const first = buf.toString("utf8", 0, bytesRead).split("\n")[0];
    const o = JSON.parse(first);
    return o?.type === "title" && typeof o.title === "string" && o.title.trim()
      ? o.title.trim().slice(0, 200)
      : null;
  } catch {
    return null;
  } finally {
    await fh?.close();
  }
}

const getSeen = db.prepare("SELECT * FROM omp_sessions WHERE path = ?");
const upsertSeen = db.prepare(
  `INSERT INTO omp_sessions (path, session_id, agent, size, mtime_ms, offset, state, title, updated_at)
   VALUES (@path, @sessionId, @agent, @size, @mtime, @offset, @state, @title,
           strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
   ON CONFLICT(path) DO UPDATE SET session_id = excluded.session_id, agent = excluded.agent,
     size = excluded.size, mtime_ms = excluded.mtime_ms, offset = excluded.offset,
     state = excluded.state, title = excluded.title, updated_at = excluded.updated_at`
);
const filesOf = db.prepare("SELECT * FROM omp_sessions WHERE session_id = ? ORDER BY agent");
const activeOmp = db.prepare("SELECT id FROM sessions WHERE tool = 'omp' AND status = 'active'");
const heldByPane = db.prepare(
  "SELECT 1 FROM dashboard_runs WHERE session_id = ? AND status IN ('running', 'spawning') LIMIT 1"
);
const insertSession = db.prepare(
  `INSERT INTO sessions (id, name, status, cwd, model, started_at, ended_at, updated_at, metadata, tool)
   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'omp')`
);
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 = 'omp'`
);
const clearUsage = db.prepare("DELETE FROM token_usage WHERE session_id = ?");
const agentTimes = db.prepare(
  "UPDATE agents SET started_at = COALESCE(?, started_at), ended_at = ? WHERE id = ?"
);

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

const parseState = (row) => {
  try {
    return { ...emptyState(), ...JSON.parse(row.state) };
  } catch {
    return emptyState();
  }
};

/** Read what one journal wrote since the last pass into its `omp_sessions` row. */
async function readJournal(f) {
  const seen = getSeen.get(f.path);
  let state = seen && seen.session_id === f.sessionId ? parseState(seen) : emptyState();
  let offset = seen && seen.session_id === f.sessionId ? seen.offset : 0;
  // Smaller than what was read: rewritten, so read again from the top.
  if (f.size < offset) {
    state = emptyState();
    offset = 0;
  }
  const got = await readNew(f.path, offset, f.size);
  for (const line of got.lines) {
    if (!line) continue;
    try {
      applyEntry(state, JSON.parse(line));
    } catch {
      /* a line that is not JSON: skipped */
    }
  }
  const title = f.agent ? null : await readTitle(f.path);
  upsertSeen.run({
    path: f.path,
    sessionId: f.sessionId,
    agent: f.agent,
    size: f.size,
    // A journal read only part-way is re-read next pass even if unchanged.
    mtime: got.offset < f.size ? -1 : f.mtime,
    offset: got.offset,
    state: JSON.stringify(state),
    title,
  });
}

/**
 * Rewrite one session's row, token usage and subagents from its journals'
 * rows. @returns the session's broadcast type, or null (nothing written).
 */
function writeSession(sessionId, now, broadcast) {
  const files = filesOf.all(sessionId);
  const main = files.find((r) => r.agent === "");
  if (!main) return null;
  const existing = stmts.getSession.get(sessionId);
  if (existing && existing.tool !== "omp") return null;
  const ms = parseState(main);
  const runAs = require("./run-as");
  // S3: a member's journal names only their own folder and session.
  if (!runAs.sameOwner(main.path, ms.cwd, existing?.cwd)) return null;
  const lastMtime = Math.max(...files.map((r) => r.mtime_ms));
  const lastWrite = new Date(lastMtime > 0 ? lastMtime : now).toISOString();
  const active = !ms.exited && (fresh(lastMtime, now) || !!heldByPane.get(sessionId));
  const status = active ? "active" : "completed";
  const endedAt = active ? null : lastWrite;
  const firstLine = ms.firstPrompt ? ms.firstPrompt.split("\n")[0].slice(0, 60) : null;
  const name = main.title || firstLine || existing?.name || `OMP ${sessionId.slice(0, 8)}`;
  const usage = {};
  for (const r of files) {
    for (const [m, u] of Object.entries(parseState(r).usage)) {
      const t = (usage[m] ||= { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 });
      t.input += u.input;
      t.output += u.output;
      t.cacheRead += u.cacheRead;
      t.cacheWrite += u.cacheWrite;
    }
  }
  const agentEvents = [];
  db.transaction(() => {
    if (existing) {
      updateSession.run(name, status, ms.model, endedAt, lastWrite, ms.cwd, ms.startedAt, sessionId);
    } else {
      insertSession.run(
        sessionId,
        name,
        status,
        ms.cwd,
        ms.model,
        ms.startedAt || lastWrite,
        endedAt,
        lastWrite,
        JSON.stringify({ source: "omp-journal" })
      );
    }
    claimToolSessions(sessionId); // its workspace, before the broadcast
    // Every usage row of an omp session is this module's: rewritten whole.
    clearUsage.run(sessionId);
    for (const [m, t] of Object.entries(usage)) {
      stmts.setTokenUsage.run(
        sessionId,
        m,
        "standard",
        "global",
        "standard",
        t.input,
        t.output,
        t.cacheRead,
        t.cacheWrite,
        0,
        0,
        0,
        0
      );
    }
    for (const r of files) {
      if (!r.agent || !runAs.sameOwner(main.path, r.path)) continue;
      const s = parseState(r);
      const id = agentRowId(sessionId, r.agent);
      const done = !active || s.exited || !fresh(r.mtime_ms, now);
      const agentStatus = done ? "completed" : "working";
      const ended = done ? new Date(r.mtime_ms > 0 ? r.mtime_ms : now).toISOString() : null;
      const label = r.agent.split(".").pop() || r.agent;
      const prior = stmts.getAgent.get(id);
      if (prior && prior.session_id !== sessionId) continue;
      if (!prior) {
        stmts.insertAgent.run(
          id,
          sessionId,
          label,
          "subagent",
          "task",
          agentStatus,
          s.firstPrompt,
          null,
          JSON.stringify({ source: "omp-journal", transcript: r.agent })
        );
      } else if (prior.status !== agentStatus || (s.firstPrompt && prior.task !== s.firstPrompt)) {
        stmts.updateAgent.run(null, agentStatus, s.firstPrompt, null, ended, null, id);
      } else continue;
      agentTimes.run(s.startedAt, ended, id);
      agentEvents.push([prior ? "agent_updated" : "agent_created", id]);
    }
  })();
  for (const [type, id] of agentEvents) broadcast(type, stmts.getAgent.get(id));
  return existing ? "session_updated" : "session_created";
}

let running = null;
// Fresh sessions the last pass left for a later one to settle (the poll keeps
// its base pace while there are any).
let lastPending = 0;

/**
 * One pass: every changed journal is read on from where it stopped, then each
 * session it touched — and each one still `active` — is rewritten.
 * @returns {Promise<{files: number, sessions: number}>}
 */
function ingestOmpUsage({ broadcast = () => {}, now = Date.now() } = {}) {
  if (running) return running;
  running = (async () => {
    let read = 0;
    let written = 0;
    const touched = new Set();
    for (const f of journalFiles()) {
      const seen = getSeen.get(f.path);
      if (seen && seen.mtime_ms === f.mtime && seen.size === f.size) continue;
      try {
        await readJournal(f);
        read += 1;
        touched.add(f.sessionId);
      } catch (err) {
        console.warn(`[omp-usage] skipped ${f.path}: ${err.message}`);
      }
      await new Promise((r) => setImmediate(r));
    }
    for (const { id } of activeOmp.all()) touched.add(id);
    let pending = 0;
    for (const id of touched) {
      try {
        const type = writeSession(id, now, broadcast);
        if (!type) continue;
        const s = stmts.getSession.get(id);
        if (s?.status === "active") pending += 1;
        written += 1;
        broadcast(type, s);
      } catch (err) {
        console.warn(`[omp-usage] session ${id}: ${err.message}`);
      }
    }
    lastPending = pending;
    return { files: read, sessions: written };
  })().finally(() => {
    running = null;
  });
  return running;
}

/**
 * The journal behind a session's transcript: its own (`agent` '' / "main"),
 * or one subagent's by key. Null when unknown.
 */
function journalPath(sessionId, agent = "") {
  const key = agent === "main" ? "" : agent || "";
  const row = db
    .prepare("SELECT path FROM omp_sessions WHERE session_id = ? AND agent = ?")
    .get(sessionId, key);
  return row?.path || null;
}

/** A session's subagent journals: `[{agent, name, task, startedAt}]`. */
function subagentJournals(sessionId) {
  return filesOf
    .all(sessionId)
    .filter((r) => r.agent)
    .map((r) => {
      const s = parseState(r);
      return {
        agent: r.agent,
        name: r.agent.split(".").pop() || r.agent,
        task: s.firstPrompt,
        startedAt: s.startedAt,
        rowId: agentRowId(sessionId, r.agent),
      };
    });
}

/**
 * First pass at once (in the background), then every DASHBOARD_OMP_POLL_MS
 * (default 15 s; 0 disables) — slowed to once a minute after a pass that read
 * nothing, wrote nothing and left no active session to settle.
 */
function startOmpUsagePoll(broadcast) {
  const POLL_MS = process.env.DASHBOARD_OMP_POLL_MS
    ? Number(process.env.DASHBOARD_OMP_POLL_MS)
    : 15_000;
  if (!Number.isFinite(POLL_MS) || POLL_MS <= 0) return;
  require("./backoff-poll").startBackoffPoll(
    () =>
      ingestOmpUsage({ broadcast }).then((r) => {
        if (r.sessions > 0 && pricingSync.syncEnabled()) void pricingSync.syncPricing(db);
        return r.files + r.sessions + lastPending > 0;
      }),
    { baseMs: POLL_MS }
  );
}

module.exports = {
  getOmpDir,
  agentKey,
  applyEntry,
  emptyState,
  ingestOmpUsage,
  journalPath,
  subagentJournals,
  startOmpUsagePoll,
};