Cache Warm

1import type { EngineInterface, ModelForkResult, ModelUsage, Register, Timer, TurnStepInput } from 'claude-code'
2
3type Ttl = '5m' | '1h'
4type CacheMode = Ttl | 'off' | 'unknown'
5type Evidence = { mode: Ttl | 'off'; model: string; messageId: string; timestamp: number }
6type Detection = { mode: CacheMode; reason: string; source?: Evidence; detail?: string }
7type MainRequest = {
8  generation: number
9  startedAt: number
10  model: string
11  completed: boolean
12}
13type PingOutcome = {
14  sessionId: string
15  model: string
16  mode: Ttl
17  startedAt: number
18  endedAt: number | null
19  usage: ModelUsage | null
20  result: 'pending' | 'hit' | 'miss' | 'failed' | 'stale'
21  response: string | null
22}
23
24const TTL_MS: Record<Ttl, number> = { '5m': 5 * 60_000, '1h': 60 * 60_000 }
25// Each idle period permits at most this many successful refreshes.
26const MAX_PINGS = 30
27const PING_AT = 0.9
28const MIN_HIT_RATIO = 0.9
29const TRANSCRIPT_TAIL_BYTES = 1_000_000
30const TRANSCRIPT_WRITE_WAIT_MS = 5_000
31const TRANSCRIPT_CHECKS = 3
32const RECENT_PINGS = MAX_PINGS
33const PING_PROMPT = 'Prompt cache keep-alive ping. Do not use tools. Reply with exactly: ok'
34
35/** Checks a decoded JSON object before reading its fields. */
36function isRecord(value: unknown): value is Record<string, unknown> {
37  return typeof value === 'object' && value !== null && !Array.isArray(value)
38}
39
40/** Requires explicit, nonnegative token totals; missing counters are unknown. */
41function validUsage(value: unknown): value is ModelUsage {
42  return isRecord(value) && ['input_tokens', 'output_tokens', 'cache_read_input_tokens', 'cache_creation_input_tokens']
43    .every(key => Number.isSafeInteger(value[key]) && (value[key] as number) >= 0)
44}
45
46/** Reads TTL writes, an explicit cache-free response, or a read needing evidence. */
47function usageMode(value: unknown): CacheMode | 'read' {
48  if (!isRecord(value)) return 'unknown'
49  const split = value.cache_creation
50  if (!validUsage(value)) return 'unknown'
51  let write5m = 0
52  let write1h = 0
53  if (split !== undefined && split !== null) {
54    if (!isRecord(split)) return 'unknown'
55    const five = split.ephemeral_5m_input_tokens
56    const hour = split.ephemeral_1h_input_tokens
57    if (!Number.isSafeInteger(five) || (five as number) < 0 ||
58        !Number.isSafeInteger(hour) || (hour as number) < 0) return 'unknown'
59    write5m = five as number
60    write1h = hour as number
61  } else if (value.cache_creation_input_tokens !== 0) {
62    return 'unknown'
63  }
64  if (write5m + write1h !== value.cache_creation_input_tokens) return 'unknown'
65  // With mixed TTLs, the shortest prefix lifetime determines the schedule.
66  if (write5m > 0) return '5m'
67  if (write1h > 0) return '1h'
68  if (value.cache_read_input_tokens > 0) return 'read'
69  return value.input_tokens > 0 ? 'off' : 'unknown'
70}
71
72/** Selects current main-session usage and respects intervening cache/model boundaries. */
73function cacheModeFromTail(
74  tail: string,
75  cut: boolean,
76  sessionId: string,
77  request: MainRequest,
78  epochAt: number,
79  prior: Evidence | undefined,
80  now: number,
81): Detection {
82  const lines = tail.split('\n')
83  if (cut) lines.shift()
84  // A row without its terminating newline may still be in flight.
85  if (!tail.endsWith('\n')) lines.pop()
86  const messages = new Map<string, { model: string; messageId: string; timestamp: number; order: number; usage: unknown }>()
87  for (const [order, line] of lines.entries()) {
88    if (line.trim() === '') continue
89    let row: unknown
90    try {
91      row = JSON.parse(line)
92    } catch (error) {
93      return { mode: 'unknown', reason: 'transcript-parse-error', detail: `row ${order + 1}: ${String(error)}` }
94    }
95    if (!isRecord(row) || row.type !== 'assistant' || row.sessionId !== sessionId ||
96        row.isSidechain === true || row.agentId != null || row.agent_id != null ||
97        row.isApiErrorMessage === true || row.error != null) continue
98    const message = row.message
99    if (!isRecord(message) || message.role !== 'assistant') continue
100    if (typeof message.model === 'string' && message.model.startsWith('<')) continue
101    const timestamp = typeof row.timestamp === 'string' ? Date.parse(row.timestamp) : NaN
102    if (!Number.isFinite(timestamp) || timestamp > now || typeof message.id !== 'string' ||
103        message.id === '' || typeof message.model !== 'string' || message.model === '') {
104      return { mode: 'unknown', reason: 'invalid-assistant-metadata' }
105    }
106    if (timestamp < epochAt) continue
107    // Streaming/content-block rows share a message id: the last row wins.
108    messages.set(message.id, { model: message.model, messageId: message.id, timestamp, order, usage: message.usage })
109  }
110  const rows = [...messages.values()].sort((a, b) => b.timestamp - a.timestamp || b.order - a.order)
111  const latest = rows[0]
112  if (latest === undefined || latest.timestamp < request.startedAt) return { mode: 'unknown', reason: 'no-current-usage' }
113  if (latest.model !== request.model) return { mode: 'unknown', reason: 'usage-model-mismatch' }
114  const mode = usageMode(latest.usage)
115  if (mode === 'unknown') return { mode, reason: 'invalid-usage' }
116  if (mode !== 'read') {
117    return {
118      mode, reason: mode === 'off' ? 'cache-off' : 'ttl-from-write',
119      source: { mode, model: latest.model, messageId: latest.messageId, timestamp: latest.timestamp },
120    }
121  }
122  const inherited = prior !== undefined && prior.mode !== 'off' && prior.model === request.model &&
123    prior.timestamp >= epochAt && prior.timestamp <= latest.timestamp ? prior : undefined
124  for (const row of rows.slice(1)) {
125    if (inherited !== undefined && row.timestamp < inherited.timestamp) break
126    if (row.model !== request.model) return { mode: 'unknown', reason: 'ttl-boundary' }
127    const older = usageMode(row.usage)
128    if (older === 'unknown') return { mode: 'unknown', reason: 'invalid-usage' }
129    if (older === 'off') return { mode: 'unknown', reason: 'ttl-boundary' }
130    if (older !== 'read') {
131      return {
132        mode: older, reason: 'ttl-from-prior-write',
133        source: { mode: older, model: row.model, messageId: row.messageId, timestamp: row.timestamp },
134      }
135    }
136  }
137  return inherited === undefined ? { mode: 'unknown', reason: 'ttl-unavailable' } :
138    { mode: inherited.mode, reason: 'ttl-from-prior-write', source: inherited }
139}
140
141/** Limits state filenames to one safe path component. */
142function safeName(value: string | undefined): value is string {
143  return value !== undefined && value !== '.' && value !== '..' && /^[A-Za-z0-9_.-]{1,128}$/.test(value)
144}
145
146/** Accepts local absolute main-transcript paths passed by the engine. */
147function safeTranscript(path: string): boolean {
148  return path.startsWith('/') && !path.startsWith('//') && path.endsWith('.jsonl') &&
149    !/[\x00-\x1f\x7f]/.test(path) && !path.split('/').includes('..') && !path.includes('/subagents/')
150}
151
152// One module environment belongs to one process; session changes reset this state.
153let maxIdleMs = 0
154const state = {
155  sessionId: null as string | null,
156  transcriptPath: null as string | null,
157  model: null as string | null,
158  effort: null as TurnStepInput['effort'] | null,
159  mode: 'unknown' as CacheMode,
160  status: 'paused' as 'paused' | 'warming' | 'pinging' | 'stopped',
161  reason: 'awaiting-main-request',
162  detail: null as string | null,
163  requestAt: null as number | null,
164  refreshedAt: null as number | null,
165  nextPingAt: null as number | null,
166  ttlSource: null as Evidence | null,
167  detectedAt: null as number | null,
168  pings: 0,
169  recentPings: [] as PingOutcome[],
170}
171let generation = 0
172let readSequence = 0
173let epochAt = 0
174let request: MainRequest | undefined
175let evidence: Evidence | undefined
176let inputModel: string | undefined
177let answeredModel: string | undefined
178let allowDetection = false
179let forkInFlight = false
180let timer: Timer | undefined
181let recheck: Timer | undefined
182let publishes = Promise.resolve()
183
184/** Serializes complete snapshots so an older write cannot finish last. */
185function publish($: EngineInterface): Promise<void> {
186  const sessionId = state.sessionId
187  if (sessionId === null) return publishes
188  const ttl = state.mode === '1h' || state.mode === '5m' ? state.mode : undefined
189  const text = JSON.stringify({
190    ...state, generation,
191    expiresAt: ttl === undefined || state.refreshedAt === null ? null : state.refreshedAt + TTL_MS[ttl],
192  }) + '\n'
193  publishes = publishes.then(async () => {
194    const user = await $.env.get('USER')
195    if (!safeName(user) || !safeName(sessionId)) throw new Error('unsafe USER or session id for state filename')
196    await $.fs.write(`/tmp/claude_cachewarm_${user}_${sessionId}.state`, text)
197  }).catch(error => {
198    $.ui.log(`cache-warm: state file not written: ${String(error)}`, { to: 'debug' })
199  })
200  return publishes
201}
202
203/** Cancels both kinds of timers and invalidates asynchronous work. */
204function invalidate() {
205  generation += 1
206  readSequence += 1
207  timer?.cancel()
208  recheck?.cancel()
209  timer = undefined
210  recheck = undefined
211  state.nextPingAt = null
212}
213
214/** Stops this request until another real main request explicitly starts it. */
215function stop($: EngineInterface, reason: string, clear = false, paused = false, detail: string | null = null): Promise<void> {
216  invalidate()
217  allowDetection = false
218  state.status = paused ? 'paused' : 'stopped'
219  state.reason = reason
220  state.detail = detail
221  if (clear) {
222    request = undefined
223    evidence = undefined
224    inputModel = undefined
225    answeredModel = undefined
226    state.mode = 'unknown'
227    state.refreshedAt = null
228    state.ttlSource = null
229    state.detectedAt = null
230  }
231  $.ui.status(undefined)
232  return publish($)
233}
234
235/** Starts a new session identity without retaining its predecessor's cache. */
236function useSession($: EngineInterface, sessionId: string, now: number) {
237  if (state.sessionId === sessionId) return
238  const first = state.sessionId === null
239  if (!first) void stop($, 'session-changed', true)
240  invalidate()
241  state.sessionId = sessionId
242  state.transcriptPath = null
243  state.model = null
244  state.effort = null
245  state.requestAt = null
246  state.refreshedAt = null
247  state.ttlSource = null
248  state.detectedAt = null
249  state.pings = 0
250  state.mode = 'unknown'
251  state.status = 'paused'
252  state.reason = 'awaiting-main-request'
253  state.detail = null
254  request = undefined
255  evidence = undefined
256  inputModel = undefined
257  answeredModel = undefined
258  allowDetection = false
259  epochAt = first ? 0 : now
260}
261
262/** Makes read failures visible; no configured or guessed TTL is substituted. */
263async function readMode($: EngineInterface, current: MainRequest): Promise<Detection> {
264  const path = state.transcriptPath
265  if (path === null) return { mode: 'unknown', reason: 'no-transcript-path' }
266  let size: number
267  let tail: string
268  try {
269    const stat = await $.fs.stat(path)
270    if (stat.kind !== 'file') return { mode: 'unknown', reason: 'transcript-read-failed', detail: 'not a regular file' }
271    size = stat.size
272    const result = await $.process.run(['tail', '-c', String(TRANSCRIPT_TAIL_BYTES), '--', path], { timeoutMs: 5_000 })
273    if (result.exitCode !== 0 || result.isStdoutTruncated) {
274      return { mode: 'unknown', reason: 'transcript-read-failed', detail: `tail exit ${result.exitCode}: ${result.stderr.slice(0, 500)}` }
275    }
276    tail = result.stdout
277  } catch (error) {
278    return { mode: 'unknown', reason: 'transcript-read-failed', detail: String(error) }
279  }
280  const now = await $.clock.now()
281  // The file may have grown between stat and tail.
282  const cut = size >= TRANSCRIPT_TAIL_BYTES || new TextEncoder().encode(tail).length >= TRANSCRIPT_TAIL_BYTES
283  return cacheModeFromTail(tail, cut, state.sessionId!, current, epochAt, evidence, now)
284}
285
286/** Reports unexpected callback failures without changing a newer request. */
287function runCheck($: EngineInterface, token: number, remaining: number) {
288  void evaluate($, token, remaining).catch(error => {
289    $.ui.log(`cache-warm: check failed: ${String(error)}`, { to: 'debug' })
290    if (token === generation) void stop($, 'check-failed', false, true, String(error))
291  })
292}
293
294/** Bounds transcript-write retries; later relevant events can try again. */
295function checkLater($: EngineInterface, token: number, remaining: number) {
296  recheck?.cancel()
297  recheck = undefined
298  if (remaining === 0) return
299  recheck = $.clock.after(TRANSCRIPT_WRITE_WAIT_MS, () => {
300    if (token !== generation) return
301    recheck = undefined
302    runCheck($, token, remaining - 1)
303  })
304}
305
306/** Plans from the real request start, with expiry, idle, and paid-ping limits. */
307async function plan($: EngineInterface, token: number, now: number): Promise<void> {
308  if (token !== generation || !allowDetection || request === undefined || state.refreshedAt === null ||
309      (state.mode !== '1h' && state.mode !== '5m')) return
310  const ttl = state.mode
311  const deadline = request.startedAt + maxIdleMs
312  if (now >= deadline) return stop($, 'idle-limit')
313  if (state.pings >= MAX_PINGS) return stop($, 'ping-cap')
314  if (now >= state.refreshedAt + TTL_MS[ttl]) return stop($, 'cache-expired', false, true)
315  const due = state.refreshedAt + TTL_MS[ttl] * PING_AT
316  if (now >= due) return ping($, token, ttl, now)
317  timer?.cancel()
318  state.nextPingAt = due < deadline ? due : null
319  state.status = 'warming'
320  state.reason = due < deadline ? 'scheduled' : 'idle-limit-before-next-ping'
321  timer = $.clock.after(Math.min(due, deadline) - now, () => {
322    if (token !== generation) return
323    timer = undefined
324    runCheck($, token, TRANSCRIPT_CHECKS)
325  })
326  await publish($)
327}
328
329/** Adopts only current usage and guards both generations and concurrent reads. */
330async function evaluate($: EngineInterface, token: number, remaining: number): Promise<void> {
331  const current = request
332  if (token !== generation || !allowDetection || current === undefined || !current.completed) return
333  const sequence = ++readSequence
334  recheck?.cancel()
335  recheck = undefined
336  const checkedAt = await $.clock.now()
337  if (token !== generation || sequence !== readSequence) return
338  if (checkedAt >= current.startedAt + maxIdleMs) return stop($, 'idle-limit')
339  const detection = await readMode($, current)
340  const now = await $.clock.now()
341  if (token !== generation || sequence !== readSequence) return
342  state.mode = detection.mode
343  state.reason = detection.reason
344  state.detail = detection.detail ?? null
345  state.ttlSource = detection.source ?? null
346  state.detectedAt = now
347  if (detection.mode === 'unknown') {
348    timer?.cancel()
349    timer = undefined
350    state.nextPingAt = null
351    state.status = 'paused'
352    if (['ttl-boundary', 'usage-model-mismatch', 'invalid-usage', 'transcript-parse-error', 'invalid-assistant-metadata'].includes(detection.reason)) {
353      evidence = undefined
354    }
355    checkLater($, token, remaining)
356    await publish($)
357    return
358  }
359  if (detection.mode === 'off') {
360    evidence = undefined
361    state.refreshedAt = null
362    return stop($, 'cache-off')
363  }
364  evidence = detection.source
365  state.refreshedAt ??= current.startedAt
366  await plan($, token, now)
367}
368
369/** Records each fork, including failures and completions superseded by real work. */
370async function ping($: EngineInterface, token: number, ttl: Ttl, startedAt: number): Promise<void> {
371  if (forkInFlight) return stop($, 'ping-in-flight', false, true)
372  forkInFlight = true
373  const outcome: PingOutcome = {
374    sessionId: state.sessionId!, model: state.model!, mode: ttl, startedAt, endedAt: null,
375    usage: null, result: 'pending', response: null,
376  }
377  state.recentPings.push(outcome)
378  state.recentPings = state.recentPings.slice(-RECENT_PINGS)
379  state.status = 'pinging'
380  state.reason = 'refreshing'
381  state.nextPingAt = null
382  void publish($)
383  let result: ModelForkResult | undefined
384  let failure: string | undefined
385  try {
386    result = await $.model.fork({ prompt: PING_PROMPT })
387  } catch (error) {
388    failure = String(error)
389  } finally {
390    forkInFlight = false
391  }
392  const endedAt = await $.clock.now()
393  outcome.endedAt = endedAt
394  if (result !== undefined) {
395    outcome.response = result.isAnswered ? 'answered' : result.reason
396    if ('usage' in result) outcome.usage = result.usage
397    if (!result.isAnswered && result.reason === 'api-error') failure = `api-error: ${result.error} (${result.status})`
398  } else {
399    outcome.response = failure ?? 'fork failed'
400  }
401  if (token !== generation) {
402    outcome.result = 'stale'
403    await publish($)
404    return
405  }
406  if (result === undefined || (!result.isAnswered && result.reason !== 'empty-reply')) {
407    outcome.result = 'failed'
408    return stop($, outcome.response === 'nothing-to-fork' ? 'fork-unavailable' : 'ping-failed', false, true, failure ?? outcome.response)
409  }
410  if (!validUsage(outcome.usage)) {
411    outcome.result = 'failed'
412    return stop($, 'ping-usage-invalid', false, true)
413  }
414  const usage = outcome.usage
415  const total = usage.input_tokens + usage.cache_read_input_tokens + usage.cache_creation_input_tokens
416  if (usage.cache_read_input_tokens + usage.cache_creation_input_tokens > 0) state.refreshedAt = startedAt
417  if (total === 0 || usage.cache_read_input_tokens / total < MIN_HIT_RATIO) {
418    outcome.result = 'miss'
419    return stop($, 'cache-miss')
420  }
421  outcome.result = 'hit'
422  state.pings += 1
423  $.ui.status(`cache warm (${ttl}): ${state.pings} ping${state.pings === 1 ? '' : 's'}`)
424  // The response may itself have consumed the TTL. Schedule using its end time.
425  await plan($, token, endedAt)
426}
427
428export const register: Register = (on, options) => {
429  const minutes = options.maxIdleMinutes
430  if (typeof minutes !== 'number' || !Number.isFinite(minutes) || minutes < 0 || !Number.isFinite(minutes * 60_000)) {
431    throw new RangeError('cache-warm: maxIdleMinutes must be a finite nonnegative number')
432  }
433  maxIdleMs = minutes * 60_000
434
435  // MessageDisplay fires for each streaming batch; it needs no path RPC or write.
436  const classicEvents = new Set([
437    'classic.SessionStart', 'classic.Stop', 'classic.StopFailure', 'classic.UserPromptSubmit',
438    'classic.PostToolUse', 'classic.PostToolUseFailure', 'classic.PostToolBatch', 'classic.PostCompact',
439    'classic.PostModelSwitch', 'classic.SessionEnd',
440  ])
441  on('classic.*', async ($, e, next) => {
442    if (!classicEvents.has(next.event) || next.is('classic.PreToolUse', e) || e.agent_id !== undefined) return next(e)
443    const sessionId = await $.session.id()
444    if (e.session_id !== sessionId) return next(e)
445    const now = await $.clock.now()
446    let changed = state.sessionId !== sessionId
447    useSession($, sessionId, now)
448    if (next.is('classic.SessionStart', e) && e.source !== 'compact') {
449      // A first, late startup path can complete an already observed request.
450      const lateStartup = e.source === 'startup' && state.transcriptPath === null && request !== undefined
451      if (!lateStartup) {
452        changed = true
453        void stop($, 'awaiting-main-request', true)
454        state.requestAt = null
455        state.pings = 0
456        epochAt = now
457      }
458    }
459    const path = safeTranscript(e.transcript_path) ? e.transcript_path : null
460    if (path !== state.transcriptPath) {
461      changed = true
462      if (state.transcriptPath === null && path !== null) {
463        invalidate()
464        if (request !== undefined) request.generation = generation
465      } else {
466        void stop($, 'transcript-changed', true)
467        epochAt = now
468      }
469      state.transcriptPath = path
470    }
471    if (next.is('classic.PostModelSwitch', e)) {
472      state.model = e.to_model
473      state.effort = null
474      epochAt = now
475      await stop($, 'model-switched', true)
476    } else if (next.is('classic.SessionEnd', e)) {
477      await stop($, 'session-ended', true)
478    } else if (allowDetection && request?.completed && state.mode === 'unknown') {
479      await evaluate($, generation, TRANSCRIPT_CHECKS)
480    } else if (changed) {
481      await publish($)
482    }
483    return next(e)
484  })
485
486  on('turn.step', async function* ($, e, next) {
487    // Forks and embedded/in-process agents have an agentId. A tmux teammate's
488    // own main loop does not. No global ping flag may suppress real main work.
489    if (e.agentId !== undefined) return yield* next(e)
490    const sessionId = await $.session.id()
491    const now = await $.clock.now()
492    useSession($, sessionId, now)
493    if (inputModel !== undefined && (inputModel !== e.model || state.effort !== (e.effort ?? null))) {
494      epochAt = now
495      evidence = undefined
496      answeredModel = undefined
497    }
498    inputModel = e.model
499    invalidate()
500    const current: MainRequest = { generation, startedAt: now, model: answeredModel ?? e.model, completed: false }
501    request = current
502    allowDetection = maxIdleMs > 0
503    state.model = current.model
504    state.effort = e.effort ?? null
505    state.requestAt = now
506    state.refreshedAt = null
507    state.pings = 0
508    state.mode = 'unknown'
509    state.ttlSource = null
510    state.detectedAt = null
511    state.detail = null
512    state.status = allowDetection ? 'paused' : 'stopped'
513    state.reason = !allowDetection ? 'disabled' : state.transcriptPath === null ? 'no-transcript-path' : 'main-request-active'
514    $.ui.status(undefined)
515    await publish($)
516    let responseModel: string | undefined
517    try {
518      const result = yield* next(e)
519      responseModel = result.usage?.model
520      return result
521    } finally {
522      if (request === current && current.generation === generation) {
523        if (responseModel !== undefined && responseModel !== current.model) {
524          invalidate()
525          current.generation = generation
526          current.model = responseModel
527          state.model = responseModel
528          // Aliases and context suffixes may differ from the answered model.
529          // Only an actual response-model change adds another epoch boundary.
530          if (answeredModel !== undefined && answeredModel !== responseModel) {
531            evidence = undefined
532            epochAt = current.startedAt
533          }
534        }
535        if (responseModel !== undefined) answeredModel = responseModel
536        current.completed = true
537        if (allowDetection) {
538          state.reason = state.transcriptPath === null ? 'no-transcript-path' : 'awaiting-transcript'
539          checkLater($, generation, TRANSCRIPT_CHECKS)
540        }
541        await publish($)
542      }
543    }
544  })
545
546  on('session.end', async ($, e, next) => {
547    if (state.sessionId === null) state.sessionId = e.sessionId
548    if (state.sessionId === e.sessionId) await stop($, 'session-ended', true)
549    return next(e)
550  })
551}
552