inline-headroom

1import type { EngineInterface, Register } from "claude-code";
2import { RETAIN_DAYS, VIEWS, cacheHitRatio, clampEffort, dayKey, findVolatile, foldPending, isCacheDrop, isToolError, toCount, viewTitle, windowStart, zeroCounters } from "./policy.mjs";
3import type { Counters, View, VolatileFinding } from "./policy.mjs";
4
5// Module state resets on hot reload and on an options change (the engine reloads the module).
6// Persisted totals survive: each load writes its own rows under a new WRITER.
7let toolErrored = false; // any main-loop tool error since the last main-loop step
8// The Session view: this session's counters, in memory.
9const stats = {
10  steps: 0,
11  clamped: 0,
12  cacheDrops: 0,
13  lastHit: undefined as number | undefined,
14  volatile: [] as VolatileFinding[],
15};
16const WRITER = Math.random().toString(36).slice(2).padEnd(8, "0"); // this module instance's stats rows
17const pending: Counters = zeroCounters(); // counter deltas since the last fold
18const days: Record<string, Counters> = {}; // this writer's per-day totals that may still need writing
19let flushing: Promise<void> = Promise.resolve(); // serializes stats_put: a newer snapshot always lands after an older one
20let view: View = "session"; // the pane's view; every /headroom resets it
21type Sums = { view: View; today?: string; data?: Counters; error?: string };
22let sums: Sums | undefined; // the last aggregate view's totals, fetched outside render
23
24const pct = (n: number | undefined): string => (n === undefined ? "–" : `${Math.round(n * 100)}%`);
25const PANE = "headroom"; // the /headroom pane's id (1-64 of letters, digits, _ and -)
26const message = (err: unknown): string => (err instanceof Error ? err.message : String(err));
27
28// The engine refuses $.<noun> as a bare value, so same-file helpers take the whole $.
29// Resolves the op's result; rejects with the server's message on a refusal or an error result.
30const callStorage = async ($: EngineInterface, tool: string, args: Record<string, unknown>): Promise<unknown> => {
31  const c = await $.mcp.connect("storage"); // the key in this plugin's .mcp.json; the engine namespaces it and returns the namespaced `server` that `call` takes
32  if (!c.isConnected) throw new Error(c.message);
33  const r = await $.mcp.call(c.server, tool, args);
34  if (r.isError) throw new Error(r.content[0]?.text ?? "storage call failed");
35  // structuredContent is documented only for tools with an outputSchema; the text block carries the same JSON.
36  return r.structuredContent ?? JSON.parse(r.content[0]?.text ?? "null");
37};
38
39// Fetches one aggregate view's sums into `sums`; a result that arrives after the view changed is dropped.
40const loadSums = async ($: EngineInterface, now: number, v: View): Promise<void> => {
41  const today = dayKey(now);
42  if (sums?.view !== v) sums = { view: v, today }; // loading (a refresh keeps the old numbers on screen)
43  let next: Sums;
44  try {
45    const span = VIEWS.find((x) => x.id === v)?.days ?? 1;
46    next = { view: v, today, data: (await callStorage($, "stats_sum", { since: windowStart(today, span) })) as Counters };
47  } catch (err) {
48    next = { view: v, today, error: message(err) };
49  }
50  if (view === v) sums = next;
51};
52
53export const register: Register = (on, options) => {
54  const effortOn = options.effort_routing_enabled !== false;
55  const cacheOn = options.cache_aligner_enabled !== false;
56  // Unset means the manifest default (true), like the sibling toggles; the server stays fail-closed.
57  const storageOn = options.storage_enabled !== false;
58
59  on("session.start", async ($, e, next) => {
60    await $.command.register({
61      name: "headroom",
62      description: "inline-headroom stats: effort clamps, cache-hit drops, volatile prompt values",
63    });
64    return next(e);
65  });
66
67  // One row per finding, so a long list wraps per row instead of one clipped line.
68  const sessionLines = (): string[] => [
69    `effort routing: ${effortOn ? "on" : "off"} · main-loop steps ${stats.steps} · clamped ${stats.clamped}`,
70    `cache aligner: ${cacheOn ? "on" : "off"} · last hit ${pct(stats.lastHit)} · drops ${stats.cacheDrops}`,
71    stats.volatile.length ? "volatile shared values:" : "volatile shared values: none",
72    ...stats.volatile.map((v) => `  ${v.id} ${v.kind} ${v.sample}`),
73  ];
74
75  // Today / 7 days / 30 days: totals across every session on this host, read from storage.
76  const aggregateLines = (): string[] => {
77    if (!storageOn) return ["storage is off (storage_enabled is not true): only the Session view is kept"];
78    if (sums?.view !== view) return ["loading…"];
79    if (sums.error !== undefined) return [`storage unavailable: ${sums.error}`];
80    if (!sums.data) return ["loading…"];
81    const d = sums.data;
82    return [`main-loop steps ${d.steps} · clamped ${d.clamped}`, `cache hit ${pct(cacheHitRatio(d))} · drops ${d.cache_drops}`];
83  };
84
85  on("command.run", { command: "headroom" }, async ($) => {
86    view = "session"; // every /headroom opens on the default view
87    $.ui.invalidate("ui.render"); // an already-open pane redraws on it
88    const r = await $.ui.open({ id: PANE, title: "Headroom", focus: true, closeOnEscape: true });
89    // Pane placed: print nothing (no transcript line, nothing in the model's context).
90    // Not placed (headless/SDK, narrow terminal): fall back to the plain text, always the Session view.
91    return r.isPlaced ? {} : { text: [`showing: ${viewTitle("session")}`, ...sessionLines()].join("\n") };
92  });
93
94  on("ui.render", { component: "Pane", requestId: PANE }, async ($, e) => {
95    const { Box, Button, Text } = $.ui.resolve(e);
96    // Every way to press a Button raises the same onPress; the active view's label is drawn at full strength.
97    const switcher = Box({
98      flexDirection: "row",
99      gap: 2,
100      flexWrap: "wrap",
101      children: VIEWS.map((v, i) =>
102        Button({
103          key: v.id,
104          label: v.label,
105          hotkey: String(i + 1),
106          plain: true,
107          dimColor: v.id !== view,
108          onPress: async () => {
109            // Nobody awaits this handler, so it must never reject (an unhandled rejection can be fatal).
110            try {
111              view = v.id;
112              $.ui.invalidate("ui.render");
113              if (v.id === "session" || !storageOn) return;
114              await loadSums($, await $.clock.now(), v.id);
115              $.ui.invalidate("ui.render");
116            } catch (err) {
117              // The clock or an invalidate failed outside loadSums' own error handling: say why instead of "loading…".
118              if (view === v.id) sums = { view: v.id, error: message(err) };
119              try {
120                $.ui.invalidate("ui.render");
121              } catch {
122                // $ itself is refused: nothing is left to redraw with
123              }
124            }
125          },
126        }),
127      ),
128    });
129    const rows = [`showing: ${viewTitle(view, sums?.view === view ? sums.today : undefined)}`, ...(view === "session" ? sessionLines() : aggregateLines())];
130    return Box({ flexDirection: "column", children: [switcher, ...rows.map((s) => Text({ children: [s] }))] });
131  });
132
133  if (effortOn) {
134    // Observe AFTER the tool ran. Subagent calls are ignored; under parallel
135    // calls the flag is sticky, so a single error disables the next clamp.
136    on("tool.call", async (_$, e, next) => {
137      try {
138        const r = await next(e);
139        if (!e.agentId && isToolError(r)) toolErrored = true;
140        return r;
141      } catch (err) {
142        // A tool that throws instead of returning an error result is still a failure.
143        if (!e.agentId) toolErrored = true;
144        throw err;
145      }
146    });
147  }
148
149  on("turn.step", async function* ($, e, next) {
150    if (e.agentId) return yield* next(e);
151    stats.steps += 1;
152    pending.steps += 1;
153    try {
154      let ev = e;
155      // ponytail: mid-turn user input (a queued command) arriving at index > 0 is
156      // still treated as mechanical; detect it via a prompt/queue event if that matters.
157      if (effortOn && e.index > 0 && !toolErrored) {
158        const to = clampEffort(e.effort);
159        if (to !== undefined) {
160          ev = { ...e, effort: to };
161          stats.clamped += 1;
162          pending.clamped += 1;
163        }
164      }
165      toolErrored = false; // consumed per step; index 0 resets it too
166      const result = yield* next(ev);
167      if (cacheOn && result?.usage) {
168        const u = result.usage;
169        pending.input_tokens += toCount(u.input_tokens);
170        pending.cache_read_input_tokens += toCount(u.cache_read_input_tokens);
171        pending.cache_creation_input_tokens += toCount(u.cache_creation_input_tokens);
172        const hit = cacheHitRatio(result.usage);
173        if (hit !== undefined) {
174          if (isCacheDrop(stats.lastHit, hit)) {
175            stats.cacheDrops += 1;
176            pending.cache_drops += 1;
177            const ids = [...new Set(stats.volatile.map((v) => v.id))];
178            $.ui.log(`cache drop ${pct(stats.lastHit)} → ${pct(hit)} (wrote ${result.usage.cache_creation_input_tokens} tok)` + (ids.length ? ` · volatile: ${ids.join(", ")}` : ""));
179          }
180          stats.lastHit = hit;
181        }
182      }
183      return result;
184    } finally {
185      // also on a failed or aborted step: steps/clamped moved before it
186      $.ui.invalidate("ui.render"); // redraw an open /headroom pane with the new stats
187    }
188  });
189
190  if (storageOn) {
191    on("turn.complete", async ($, e, next) => {
192      const r = await next(e);
193      if (e.agentId) return r; // subagent turns carry no main-loop steps
194      let now: number;
195      let today: string;
196      let purgeBefore: string;
197      let rows: ReturnType<typeof foldPending>;
198      try {
199        now = await $.clock.now();
200        today = dayKey(now);
201        purgeBefore = windowStart(today, RETAIN_DAYS);
202        rows = foldPending(days, pending, today, purgeBefore);
203      } catch {
204        return r; // stats are a side feature: a bad clock never costs the turn's result
205      }
206      // Not awaited: a cold service start (up to 3 s) never delays the turn's end.
207      flushing = flushing
208        .then(async () => {
209          try {
210            await callStorage($, "stats_put", { writer: WRITER, rows, purgeBefore });
211          } catch {
212            return; // days keeps every total: the next turn rewrites them
213          }
214          for (const d of Object.keys(days)) if (d < today) delete days[d]; // final rows, written
215          if (view === "session") return;
216          await loadSums($, now, view);
217          $.ui.invalidate("ui.render");
218        })
219        // The chain must never reject: a rejected `flushing` would skip every later flush for the
220        // module's life. A throw after the write (loadSums, invalidate, or `$` refused once the
221        // hook has returned) only loses this one pane refresh.
222        .catch(() => {});
223      return r;
224    });
225  }
226
227  // /clear and resume go on in this process under a new session id: the Session view starts over.
228  // Persisted totals and pending deltas are session-agnostic and carry on.
229  on("session.end", async ($, e, next) => {
230    Object.assign(stats, { steps: 0, clamped: 0, cacheDrops: 0, lastHit: undefined, volatile: [] });
231    $.ui.invalidate("ui.render");
232    const r = await next(e);
233    // A headless run exits after this chain: let a started stats_put finish. `flushing` never
234    // rejects, and the engine's ~1.5 s end bound cuts the wait; core's end step already ran.
235    await flushing;
236    return r;
237  });
238
239  if (cacheOn) {
240    // Detector only (upstream CacheAligner): the composed prompt is returned untouched.
241    on("prompt.compose", async ($, e, next) => {
242      const r = await next(e);
243      stats.volatile = findVolatile(r.sections);
244      $.ui.invalidate("ui.render");
245      return r;
246    });
247  }
248};
249

1// Pure policy for inline-headroom. Never reference the mods `$` here: the
2// engine's validation rejects passing the mods API into functions imported
3// from another file. All `$` use lives inside hooks in register.ts. No clock
4// reads either: time comes in as arguments.
5
6/** @typedef {'low'|'medium'|'high'|'xhigh'|'max'} Effort */
7/** @typedef {'uuid'|'iso8601'|'jwt'|'hex_hash'} VolatileKind */
8/** @typedef {{id: string, kind: VolatileKind, sample: string}} VolatileFinding */
9/** @typedef {'session'|'day'|'7d'|'30d'} View */
10/** @typedef {{steps: number, clamped: number, cache_drops: number, input_tokens: number, cache_read_input_tokens: number, cache_creation_input_tokens: number}} Counters */
11
12export const EFFORT_ORDER = /** @type {const} */ (["low", "medium", "high", "xhigh", "max"]);
13export const CACHE_DROP_THRESHOLD = 0.6;
14export const MAX_FINDINGS = 10;
15/** Pane views in hotkey order (1-4). `days`: the rolling window, today included; 0 = this session, in memory. */
16export const VIEWS = /** @type {const} */ ([
17  { id: "session", label: "Session", days: 0 },
18  { id: "day", label: "Today", days: 1 },
19  { id: "7d", label: "7 days", days: 7 },
20  { id: "30d", label: "30 days", days: 30 },
21]);
22/** Persisted days kept: the longest view's window, never less. */
23export const RETAIN_DAYS = 30;
24
25const TOKEN_SPLIT = /[\s"'`()<>[\]{},;]+/;
26const UUID_RE = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i;
27const ISO_RE = /^\d{4}-\d{2}-\d{2}(T\d{2}:\d{2}.*)?$/;
28const JWT_RE = /^[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+$/;
29const HEX_RE = /^[0-9a-f]+$/i;
30const HEX_LENGTHS = new Set([32, 40, 64]);
31// The persisted counters in mcp/server.mjs STATS_KEYS order. That file stays import-free, so the list is
32// repeated on purpose; test/inline-headroom/storage.test.mjs pins both orders.
33const COUNTER_KEYS = /** @type {const} */ (["steps", "clamped", "cache_drops", "input_tokens", "cache_read_input_tokens", "cache_creation_input_tokens"]);
34const DAY_MS = 86400000;
35
36/**
37 * Clamp-only: returns `target` only when `current` is a known effort level
38 * strictly above it; otherwise undefined (absent, numeric, unknown, at/below).
39 * @param {unknown} current
40 * @param {Effort} [target='low']
41 * @returns {Effort|undefined}
42 */
43export function clampEffort(current, target = "low") {
44  if (typeof current !== "string") return undefined;
45  const c = EFFORT_ORDER.indexOf(/** @type {Effort} */ (current));
46  return c > EFFORT_ORDER.indexOf(target) ? target : undefined;
47}
48
49/**
50 * A tool.call result is a failure when it is errored or denied.
51 * @param {Record<string, unknown>} result
52 * @returns {boolean}
53 */
54export function isToolError(result) {
55  return result.isError === true || typeof result.deny === "string";
56}
57
58/**
59 * Anchored per-token classification (upstream CacheAligner's four classes).
60 * @param {string} tok
61 * @returns {VolatileKind|undefined}
62 */
63function classify(tok) {
64  if (UUID_RE.test(tok)) return "uuid";
65  if (ISO_RE.test(tok) && !Number.isNaN(Date.parse(tok))) return "iso8601";
66  if (tok.startsWith("eyJ") && JWT_RE.test(tok)) return "jwt";
67  if (HEX_LENGTHS.has(tok.length) && HEX_RE.test(tok)) return "hex_hash";
68  return undefined;
69}
70
71/**
72 * Flags volatile values in the cacheable (`shared`) system-prompt sections.
73 * @param {readonly {id: string, text: string, scope: 'shared'|'session'}[]} sections
74 * @returns {VolatileFinding[]}
75 */
76export function findVolatile(sections) {
77  /** @type {VolatileFinding[]} */
78  const out = [];
79  for (const s of sections) {
80    if (s.scope !== "shared") continue;
81    for (const raw of s.text.split(TOKEN_SPLIT)) {
82      const tok = raw.replace(/[.:]+$/, "");
83      const kind = classify(tok);
84      if (kind === undefined) continue;
85      out.push({ id: s.id, kind, sample: tok.slice(0, 12) + "…" });
86      if (out.length >= MAX_FINDINGS) return out;
87    }
88  }
89  return out;
90}
91
92/**
93 * @param {{input_tokens: number, cache_read_input_tokens: number, cache_creation_input_tokens: number}} usage
94 * @returns {number|undefined}
95 */
96export function cacheHitRatio(usage) {
97  const read = usage.cache_read_input_tokens;
98  const total = usage.input_tokens + read + usage.cache_creation_input_tokens;
99  return total > 0 ? read / total : undefined;
100}
101
102/**
103 * @param {number|undefined} prev
104 * @param {number} hit
105 * @returns {boolean}
106 */
107export function isCacheDrop(prev, hit) {
108  return prev !== undefined && prev >= CACHE_DROP_THRESHOLD && hit < CACHE_DROP_THRESHOLD;
109}
110
111/**
112 * A usage field as a persistable counter delta: a non-negative safe integer, else 0. One missing or
113 * odd field (undefined, NaN, a fraction) must not poison the running totals, because the server
114 * rejects every stats_put that carries such a row.
115 * @param {unknown} n
116 * @returns {number}
117 */
118export function toCount(n) {
119  return typeof n === "number" && Number.isSafeInteger(n) && n >= 0 ? n : 0;
120}
121
122/** @returns {Counters} all six counters at 0 */
123export function zeroCounters() {
124  return /** @type {Counters} */ (Object.fromEntries(COUNTER_KEYS.map((k) => [k, 0])));
125}
126
127/**
128 * Local calendar day of an epoch-ms instant as YYYY-MM-DD.
129 * @param {number} ms
130 * @param {number} [offsetMin] minutes UTC minus local (Date#getTimezoneOffset); defaults to this runtime's offset at `ms`
131 * @returns {string}
132 */
133export function dayKey(ms, offsetMin = new Date(ms).getTimezoneOffset()) {
134  return new Date(ms - offsetMin * 60000).toISOString().slice(0, 10);
135}
136
137/**
138 * First day of a rolling window of `days` calendar days ending with (and including) `day`.
139 * Calendar arithmetic in UTC, so DST never shifts it.
140 * @param {string} day YYYY-MM-DD
141 * @param {number} days >= 1
142 * @returns {string}
143 */
144export function windowStart(day, days) {
145  return new Date(Date.parse(`${day}T00:00:00Z`) - (days - 1) * DAY_MS).toISOString().slice(0, 10);
146}
147
148/**
149 * The pane header's view name: "Session"; "Today (2026-10-03)"; "7 days (2026-09-27 – 2026-10-03)".
150 * Without `today` an aggregate view reads as its label alone.
151 * @param {View} view
152 * @param {string} [today] YYYY-MM-DD
153 * @returns {string}
154 */
155export function viewTitle(view, today) {
156  const { label, days } = VIEWS.find((v) => v.id === view) ?? VIEWS[0];
157  // days 0 first: windowStart(today, 0) would be tomorrow
158  if (days === 0 || today === undefined) return label;
159  if (days === 1) return `${label} (${today})`;
160  return `${label} (${windowStart(today, days)} – ${today})`;
161}
162
163/**
164 * Adds `pending` to `days[today]` (creating it), zeroes `pending` in place, drops every day outside
165 * [purgeBefore, today], and returns each remaining day as a stats_put row (copies, so later folds
166 * never change a queued write). Mutates `days` and `pending`.
167 * @param {Record<string, Counters>} days
168 * @param {Counters} pending
169 * @param {string} today YYYY-MM-DD
170 * @param {string} purgeBefore YYYY-MM-DD
171 * @returns {(Counters & {day: string})[]}
172 */
173export function foldPending(days, pending, today, purgeBefore) {
174  const total = (days[today] ??= zeroCounters());
175  for (const k of COUNTER_KEYS) {
176    total[k] += pending[k];
177    pending[k] = 0; // in place: register.ts holds pending in a const
178  }
179  // Days after today go too, so a clock moved backwards never grows a stats_put past the window.
180  for (const day of Object.keys(days)) if (day < purgeBefore || day > today) delete days[day];
181  return Object.entries(days).map(([day, c]) => ({ day, ...c }));
182}
183