cache-tax

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { CacheSettings } from '../types'
5
6const settings = atom(
7  { plugin: 'cache-tax', key: 'settings' } as const,
8  { ttlMinutes: 60, keepWarm: false, maxPings: 8 } as CacheSettings,
9)
10
11const PING = '[keep-warm ping] Reply with only the word: ok'
12// Below this many context tokens a cold start is cheap; don't interrupt.
13const MIN_TOKENS = 20_000
14const CONFIRM_WINDOW_MS = 3 * 60_000
15
16// USD per million input tokens, API list prices. An estimate only.
17const INPUT_PRICE: Array<[RegExp, number]> = [
18  [/haiku/i, 1],
19  [/sonnet/i, 3],
20  [/opus|fable/i, 5],
21]
22const priceOf = (model: string) => INPUT_PRICE.find(([re]) => re.test(model))?.[1] ?? 3
23
24const fmtTok = (n: number) => (n >= 1e6 ? `${(n / 1e6).toFixed(2)}M` : `${Math.round(n / 1000)}k`)
25
26let lastActivity = 0
27let lastContext = 0
28let isWorking = false
29let pings = 0
30let pending: { text: string; at: number } | undefined
31
32async function save($: EngineInterface) {
33  await $.store.set('settings', await read($, settings))
34}
35
36async function refreshStatus($: EngineInterface) {
37  if (lastActivity === 0) return $.ui.status(undefined)
38  const s = await read($, settings)
39  const left = s.ttlMinutes * 60_000 - ((await $.clock.now()) - lastActivity)
40  const warm = s.keepWarm ? ' · keepwarm' : ''
41  $.ui.status(left > 0 ? `cache warm · ${Math.ceil(left / 60_000)}m left${warm}` : `cache cold${warm}`)
42}
43
44export const register: Register = on => {
45  on('session.start', async ($, e, next) => {
46    const saved = (await $.store.get('settings')) as Partial<CacheSettings> | undefined
47    if (saved) await update($, settings, s => ({ ...s, ...saved }))
48    await $.command.register({
49      name: 'keepwarm',
50      description: 'Cache Tax: ping the session just before the prompt cache goes cold (on/off)',
51      argumentHint: 'on | off',
52    })
53    await $.command.register({
54      name: 'cachetax',
55      description: 'Cache Tax: show cache state, or set the TTL: /cachetax ttl <minutes>',
56      argumentHint: 'ttl <minutes> | pings <n>',
57    })
58    $.clock.every(60_000, () => {
59      void (async () => {
60        await refreshStatus($)
61        const s = await read($, settings)
62        if (!s.keepWarm || isWorking || lastActivity === 0 || pings >= s.maxPings) return
63        const idle = (await $.clock.now()) - lastActivity
64        const ttl = s.ttlMinutes * 60_000
65        // Fire inside the last four minutes before expiry.
66        if (idle > ttl - 4 * 60_000 && idle < ttl) {
67          pings += 1
68          await $.prompt.submit({ text: PING })
69        }
70      })()
71    })
72    return next(e)
73  })
74
75  on('command.run', { command: 'keepwarm' }, async ($, e) => {
76    const arg = e.args.trim().toLowerCase()
77    const current = await read($, settings)
78    const keepWarm = arg === 'off' ? false : arg === 'on' ? true : !current.keepWarm
79    await update($, settings, s => ({ ...s, keepWarm }))
80    await save($)
81    pings = 0
82    await refreshStatus($)
83    return {
84      text: keepWarm
85        ? `Keep-warm on: up to ${current.maxPings} tiny pings, each just before the ${current.ttlMinutes}-minute cache expires.`
86        : 'Keep-warm off.',
87    }
88  })
89
90  on('command.run', { command: 'cachetax' }, async ($, e) => {
91    const [key, value] = e.args.trim().split(/\s+/)
92    const n = Number(value)
93    if (key === 'ttl' && n > 0) await update($, settings, s => ({ ...s, ttlMinutes: n }))
94    else if (key === 'pings' && n >= 0) await update($, settings, s => ({ ...s, maxPings: n }))
95    await save($)
96    const s = await read($, settings)
97    const idle = lastActivity ? Math.round(((await $.clock.now()) - lastActivity) / 60_000) : 0
98    return {
99      text: `TTL ${s.ttlMinutes}m · keep-warm ${s.keepWarm ? 'on' : 'off'} (max ${s.maxPings} pings) · idle ${idle}m · context ≈${fmtTok(lastContext)} tokens`,
100    }
101  })
102
103  on('turn.start', async ($, e, next) => {
104    isWorking = true
105    if (e.text !== PING) pings = 0
106    return next(e)
107  })
108
109  on('turn.step', async function* ($, e, next) {
110    const result = yield* next(e)
111    if (!e.agentId && result.usage) {
112      const u = result.usage
113      lastContext = u.input_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens + u.output_tokens
114      lastActivity = await $.clock.now()
115    }
116    return result
117  })
118
119  on('turn.complete', async ($, e, next) => {
120    if (!e.agentId) {
121      isWorking = false
122      lastActivity = await $.clock.now()
123      await refreshStatus($)
124    }
125    return next(e)
126  })
127
128  on('prompt.submit', async ($, e, next) => {
129    const kind = e.origin.kind
130    if (e.text === PING || (kind !== 'composer' && kind !== 'sdk' && kind !== 'bridge') || lastActivity === 0) return next(e)
131    const s = await read($, settings)
132    const now = await $.clock.now()
133    const idle = now - lastActivity
134    if (idle < s.ttlMinutes * 60_000) return next(e)
135
136    const usage = await $.session.usage().catch(() => undefined)
137    const tokens = usage?.context.tokens ?? lastContext
138    if (tokens < MIN_TOKENS) return next(e)
139
140    if (pending && pending.text === e.text && now - pending.at < CONFIRM_WINDOW_MS) {
141      pending = undefined
142      return next(e)
143    }
144
145    const price = priceOf(await $.session.model())
146    const writeRate = s.ttlMinutes >= 60 ? 2 : 1.25
147    const cold = (tokens * price * writeRate) / 1e6
148    const warm = (tokens * price * 0.1) / 1e6
149    const minutes = Math.round(idle / 60_000)
150    pending = { text: e.text, at: now }
151    const text = e.text
152    // Put the prompt back in the composer so Enter sends it.
153    $.clock.after(30, () => void $.prompt.fill({ text, mode: 'replace' }))
154    // The drop reason sits behind "View details" in the desktop app; say it where it is seen.
155    $.ui.toast(
156      `Cache Tax held your prompt: the cache went cold ${minutes}m ago, so sending re-reads ~${fmtTok(tokens)} tokens (≈$${cold.toFixed(2)}). Press Enter again to send it.`,
157      { timeoutMs: 12000 },
158    )
159    $.ui.status('cache cold · held one prompt · press Enter again to send')
160    return {
161      drop:
162        `Cache Tax: the prompt cache went cold ${minutes}m ago. Sending re-reads ~${fmtTok(tokens)} tokens ` +
163        `(≈$${cold.toFixed(2)} API-equivalent vs ≈$${warm.toFixed(2)} warm, and more of your usage limit). ` +
164        `Press Enter again within 3 minutes to send anyway, or /compact first to shrink it. Tip: /keepwarm.`,
165    }
166  }).catch(($, e, next) => next(e))
167}
168

1export type CacheSettings = { ttlMinutes: number; keepWarm: boolean; maxPings: number }
2
3declare module 'claude-code' {
4  interface PluginState {
5    'cache-tax': { settings: CacheSettings }
6  }
7}
8