stack/packages/bus/src/usage.mjs

// REQ-LIMIT-4: the host calls sources only through the registry loader.
import { createLoader } from '../../plugins/src/loader.mjs';
const KINDS = ['input', 'output', 'cacheWrite', 'cacheRead'];
const zero = () => ({ input: 0, output: 0, cacheWrite: 0, cacheRead: 0 });
const unknown = code => ({ status: 'unknown', code });
function snapshot(value) {
  if (value?.status === 'unknown') return unknown(/^transcript-(path|size|line|entries|records|changed|parse|usage|harness)$/.test(value.code) ? value.code : 'source-unknown');
  if (value?.status !== 'known' || !value.entries || typeof value.entries !== 'object' || Array.isArray(value.entries))
    return unknown('source-invalid');
  const entries = Object.entries(value.entries);
  if (entries.length > 20000) return unknown('source-invalid');
  for (const [id, v] of entries) {
    if (id.length > 512 || !v || Object.keys(v).length !== 4 ||
        KINDS.some(k => !Number.isSafeInteger(v[k]) || v[k] < 0)) return unknown('source-invalid');
  }
  return { status: 'known', entries: value.entries };
}
function delta(before, after) {
  if (!before || before.status !== 'known') return unknown('baseline-unknown');
  if (after.status !== 'known') return after;
  for (const [id, v] of Object.entries(before.entries)) {
    if (!Object.hasOwn(after.entries, id) || KINDS.some(k => after.entries[id][k] !== v[k])) return unknown('transcript-rewritten');
  }
  const sum = zero(); let added = false;
  for (const [id, v] of Object.entries(after.entries)) {
    if (Object.hasOwn(before.entries, id)) continue;
    added = true;
    for (const k of KINDS) {
      sum[k] += v[k];
      if (!Number.isSafeInteger(sum[k])) return unknown('usage-overflow');
    }
  }
  return added ? { status: 'known', ...sum } : unknown('no-new-usage');
}
// One record per (run, turn, source): the latest row, an end row winning. A turn is known for a
// source only when that row is an end row with status known; a begin row alone is unknown.
// view() and the limit checks (row 125, packages/bus/src/limits.mjs) count the same records.
export function turnRecords(store, business) {
  const turns = new Map();
  for (const row of store.all(`SELECT * FROM usage_records WHERE business=? ORDER BY seq`, business)) {
    const key = JSON.stringify([row.run, row.turn, row.source]);
    if (row.phase === 'end' || !turns.has(key)) turns.set(key, row);
  }
  return [...turns.values()].map(row => {
    const known = row.phase === 'end' && row.status === 'known';
    return {
      at: row.at, run: row.run, turn: row.turn, source: row.source, instance: row.instance,
      provider: row.provider, model: row.model, account: row.account, status: known ? 'known' : 'unknown',
      input: known ? row.input : null, output: known ? row.output : null,
      cacheWrite: known ? row.cache_write : null, cacheRead: known ? row.cache_read : null,
    };
  });
}
export async function createUsage({ store, businesses, log }) {
  const loaders = new Map();
  for (const [id, b] of Object.entries(businesses)) {
    const loader = createLoader({ log });
    await loader.load(b.plugins ?? {});
    loaders.set(id, loader);
  }
  // This map serializes retries per run, without keeping a SQL transaction across source awaits.
  const pending = new Map();
  async function perform(session, phase, turn) {
    const { business, run, role: instance, harness, provider = null, model = null, account = null } = session;
    const existing = store.all('SELECT * FROM usage_records WHERE business=? AND run=? AND turn=? AND phase=?', business, run, turn, phase);
    if (existing.length) return { recorded: true };
    const loader = loaders.get(business);
    const modules = loader.modules('usage-source'), names = loader.loaded()['usage-source'];
    const bases = store.all("SELECT * FROM usage_records WHERE business=? AND run=? AND turn=? AND phase='begin'", business, run, turn);
    const sources = modules.length ? modules.map((module, i) => ({ module, name: names[i] })) : [{ name: '@none' }];
    // Preserve a failed/disabled baseline even if configuration changes between calls/restarts.
    for (const base of bases) if (!sources.some(s => s.name === base.source)) sources.push({ name: base.source });
    const rows = [];
    // Row 65 (#1558): where the run's harness keeps its session files is the module's transcript(),
    // which the launcher binds into the run record. A run bound without one has no reading.
    const root = session.transcript ?? null;
    for (const { module, name } of sources) {
      let sample = unknown('no-source');
      if (module && root === null) sample = unknown('no-transcript');
      else if (module) {
        try { sample = snapshot(await module.readTurn({ harness, root })); }
        catch { sample = unknown('source-failed'); }
      }
      const base = bases.find(row => row.source === name);
      const result = phase === 'begin' ? sample : delta(base?.cursor ? JSON.parse(base.cursor) : null, sample);
      rows.push({ name, result });
    }
    store.transaction(() => {
      for (const { name, result } of rows) {
        const known = result.status === 'known';
        const counts = phase === 'begin' ? zero() : result;
        store.run(`INSERT INTO usage_records(at,business,run,instance,provider,model,account,turn,source,phase,status,code,cursor,input,output,cache_write,cache_read)
          VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`, new Date().toISOString(), business, run, instance,
          provider, model, account, turn, name, phase, result.status, result.code ?? null,
          phase === 'begin' ? JSON.stringify(result) : null,
          known ? counts.input : null, known ? counts.output : null, known ? counts.cacheWrite : null, known ? counts.cacheRead : null);
      }
    });
    for (const { name, result } of rows) if (phase === 'end' && result.status === 'unknown') {
      log({ event: 'usage-unknown', business, run, instance, turn, source: name, code: result.code });
    }
    return { recorded: true };
  }
  function collect(session, phase, turn) {
    const key = session.business + ':' + session.run;
    const op = (pending.get(key) ?? Promise.resolve()).then(() => perform(session, phase, turn));
    const tail = op.catch(() => {});
    pending.set(key, tail);
    return op.finally(() => { if (pending.get(key) === tail) pending.delete(key); });
  }
  function view(business) {
    const groups = { byRun: new Map(), byInstance: new Map(), byProvider: new Map(), byModel: new Map(), byAccount: new Map() };
    const total = { ...zero(), unknown: new Set() };
    const add = (target, row) => {
      if (row.status !== 'known') target.unknown.add(row.run + ':' + row.turn);
      else for (const key of KINDS) target[key] = BigInt(target[key]) + BigInt(row[key]);
    };
    for (const row of turnRecords(store, business)) {
      add(total, row);
      for (const [name, identity] of Object.entries({ byRun: { run: row.run }, byInstance: { instance: row.instance },
        byProvider: { provider: row.provider }, byModel: { provider: row.provider, model: row.model }, byAccount: { account: row.account } })) {
        const key = JSON.stringify(identity);
        if (!groups[name].has(key)) groups[name].set(key, { ...identity, ...zero(), unknown: new Set() });
        add(groups[name].get(key), row);
      }
    }
    const finish = value => {
      const { unknown, ...rest } = value;
      for (const kind of KINDS) rest[kind] = BigInt(rest[kind]) <= BigInt(Number.MAX_SAFE_INTEGER) ? Number(rest[kind]) : String(rest[kind]);
      return { ...rest, unknownTurns: unknown.size };
    };
    return { total: finish(total), ...Object.fromEntries(Object.entries(groups).map(([key, map]) => [key, [...map.values()].map(finish)])) };
  }
  return { collect, view };
}