AgentHub/server/lib/hourly-usage.js

/**
 * @file Per-hour token usage (`token_usage_hourly`): written from a Claude
 * transcript's `tokensByHour` (`server/lib/transcript-cache.js`) by the hook
 * writer, back-filled once for sessions recorded before the table existed, and
 * read only by the daily cost split in `GET /api/pricing/cost` — so a session
 * that runs for days books each day's spend to that day instead of its start.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const { db, stmts } = require("../db");
const TranscriptCache = require("./transcript-cache");
const { longArgs } = require("./token-usage");

// What this process last wrote per session and hour bucket: a long session has
// one bucket per active hour per model, and every hook re-sent all of them.
// The upsert keeps the MAX, so skipping a bucket whose numbers did not move
// writes exactly the same rows. Bounded (oldest session dropped first); a miss
// only means one full write, never a lost one.
const WRITTEN_MAX = 500;
const written = new Map();

/** Write the hours of one transcript parse that changed (MAX per hour, see the statement). */
function writeHourly(sessionId, tokensByHour) {
  let seen = written.get(sessionId);
  if (seen) written.delete(sessionId);
  else seen = new Map();
  written.set(sessionId, seen);
  if (written.size > WRITTEN_MAX) written.delete(written.keys().next().value);
  for (const [key, t] of Object.entries(tokensByHour || {})) {
    const sig = JSON.stringify(t);
    if (seen.get(key) === sig) continue;
    seen.set(key, sig);
    stmts.upsertHourlyUsage.run(
      sessionId,
      t.hour,
      t.model,
      t.speed,
      t.geo,
      t.tier,
      t.input,
      t.output,
      t.cacheRead,
      t.cacheWrite,
      t.cacheWrite1h,
      t.webSearch,
      t.webFetch,
      t.codeExec
    );
    const long = longArgs(t);
    if (long)
      stmts.raiseLongHourly.run(...long, sessionId, t.hour, t.model, t.speed, t.geo, t.tier);
  }
}

/**
 * Fill the table for Claude sessions that have tokens but no hourly rows yet,
 * one transcript per tick so a large history never blocks the event loop.
 * Resolves with the number of sessions written.
 */
function backfillHourly() {
  const rows = db
    .prepare(
      `SELECT s.id, s.transcript_path FROM sessions s
       WHERE s.transcript_path IS NOT NULL AND s.transcript_path != ''
         AND EXISTS (SELECT 1 FROM token_usage tu WHERE tu.session_id = s.id)
         AND NOT EXISTS (SELECT 1 FROM token_usage_hourly h WHERE h.session_id = s.id)`
    )
    .all();
  const cache = new TranscriptCache(1);
  let written = 0;
  return new Promise((resolve) => {
    const next = (i) => {
      if (i >= rows.length) return resolve(written);
      const { id, transcript_path } = rows[i];
      try {
        const result = cache.extract(transcript_path);
        if (result?.tokensByHour && Object.keys(result.tokensByHour).length > 0) {
          db.transaction(() => writeHourly(id, result.tokensByHour))();
          written++;
        }
      } catch {
        /* an unreadable transcript keeps its start-date booking */
      }
      setImmediate(() => next(i + 1));
    };
    next(0);
  });
}

/** Forget what was written — for `ids`, or every session. Whatever deletes
 *  hourly rows calls it, or the next hooks would skip re-writing them. */
function forgetHourly(ids = null) {
  if (!ids) return written.clear();
  for (const id of ids) written.delete(id);
}

module.exports = { writeHourly, backfillHourly, forgetHourly };