context-usage

1import { atom, read, update } from 'claude-code'
2import type {
3  Color,
4  Elements,
5  EngineInterface,
6  Register,
7  SessionContextUsage,
8} from 'claude-code'
9
10import type { ContextFill, Reading, UsageWindow } from '../types'
11
12const OFF = 0
13const STEPS = [OFF, 97, 90, 70, 60, 50]
14const stepLabel = (step: number) => (step === OFF ? 'Off' : `${step}%`)
15const OPTIONS = STEPS.map(step => ({ value: String(step), label: stepLabel(step) }))
16const DEFAULT_THRESHOLD = 97
17const THRESHOLD_KEY = 'threshold'
18const READING_KEY = 'reading'
19
20const LABELS: Record<string, string> = {
21  // Padded so the hours line up with a two-digit window such as 24h.
22  five_hour: ' 5h',
23  seven_day: 'Week',
24  spend_limit: 'Spend',
25}
26const ORDER = ['five_hour', 'seven_day', 'spend_limit']
27const ALWAYS = ['five_hour', 'seven_day']
28const MINUTE = 60_000
29const HOUR = 60 * MINUTE
30const DAY = 24 * HOUR
31
32// Shared with cache-meter, context-meter and usage-limits so the gauges of all rows line up.
33const LABEL_WIDTH = 7
34const SEGMENTS = 5
35const PERCENT_WIDTH = 4
36const DETAIL_WIDTH = 10
37
38type Tags = Pick<Elements['terminal'], 'Box' | 'Text'>
39type Layout = { hasGauge: boolean; hasDetails: boolean; isUsageAbove: boolean }
40
41const DESKTOP_BUTTON_ROWS = 2
42
43type Cell = {
44  label: string
45  percent: number | null
46  tone: Color | undefined
47  detail: string
48}
49
50const context = atom({ plugin: 'context-usage', key: 'context' } as const, null)
51const threshold = atom(
52  { plugin: 'context-usage', key: 'threshold' } as const,
53  DEFAULT_THRESHOLD,
54)
55const windows = atom({ plugin: 'context-usage', key: 'windows' } as const, [])
56const now = atom({ plugin: 'context-usage', key: 'now' } as const, 0)
57
58const toFill = (live: SessionContextUsage): ContextFill | null =>
59  live.percent === undefined || live.tokens === undefined
60    ? null
61    : { percent: live.percent, tokens: live.tokens, window: live.window }
62
63const tokens = (count: number) => {
64  const thousands = Math.round(count / 1000)
65
66  if (count < 1000) {
67    return String(count)
68  }
69
70  return thousands < 1000
71    ? `${thousands}k`
72    : `${Number((count / 1_000_000).toFixed(1))}M`
73}
74
75const contextTone = (percent: number, threshold: number): Color | undefined => {
76  const limit = threshold === OFF ? 100 : threshold
77
78  return percent >= limit ? 'error' : percent >= limit - 15 ? 'warning' : undefined
79}
80
81const usageTone = (percent: number): Color | undefined =>
82  percent >= 90 ? 'error' : percent >= 75 ? 'warning' : undefined
83
84const labelOf = (kind: string) => LABELS[kind] ?? kind.replaceAll('_', ' ')
85
86const rankOf = (kind: string) => {
87  const rank = ORDER.indexOf(kind)
88
89  return rank === -1 ? ORDER.length : rank
90}
91
92const span = (duration: number) => {
93  const days = Math.floor(duration / DAY)
94  const hours = Math.floor((duration % DAY) / HOUR)
95  const minutes = Math.max(1, Math.floor((duration % HOUR) / MINUTE))
96
97  if (days > 0) {
98    return `${days}d ${hours}h`
99  }
100
101  return hours > 0 ? `${hours}h ${minutes}m` : `${minutes}m`
102}
103
104const untilReset = (resetsAt: string | undefined, at: number) => {
105  const left = resetsAt === undefined ? Number.NaN : Date.parse(resetsAt) - at
106
107  return Number.isNaN(left) ? '' : `(${span(left)})`
108}
109
110const hasReset = (one: UsageWindow, at: number) =>
111  one.resetsAt !== undefined && Date.parse(one.resetsAt) <= at
112
113const cellsOf = (list: readonly UsageWindow[], at: number): Cell[] =>
114  [...new Set([...ALWAYS, ...list.map(one => one.kind)])]
115    .sort((a, b) => rankOf(a) - rankOf(b))
116    .map(kind => {
117      const one = list.find(held => held.kind === kind)
118
119      if (one === undefined) {
120        return { label: labelOf(kind), percent: null, tone: undefined, detail: '' }
121      }
122
123      // A window past its reset has started over, but the engine reports that only with the next response.
124      if (hasReset(one, at)) {
125        return { label: labelOf(kind), percent: 0, tone: undefined, detail: '' }
126      }
127
128      return {
129        label: labelOf(kind),
130        percent: one.percentUsed,
131        tone: usageTone(one.percentUsed),
132        detail: untilReset(one.resetsAt, at),
133      }
134    })
135
136const meter = ({ Box, Text }: Tags, layout: Layout, one: Cell) => {
137  const filled =
138    one.percent === null
139      ? 0
140      : Math.max(
141          0,
142          Math.min(SEGMENTS, Math.round((one.percent / 100) * SEGMENTS)),
143        )
144
145  return (
146    <Box flexDirection="row" flexShrink={0}>
147      <Box minWidth={LABEL_WIDTH} marginRight={1} flexShrink={0}>
148        <Text dimColor>{one.label}</Text>
149      </Box>
150      {layout.hasGauge && (
151        <Box minWidth={SEGMENTS} marginRight={1} flexShrink={0}>
152          <Text color={one.tone}>{'▮'.repeat(filled)}</Text>
153          <Text dimColor>{'▯'.repeat(SEGMENTS - filled)}</Text>
154        </Box>
155      )}
156      <Box
157        minWidth={PERCENT_WIDTH}
158        marginRight={1}
159        flexShrink={0}
160        justifyContent="flex-end"
161      >
162        <Text color={one.percent === null ? 'inactive' : one.tone}>
163          {one.percent === null ? '–' : `${Math.round(one.percent)}%`}
164        </Text>
165      </Box>
166      {layout.hasDetails && (
167        <Box minWidth={DETAIL_WIDTH} marginRight={1} flexShrink={0}>
168          <Text dimColor>{one.detail}</Text>
169        </Box>
170      )}
171    </Box>
172  )
173}
174
175const runCommand = ($: EngineInterface, command: string) => {
176  $.command
177    .run({ command })
178    .catch(() => $.ui.toast(`context-usage: /${command} could not run`))
179}
180
181const cycleThreshold = async ($: EngineInterface) => {
182  const value = await update(
183    $,
184    threshold,
185    held => STEPS[(STEPS.indexOf(held) + 1) % STEPS.length] ?? DEFAULT_THRESHOLD,
186  )
187  await $.store.set(THRESHOLD_KEY, value)
188}
189
190const setThreshold = async ($: EngineInterface, value: number) => {
191  if (STEPS.includes(value)) {
192    await update($, threshold, () => value)
193    await $.store.set(THRESHOLD_KEY, value)
194  }
195}
196
197const restoreContext = async ($: EngineInterface) => {
198  const stored = await $.store.get(THRESHOLD_KEY)
199
200  if (typeof stored === 'number' && STEPS.includes(stored)) {
201    await update($, threshold, () => stored)
202  }
203
204  // The engine has no reading before the first response of a session, so an empty one counts as nothing sent yet.
205  const usage = await $.session.usage()
206  await update(
207    $,
208    context,
209    () => toFill(usage.context) ?? { percent: 0, tokens: 0, window: usage.context.window },
210  )
211}
212
213const restoreUsage = async ($: EngineInterface) => {
214  const at = await $.clock.now()
215  const stored = (await $.store.get(READING_KEY)) as Reading | undefined
216  await update($, now, () => at)
217
218  if (stored !== undefined) {
219    await update($, windows, () => stored.windows)
220  }
221}
222
223// The next window starts no earlier than this one resets, so it resets at least a whole window later.
224const SAME_WINDOW = HOUR
225
226const newer = (first: UsageWindow, second: UsageWindow) => {
227  const gap = Date.parse(first.resetsAt ?? '') - Date.parse(second.resetsAt ?? '')
228
229  // Without a reset time there is no telling which is newer, so the session's own reading stays.
230  if (Number.isNaN(gap)) {
231    return first
232  }
233
234  if (Math.abs(gap) > SAME_WINDOW) {
235    return gap > 0 ? first : second
236  }
237
238  return second.percentUsed > first.percentUsed ? second : first
239}
240
241// Every session writes the same store, so the readings of all of them meet there.
242const merge = (...lists: (readonly UsageWindow[])[]) => {
243  const byKind = new Map<string, UsageWindow>()
244
245  for (const one of lists.flat()) {
246    const held = byKind.get(one.kind)
247    byKind.set(one.kind, held === undefined ? one : newer(held, one))
248  }
249
250  return [...byKind.values()].sort((a, b) => rankOf(a.kind) - rankOf(b.kind))
251}
252
253const refreshUsage = async ($: EngineInterface, list?: readonly UsageWindow[]) => {
254  // The engine has no reading before the first response of a session.
255  if (list !== undefined && list.length === 0) {
256    return restoreUsage($)
257  }
258
259  const at = await $.clock.now()
260  const stored = (await $.store.get(READING_KEY)) as Reading | undefined
261  await update($, now, () => at)
262
263  // Without a measurement this only takes in what other sessions stored since.
264  if (list === undefined) {
265    await update($, windows, held => merge(held, stored?.windows ?? []))
266
267    return
268  }
269
270  // A measurement can leave a window out, which then keeps its last reading.
271  const merged = merge(list, stored?.windows ?? [])
272  await update($, windows, () => merged)
273  await $.store.set(READING_KEY, { windows: merged })
274}
275
276export const register: Register = (on, options) => {
277  const layout: Layout = {
278    hasGauge: options.display !== 'text',
279    hasDetails: options.details !== false,
280    isUsageAbove: options.usagePosition === 'above',
281  }
282  // Re-armed only once the fill drops below the threshold, so a compaction that lands above it does not loop.
283  let isArmed = true
284
285  on('session.start', async ($, e, next) => {
286    await restoreContext($)
287    const usage = await $.session.usage()
288    await refreshUsage($, usage.rateLimits)
289
290    $.clock.every(MINUTE, () => void refreshUsage($))
291
292    return next(e)
293  })
294
295  // A /clear wipes every plugin's state after session.end and raises no session.start, only this.
296  on('classic.SessionStart', async ($, e, next) => {
297    if (e.source === 'clear') {
298      await restoreContext($).catch(() => undefined)
299      await restoreUsage($).catch(() => undefined)
300    }
301
302    return next(e)
303  })
304
305  on('session.end', async ($, e, next) => {
306    if (e.reason === 'clear') {
307      isArmed = true
308      await update($, context, () => null)
309    }
310
311    return next(e)
312  })
313
314  on('session.measure', async ($, e, next) => {
315    await refreshUsage($, e.rateLimits)
316
317    if (e.changed.includes('context')) {
318      const fill = toFill(e.context)
319      const limit = await read($, threshold)
320      await update($, context, () => fill)
321
322      if (fill !== null && (limit === OFF || fill.percent < limit)) {
323        isArmed = true
324      } else if (fill !== null && isArmed) {
325        isArmed = false
326        // A compaction is refused while the turn that raised this measurement still runs.
327        $.clock.after(250, () => runCommand($, 'compact'))
328      }
329    }
330
331    return next(e)
332  })
333
334  on('session.compact', async ($, e, next) => {
335    const compacted = await next(e)
336    const isLive = e.agentId === undefined && e.trigger !== 'precompute'
337
338    if (isLive && compacted.skip === undefined) {
339      const after = compacted.tokensAfter
340      await update($, context, held =>
341        held === null || after === undefined
342          ? null
343          : {
344              tokens: after,
345              window: held.window,
346              percent: Math.round((after / held.window) * 100),
347            },
348      )
349    }
350
351    return compacted
352  }).catch(($, e, next) => next(e))
353
354  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
355    const below = await next(e)
356
357    if (e.props.hasSurvey) {
358      return below
359    }
360
361    const { Box, Button, Text } = $.ui.resolve(e)
362    const fill = await read($, context)
363    const limit = await read($, threshold)
364    const list = await read($, windows)
365    const at = await read($, now)
366    let picker
367
368    // The mobile surface draws no Select.
369    if (e.surface === 'mobile') {
370      picker = (
371        <Button
372          key="threshold"
373          label={`Auto ${stepLabel(limit)}`}
374          onPress={() => void cycleThreshold($)}
375        />
376      )
377    } else {
378      const { Select } = $.ui.resolve(e)
379      picker = (
380        <Select
381          key="threshold"
382          label="Auto"
383          options={OPTIONS}
384          value={String(limit)}
385          onSelect={value => void setThreshold($, Number(value))}
386        />
387      )
388    }
389
390    const contextRow = (
391      <Box
392        key="context"
393        flexDirection="row"
394        flexWrap="wrap"
395        alignItems="center"
396        // The desktop's native Button is taller than one text row and spills onto the row below.
397        minHeight={e.surface === 'desktop' ? DESKTOP_BUTTON_ROWS : undefined}
398      >
399        {meter({ Box, Text }, layout, {
400          label: 'Context',
401          percent: fill === null ? null : fill.percent,
402          tone: contextTone(fill === null ? 0 : fill.percent, limit),
403          detail: fill === null ? '' : `(${tokens(fill.tokens)}/${tokens(fill.window)})`,
404        })}
405        <Box flexDirection="row" gap={1} flexShrink={0}>
406          {picker}
407          <Button key="compact" label="Compact" onPress={() => runCommand($, 'compact')} />
408          <Button key="clear" label="Clear" onPress={() => runCommand($, 'clear')} />
409        </Box>
410      </Box>
411    )
412    const usageRow = (
413      <Box key="usage" flexDirection="row" flexWrap="wrap">
414        {cellsOf(list, at).map(one => meter({ Box, Text }, layout, one))}
415      </Box>
416    )
417
418    return (
419      <Box flexDirection="column">
420        {layout.isUsageAbove ? [usageRow, contextRow] : [contextRow, usageRow]}
421        {below}
422      </Box>
423    )
424  })
425}
426

1export type ContextFill = { percent: number; tokens: number; window: number }
2
3export type UsageWindow = {
4  kind: string
5  percentUsed: number
6  resetsAt?: string
7}
8
9export type Reading = { windows: UsageWindow[] }
10
11declare module 'claude-code' {
12  interface PluginState {
13    'context-usage': {
14      context: ContextFill | null
15      threshold: number
16      windows: UsageWindow[]
17      now: number
18    }
19  }
20}
21