AgentHub/server/lib/pool-usage.js

/**
 * @file pool-usage.js
 * @description $ and tokens one account-pool account spent in the last 5 hours
 * and 7 days, for the balancer and `GET /api/accounts` (claude-pool.js,
 * codex-pool.js). Tokens are booked to hours by `hourlyTokenRows` — the same
 * booking as the daily cost split, so sub-agent buckets and pre-compaction
 * baselines count — and priced by `calculateCost`, never a second formula.
 * A Claude session counts only for the hours one of THIS account's runs held
 * it (start hour to end hour), so a session resumed on another account is
 * not booked twice; Codex has no hourly rows, so its session counts whole
 * once a run of the account started in the window, as before. With `ownerId`
 * only the runs of that user's workspaces count — one user's share of the
 * account, for the Administration page's user detail. A remote client's
 * company-account session (`remote_sessions`, server/hub/) holds its account
 * like a run does, from its registration to its end.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const WINDOWS = { "5h": 5, "7d": 24 * 7 };
const hourAgo = (hours, now) => new Date(now - hours * 3600e3).toISOString().slice(0, 13);
const tokensOf = (r) =>
  r.input_tokens + r.output_tokens + r.cache_read_tokens + r.cache_write_tokens;

/**
 * `{spend5h, spend7d, tokens5h, tokens7d}` of `account` on `tool`. With
 * `wholeSession`, a session's every row counts in a window its run started in;
 * with `ownerId`, only runs on that user's workspaces.
 */
function accountUsage(
  tool,
  account,
  { wholeSession = false, now = Date.now(), ownerId = null } = {}
) {
  const { db, stmts } = require("../db");
  const { calculateCost, hourlyTokenRows } = require("../routes/pricing");
  const since7d = new Date(now - WINDOWS["7d"] * 3600e3).toISOString();
  const owned = ownerId == null ? "" : "AND lane_id IN (SELECT id FROM lanes WHERE owner_id = ?)";
  const ownedArgs = ownerId == null ? [] : [ownerId];
  const remoteOwned = ownerId == null ? "" : "AND user_id = ?";
  // Every holder of the account: the dashboard's runs, then remote clients'
  // sessions (a run's columns under the same names).
  const HOLDERS = `SELECT session_id, started_at, ended_at FROM dashboard_runs
        WHERE tool = ? AND account = ? AND session_id IS NOT NULL
          AND COALESCE(ended_at, '9999') >= ? ${owned}
     UNION ALL
     SELECT session_id, registered_at, ended_at FROM remote_sessions
        WHERE tool = ? AND account = ? AND COALESCE(ended_at, '9999') >= ? ${remoteOwned}`;
  const holderArgs = [tool, account, since7d, ...ownedArgs, tool, account, since7d, ...ownedArgs];
  const runs = db
    .prepare(
      `SELECT session_id, substr(started_at, 1, 13) AS start_hour, started_at,
              substr(COALESCE(ended_at, '9999'), 1, 13) AS end_hour
         FROM (${HOLDERS})`
    )
    .all(...holderArgs);
  const out = { spend5h: 0, spend7d: 0, tokens5h: 0, tokens7d: 0 };
  if (runs.length === 0) return out;
  const runsOf = new Map();
  for (const r of runs) runsOf.set(r.session_id, [...(runsOf.get(r.session_id) || []), r]);
  const rows = hourlyTokenRows(`WHERE s.id IN (SELECT session_id FROM (${HOLDERS}))`, holderArgs);
  const rules = stmts.listPricing.all();
  for (const [label, hours] of Object.entries(WINDOWS)) {
    const sinceHour = hourAgo(hours, now);
    const sinceIso = new Date(now - hours * 3600e3).toISOString();
    const held = rows.filter((row) =>
      (runsOf.get(row.session_id) || []).some((run) =>
        wholeSession
          ? run.started_at >= sinceIso
          : row.date >= sinceHour && row.date >= run.start_hour && row.date <= run.end_hour
      )
    );
    out[`spend${label}`] = held.length ? calculateCost(held, rules).total_cost : 0;
    out[`tokens${label}`] = Math.round(held.reduce((n, r) => n + tokensOf(r), 0));
  }
  return out;
}

module.exports = { accountUsage };