stack/packages/bus/src/limits.mjs

// Row 125 (#1619, REQ-LIMIT-3, REQ-LIMIT-4): core's side of the limit checks. It loads only the
// `limit-check` modules the organization file switches on, one loader per business, and asks
// them through the loader alone: mayLaunch before a launch, mustStop after each turn's usage end,
// report for Settings and `mosaic limits`. Each check gets a copy of the business's launch limits
// and its per-turn usage records (usage.mjs turnRecords).
//
// Fail closed (lead decision 150). A check that throws, answers in the wrong shape, answers with
// a promise, or is switched on but not loaded (a refused load keeps the old modules) refuses the
// launch, stops the session at the end of its turn, and shows every limit as `check failed`. An
// organization file that can't be read refuses and stops the same way, with check null; report
// then throws, as the verb's own read does.
//
// refresh, mayLaunch and mustStop are async: they re-read the file and reload first. report is
// sync, so the organization verbs stay sync; it uses the modules the last refresh loaded.
import { createLoader, isDailyModelQuota } from '../../plugins/src/loader.mjs';
import { turnRecords } from './usage.mjs';

const KIND = 'limit-check';
const KINDS = ['input', 'output', 'cacheWrite', 'cacheRead'];
const CODE = /^[a-z][a-z0-9-]*$/;
const own = (value, key) => value !== null && typeof value === 'object' && !Array.isArray(value) && Object.hasOwn(value, key);
const plainObject = value => value !== null && typeof value === 'object' && Object.getPrototypeOf(value) === Object.prototype;
const exact = (value, keys) => plainObject(value) && Object.keys(value).length === keys.length && keys.every(k => Object.hasOwn(value, k));
const amount = value => (Number.isSafeInteger(value) && value >= 0) || (typeof value === 'string' && /^(0|[1-9]\d*)$/.test(value));
const turns = value => Number.isSafeInteger(value) && value >= 0;
const same = (a, b) => JSON.stringify(a) === JSON.stringify(b);
const HEX12 = /^[a-f0-9]{12}$/;
// The key lookup's own refusals pass through by name; any other throw is proxy-key-refused.
const LOOKUP_CODES = ['proxy-key-refused', 'key-quota-no-model-key'];
// The ids that name a daily per-key quota (row 333), in file order.
const dailyKeyIds = limits => Object.entries(limits?.tokens?.keys ?? {}).filter(([k, entry]) => isDailyModelQuota(k, entry)).map(([k]) => k);
// Row 333 R3 (#1831): one array of key ids per hired agent, no id in two.
const validLineages = value => {
  if (!Array.isArray(value) || value.some(l => !Array.isArray(l) || !l.length || l.some(id => typeof id !== 'string' || !HEX12.test(id)))) return false;
  const ids = value.flat();
  return new Set(ids).size === ids.length;
};

// Every token limit in launch.limits, in file order, at its schema depth only: tokens.instances.X,
// tokens.providers.X, tokens.keys.X, tokens.models.P.M and tokens.daily.A (row 126). The agent limits (total, providers,
// models) are never checks' limits, whatever their names, so a provider called `limit` there isn't
// taken for one. The settings writer (business launch-settings.mjs) walks the same depths.
const DEPTH = { instances: 1, providers: 1, keys: 1, models: 2, daily: 1 };
export function limitLeaves(limits) {
  const out = [];
  const walk = (value, depth, path) => {
    if (!plainObject(value)) return;
    for (const [key, entry] of Object.entries(value)) {
      if (depth > 1) walk(entry, depth - 1, [...path, key]);
      else if (plainObject(entry)) out.push({ path: [...path, key], limit: entry.limit, since: entry.since });
    }
  };
  const tokens = plainObject(limits) ? limits.tokens : undefined;
  if (plainObject(tokens)) {
    for (const [kind, value] of Object.entries(tokens)) if (Object.hasOwn(DEPTH, kind)) walk(value, DEPTH[kind], ['tokens', kind]);
  }
  return out;
}

class Failure extends Error {
  constructor(code) { super(code); this.code = code; }
}

// makeLoader is the plug-in loader's factory; tests pass a scratch tree's copy to load fixture checks.
// Row 333: keyForSubject names the subject's model key id: the agent's current key for a launch,
// the run's bound key for mustStop. R3: keyLineages lists the business's hired agents' key
// lineages, each the current, retiring and earlier ids of one agent. A daily key quota follows
// the lineage that holds its id, so it outlives rotation and re-issue. Both run only while a
// daily key quota is set and a check is switched on. Without keyLineages they refuse.
export function createLimits({ store, read, pluginDataRoot = null, log, makeLoader = createLoader, keyForSubject = () => null,
  keyLineages = () => { throw new Failure('proxy-key-refused'); } }) {
  if (typeof read !== 'function' || typeof log !== 'function') throw new TypeError('createLimits needs read and log');
  const loaders = new Map();
  // Row 333 R3: an orphaned key quota names an id in no hired agent's lineage. It limits nothing
  // and refuses nothing. One value-free log line per leaf when it becomes orphaned, not per launch.
  // R4: every refresh observes too, and a business with no daily key quota left remembers none, so a
  // leaf removed and added back logs again. The set lives in memory: a bus start is a fresh
  // observation and may log an unchanged orphan once more.
  const orphansLogged = new Map();
  function noteOrphans(business, limits, lineages) {
    const ids = dailyKeyIds(limits).filter(id => !lineages.some(l => l.includes(id)));
    const before = orphansLogged.get(business) ?? new Set();
    for (const id of ids) if (!before.has(id)) log({ event: 'limit-key-orphaned', business, path: ['tokens', 'keys', id] });
    orphansLogged.set(business, new Set(ids));
    return ids;
  }
  // The orphaned ids now, [] when no daily key quota is set, or null when the lineages can't be
  // read. Then the remembered set keeps only the ids the document still holds and gains none: an
  // unreadable lookup proves no leaf orphaned, but the document does prove which leaves were removed.
  function observeOrphans(business, limits) {
    const ids = dailyKeyIds(limits);
    if (!ids.length) { orphansLogged.set(business, new Set()); return []; }
    const lineages = lineagesFor(business);
    if (lineages) return noteOrphans(business, limits, lineages);
    const before = orphansLogged.get(business);
    if (before) orphansLogged.set(business, new Set([...before].filter(id => ids.includes(id))));
    return null;
  }
  // The business's lineages, or null when they can't be read or are malformed.
  function lineagesFor(business) {
    let value;
    try { value = keyLineages(business); } catch { return null; }
    return validLineages(value) ? value : null;
  }
  const loaderFor = business => {
    if (!loaders.has(business)) loaders.set(business, makeLoader({ log, dataRoot:pluginDataRoot }));
    return loaders.get(business);
  };

  // What one business has now: its limits, its records, and its switched-on checks. A check the
  // file switches on that the loader doesn't have (a refused load) is not loaded.
  function state(business, b = undefined) {
    if (b === undefined) {
      try { b = read(business); } catch { throw new Failure('organization-unreadable'); }
    }
    const setting = b?.plugins?.[KIND] ?? [];
    const loader = loaderFor(business);
    // The loader seeds every kind with an empty list, so a kind never loaded, or refused on its
    // first load, has nothing switched on.
    const names = loader.loaded()[KIND];
    const modules = loader.modules(KIND);
    let records;
    try { records = turnRecords(store, business); } catch { throw new Failure('records-unreadable'); }
    return {
      limits: b?.launch?.limits ?? {},
      records,
      checks: same(setting, names) ? names.map((name, i) => ({ name, module: modules[i] })) : null,
      missing: setting.find(name => !names.includes(name)) ?? setting[0] ?? null,
    };
  }

  async function refresh(business) {
    let b;
    try { b = read(business); } catch { return; }
    observeOrphans(business, b?.launch?.limits ?? {});
    await loaderFor(business).load({ [KIND]: b?.plugins?.[KIND] ?? [] });
  }

  // Ask one check one export. The answer comes back, or a code: the check's own code if it threw
  // one in the plug-in code shape, check-threw for any other throw, malformed-answer otherwise.
  function ask(business, check, name, args, valid) {
    try {
      const answer = check.module[name](structuredClone(args));
      if (answer !== null && (typeof answer === 'object' || typeof answer === 'function') && typeof answer.then === 'function') {
        Promise.resolve(answer).catch(() => {});
        throw new Failure('malformed-answer');
      }
      if (!valid(answer)) throw new Failure('malformed-answer');
      return { answer };
    } catch (e) {
      const code = e instanceof Failure ? e.code : typeof e?.code === 'string' && CODE.test(e.code) ? e.code : 'check-threw';
      log({ event: 'limit-check-failed', business, check: check.name, export: name, code });
      return { code };
    }
  }

  const validHit = (hit, leaves) =>
    exact(hit, ['path', 'limit', 'count', 'unknownTurns']) &&
    leaves.some(l => same(l.path, hit.path) && l.limit === hit.limit) && amount(hit.count) && turns(hit.unknownTurns);

  // mayLaunch and mustStop: every switched-on check is asked in order; the first refusal or
  // stop wins. key is the answer's boolean (allow, stop) and go its value when nothing is hit.
  async function decide(business, subject, name, key, go, halt, context = {}) {
    await refresh(business);
    let st;
    try { st = state(business); } catch (e) {
      log({ event: 'limit-check-failed', business, check: null, export: name, code: e.code });
      return halt({ check: null, code: e.code });
    }
    if (!st.checks) {
      log({ event: 'limit-check-failed', business, check: st.missing, export: name, code: 'not-loaded' });
      return halt({ check: st.missing, code: 'not-loaded' });
    }
    if (!st.checks.length) return { [key]: go };
    const leaves = limitLeaves(st.limits);
    let lineages;
    if (dailyKeyIds(st.limits).length) {
      let modelKeyId, code = 'proxy-key-refused';
      try {
        if (name === 'mustStop' && (typeof context?.run !== 'string' || !context.run)) throw new Failure('proxy-key-refused');
        modelKeyId = keyForSubject(business, { ...subject, ...context });
      } catch (e) { if (LOOKUP_CODES.includes(e?.code)) code = e.code; }
      if (typeof modelKeyId === 'string' && HEX12.test(modelKeyId)) lineages = lineagesFor(business);
      // The subject's key must sit in exactly one hired agent's lineage, read in the same tick.
      if (!lineages?.some(l => l.includes(modelKeyId))) {
        log({ event: 'limit-check-failed', business, check: null, export: name, code });
        return halt({ check: null, code });
      }
      subject = { ...subject, modelKeyId };
      noteOrphans(business, st.limits, lineages);
    }
    const args = { business, limits: st.limits, subject, records: st.records, now: new Date().toISOString(), ...(lineages ? { lineages } : {}) };
    for (const check of st.checks) {
      const { answer, code } = ask(business, check, name, args, a =>
        exact(a, [key]) ? a[key] === go : exact(a, [key, 'hit']) && a[key] === !go && validHit(a.hit, leaves));
      if (code) return halt({ check: check.name, code });
      if (answer[key] !== go) {
        log({ event: 'limit-reached', business, check: check.name, export: name, subject, ...answer.hit });
        return halt({ check: check.name, hit: answer.hit });
      }
    }
    return { [key]: go };
  }

  const subjectOf = ({ instance, provider = null, model = null, account = null }) =>
    ({ instance, provider, model, account });

  // { allow: true } | { allow: false, check, hit } | { allow: false, check, code }
  const mayLaunch = (business, subject) =>
    decide(business, subjectOf(subject), 'mayLaunch', 'allow', true, why => ({ allow: false, ...why }));
  // { stop: false } | { stop: true, check, hit } | { stop: true, check, code }
  const mustStop = (business, subject, context = {}) =>
    decide(business, subjectOf(subject), 'mustStop', 'stop', false, why => ({ stop: true, ...why }), context);

  const validEntry = (e, leaves) => {
    const leaf = plainObject(e) && Array.isArray(e.path) && leaves.find(l => same(l.path, e.path));
    if (!leaf || e.limit !== leaf.limit || e.since !== leaf.since) return false;
    if (exact(e, ['path', 'limit', 'since', 'notEnforced'])) return typeof e.notEnforced === 'string' && e.notEnforced.length > 0;
    return exact(e, ['path', 'limit', 'since', 'count', 'kinds', 'unknownTurns']) && amount(e.count) &&
      exact(e.kinds, KINDS) && KINDS.every(k => amount(e.kinds[k])) && turns(e.unknownTurns);
  };

  // One entry per limit, in file order:
  //   { path, limit, since, status: 'counted', check, count, kinds, unknownTurns }
  //   { path, limit, since, status: 'not-enforced', reason }   (reason null when no check names it)
  //   { path, limit, since, status: 'check-failed', check, code }
  // A failed check fails every limit, so a count never shows beside a check that couldn't answer.
  function report(business) {
    // An unreadable file fails the report, as it fails the verb's own read.
    const b = read(business);
    const leaves = limitLeaves(b?.launch?.limits ?? {});
    const failedAll = (check, code) => leaves.map(({ path, limit, since }) => ({ path, limit, since, status: 'check-failed', check, code }));
    let st;
    try { st = state(business, b); } catch (e) { return failedAll(null, e.code); }
    if (!st.checks) return failedAll(st.missing, 'not-loaded');
    // Row 333 R3: a key quota's count is its whole lineage's, so a report needs the lineages too.
    let lineages;
    if (st.checks.length && dailyKeyIds(st.limits).length) {
      lineages = lineagesFor(business);
      if (!lineages) return failedAll(null, 'proxy-key-refused');
      noteOrphans(business, st.limits, lineages);
    }
    const answers = [];
    for (const check of st.checks) {
      const { answer, code } = ask(business, check, 'report', { business, limits: st.limits, records: st.records, now: new Date().toISOString(), ...(lineages ? { lineages } : {}) },
        a => Array.isArray(a) && a.every(e => validEntry(e, leaves)) && new Set(a.map(e => JSON.stringify(e.path))).size === a.length);
      if (code) return failedAll(check.name, code);
      answers.push({ check: check.name, entries: answer });
    }
    return leaves.map(({ path, limit, since }) => {
      const found = answers.flatMap(({ check, entries }) => entries.filter(e => same(e.path, path)).map(e => ({ check, e })));
      const counted = found.find(({ e }) => Object.hasOwn(e, 'count'));
      if (counted) {
        const { count, kinds, unknownTurns } = counted.e;
        return { path, limit, since, status: 'counted', check: counted.check, count, kinds, unknownTurns };
      }
      const named = found.find(({ e }) => Object.hasOwn(e, 'notEnforced'));
      return { path, limit, since, status: 'not-enforced', reason: named ? named.e.notEnforced : null };
    });
  }

  // Row 333 R3: the daily key quotas that match no hired agent, for the Console usage view. Null
  // when the file or the lineages can't be read; the view then says nothing either way.
  function orphanedKeys(business) {
    let b;
    try { b = read(business); } catch { return null; }
    return observeOrphans(business, b?.launch?.limits ?? {});
  }

  return { refresh, mayLaunch, mustStop, report, orphanedKeys };
}

// Host metadata adapters, extracted from runtime so ownership cases are testable. records is
// Credentials.status(): metadata only, never a key.
// Row 333 R3: an agent's lineage is its current key id, its retiring ids and its earlier ids.
const lineageOf = record => [record.keyId, ...(record.retiring ?? []).map(r => r.keyId), ...(record.previous ?? []).map(r => r.keyId)];
export function modelUsageKey(records, runs, business, subject) {
  const fail = (code = 'proxy-key-refused') => { throw new Failure(code); };
  const instance = business + '/' + subject.instance;
  const owners = records.filter(r => r.service === 'model' && r.instance === instance);
  // Row 333 R3 ruling 2: an agent with no model key (no proxy, or not issued) can't be matched to a
  // key quota. Still a refusal, under its own name.
  if (!owners.length || (owners.length === 1 && owners[0].keyId === undefined)) fail('key-quota-no-model-key');
  if (owners.length !== 1) fail();
  const record = owners[0], ids = lineageOf(record);
  if (ids.some(id => typeof id !== 'string' || !HEX12.test(id)) || new Set(ids).size !== ids.length) fail();
  // A key ID must identify exactly this owner, over every generation either agent ever held.
  for (const other of records.filter(r => r.service === 'model' && r !== record))
    if (lineageOf(other).some(id => ids.includes(id))) fail();
  if (!subject.run) return record.keyId;
  // A run is bound to a key the proxy still accepts: current or retiring, never an earlier one.
  const run = runs.get(business + ':' + subject.run);
  if (!run || run.instance !== instance || ![record.keyId, ...(record.retiring ?? []).map(r => r.keyId)].includes(run.keyId)) fail();
  return run.keyId;
}
// Every hired agent's lineage in one business. Any malformed id, or an id in two lineages
// anywhere on the host, refuses.
export function modelKeyLineages(records, business) {
  const seen = new Set(), out = [];
  for (const record of records.filter(r => r.service === 'model' && r.keyId !== undefined)) {
    const ids = lineageOf(record);
    for (const id of ids) {
      if (typeof id !== 'string' || !HEX12.test(id) || seen.has(id)) throw new Failure('proxy-key-refused');
      seen.add(id);
    }
    if (typeof record.instance === 'string' && record.instance.startsWith(business + '/')) out.push(ids);
  }
  return out;
}
// Row 333 R4 (#1831): the settings writer's lookup. The lineage in `business` that holds keyId, or
// [] when none does. Ids only, from the bus's own records; the same checks refuse.
export function modelKeyLineage(records, business, keyId) {
  return modelKeyLineages(records, business).find(l => l.includes(keyId)) ?? [];
}