cache-warmup

1import type {
2  EngineInterface,
3  ModelForkResult,
4  PluginOptions,
5  Register,
6  Timer
7} from "claude-code";
8
9// why: keeps the prompt cache warm with no band: one short request over the conversation shortly before
10// the entry lapses. Its cache read restarts the entry's life. Logic ported from vayaan-labs/cache-maxxer (MIT).
11
12const COMMAND = "cache-warmup";
13// why: one word keeps output cost near zero; the request exists only to read the cached prefix
14const PROMPT = "Reply with exactly one word: ok. Do not use any tools.";
15
16const HOUR_MS = 3_600_000;
17const FIVE_MIN_MS = 300_000;
18const TAIL_BYTES = 262_144;
19const RECHECK_MS = 10 * 60_000;
20const RETRY_MS = 15_000;
21const LEADS: Record<string, number> = {
22  "1m": 60_000,
23  "2m": 120_000,
24  "4m": 240_000,
25  "8m": 480_000
26};
27const CAPS: Record<string, number> = {
28  "1h": HOUR_MS,
29  "3h": 3 * HOUR_MS,
30  "8h": 8 * HOUR_MS
31};
32
33type Cfg = { ttl: string; lead: string; idleCap: string };
34
35// invariant: module state; a hot reload starts it over and the next main request re-arms it
36const rt = {
37  isEnabled: true,
38  startedAt: 0,
39  learnedTtlMs: null as number | null,
40  lastActivity: 0,
41  isTurnRunning: false,
42  isPinging: false,
43  // why: the entry whose ping failed, so a failure is not retried on every tick until expiry
44  skippedFor: -1,
45  timer: null as Timer | null,
46  lookup: { at: -Infinity, isRunning: false, path: "" }
47};
48
49const ttlMs = (cfg: Cfg) =>
50  cfg.ttl === "5m"
51    ? FIVE_MIN_MS
52    : cfg.ttl === "1h"
53      ? HOUR_MS
54      : (rt.learnedTtlMs ?? HOUR_MS);
55
56// why: never more than half the entry's life, so a short lead setting cannot ping right after a request
57const leadMs = (cfg: Cfg, ttl: number) =>
58  Math.min(LEADS[cfg.lead] ?? (ttl >= 600_000 ? 240_000 : 40_000), ttl / 2);
59
60type Split = {
61  ephemeral_1h_input_tokens?: unknown;
62  ephemeral_5m_input_tokens?: unknown;
63};
64const count = (n: unknown) => (typeof n === "number" && n > 0 ? n : 0);
65
66// invariant: the newest cache write's split in a transcript tail: the entry lives as long as the larger part.
67export function parseTtl(tail: string): number | null {
68  const lines = tail.split("\n");
69  for (let i = lines.length - 1; i >= 0; i--) {
70    const line = lines[i];
71    if (!line || !line.includes("cache_creation")) continue;
72    try {
73      const entry = JSON.parse(line) as {
74        message?: { usage?: { cache_creation?: Split } };
75        usage?: { cache_creation?: Split };
76      };
77      const split =
78        entry.message?.usage?.cache_creation ?? entry.usage?.cache_creation;
79      if (!split) continue;
80      const h = count(split.ephemeral_1h_input_tokens);
81      const m = count(split.ephemeral_5m_input_tokens);
82      if (h + m > 0) return m > h ? FIVE_MIN_MS : HOUR_MS;
83    } catch {
84      // why: the tail's first line is usually cut mid-JSON
85    }
86  }
87  return null;
88}
89
90async function learnTtl($: EngineInterface, cfg: Cfg) {
91  if (cfg.ttl !== "auto" || rt.lookup.isRunning) return;
92  const now = await $.clock.now();
93  if (now - rt.lookup.at < (rt.learnedTtlMs === null ? RETRY_MS : RECHECK_MS))
94    return;
95  rt.lookup.isRunning = true;
96  rt.lookup.at = now;
97  try {
98    if (rt.lookup.path === "") {
99      const id = await $.session.id();
100      const root =
101        (await $.env.get("CLAUDE_CONFIG_DIR")) ??
102        `${(await $.env.get("HOME")) ?? ""}/.claude`;
103      const found = await $.process.run(
104        [
105          "/usr/bin/find",
106          `${root}/projects`,
107          "-maxdepth",
108          "2",
109          "-name",
110          `${id}.jsonl`
111        ],
112        { timeoutMs: 15_000 }
113      );
114      rt.lookup.path =
115        found.stdout.split("\n").find((line) => line.trim() !== "") ?? "";
116    }
117    if (rt.lookup.path === "") return;
118    // why: a long session's transcript is large; the newest write is near the end
119    const tail = await $.process.run(
120      ["/usr/bin/tail", "-c", String(TAIL_BYTES), rt.lookup.path],
121      { timeoutMs: 15_000 }
122    );
123    const ms = tail.exitCode === 0 ? parseTtl(tail.stdout) : null;
124    if (ms !== null) rt.learnedTtlMs = ms;
125  } catch {
126    // why: the length stays assumed (1h) until a later look succeeds
127  } finally {
128    rt.lookup.isRunning = false;
129  }
130}
131
132function arm($: EngineInterface, cfg: Cfg) {
133  rt.timer ??= $.clock.every(1000, () => void tick($, cfg));
134}
135
136function disarm() {
137  rt.timer?.cancel();
138  rt.timer = null;
139}
140
141async function tick($: EngineInterface, cfg: Cfg) {
142  if (!rt.isEnabled || rt.startedAt === 0) return disarm();
143  if (cfg.ttl === "auto" && rt.learnedTtlMs === null) void learnTtl($, cfg);
144  const now = await $.clock.now();
145  const ttl = ttlMs(cfg);
146  const left = rt.startedAt + ttl - now;
147  // why: a ping after expiry rewrites the whole prefix at write price, which is what warming avoids
148  if (left <= 0) return disarm();
149  if (
150    left > leadMs(cfg, ttl) ||
151    rt.isTurnRunning ||
152    rt.isPinging ||
153    rt.skippedFor === rt.startedAt
154  )
155    return;
156  const cap = CAPS[cfg.idleCap];
157  if (cap !== undefined && now - rt.lastActivity > cap) return disarm();
158  const r = await ping($, now);
159  if (!r.ok) {
160    rt.skippedFor = rt.startedAt;
161    $.ui.toast(`Cache warmup skipped: ${r.reason}`);
162    return;
163  }
164  try {
165    await $.ui.notify(
166      `Cache warmed · ${Math.round(r.read / 1000)}K tokens read`,
167      { title: "Cache warmup" }
168    );
169  } catch {
170    // why: notify rejects on a hook's deny or a failed write; the cache is warm regardless
171  }
172}
173
174type PingResult = { ok: true; read: number } | { ok: false; reason: string };
175
176async function ping(
177  $: EngineInterface,
178  startedAt: number
179): Promise<PingResult> {
180  rt.isPinging = true;
181  try {
182    const r: ModelForkResult = await $.model.fork({ prompt: PROMPT });
183    if (!r.isAnswered) return { ok: false, reason: failureOf(r) };
184    if (r.usage.cache_read_input_tokens === 0)
185      return { ok: false, reason: "the request did not read the cache" };
186    if (startedAt > rt.startedAt) rt.startedAt = startedAt;
187    return { ok: true, read: r.usage.cache_read_input_tokens };
188  } catch (error) {
189    return {
190      ok: false,
191      reason: error instanceof Error ? error.message : String(error)
192    };
193  } finally {
194    rt.isPinging = false;
195  }
196}
197
198function failureOf(r: Exclude<ModelForkResult, { isAnswered: true }>): string {
199  if (r.reason === "api-error")
200    return `the API answered ${r.status ?? "with an error"} (${r.error.replace(/_/g, " ")})`;
201  if (r.reason === "nothing-to-fork") return "nothing to fork yet";
202  if (r.reason === "aborted") return "the request was cut short";
203  return "the reply was empty";
204}
205
206async function statusText($: EngineInterface, cfg: Cfg) {
207  if (!rt.isEnabled)
208    return "Cache warmup is off. /cache-warmup on to turn it on.";
209  if (rt.startedAt === 0 || !rt.timer)
210    return "Cache warmup is on; it arms on your next message.";
211  const ttl = ttlMs(cfg);
212  const left = rt.startedAt + ttl - (await $.clock.now());
213  const nextPing = Math.max(0, Math.round((left - leadMs(cfg, ttl)) / 1000));
214  return `Cache warmup is on · ${ttl === HOUR_MS ? "1h" : "5m"} cache · next ping in ${Math.floor(nextPing / 60)}m${nextPing % 60}s`;
215}
216
217async function setEnabled($: EngineInterface, cfg: Cfg, isEnabled: boolean) {
218  rt.isEnabled = isEnabled;
219  await $.store.set("enabled", isEnabled);
220  if (!isEnabled) return disarm();
221  rt.lastActivity = await $.clock.now();
222  if (rt.startedAt > 0) arm($, cfg);
223}
224
225export const register: Register = (on, options: PluginOptions) => {
226  const cfg: Cfg = {
227    ttl: String(options.ttl ?? "auto"),
228    lead: String(options.lead ?? "auto"),
229    idleCap: String(options.idle_cap ?? "3h")
230  };
231
232  on("session.start", async ($, e, next) => {
233    const started = await next(e);
234    rt.isEnabled = (await $.store.get("enabled")) !== false;
235    rt.lastActivity = await $.clock.now();
236    try {
237      await $.command.register({
238        name: COMMAND,
239        description: "Keep the prompt cache warm: on, off, or show status",
240        argumentHint: "[on | off]",
241        immediate: true
242      });
243    } catch {
244      // why: another mod may own the name; warming still works without the command
245    }
246    return started;
247  });
248
249  // why: a resumed conversation's cache may still be live; restore its start so it is kept warm while idle
250  on(
251    "classic.SessionStart",
252    { source: ["clear", "resume", "fork"] },
253    async ($, e, next) => {
254      rt.startedAt = 0;
255      rt.lookup.path = "";
256      if (
257        e.source !== "clear" &&
258        e.seconds_since_last_response !== undefined &&
259        !e.prompt_cache_likely_expired
260      ) {
261        rt.startedAt =
262          (await $.clock.now()) - e.seconds_since_last_response * 1000;
263        arm($, cfg);
264      }
265      return next(e);
266    }
267  );
268
269  on("session.end", async ($, e, next) => {
270    rt.isTurnRunning = false;
271    if (e.reason === "clear") rt.startedAt = 0;
272    return next(e);
273  });
274
275  on("turn.start", async ($, e, next) => {
276    rt.isTurnRunning = true;
277    rt.lastActivity = await $.clock.now();
278    return next(e);
279  });
280
281  on("turn.complete", async ($, e, next) => {
282    const done = await next(e);
283    if (!e.agentId) rt.isTurnRunning = false;
284    return done;
285  });
286
287  // invariant: only main-conversation requests move the entry; subagents cache prefixes of their own
288  on("turn.step", async function* ($, e, next) {
289    const startedAt = e.agentId ? 0 : await $.clock.now();
290    const result = yield* next(e);
291    const u = result.usage;
292    if (
293      !e.agentId &&
294      u &&
295      u.cache_read_input_tokens + u.cache_creation_input_tokens > 0
296    ) {
297      rt.startedAt = startedAt;
298      rt.skippedFor = -1;
299      if (rt.isEnabled) arm($, cfg);
300      void learnTtl($, cfg);
301    }
302    return result;
303  });
304
305  on("command.run", { command: COMMAND }, async ($, e) => {
306    const word = e.args.trim().toLowerCase();
307    if (word === "on" || word === "off")
308      await setEnabled($, cfg, word === "on");
309    return { text: await statusText($, cfg) };
310  });
311};
312