cache-warmer

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, PluginOptions, Register, Timer } from 'claude-code'
3
4import type { WarmPhase, WarmState } from '../types'
5
6// cache-warmer
7//
8// After each main-thread turn the session's prompt cache is warm. If the
9// person goes idle, this mod touches that cache with a silent fork request
10// every `pingIntervalMinutes` (under the cache lifetime), so returning inside
11// the warm window costs a cache read, not a full re-send. Once the idle time
12// passes `warmHours`, pinging stops: a session above `minCompactContextK` is
13// compacted once while still warm, a smaller one is simply left to cool.
14//
15// It never pings a cache that has already lapsed (laptop asleep, timer late),
16// and it stops for the day once `dailyBudgetK` keep-alive spend is reached
17// across every session, counted in full-price token equivalents.
18
19type Engine = EngineInterface
20
21type Config = {
22  pingIntervalMs: number
23  cacheTtlMs: number
24  warmWindowMs: number
25  minPingTokens: number
26  minCompactTokens: number
27  dailyBudget: number
28}
29
30const INITIAL: WarmState = {
31  phase: 'idle',
32  lastActivityAt: 0,
33  lastPingAt: 0,
34  nextPingAt: 0,
35  pings: 0,
36  pingTokens: 0,
37  contextTokens: 0,
38  errors: 0,
39  note: '',
40}
41
42const warm = atom({ plugin: 'cache-warmer', key: 'warm' } as const, INITIAL)
43
44const PING_PROMPT = 'Reply with the single word: ok'
45
46const COMPACT_INSTRUCTIONS =
47  'The user is resuming this session after a long break. Keep: the current task and its goal, ' +
48  'the plan and exactly where it stands, every file touched and what changed in each, decisions ' +
49  'made and why, open items and next steps, and any commands, paths or URLs needed to continue.'
50
51// Full-price token equivalents per token, by kind (base input price = 1).
52const WEIGHT = { input: 1, cacheWrite: 1.25, cacheRead: 0.1, output: 5 }
53
54const RETRY_MS = 5 * 60_000
55const MAX_ERRORS = 3
56
57// Module state: set once per load by register(), read by the functions below.
58// A hot reload starts a fresh environment, so the pending timer is dropped
59// with it; session.start re-arms from $.state, which the host keeps.
60let cfg: Config = configFrom({})
61let timer: Timer | null = null
62
63function num(value: unknown, fallback: number): number {
64  return typeof value === 'number' && Number.isFinite(value) && value > 0 ? value : fallback
65}
66
67function configFrom(options: PluginOptions): Config {
68  return {
69    pingIntervalMs: num(options.pingIntervalMinutes, 50) * 60_000,
70    cacheTtlMs: num(options.cacheTtlMinutes, 60) * 60_000,
71    warmWindowMs: num(options.warmHours, 10) * 3_600_000,
72    minPingTokens: num(options.minPingContextK, 100) * 1000,
73    minCompactTokens: num(options.minCompactContextK, 500) * 1000,
74    dailyBudget: num(options.dailyBudgetK, 2000) * 1000,
75  }
76}
77
78function pad(n: number): string {
79  return String(n).padStart(2, '0')
80}
81
82function hhmm(ms: number): string {
83  const d = new Date(ms)
84  return `${pad(d.getHours())}:${pad(d.getMinutes())}`
85}
86
87function dateKey(ms: number): string {
88  const d = new Date(ms)
89  return `${d.getFullYear()}-${pad(d.getMonth() + 1)}-${pad(d.getDate())}`
90}
91
92function k(n: number): string {
93  return `${Math.round(n / 1000)}k`
94}
95
96function span(ms: number): string {
97  const m = Math.max(0, Math.round(ms / 60_000))
98  return m < 60 ? `${m}m` : `${Math.floor(m / 60)}h${pad(m % 60)}`
99}
100
101function budgetKey(ms: number): string {
102  return `budget:${dateKey(ms)}`
103}
104
105function offKey(ms: number): string {
106  return `off:${dateKey(ms)}`
107}
108
109function cancel(): void {
110  timer?.cancel()
111  timer = null
112}
113
114function paint($: Engine, s: WarmState): void {
115  switch (s.phase) {
116    case 'warming':
117      $.ui.status(
118        `warm · ${s.pings} ping${s.pings === 1 ? '' : 's'} · next ${hhmm(s.nextPingAt)} · ${k(s.contextTokens)}`,
119      )
120      break
121    case 'compacted':
122      $.ui.status(`compacted · ${s.note}`)
123      break
124    case 'cold':
125      $.ui.status(`cold · ${s.note}`)
126      break
127    case 'paused':
128      $.ui.status('warm off for this session · /warm on')
129      break
130    case 'stopped':
131      $.ui.status(`warm stopped · ${s.note}`)
132      break
133    default:
134      $.ui.status(undefined)
135  }
136}
137
138async function setPhase($: Engine, phase: WarmPhase, note = ''): Promise<WarmState> {
139  const s = await update($, warm, (prev): WarmState => ({ ...prev, phase, note, nextPingAt: 0 }))
140  paint($, s)
141  return s
142}
143
144function log($: Engine, text: string): void {
145  $.ui.log(`cache-warmer: ${text}`, { to: 'debug' })
146}
147
148async function spentToday($: Engine, now: number): Promise<number> {
149  return Number((await $.store.get(budgetKey(now))) ?? 0)
150}
151
152async function addSpend($: Engine, now: number, equivalents: number): Promise<number> {
153  const total = (await spentToday($, now)) + equivalents
154  await $.store.set(budgetKey(now), Math.round(total))
155  return total
156}
157
158async function isOffToday($: Engine, now: number): Promise<boolean> {
159  return (await $.store.get(offKey(now))) === true
160}
161
162async function schedule($: Engine, delayMs: number): Promise<void> {
163  cancel()
164  const now = await $.clock.now()
165  const wait = Math.max(0, delayMs)
166  const s = await update($, warm, (prev): WarmState => ({ ...prev, phase: 'warming', nextPingAt: now + wait }))
167  timer = $.clock.after(wait, () => {
168    void tick($)
169  })
170  paint($, s)
171}
172
173// Arm the first ping relative to the last time the cache was touched.
174async function armFrom($: Engine, touchedAt: number): Promise<void> {
175  const now = await $.clock.now()
176  const sinceTouch = now - touchedAt
177  if (sinceTouch >= cfg.cacheTtlMs) {
178    await setPhase($, 'cold', `cache lapsed ${span(sinceTouch - cfg.cacheTtlMs)} ago, nothing to keep warm`)
179    return
180  }
181  await schedule($, cfg.pingIntervalMs - sinceTouch)
182}
183
184async function compactWarm($: Engine, s: WarmState): Promise<void> {
185  try {
186    const r = await $.session.compact({ instructions: COMPACT_INSTRUCTIONS })
187    if (r.skip !== undefined) {
188      await setPhase($, 'stopped', `compaction skipped: ${r.skip}`)
189      return
190    }
191    const before = r.tokensBefore ?? s.contextTokens
192    const after = r.tokensAfter
193    const now = await $.clock.now()
194    const note = `${hhmm(now)} · ${k(before)}${after !== undefined ? ` → ${k(after)}` : ''} while warm`
195    await setPhase($, 'compacted', note)
196    $.ui.toast(`cache-warmer: compacted the session while its cache was warm (${k(before)})`)
197  } catch (error) {
198    await setPhase($, 'stopped', `compaction failed: ${error instanceof Error ? error.message : String(error)}`)
199  }
200}
201
202async function tick($: Engine): Promise<void> {
203  timer = null
204  try {
205    const s = await read($, warm)
206    if (s.phase !== 'warming') return
207
208    const now = await $.clock.now()
209    const idleFor = now - s.lastActivityAt
210    const sinceTouch = now - Math.max(s.lastActivityAt, s.lastPingAt)
211
212    // 1. The cache has already lapsed (the machine slept, the timer ran late).
213    //    Pinging now would pay full price to warm a cache nobody is using.
214    if (sinceTouch >= cfg.cacheTtlMs) {
215      await setPhase($, 'cold', `cache lapsed after ${span(sinceTouch)} untouched, not re-warmed`)
216      return
217    }
218
219    // 2. The warm window is over: compact a large session once, let a small one cool.
220    if (idleFor >= cfg.warmWindowMs) {
221      if (s.contextTokens >= cfg.minCompactTokens) {
222        await compactWarm($, s)
223      } else {
224        await setPhase(
225          $,
226          'stopped',
227          `${span(idleFor)} idle, ${k(s.contextTokens)} under the ${k(cfg.minCompactTokens)} compact line, left to cool`,
228        )
229      }
230      return
231    }
232
233    // 3. Switched off for the day, or the shared daily budget is spent.
234    if (await isOffToday($, now)) {
235      await setPhase($, 'stopped', 'off in every session for today (/warm off all)')
236      return
237    }
238    const spent = await spentToday($, now)
239    if (spent >= cfg.dailyBudget) {
240      await setPhase($, 'stopped', `daily keep-alive budget reached (${k(spent)} of ${k(cfg.dailyBudget)})`)
241      $.ui.toast('cache-warmer: daily keep-alive budget reached, pinging stopped for today')
242      return
243    }
244
245    // 4. Touch the cache: one tool-less request over the session's own prefix.
246    const r = await $.model.fork({ prompt: PING_PROMPT })
247    if (!r.isAnswered && r.reason === 'nothing-to-fork') {
248      await setPhase($, 'stopped', 'nothing to keep warm yet')
249      return
250    }
251    const u = r.usage
252    const equivalents =
253      u.input_tokens * WEIGHT.input +
254      u.cache_creation_input_tokens * WEIGHT.cacheWrite +
255      u.cache_read_input_tokens * WEIGHT.cacheRead +
256      u.output_tokens * WEIGHT.output
257    if (equivalents > 0) await addSpend($, now, equivalents)
258
259    if (!r.isAnswered) {
260      const errors = s.errors + 1
261      await update($, warm, (prev): WarmState => ({ ...prev, errors }))
262      log($, `ping failed (${r.reason}), attempt ${errors}`)
263      if (errors >= MAX_ERRORS) {
264        await setPhase($, 'stopped', `${errors} pings failed in a row (${r.reason})`)
265        return
266      }
267      await schedule($, RETRY_MS)
268      return
269    }
270
271    const total = u.input_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens
272    const hit = total > 0 ? u.cache_read_input_tokens / total : 0
273
274    // 5. The ping found the cache cold: paid once, stop here.
275    if (hit < 0.5) {
276      await setPhase($, 'cold', `ping at ${hhmm(now)} found the cache cold (${Math.round(hit * 100)}% served), stopped`)
277      $.ui.toast('cache-warmer: the cache was already cold, keep-alive stopped for this idle period')
278      return
279    }
280
281    await update($, warm, (prev): WarmState => ({
282      ...prev,
283      pings: prev.pings + 1,
284      pingTokens: prev.pingTokens + u.cache_read_input_tokens,
285      lastPingAt: now,
286      contextTokens: total > 0 ? total : prev.contextTokens,
287      errors: 0,
288    }))
289    log($, `ping ${s.pings + 1}: ${k(u.cache_read_input_tokens)} served from cache, ${Math.round(hit * 100)}% hit`)
290    await schedule($, cfg.pingIntervalMs)
291  } catch (error) {
292    log($, `tick failed: ${error instanceof Error ? error.message : String(error)}`)
293  }
294}
295
296async function statusText($: Engine): Promise<string> {
297  const s = await read($, warm)
298  const now = await $.clock.now()
299  const spent = await spentToday($, now)
300  const off = await isOffToday($, now)
301  const lines: (string | null)[] = [
302    `cache-warmer · ${s.phase}${s.note ? ` · ${s.note}` : ''}`,
303    `  context: ${k(s.contextTokens)} tokens`,
304    s.lastActivityAt > 0
305      ? `  last activity: ${hhmm(s.lastActivityAt)} (${span(now - s.lastActivityAt)} ago)`
306      : '  last activity: none yet',
307    `  pings this idle period: ${s.pings} (${k(s.pingTokens)} served from cache)`,
308    s.phase === 'warming' ? `  next ping: ${hhmm(s.nextPingAt)}` : null,
309    `  today's keep-alive spend: ${k(spent)} of ${k(cfg.dailyBudget)} equivalents${off ? ' · off in every session today' : ''}`,
310    `  settings: ping every ${span(cfg.pingIntervalMs)} · cache lifetime ${span(cfg.cacheTtlMs)} · warm window ${span(cfg.warmWindowMs)} · ping ≥ ${k(cfg.minPingTokens)} · compact ≥ ${k(cfg.minCompactTokens)}`,
311    '  commands: /warm · /warm off · /warm on · /warm off all · /warm on all · /warm now',
312  ]
313  return lines.filter((line): line is string => line !== null).join('\n')
314}
315
316async function onSessionStart($: Engine): Promise<void> {
317  await $.command.register({
318    name: 'warm',
319    description: 'Cache warmer: status, off/on for this session, off all/on all (every session, today), now (ping)',
320    argumentHint: '[off|on|off all|on all|now]',
321  })
322
323  // After a hot reload or a restart: pick the schedule back up from state.
324  const s = await read($, warm)
325  if (s.phase === 'warming' && s.lastActivityAt > 0) {
326    await armFrom($, Math.max(s.lastActivityAt, s.lastPingAt))
327  } else {
328    paint($, s)
329  }
330}
331
332async function onPromptSubmit($: Engine): Promise<void> {
333  cancel()
334  const s = await update($, warm, (prev): WarmState => ({
335    ...prev,
336    phase: prev.phase === 'paused' ? 'paused' : 'active',
337    nextPingAt: 0,
338    note: '',
339  }))
340  paint($, s)
341}
342
343async function onTurnComplete($: Engine): Promise<void> {
344  const now = await $.clock.now()
345  const { context } = await $.session.usage()
346  const tokens = context.tokens ?? 0
347
348  const s = await update($, warm, (prev): WarmState => ({
349    ...prev,
350    lastActivityAt: now,
351    lastPingAt: 0,
352    nextPingAt: 0,
353    pings: 0,
354    pingTokens: 0,
355    errors: 0,
356    contextTokens: tokens,
357    note: '',
358    phase: prev.phase === 'paused' ? 'paused' : tokens >= cfg.minPingTokens ? 'warming' : 'small',
359  }))
360
361  if (s.phase === 'warming') {
362    await schedule($, cfg.pingIntervalMs)
363  } else {
364    paint($, s)
365  }
366}
367
368async function resumeHere($: Engine): Promise<string> {
369  const s = await update($, warm, (prev): WarmState => ({ ...prev, note: '' })) // drop the "off" note
370  if (s.lastActivityAt === 0) {
371    await setPhase($, 'idle')
372    return 'cache-warmer is on for this session. It arms after the next turn.'
373  }
374  const { context } = await $.session.usage()
375  const tokens = context.tokens ?? s.contextTokens
376  await update($, warm, (prev): WarmState => ({ ...prev, contextTokens: tokens }))
377  if (tokens < cfg.minPingTokens) {
378    await setPhase($, 'small')
379    return `cache-warmer is on for this session. Context is ${k(tokens)}, under the ${k(cfg.minPingTokens)} ping line, so nothing to do yet.`
380  }
381  await armFrom($, Math.max(s.lastActivityAt, s.lastPingAt))
382  return statusText($)
383}
384
385async function onWarmCommand($: Engine, args: string): Promise<string> {
386  const arg = args.trim().toLowerCase().replace(/\s+/g, ' ')
387  const now = await $.clock.now()
388
389  // This session only. 'paused' survives later turns, so it stays off until /warm on.
390  if (arg === 'off' || arg === 'pause') {
391    cancel()
392    await setPhase($, 'paused', 'off for this session (/warm off)')
393    return 'cache-warmer is off for this session only. Other sessions keep warming. /warm on turns it back on here.'
394  }
395
396  if (arg === 'on' || arg === 'resume') {
397    const note = (await isOffToday($, now))
398      ? ` Note: /warm off all is still in effect for every session today, so nothing will ping until /warm on all.`
399      : ''
400    return (await resumeHere($)) + note
401  }
402
403  // Every session, for the rest of today.
404  if (arg === 'off all') {
405    cancel()
406    await $.store.set(offKey(now), true)
407    await setPhase($, 'stopped', 'off in every session for today (/warm off all)')
408    return `cache-warmer is off in every session for the rest of ${dateKey(now)}. /warm on all lifts it.`
409  }
410
411  if (arg === 'on all') {
412    await $.store.delete(offKey(now))
413    return `Lifted /warm off all: every session can warm again.\n${await resumeHere($)}`
414  }
415
416  if (arg === 'now') {
417    const s = await read($, warm)
418    if (s.lastActivityAt === 0) return 'Nothing to ping yet: no turn has completed in this session.'
419    await update($, warm, (prev): WarmState => ({ ...prev, phase: 'warming' }))
420    await tick($)
421    return statusText($)
422  }
423
424  return statusText($)
425}
426
427export const register: Register = (on, options) => {
428  cfg = configFrom(options)
429
430  on('session.start', async ($, e, next) => {
431    await onSessionStart($)
432    return next(e)
433  })
434
435  on('prompt.submit', async ($, e, next) => {
436    await onPromptSubmit($)
437    return next(e)
438  })
439
440  on('turn.complete', async ($, e, next) => {
441    const done = await next(e)
442    if (e.agentId === undefined) await onTurnComplete($) // a subagent's turn leaves the main cache as it was
443    return done
444  })
445
446  on('session.end', ($, e, next) => {
447    cancel()
448    return next(e)
449  })
450
451  on('command.run', { command: 'warm' }, async ($, e) => ({ text: await onWarmCommand($, e.args) }))
452}
453

1export type WarmPhase =
2  | 'idle'
3  | 'active'
4  | 'small'
5  | 'warming'
6  | 'compacted'
7  | 'cold'
8  | 'paused'
9  | 'stopped'
10
11export type WarmState = {
12  phase: WarmPhase
13  lastActivityAt: number
14  lastPingAt: number
15  nextPingAt: number
16  pings: number
17  pingTokens: number
18  contextTokens: number
19  errors: number
20  note: string
21}
22
23declare module 'claude-code' {
24  interface PluginState {
25    'cache-warmer': { warm: WarmState }
26  }
27}
28