prompt-cache-tracker

1import type { EngineInterface, Register } from 'claude-code'
2
3type Ttl = '5m' | '1h'
4type TtlMode = 'auto' | Ttl
5
6const MINUTE = 60_000
7const TTL_MS: Record<Ttl, number> = { '5m': 5 * MINUTE, '1h': 60 * MINUTE }
8// Redraws: under a minute, so no minute of the label is skipped; faster on 5m
9// so the bar doesn't jump. Each divides its tier's lifetime and 10% of it.
10const TICK: Record<Ttl, number> = { '5m': 10_000, '1h': 45_000 }
11// Renew at 10% left.
12const RENEW_SHARE = 0.1
13// Keep warm stops after this long with no message from the person.
14const IDLE_LIMIT: Record<Ttl, number> = { '5m': 30 * MINUTE, '1h': 12 * 60 * MINUTE }
15// When $.fs.read refuses the transcript (too big), only its last TAIL_BYTES are read.
16const TAIL_BYTES = 2_000_000
17
18const BAR_CELLS = 24
19// An emoji keeps its own colors, so an expired cache shows a blank of the
20// same width instead of a lit flame.
21const FLAME = '🔥'
22const FLAME_BLANK = '  '
23const EXPIRED = '#6b6a64'
24const TRACK = '#3a3a37'
25
26const PING = 'Warming up (just reply ok)'
27
28// Emerald at full, Leaf at 75%, yellow at half, red when empty.
29const STOPS: [number, number, number, number][] = [
30  [1, 158, 64, 52],
31  [0.75, 110, 70, 50],
32  [0.5, 48, 96, 53],
33  [0, 0, 84, 60],
34]
35
36export function colorAt(fraction: number): string {
37  const p = Math.min(1, Math.max(0, fraction))
38  for (let i = 0; i < STOPS.length - 1; i++) {
39    const [hi, h1, s1, l1] = STOPS[i]!
40    const [lo, h2, s2, l2] = STOPS[i + 1]!
41    if (p >= lo) {
42      const k = (p - lo) / (hi - lo)
43      return hslToHex(h2 + (h1 - h2) * k, s2 + (s1 - s2) * k, l2 + (l1 - l2) * k)
44    }
45  }
46  const [, h, s, l] = STOPS[STOPS.length - 1]!
47  return hslToHex(h, s, l)
48}
49
50function hslToHex(h: number, s: number, l: number): string {
51  const sat = s / 100
52  const lig = l / 100
53  const a = sat * Math.min(lig, 1 - lig)
54  const channel = (n: number) => {
55    const k = (n + h / 30) % 12
56    const v = lig - a * Math.max(-1, Math.min(k - 3, 9 - k, 1))
57    return Math.round(v * 255)
58      .toString(16)
59      .padStart(2, '0')
60  }
61  return `#${channel(0)}${channel(8)}${channel(4)}`
62}
63
64// How long before expiry the ping goes out.
65export function renewAt(tier: Ttl): number {
66  return TTL_MS[tier] * RENEW_SHARE
67}
68
69function label(remainingMs: number): string {
70  if (remainingMs <= 0) return 'expired'
71  return `${Math.ceil(remainingMs / MINUTE)}m left`
72}
73
74function flameSvg(color: string, isExpired: boolean): string {
75  const expired = isExpired ? '<path d="M3 3l18 18"/>' : ''
76  return (
77    `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="${color}" ` +
78    `stroke-width="2" stroke-linecap="round" stroke-linejoin="round">` +
79    `<path d="M12 12c2-2.96 0-7-1-8 0 3.038-1.773 4.741-3 6-1.226 1.26-2 3.24-2 5a6 6 0 1 0 12 0c0-1.532-1.056-3.94-2-5-1.786 3-2.791 3-4 2z"/>` +
80    `${expired}</svg>`
81  )
82}
83
84function barSvg(color: string, fraction: number): string {
85  const w = 200
86  const filled = Math.round(w * Math.min(1, Math.max(0, fraction)))
87  return (
88    `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 ${w} 6">` +
89    `<rect width="${w}" height="6" rx="3" fill="${TRACK}"/>` +
90    (filled > 0 ? `<rect width="${filled}" height="6" rx="3" fill="${color}"/>` : '') +
91    `</svg>`
92  )
93}
94
95// Module state: starts over whenever the engine loads the module again (a
96// reload while developing, or a lost hooks worker), like a new session.
97let anchorAt: number | null = null
98let ttlMode: TtlMode = 'auto'
99// What the transcript says the last cache write used; null until it says.
100let writtenTier: Ttl | null = null
101let keepWarm = false
102let mainBusy = false
103// A person's prompt has gone in and its turn hasn't started.
104let isPromptPending = false
105let lastHumanAt = 0
106let pingedFor: number | null = null
107// The session's log file, found by the session id it is named for; /clear
108// starts a new id and so a new file.
109let transcriptPath = ''
110let transcriptFor = ''
111let tick: { cancel: () => void } | null = null
112
113// The tier the countdown follows: the person's pick, else the transcript's.
114function tierNow(): Ttl | null {
115  return ttlMode !== 'auto' ? ttlMode : writtenTier
116}
117
118// A fresh conversation: no cache yet, bar hidden until a reply says which
119// tier the new cache is on. Keep warm stays as the person left it; it starts
120// off and only /keep-warm turns it on.
121function resetConversation() {
122  anchorAt = null
123  writtenTier = null
124  pingedFor = null
125  mainBusy = false
126  isPromptPending = false
127}
128
129function toggleKeepWarm($: EngineInterface) {
130  keepWarm = !keepWarm
131  $.ui.invalidate('ui.render')
132}
133
134const CACHE_NAME: Record<Ttl, string> = { '5m': '5-minute', '1h': '1-hour' }
135
136// Schedules the next redraw and check on the cache refresh plus whole ticks,
137// so one lands exactly on the renew point and the last on expiry. Nothing to
138// count (no bar, or expired): no timer. While the bar shows it keeps
139// redrawing until expiry, even once pinging is done.
140async function startTick($: EngineInterface) {
141  const t = await $.clock.now()
142  tick?.cancel()
143  tick = null
144  const anchor = anchorAt
145  const tier = tierNow()
146  if (anchor === null || tier === null) return
147  const next = anchor + (Math.floor(Math.max(0, t - anchor) / TICK[tier]) + 1) * TICK[tier]
148  if (next > anchor + TTL_MS[tier]) return
149  tick = $.clock.after(next - t, async () => {
150    await startTick($)
151    await renewIfDue($)
152  })
153}
154
155// No cache to count until the next reply: hide the bar and stop the timer.
156async function hideBar($: EngineInterface) {
157  anchorAt = null
158  $.ui.invalidate('ui.render')
159  await startTick($)
160}
161
162async function setWrittenTier($: EngineInterface, tier: Ttl) {
163  if (tier === writtenTier) return
164  writtenTier = tier
165  $.ui.invalidate('ui.render')
166  await startTick($)
167}
168
169// The newest main-thread cache write in the transcript says which tier the
170// request asked for: Claude Code splits it into 1h and 5m tokens. 1-hour
171// entries must come before 5-minute ones in a request, so any 5-minute write
172// is the conversation's end, which expires first.
173export function tierFromTranscript(text: string): Ttl | null {
174  const lines = text.split('\n')
175  for (let i = lines.length - 1; i >= 0; i--) {
176    const line = lines[i]!
177    if (!line.includes('"cache_creation"')) continue
178    let row: any
179    try {
180      row = JSON.parse(line)
181    } catch {
182      continue // The tail can start mid-line.
183    }
184    if (row?.type !== 'assistant' || row.isSidechain) continue
185    if (row.message?.model === '<synthetic>') continue
186    const written = row.message?.usage?.cache_creation
187    if (Number(written?.ephemeral_5m_input_tokens) > 0) return '5m'
188    if (Number(written?.ephemeral_1h_input_tokens) > 0) return '1h'
189  }
190  return null
191}
192
193async function readTranscriptTier($: EngineInterface, path: string): Promise<Ttl | null> {
194  if (!path) return null
195  let text: string
196  try {
197    text = await $.fs.read(path)
198  } catch {
199    // Too big for one read (or gone): read only its end.
200    try {
201      const tail = await $.process.run(['tail', '-c', String(TAIL_BYTES), path])
202      if (tail.exitCode !== 0) return null
203      text = tail.stdout
204    } catch {
205      return null
206    }
207  }
208  return tierFromTranscript(text)
209}
210
211// Where Claude Code keeps the session's log: projects/<folder>/<session id>.jsonl.
212// The same id in two folders (a session resumed elsewhere): the newest file.
213const FIND_TRANSCRIPT = 'ls -t "${CLAUDE_CONFIG_DIR:-$HOME/.claude}"/projects/*/"$1".jsonl 2>/dev/null | head -n 1'
214
215async function findTranscript($: EngineInterface): Promise<string> {
216  const id = await $.session.id()
217  if (id === transcriptFor) return transcriptPath
218  const found = await $.process.run(['sh', '-c', FIND_TRANSCRIPT, 'sh', id])
219  const path = found.exitCode === 0 ? found.stdout.trim() : ''
220  // Not written yet: look again next time.
221  if (path) {
222    transcriptPath = path
223    transcriptFor = id
224  }
225  return path
226}
227
228// Reads the transcript and follows its tier. A read that finds nothing keeps
229// a tier already known; with none known the bar stays hidden.
230async function checkTranscript($: EngineInterface) {
231  const tier = await readTranscriptTier($, await findTranscript($))
232  if (tier !== null) await setWrittenTier($, tier)
233}
234
235async function renewIfDue($: EngineInterface) {
236  const t = await $.clock.now()
237  $.ui.invalidate('ui.render')
238  const anchor = anchorAt
239  const tier = tierNow()
240  if (anchor === null || tier === null) return
241  if (mainBusy || isPromptPending || !keepWarm || anchor === pingedFor) return
242  if (t - lastHumanAt > IDLE_LIMIT[tier]) return
243  const remaining = anchor + TTL_MS[tier] - t
244  if (remaining <= 0 || remaining > renewAt(tier)) return
245
246  pingedFor = anchor
247  let result
248  try {
249    result = await $.prompt.submit({ text: PING })
250  } catch {
251    result = { drop: 'failed' }
252  }
253  if (result.drop !== undefined) {
254    // Refused or failed to start: let the next tick try again.
255    if (pingedFor === anchor) pingedFor = null
256  }
257}
258
259export const register: Register = on => {
260  on('session.start', async ($, e, next) => {
261    const stored = await $.store.get('ttlMode')
262    if (stored === '5m' || stored === '1h' || stored === 'auto') ttlMode = stored
263    resetConversation()
264    lastHumanAt = await $.clock.now()
265    await startTick($)
266
267    await $.command.register({
268      name: 'keep-warm',
269      description: 'Turn automatic cache refresh on or off',
270    })
271    await $.command.register({
272      name: 'cache-ttl',
273      description: 'Choose a 5-minute or 1-hour cache for the progress bar, or detect it automatically',
274      argumentHint: 'auto | 5m | 1h',
275    })
276
277    return next(e)
278  })
279
280  on('prompt.submit', async ($, e, next) => {
281    const kind = e.origin?.kind
282    const isPerson = kind === 'composer' || kind === 'bridge' || kind === 'sdk'
283    if (isPerson) {
284      lastHumanAt = await $.clock.now()
285      // Hold off a ping from Enter until the turn starts.
286      if (!e.command) isPromptPending = true
287    }
288    let result
289    try {
290      result = await next(e)
291    } catch (error) {
292      if (isPerson) isPromptPending = false
293      throw error
294    }
295    // A prompt something dropped starts no turn.
296    if (isPerson && result.drop !== undefined) isPromptPending = false
297    return result
298  })
299
300  on('session.end', async ($, e, next) => {
301    if (e.reason === 'clear') {
302      resetConversation()
303      await hideBar($)
304    }
305    return next(e)
306  })
307
308  on('session.compact', async ($, e, next) => {
309    const result = await next(e)
310    // The compacted conversation no longer uses the old cache; the next
311    // request starts a new one.
312    if (e.agentId === undefined && e.trigger !== 'precompute' && !('skip' in result && result.skip)) {
313      pingedFor = null
314      await hideBar($)
315    }
316    return result
317  })
318
319  on('turn.start', ($, e, next) => {
320    if (e.agentId === undefined) {
321      mainBusy = true
322      isPromptPending = false
323    }
324    return next(e)
325  })
326
327  on('turn.complete', async ($, e, next) => {
328    if (e.agentId === undefined) {
329      mainBusy = false
330      // A queued prompt removed before its turn started leaves no turn of its
331      // own; the turn it waited behind ends here.
332      isPromptPending = false
333      // The turn's rows are in the log by now. The turn's end waits on this,
334      // so a failed read must never hold it.
335      await checkTranscript($).catch(() => undefined)
336    }
337    return next(e)
338  })
339
340  on('turn.step', async function* ($, e, next) {
341    if (e.agentId !== undefined) return yield* next(e)
342    const stepStartedAt = await $.clock.now()
343    const result = yield* next(e)
344    const usage = result.usage
345    if (usage === null) return result
346
347    // No cache touched (caching off, or too short to cache): nothing to count.
348    if (usage.cache_read_input_tokens + usage.cache_creation_input_tokens === 0) return result
349
350    // The cache entry was refreshed when the request went out, not when the
351    // reply finished streaming.
352    anchorAt = stepStartedAt
353    $.ui.invalidate('ui.render')
354    await startTick($)
355    void checkTranscript($).catch(() => undefined)
356    return result
357  })
358
359  on('command.run', { command: 'keep-warm' }, async $ => {
360    toggleKeepWarm($)
361    return { text: `Keep warm is ${keepWarm ? 'on' : 'off'}.` }
362  })
363
364  on('command.run', { command: 'cache-ttl' }, async ($, e) => {
365    const arg = e.args.trim()
366    if (arg !== 'auto' && arg !== '5m' && arg !== '1h') {
367      const current = tierNow()
368      return {
369        text: current
370          ? `Tracking the ${CACHE_NAME[current]} cache. Change it with /cache-ttl 5m, 1h or auto.`
371          : 'Waiting for your first reply to detect the cache.',
372      }
373    }
374    ttlMode = arg
375    await $.store.set('ttlMode', arg)
376    $.ui.invalidate('ui.render')
377    await startTick($)
378    return { text: arg === 'auto' ? 'Detecting the cache automatically.' : `Tracking the ${CACHE_NAME[arg]} cache.` }
379  })
380
381  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
382    // Shown only once a reply has made a cache and its tier is known.
383    const anchor = anchorAt
384    const tier = tierNow()
385    if (e.props.hasSurvey || anchor === null || tier === null) return next(e)
386
387    const toggle = () => toggleKeepWarm($)
388    const t = Math.max(await $.clock.now(), anchor)
389    const check = `${keepWarm ? '[✓]' : '[ ]'} keep warm`
390    const total = TTL_MS[tier]
391    const remaining = anchor + total - t
392    const fraction = remaining / total
393    const isExpired = remaining <= 0
394    const color = isExpired ? EXPIRED : colorAt(fraction)
395    const text = label(remaining)
396
397    if (e.surface === 'terminal') {
398      const { Box, Text, Button } = $.ui.resolve(e)
399      const filled = isExpired ? 0 : Math.max(1, Math.round(BAR_CELLS * fraction))
400      return (
401        // One empty row above, to match the engine's row under the band.
402        <Box flexDirection="row" width={e.props.bodyColumns} marginTop={1}>
403          <Text>{isExpired ? FLAME_BLANK : FLAME}{' '}</Text>
404          <Text color={color}>{'━'.repeat(filled)}</Text>
405          <Text color={TRACK}>{'━'.repeat(BAR_CELLS - filled)}</Text>
406          <Text color={color}>  {text}</Text>
407          <Box flexGrow={1} />
408          <Button key="keep-warm" plain dimColor={!keepWarm} onPress={toggle}>
409            {check}
410          </Button>
411        </Box>
412      )
413    }
414
415    const { Box, Text, Button, Svg } = $.ui.resolve(e)
416    return (
417      <Box flexDirection="row" alignItems="center" gap={1}>
418        {/* A row of its own, so the desktop centers the icon on the text. */}
419        <Box flexDirection="row" alignItems="center" gap={1}>
420          <Svg source={flameSvg(color, isExpired)} alt="prompt cache" width={16} height={16} />
421          <Svg source={barSvg(color, isExpired ? 0 : fraction)} alt={text} width={200} height={6} />
422          <Text color={color}>{text}</Text>
423        </Box>
424        <Box flexGrow={1} />
425        <Button key="keep-warm" plain dimColor={!keepWarm} onPress={toggle}>
426          {check}
427        </Button>
428      </Box>
429    )
430  })
431}
432