cache-keepalive

1import type {EngineInterface, Register, Timer} from "claude-code";
2
3const MINUTE_MS = 60_000;
4const POLL_MS = 30_000;
5// A refresh goes out this long before the cache would expire.
6const LEAD_MS = {"1h": 5 * MINUTE_MS, "5m": MINUTE_MS};
7const TTL_MS = {"1h": 60 * MINUTE_MS, "5m": 5 * MINUTE_MS};
8// Rows read back from the transcript's end when looking for the cache tier.
9const TIER_SCAN_ROWS = 400;
10// A refresh that read less than this from the cache found it already cold.
11const MIN_CACHE_READ_TOKENS = 1_000;
12const PING_PROMPT = "Keep-alive ping. Reply with exactly: ok";
13
14type CacheTier = "1h" | "5m";
15
16const config = {
17  maxIdleMs: 3 * 60 * MINUTE_MS,
18  // Overrides the tier's TTL so a live test can shorten it.
19  ttlOverrideMs: undefined as number | undefined,
20};
21
22// The module's own, so a reload starts it over.
23const idle = {
24  isDisabled: false,
25  isInTurn: false,
26  isPinging: false,
27  isStopped: true,
28  // The last time the cache was read: the turn's end or the latest refresh.
29  lastReadAt: undefined as number | undefined,
30  lastTurnEndAt: undefined as number | undefined,
31  refreshes: 0,
32  tier: "1h" as CacheTier,
33  transcriptPath: undefined as string | undefined,
34};
35
36let poller: Timer | undefined;
37let lastPollAt = 0;
38
39const hoursToMs = (value: unknown, fallbackMs: number) => {
40  const hours = Number(value);
41  return Number.isFinite(hours) && hours > 0
42    ? hours * 60 * MINUTE_MS
43    : fallbackMs;
44};
45
46// The tier of the newest main-thread cache write; a turn that only read the
47// cache writes nothing, so older rows are searched until one wrote.
48const cacheTier = (transcript: string): CacheTier => {
49  const rows = transcript.split("\n");
50  for (
51    let i = rows.length - 1;
52    i >= 0 && i >= rows.length - TIER_SCAN_ROWS;
53    i--
54  ) {
55    const row = rows[i];
56    if (row === undefined || !row.includes('"cache_creation"')) continue;
57    try {
58      const entry = JSON.parse(row);
59      if (entry.type !== "assistant" || entry.isSidechain) continue;
60      const written = entry.message?.usage?.cache_creation;
61      if (written?.ephemeral_1h_input_tokens > 0) return "1h";
62      if (written?.ephemeral_5m_input_tokens > 0) return "5m";
63    } catch {
64      // a row still being written
65    }
66  }
67  // Unknown: a 5m cache found cold at the first refresh stops the mod anyway.
68  return "1h";
69};
70
71const transcriptPath = async ($: EngineInterface) => {
72  if (idle.transcriptPath !== undefined) return idle.transcriptPath;
73  const home =
74    (await $.env.get("USERPROFILE")) ?? (await $.env.get("HOME")) ?? "";
75  const project = (await $.session.cwd()).replace(/[^a-zA-Z0-9]/g, "-");
76  return `${home}/.claude/projects/${project}/${await $.session.id()}.jsonl`;
77};
78
79const readTier = async ($: EngineInterface): Promise<CacheTier> => {
80  const path = await transcriptPath($);
81  // A long transcript is past the 4 MiB that $.fs.read takes, so read its tail.
82  try {
83    const tail = await $.process.run([
84      "tail",
85      "-n",
86      String(TIER_SCAN_ROWS),
87      path,
88    ]);
89    if (tail.exitCode === 0) return cacheTier(tail.stdout);
90  } catch {
91    // no tail on PATH
92  }
93  try {
94    return cacheTier(String(await $.fs.read(path)));
95  } catch {
96    return "1h";
97  }
98};
99
100const ttlMs = () => config.ttlOverrideMs ?? TTL_MS[idle.tier];
101const leadMs = () => Math.min(LEAD_MS[idle.tier], ttlMs() / 2);
102
103const stop = ($: EngineInterface, why: string) => {
104  idle.isStopped = true;
105  $.ui.log(
106    `cache-keepalive: stopped after ${idle.refreshes} refreshes, ${why}`,
107  );
108};
109
110const refreshIfDue = async ($: EngineInterface, now: number) => {
111  if (
112    idle.isStopped ||
113    idle.lastReadAt === undefined ||
114    idle.lastTurnEndAt === undefined
115  )
116    return;
117  if (now - idle.lastReadAt < ttlMs() - leadMs()) return;
118  // Past the TTL (the machine slept) a refresh would re-cache the whole
119  // transcript, which is what the person's next prompt does anyway.
120  if (now - idle.lastReadAt >= ttlMs())
121    return stop($, "the cache was already cold");
122  if (now - idle.lastTurnEndAt >= config.maxIdleMs)
123    return stop($, "the idle limit was reached");
124  idle.isPinging = true;
125  try {
126    const reply = await $.model.fork({prompt: PING_PROMPT});
127    if (!reply.isAnswered)
128      return stop($, `the refresh failed (${reply.reason})`);
129    if (reply.usage.cache_read_input_tokens < MIN_CACHE_READ_TOKENS) {
130      return stop($, "the refresh found the cache cold");
131    }
132    idle.lastReadAt = now;
133    idle.refreshes++;
134    $.ui.log(
135      `cache-keepalive: refreshed, ${reply.usage.cache_read_input_tokens} tokens read from cache`,
136    );
137  } finally {
138    idle.isPinging = false;
139  }
140};
141
142const poll = async ($: EngineInterface) => {
143  const now = await $.clock.now();
144  lastPollAt = now;
145  if (idle.isInTurn || idle.isPinging || idle.isDisabled) return;
146  await refreshIfDue($, now);
147};
148
149// Starts the poll, or starts it again when it has gone quiet: no session.start
150// fires after /clear, so a turn is the next chance to notice.
151const ensurePoller = async ($: EngineInterface) => {
152  const now = await $.clock.now();
153  if (poller !== undefined && now - lastPollAt < 3 * POLL_MS) return;
154  poller?.cancel();
155  lastPollAt = now;
156  poller = $.clock.every(POLL_MS, () => void poll($));
157};
158
159const keepaliveCommand = (args: string) => {
160  const word = args.trim().toLowerCase();
161  if (word === "off") {
162    idle.isDisabled = true;
163    return "cache-keepalive is off for this session.";
164  }
165  if (word === "on") {
166    idle.isDisabled = false;
167    return "cache-keepalive is on; it starts when the next turn ends.";
168  }
169  const state = idle.isDisabled
170    ? "off"
171    : idle.isStopped
172      ? "waiting for the next turn"
173      : "keeping warm";
174  return `cache-keepalive: ${state}, ${idle.refreshes} refreshes this idle gap, tier ${idle.tier}. Usage: /keepalive off | on`;
175};
176
177export const register: Register = (on, options) => {
178  config.maxIdleMs = hoursToMs(options.maxIdleHours, 3 * 60 * MINUTE_MS);
179
180  on("session.start", async ($, e, next) => {
181    const ttlMinutes = Number(await $.env.get("CACHE_KEEPALIVE_TTL_MIN"));
182    config.ttlOverrideMs = ttlMinutes > 0 ? ttlMinutes * MINUTE_MS : undefined;
183    await $.command.register({
184      argumentHint: "off | on",
185      description: "Cache keep-alive: off or on for this session, or its state",
186      name: "keepalive",
187    });
188    await ensurePoller($);
189    return next(e);
190  });
191
192  on("command.run", {command: "keepalive"}, async (_$, e) => ({
193    text: keepaliveCommand(e.args),
194  }));
195
196  // Turns raised while the refresh forks are the mod's own, never the person's.
197  on("turn.start", (_$, e, next) => {
198    if (e.agentId === undefined && !idle.isPinging) idle.isInTurn = true;
199    return next(e);
200  });
201
202  on("turn.complete", async ($, e, next) => {
203    if (e.agentId === undefined && !idle.isPinging) {
204      const now = await $.clock.now();
205      idle.isInTurn = false;
206      idle.lastTurnEndAt = now;
207      idle.lastReadAt = now;
208      idle.refreshes = 0;
209      idle.isStopped = false;
210      // The turn's own cache write may reach the transcript a moment later;
211      // an older row of the same session carries the same tier.
212      idle.tier = await readTier($);
213      await ensurePoller($);
214    }
215    return next(e);
216  });
217
218  on("classic.Stop", (_$, e, next) => {
219    idle.transcriptPath = e.transcript_path;
220    return next(e);
221  }).catch((_$, e, next) => next(e));
222};
223