pilot-light

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, ModelUsage, Register } from 'claude-code'
3
4import type { Snap } from '../types'
5
6const MINUTE = 60_000
7const HOUR = 60 * MINUTE
8const WARN_LEAD_MS = 5 * MINUTE
9const PING_LEAD_MS = 20_000
10
11type Config = { configuredTtl: number; keepAlive: boolean }
12
13const INITIAL: Snap = {
14  lastResponseAt: 0,
15  lastRequestAt: 0,
16  ttlMs: null,
17  contextTokens: 0,
18  window: 0,
19  rateLimits: [],
20  costUsd: 0,
21  model: '',
22  isWorking: false,
23  pinging: false,
24  cold: false,
25  keepAlives: 0,
26  pingUsd: 0,
27  lastPingUsd: 0,
28  pingFailed: false,
29  warned: false,
30  savedUsd: 0,
31}
32// The shape tag versions the stored value: bump it whenever Snap gains a field,
33// or a reload reads the previous version's snapshot and arithmetic on its
34// missing fields lands null in the store.
35const snap = atom({ plugin: 'pilot-light', key: 'snap' } as const, INITIAL, { shape: 'v3' })
36
37// What a fresh stretch of pings starts from; also what a person's prompt restores.
38const RESTED = { pinging: false, keepAlives: 0, pingUsd: 0, lastPingUsd: 0, pingFailed: false, warned: false }
39
40// $/MTok per model id fragment: [fragment, input, cache read, output]. First
41// match wins, so the longer ids come first. A cache write is 1.25x input on
42// the 5-minute TTL and 2x on the 1-hour TTL. Unknown models price as Opus
43// 5.5, the engine's own default tier.
44const PRICES: readonly (readonly [string, number, number, number])[] = [
45  ['fable-5-1', 10, 0.25, 50],
46  ['mythos-5-1', 10, 0.25, 50],
47  ['fable-5', 10, 1, 50],
48  ['mythos-5', 10, 1, 50],
49  ['opus-5-5', 4, 0.2, 20],
50  ['opus-5', 5, 0.5, 25],
51  ['opus-4', 5, 0.5, 25],
52  ['sonnet-5', 2, 0.2, 10],
53  ['sonnet-4', 3, 0.3, 15],
54  ['haiku-5', 0.1, 0.01, 0.5],
55  ['haiku-4', 1, 0.1, 5],
56]
57const priceOf = (model: string) => PRICES.find(([id]) => model.includes(id)) ?? PRICES[4]!
58const writeMultOf = (ttl: number) => (ttl >= HOUR ? 2 : 1.25)
59const coldUsdOf = (s: Snap, ttl: number) => (s.contextTokens * priceOf(s.model)[1] * writeMultOf(ttl)) / 1e6
60const usdOf = (u: ModelUsage, model: string, ttl: number) => {
61  const [, input, cacheRead, output] = priceOf(model)
62  return (
63    (u.input_tokens * input +
64      u.cache_read_input_tokens * cacheRead +
65      u.cache_creation_input_tokens * input * writeMultOf(ttl) +
66      u.output_tokens * output) /
67    1e6
68  )
69}
70
71const PING_PROMPT = 'Reply with exactly: ok.'
72
73const HANDOFF_PROMPT =
74  'Write a handoff brief for a fresh session that has none of this context. ' +
75  'Markdown, under 400 words, with these sections: Goal; Current state (done and in progress); ' +
76  'Key decisions and why; Files touched (paths); Open questions and blockers; Next steps. ' +
77  'Write it so the next agent can continue without asking. Output the brief only.'
78
79const fmtK = (n: number) =>
80  n >= 1e6 ? `${(n / 1e6).toFixed(1)}M` : n >= 1e3 ? `${Math.round(n / 1e3)}k` : `${n}`
81const fmtUsd = (n: number) => (n > 0 && n < 0.005 ? '<$0.01' : `$${n.toFixed(2)}`)
82const fmtClock = (ms: number) => {
83  const s = Math.max(0, Math.ceil(ms / 1000))
84  const h = Math.floor(s / 3600)
85  const m = Math.floor((s % 3600) / 60)
86  const ss = `${s % 60}`.padStart(2, '0')
87  return h ? `${h}:${`${m}`.padStart(2, '0')}:${ss}` : `${m}:${ss}`
88}
89const ttlOf = (s: Snap, cfg: Config) => s.ttlMs ?? cfg.configuredTtl
90const warnLeadOf = (ttl: number) => Math.min(WARN_LEAD_MS, ttl / 2)
91
92async function measure($: EngineInterface) {
93  const u = await $.session.usage()
94  const model = await $.session.model()
95  const savedUsd = Number(await $.store.get('savedUsd')) || 0
96  await update($, snap, s => ({
97    ...s,
98    model,
99    savedUsd,
100    contextTokens: u.context.tokens ?? s.contextTokens,
101    window: u.context.window,
102    rateLimits: u.rateLimits.map(r => ({ kind: r.kind, percentUsed: r.percentUsed })),
103    costUsd: u.cost?.usd ?? s.costUsd,
104  }))
105}
106
107// The API reports which TTL it wrote under usage.cache_creation; the engine
108// keeps that only in the transcript file, so the last assistant line is the
109// one ground truth for the window this session actually runs on.
110async function detectTtl($: EngineInterface, transcriptPath: string) {
111  const { stdout } = await $.process.run(['tail', '-c', '1000000', transcriptPath])
112  const last = stdout.match(/"cache_creation":\{[^}]*\}/g)?.at(-1)
113  if (!last) return
114  const count = (key: string) => Number(last.match(new RegExp(`"${key}":(\\d+)`))?.[1] ?? 0)
115  const h1 = count('ephemeral_1h_input_tokens')
116  const m5 = count('ephemeral_5m_input_tokens')
117  if (!h1 && !m5) return
118  await update($, snap, s => ({ ...s, ttlMs: h1 ? HOUR : 5 * MINUTE }))
119}
120
121// A fork re-sends the main thread's last request plus one user message, so
122// the API reads the cached prefix and restarts its timer: no turn, nothing in
123// the transcript. Pinging stops once a stretch has cost as much as the cold
124// re-send it is avoiding, since from there re-caching is the cheaper path.
125async function ping($: EngineInterface, s: Snap, ttl: number) {
126  const cold = coldUsdOf(s, ttl)
127  if (s.pingUsd + s.lastPingUsd > cold) return
128  // Each model has its own cache: after /model the old entry is nothing to keep.
129  if ((await $.session.model()) !== s.model) {
130    await update($, snap, x => ({ ...x, lastResponseAt: 0 }))
131    return
132  }
133  // Claim the flag in the write that reads it, so two ticks overlapping a
134  // slow host round trip fork once.
135  let claimed = false
136  await update($, snap, x => {
137    claimed = !x.pinging
138    return claimed ? { ...x, pinging: true } : x
139  })
140  if (!claimed) return
141  const startedAt = await $.clock.now()
142  let reply
143  try {
144    reply = await $.model.fork({ prompt: PING_PROMPT })
145  } catch (error) {
146    await update($, snap, x => ({ ...x, pinging: false, pingFailed: true }))
147    $.ui.log(`pilot-light: keep-alive ping refused (${String(error)}); no more until you prompt`)
148    return
149  }
150  const served = reply.isAnswered || reply.reason === 'empty-reply'
151  const usd = 'usage' in reply ? usdOf(reply.usage, s.model, ttl) : 0
152  // A prompt or /clear that landed meanwhile already reset the stretch
153  // (pinging false): this ping belonged to the old one and is dropped.
154  await update($, snap, x =>
155    x.pinging
156      ? {
157          ...x,
158          pinging: false,
159          pingFailed: !served,
160          keepAlives: served ? x.keepAlives + 1 : x.keepAlives,
161          pingUsd: x.pingUsd + usd,
162          lastPingUsd: served ? usd : x.lastPingUsd,
163          lastResponseAt: served ? Math.max(x.lastResponseAt, startedAt) : x.lastResponseAt,
164          cold: served ? false : x.cold,
165        }
166      : x,
167  )
168  $.ui.log(
169    served
170      ? `pilot-light: keep-alive ping ${s.keepAlives + 1} cost ${fmtUsd(usd)}, ${fmtUsd(s.pingUsd + usd)} of ${fmtUsd(cold)} this stretch`
171      : `pilot-light: keep-alive ping failed (${reply.isAnswered ? 'unknown' : reply.reason}); no more until you prompt`,
172  )
173}
174
175async function tick($: EngineInterface, cfg: Config) {
176  const s = await read($, snap)
177  if (!s.lastResponseAt || s.isWorking || s.pinging) return
178  const ttl = ttlOf(s, cfg)
179  const remaining = s.lastResponseAt + ttl - (await $.clock.now())
180  if (remaining <= 0) {
181    // The state write redraws the band once (its render reads snap), whatever
182    // the clock did meanwhile; then the band rests until the next request.
183    if (!s.cold) await update($, snap, x => ({ ...x, cold: true }))
184    return
185  }
186  $.ui.invalidate('ui.render')
187  if (!s.warned && remaining <= warnLeadOf(ttl)) {
188    await update($, snap, x => ({ ...x, warned: true }))
189    const text = `Prompt cache goes cold in ${fmtClock(remaining)}. Send a message to keep it warm.`
190    $.ui.toast(text)
191    void $.ui.notify(text, { title: 'Pilot Light' }).catch(() => undefined)
192  }
193  if (cfg.keepAlive && !s.pingFailed && remaining <= PING_LEAD_MS) await ping($, s, ttl)
194}
195
196async function handoff($: EngineInterface) {
197  $.ui.toast('Writing handoff brief...')
198  const reply = await $.model.fork({ prompt: HANDOFF_PROMPT })
199  if (!reply.isAnswered) {
200    $.ui.toast(`Handoff skipped: ${reply.reason}`)
201    return
202  }
203  const brief = reply.text.trim()
204  const draft = (await $.prompt.read()).text.trim()
205  const home = (await $.env.get('HOME')) ?? (await $.session.cwd())
206  const stamp = new Date(await $.clock.now()).toISOString().replace(/[:.]/g, '-')
207  const path = `${home}/.claude/pilot-light/handoff-${stamp}.md`
208  await $.fs.write(path, `${brief}\n`)
209  await $.command.run({ command: 'clear' })
210  const text =
211    `Continue from this handoff brief (saved at ${path}):\n\n${brief}` +
212    (draft ? `\n\n## Pending draft\n\n${draft}` : '')
213  const filled = await $.prompt.fill({ text })
214  $.ui.toast(filled.isFilled ? `Handoff brief pasted into the prompt; saved at ${path}` : `Prompt box refused the paste; brief saved at ${path}`)
215}
216
217export const register: Register = (on, options) => {
218  const cfg: Config = {
219    configuredTtl: (Number(options.cacheWindowMinutes) || 60) * MINUTE,
220    keepAlive: options.keepAlive !== false,
221  }
222
223  on('session.start', async ($, e, next) => {
224    await measure($)
225    // A reload mid-fork drops the continuation that would have cleared the flag.
226    await update($, snap, s => ({ ...s, pinging: false }))
227    if (e.isInteractive) {
228      $.clock.every(1000, () => void tick($, cfg).catch(error => $.ui.log(`pilot-light: tick failed: ${String(error)}`)))
229    }
230    return next(e)
231  })
232
233  on('session.measure', async ($, e, next) => {
234    await update($, snap, s => ({
235      ...s,
236      contextTokens: e.context.tokens ?? s.contextTokens,
237      window: e.context.window,
238      rateLimits: e.rateLimits.map(r => ({ kind: r.kind, percentUsed: r.percentUsed })),
239      costUsd: e.cost?.usd ?? s.costUsd,
240    }))
241    return next(e)
242  })
243
244  on('turn.start', async ($, e, next) => {
245    if (e.agentId === undefined) await update($, snap, s => ({ ...s, isWorking: true, lastRequestAt: 0 }))
246    return next(e)
247  })
248
249  // The cache TTL runs from the start of the request that read it, not from
250  // the end of the response, so the countdown anchors on the last request.
251  on('turn.step', async function* ($, e, next) {
252    if (e.agentId === undefined) {
253      const at = await $.clock.now()
254      await update($, snap, s => ({ ...s, lastRequestAt: at }))
255    }
256    return yield* next(e)
257  })
258
259  on('turn.complete', async ($, e, next) => {
260    if (e.agentId !== undefined) return next(e)
261    const now = await $.clock.now()
262    const u = e.usage
263    await update($, snap, s => ({
264      ...s,
265      isWorking: false,
266      cold: u ? false : s.cold,
267      lastResponseAt: u ? s.lastRequestAt || now : s.lastResponseAt,
268      model: u?.model ?? s.model,
269    }))
270    return next(e)
271  })
272
273  on('classic.Stop', async ($, e, next) => {
274    if (e.agent_id === undefined) void detectTtl($, e.transcript_path).catch(() => undefined)
275    return next(e)
276  }).catch(($, e, next) => next(e))
277
278  on('prompt.submit', async ($, e, next) => {
279    // Only the person's own Enter re-arms the pings: a /loop, a task
280    // notification or a peer message arriving while they are away must not.
281    if (e.origin.kind === 'composer' || e.origin.kind === 'bridge') {
282      // Coming back to a cache the pings kept warm: the cold re-send that
283      // did not happen, less what the pings cost, is what the stretch saved.
284      const s = await read($, snap)
285      const ttl = ttlOf(s, cfg)
286      const warm = s.lastResponseAt > 0 && s.lastResponseAt + ttl > (await $.clock.now())
287      const saved = s.keepAlives && warm ? Math.max(0, coldUsdOf(s, ttl) - s.pingUsd) : 0
288      const savedUsd = s.savedUsd + saved
289      if (saved) await $.store.set('savedUsd', savedUsd)
290      await update($, snap, x => ({ ...x, ...RESTED, savedUsd }))
291    }
292    return next(e)
293  }).catch(($, e, next) => next(e))
294
295  on('session.end', async ($, e, next) => {
296    if (e.reason === 'clear') {
297      await update($, snap, s => ({ ...s, ...RESTED, lastResponseAt: 0, contextTokens: 0, costUsd: 0 }))
298    }
299    return next(e)
300  })
301
302  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
303    if (e.props.hasSurvey) return next(e)
304    const s = await read($, snap)
305    const { Box, Button, Text } = $.ui.resolve(e)
306
307    const ttl = ttlOf(s, cfg)
308    const remaining = s.lastResponseAt ? s.lastResponseAt + ttl - (await $.clock.now()) : null
309    const ttlTag = s.ttlMs === null ? `${Math.round(ttl / MINUTE)}m?` : ttl >= HOUR ? '1h' : '5m'
310    const state = s.isWorking
311      ? 'working'
312      : remaining === null
313        ? 'no cache yet'
314        : remaining <= 0
315          ? 'cold'
316          : fmtClock(remaining)
317    const cold = coldUsdOf(s, ttl)
318    const pings = s.keepAlives ? ` · 🔁 ${s.keepAlives} ${fmtUsd(s.pingUsd)}/${fmtUsd(cold)}` : ''
319    const failed = s.pingFailed ? ' · ping failed' : ''
320    const tone =
321      remaining === null || s.isWorking || remaining > warnLeadOf(ttl)
322        ? {}
323        : remaining <= 0
324          ? { dimColor: true }
325          : { color: 'yellow' }
326
327    const warm = (s.contextTokens * priceOf(s.model)[2]) / 1e6
328    const pct = s.window ? ` (${Math.round((100 * s.contextTokens) / s.window)}%)` : ''
329    const limit = (kind: string, label: string) => {
330      const r = s.rateLimits.find(r => r.kind === kind)
331      return r ? `${label} ${Math.round(r.percentUsed)}%` : `${label} n/a`
332    }
333
334    return (
335      <Box flexWrap="wrap" gap={2}>
336        <Text {...tone}>{remaining !== null && remaining <= 0 ? '🧊' : '🔥'} {state} /{ttlTag}{pings}{failed}</Text>
337        <Text dimColor>🧠 {fmtK(s.contextTokens)}/{fmtK(s.window)}{pct}</Text>
338        <Text dimColor>📊 {limit('five_hour', '5h')} · {limit('seven_day', '7d')}</Text>
339        <Text dimColor>💰 {fmtUsd(s.costUsd)}</Text>
340        <Text dimColor>🧊 next {fmtUsd(cold)} (warm {fmtUsd(warm)})</Text>
341        <Text dimColor>💸 saved {fmtUsd(s.savedUsd)}</Text>
342        <Button key="handoff" label="📋🧹" onPress={() => handoff($)} />
343      </Box>
344    )
345  })
346}
347

1export type RateLimit = { kind: string; percentUsed: number }
2
3export type Snap = {
4  /** `$.clock.now()` of the request that last read the cache; 0 before the first. */
5  lastResponseAt: number
6  /** Start of the latest main-loop request of the running turn; 0 between turns. */
7  lastRequestAt: number
8  /** Cache lifetime read from API usage; null until detected (config applies). */
9  ttlMs: number | null
10  /** Prompt tokens the next request re-sends. */
11  contextTokens: number
12  window: number
13  rateLimits: RateLimit[]
14  costUsd: number
15  model: string
16  isWorking: boolean
17  /** A keep-alive fork is in flight. */
18  pinging: boolean
19  /** The band has been drawn 'cold' for the current anchor. */
20  cold: boolean
21  /** Keep-alive pings sent since the person last prompted. */
22  keepAlives: number
23  /** What those pings cost, and what the latest one cost. */
24  pingUsd: number
25  lastPingUsd: number
26  /** A ping failed this stretch; no more are tried until the person prompts. */
27  pingFailed: boolean
28  /** A cold-cache notification went out since the person last prompted. */
29  warned: boolean
30  /** Lifetime: cold re-sends the pings avoided, less their cost; mirrored from $.store. */
31  savedUsd: number
32}
33
34declare module 'claude-code' {
35  interface PluginState {
36    'pilot-light': { snap: Shaped<Snap> }
37  }
38}
39