keep-warm

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface as Engine, Register } from 'claude-code'
3
4import type { Cache } from '../types'
5
6// Just under the 5-minute cache TTL, so each ping lands before the entry lapses.
7const PING_EVERY_MS = 290_000
8// Past this much idle time, one cache rewrite is cheaper than more pings.
9const MAX_IDLE_MS = 30 * 60_000
10const CACHE_TTL_MS = 5 * 60_000
11const PING_PROMPT = 'Keep-alive ping. Reply with exactly: ok'
12
13const cache = atom({ plugin: 'keep-warm', key: 'cache' } as const, null as Cache | null)
14
15const clockText = (ms: number) => {
16  const seconds = Math.max(0, Math.ceil(ms / 1000))
17
18  return `${Math.floor(seconds / 60)}:${String(seconds % 60).padStart(2, '0')}`
19}
20
21type Timer = { cancel: () => void }
22
23type Warm = {
24  timer?: Timer
25  tick?: Timer
26  idleSince: number
27  pings: number
28  // Thousands of tokens the last ping read from the cache.
29  cachedK: number
30  // Bumped by each keep-warm start and halt, so a pending ping of a retired run does nothing.
31  runs: number
32}
33
34const stopTimers = (w: Warm) => {
35  w.timer?.cancel()
36  w.timer = undefined
37  w.tick?.cancel()
38  w.tick = undefined
39}
40
41// Starts pinging from `anchor`, retiring any run before it.
42async function keepWarm($: Engine, w: Warm, anchor: number) {
43  const run = ++w.runs
44  stopTimers(w)
45  w.idleSince = anchor
46  w.pings = 0
47  await update($, cache, () => ({
48    warmUntil: anchor + CACHE_TTL_MS,
49    now: anchor,
50    keepWarmUntil: anchor + MAX_IDLE_MS,
51  }))
52
53  const endKeepWarm = async (status: string) => {
54    w.timer?.cancel()
55    w.timer = undefined
56    $.ui.status(status)
57    await update($, cache, c => c && { ...c, keepWarmUntil: null })
58  }
59
60  const ping = async () => {
61    if (run !== w.runs) {
62      return
63    }
64
65    if ((await $.clock.now()) - w.idleSince >= MAX_IDLE_MS) {
66      await endKeepWarm(`keep-warm: idle ${MAX_IDLE_MS / 60_000}m, stopped after ${w.pings} pings`)
67
68      return
69    }
70
71    const sentAt = await $.clock.now()
72    const reply = await $.model.fork({ prompt: PING_PROMPT })
73
74    if (run !== w.runs) {
75      return
76    }
77
78    if (!reply.isAnswered && reply.reason !== 'empty-reply') {
79      await endKeepWarm(`keep-warm: stopped (${reply.reason})`)
80
81      return
82    }
83
84    w.pings += 1
85    const now = await $.clock.now()
86    await update($, cache, c => c && { ...c, warmUntil: sentAt + CACHE_TTL_MS, now })
87    w.cachedK = Math.round(reply.usage.cache_read_input_tokens / 1000)
88    $.ui.status(`keep-warm: ${w.pings} pings, ${w.cachedK}k cached`)
89    schedule(sentAt)
90  }
91
92  // Each ping lands PING_EVERY_MS after the last request was sent, whenever that was.
93  const schedule = (from: number) => {
94    void $.clock.now().then(now => {
95      if (run === w.runs) {
96        w.timer = $.clock.after(Math.max(0, from + PING_EVERY_MS - now), () => void ping())
97      }
98    })
99  }
100
101  // Redraws the countdown each second; once the cache is cold it no longer changes.
102  const advance = async () => {
103    const now = await $.clock.now()
104    await update($, cache, c => c && { ...c, now })
105    const c = await read($, cache)
106
107    if (c === null || (c.keepWarmUntil === null && now >= c.warmUntil)) {
108      w.tick?.cancel()
109      w.tick = undefined
110    }
111  }
112
113  schedule(anchor)
114  w.tick = $.clock.every(1000, () => void advance())
115}
116
117// The status line is hidden behind a question or permission dialog, so the transcript records the pings.
118async function halt($: Engine, w: Warm) {
119  if (w.pings > 0) {
120    $.ui.log(`kept the cache warm with ${w.pings} pings, ${w.cachedK}k cached`)
121    w.pings = 0
122  }
123
124  w.runs += 1
125  stopTimers(w)
126  $.ui.status(undefined)
127  await update($, cache, () => null)
128}
129
130// Lets the cache lapse: retires the pending ping and keeps only the countdown running.
131async function letCool($: Engine, w: Warm) {
132  w.runs += 1
133  w.timer?.cancel()
134  w.timer = undefined
135  $.ui.status(undefined)
136  await update($, cache, c => c && { ...c, keepWarmUntil: null })
137}
138
139export const register: Register = on => {
140  const w: Warm = { idleSince: 0, pings: 0, cachedK: 0, runs: 0 }
141  // When the main loop last sent a model request; the cache TTL counts from here.
142  let lastRequestAt: number | undefined
143  // Main-loop tool calls in flight; parallel calls share one keep-warm run.
144  let toolsRunning = 0
145  // Set by /let-cool; the next message from the person clears it and keep-warm resumes.
146  let isCooling = false
147
148  on('session.start', async ($, e, next) => {
149    await $.command.register({
150      name: 'let-cool',
151      description: 'Stop keep-warm pings; the cache goes cold, and the next message resumes keep-warm',
152    })
153
154    return next(e)
155  })
156
157  on('command.run', { command: 'let-cool' }, async $ => {
158    isCooling = true
159    await letCool($, w)
160
161    return { text: 'Keep-warm paused. The cache goes cold within 5 minutes; your next message resumes it.' }
162  })
163
164  on('turn.start', async ($, e, next) => {
165    if (e.agentId === undefined) {
166      isCooling = false
167      toolsRunning = 0
168      await halt($, w)
169    }
170
171    return next(e)
172  })
173
174  on('turn.step', async function* ($, e, next) {
175    if (e.agentId === undefined) {
176      lastRequestAt = await $.clock.now()
177    }
178
179    return yield* next(e)
180  })
181
182  on('turn.complete', async ($, e, next) => {
183    const done = await next(e)
184
185    if (e.agentId === undefined && !isCooling) {
186      await keepWarm($, w, lastRequestAt ?? (await $.clock.now()))
187    }
188
189    return done
190  })
191
192  // A turn waits inside a tool call (a question to the person, a permission prompt, a long
193  // tool or subagent) without sending requests, so keep warm there as when the turn is done.
194  on('tool.call', async ($, e, next) => {
195    if (e.agentId !== undefined) {
196      return next(e)
197    }
198
199    toolsRunning += 1
200
201    if (toolsRunning === 1 && !isCooling) {
202      await keepWarm($, w, lastRequestAt ?? (await $.clock.now()))
203    }
204
205    try {
206      return await next(e)
207    } finally {
208      toolsRunning -= 1
209
210      if (toolsRunning === 0) {
211        await halt($, w)
212      }
213    }
214  })
215
216  on('session.end', ($, e, next) => {
217    w.runs += 1
218    stopTimers(w)
219
220    return next(e)
221  })
222
223  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
224    const c = await read($, cache)
225
226    if (e.props.hasSurvey || e.props.isWorking || c === null) {
227      return next(e)
228    }
229
230    const { Box, Text } = $.ui.resolve(e)
231    const left = c.warmUntil - c.now
232
233    if (left <= 0) {
234      return (
235        <Box justifyContent="flex-end">
236          <Text dimColor>🧊 cache cold</Text>
237        </Box>
238      )
239    }
240
241    const [dot, color] =
242      left > 3 * 60_000 ? ['🟢', 'success'] : left > 60_000 ? ['🟡', 'warning'] : ['🔴', 'error']
243
244    return (
245      <Box justifyContent="flex-end">
246        <Text color={color}>
247          {dot} ⏳ cache {clockText(left)}
248        </Text>
249        {c.keepWarmUntil !== null ? (
250          <Text dimColor> · keep-warm {clockText(c.keepWarmUntil - c.now)}</Text>
251        ) : null}
252      </Box>
253    )
254  })
255}
256

1export type Cache = { warmUntil: number; now: number; keepWarmUntil: number | null }
2
3declare module 'claude-code' {
4  interface PluginState {
5    'keep-warm': { cache: Cache | null }
6  }
7}
8