context-meter

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Limit, Meter, Resume } from '../types'
5
6const empty: Meter = {
7  tokens: null,
8  window: 0,
9  compactAt: null,
10  growth: [],
11  lastRequestAt: null,
12  cacheRead: 0,
13  cacheWrite: 0,
14  uncached: 0,
15  limits: [],
16}
17
18/** State kept by an older version survives a reload; fill in fields it lacks. */
19export const current = (m: Partial<Meter> | null | undefined): Meter => {
20  const known = Object.fromEntries(Object.entries(m ?? {}).filter(([, v]) => v !== undefined))
21  return { ...empty, ...known }
22}
23
24const meter = atom({ plugin: 'context-meter', key: 'meter' } as const, empty)
25const now = atom({ plugin: 'context-meter', key: 'now' } as const, 0)
26const warned = atom({ plugin: 'context-meter', key: 'warned' } as const, false)
27const resume = atom({ plugin: 'context-meter', key: 'resume' } as const, null as Resume | null)
28/** `<kind>:<threshold>:<reset>` already toasted, so each crossing warns once per window. */
29const limitWarned = atom({ plugin: 'context-meter', key: 'limitWarned' } as const, [] as string[])
30
31const WARN_MS = 5 * 60_000
32/** Below this a cold re-cache is cheap enough not to hold the first message. */
33const GUARD_TOKENS = 50_000
34/** Turns left and hit rate only earn their place below these. */
35const SHOW_TURNS_BELOW = 30
36const SHOW_HIT_BELOW = 0.8
37const LIMIT_THRESHOLDS = [80, 95]
38
39// Claude Code theme keys, so the band follows /theme.
40const DIM = 'inactive'
41const ACCENT = 'claude'
42const WARN = 'warning'
43const ERR = 'error'
44
45const k = (n: number) => (n >= 1_000_000 ? `${(n / 1_000_000).toFixed(1)}M` : `${Math.round(n / 1000)}k`)
46
47const duration = (ms: number) => {
48  const m = Math.ceil(ms / 60_000)
49  return m >= 60 ? `${Math.floor(m / 60)}h${String(m % 60).padStart(2, '0')}m` : `${m}m`
50}
51
52const usd = (n: number | null) => (n === null ? '' : ` (~$${n.toFixed(2)})`)
53
54const severity = (fraction: number) => (fraction >= 0.9 ? ERR : fraction >= 0.75 ? WARN : undefined)
55
56const LIMIT_LABEL: Record<string, string> = { five_hour: '5h', seven_day: 'week', spend_limit: 'spend' }
57
58/** Reset time as the person reads it: a clock time within the day, a weekday further out. */
59export const resetLabel = (iso: string | undefined, t: number): string => {
60  if (!iso) return ''
61  const at = new Date(iso)
62  if (Number.isNaN(at.getTime())) return ''
63  const hm = `${String(at.getHours()).padStart(2, '0')}:${String(at.getMinutes()).padStart(2, '0')}`
64  if (at.getTime() - t < 20 * 3600_000) return hm
65  return `${['Sun', 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat'][at.getDay()]} ${hm}`
66}
67
68/** Turns left before auto-compaction, from the average growth of the last few turns. */
69export const turnsLeft = (m: Meter): number | null => {
70  if (m.tokens === null || m.compactAt === null) return null
71  const recent = m.growth.slice(-5)
72  if (recent.length === 0) return null
73  const avg = recent.reduce((a, b) => a + b, 0) / recent.length
74  if (avg <= 0) return null
75  return Math.max(0, Math.floor((m.compactAt - m.tokens) / avg))
76}
77
78export const hitRate = (m: Meter): number | null => {
79  const total = m.cacheRead + m.cacheWrite + m.uncached
80  return total === 0 ? null : m.cacheRead / total
81}
82
83/** Reads the auto-compact threshold from the locally estimated /context breakdown (no API request). */
84async function refreshCompactAt($: EngineInterface) {
85  try {
86    const { context, rateLimits } = await $.session.usage({ breakdown: 'summary' })
87    const threshold = context.breakdown?.isAutoCompactEnabled ? context.breakdown.autoCompactThreshold : undefined
88    await update($, meter, old => {
89      const m = current(old)
90      return { ...m, window: context.window, compactAt: threshold ?? null, limits: rateLimits.length > 0 ? rateLimits : m.limits }
91    })
92  } catch {
93    // No session bound yet; the next measurement retries.
94  }
95}
96
97/** Toasts once per window when a usage limit crosses 80% and 95%. */
98async function warnLimits($: EngineInterface, limits: readonly Limit[]) {
99  const seen = await read($, limitWarned)
100  const t = await $.clock.now()
101  const fresh: string[] = []
102  for (const limit of limits) {
103    const crossed = LIMIT_THRESHOLDS.filter(x => limit.percentUsed >= x).pop()
104    if (crossed === undefined) continue
105    const id = `${limit.kind}:${crossed}:${limit.resetsAt ?? ''}`
106    if (seen.includes(id)) continue
107    fresh.push(id)
108    const label = LIMIT_LABEL[limit.kind] ?? limit.kind
109    const reset = resetLabel(limit.resetsAt, t)
110    $.ui.toast(`${label} usage at ${Math.round(limit.percentUsed)}%${reset ? `, resets ${reset}` : ''}`)
111  }
112  if (fresh.length > 0) await update($, limitWarned, list => [...list, ...fresh].slice(-20))
113}
114
115/** Compacting reads the cold context once, then later messages stay small. */
116async function compactNow($: EngineInterface) {
117  await update($, resume, () => null)
118  await $.session.compact()
119}
120
121export const register: Register = (on, options) => {
122  const ttlMs = Number(options.cacheTtlMinutes ?? '60') * 60_000
123  const isGuarding = options.resumeGuard !== false
124
125  // A resume reports the old context and how long it sat; show both before the first message.
126  on('classic.SessionStart', async ($, e, next) => {
127    if ((e.source === 'resume' || e.source === 'fork') && e.context_tokens) {
128      const t = await $.clock.now()
129      const idleMs = (e.seconds_since_last_response ?? 0) * 1000
130      const tokens = e.context_tokens
131      await update($, now, () => t)
132      await update($, meter, old => ({ ...current(old), tokens, growth: [], lastRequestAt: t - idleMs }))
133      void refreshCompactAt($)
134      if (e.prompt_cache_likely_expired) {
135        const isHeld = isGuarding && tokens >= GUARD_TOKENS
136        await update($, resume, () => ({ tokens, idleMs, usd: e.estimated_cache_write_usd ?? null, isHeld }))
137      }
138    }
139
140    return next(e)
141  })
142
143  // Hold the first prompt once when it would re-cache a big cold context.
144  on('prompt.submit', async ($, e, next) => {
145    const r = await read($, resume)
146    if (!r?.isHeld || e.origin || e.text.trimStart().startsWith('/')) return next(e)
147    await update($, resume, x => (x ? { ...x, isHeld: false } : x))
148
149    return {
150      drop:
151        `Prompt cache is cold (idle ${duration(r.idleMs)}): this message would re-cache ${k(r.tokens)} tokens${usd(r.usd)}. ` +
152        'Send it again to go ahead, or /compact or start a new session first.',
153    }
154  })
155
156  on('session.start', async ($, e, next) => {
157    await update($, now, () => 0)
158    void refreshCompactAt($)
159    $.clock.every(15_000, async () => {
160      const t = await $.clock.now()
161      await update($, now, () => t)
162      const m = current(await read($, meter))
163      if (m.lastRequestAt === null) return
164      const left = m.lastRequestAt + ttlMs - t
165      if (left > 0 && left <= WARN_MS && !(await read($, warned))) {
166        await update($, warned, () => true)
167        $.ui.toast(`Prompt cache expires in ${duration(left)}: send something to keep it warm`)
168      }
169    })
170
171    return next(e)
172  })
173
174  // Every model request rewrites the cache's expiry and reports cache usage.
175  on('turn.step', async function* ($, e, next) {
176    const result = yield* next(e)
177    const usage = result.usage
178    if (usage) {
179      const t = await $.clock.now()
180      await update($, warned, () => false)
181      await update($, resume, () => null)
182      await update($, now, () => t)
183      await update($, meter, old => ({
184        ...current(old),
185        lastRequestAt: t,
186        cacheRead: usage.cache_read_input_tokens,
187        cacheWrite: usage.cache_creation_input_tokens,
188        uncached: usage.input_tokens,
189      }))
190    }
191    return result
192  })
193
194  on('session.measure', async ($, e, next) => {
195    const tokens = e.context.tokens ?? null
196    await update($, meter, old => {
197      const m = current(old)
198      const delta = tokens !== null && m.tokens !== null ? tokens - m.tokens : 0
199      return {
200        ...m,
201        tokens: tokens ?? m.tokens,
202        window: e.context.window,
203        // A drop means a compaction or /clear: start the growth average over.
204        growth: delta < 0 ? [] : delta > 0 ? [...m.growth, delta].slice(-10) : m.growth,
205        limits: e.rateLimits.length > 0 ? e.rateLimits : m.limits,
206      }
207    })
208    if (e.changed.includes('context')) void refreshCompactAt($)
209    if (e.changed.includes('rateLimits')) await warnLimits($, e.rateLimits)
210
211    return next(e)
212  })
213
214  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
215    const m = current(await read($, meter))
216    if (e.props.hasSurvey || (m.tokens === null && m.limits.length === 0)) return next(e)
217
218    const { Box, Button, Text } = $.ui.resolve(e)
219    const t = (await read($, now)) || (await $.clock.now())
220    const r = await read($, resume)
221
222    // Context: a slim line, accent-filled, warning colours near the compact point.
223    const limit = m.compactAt ?? (m.window || null)
224    const fill = m.tokens === null || limit === null ? null : m.tokens / limit
225    const width = 12
226    const filled = fill === null ? 0 : Math.max(0, Math.min(width, Math.round(fill * width)))
227    const turns = turnsLeft(m)
228
229    // Cache: just the time left; the hit rate only when it is poor.
230    const rate = hitRate(m)
231    const cacheLeft = m.lastRequestAt === null ? null : m.lastRequestAt + ttlMs - t
232
233    const context =
234      m.tokens === null ? null : (
235        <Box key="context">
236          <Text color={DIM}>context </Text>
237          {fill !== null && <Text color={severity(fill) ?? ACCENT}>{'━'.repeat(filled)}</Text>}
238          {fill !== null && <Text color={DIM}>{'─'.repeat(width - filled)} </Text>}
239          {fill !== null && <Text color={severity(fill) ?? DIM}>{Math.round(fill * 100)}%</Text>}
240          <Text color={DIM}>
241            {fill !== null ? ' ' : ''}
242            {k(m.tokens)}
243            {limit !== null ? `/${k(limit)}` : ''}
244          </Text>
245          {turns !== null && turns < SHOW_TURNS_BELOW && (
246            <Text color={turns < 5 ? ERR : WARN}> ~{turns} turns to compact</Text>
247          )}
248        </Box>
249      )
250
251    const cache =
252      cacheLeft === null ? null : (
253        <Box key="cache">
254          <Text color={DIM}>cache </Text>
255          {cacheLeft <= 0 ? (
256            <Text color={ERR}>cold · next message re-reads {k(m.tokens ?? 0)}</Text>
257          ) : (
258            <Text color={cacheLeft <= WARN_MS ? WARN : DIM}>{duration(cacheLeft)}</Text>
259          )}
260          {cacheLeft > 0 && rate !== null && rate < SHOW_HIT_BELOW && (
261            <Text color={WARN}> {Math.round(rate * 100)}% hit</Text>
262          )}
263        </Box>
264      )
265
266    const usage =
267      m.limits.length === 0 ? null : (
268        <Box key="usage">
269          <Text color={DIM}>usage </Text>
270          {m.limits.map((l, i) => {
271            const reset = resetLabel(l.resetsAt, t)
272            return (
273              <Box key={l.kind}>
274                <Text color={DIM}>
275                  {i > 0 ? ' · ' : ''}
276                  {LIMIT_LABEL[l.kind] ?? l.kind}{' '}
277                </Text>
278                <Text color={severity(l.percentUsed / 100) ?? DIM}>{Math.round(l.percentUsed)}%</Text>
279                {reset !== '' && <Text color={DIM}> ↻ {reset}</Text>}
280              </Box>
281            )
282          })}
283        </Box>
284      )
285
286    const parts = [context, cache, usage].filter(x => x !== null)
287    const row = (
288      <Box>
289        {parts.map((part, i) => (
290          <Box key={`part-${i}`}>
291            {i > 0 && <Text color={DIM}> · </Text>}
292            {part}
293          </Box>
294        ))}
295      </Box>
296    )
297
298    if (r === null) return row
299
300    return (
301      <Box flexDirection="column">
302        <Box>
303          <Text color={WARN}>
304            Resumed {k(r.tokens)} after {duration(r.idleMs)} idle · cache expired, next message re-caches it{usd(r.usd)}
305            {r.isHeld ? ' · first send is held once' : ''}{' '}
306          </Text>
307          <Button key="compact" label="Compact" onPress={() => compactNow($)} />
308          <Text> </Text>
309          <Button key="dismiss" label="Dismiss" onPress={() => update($, resume, () => null)} />
310        </Box>
311        {row}
312      </Box>
313    )
314  })
315}
316

1export type Meter = {
2  /** Input tokens of the last response (uncached + cache read + cache write). */
3  tokens: number | null
4  window: number
5  /** Token count where auto-compaction runs; null when off or unknown. */
6  compactAt: number | null
7  /** Context growth per turn, most recent last. */
8  growth: number[]
9  /** When the last model request refreshed the cache, clock ms; null before one. */
10  lastRequestAt: number | null
11  cacheRead: number
12  cacheWrite: number
13  uncached: number
14  /** Account usage windows (five_hour, seven_day, ...), as the last response reported them. */
15  limits: Limit[]
16}
17
18export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
19
20/** A resumed session whose cache had expired: what its first message would re-cache. */
21export type Resume = {
22  tokens: number
23  idleMs: number
24  usd: number | null
25  /** True until the person sends anyway, compacts, or a request goes through. */
26  isHeld: boolean
27}
28
29declare module 'claude-code' {
30  interface PluginState {
31    'context-meter': { meter: Meter; now: number; warned: boolean; resume: Resume | null; limitWarned: string[] }
32  }
33}
34