talon-agent 5.4.1 → 5.5.1

package/src/core/engine/backend-router/headroom.ts

@@ -0,0 +1,267 @@
+/**
+ * Headroom — "how much of this backend is left?", answered the same way for
+ * every backend regardless of what it can tell us about itself.
+ *
+ * Two sources, in precedence order:
+ *
+ *   - `plan`   — the backend's own subscription windows
+ *                (`UsageTelemetry.getPlanUsage`). Authoritative: it is the
+ *                provider's own count, including spend from outside Talon.
+ *   - `ledger` — Talon's local rolling token count (see `ledger.ts`) against
+ *                the operator's soft budget (`config.backendBudgets`). The
+ *                fallback for backends with no account API.
+ *
+ * A backend with neither reports `source: "none"` and headroom 1. That is a
+ * deliberate "no evidence of pressure", not a claim of capacity — the
+ * router's comparator ranks it *below* any backend with real telemetry at
+ * the same headroom, so an unmeasured backend never outranks a measured one
+ * it is tied with.
+ *
+ * Reads are cached for 60s per backend: `/usage`, the router and the
+ * `plan_usage` tool all ask, and a plan lookup can be a subprocess spawn. A
+ * failed refresh keeps the last good value and flags it `stale` rather than
+ * pretending the backend emptied.
+ */
+
+import type { PlanUsage } from "../../agent-runtime/capabilities.js";
+import type { TalonConfig } from "../../config/index.js";
+import {
+  getPooledBackend,
+  listAvailableBackends,
+} from "../backend-controller/index.js";
+import { ledgerUsage } from "./ledger.js";
+
+/** How long a headroom reading is reused before the source is asked again. */
+export const HEADROOM_CACHE_MS = 60_000;
+
+/** Where a headroom figure came from. Also its ranking priority. */
+export type HeadroomSource = "plan" | "ledger" | "none";
+
+/** The window that is closest to its limit — what the ceiling is judged on. */
+export interface LimitingWindow {
+  readonly label: string;
+  /** 0-100. */
+  readonly percent: number;
+  readonly resetsAt?: string;
+}
+
+export interface BackendHeadroom {
+  readonly id: string;
+  readonly label: string;
+  /** 0..1 — 1 is empty, 0 is at the limit. */
+  readonly headroom: number;
+  readonly limiting?: LimitingWindow;
+  readonly source: HeadroomSource;
+  /** Epoch ms of the underlying read. */
+  readonly fetchedAt: number;
+  /** True when the last refresh failed and this is the previous value. */
+  readonly stale?: boolean;
+  /**
+   * The raw plan reading this came from, when `source` is `"plan"`. Carried
+   * so `/usage` and the `plan_usage` tool render the real windows off the
+   * same (cached) fetch the router ranked on, instead of asking twice.
+   */
+  readonly plan?: PlanUsage;
+}
+
+interface CacheEntry {
+  value: BackendHeadroom;
+  /** Epoch ms the value was computed (not the same as a stale `fetchedAt`). */
+  cachedAt: number;
+}
+
+const cache = new Map<string, CacheEntry>();
+
+function clampPercent(percent: number): number {
+  if (!Number.isFinite(percent)) return 0;
+  return Math.max(0, Math.min(100, percent));
+}
+
+/** headroom = 1 − (tightest window) / 100. */
+function headroomFor(percent: number): number {
+  return Math.max(0, Math.min(1, 1 - clampPercent(percent) / 100));
+}
+
+/**
+ * The tightest of a plan's windows. `undefined` when the plan reports none,
+ * which is how a backend that answers but has nothing to say is told apart
+ * from one that never answered.
+ */
+export function limitingWindowOf(
+  usage: PlanUsage | undefined,
+): LimitingWindow | undefined {
+  if (!usage || usage.windows.length === 0) return undefined;
+  let worst = usage.windows[0] as NonNullable<(typeof usage.windows)[0]>;
+  for (const window of usage.windows) {
+    if (window.percent > worst.percent) worst = window;
+  }
+  return {
+    label: worst.label,
+    percent: clampPercent(worst.percent),
+    ...(worst.resetsAt ? { resetsAt: worst.resetsAt } : {}),
+  };
+}
+
+/** Headroom from a `PlanUsage`, or `undefined` when it carries no windows. */
+export function headroomFromPlan(
+  id: string,
+  label: string,
+  usage: PlanUsage | undefined,
+): BackendHeadroom | undefined {
+  const limiting = limitingWindowOf(usage);
+  if (!limiting || !usage) return undefined;
+  return {
+    id,
+    label,
+    headroom: headroomFor(limiting.percent),
+    limiting,
+    source: "plan",
+    fetchedAt: usage.fetchedAt,
+    plan: usage,
+  };
+}
+
+/** The soft budget an operator declared for a backend, if any. */
+function budgetFor(
+  config: TalonConfig | undefined,
+  id: string,
+): { tokensPer5h?: number; tokensPerDay?: number } | undefined {
+  const budget = config?.backendBudgets?.[id];
+  if (!budget) return undefined;
+  if (budget.tokensPer5h === undefined && budget.tokensPerDay === undefined) {
+    return undefined;
+  }
+  return budget;
+}
+
+/** Whether the operator gave this backend a local budget to measure against. */
+export function hasBudget(
+  config: TalonConfig | undefined,
+  id: string,
+): boolean {
+  return budgetFor(config, id) !== undefined;
+}
+
+/**
+ * Headroom from the local ledger. The tighter of the two configured windows
+ * wins, so a backend that is fine on the day but has just burned its 5h
+ * allowance still reads as full.
+ */
+export function headroomFromLedger(
+  id: string,
+  label: string,
+  config: TalonConfig | undefined,
+  now = Date.now(),
+): BackendHeadroom | undefined {
+  const budget = budgetFor(config, id);
+  if (!budget) return undefined;
+  const used = ledgerUsage(id, now);
+  const windows: LimitingWindow[] = [];
+  if (budget.tokensPer5h !== undefined) {
+    windows.push({
+      label: "5h (local budget)",
+      percent: clampPercent((used.tokens5h / budget.tokensPer5h) * 100),
+    });
+  }
+  if (budget.tokensPerDay !== undefined) {
+    windows.push({
+      label: "24h (local budget)",
+      percent: clampPercent((used.tokensDay / budget.tokensPerDay) * 100),
+    });
+  }
+  let worst = windows[0] as LimitingWindow;
+  for (const window of windows)
+    if (window.percent > worst.percent) worst = window;
+  return {
+    id,
+    label,
+    headroom: headroomFor(worst.percent),
+    limiting: worst,
+    source: "ledger",
+    fetchedAt: now,
+  };
+}
+
+/** The "nothing to measure" reading. Headroom 1, but lowest ranking source. */
+function unknownHeadroom(
+  id: string,
+  label: string,
+  now: number,
+): BackendHeadroom {
+  return { id, label, headroom: 1, source: "none", fetchedAt: now };
+}
+
+/** Ask a pooled backend for its plan windows. Rejects like the backend does. */
+async function readPlanUsage(id: string): Promise<PlanUsage | undefined> {
+  const backend = getPooledBackend(id);
+  const read = backend?.usage?.getPlanUsage;
+  if (!read || !backend?.usage) return undefined;
+  return read.call(backend.usage);
+}
+
+/**
+ * Headroom for one backend, cached for {@link HEADROOM_CACHE_MS}.
+ *
+ * `force` skips the cache — the `plan_usage` tool asks for a fresh read
+ * because the operator is looking at the number right now.
+ */
+export async function getBackendHeadroom(
+  id: string,
+  label: string,
+  config: TalonConfig | undefined,
+  options?: { force?: boolean; now?: number },
+): Promise<BackendHeadroom> {
+  const now = options?.now ?? Date.now();
+  const cached = cache.get(id);
+  if (!options?.force && cached && now - cached.cachedAt < HEADROOM_CACHE_MS) {
+    return cached.value;
+  }
+
+  let value: BackendHeadroom;
+  try {
+    const plan = headroomFromPlan(id, label, await readPlanUsage(id));
+    value =
+      plan ??
+      headroomFromLedger(id, label, config, now) ??
+      unknownHeadroom(id, label, now);
+  } catch {
+    // The source is unreachable this minute. Keeping the last good reading
+    // is the conservative answer: forgetting it would read as "empty" and
+    // send the next background run straight at a backend near its ceiling.
+    value = cached
+      ? { ...cached.value, stale: true }
+      : unknownHeadroom(id, label, now);
+  }
+  cache.set(id, { value, cachedAt: now });
+  return value;
+}
+
+/** Headroom for every backend the config exposes, in config order. */
+export async function collectBackendHeadroom(
+  config: TalonConfig | undefined,
+  options?: { force?: boolean; now?: number },
+): Promise<BackendHeadroom[]> {
+  const backends = listAvailableBackends(config);
+  return Promise.all(
+    backends.map(({ id, label }) =>
+      getBackendHeadroom(id, label, config, options),
+    ),
+  );
+}
+
+/** One-line rendering shared by `/usage`, `plan_usage` and the router log. */
+export function formatHeadroom(entry: BackendHeadroom): string {
+  const pct = `${Math.round(entry.headroom * 100)}%`;
+  const detail =
+    entry.source === "none"
+      ? "no usage signal"
+      : `${entry.limiting?.label ?? "window"} ${Math.round(entry.limiting?.percent ?? 0)}% used`;
+  const tag = entry.source === "ledger" ? " (local budget)" : "";
+  const stale = entry.stale ? " (stale)" : "";
+  return `${pct} — ${detail}${tag}${stale}`;
+}
+
+/** Test seam — drop every cached reading. */
+export function resetHeadroomCacheForTest(): void {
+  cache.clear();
+}

package/src/core/engine/backend-router/index.ts

@@ -0,0 +1,52 @@
+/**
+ * Plan-aware backend router.
+ *
+ *   - `ledger`   — the local rolling token count that gives budget-only
+ *                  backends a headroom signal.
+ *   - `headroom` — one comparable "how much is left" per backend, from the
+ *                  plan API where there is one and the ledger where there
+ *                  is not.
+ *   - `router`   — the decision: who runs this background job.
+ */
+
+export {
+  flushBackendLedger,
+  ledgerUsage,
+  loadBackendLedger,
+  recordBackendRunUsage,
+  recordBackendUsage,
+  resetBackendLedgerForTest,
+  tokensInWindow,
+  LEDGER_RETENTION_MS,
+  LEDGER_SHORT_WINDOW_MS,
+} from "./ledger.js";
+export {
+  collectBackendHeadroom,
+  formatHeadroom,
+  getBackendHeadroom,
+  hasBudget,
+  headroomFromLedger,
+  headroomFromPlan,
+  limitingWindowOf,
+  resetHeadroomCacheForTest,
+  HEADROOM_CACHE_MS,
+  type BackendHeadroom,
+  type HeadroomSource,
+  type LimitingWindow,
+} from "./headroom.js";
+export {
+  collectBackendUsage,
+  leadWith,
+  type BackendUsageSnapshot,
+} from "./usage.js";
+export {
+  chooseBackend,
+  resolveRoutedModel,
+  taskClassForEffort,
+  DEFAULT_CEILING_PERCENT,
+  type RouteDecision,
+  type RouteHints,
+  type RoutePurpose,
+  type RouteRequest,
+  type TaskClass,
+} from "./router.js";

package/src/core/engine/backend-router/ledger.ts

@@ -0,0 +1,248 @@
+/**
+ * Local rolling token ledger — the headroom signal for backends with no
+ * account usage API.
+ *
+ * Claude and Codex report subscription windows; `agy` has no account
+ * endpoint and `openai-agents` has no plan at all. Without a second signal
+ * the router would treat those as infinitely fresh and pile every background
+ * run onto them. So Talon counts what it spends itself: every chat turn,
+ * one-shot and sub-agent folds its token total into a per-backend ledger,
+ * and `headroom.ts` reads that against the operator's soft budget
+ * (`config.backendBudgets`).
+ *
+ * The ledger is deliberately a *local estimate*, not accounting: it only
+ * sees what this daemon ran, and a plan API always takes precedence when
+ * one exists. Entries older than the widest window (24h) are pruned on
+ * every read and write, which also bounds the file.
+ *
+ * Persistence is `~/.talon/data/backend-ledger.json`, written atomically so
+ * a restart mid-write can't leave a truncated file — a zeroed ledger would
+ * silently hand a spent backend a full headroom score.
+ */
+
+import { readFile } from "node:fs/promises";
+import { resolve } from "node:path";
+import { dirs } from "../../../util/paths.js";
+import { logWarn } from "../../../util/log.js";
+import { writePrivateJson } from "../../mesh/persist.js";
+
+/** Widest window the ledger answers for; everything older is dropped. */
+export const LEDGER_RETENTION_MS = 24 * 60 * 60_000;
+/** The short window, as `backendBudgets.tokensPer5h` measures it. */
+export const LEDGER_SHORT_WINDOW_MS = 5 * 60 * 60_000;
+
+/** How long a write is coalesced for, so a burst of turns is one fsync. */
+const FLUSH_DEBOUNCE_MS = 2_000;
+
+/** One recorded spend: when it happened and how many tokens it cost. */
+interface LedgerEntry {
+  /** Epoch ms. */
+  readonly t: number;
+  /** Total tokens (input + output + cache + thinking, as the backend counts). */
+  readonly n: number;
+}
+
+interface LedgerFile {
+  readonly version: 1;
+  readonly backends: Record<string, LedgerEntry[]>;
+}
+
+interface LedgerState {
+  entries: Map<string, LedgerEntry[]>;
+  loaded: boolean;
+  loading: Promise<void> | null;
+  flushTimer: ReturnType<typeof setTimeout> | null;
+  dirty: boolean;
+  /** Overridable for tests; resolved lazily so TALON_HOME changes are seen. */
+  path: string | null;
+}
+
+const state: LedgerState = {
+  entries: new Map(),
+  loaded: false,
+  loading: null,
+  flushTimer: null,
+  dirty: false,
+  path: null,
+};
+
+function ledgerPath(): string {
+  return state.path ?? resolve(dirs.data, "backend-ledger.json");
+}
+
+/** Drop entries that have aged out of the widest window. */
+function prune(list: LedgerEntry[], now: number): LedgerEntry[] {
+  const floor = now - LEDGER_RETENTION_MS;
+  // Entries are appended in time order, so the survivors are a suffix —
+  // but a clock step backwards can break that, hence a filter not a slice.
+  return list.filter((e) => e.t > floor);
+}
+
+function parseLedger(raw: unknown): Map<string, LedgerEntry[]> {
+  const out = new Map<string, LedgerEntry[]>();
+  if (!raw || typeof raw !== "object") return out;
+  const file = raw as Partial<LedgerFile>;
+  if (file.version !== 1 || !file.backends) return out;
+  const now = Date.now();
+  for (const [id, list] of Object.entries(file.backends)) {
+    if (!Array.isArray(list)) continue;
+    const clean = list.filter(
+      (e): e is LedgerEntry =>
+        Boolean(e) &&
+        typeof (e as LedgerEntry).t === "number" &&
+        typeof (e as LedgerEntry).n === "number" &&
+        Number.isFinite((e as LedgerEntry).t) &&
+        Number.isFinite((e as LedgerEntry).n),
+    );
+    const pruned = prune(clean, now);
+    if (pruned.length > 0) out.set(id, pruned);
+  }
+  return out;
+}
+
+/**
+ * Load the persisted ledger once per process. Idempotent and safe to call
+ * concurrently — every caller awaits the same read. A missing or corrupt
+ * file starts an empty ledger rather than failing a routing decision.
+ */
+export async function loadBackendLedger(): Promise<void> {
+  if (state.loaded) return;
+  if (state.loading) return state.loading;
+  state.loading = (async () => {
+    try {
+      const raw = await readFile(ledgerPath(), "utf8");
+      const parsed = parseLedger(JSON.parse(raw));
+      // In-process records taken while the read was in flight win: merge
+      // rather than replace, so a turn that landed during boot isn't lost.
+      for (const [id, list] of parsed) {
+        state.entries.set(id, [...list, ...(state.entries.get(id) ?? [])]);
+      }
+    } catch {
+      /* no ledger yet, or unreadable — start empty */
+    }
+    state.loaded = true;
+    state.loading = null;
+  })();
+  return state.loading;
+}
+
+function scheduleFlush(): void {
+  state.dirty = true;
+  if (state.flushTimer) return;
+  const timer = setTimeout(() => {
+    state.flushTimer = null;
+    void flushBackendLedger();
+  }, FLUSH_DEBOUNCE_MS);
+  timer.unref?.();
+  state.flushTimer = timer;
+}
+
+/** Write the ledger out now. Exported so shutdown and tests can force it. */
+export async function flushBackendLedger(): Promise<void> {
+  if (!state.dirty) return;
+  state.dirty = false;
+  const now = Date.now();
+  const backends: Record<string, LedgerEntry[]> = {};
+  for (const [id, list] of state.entries) {
+    const pruned = prune(list, now);
+    state.entries.set(id, pruned);
+    if (pruned.length > 0) backends[id] = pruned;
+  }
+  try {
+    await writePrivateJson(ledgerPath(), {
+      version: 1,
+      backends,
+    } satisfies LedgerFile);
+  } catch (err) {
+    logWarn(
+      "router",
+      `Could not persist the backend ledger: ${err instanceof Error ? err.message : String(err)}`,
+    );
+  }
+}
+
+/**
+ * Fold one run's token spend into a backend's ledger.
+ *
+ * Called from the shared turn accounting and from every one-shot completion
+ * path, so *all* backends accumulate a ledger — the plan API simply wins over
+ * it where one exists. Cheap and synchronous: the disk write is debounced.
+ */
+export function recordBackendUsage(
+  backendId: string,
+  tokens: number,
+  at = Date.now(),
+): void {
+  if (!backendId) return;
+  if (!Number.isFinite(tokens) || tokens <= 0) return;
+  // A first record before the file has been read would be overwritten by the
+  // load; kick the read off here so the merge in loadBackendLedger keeps it.
+  if (!state.loaded && !state.loading) void loadBackendLedger();
+  const list = state.entries.get(backendId) ?? [];
+  list.push({ t: at, n: Math.round(tokens) });
+  state.entries.set(backendId, prune(list, at));
+  scheduleFlush();
+}
+
+/**
+ * Fold a completed run's usage into the ledger. The same four fields every
+ * background path reports (`OneShotUsage`, `TaskUsage`), summed — cache
+ * reads included, because they still count against a subscription window.
+ */
+export function recordBackendRunUsage(
+  backendId: string,
+  usage:
+    | {
+        inputTokens?: number;
+        outputTokens?: number;
+        cacheRead?: number;
+        cacheWrite?: number;
+      }
+    | undefined
+    | null,
+  at = Date.now(),
+): void {
+  if (!usage) return;
+  const total =
+    (usage.inputTokens ?? 0) +
+    (usage.outputTokens ?? 0) +
+    (usage.cacheRead ?? 0) +
+    (usage.cacheWrite ?? 0);
+  recordBackendUsage(backendId, total, at);
+}
+
+/** Tokens a backend spent inside a window ending now. */
+export function tokensInWindow(
+  backendId: string,
+  windowMs: number,
+  now = Date.now(),
+): number {
+  const list = state.entries.get(backendId);
+  if (!list || list.length === 0) return 0;
+  const floor = now - windowMs;
+  let total = 0;
+  for (const entry of list) if (entry.t > floor) total += entry.n;
+  return total;
+}
+
+/** Both windows the budget schema knows about, for one backend. */
+export function ledgerUsage(
+  backendId: string,
+  now = Date.now(),
+): { tokens5h: number; tokensDay: number } {
+  return {
+    tokens5h: tokensInWindow(backendId, LEDGER_SHORT_WINDOW_MS, now),
+    tokensDay: tokensInWindow(backendId, LEDGER_RETENTION_MS, now),
+  };
+}
+
+/** Test seam — point the ledger at a temp file and start from empty. */
+export function resetBackendLedgerForTest(path?: string): void {
+  if (state.flushTimer) clearTimeout(state.flushTimer);
+  state.entries = new Map();
+  state.loaded = false;
+  state.loading = null;
+  state.flushTimer = null;
+  state.dirty = false;
+  state.path = path ?? null;
+}