codex-market-app/electron/lib/token_usage.js

// Read Codex's own cumulative token snapshots from ~/.codex/sessions.
"use strict";

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

const CODEX_HOME = process.env.CODEX_HOME || path.join(os.homedir(), ".codex");
const SESSIONS_ROOT = path.join(CODEX_HOME, "sessions");
const RESET_PATH = path.join(CODEX_HOME, ".tokendun-token-reset.json");
const fileCache = new Map();

function clean(n) {
  n = Number(n);
  return Number.isFinite(n) && n > 0 ? Math.round(n) : 0;
}

function dayKey(date = new Date()) {
  const y = date.getFullYear();
  const m = String(date.getMonth() + 1).padStart(2, "0");
  const d = String(date.getDate()).padStart(2, "0");
  return `${y}-${m}-${d}`;
}

function blank(date = dayKey()) {
  return {
    available: false,
    date,
    total: 0,
    input: 0,
    output: 0,
    cached: 0,
    reasoning: 0,
    unknown: 0,
    lastSource: "session",
  };
}

function snapshotUsage(info) {
  const u = info && typeof info === "object" ? info : {};
  const input = clean(u.input_tokens || u.prompt_tokens || u.inputTokenCount || u.promptTokenCount);
  const cached = clean(
    u.cached_input_tokens ||
      u.cache_read_input_tokens ||
      u.cached_tokens ||
      (u.input_tokens_details && u.input_tokens_details.cached_tokens)
  );
  const output = clean(u.output_tokens || u.completion_tokens || u.candidatesTokenCount);
  const reasoning = clean(u.reasoning_output_tokens || u.reasoning_tokens || u.thoughtsTokenCount);
  const total = clean(u.total_tokens || u.totalTokenCount || u.total) || input + output;
  return { total, input, output, cached, reasoning };
}

function snapshotAt(line) {
  let item;
  try { item = JSON.parse(line); } catch { return null; }
  if (item?.type !== "event_msg" || item?.payload?.type !== "token_count") return null;
  const timestamp = new Date(item.timestamp);
  if (!Number.isFinite(timestamp.getTime())) return null;
  const usage = snapshotUsage(item.payload.info?.total_token_usage);
  return usage.total > 0 ? { timestamp, usage } : null;
}

function parseSessionFile(file) {
  try {
    const stat = fs.statSync(file);
    const cached = fileCache.get(file);
    if (cached && cached.size === stat.size && cached.mtimeMs === stat.mtimeMs) return cached.value;
    const snapshots = [];
    const text = fs.readFileSync(file, "utf8");
    for (const line of text.split("\n")) {
      const snapshot = snapshotAt(line);
      if (snapshot) snapshots.push(snapshot);
    }
    const value = { snapshots };
    fileCache.set(file, { size: stat.size, mtimeMs: stat.mtimeMs, value });
    return value;
  } catch {
    return { snapshots: [] };
  }
}

function utcDatePath(date) {
  return path.join(
    SESSIONS_ROOT,
    String(date.getUTCFullYear()).padStart(4, "0"),
    String(date.getUTCMonth() + 1).padStart(2, "0"),
    String(date.getUTCDate()).padStart(2, "0")
  );
}

function sessionFilesForDay(start, end) {
  const dirs = new Set();
  for (const cursor = new Date(start); cursor < end; cursor.setDate(cursor.getDate() + 1)) {
    dirs.add(utcDatePath(cursor));
  }
  dirs.add(utcDatePath(new Date(end.getTime() - 1)));
  const files = [];
  for (const dir of dirs) {
    try {
      for (const entry of fs.readdirSync(dir, { withFileTypes: true })) {
        if (entry.isFile() && entry.name.endsWith(".jsonl")) files.push(path.join(dir, entry.name));
      }
    } catch {}
  }
  return files;
}

function delta(current, before) {
  return Math.max(0, clean(current) - clean(before));
}

function readReset() {
  try {
    const value = JSON.parse(fs.readFileSync(RESET_PATH, "utf8"));
    return value && typeof value === "object" ? value : null;
  } catch {
    return null;
  }
}

function applyReset(value) {
  const reset = readReset();
  if (!reset || reset.date !== value.date) return value;
  return {
    ...value,
    total: delta(value.total, reset.total),
    input: delta(value.input, reset.input),
    output: delta(value.output, reset.output),
    cached: delta(value.cached, reset.cached),
    reasoning: delta(value.reasoning, reset.reasoning),
  };
}

function readDailyRaw(now = new Date()) {
  const date = dayKey(now);
  const start = new Date(now.getFullYear(), now.getMonth(), now.getDate());
  const end = new Date(now.getFullYear(), now.getMonth(), now.getDate() + 1);
  const result = blank(date);
  result.available = fs.existsSync(SESSIONS_ROOT);

  for (const file of sessionFilesForDay(start, end)) {
    const snapshots = parseSessionFile(file).snapshots;
    let before = null;
    let current = null;
    for (const snapshot of snapshots) {
      if (snapshot.timestamp < start) before = snapshot;
      else if (snapshot.timestamp < end && (!current || snapshot.timestamp > current.timestamp)) current = snapshot;
    }
    if (!current) continue;
    const b = before?.usage || {};
    const u = current.usage;
    result.total += delta(u.total, b.total);
    result.input += delta(u.input, b.input);
    result.output += delta(u.output, b.output);
    result.cached += delta(u.cached, b.cached);
    result.reasoning += delta(u.reasoning, b.reasoning);
  }
  return result;
}

function readDaily(now = new Date()) {
  return applyReset(readDailyRaw(now));
}

function reset(now = new Date()) {
  const value = readDailyRaw(now);
  fs.mkdirSync(path.dirname(RESET_PATH), { recursive: true });
  fs.writeFileSync(RESET_PATH, JSON.stringify({ ...value, date: dayKey(now) }), "utf8");
  return blank(dayKey(now));
}

module.exports = {
  readDaily,
  reset,
  _snapshotUsage: snapshotUsage,
  _sessionFilesForDay: sessionFilesForDay,
};