context-meter

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { MeterView } from '../types'
5import { SETTINGS_FILE, fit, pieces, resolveSettings, startSizeOf, viewOf, type Source } from './meter'
6
7const view = atom({ plugin: 'context-meter', key: 'view' } as const, null as MeterView | null)
8
9const SAFE_SESSION_ID = /^[A-Za-z0-9_-]{1,120}$/
10
11async function refresh($: EngineInterface): Promise<void> {
12  const fileText = async (path: string): Promise<Source> => {
13    if (!(await $.fs.exists(path))) return { kind: 'missing' }
14    try {
15      return { kind: 'text', text: String(await $.fs.read(path)) }
16    } catch {
17      return { kind: 'unreadable' }
18    }
19  }
20
21  const { context } = await $.session.usage()
22  const project = (await $.env.get('CLAUDE_PROJECT_DIR')) || (await $.session.cwd())
23  const home = await $.env.get('HOME')
24  const sources: Source[] = [await fileText(`${project}/.claude/${SETTINGS_FILE}`)]
25  if (home) sources.push(await fileText(`${home}/.claude/${SETTINGS_FILE}`))
26
27  const sessionId = await $.session.id()
28  const state: Source = SAFE_SESSION_ID.test(sessionId)
29    ? await fileText(`${project}/.claude/rotation/state/${sessionId}.enforcement.json`)
30    : { kind: 'missing' }
31  const startSize = startSizeOf(state.kind === 'text' ? state.text : undefined)
32
33  await update($, view, () => viewOf(context, resolveSettings(sources), startSize))
34}
35
36export const register: Register = on => {
37  on('session.start', async ($, e, next) => {
38    await refresh($)
39    return next(e)
40  })
41
42  on('session.measure', async ($, e, next) => {
43    await refresh($)
44    return next(e)
45  })
46
47  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
48    const current = await read($, view)
49    const line = current === null || e.props.hasSurvey ? null : pieces(current)
50    if (line === null) return next(e)
51
52    const { Box, Text } = $.ui.resolve(e)
53
54    return (
55      <Box>
56        {fit(line, e.props.bodyColumns).map(piece => (
57          <Text color={piece.tone} dimColor={piece.isDim}>
58            {piece.text}
59          </Text>
60        ))}
61      </Box>
62    )
63  })
64}
65

1// The settings and the bands of the context-rotation guard, ported from its `lib/context_budget.py`
2// and `lib/band_guard.py`, and the one line the meter draws from them. Pure: no `$` in this file.
3// tests/fixtures/meter_parity.json in the guard's folder holds the cases both sides must agree on.
4import { isNum, parseJson, type Json } from './json'
5import type { MeterView } from '../types'
6
7export const SETTINGS_FILE = 'context-rotation.json'
8export const BAR_CELLS = 8
9
10export type Band = 'below' | 'warn' | 'force' | 'hard'
11export type Tone = 'success' | 'warning' | 'error'
12export type Source = { kind: 'missing' } | { kind: 'text'; text: string } | { kind: 'unreadable' }
13export type Triple = [number, number, number]
14
15export type Settings = {
16  enabled: boolean
17  showMeter: boolean
18  windowTokens: number
19  minWorkTokens: number
20  warnPct: number
21  forcePct: number
22  hardPct: number
23  tokens: Triple | null
24}
25
26const DEFAULTS: Settings = {
27  enabled: true,
28  showMeter: true,
29  windowTokens: 200000,
30  minWorkTokens: 20000,
31  warnPct: 55,
32  forcePct: 65,
33  hardPct: 75,
34  tokens: null,
35}
36
37type Raw = { [key: string]: Json }
38
39function isPlainObject(value: Json): value is Raw {
40  return typeof value === 'object' && value !== null && !Array.isArray(value) && !isNum(value)
41}
42
43const has = (raw: Raw, key: string): boolean => key in raw
44
45function boolOf(raw: Raw, key: string, fallback: boolean): boolean {
46  const value = raw[key]
47  return has(raw, key) && typeof value === 'boolean' ? value : fallback
48}
49
50function positiveIntOf(raw: Raw, key: string, fallback: number): number {
51  const value = raw[key]
52  return has(raw, key) && isNum(value) && value.isInt && value.value > 0 ? value.value : fallback
53}
54
55function numberOf(value: Json | undefined): number | null {
56  return isNum(value) ? value.value : null
57}
58
59function percentages(raw: Raw): Triple | null {
60  if (!['warn_pct', 'force_pct', 'hard_pct'].some(key => has(raw, key))) return null
61  const warn = has(raw, 'warn_pct') ? numberOf(raw.warn_pct) : DEFAULTS.warnPct
62  const force = has(raw, 'force_pct') ? numberOf(raw.force_pct) : DEFAULTS.forcePct
63  const hard = has(raw, 'hard_pct') ? numberOf(raw.hard_pct) : DEFAULTS.hardPct
64  if (warn === null || force === null || hard === null) return null
65  return 0 < warn && warn < force && force < hard && hard <= 100 ? [warn, force, hard] : null
66}
67
68function tokenOf(raw: Raw, key: string): number | null {
69  const value = raw[key]
70  return has(raw, key) && isNum(value) && value.isInt && value.value > 0 ? value.value : null
71}
72
73function absolute(raw: Raw): Triple | null {
74  if (!['warn_tokens', 'force_tokens', 'hard_tokens'].some(key => has(raw, key))) return null
75  const warn = tokenOf(raw, 'warn_tokens')
76  const force = tokenOf(raw, 'force_tokens')
77  const hard = tokenOf(raw, 'hard_tokens')
78  if (warn === null || force === null || hard === null) return null
79  return warn < force && force < hard ? [warn, force, hard] : null
80}
81
82/** `_apply`: every key falls back to its own default, the percentage triple and the token triple as one. */
83export function applySettings(raw: Raw): Settings {
84  const pct = percentages(raw)
85  return {
86    enabled: boolOf(raw, 'enabled', DEFAULTS.enabled),
87    showMeter: boolOf(raw, 'show_meter', DEFAULTS.showMeter),
88    windowTokens: positiveIntOf(raw, 'window_tokens', DEFAULTS.windowTokens),
89    minWorkTokens: positiveIntOf(raw, 'min_work_tokens', DEFAULTS.minWorkTokens),
90    warnPct: pct ? pct[0] : DEFAULTS.warnPct,
91    forcePct: pct ? pct[1] : DEFAULTS.forcePct,
92    hardPct: pct ? pct[2] : DEFAULTS.hardPct,
93    tokens: absolute(raw),
94  }
95}
96
97/** `load_settings`: the first file that exists decides, however bad it is; none means the defaults. */
98export function resolveSettings(sources: Source[]): Settings {
99  for (const source of sources) {
100    if (source.kind === 'missing') continue
101    if (source.kind === 'unreadable') return applySettings(Object.create(null) as Raw)
102    try {
103      const data = parseJson(source.text)
104      return applySettings(isPlainObject(data) ? data : (Object.create(null) as Raw))
105    } catch {
106      return applySettings(Object.create(null) as Raw)
107    }
108  }
109  return { ...DEFAULTS }
110}
111
112/** `thresholds`: the absolute values when all three are valid, else the share of the window. */
113export function thresholds(settings: Settings): Triple {
114  if (settings.tokens) return settings.tokens
115  const window = settings.windowTokens
116  return [(window * settings.warnPct) / 100, (window * settings.forcePct) / 100, (window * settings.hardPct) / 100]
117}
118
119/** `band`: below, warn, force or hard for a measured size. */
120export function bandOf(size: number, [warn, force, hard]: Triple): Band {
121  if (size >= hard) return 'hard'
122  if (size >= force) return 'force'
123  return size >= warn ? 'warn' : 'below'
124}
125
126/** True while an injected session has not grown `min_work_tokens` past its start size. */
127export function isHeld(size: number, startSize: number | null, minWorkTokens: number): boolean {
128  return startSize !== null && size - startSize < minWorkTokens
129}
130
131/**
132 * `band_guard._loop_level`: the level the guard enforces. A session continued from an injected handover
133 * is held below force and hard until it grew `min_work_tokens` past its start size.
134 */
135export function guardLevel(size: number, startSize: number | null, settings: Settings): 'low' | 'force' | 'hard' {
136  const band = bandOf(size, thresholds(settings))
137  if ((band === 'force' || band === 'hard') && isHeld(size, startSize, settings.minWorkTokens)) return 'low'
138  return band === 'force' || band === 'hard' ? band : 'low'
139}
140
141/** What a session's view needs from the settings and the guard's own state, as plain data. */
142export function viewOf(
143  context: { tokens?: number; window: number },
144  settings: Settings,
145  startSize: number | null,
146): MeterView {
147  const [warn, force, hard] = thresholds(settings)
148  return {
149    tokens: typeof context.tokens === 'number' ? context.tokens : null,
150    window: context.window,
151    warn,
152    force,
153    hard,
154    fromPct: settings.tokens === null,
155    windowTokens: settings.windowTokens,
156    enabled: settings.enabled,
157    showMeter: settings.showMeter,
158    minWorkTokens: settings.minWorkTokens,
159    startSize,
160  }
161}
162
163/** The guard's enforcement state `start_size`, a non-negative integer, or null. */
164export function startSizeOf(text: string | undefined): number | null {
165  if (text === undefined) return null
166  try {
167    const state = parseJson(text)
168    if (!isPlainObject(state)) return null
169    const size = state.start_size
170    return isNum(size) && size.isInt && size.value >= 0 ? size.value : null
171  } catch {
172    return null
173  }
174}
175
176/** A token count as the line shows it: 850, 3.2K, 412K, 1.2M. */
177export function tokenText(tokens: number): string {
178  const shown = (value: number, unit: string, digits: number) => String(Number(value.toFixed(digits))) + unit
179  if (tokens < 1000) return String(Math.round(tokens))
180  if (tokens < 10000) return shown(tokens / 1000, 'K', 1)
181  if (tokens < 1000000) return shown(Math.round(tokens / 1000), 'K', 0)
182  return shown(tokens / 1000000, 'M', 1)
183}
184
185export function barOf(size: number, hard: number): string {
186  const filled = Math.max(0, Math.min(BAR_CELLS, Math.floor((size / hard) * BAR_CELLS)))
187  return '▓'.repeat(filled) + '░'.repeat(BAR_CELLS - filled)
188}
189
190export type Piece = { text: string; tone?: Tone; isDim?: boolean }
191
192const ACTION: Record<Band, string> = {
193  below: '',
194  warn: 'wrap up',
195  force: 'write the handover',
196  hard: 'only reads are allowed',
197}
198const TONE: Record<Band, Tone> = { below: 'success', warn: 'warning', force: 'error', hard: 'error' }
199const BAND_WORD: Record<Band, string> = { below: 'ok', warn: 'warn', force: 'force', hard: 'hard' }
200
201function nextThreshold(view: MeterView, band: Band): string | null {
202  if (band === 'hard') return null
203  if (band === 'force') return 'hard at ' + tokenText(view.hard)
204  const afterStart = view.startSize === null ? 0 : view.startSize + view.minWorkTokens
205  if (afterStart > view.force) return 'force from ' + tokenText(afterStart)
206  return band === 'warn' ? 'force at ' + tokenText(view.force) : 'warn at ' + tokenText(view.warn)
207}
208
209function notes(view: MeterView): string[] {
210  const found: string[] = []
211  if (view.fromPct && view.window !== view.windowTokens) {
212    found.push('window ' + tokenText(view.window) + ', bands assume ' + tokenText(view.windowTokens))
213  }
214  if (view.hard > view.window) {
215    found.push('hard ' + tokenText(view.hard) + ' is above the ' + tokenText(view.window) + ' window')
216  }
217  return found
218}
219
220/** What the line says for a view, as pieces in the order they are dropped from the end. Null: nothing to draw. */
221export function pieces(view: MeterView): Piece[] | null {
222  if (!view.showMeter) return null
223  if (view.tokens === null) return [{ text: 'context: no reading yet', isDim: true }]
224  const head = 'context ' + barOf(view.tokens, view.hard) + ' ' + tokenText(view.tokens)
225  if (!view.enabled) return [{ text: head, isDim: true }, { text: ' · rotation off', isDim: true }]
226
227  const raw = bandOf(view.tokens, [view.warn, view.force, view.hard])
228  const held = (raw === 'force' || raw === 'hard') && isHeld(view.tokens, view.startSize, view.minWorkTokens)
229  const band: Band = held ? 'warn' : raw
230  const tone = TONE[band]
231  const out: Piece[] = [{ text: head, tone }, { text: ' · ' + BAND_WORD[band], tone }]
232  if (ACTION[band]) out.push({ text: ': ' + ACTION[band], tone })
233  const next = nextThreshold(view, band)
234  if (next) out.push({ text: ' · ' + next, tone })
235  for (const note of notes(view)) out.push({ text: ' · ' + note, tone: 'warning' })
236  return out
237}
238
239const cells = (list: Piece[]) => list.reduce((sum, piece) => sum + [...piece.text].length, 0)
240
241/** Fits the line to `columns` cells: trailing pieces go first, then the bar, never the band word. */
242export function fit(list: Piece[], columns: number): Piece[] {
243  const [first] = list
244  if (first === undefined) return list
245  const floor = first.text.startsWith('context:') ? 1 : 2
246  let kept = list
247  while (kept.length > floor && cells(kept) > columns) kept = kept.slice(0, -1)
248  if (cells(kept) > columns) {
249    kept = [{ ...first, text: first.text.replace(/ [▓░]+/, '') }, ...kept.slice(1)]
250  }
251  return kept
252}
253

1// A JSON reader that keeps what JSON.parse drops: whether a number was written as an integer.
2// The guard's validation (`context_budget.py`) tells 2000 from 2000.0, so the port has to as well.
3// It accepts what Python's json module accepts, NaN and the infinities included.
4
5export type Num = { readonly isNumber: true; readonly value: number; readonly isInt: boolean }
6export type Json = null | boolean | string | Num | Json[] | { [key: string]: Json }
7
8export class JsonError extends Error {}
9
10const WHITESPACE = ' \t\n\r'
11const NUMBER = /^-?(0|[1-9][0-9]*)(\.[0-9]+)?([eE][+-]?[0-9]+)?/
12const ESCAPES: Record<string, string> = { '"': '"', '\\': '\\', '/': '/', b: '\b', f: '\f', n: '\n', r: '\r', t: '\t' }
13
14export function isNum(value: Json | undefined): value is Num {
15  return typeof value === 'object' && value !== null && !Array.isArray(value) && (value as Num).isNumber === true
16}
17
18export function parseJson(text: string): Json {
19  let at = 0
20
21  const skip = () => {
22    while (at < text.length && WHITESPACE.includes(text.charAt(at))) at += 1
23  }
24
25  const literal = (word: string, value: Json): Json => {
26    if (!text.startsWith(word, at)) throw new JsonError('unexpected token at ' + at)
27    at += word.length
28    return value
29  }
30
31  const num = (value: number, isInt: boolean): Num => ({ isNumber: true, value, isInt })
32
33  const string = (): string => {
34    at += 1
35    let out = ''
36    while (at < text.length) {
37      const ch = text.charAt(at)
38      if (ch === '"') {
39        at += 1
40        return out
41      }
42      if (ch === '\\') {
43        const next = text.charAt(at + 1)
44        if (next === 'u') {
45          const hex = text.slice(at + 2, at + 6)
46          if (!/^[0-9a-fA-F]{4}$/.test(hex)) throw new JsonError('bad escape at ' + at)
47          out += String.fromCharCode(parseInt(hex, 16))
48          at += 6
49        } else if (next in ESCAPES) {
50          out += ESCAPES[next]
51          at += 2
52        } else {
53          throw new JsonError('bad escape at ' + at)
54        }
55      } else if (ch < ' ') {
56        throw new JsonError('control character in a string at ' + at)
57      } else {
58        out += ch
59        at += 1
60      }
61    }
62    throw new JsonError('unterminated string')
63  }
64
65  const value = (): Json => {
66    skip()
67    const ch = text[at]
68    if (ch === undefined) throw new JsonError('unexpected end')
69    if (ch === '{') return object()
70    if (ch === '[') return array()
71    if (ch === '"') return string()
72    if (ch === 't') return literal('true', true)
73    if (ch === 'f') return literal('false', false)
74    if (ch === 'n') return literal('null', null)
75    if (ch === 'N') return literal('NaN', num(NaN, false))
76    if (text.startsWith('Infinity', at)) return literal('Infinity', num(Infinity, false))
77    if (text.startsWith('-Infinity', at)) return literal('-Infinity', num(-Infinity, false))
78    const found = NUMBER.exec(text.slice(at))
79    if (!found) throw new JsonError('unexpected token at ' + at)
80    at += found[0].length
81    return num(Number(found[0]), found[2] === undefined && found[3] === undefined)
82  }
83
84  const array = (): Json => {
85    at += 1
86    const items: Json[] = []
87    skip()
88    if (text[at] === ']') {
89      at += 1
90      return items
91    }
92    for (;;) {
93      items.push(value())
94      skip()
95      if (text[at] === ',') at += 1
96      else if (text[at] === ']') {
97        at += 1
98        return items
99      } else throw new JsonError('expected , or ] at ' + at)
100    }
101  }
102
103  const object = (): Json => {
104    at += 1
105    const members = Object.create(null) as { [key: string]: Json }
106    skip()
107    if (text[at] === '}') {
108      at += 1
109      return members
110    }
111    for (;;) {
112      skip()
113      if (text[at] !== '"') throw new JsonError('expected a key at ' + at)
114      const key = string()
115      skip()
116      if (text[at] !== ':') throw new JsonError('expected : at ' + at)
117      at += 1
118      members[key] = value()
119      skip()
120      if (text[at] === ',') at += 1
121      else if (text[at] === '}') {
122        at += 1
123        return members
124      } else throw new JsonError('expected , or } at ' + at)
125    }
126  }
127
128  const result = value()
129  skip()
130  if (at !== text.length) throw new JsonError('extra data at ' + at)
131  return result
132}
133

1export type MeterView = {
2  /** Tokens of the last response, or null while there is no reading. */
3  tokens: number | null
4  /** The context window of the session's model, in tokens. */
5  window: number
6  /** The warn, force and hard thresholds in tokens, as the guard computes them. */
7  warn: number
8  force: number
9  hard: number
10  /** True when the thresholds come from the percentages of `window_tokens`. */
11  fromPct: boolean
12  windowTokens: number
13  enabled: boolean
14  showMeter: boolean
15  minWorkTokens: number
16  /** The guard's recorded start size of this session, or null. */
17  startSize: number | null
18}
19
20declare module 'claude-code' {
21  interface PluginState {
22    'context-meter': { view: MeterView | null }
23  }
24}
25