stack/packages/plugins/limit-check/token-limit/index.mjs

// Row 125 (#1619, REQ-LIMIT-3, REQ-LIMIT-4): token limits per role instance,
// provider, provider and model, and key. Pure: core passes the business's
// launch limits and its per-turn usage records, and nothing here reads a file
// or a store. A count is every kind summed over the known records since the
// limit's `since`; an unknown turn adds nothing and is counted on its own.
import { isDailyModelQuota } from '../../src/model-quota.mjs';
export const interfaceVersion = 2;

const KINDS = ['input', 'output', 'cacheWrite', 'cacheRead'];
const SAFE = BigInt(Number.MAX_SAFE_INTEGER);
const refuse = code => { throw Object.assign(new Error(code), { code }); };
const plain = n => n <= SAFE ? Number(n) : String(n);
const time = value => {
  const t = typeof value === 'string' ? Date.parse(value) : NaN;
  if (Number.isNaN(t)) refuse('bad-time');
  return t;
};
const amount = value => Number.isSafeInteger(value) && value >= 0 ? BigInt(value) : refuse('bad-record');

// Every token limit in launch.limits.tokens, in a fixed order, with the rule
// that says which records and subjects it covers.
function leaves(limits) {
  const tokens = limits?.tokens ?? {};
  const out = [];
  for (const [n, entry] of Object.entries(tokens.instances ?? {}))
    out.push({ path: ['tokens', 'instances', n], entry, match: x => x.instance === n });
  for (const [p, entry] of Object.entries(tokens.providers ?? {}))
    out.push({ path: ['tokens', 'providers', p], entry, match: x => x.provider === p });
  for (const [p, models] of Object.entries(tokens.models ?? {}))
    for (const [m, entry] of Object.entries(models))
      out.push({ path: ['tokens', 'models', p, m], entry, match: x => x.provider === p && x.model === m });
  // A key is the record's account. Accounts are null until row 126 names them,
  // so a key limit matches nothing until then.
  for (const [k, entry] of Object.entries(tokens.keys ?? {}))
    if (!isDailyModelQuota(k, entry)) out.push({ path: ['tokens', 'keys', k], entry,
      match: x => x.account === k, key: true });
  return out;
}

function tally(leaf, records) {
  const since = time(leaf.entry.since);
  const kinds = Object.fromEntries(KINDS.map(k => [k, 0n]));
  const unknown = new Set();
  for (const r of records) {
    if (!leaf.match(r) || time(r.at) < since) continue;
    if (r.status === 'known') for (const k of KINDS) kinds[k] += amount(r[k]);
    else unknown.add(JSON.stringify([r.run, r.turn]));
  }
  return { count: KINDS.reduce((sum, k) => sum + kinds[k], 0n), kinds, unknownTurns: unknown.size };
}

// The first limit covering the subject whose count is at or over it.
function hitFor({ limits, subject, records }) {
  if (!Array.isArray(records)) refuse('bad-records');
  for (const leaf of leaves(limits)) {
    if (!leaf.match(subject)) continue;
    const { count, unknownTurns } = tally(leaf, records);
    if (count >= BigInt(leaf.entry.limit)) return { path: leaf.path, limit: leaf.entry.limit, count: plain(count), unknownTurns };
  }
  return null;
}

export function mayLaunch(args) {
  const hit = hitFor(args);
  return hit ? { allow: false, hit } : { allow: true };
}

export function mustStop(args) {
  const hit = hitFor(args);
  return hit ? { stop: true, hit } : { stop: false };
}

// Every token limit with its count. A key limit is reported as not enforced
// while no record names an account (lead decision 150, Q1).
export function report({ limits, records }) {
  if (!Array.isArray(records)) refuse('bad-records');
  const named = records.some(r => r.account != null);
  return leaves(limits).map(leaf => {
    const { path, entry: { limit, since } } = leaf;
    if (leaf.key && !named) return { path, limit, since, notEnforced: 'no records name a key' };
    const { count, kinds, unknownTurns } = tally(leaf, records);
    return { path, limit, since, count: plain(count), kinds: Object.fromEntries(KINDS.map(k => [k, plain(kinds[k])])), unknownTurns };
  });
}