cache-keepalive

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, Timer } from 'claude-code'
3
4const MINUTE = 60_000
5// The 1-hour cache entry lapses 60 minutes after the request that last read or
6// wrote it, counted from when that request was sent, not from when its
7// response or the turn finished.
8const TTL_MS = 60 * MINUTE
9// Pings due later than this are skipped: the entry is probably gone, and a
10// ping would pay a full cache write for nothing. The minute of slack covers
11// timer drift and the time the ping request takes to reach the API.
12const LATEST_PING_MS = TTL_MS - MINUTE
13const PING_PROMPT =
14  'Automated prompt-cache keep-alive. Do not use tools. Reply with exactly: ok'
15
16const lastWarmAt = atom({ plugin: 'cache-keepalive', key: 'lastWarmAt' } as const, null)
17const pings = atom({ plugin: 'cache-keepalive', key: 'pings' } as const, 0)
18const isPaused = atom({ plugin: 'cache-keepalive', key: 'isPaused' } as const, false)
19
20// The module's own variables start over on every load; what must survive a
21// reload lives in the atoms above.
22let idleMs = 55 * MINUTE
23let maxPings = 6
24let timer: Timer | undefined
25let dueAt: number | undefined
26let isTurnRunning = false
27let isPinging = false
28
29const formatTokens = (n: number) => (n >= 1000 ? `${Math.round(n / 1000)}k` : `${n}`)
30
31function disarm() {
32  timer?.cancel()
33  timer = undefined
34  dueAt = undefined
35}
36
37// Schedules the next ping for idleMs after the cache was last warmed.
38async function arm($: EngineInterface) {
39  disarm()
40  const warmedAt = await read($, lastWarmAt)
41  if (warmedAt === null || isTurnRunning || (await read($, isPaused))) {
42    return
43  }
44  const sent = await read($, pings)
45  if (maxPings > 0 && sent >= maxPings) {
46    $.ui.status(`cache keep-alive: stopped after ${sent} pings`)
47    return
48  }
49  const now = await $.clock.now()
50  if (now - warmedAt >= LATEST_PING_MS) {
51    $.ui.status('cache keep-alive: cache likely cold')
52    return
53  }
54  const wait = Math.max(warmedAt + idleMs - now, 0)
55  dueAt = now + wait
56  timer = $.clock.after(wait, () => {
57    ping($).catch(err => $.ui.log(`cache-keepalive: ping threw: ${err}`, { to: 'debug' }))
58  })
59  $.ui.status(
60    maxPings > 0 ? `cache keep-alive: armed (${sent}/${maxPings})` : 'cache keep-alive: armed',
61  )
62}
63
64async function ping($: EngineInterface) {
65  timer = undefined
66  dueAt = undefined
67  if (isTurnRunning || isPinging || (await read($, isPaused))) {
68    return
69  }
70  const warmedAt = await read($, lastWarmAt)
71  const now = await $.clock.now()
72  // A late timer (the machine slept) finds the entry already gone: skip
73  // rather than pay to rebuild a cache nobody may come back to.
74  if (warmedAt === null || now - warmedAt >= LATEST_PING_MS) {
75    $.ui.status('cache keep-alive: cache likely cold')
76    return
77  }
78
79  isPinging = true
80  const r = await $.model.fork({ prompt: PING_PROMPT }).finally(() => {
81    isPinging = false
82  })
83  if (isTurnRunning) {
84    return
85  }
86
87  if (!r.isAnswered && r.reason === 'nothing-to-fork') {
88    await update($, lastWarmAt, () => null)
89    $.ui.status(undefined)
90    return
91  }
92  if (!r.isAnswered && r.reason !== 'empty-reply') {
93    $.ui.log(`cache-keepalive: ping failed (${r.reason})`, { to: 'debug' })
94    $.ui.status('cache keep-alive: ping failed')
95    return
96  }
97
98  const { cache_read_input_tokens: hit, cache_creation_input_tokens: wrote } = r.usage
99  if (hit === 0 && wrote > 0) {
100    // Nothing was served from cache: the TTL is shorter than assumed (the
101    // 5-minute cache) or the model changed. Pinging again would only repeat
102    // a full cache write, so stand down until the next turn.
103    await update($, lastWarmAt, () => null)
104    $.ui.status('cache keep-alive: no cache hit, paused until next turn')
105    $.ui.toast('cache-keepalive: ping found the cache cold; is the 1-hour TTL on?')
106    return
107  }
108
109  // `now` was read before the fork was sent, so the countdown runs from no
110  // later than the request that refreshed the entry.
111  await update($, lastWarmAt, () => now)
112  const sent = await update($, pings, n => n + 1)
113  $.ui.log(`cache-keepalive: refreshed ${formatTokens(hit)} cached tokens (ping ${sent})`, {
114    to: 'debug',
115  })
116  await arm($)
117}
118
119async function describe($: EngineInterface) {
120  if (await read($, isPaused)) {
121    return 'Cache keep-alive is off. Run /cache-keepalive on to resume.'
122  }
123  const sent = await read($, pings)
124  const limit = maxPings > 0 ? `/${maxPings}` : ''
125  if (dueAt === undefined) {
126    return `Cache keep-alive is idle (${sent}${limit} pings sent). It arms after the next response.`
127  }
128  const minutes = Math.max(Math.round((dueAt - (await $.clock.now())) / MINUTE), 0)
129  return `Next keep-alive in ${minutes} min (${sent}${limit} pings sent since your last turn).`
130}
131
132export const register: Register = (on, options) => {
133  // At most 58, so a ping is always due before LATEST_PING_MS.
134  idleMs = Math.min(Math.max(Number(options.idleMinutes ?? 55), 1), 58) * MINUTE
135  maxPings = Math.max(Number(options.maxPings ?? 6), 0)
136
137  on('session.start', async ($, e, next) => {
138    const result = await next(e)
139    await $.command.register({
140      name: 'cache-keepalive',
141      description: 'Prompt-cache keep-alive: status | on | off | now',
142    })
143    // A reload keeps $.state, so pick up where the previous load left off.
144    await arm($)
145    return result
146  })
147
148  on('turn.start', ($, e, next) => {
149    isTurnRunning = true
150    disarm()
151    return next(e)
152  })
153
154  // Each main-thread request reads or writes the cache when it is sent, so its
155  // send time (taken before the response streams) is when the TTL restarts.
156  // A long response or a long turn therefore can't push the ping past expiry.
157  // Subagent steps send their own prefix and don't warm the main thread's.
158  on('turn.step', async function* ($, e, next) {
159    if (e.agentId !== undefined) {
160      return yield* next(e)
161    }
162    const sentAt = await $.clock.now()
163    const result = yield* next(e)
164    const usage = result.usage
165    // No usage means no response arrived, so the entry may not have been
166    // touched; keep the older, safer timestamp.
167    if (usage !== null && usage.cache_read_input_tokens + usage.cache_creation_input_tokens > 0) {
168      await update($, lastWarmAt, () => sentAt)
169    }
170    return result
171  })
172
173  on('turn.complete', async ($, e, next) => {
174    const result = await next(e)
175    if (e.agentId === undefined) {
176      isTurnRunning = false
177      await update($, pings, () => 0)
178      await arm($)
179    }
180    return result
181  })
182
183  on('command.run', { command: 'cache-keepalive' }, async ($, e) => {
184    const arg = e.args.trim().toLowerCase()
185
186    if (arg === 'off') {
187      await update($, isPaused, () => true)
188      disarm()
189      $.ui.status(undefined)
190      return { text: 'Cache keep-alive is off for this session.' }
191    }
192    if (arg === 'on') {
193      await update($, isPaused, () => false)
194      await update($, pings, () => 0)
195      await arm($)
196      return { text: await describe($) }
197    }
198    if (arg === 'now') {
199      await ping($)
200      return { text: `Keep-alive attempted. ${await describe($)}` }
201    }
202    return { text: await describe($) }
203  })
204}
205

1export type WarmedAt = number | null
2
3declare module 'claude-code' {
4  interface PluginState {
5    'cache-keepalive': {
6      /** Epoch ms when the last main-thread request or keep-alive that read or wrote the cache was sent; null before the first. */
7      lastWarmAt: WarmedAt
8      /** Keep-alives sent since the person last ran a turn. */
9      pings: number
10      /** Turned off with /cache-keepalive off. */
11      isPaused: boolean
12    }
13  }
14}
15