context-meter

1import { atom, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { Meter } from '../types'
5
6const INITIAL: Meter = {
7  contextTokens: 0,
8  contextWindow: 200000,
9  expect: null,
10  partial: 0,
11  rewritten: 0,
12  full: 0,
13  level: 0,
14}
15
16const meter = atom({ plugin: 'context-meter', key: 'meter' } as const, INITIAL)
17
18// Absolute context size, not a percentage of the window.
19const THRESHOLDS = [100000, 150000, 200000]
20const LEVEL_HINT = ['', 'Keep an eye on it.', 'Consider compacting soon.', 'Time to compact or start fresh.']
21
22export function contextLevel(tokens: number): number {
23  return THRESHOLDS.filter(t => tokens >= t).length
24}
25
26export function fmtK(n: number): string {
27  if (n >= 999950) return `${(n / 1000000).toFixed(1)}M`
28  return n >= 10000 ? `${Math.round(n / 1000)}k` : `${(n / 1000).toFixed(1)}k`
29}
30
31export type Miss = { kind: 'partial' | 'full'; tokens: number } | null
32
33// A call should read back what the previous call left in the cache: its read
34// should equal the previous read + write. A read of 0 is a full miss, a smaller
35// read a partial miss; the shortfall is the number of tokens cached again.
36export function recordCall(
37  m: Meter,
38  usage: { cache_read_input_tokens: number; cache_creation_input_tokens: number; input_tokens: number },
39): { meter: Meter; miss: Miss } {
40  const cr = usage.cache_read_input_tokens
41  const cw = usage.cache_creation_input_tokens
42  const next = { ...m, contextTokens: usage.input_tokens + cw + cr }
43  let miss: Miss = null
44  if (m.expect !== null) {
45    if (cr === 0 && m.expect > 0) {
46      next.full += 1
47      miss = { kind: 'full', tokens: m.expect }
48    } else if (cr < m.expect) {
49      next.partial += 1
50      next.rewritten += m.expect - cr
51      miss = { kind: 'partial', tokens: m.expect - cr }
52    }
53  }
54  next.expect = cr + cw
55  return { meter: next, miss }
56}
57
58export const register: Register = on => {
59  // Seed the context fill from what the engine already knows, so a resumed session
60  // doesn't re-announce thresholds it is already past. Only crossings announce.
61  on('session.start', async ($, e, next) => {
62    const result = await next(e)
63    const { context } = await $.session.usage()
64    await update($, meter, m => {
65      const contextTokens = context.tokens ?? m.contextTokens
66      return { ...m, contextWindow: context.window, contextTokens, level: contextLevel(contextTokens) }
67    })
68    return result
69  })
70
71  on('turn.step', async function* ($, e, next) {
72    const result = yield* next(e)
73    // Subagent calls have their own cache chain; leave them out.
74    if (e.agentId !== undefined || !result.usage) return result
75
76    const usage = result.usage
77    let miss: Miss = null
78    let crossed = 0
79    await update($, meter, m => {
80      const r = recordCall(m, usage)
81      miss = r.miss
82      const level = contextLevel(r.meter.contextTokens)
83      crossed = level > m.level ? level : 0
84      return { ...r.meter, level }
85    })
86
87    if (crossed > 0) {
88      $.ui.toast(`Context passed ${THRESHOLDS[crossed - 1] / 1000}k. ${LEVEL_HINT[crossed]}`, { timeoutMs: 6000 })
89    }
90    const m = miss as Miss
91    if (m !== null) {
92      $.ui.toast(
93        m.kind === 'full'
94          ? `Cache miss: the whole prompt (${fmtK(m.tokens)}) was cached again.`
95          : `Partial cache miss: ${fmtK(m.tokens)} cached again.`,
96        { timeoutMs: 6000 },
97      )
98    }
99    return result
100  })
101
102  on('session.measure', async ($, e, next) => {
103    const { tokens, window } = e.context
104    await update($, meter, m => ({ ...m, contextWindow: window, contextTokens: tokens ?? m.contextTokens }))
105    return next(e)
106  })
107}
108

1export type Meter = {
2  /** Input tokens of the last main-thread call (uncached + cache write + cache read). */
3  contextTokens: number
4  /** The model's context window size. */
5  contextWindow: number
6  /** What the next call should read from cache: the previous read + write; null before the first call. */
7  expect: number | null
8  /** Calls that read less than expected, but not zero. */
9  partial: number
10  /** Tokens that had to be cached again across the partial misses. */
11  rewritten: number
12  /** Calls that read nothing although something was cached. */
13  full: number
14  /** The highest threshold (0 none, 1 = 100k, 2 = 150k, 3 = 200k) already announced; drops when context shrinks. */
15  level: number
16}
17
18declare module 'claude-code' {
19  interface PluginState {
20    'context-meter': { meter: Meter }
21  }
22}
23