context-guard

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { Gauge } from '../types'
5
6const gauge = atom({ plugin: 'context-guard', key: 'gauge' } as const, null)
7const snooze = atom({ plugin: 'context-guard', key: 'snooze' } as const, 0)
8const samples = atom({ plugin: 'context-guard', key: 'samples' } as const, [])
9const lastTurnTokens = atom({ plugin: 'context-guard', key: 'lastTurnTokens' } as const, null)
10const autoAt = atom({ plugin: 'context-guard', key: 'autoAt' } as const, null)
11const eater = atom({ plugin: 'context-guard', key: 'eater' } as const, null)
12const isPast = atom({ plugin: 'context-guard', key: 'isPast' } as const, false)
13const isQueued = atom({ plugin: 'context-guard', key: 'isQueued' } as const, false)
14const isHandoffPending = atom({ plugin: 'context-guard', key: 'isHandoffPending' } as const, false)
15const handoffPath = atom({ plugin: 'context-guard', key: 'handoffPath' } as const, null)
16
17const GREEN = '#1D9E75'
18const AMBER = '#EF9F27'
19const CORAL = '#D85A30'
20const RED = '#E24B4A'
21const TRACK = '#888780'
22const LOG_KEY = 'log'
23const HANDOFF_RE = /\.claude\/handoffs\/[\w.-]+\.md/
24
25const k = (n: number) => `${Math.round(n / 1000)}k`
26
27type Opts = { line: number; margin: number; step: number; instructions: string; isMacNotify: boolean; isHeadless: boolean }
28
29type Runway = { turns: number | null; target: 'line' | 'auto' }
30
31// `before` is the gauge read ahead of a compact, since the compact itself clears the gauge
32async function log($: any, action: string, detail?: string, before?: Gauge | null) {
33  const g: Gauge | null = before ?? (await read($, gauge))
34  const prior = ((await $.store.get(LOG_KEY)) as unknown[] | undefined) ?? []
35  const row = { at: await $.clock.now(), action, tokens: g?.tokens ?? null, line: g?.line ?? null, ...(detail ? { detail } : {}) }
36  await $.store.set(LOG_KEY, [...prior, row].slice(-300))
37}
38
39async function learnAutoAt($: any) {
40  const usage = await $.session.usage({ breakdown: 'summary' })
41  const b = usage.context.breakdown
42  const at = b && b.isAutoCompactEnabled && b.autoCompactThreshold ? b.autoCompactThreshold : null
43  // 0 records "auto-compact is off" so it is not looked up again on every response
44  await update($, autoAt, () => at ?? 0)
45
46  return at
47}
48
49async function learnEater($: any) {
50  const usage = await $.session.usage({ breakdown: 'summary' })
51  const b = usage.context.breakdown
52  if (!b || b.totalTokens <= 0) return
53  const used = b.categories.filter((c: any) => c.kind === 'used')
54  const top = used.reduce((a: any, c: any) => (c.tokens > a.tokens ? c : a), used[0])
55  if (!top) return
56  await update($, eater, () => `${top.name} ${Math.round((top.tokens / b.totalTokens) * 100)}%`)
57}
58
59async function runway($: any, g: Gauge): Promise<Runway> {
60  const target = g.tokens < g.line ? 'line' : 'auto'
61  const goal = target === 'line' ? g.line : g.autoAt
62  const list: number[] = await read($, samples)
63  if (goal === null || list.length < 2) return { turns: null, target }
64  const burn = list.reduce((a, b) => a + b, 0) / list.length
65  if (burn <= 0) return { turns: null, target }
66
67  return { turns: Math.max(0, Math.floor((goal - g.tokens) / burn)), target }
68}
69
70function statusText(g: Gauge | null) {
71  if (!g) return 'ctx —'
72  if (g.tokens >= g.line) return `▲ ctx ${k(g.tokens)} · past your ${k(g.line)} line`
73  if (g.tokens >= g.warnFrom) return `◆ ctx ${k(g.tokens)} · ${k(g.line - g.tokens)} to your line`
74  const filled = Math.min(8, Math.round((g.tokens / g.line) * 8))
75
76  return `ctx ${'▰'.repeat(filled)}${'▱'.repeat(8 - filled)} ${k(g.tokens)}`
77}
78
79async function refresh($: any, opts: Opts, tokens: number | undefined) {
80  if (tokens === undefined) {
81    await update($, gauge, () => null)
82    $.ui.status(statusText(null))
83    return
84  }
85  const known: number | null = await read($, autoAt)
86  const at = known === null ? await learnAutoAt($) : known || null
87  const base = at === null ? opts.line : Math.min(opts.line, at - opts.margin)
88  const line = base + (await read($, snooze))
89  const g: Gauge = { tokens, line, warnFrom: line - opts.margin, autoAt: at }
90  await update($, gauge, () => g)
91  $.ui.status(statusText(g))
92
93  const wasPast = await read($, isPast)
94  const nowPast = tokens >= line
95  if (nowPast && !wasPast) {
96    const message = `Context ${k(tokens)}, past your ${k(line)} line`
97    $.ui.toast(message, { timeoutMs: 10000 })
98    if (opts.isMacNotify) void notifyMac($, message)
99    await learnEater($)
100    await log($, 'crossed')
101  }
102  if (nowPast !== wasPast) await update($, isPast, () => nowPast)
103}
104
105// A macOS banner reaches you when Claude Code is in the background; elsewhere osascript is absent and this does nothing
106async function notifyMac($: any, text: string) {
107  try {
108    await $.process.run(['osascript', '-e', `display notification ${JSON.stringify(text)} with title "context-guard"`], { timeoutMs: 5000 })
109  } catch {}
110}
111
112async function resetAfterCompact($: any) {
113  await update($, samples, () => [])
114  await update($, lastTurnTokens, () => null)
115  await update($, snooze, () => 0)
116  await update($, eater, () => null)
117  await update($, isPast, () => false)
118  await update($, isQueued, () => false)
119  await update($, autoAt, () => null)
120}
121
122const errorText = (err: unknown) => (err instanceof Error ? err.message : String(err)).slice(0, 200)
123
124// One attempt; true when it compacted or a hook vetoed it, false when the engine refused
125async function tryCompact($: any, opts: Opts): Promise<{ isDone: boolean; error?: string }> {
126  const before: Gauge | null = await read($, gauge)
127  try {
128    const done = await $.session.compact({ instructions: opts.instructions })
129    if (done.skip) $.ui.toast('Compact was vetoed by another hook')
130    await log($, done.skip ? 'compact-vetoed' : 'compact', undefined, before)
131
132    return { isDone: true }
133  } catch (err) {
134    return { isDone: false, error: errorText(err) }
135  }
136}
137
138// SDK sessions (the desktop app, -p) can't compact from a plugin; there it runs as a /compact command,
139// which the engine itself holds until the session is idle
140async function compactByCommand($: any, opts: Opts) {
141  const before: Gauge | null = await read($, gauge)
142  try {
143    await $.command.run({ command: 'compact', args: opts.instructions })
144    await log($, 'compact-command', undefined, before)
145  } catch (err) {
146    $.ui.toast(`Compact refused: ${errorText(err)}`, { timeoutMs: 10000 })
147    await log($, 'compact-error', errorText(err))
148  }
149}
150
151// Pressed mid-turn: queue it for the end of the turn. Pressed while idle: compact now, and show the real reason if refused.
152async function compact($: any, opts: Opts, isWorking: boolean) {
153  if (opts.isHeadless) return compactByCommand($, opts)
154  if (isWorking) {
155    await update($, isQueued, () => true)
156    await log($, 'compact-queued')
157    return
158  }
159  const { isDone, error } = await tryCompact($, opts)
160  if (!isDone && /headless|SDK/i.test(error ?? '')) {
161    opts.isHeadless = true
162    return compactByCommand($, opts)
163  }
164  if (!isDone) {
165    $.ui.toast(`Compact refused: ${error}`, { timeoutMs: 10000 })
166    await log($, 'compact-error', error)
167  }
168}
169
170const RETRY_DELAYS = [1000, 2000, 4000, 8000]
171
172// A queued compact waits past the end of the turn, then retries with backoff until the session accepts it
173async function runQueued($: any, opts: Opts, attempt: number) {
174  if (!(await read($, isQueued))) return
175  const { isDone, error } = await tryCompact($, opts)
176  if (isDone) {
177    await update($, isQueued, () => false)
178    return
179  }
180  if (attempt + 1 < RETRY_DELAYS.length) {
181    await log($, 'compact-retry', error)
182    $.clock.after(RETRY_DELAYS[attempt + 1], () => void runQueued($, opts, attempt + 1))
183    return
184  }
185  await update($, isQueued, () => false)
186  $.ui.toast(`Compact refused: ${error}`, { timeoutMs: 10000 })
187  await log($, 'compact-error', error)
188}
189
190async function handoff($: any) {
191  try {
192    await $.command.run({ command: 'creating-handoffs' })
193    await update($, isHandoffPending, () => true)
194    await log($, 'handoff')
195  } catch {
196    $.ui.toast('Could not start /creating-handoffs')
197  }
198}
199
200export const register: Register = (on, options) => {
201  const opts: Opts = {
202    line: Number(options.lineTokens ?? 350000),
203    margin: Number(options.marginTokens ?? 50000),
204    step: Number(options.snoozeTokens ?? 50000),
205    instructions: String(options.compactInstructions ?? ''),
206    isMacNotify: options.macNotification !== false,
207    isHeadless: false,
208  }
209
210  on('session.start', async ($, e, next) => {
211    opts.isHeadless = !e.isInteractive
212    await $.command.register({ name: 'context-guard-log', description: 'Show the last context-guard actions (for tuning the line)' })
213    const { context } = await $.session.usage()
214    await refresh($, opts, context.tokens)
215    // A compact queued before a reload or resume would otherwise wait for the next turn to end
216    if (await read($, isQueued)) $.clock.after(RETRY_DELAYS[0], () => void runQueued($, opts, 0))
217
218    return next(e)
219  })
220
221  on('command.run', { command: 'context-guard-log' }, async $ => {
222    const rows = ((await $.store.get(LOG_KEY)) as any[] | undefined) ?? []
223    if (rows.length === 0) return { text: 'No context-guard actions logged yet.' }
224    const lines = rows.slice(-20).map(r => `${new Date(r.at).toISOString().slice(0, 16)}  ${r.action.padEnd(15)} ${r.tokens === null ? '—' : k(r.tokens)} (line ${r.line === null ? '—' : k(r.line)})${r.detail ? `  ${r.detail}` : ''}`)
225
226    return { text: lines.join('\n') }
227  })
228
229  on('session.measure', async ($, e, next) => {
230    const result = await next(e)
231    await refresh($, opts, e.context.tokens)
232
233    return result
234  })
235
236  on('turn.complete', async ($, e, next) => {
237    const result = await next(e)
238    const { context } = await $.session.usage()
239    const now = context.tokens
240    const prev = await read($, lastTurnTokens)
241    if (now !== undefined && prev !== null && now > prev) {
242      await update($, samples, list => [...list, now - prev].slice(-5))
243    }
244    await update($, lastTurnTokens, () => now ?? null)
245    if (await read($, isQueued)) {
246      $.clock.after(RETRY_DELAYS[0], () => void runQueued($, opts, 0))
247    }
248
249    return result
250  })
251
252  on('session.compact', async ($, e, next) => {
253    const result = await next(e)
254    if (!result.skip && e.trigger !== 'precompute' && !e.agentId) {
255      await resetAfterCompact($)
256      await update($, gauge, () => null)
257      $.ui.status(statusText(null))
258    }
259
260    return result
261  })
262
263  on('tool.call', async ($, e, next) => {
264    const result = await next(e)
265    if (await read($, isHandoffPending)) {
266      let seen = ''
267      try { seen = JSON.stringify(e) + JSON.stringify(result) } catch {}
268      const hit = seen.match(HANDOFF_RE)
269      if (hit) {
270        await update($, handoffPath, () => hit[0])
271        await update($, isHandoffPending, () => false)
272      }
273    }
274
275    return result
276  })
277
278  on('session.end', async ($, e, next) => {
279    if (e.reason === 'clear') {
280      await resetAfterCompact($)
281      await update($, handoffPath, () => null)
282      await update($, isHandoffPending, () => false)
283      await update($, gauge, () => null)
284      $.ui.status(statusText(null))
285    }
286
287    return next(e)
288  })
289
290  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
291    if (e.props.hasSurvey) return next(e)
292    const g: Gauge | null = await read($, gauge)
293    const saved: string | null = await read($, handoffPath)
294    const queued = await read($, isQueued)
295    const { Box, Text, Button, Svg } = $.ui.resolve(e) as any
296    const isDesktop = e.surface !== 'terminal'
297    // Draw beside other plugins' content in this slot instead of replacing it
298    const stack = async (mine: unknown) => (
299      <Box flexDirection="column">
300        {mine}
301        {await next(e)}
302      </Box>
303    )
304
305    if (saved && e.props.isWorking) {
306      return stack(<Text dimColor>Writing handoff: {saved}</Text>)
307    }
308
309    if (saved) {
310      return stack(
311        <Box flexDirection="row" gap={1} flexWrap="wrap">
312          <Text color={GREEN}>✓ Handoff saved: {saved}</Text>
313          <Button key="clear" label="Clear" hotkey="x" onPress={async () => { await log($, 'clear'); await $.command.run({ command: 'clear' }) }} />
314          <Button key="dismiss" label="Dismiss" role="dismiss" onPress={() => update($, handoffPath, () => null)} />
315        </Box>
316      )
317    }
318
319    if (queued) {
320      return stack(
321        <Box flexDirection="row" gap={1}>
322          <Text dimColor>{e.props.isWorking ? 'Compact queued. Runs when this turn ends.' : 'Compacting once the session is free…'}</Text>
323          <Button key="cancel" label="Cancel" hotkey="x" onPress={async () => { await update($, isQueued, () => false); await log($, 'compact-cancel') }} />
324        </Box>
325      )
326    }
327
328    if (!g || g.tokens < g.warnFrom) return next(e)
329
330    const past = g.tokens >= g.line
331    const run = await runway($, g)
332    const width = Math.max(12, Math.min(40, e.props.bodyColumns - (past ? 62 : 40)))
333    const bar = isDesktop ? <Svg source={svgGauge(g)} alt={`Context ${k(g.tokens)} of ${k(g.line)} line`} height={18} /> : textGauge(Text, g, width)
334    // A far-off estimate (e.g. ~119 turns) is noise; only show it when it is close
335    const turns = run.turns === null || run.turns > 30 ? null : `~${run.turns} turn${run.turns === 1 ? '' : 's'}`
336
337    if (!past) {
338      return stack(
339        <Box flexDirection="row" alignItems="center" gap={1}>
340          <Box flexShrink={0}><Text color={AMBER} bold wrap="truncate">◆ {k(g.tokens)}</Text></Box>
341          <Box flexGrow={1} flexShrink={1} minWidth={0}>{bar}</Box>
342          <Box flexShrink={0}><Text dimColor wrap="truncate">{k(g.line - g.tokens)} to your {k(g.line)} line{turns ? ` · ${turns}` : ''}</Text></Box>
343        </Box>
344      )
345    }
346
347    const who: string | null = await read($, eater)
348    const detail = [
349      `${k(g.tokens - g.line)} past your ${k(g.line)} line`,
350      turns ? `${turns} to auto-compact` : null,
351      g.autoAt === null ? null : `auto-compact at ${k(g.autoAt)}`,
352      who ? `biggest: ${who}` : null,
353    ].filter(Boolean).join(' · ')
354
355    return stack(
356      <Box flexDirection="column">
357        <Box flexDirection="row" alignItems="center" gap={1}>
358          <Box flexShrink={0}><Text color={CORAL} bold wrap="truncate">▲ {k(g.tokens)}</Text></Box>
359          <Box flexGrow={1} flexShrink={1} minWidth={0}>{bar}</Box>
360          <Box flexDirection="row" gap={1} flexShrink={0}>
361            <Button key="compact" label="Compact" hotkey="c" onPress={() => compact($, opts, e.props.isWorking)} />
362            <Button key="handoff" label="Hand off" hotkey="h" onPress={() => handoff($)} />
363            <Button key="snooze" label={`Remind at ${k(g.line + opts.step)}`} hotkey="s" onPress={async () => {
364              await update($, snooze, n => n + opts.step)
365              await log($, 'snooze')
366              await refresh($, opts, g.tokens)
367            }} />
368          </Box>
369        </Box>
370        <Text dimColor wrap="truncate">{detail}</Text>
371      </Box>,
372    )
373  })
374}
375
376const scaleOf = (g: Gauge) => g.autoAt ?? Math.max(g.line * 1.15, g.tokens * 1.05)
377
378const colorAt = (g: Gauge, t: number) => (t < g.warnFrom ? GREEN : t < g.line ? AMBER : CORAL)
379
380const textGauge = (Text: any, g: Gauge, width: number) => {
381  const max = scaleOf(g)
382  const cells: { ch: string; color: string }[] = []
383  const lineAt = Math.min(width - 1, Math.round((g.line / max) * width))
384  const fillTo = Math.min(width, Math.round((g.tokens / max) * width))
385  for (let i = 0; i < width; i++) {
386    const t = ((i + 0.5) / width) * max
387    if (i === lineAt) cells.push({ ch: '┃', color: '' })
388    else if (i === fillTo - 1 && g.tokens >= g.line) cells.push({ ch: '●', color: CORAL })
389    else if (i < fillTo) cells.push({ ch: '━', color: colorAt(g, t) })
390    else cells.push({ ch: '━', color: TRACK })
391  }
392  if (g.autoAt !== null) cells.push({ ch: '▌', color: RED })
393  const runs: { text: string; color: string }[] = []
394  for (const c of cells) {
395    const last = runs[runs.length - 1]
396    if (last && last.color === c.color) last.text += c.ch
397    else runs.push({ text: c.ch, color: c.color })
398  }
399
400  return (
401    <Text>
402      {runs.map(r => (r.color === TRACK ? <Text dimColor>{r.text}</Text> : r.color ? <Text color={r.color}>{r.text}</Text> : <Text>{r.text}</Text>))}
403    </Text>
404  )
405}
406
407const MARK = '#F1EFE8'
408const MARK_EDGE = '#2C2C2A'
409
410const svgGauge = (g: Gauge) => {
411  const W = 1000
412  const max = scaleOf(g)
413  const x = (t: number) => Math.min(W, Math.max(0, (t / max) * W))
414  const seg = (from: number, to: number, color: string) => (to > from ? `<rect x="${x(from)}" y="9" width="${x(to) - x(from)}" height="6" fill="${color}"/>` : '')
415  const edge = `stroke="${MARK_EDGE}" stroke-width="1" vector-effect="non-scaling-stroke"`
416  const now = g.tokens
417  const bar = [
418    `<rect x="0" y="9" width="${W}" height="6" fill="${TRACK}" fill-opacity="0.22"/>`,
419    seg(0, Math.min(now, g.warnFrom), GREEN),
420    seg(g.warnFrom, Math.min(now, g.line), AMBER),
421    seg(g.line, now, CORAL),
422  ].join('')
423  // Your line: a caret above the bar and a notch through it, light with a dark edge so it reads on any fill or theme
424  const lx = Math.min(W - 11, Math.max(11, x(g.line)))
425  const line = `<polygon points="${lx - 11},0 ${lx + 11},0 ${lx},7" fill="${MARK}" ${edge}/>`
426    + `<rect x="${lx - 1.5}" y="5" width="3" height="13" fill="${MARK}" ${edge}/>`
427  const auto = g.autoAt === null ? '' : `<rect x="${Math.min(W - 4, x(g.autoAt) - 2)}" y="7" width="4" height="10" rx="1" fill="${RED}"/>`
428
429  // width="4000" asks for more than any slot, so the surface caps it to the full row
430  return `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 ${W} 18" preserveAspectRatio="none" width="4000" height="18">`
431    + `<defs><clipPath id="round"><rect x="0" y="9" width="${W}" height="6" rx="3"/></clipPath></defs>`
432    + `<g clip-path="url(#round)">${bar}</g>${auto}${line}</svg>`
433}
434

1export type Gauge = {
2  tokens: number
3  line: number
4  warnFrom: number
5  autoAt: number | null
6}
7
8declare module 'claude-code' {
9  interface PluginState {
10    'context-guard': {
11      gauge: Gauge | null
12      snooze: number
13      samples: number[]
14      lastTurnTokens: number | null
15      autoAt: number | null
16      eater: string | null
17      isPast: boolean
18      isQueued: boolean
19      isHandoffPending: boolean
20      handoffPath: string | null
21    }
22  }
23}
24