cold-cache-guard

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, PromptOrigin, Register } from 'claude-code'
3
4const lastAnswerAtMs = atom({ plugin: 'cold-cache-guard', key: 'lastAnswerAtMs' } as const, null)
5const resumedTokens = atom({ plugin: 'cold-cache-guard', key: 'resumedTokens' } as const, null)
6const isSettled = atom({ plugin: 'cold-cache-guard', key: 'isSettled' } as const, false)
7
8type Cold = { idleMs: number; tokens: number; usd?: number }
9
10const TEXTS = {
11  en: {
12    question: (cold: Cold) =>
13      `Cold prompt cache: last answer ${durationOf(cold.idleMs)} ago, ` +
14      `~${Math.round(cold.tokens / 1000)}k tokens to cache again` +
15      (cold.usd === undefined ? '' : ` (~$${cold.usd.toFixed(2)})`) +
16      '. What now?',
17    keep: 'Keep the session as is',
18    compactNow: 'Compact now',
19    send: 'Send as is',
20    compactFirst: 'Compact first',
21    clear: 'Start over (/clear)',
22    cancel: 'Cancel',
23    compacting: 'Compacting; your prompt will be back in the box.',
24    cleared: 'Session cleared; your prompt is back in the box.',
25    cancelled: 'Not sent; your prompt is back in the box.',
26  },
27  fr: {
28    question: (cold: Cold) =>
29      `Cache froid : dernière réponse il y a ${durationOf(cold.idleMs)}, ` +
30      `~${Math.round(cold.tokens / 1000)}k tokens à remettre en cache` +
31      (cold.usd === undefined ? '' : ` (~${cold.usd.toFixed(2)} $)`) +
32      '. Que faire ?',
33    keep: 'Garder la session telle quelle',
34    compactNow: 'Compacter maintenant',
35    send: 'Envoyer tel quel',
36    compactFirst: "Compacter d'abord",
37    clear: 'Repartir à vide (/clear)',
38    cancel: 'Annuler',
39    compacting: 'Compaction en cours, le prompt reviendra dans le champ.',
40    cleared: 'Session vidée, prompt remis dans le champ.',
41    cancelled: 'Envoi annulé, prompt remis dans le champ.',
42  },
43}
44
45type Texts = (typeof TEXTS)['en']
46
47function durationOf(ms: number): string {
48  const minutes = Math.floor(ms / 60_000)
49  return minutes < 60
50    ? `${minutes} min`
51    : `${Math.floor(minutes / 60)}h${String(minutes % 60).padStart(2, '0')}`
52}
53
54function isPerson(origin: PromptOrigin): boolean {
55  return origin.kind === 'composer' || origin.kind === 'bridge'
56}
57
58async function textsOf($: EngineInterface): Promise<Texts> {
59  const { language } = await $.settings.read()
60  return typeof language === 'string' && /^(fr|french|fran[cç]ais)/i.test(language)
61    ? TEXTS.fr
62    : TEXTS.en
63}
64
65/**
66 * The dialog of a cold resume, before anything is typed.
67 */
68async function askAtResume($: EngineInterface, cold: Cold): Promise<void> {
69  const texts = await textsOf($)
70  const choice = await $.ui
71    .ask(texts.question(cold), {
72      header: 'Cache',
73      options: [texts.keep, texts.compactNow, texts.clear],
74    })
75    .catch(() => undefined)
76  if (choice === undefined) {
77    return
78  }
79  await update($, isSettled, () => true)
80  if (choice === texts.compactNow) {
81    await $.session.compact()
82  } else if (choice === texts.clear) {
83    await $.command.run({ command: 'clear' })
84  }
85}
86
87// A prompt.submit hook cannot compact or run a command itself: both would wait
88// on the turn it holds. These run once its dispatch has ended.
89
90async function compactThenRefill($: EngineInterface, text: string): Promise<void> {
91  await $.session.compact()
92  // Refilled rather than resubmitted: the model would read a plugin's prompt
93  // as the plugin's message, not the person's.
94  await $.prompt.fill({ text })
95}
96
97async function clearThenRefill($: EngineInterface, text: string): Promise<void> {
98  await $.command.run({ command: 'clear' })
99  await $.prompt.fill({ text })
100}
101
102export const register: Register = (on, options) => {
103  const ttlMs = Number(options.ttlMinutes ?? 60) * 60_000
104  const minTokens = Number(options.minTokens ?? 30_000)
105
106  // Claude Code measures a resumed transcript itself, with the session's real
107  // cache lifetime; that estimate beats the configured one.
108  on('classic.SessionStart', async ($, e, next) => {
109    const result = await next(e)
110    const seconds = e.seconds_since_last_response
111    if (seconds === undefined) {
112      return result
113    }
114    const now = await $.clock.now()
115    await update($, lastAnswerAtMs, () => now - seconds * 1000)
116    await update($, resumedTokens, () => e.context_tokens ?? null)
117    await update($, isSettled, () => false)
118    const tokens = e.context_tokens ?? 0
119    if (e.prompt_cache_likely_expired && tokens >= minTokens) {
120      // Not awaited: the session goes on loading while the dialog waits.
121      void askAtResume($, {
122        idleMs: seconds * 1000,
123        tokens,
124        usd: e.estimated_cache_write_usd,
125      }).catch(error => $.ui.log(`resume question failed: ${error}`, { to: 'debug' }))
126    }
127    return result
128  })
129
130  on('turn.complete', async ($, e, next) => {
131    const result = await next(e)
132    const now = await $.clock.now()
133    await update($, lastAnswerAtMs, () => now)
134    await update($, resumedTokens, () => null)
135    await update($, isSettled, () => false)
136    return result
137  })
138
139  on('prompt.submit', async ($, e, next) => {
140    if (!isPerson(e.origin) || e.turnId !== undefined || (await read($, isSettled))) {
141      return next(e)
142    }
143    const last = await read($, lastAnswerAtMs)
144    if (last === null) {
145      return next(e)
146    }
147    const idleMs = (await $.clock.now()) - last
148    const tokens =
149      (await $.session.usage()).context.tokens ?? (await read($, resumedTokens)) ?? 0
150    if (idleMs < ttlMs || tokens < minTokens) {
151      return next(e)
152    }
153
154    const texts = await textsOf($)
155    const choice = await $.ui
156      .ask(texts.question({ idleMs, tokens }), {
157        header: 'Cache',
158        options: [texts.send, texts.compactFirst, texts.clear, texts.cancel],
159      })
160      .catch(() => texts.cancel)
161
162    if (choice === texts.send) {
163      await update($, isSettled, () => true)
164      return next(e)
165    }
166    if (choice === texts.compactFirst) {
167      await update($, isSettled, () => true)
168      $.clock.after(0, () => void compactThenRefill($, e.text).catch(() => {}))
169      return { drop: texts.compacting }
170    }
171    if (choice === texts.clear) {
172      $.clock.after(0, () => void clearThenRefill($, e.text).catch(() => {}))
173      return { drop: texts.cleared }
174    }
175    await $.prompt.fill({ text: e.text })
176    return { drop: texts.cancelled }
177  })
178}
179

1declare module 'claude-code' {
2  interface PluginState {
3    'cold-cache-guard': {
4      /** Epoch ms of the session's last main-thread answer; null before the first one. */
5      lastAnswerAtMs: number | null
6      /** Tokens a resumed transcript's next request re-sends, until a live figure exists. */
7      resumedTokens: number | null
8      /** The person already chose for the current cold spell; the next answer clears it. */
9      isSettled: boolean
10    }
11  }
12}
13