context-alert

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, ModelUsage, Register } from 'claude-code'
3
4import type { ContextAlertTier } from '../types'
5import { assess, claimable } from './assess'
6
7// Keys `<session>/<loop>`: a /clear starts fresh claims. Each win is saved per session in
8// $.store ⇒ a resumed process stays quiet on tiers it already announced.
9const claims = atom({ plugin: 'context-alert', key: 'claims' } as const, {})
10const KEEP_MS = 7 * 86_400_000
11// Last request per loop. MAIN's input side comes from the engine: an advisor turn's
12// step usage sums every executor pass, the engine's figure = the last pass alone.
13const agentUsed = new Map<string, number>()
14let mainOutput = 0
15// A tool runs while its step still streams: usage lands at the stream's end, so a fast
16// tool waits for it. The next request needs that end anyway ⇒ the wait costs no turn time.
17const streaming = new Map<string, Promise<void>>()
18const STEP_WAIT_MS = 5_000
19// Saves run one at a time, each from the latest host state: a slow save never lands a stale snapshot.
20let saving: Promise<void> = Promise.resolve()
21
22function total(u: ModelUsage): number {
23  return u.input_tokens + u.cache_creation_input_tokens + u.cache_read_input_tokens + u.output_tokens
24}
25
26function tierOf(v: unknown): ContextAlertTier | undefined {
27  return v === 'notice' || v === 'final' ? v : undefined
28}
29
30async function claim($: EngineInterface, loop: string, tier: ContextAlertTier): Promise<boolean> {
31  const session = await $.session.id()
32  const key = `${session}/${loop}`
33  let had = (await read($, claims))[key]
34  if (had === undefined) {
35    const saved = (await $.store.get(`claims:${session}`)) as { claims?: Record<string, unknown> } | undefined
36    had = tierOf(saved?.claims?.[loop])
37  }
38  let won = false
39  await update($, claims, m => {
40    const prior = m[key] ?? had
41    won = claimable(prior, tier)
42    const kept = won ? tier : prior
43    return kept === undefined || m[key] === kept ? m : { ...m, [key]: kept }
44  })
45  if (won) {
46    const save = saving.then(async () => {
47      const prefix = `${session}/`
48      const all = await read($, claims)
49      const mine = Object.fromEntries(Object.entries(all).filter(([k]) => k.startsWith(prefix)).map(([k, v]) => [k.slice(prefix.length), v]))
50      await $.store.set(`claims:${session}`, { at: await $.clock.now(), claims: mine })
51    })
52    saving = save.catch(() => {})
53    await save
54  }
55  return won
56}
57
58async function alert($: EngineInterface, agentId: string | undefined, signal: AbortSignal): Promise<string | undefined> {
59  const step = streaming.get(agentId ?? 'main')
60  if (step && !signal.aborted) {
61    const stop = new AbortController()
62    const quit = () => stop.abort()
63    signal.addEventListener('abort', quit, { once: true })
64    await Promise.race([step, $.clock.sleep(STEP_WAIT_MS, { signal: stop.signal }).catch(() => {})])
65    signal.removeEventListener('abort', quit)
66    stop.abort()
67  }
68  if (signal.aborted) return undefined
69  let used: number | undefined
70  let window: number
71  if (agentId === undefined) {
72    const { context } = await $.session.usage()
73    used = context.tokens === undefined ? undefined : context.tokens + mainOutput
74    window = context.window
75  } else {
76    used = agentUsed.get(agentId)
77    window = Number(await $.env.get('CLAUDE_CODE_MAX_CONTEXT_TOKENS')) || 200_000
78  }
79  if (used === undefined) return undefined
80  const acw = Number(await $.env.get('CLAUDE_CODE_AUTO_COMPACT_WINDOW')) || 0
81  const found = assess(used, window, acw, agentId !== undefined)
82  if (!found) return undefined
83  return (await claim($, agentId ?? 'main', found.tier)) ? found.text : undefined
84}
85
86export const register: Register = on => {
87  on('session.start', async ($, e, next) => {
88    const now = await $.clock.now()
89    for (const key of await $.store.keys()) {
90      if (!key.startsWith('claims:')) continue
91      const saved = (await $.store.get(key)) as { at?: unknown } | undefined
92      if (typeof saved?.at !== 'number' || now - saved.at > KEEP_MS) await $.store.delete(key)
93    }
94    return next(e)
95  })
96
97  on('turn.step', async function* (_$, e, next) {
98    const loop = e.agentId ?? 'main'
99    let done = () => {}
100    streaming.set(loop, new Promise<void>(resolve => (done = resolve)))
101    try {
102      const r = yield* next(e)
103      if (r.usage) {
104        if (e.agentId === undefined) mainOutput = r.usage.output_tokens
105        else agentUsed.set(e.agentId, total(r.usage))
106      }
107      return r
108    } finally {
109      streaming.delete(loop)
110      done()
111    }
112  })
113
114  on('tool.call', async ($, e, next) => {
115    const r = await next(e)
116    if (r.deny !== undefined) return r
117    const text = await alert($, e.agentId, next.signal)
118    return text === undefined ? r : { ...r, context: [...(r.context ?? []), text] }
119  })
120}
121

1import type { ContextAlertTier } from '../types'
2
3// Offsets below the compaction trigger = slowest observed response to the tier + p99 step growth
4// ⇒ that response still fits when the alert fires one step past its threshold; `agent-census` alerts derives them.
5const OFFSETS = {
6  main: { notice: 65_000, final: 27_000 },
7  agent: { notice: 82_000, final: 35_000 },
8} as const
9// CC's compaction trigger for every loop = window − 33K.
10const RESERVE = 33_000
11
12const ADVICE = {
13  main: {
14    notice: 'Checkpoint ahead of compaction: land the current unit on a commit with its checklist current (.agent/spec.md Tasks or .scratch/tasks.md), then continue.',
15    final: 'Compaction is imminent: commit now with the checklist current; after compaction, reorient from the checklist + git log and continue.',
16  },
17  agent: {
18    notice: 'Hand off: bring your deliverable current on disk, then return its path, completion status + next action.',
19    final: 'Compaction is imminent: save progress to your deliverable now, then return its path, completion status + next action.',
20  },
21} as const
22
23export type Assessment = { tier: ContextAlertTier; text: string }
24
25/** `used` = last request's input + cache + output of the loop; `window` = its raw window, clamped to ACW when set. */
26export function assess(used: number, window: number, acw: number, isAgent: boolean): Assessment | undefined {
27  const trigger = (acw > 0 ? Math.min(window, acw) : window) - RESERVE
28  const { notice, final } = OFFSETS[isAgent ? 'agent' : 'main']
29  const tier: ContextAlertTier | undefined = used >= trigger - final ? 'final' : used >= trigger - notice ? 'notice' : undefined
30  if (!tier) return undefined
31  const left = Math.max(0, Math.trunc((trigger - used) / 1000))
32  const label = trigger % 1_000_000 === 0 ? `${trigger / 1_000_000}M` : `${Math.trunc(trigger / 1000)}K`
33  const advice = ADVICE[isAgent ? 'agent' : 'main'][tier]
34  return { tier, text: `Context ${Math.trunc(used / 1000)}K/${label} — ${left}K left. ${advice}` }
35}
36
37/** final covers notice ⇒ regrowth after compaction stays quiet. */
38export function claimable(had: ContextAlertTier | undefined, tier: ContextAlertTier): boolean {
39  return had !== 'final' && had !== tier
40}
41

1export type ContextAlertTier = 'notice' | 'final'
2
3declare module 'claude-code' {
4  interface PluginState {
5    'context-alert': { claims: Record<string, ContextAlertTier> }
6  }
7}
8