AgentHub/server/lib/budgets.js

/**
 * @file Spend budgets. A budget is an alert rule of type `budget` (config:
 * `period` day | week | month, `limit_usd`, `thresholds` in percent, optional
 * `tool` and `lane` scope, `tz_offset` = the browser's getTimezoneOffset() so a
 * period starts at the user's local midnight). This module computes a budget's
 * spend in its current period — priced exactly like GET /api/pricing/cost, by
 * the day the tokens were spent — the pace-based projection to the period's
 * end, and the status list behind GET /api/budgets/status. Firing lives in
 * `alerts.js`, once per threshold per period. Budgets are org-wide; a server
 * mode member's status shows each one's spend over their own sessions only.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const { stmts } = require("../db");
const { memberId, ownedOnlySessionWhere } = require("./scope");
const { calculateCost, dailyTokenRows, costWhere } = require("../routes/pricing");

const PERIODS = ["day", "week", "month"];
const DAY_MS = 86_400_000;

const ymd = (ms) => new Date(ms).toISOString().slice(0, 10);

/**
 * The period containing `now`, in the user's local time.
 * @param {"day"|"week"|"month"} period
 * @param {number} now epoch ms
 * @param {number} tzOffset minutes, as getTimezoneOffset() (UTC+7 → -420)
 * @param {number} [back] periods before the current one (0 = current)
 * @returns {{start: string, end: string, startUtc: string, days: number, elapsed: number}}
 *   `start`/`end` local dates (inclusive), `startUtc` the start as an ISO
 *   instant, `elapsed` the fraction of the period gone (1 for a past period).
 */
function periodBounds(period, now, tzOffset = 0, back = 0) {
  const local = new Date(now - tzOffset * 60_000); // read with getUTC*
  let startLocal;
  let endLocal; // exclusive
  if (period === "day") {
    startLocal = Date.UTC(local.getUTCFullYear(), local.getUTCMonth(), local.getUTCDate() - back);
    endLocal = startLocal + DAY_MS;
  } else if (period === "week") {
    const dow = (local.getUTCDay() + 6) % 7; // Monday = 0
    startLocal = Date.UTC(
      local.getUTCFullYear(),
      local.getUTCMonth(),
      local.getUTCDate() - dow - 7 * back
    );
    endLocal = startLocal + 7 * DAY_MS;
  } else {
    startLocal = Date.UTC(local.getUTCFullYear(), local.getUTCMonth() - back, 1);
    endLocal = Date.UTC(local.getUTCFullYear(), local.getUTCMonth() - back + 1, 1);
  }
  const span = endLocal - startLocal;
  return {
    start: ymd(startLocal),
    end: ymd(endLocal - DAY_MS),
    startUtc: new Date(startLocal + tzOffset * 60_000).toISOString(),
    days: Math.round(span / DAY_MS),
    elapsed: back > 0 ? 1 : Math.min(1, Math.max(0, (local.getTime() - startLocal) / span)),
  };
}

/**
 * Spend of a budget config in one period (current by default); with a member's
 * `req`, over that member's sessions only.
 * @returns {{period_start: string, period_end: string, spend: number,
 *   projected: number, pct: number}}
 */
function budgetSpend(cfg, { now = Date.now(), back = 0, rules, req } = {}) {
  const tz = Number.isFinite(cfg.tz_offset) ? cfg.tz_offset : 0;
  const b = periodBounds(cfg.period, now, tz, back);
  // A session that spent in the period was touched in it, so older ones are
  // skipped before their hours are read.
  const extra = [{ clause: "s.updated_at >= ?", params: [b.startUtc] }];
  const scope = cfg.tool ? { where: "tool = ?", params: [cfg.tool] } : null;
  let paid = [];
  if (memberId(req) != null) {
    // Their OWN money: their workspaces' sessions, by the hours they paid
    // where a collaborator ran one — and the hours they paid for in
    // workspaces shared with them (server/lib/session-payers.js).
    const own = ownedOnlySessionWhere(req, "s.id", "s.lane_id");
    const payers = require("./session-payers");
    extra.push({ clause: `(${own.where}) AND s.id NOT IN (${payers.SPLIT})`, params: own.params });
    const split = costWhere({
      scope,
      lane: cfg.lane,
      extra: [
        { clause: "s.updated_at >= ?", params: [b.startUtc] },
        { clause: `s.id IN (${payers.SPLIT})`, params: [] },
      ],
    });
    paid = payers.paidDailyRows(req.user.id, split.whereClause, split.params, `${-tz} minutes`);
  }
  const where = costWhere({ scope, lane: cfg.lane, extra });
  const rows = dailyTokenRows(`${-tz} minutes`, where.whereClause, where.params)
    .concat(paid)
    .filter((r) => r.date >= b.start && r.date <= b.end);
  const spend = calculateCost(rows, rules || stmts.listPricing.all()).total_cost;
  // The pace so far, carried to the period's end; at least an hour of pace so
  // the first minutes of a period do not project a fortune.
  const projected = spend / Math.max(b.elapsed, 1 / (b.days * 24));
  const round = (n) => Math.round(n * 10000) / 10000;
  return {
    period_start: b.start,
    period_end: b.end,
    spend: round(spend),
    projected: round(back > 0 ? spend : projected),
    pct: cfg.limit_usd > 0 ? round((spend / cfg.limit_usd) * 100) : 0,
  };
}

/** The thresholds a spend has reached, lowest first. */
function crossed(cfg, pct) {
  return [...cfg.thresholds].sort((a, b) => a - b).filter((t) => pct >= t);
}

/**
 * Every enabled budget rule with its current-period spend (`req`: see budgetSpend).
 * @returns {Array<object>} rule id/name/config + `budgetSpend` + `crossed`.
 */
function budgetStatus({ now = Date.now(), req } = {}) {
  const rules = stmts.listPricing.all();
  return stmts.listEnabledAlertRules
    .all()
    .filter((r) => r.rule_type === "budget")
    .map((r) => {
      let cfg;
      try {
        cfg = JSON.parse(r.config || "{}");
      } catch {
        return null;
      }
      if (!PERIODS.includes(cfg.period) || !(cfg.limit_usd > 0)) return null;
      const s = budgetSpend(cfg, { now, rules, req });
      return { id: r.id, name: r.name, config: cfg, ...s, crossed: crossed(cfg, s.pct) };
    })
    .filter(Boolean);
}

module.exports = { PERIODS, periodBounds, budgetSpend, budgetStatus, crossed };