sprag-cli 3.46.2 → 3.48.0

package/src/model-alias.js

@@ -18,22 +18,30 @@
  * all that is missing is ARN → alias.
  *
  * Resolution order, cheapest first:
- *   1. not an ARN            → return the input unchanged (direct-API users
+ *   1. user override         → `modelAliases` in profile-map.json
+ *   2. gateway declaration   → `gateway.aliases` in profile-map.json, fetched
+ *                              from the LiteLLM gateway's own /model/info;
+ *                              consulted ahead of learned votes because it is
+ *                              the gateway's own declared mapping, not an
+ *                              inference drawn from transcripts
+ *   3. not an ARN            → return the input unchanged (direct-API users
  *                              must keep their existing behaviour)
- *   2. user override         → `modelAliases` in profile-map.json
- *   3. learned mapping       → profile id → role, learned from transcripts
- *   4. otherwise             → 'unknown' (never a silent Sonnet guess)
+ *   4. self-describing id    → a profile id that already spells out a family
+ *                              (`foundation-model` ARNs, cross-region ids)
+ *   5. learned mapping       → profile id → role, learned from transcripts
+ *   6. otherwise             → 'unknown' (never a silent Sonnet guess)
  *
  * A profile id is never hardcoded here. Ids differ per account and change
  * with gateway config, and the ARN embeds a 12-digit AWS account id — this
  * package is published to npm, so neither may live in the source.
  */
 
-import { readFileSync, writeFileSync, mkdirSync, existsSync, createReadStream } from 'node:fs';
+import { readFileSync, writeFileSync, mkdirSync, existsSync, createReadStream, statSync, renameSync } from 'node:fs';
 import { readFile, readdir } from 'node:fs/promises';
 import { createInterface } from 'node:readline';
 import { join, dirname, basename } from 'node:path';
 import { userDataDir, claudeUserDir } from './paths.js';
+import { gatewayBase } from './gateway-auth.js';
 
 /** Marker returned when a gateway id could not be resolved to a model. */
 export const UNKNOWN_MODEL = 'unknown';
@@ -109,9 +117,19 @@ export function loadProfileMap() {
   try {
     data = JSON.parse(readFileSync(profileMapPath(), 'utf8'));
   } catch { /* absent on first run, and unreadable is not fatal */ }
+  // `gateway` is written by `sprag profile-map --refresh` (src/litellm-models.js),
+  // never by this module. Shape it defensively: a half-written or
+  // hand-edited file must fall back to "no gateway data" rather than throw
+  // or feed garbage into gatewayAlias() below.
+  const gw = data.gateway;
+  const gateway = (gw && typeof gw === 'object'
+    && typeof gw.base === 'string' && gw.aliases && typeof gw.aliases === 'object')
+    ? { ...gw, skipped: Array.isArray(gw.skipped) ? gw.skipped : [] }
+    : null;
   cached = {
     version: 1,
     modelAliases: data.modelAliases && typeof data.modelAliases === 'object' ? data.modelAliases : {},
+    gateway,
     learned: data.learned && typeof data.learned === 'object' ? data.learned : {},
     learnedAt: data.learnedAt || null,
     scannedSessions: data.scannedSessions || 0,
@@ -120,10 +138,53 @@ export function loadProfileMap() {
 }
 
 export function saveProfileMap(map) {
+  writeProfileMapFile(map);
+  cached = map;
+}
+
+/**
+ * Write the file so a reader never sees a half-written one: serialize to a
+ * sibling temp path, then rename, which is atomic within a directory. The old
+ * in-place `writeFileSync` could be interrupted mid-write, and `loadProfileMap`
+ * treats a corrupt file as an empty map — so a crash at the wrong moment
+ * silently discarded the user's hand-written `modelAliases`.
+ */
+function writeProfileMapFile(map) {
   const dir = userDataDir();
   mkdirSync(dir, { recursive: true });
-  writeFileSync(profileMapPath(), JSON.stringify(map, null, 2));
-  cached = map;
+  const target = profileMapPath();
+  const tmp = `${target}.tmp-${process.pid}`;
+  writeFileSync(tmp, JSON.stringify(map, null, 2));
+  renameSync(tmp, target);
+}
+
+/**
+ * Merge `patch`'s top-level keys into whatever is on disk RIGHT NOW, then
+ * write.
+ *
+ * Two producers update this file from separate detached processes, and each
+ * owns different keys: route-scan's learner writes `learned`/`learnedAt`/
+ * `scannedSessions`, and the gateway refresh writes `gateway`. Both used to
+ * spread a snapshot taken when their process started (`{...loadProfileMap()}`)
+ * and write the whole document back, so whichever finished last erased the
+ * other's section. A scan takes about half a second and the gateway refresh
+ * makes a network call inside the same window, so the overlap was not
+ * theoretical — and the damage is silent: resolution simply falls back to
+ * 'unknown' with nothing in any log to say why.
+ *
+ * Re-reading immediately before the rename narrows the window from the length
+ * of the whole operation to the few microseconds between read and rename. That
+ * is not a lock, and a true one would need an O_EXCL lockfile plus stale-lock
+ * recovery; the remaining race can only lose a section that was written inside
+ * those microseconds, and the next refresh restores it.
+ */
+export function updateProfileMap(patch) {
+  cached = null;
+  const fresh = loadProfileMap();
+  const next = { ...fresh, ...patch };
+  writeProfileMapFile(next);
+  cached = next;
+  return next;
 }
 
 /** Drop the memoized map. Tests use this after pointing paths elsewhere. */
@@ -154,6 +215,58 @@ function overrideAlias(model, map) {
   return null;
 }
 
+/**
+ * Alias declared by the gateway itself (LiteLLM's `/model/info`, cached
+ * under `gateway` in profile-map.json) rather than inferred from transcript
+ * votes. Consulted ahead of the learned-vote path for exactly that reason:
+ * this is the gateway telling us its own configuration, not a guess built
+ * from run counts.
+ *
+ * Isolated per gateway address the same way the budget cache is (see
+ * litellm-budget.js) — a table fetched from one `ANTHROPIC_BASE_URL` must
+ * never answer for a different one, so a base mismatch is treated the same
+ * as no gateway data at all.
+ *
+ * Looked up three ways, because the string this function receives is not
+ * always the string `/model/info` reported: transcripts append a `[1m]`
+ * context-window suffix that the gateway's response never carries, and the
+ * input here is sometimes a full ARN where `/model/info` only ever named
+ * the bare profile id. Fails quiet throughout — any shape mismatch (no
+ * gateway section, a different base, no matching key) returns null and the
+ * caller falls through to the resolution steps that existed before this one.
+ *
+ * Freshness matters as much as the base. This step outranks learned votes
+ * because the gateway DECLARES the mapping rather than inferring it, and that
+ * argument only holds while the declaration is current. A machine whose
+ * refresh has been failing for weeks (permissions revoked, `/model/info` now
+ * 403, the proxy moved) would otherwise let a months-old table outrank
+ * thousands of accumulating votes forever, which is the opposite of what the
+ * ordering was chosen for. Past the grace window the table steps aside and
+ * learning answers again; it is not deleted, because the next successful
+ * refresh should be able to reuse the file, and `sprag profile-map` still
+ * shows it with its age.
+ */
+const MODEL_MAP_GRACE_MS = 14 * 24 * 60 * 60 * 1000;
+
+function gatewayAlias(model, map, env, now = Date.now()) {
+  const gw = map.gateway;
+  if (!gw || !gw.aliases || gw.base !== gatewayBase(env)) return null;
+  const fetchedAt = Date.parse(gw.fetchedAt);
+  if (!Number.isFinite(fetchedAt) || now - fetchedAt > MODEL_MAP_GRACE_MS) return null;
+  const bare = model.replace(/\[[^\]]*\]$/, '');
+  const pid = profileIdFrom(model);
+  // `hasOwn` rather than a bare index: this object comes from JSON.parse, so
+  // it carries Object.prototype, and a transcript model id of 'constructor'
+  // or 'toString' would otherwise resolve to a function.
+  for (const key of [model, bare, pid]) {
+    if (key && Object.hasOwn(gw.aliases, key)) {
+      const alias = gw.aliases[key];
+      if (typeof alias === 'string' && alias) return alias;
+    }
+  }
+  return null;
+}
+
 // ── resolution ───────────────────────────────────────────────────────────
 
 /**
@@ -176,6 +289,13 @@ export function resolveModelAlias(rawModel, { env = process.env } = {}) {
   const override = overrideAlias(model, map);
   if (override) return override;
 
+  // The gateway's own declared mapping runs next, and — like the override
+  // above — for every id, not just ARNs. A house alias with no family name
+  // (`prod-large`) is never an ARN, so it would otherwise fall straight
+  // through the isGatewayModelId check below to a silent Sonnet default.
+  const gatewayHit = gatewayAlias(model, map, env);
+  if (gatewayHit) return gatewayHit;
+
   if (!isGatewayModelId(model)) return model;
 
   const pid = profileIdFrom(model);
@@ -423,6 +543,37 @@ export function tallyVotes(votes, { minVotes = MIN_VOTES, minAgreement = MIN_AGR
   return learned;
 }
 
+/**
+ * Sort transcript paths newest-first, so the `maxSessions` truncation in
+ * learnProfileMapping keeps the most RECENT sessions.
+ *
+ * This used to be a caller contract ("newest first" in the JSDoc), and the
+ * only caller broke it: `runRouteScan` hands over `discoverSessionFiles()`
+ * output, which parser.js sorts OLDEST first because chronological order is
+ * what parsing wants. Learning then read the oldest `maxSessions` files and
+ * never saw recent delegations, so a profile that only appears in the last few
+ * days stayed unresolved and every run on it dropped out of the aggregate.
+ * Measured on one 14-day window: 95 transcripts, of which learning read
+ * 09-04 through 09-15 and discarded the three most recent days — including the
+ * sessions holding the haiku evidence it needed.
+ *
+ * Sorting here rather than at the call site puts the invariant in the same
+ * function as the truncation that depends on it, so a future caller cannot
+ * reintroduce the bug. Cost is one stat() per candidate path, against a scan
+ * that already stats every file to build the window.
+ */
+function newestFirst(paths) {
+  return paths
+    .map((p) => {
+      let mtime = 0;
+      // An unreadable path sorts last and then fails harmlessly in the loop.
+      try { mtime = statSync(p).mtimeMs; } catch { /* keep 0 */ }
+      return { p, mtime };
+    })
+    .sort((a, b) => b.mtime - a.mtime)
+    .map((e) => e.p);
+}
+
 /**
  * Learn profile id → role from transcripts and persist the result.
  *
@@ -432,7 +583,8 @@ export function tallyVotes(votes, { minVotes = MIN_VOTES, minAgreement = MIN_AGR
  * id and the requested role identify each other.
  *
  * @param {object}   opts
- * @param {string[]} opts.sessionPaths transcripts to read, newest first
+ * @param {string[]} opts.sessionPaths transcripts to read, in any order — they
+ *                                     are sorted newest-first here, see below
  * @param {number}   opts.maxSessions  cap on files read (learning is a scan)
  * @returns {Promise<{learned: object, scannedSessions: number, gateway: boolean}>}
  */
@@ -442,7 +594,7 @@ export async function learnProfileMapping({ sessionPaths = [], maxSessions = 40
   let scanned = 0;
   let gateway = false;
 
-  for (const sessionPath of sessionPaths.slice(0, maxSessions)) {
+  for (const sessionPath of newestFirst(sessionPaths).slice(0, maxSessions)) {
     const requestedByToolUse = new Map();
     let sawGateway = false;
     try {
@@ -472,13 +624,15 @@ export async function learnProfileMapping({ sessionPaths = [], maxSessions = 40
 
   const learned = tallyVotes(votes);
   const map = loadProfileMap();
-  const next = {
-    ...map,
-    learned,
-    learnedAt: new Date().toISOString(),
-    scannedSessions: scanned,
-  };
   // Nothing to record on a non-gateway machine — do not create the file there.
-  if (gateway || Object.keys(map.learned || {}).length) saveProfileMap(next);
+  // Writes only this learner's own keys, so a gateway refresh that landed while
+  // this scan was running keeps its section (see updateProfileMap).
+  if (gateway || Object.keys(map.learned || {}).length) {
+    updateProfileMap({
+      learned,
+      learnedAt: new Date().toISOString(),
+      scannedSessions: scanned,
+    });
+  }
   return { learned, scannedSessions: scanned, gateway };
 }