cache-guard

1import type { Register } from 'claude-code'
2
3// Adapted from the ask guard in this repo's former cache-tax fork (karanb192/cache-tax 2.2.1, MIT):
4// the dialog and the state it needs to price a cold send, without keepwarm.
5const TTL_MS = 60 * 60 * 1000
6const BIG_TOKENS = 50000
7
8// $ per million tokens, [cache read, 1h cache write], list prices September 2026.
9// Longer family names first: a model id matches the first row it contains.
10const PRICES: Array<[string, number, number]> = [
11  ['fable-5-1', 0.25, 20],
12  ['fable-5', 1, 20],
13  ['opus-5', 0.5, 10],
14  ['opus-4', 0.5, 10],
15  ['sonnet-5', 0.2, 4],
16  ['sonnet', 0.3, 6],
17  ['haiku', 0.1, 2],
18]
19
20export type State = {
21  lastRequestAt: number
22  lastModel: string | null
23  ctx: number
24  compacted: boolean
25  ackedAt: number
26}
27
28export function freshState(): State {
29  return { lastRequestAt: 0, lastModel: null, ctx: 0, compacted: false, ackedAt: 0 }
30}
31
32function priceOf(model: string | null): [number, number] | null {
33  const m = (model ?? '').toLowerCase().replace(/[\s.]+/g, '-')
34  for (const [family, read, write] of PRICES) if (m.includes(family)) return [read, write]
35  return null
36}
37
38export function fmtDuration(ms: number): string {
39  const total = Math.max(0, Math.round(ms / 60000))
40  const h = Math.floor(total / 60)
41  const m = total % 60
42  if (h >= 48) return `${Math.floor(h / 24)}d ${h % 24}h`
43  return h > 0 ? `${h}h${String(m).padStart(2, '0')}m` : `${m}m`
44}
45
46function fmtUsd(usd: number | null): string {
47  return usd == null ? 'n/a' : '$' + (usd >= 100 ? usd.toFixed(0) : usd.toFixed(2))
48}
49
50function isCold(s: State, now: number): boolean {
51  return s.lastRequestAt > 0 && !s.compacted && now - s.lastRequestAt >= TTL_MS
52}
53
54export function guardText(s: State, now: number): string {
55  const price = priceOf(s.lastModel)
56  const rate = price ? `$${price[1]}/MTok` : 'the cache-write rate'
57  const cold = price ? fmtUsd(s.ctx * price[1] / 1e6) : 'n/a'
58  const warm = price ? ` (a warm turn would have cost ${fmtUsd(s.ctx * price[0] / 1e6)})` : ''
59  return `the prompt cache went cold ${fmtDuration(now - s.lastRequestAt - TTL_MS)} ago. Sending this re-writes ` +
60    `up to ${s.ctx.toLocaleString('en-US')} tokens at ${rate} = ${cold}${warm}.`
61}
62
63export type ResumeFields = {
64  source: string
65  model?: string
66  context_tokens?: number
67  seconds_since_last_response?: number
68}
69
70/** A resumed or forked session starts with the transcript's clock, so its first send can be cold. */
71export function seedFromResume(s: State, e: ResumeFields, now: number) {
72  if (e.source !== 'resume' && e.source !== 'fork') return
73  if (typeof e.context_tokens === 'number' && e.context_tokens > 0) s.ctx = e.context_tokens
74  if (typeof e.seconds_since_last_response === 'number') s.lastRequestAt = now - e.seconds_since_last_response * 1000
75  if (typeof e.model === 'string') s.lastModel = e.model
76  s.compacted = false
77}
78
79export const register: Register = on => {
80  let s = freshState()
81
82  on('session.start', async ($, e, next) => {
83    const r = await next(e)
84    s.ctx = (await $.session.usage()).context.tokens ?? 0
85    // cache-tax in refuse or ask mode guards the same send; two dialogs for one prompt.
86    if ((await $.command.list()).some(c => c.name === 'cache-tax')) {
87      $.ui.log('cache-tax is also installed, so a cold send can be guarded twice. Run /cache-tax guard warn, or uninstall cache-tax.')
88    }
89    return r
90  })
91
92  on('classic.SessionStart', async ($, e, next) => {
93    const r = await next(e)
94    // /clear starts a new conversation in the same process; nothing priced before it still exists.
95    if (e.source === 'clear') s = freshState()
96    else seedFromResume(s, e, await $.clock.now())
97    // The resume payload may omit the model; without it the dialog cannot price the cold write.
98    if (!s.lastModel) s.lastModel = await $.session.model()
99    return r
100  })
101
102  on('prompt.submit', async ($, e, next) => {
103    if (e.origin.kind === 'plugin') return next(e)
104    if (typeof e.text !== 'string' || e.text.trimStart().startsWith('/')) return next(e)
105    const now = await $.clock.now()
106    if (!isCold(s, now) || s.ctx < BIG_TOKENS) return next(e)
107    // The resend after a dismissed dialog goes through.
108    if (s.ackedAt === s.lastRequestAt) {
109      s.ackedAt = 0
110      return next(e)
111    }
112    s.ackedAt = s.lastRequestAt
113    const text = e.text
114    const refill = () => $.prompt.fill({ text }).catch(() => undefined)
115    let pick = ''
116    try {
117      pick = await $.ui.ask(`${guardText(s, now)} Send it anyway?`, { header: 'Cache cold', options: ['Proceed', 'Start new session'] })
118    } catch {
119      // Dismissed, or no one to ask (-p): refuse once, with the prompt back in the box.
120      void refill()
121      return { drop: `cache-guard: ${guardText(s, now)} Your prompt is back in the box; send it again to pay it.` }
122    }
123    if (pick === 'Proceed') {
124      s.ackedAt = 0
125      return next(e)
126    }
127    if (pick === 'Start new session') {
128      // command.run can't be called from this hook (it would wait on the turn the hook holds), so it runs once the hook returns.
129      $.clock.after(0, () => { void $.command.run({ command: 'clear', args: '' }).then(refill, refill) })
130      return { drop: 'cache-guard: starting fresh with /clear; your prompt is back in the box.' }
131    }
132    void refill()
133    return { drop: 'cache-guard: not sent; your prompt is back in the box.' }
134  })
135
136  on('turn.step', async function* ($, e, next) {
137    if (!e.agentId) s.lastRequestAt = await $.clock.now()
138    yield* next(e)
139  })
140
141  on('turn.complete', async ($, e, next) => {
142    const r = await next(e)
143    if (e.agentId) return r
144    const now = await $.clock.now()
145    // turn.step stamps the exact request time; when no step of this turn did, the turn's end is the floor.
146    if (now - s.lastRequestAt > e.durationMs) s.lastRequestAt = now
147    s.compacted = false
148    s.ackedAt = 0
149    const u = e.usage
150    if (u) {
151      if (u.model) s.lastModel = u.model
152      // A turn's usage sums its responses; the engine's live figure is the real window, the sum only a fallback.
153      const live = (await $.session.usage()).context.tokens
154      s.ctx = live && live > 0 ? live : u.input_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens
155    }
156    return r
157  })
158
159  on('session.compact', async ($, e, next) => {
160    const r = await next(e)
161    if (!e.agentId) {
162      s.compacted = true
163      s.ctx = 0
164      s.ackedAt = 0
165    }
166    return r
167  })
168}
169