cache-clock

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4const WARM_LIMIT_MS = 2 * 60 * 60 * 1000
5const MARGIN_MS = 20 * 1000
6
7const expiresAt = atom({ plugin: 'cache-clock', key: 'expiresAt' } as const, null, {
8  shape: 'expires-1',
9})
10const warmUntil = atom({ plugin: 'cache-clock', key: 'warmUntil' } as const, null, {
11  shape: 'warm-1',
12})
13const clockNow = atom({ plugin: 'cache-clock', key: 'now' } as const, 0, { shape: 'now-1' })
14
15const clock = (ms: number) => {
16  const s = Math.max(0, Math.ceil(ms / 1000))
17
18  return `${Math.floor(s / 60)}:${String(s % 60).padStart(2, '0')}`
19}
20
21let ttlMs = 5 * 60 * 1000
22let timer: { cancel: () => void } | undefined
23
24const stopWarming = () => {
25  timer?.cancel()
26  timer = undefined
27}
28
29async function keepWarm($: EngineInterface) {
30  const start = await $.clock.now()
31  const limit = start + WARM_LIMIT_MS
32  await update($, warmUntil, () => limit)
33
34  const schedule = async () => {
35    const at = await $.clock.now()
36    const due = (await read($, expiresAt)) ?? at + ttlMs
37    timer = $.clock.after(Math.max(1000, due - at - MARGIN_MS), async () => {
38      const t = await $.clock.now()
39      if (t >= limit || (await read($, warmUntil)) === null) {
40        await update($, warmUntil, () => null)
41        $.ui.toast('Cache Clock: 2 hours passed, warming stops.')
42        return
43      }
44      // A fork rides the session's own cached prefix and never enters the chat.
45      const pinged = await $.model.fork({ prompt: 'Reply with the single word: ok' })
46      if (!pinged.isAnswered) {
47        await update($, warmUntil, () => null)
48        $.ui.toast('Cache Clock: the ping failed, warming stops.')
49        return
50      }
51      await update($, expiresAt, () => t + ttlMs)
52      await schedule()
53    })
54  }
55  await schedule()
56}
57
58export const register: Register = (on, options) => {
59  ttlMs = options.cacheTtl === '1h' ? 60 * 60 * 1000 : 5 * 60 * 1000
60
61  on('session.start', async ($, e, next) => {
62    await $.command.register({
63      name: 'keep-warm',
64      description: 'Ping the prompt cache just before it expires, for up to 2 hours',
65    })
66    $.clock.every(1000, async () => {
67      if ((await read($, expiresAt)) !== null) await update($, clockNow, () => Date.now())
68    })
69
70    return next(e)
71  })
72
73  // The reply is in: the cache was written or read, so its clock starts over.
74  on('turn.complete', async ($, e, next) => {
75    if (e.agentId === undefined && e.usage) {
76      const at = await $.clock.now()
77      await update($, expiresAt, () => at + ttlMs)
78      await update($, clockNow, () => at)
79    }
80
81    return next(e)
82  })
83
84  // A message of the person's own ends warming; the next reply restarts the clock.
85  on('prompt.submit', async ($, e, next) => {
86    if (!e.isMeta) {
87      stopWarming()
88      await update($, warmUntil, () => null)
89      await update($, expiresAt, () => null)
90    }
91
92    return next(e)
93  })
94
95  on('command.run', { command: 'keep-warm' }, async $ => {
96    if ((await read($, warmUntil)) !== null) {
97      stopWarming()
98      await update($, warmUntil, () => null)
99
100      return { text: 'Keep Warm stopped.' }
101    }
102    if ((await read($, expiresAt)) === null) return { text: 'No cache to keep warm yet.' }
103    await keepWarm($)
104
105    return { text: 'Keeping the cache warm for up to 2 hours, or until your next message.' }
106  })
107
108  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
109    const end = await read($, expiresAt)
110    if (e.props.hasSurvey || end === null) return next(e)
111
112    const now = await read($, clockNow)
113    const warming = (await read($, warmUntil)) !== null
114    const left = end - (now === 0 ? Date.now() : now)
115    const { Box, Button, Text } = $.ui.resolve(e)
116
117    return (
118      <Box>
119        <Text color={left > 60_000 ? 'green' : left > 0 ? 'yellow' : undefined} dimColor={left <= 0}>
120          {left > 0 ? `Cache warm ${clock(left)}` : 'Cache cold'}{' '}
121        </Text>
122        <Button
123          key="warm"
124          label={warming ? 'Stop Keep Warm' : 'Keep Warm'}
125          onPress={async () => {
126            if ((await read($, warmUntil)) !== null) {
127              stopWarming()
128              await update($, warmUntil, () => null)
129            } else {
130              await keepWarm($)
131            }
132          }}
133        />
134      </Box>
135    )
136  })
137}
138

1export type CacheClockTick = number
2
3declare module 'claude-code' {
4  interface PluginState {
5    'cache-clock': {
6      expiresAt: Shaped<number | null>
7      warmUntil: Shaped<number | null>
8      now: Shaped<number>
9    }
10  }
11}
12