cache-guard

1import { atom, read, update } from 'claude-code'
2import type { AgentInfo, EngineInterface, Register, Timer, TurnUsage } from 'claude-code'
3
4import {
5  agentAlerts,
6  agentMissToast,
7  agentSegment,
8  MINUTE,
9  MISSES_KEPT,
10  asBounds,
11  asMisses,
12  effectiveTtl,
13  gateDecision,
14  gateQuestion,
15  joinSegments,
16  missToast,
17  observeStep,
18  pruneTracks,
19  segmentFile,
20  segmentFor,
21} from './rules'
22
23const last = atom({ plugin: 'cache-guard', key: 'last' } as const, null)
24const misses = atom({ plugin: 'cache-guard', key: 'misses' } as const, [])
25const agents = atom({ plugin: 'cache-guard', key: 'agents' } as const, {})
26const listed = atom({ plugin: 'cache-guard', key: 'listed' } as const, [])
27const dismissed = atom({ plugin: 'cache-guard', key: 'dismissed' } as const, [])
28
29const SEND = 'Send anyway'
30const COMPACT = 'Compact first'
31const DRAFT = 'Keep as draft'
32const SEGMENT_EVERY_MS = 30_000
33
34type Config = { defaultTtlMs: number; agentTtlMs: number; gateUsd: number }
35
36async function ttlMs($: EngineInterface, config: Config): Promise<number> {
37  return effectiveTtl(asBounds(await $.store.get('ttl')), config.defaultTtlMs)
38}
39
40async function agentTtlMs($: EngineInterface, config: Config): Promise<number> {
41  return effectiveTtl(asBounds(await $.store.get('agentTtl')), config.agentTtlMs)
42}
43
44async function segmentPath($: EngineInterface, sessionId: string): Promise<string> {
45  const state = (await $.env.get('XDG_STATE_HOME')) || `${await $.env.get('HOME')}/.local/state`
46  return `${state}/harness/status/${sessionId}/cache-guard.json`
47}
48
49async function writeSegment($: EngineInterface, config: Config): Promise<void> {
50  const now = await $.clock.now()
51  const segment = joinSegments(
52    segmentFor(await read($, last), now, await ttlMs($, config)),
53    agentSegment(await read($, agents), await read($, listed), now, await agentTtlMs($, config)),
54  )
55  await $.fs.write(await segmentPath($, await $.session.id()), segmentFile(segment, await $.clock.now()))
56}
57
58async function showSegment($: EngineInterface, config: Config): Promise<void> {
59  try {
60    await writeSegment($, config)
61  } catch (err) {
62    $.ui.log(`cache-guard: could not write the status segment: ${String(err)}`, { to: 'debug' })
63  }
64}
65
66// rmdir fails harmlessly while other mods' files remain, because the session folder is shared.
67async function removeSegment($: EngineInterface, sessionId: string): Promise<void> {
68  const path = await segmentPath($, sessionId)
69  try {
70    await $.process.run(['rm', '-f', path])
71    await $.process.run(['rmdir', path.slice(0, path.lastIndexOf('/'))])
72  } catch (err) {
73    $.ui.log(`cache-guard: could not remove ${path} (${String(err)}), emptying it`, { to: 'debug' })
74    await $.fs.write(path, '')
75  }
76}
77
78function listedAgent(info: AgentInfo) {
79  return { id: info.id, type: info.type, description: info.description, status: info.status }
80}
81
82async function refreshAgents($: EngineInterface): Promise<void> {
83  const list = (await $.agent.list()).map(listedAgent)
84  const now = await $.clock.now()
85  await update($, listed, () => list)
86  await update($, agents, tracked => pruneTracks(tracked, now))
87}
88
89async function tick($: EngineInterface, config: Config): Promise<void> {
90  try {
91    await refreshAgents($)
92  } catch (err) {
93    $.ui.log(`cache-guard: could not list agents: ${String(err)}`, { to: 'debug' })
94  }
95  await showSegment($, config)
96}
97
98async function forget($: EngineInterface, config: Config): Promise<void> {
99  await update($, last, () => null)
100  await showSegment($, config)
101}
102
103async function observe($: EngineInterface, config: Config, startedAt: number, usage: TurnUsage): Promise<void> {
104  const seen = observeStep({
105    prev: await read($, last),
106    startedAt,
107    now: await $.clock.now(),
108    usage,
109    bounds: asBounds(await $.store.get('ttl')),
110    defaultTtlMs: config.defaultTtlMs,
111  })
112  await update($, last, () => seen.last)
113  await $.store.set('ttl', seen.bounds)
114  const miss = seen.miss
115  if (!miss) return
116  const kept = await update($, misses, list => [...list, miss].slice(-MISSES_KEPT))
117  await $.store.set('misses', kept)
118  $.ui.toast(missToast(miss))
119}
120
121async function observeAgent(
122  $: EngineInterface,
123  config: Config,
124  agentId: string,
125  startedAt: number,
126  usage: TurnUsage,
127): Promise<void> {
128  const prev = (await read($, agents))[agentId] ?? null
129  const seen = observeStep({
130    prev,
131    startedAt,
132    now: await $.clock.now(),
133    usage,
134    bounds: asBounds(await $.store.get('agentTtl')),
135    defaultTtlMs: config.agentTtlMs,
136  })
137  const track = { ...seen.last, firstAt: prev?.firstAt ?? startedAt }
138  await update($, agents, tracked => ({ ...tracked, [agentId]: track }))
139  await $.store.set('agentTtl', seen.bounds)
140  if (!seen.miss) return
141  const miss = { ...seen.miss, agentId }
142  const kept = await update($, misses, list => [...list, miss].slice(-MISSES_KEPT))
143  await $.store.set('misses', kept)
144  const type = (await $.agent.list()).find(a => a.id === agentId)?.type ?? 'subagent'
145  $.ui.toast(agentMissToast(type, miss))
146}
147
148async function compactThenRestore($: EngineInterface, config: Config, text: string): Promise<void> {
149  try {
150    const compacted = await $.session.compact()
151    if (compacted.messages) await forget($, config)
152    else $.ui.toast(`Compaction skipped: ${compacted.skip}`)
153  } catch (err) {
154    $.ui.toast(`Compaction failed: ${String(err)}`)
155  }
156  await $.prompt.fill({ text })
157}
158
159export const register: Register = (on, options) => {
160  const config: Config = {
161    defaultTtlMs: Number(options.ttlMinutes ?? 60) * MINUTE,
162    agentTtlMs: Number(options.agentTtlMinutes ?? 5) * MINUTE,
163    gateUsd: Number(options.gateUsd ?? 1),
164  }
165  let ticker: Timer | undefined
166
167  on('session.start', async ($, e, next) => {
168    const stored = asMisses(await $.store.get('misses'))
169    await update($, misses, () => stored)
170    await tick($, config)
171    ticker?.cancel()
172    ticker = $.clock.every(SEGMENT_EVERY_MS, () => void tick($, config))
173    return next(e)
174  })
175
176  on('session.end', async ($, e, next) => {
177    ticker?.cancel()
178    ticker = undefined
179    if (e.reason === 'clear') await update($, last, () => null)
180    await removeSegment($, e.sessionId)
181    return next(e)
182  })
183
184  // Per-step usage, because turn.complete sums every request of the turn and a
185  // cold first request would hide behind the warm reads of the steps after it.
186  on('turn.step', async function* ($, e, next) {
187    const startedAt = await $.clock.now()
188    const response = yield* next(e)
189    if (response.usage) {
190      try {
191        if (e.agentId === undefined) await observe($, config, startedAt, response.usage)
192        else await observeAgent($, config, e.agentId, startedAt, response.usage)
193      } catch (err) {
194        $.ui.log(`cache-guard: could not record step usage: ${String(err)}`, { to: 'debug' })
195      }
196    }
197    return response
198  })
199
200  on('turn.complete', async ($, e, next) => {
201    await tick($, config)
202    return next(e)
203  })
204
205  on('session.compact', async ($, e, next) => {
206    const result = await next(e)
207    if (e.agentId === undefined && e.trigger !== 'precompute' && result.messages) await forget($, config)
208    return result
209  }).catch(($, e, next) => next(e))
210
211  // Fails open because this is a cost warning: a broken guard must never eat a prompt.
212  on('prompt.submit', async ($, e, next) => {
213    // Attachments cannot be put back in the prompt box, so a prompt carrying them is never held.
214    if (e.origin.kind !== 'composer' || e.turnId !== undefined || e.attachments) return next(e)
215    const gate = gateDecision(await read($, last), await $.clock.now(), await ttlMs($, config), config.gateUsd)
216    if (!gate) return next(e)
217
218    const choice = await $.ui
219      .ask(gateQuestion(gate), { header: 'Cold cache', options: [SEND, COMPACT, DRAFT] })
220      .catch((err: unknown) => {
221        $.ui.log(`cache-guard: cold cache question unanswered (${String(err)}), sending`, { to: 'debug' })
222        return SEND
223      })
224
225    if (choice === COMPACT) {
226      // The host refuses a compaction under the prompt this hook is holding, so it runs once the prompt is dropped.
227      $.clock.after(0, () => void compactThenRestore($, config, e.text))
228      return { drop: 'Compacting first; your prompt will be back in the box.' }
229    }
230    if (choice === DRAFT) {
231      await $.prompt.fill({ text: e.text })
232      return { drop: 'Kept as a draft; the prompt cache is still cold.' }
233    }
234    return next(e)
235  }).catch(($, e, next) => next(e))
236
237  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
238    if (e.props.hasSurvey) return next(e)
239    const alerts = agentAlerts(
240      await read($, agents),
241      await read($, listed),
242      await $.clock.now(),
243      await agentTtlMs($, config),
244      await read($, dismissed),
245    )
246    if (alerts.rows.length === 0) return next(e)
247    const { Box, Button, Text } = $.ui.resolve(e)
248    return (
249      <Box flexDirection="column">
250        {alerts.rows.map(row => (
251          <Box key={row.key} flexDirection="row" flexWrap="wrap" columnGap={1}>
252            <Text color="warning">{row.text}</Text>
253            <Button
254              key={`dismiss-${row.key}`}
255              label="Dismiss"
256              dimColor
257              role="dismiss"
258              onPress={() => update($, dismissed, keys => [...keys, row.key])}
259            />
260          </Box>
261        ))}
262        {alerts.more > 0 && <Text dimColor>+{alerts.more} more</Text>}
263        {await next(e)}
264      </Box>
265    )
266  })
267}
268

1import type { ModelUsage, TurnUsage } from 'claude-code'
2
3import type { AgentTrack, LastResponse, ListedAgent, Miss, TtlBounds } from '../types'
4
5export const MINUTE = 60_000
6export const MISSES_KEPT = 20
7const MILLION = 1_000_000
8const FIVE_MINUTES = 5 * MINUTE
9const MIN_MISS_TOKENS = 50_000
10const MIN_STATUS_CTX = 20_000
11const MIN_AGENT_ALERT_TOKENS = 50_000
12const LONG_AGENT_MS = 60 * MINUTE
13const MAX_ALERT_ROWS = 3
14const TRACK_KEEP_MS = 24 * 60 * MINUTE
15
16export type Prices = { input: number; write1h: number; write5m: number; read: number }
17
18type PriceRow = { prefix: string; prices: Prices; tier?: { above: number; factor: number } }
19
20const OPUS_5_5: PriceRow = { prefix: 'claude-opus-5-5', prices: { input: 4, write1h: 8, write5m: 5, read: 0.2 } }
21
22const PRICE_ROWS: readonly PriceRow[] = [
23  OPUS_5_5,
24  { prefix: 'claude-sonnet-5-5', prices: { input: 2, write1h: 4, write5m: 2.5, read: 0.2 } },
25  { prefix: 'claude-fable-5-1', prices: { input: 10, write1h: 20, write5m: 12.5, read: 0.25 } },
26  {
27    prefix: 'claude-haiku-5-5',
28    prices: { input: 0.1, write1h: 0.2, write5m: 0.125, read: 0.01 },
29    tier: { above: 100_000, factor: 5 },
30  },
31]
32
33// Bedrock and Vertex ids carry a region or vendor prefix, so matching starts at "claude-".
34export function priceFor(model: string, tokens: number): Prices {
35  const id = model.slice(Math.max(0, model.indexOf('claude-')))
36  const row = PRICE_ROWS.find(r => id.startsWith(r.prefix)) ?? OPUS_5_5
37  if (!row.tier || tokens <= row.tier.above) return row.prices
38  const { factor } = row.tier
39  const p = row.prices
40  return { input: p.input * factor, write1h: p.write1h * factor, write5m: p.write5m * factor, read: p.read * factor }
41}
42
43export function contextOf(usage: ModelUsage): number {
44  return usage.input_tokens + usage.cache_read_input_tokens + usage.cache_creation_input_tokens
45}
46
47// A cache that lives longer than five minutes is the 1-hour cache, which writes at the higher rate.
48export function rewriteUsd(model: string, tokens: number, ttlMs: number): number {
49  const prices = priceFor(model, tokens)
50  return (tokens * (ttlMs > FIVE_MINUTES ? prices.write1h : prices.write5m)) / MILLION
51}
52
53export function compactUsd(model: string, tokens: number): number {
54  return (tokens * priceFor(model, tokens).input) / MILLION
55}
56
57export type StepKind = 'warm' | 'cold'
58
59export function classifyStep(prevCtx: number, usage: ModelUsage): StepKind | null {
60  if (prevCtx <= 0) return null
61  if (usage.cache_read_input_tokens >= 0.5 * prevCtx) return 'warm'
62  // A cold step may still read some cache, because the system and tools prefix can stay cached on its own, a large share in a subagent.
63  if (usage.cache_creation_input_tokens >= 0.5 * prevCtx) return 'cold'
64  return null
65}
66
67export const noBounds: TtlBounds = { warmMaxMs: null, coldMinMs: null }
68
69// Invariant: warmMaxMs < coldMinMs. An observation that breaks it (a plan or TTL change) replaces the bounds.
70export function learnTtl(bounds: TtlBounds, kind: StepKind, gapMs: number): TtlBounds {
71  if (kind === 'warm') {
72    if (bounds.coldMinMs !== null && gapMs >= bounds.coldMinMs) return { warmMaxMs: gapMs, coldMinMs: null }
73    return { ...bounds, warmMaxMs: Math.max(bounds.warmMaxMs ?? 0, gapMs) }
74  }
75  if (bounds.warmMaxMs !== null && gapMs <= bounds.warmMaxMs) return { warmMaxMs: null, coldMinMs: gapMs }
76  return { ...bounds, coldMinMs: Math.min(bounds.coldMinMs ?? gapMs, gapMs) }
77}
78
79export function effectiveTtl(bounds: TtlBounds, defaultMs: number): number {
80  let ttl = defaultMs
81  if (bounds.warmMaxMs !== null) ttl = Math.max(ttl, bounds.warmMaxMs)
82  if (bounds.coldMinMs !== null) ttl = Math.min(ttl, bounds.coldMinMs)
83  return ttl
84}
85
86function finiteOrNull(value: unknown): number | null {
87  return typeof value === 'number' && Number.isFinite(value) ? value : null
88}
89
90export function asBounds(stored: unknown): TtlBounds {
91  if (typeof stored !== 'object' || stored === null) return noBounds
92  const { warmMaxMs, coldMinMs } = stored as Record<string, unknown>
93  return { warmMaxMs: finiteOrNull(warmMaxMs), coldMinMs: finiteOrNull(coldMinMs) }
94}
95
96export function asMisses(stored: unknown): Miss[] {
97  if (!Array.isArray(stored)) return []
98  return stored
99    .filter((m): m is Miss => typeof m === 'object' && m !== null && typeof (m as Miss).tokens === 'number')
100    .slice(-MISSES_KEPT)
101}
102
103export type StepObservation = {
104  prev: LastResponse | null
105  startedAt: number
106  now: number
107  usage: TurnUsage
108  bounds: TtlBounds
109  defaultTtlMs: number
110}
111
112export type Observed = { last: LastResponse; bounds: TtlBounds; miss: Miss | null }
113
114export function observeStep(o: StepObservation): Observed {
115  const last = { at: o.now, ctx: contextOf(o.usage), model: o.usage.model }
116  const unchanged = { last, bounds: o.bounds, miss: null }
117  // A new model starts its own cache, so its first step says nothing about expiry.
118  if (!o.prev || o.prev.model !== o.usage.model) return unchanged
119  const gapMs = o.startedAt - o.prev.at
120  const kind = classifyStep(o.prev.ctx, o.usage)
121  // A cold step within minutes of the last response means the prefix changed (tools, system prompt), not expiry.
122  if (kind === null || (kind === 'cold' && gapMs < FIVE_MINUTES)) return unchanged
123  const bounds = learnTtl(o.bounds, kind, gapMs)
124  const tokens = o.usage.cache_creation_input_tokens
125  if (kind === 'warm' || tokens < MIN_MISS_TOKENS) return { last, bounds, miss: null }
126  const usd = rewriteUsd(o.usage.model, tokens, effectiveTtl(bounds, o.defaultTtlMs))
127  return { last, bounds, miss: { at: o.now, gapMs, tokens, usd, model: o.usage.model } }
128}
129
130export type Gate = { idleMs: number; tokens: number; rewriteUsd: number; compactUsd: number }
131
132export function gateDecision(last: LastResponse | null, now: number, ttlMs: number, gateUsd: number): Gate | null {
133  if (!last) return null
134  const idleMs = now - last.at
135  if (idleMs < ttlMs) return null
136  const rewrite = rewriteUsd(last.model, last.ctx, ttlMs)
137  if (rewrite < gateUsd) return null
138  return { idleMs, tokens: last.ctx, rewriteUsd: rewrite, compactUsd: compactUsd(last.model, last.ctx) }
139}
140
141export function formatTokens(tokens: number): string {
142  if (tokens >= MILLION) return `${(tokens / MILLION).toFixed(1)}M`
143  if (tokens >= 1000) return `${Math.round(tokens / 1000)}k`
144  return String(tokens)
145}
146
147export function formatUsd(usd: number): string {
148  return usd < 0.005 ? '<$0.01' : `$${usd.toFixed(2)}`
149}
150
151export function formatDuration(ms: number): string {
152  if (ms < MINUTE) return `${Math.max(0, Math.floor(ms / 1000))}s`
153  return `${Math.floor(ms / MINUTE)}m`
154}
155
156export type Tone = 'normal' | 'warn'
157export type Segment = { text: string; tone: Tone }
158
159export function segmentFor(last: LastResponse | null, now: number, ttlMs: number): Segment | null {
160  if (!last || last.ctx < MIN_STATUS_CTX) return null
161  const left = last.at + ttlMs - now
162  if (left > 0) return { text: `cache ${formatDuration(left)}`, tone: 'normal' }
163  return { text: `cache cold ≈ ${formatUsd(rewriteUsd(last.model, last.ctx, ttlMs))}`, tone: 'warn' }
164}
165
166// The format must match what statusline.sh reads; tests/test_statusline.sh pins that side.
167export function segmentFile(segment: Segment | null, at: number): string {
168  return segment ? JSON.stringify({ text: segment.text, tone: segment.tone, at }) : ''
169}
170
171export function gateQuestion(gate: Gate): string {
172  return (
173    `The prompt cache has likely expired (idle ${formatDuration(gate.idleMs)}). ` +
174    `Sending re-writes ~${formatTokens(gate.tokens)} tokens (≈ ${formatUsd(gate.rewriteUsd)}); ` +
175    `compacting first reads them once (≈ ${formatUsd(gate.compactUsd)}). What now?`
176  )
177}
178
179export function missToast(miss: Miss): string {
180  return (
181    `Cache went cold after ${formatDuration(miss.gapMs)}: ` +
182    `re-wrote ${formatTokens(miss.tokens)} tokens (≈ ${formatUsd(miss.usd)} at API prices)`
183  )
184}
185
186const ENDED: ReadonlySet<ListedAgent['status']> = new Set(['completed', 'failed', 'killed'])
187const RESTING: ReadonlySet<ListedAgent['status']> = new Set(['idle', 'waiting'])
188
189function isLive(agent: ListedAgent): boolean {
190  return !ENDED.has(agent.status)
191}
192
193function isCold(agent: ListedAgent, track: AgentTrack | undefined, now: number, ttlMs: number): track is AgentTrack {
194  return track !== undefined && RESTING.has(agent.status) && now - track.at >= ttlMs
195}
196
197export type AgentAlert = { key: string; text: string }
198
199// An idle alert is keyed by the agent's last response so that a dismissal lasts only until the agent works and rests again.
200export function agentAlerts(
201  agents: Readonly<Record<string, AgentTrack>>,
202  listed: readonly ListedAgent[],
203  now: number,
204  ttlMs: number,
205  dismissed: readonly string[],
206): { rows: AgentAlert[]; more: number } {
207  const alerts: AgentAlert[] = []
208  for (const agent of listed) {
209    const track = agents[agent.id]
210    if (!track || !isLive(agent)) continue
211    const label = `${agent.type} "${agent.description}"`
212    if (isCold(agent, track, now, ttlMs) && track.ctx >= MIN_AGENT_ALERT_TOKENS) {
213      const usd = formatUsd(rewriteUsd(track.model, track.ctx, ttlMs))
214      alerts.push({
215        key: `${agent.id}@${track.at}`,
216        text: `${label} idle ${formatDuration(now - track.at)} · resuming re-writes ${formatTokens(track.ctx)} ≈ ${usd}`,
217      })
218    } else if (now - track.firstAt >= LONG_AGENT_MS) {
219      alerts.push({ key: `${agent.id}:long`, text: `${label} has run ${formatDuration(now - track.firstAt)}` })
220    }
221  }
222  const open = alerts.filter(alert => !dismissed.includes(alert.key))
223  return { rows: open.slice(0, MAX_ALERT_ROWS), more: Math.max(0, open.length - MAX_ALERT_ROWS) }
224}
225
226export function agentSegment(
227  agents: Readonly<Record<string, AgentTrack>>,
228  listed: readonly ListedAgent[],
229  now: number,
230  ttlMs: number,
231): Segment | null {
232  const live = listed.filter(isLive)
233  if (live.length === 0) return null
234  const cold = live.filter(agent => isCold(agent, agents[agent.id], now, ttlMs)).length
235  if (cold === 0) return { text: `agents ${live.length}`, tone: 'normal' }
236  return { text: `agents ${live.length} · ${cold} cold`, tone: 'warn' }
237}
238
239// Tracks age out by time, not by listing, because ended agents leave $.agent.list within a minute yet can be resumed much later.
240export function pruneTracks(tracks: Readonly<Record<string, AgentTrack>>, now: number): Record<string, AgentTrack> {
241  return Object.fromEntries(Object.entries(tracks).filter(([, track]) => now - track.at < TRACK_KEEP_MS))
242}
243
244export function joinSegments(...segments: (Segment | null)[]): Segment | null {
245  const shown = segments.filter((s): s is Segment => s !== null)
246  if (shown.length === 0) return null
247  return { text: shown.map(s => s.text).join(' · '), tone: shown.some(s => s.tone === 'warn') ? 'warn' : 'normal' }
248}
249
250export function agentMissToast(type: string, miss: Miss): string {
251  return (
252    `Agent ${type} went cold after ${formatDuration(miss.gapMs)}: ` +
253    `re-wrote ${formatTokens(miss.tokens)} tokens (≈ ${formatUsd(miss.usd)})`
254  )
255}
256

1export type LastResponse = {
2  at: number
3  ctx: number
4  model: string
5}
6
7export type TtlBounds = {
8  warmMaxMs: number | null
9  coldMinMs: number | null
10}
11
12export type Miss = {
13  at: number
14  gapMs: number
15  tokens: number
16  usd: number
17  model: string
18  agentId?: string
19}
20
21export type AgentTrack = LastResponse & {
22  firstAt: number
23}
24
25export type ListedAgent = {
26  id: string
27  type: string
28  description: string
29  // AgentStatus spelled out, because the validator requires a contract with no imports.
30  status: 'pending' | 'running' | 'waiting' | 'idle' | 'completed' | 'failed' | 'killed'
31}
32
33declare module 'claude-code' {
34  interface PluginState {
35    'cache-guard': {
36      last: LastResponse | null
37      misses: Miss[]
38      agents: Record<string, AgentTrack>
39      listed: ListedAgent[]
40      dismissed: string[]
41    }
42  }
43}
44