cache-meter

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 } from '../types'
11
12const STEPS = [97, 90, 70, 60, 50]
13const OPTIONS = STEPS.map(step => ({ value: String(step), label: `${step}%` }))
14const DEFAULT_THRESHOLD = 97
15const STORE_KEY = 'threshold'
16
17// Shared with usage-limits so the columns of both rows line up.
18const LABEL_WIDTHS = [5, 7]
19const SEGMENTS = 5
20const PERCENT_WIDTH = 4
21const DETAIL_WIDTH = 10
22
23type Tags = Pick<Elements['terminal'], 'Box' | 'Text'>
24type Layout = { hasGauge: boolean; hasDetails: boolean }
25type Cell = {
26  label: string
27  column: number
28  percent: number | null
29  tone: Color
30  detail: string
31}
32
33const cache = atom({ plugin: 'cache-meter', key: 'cache' } as const, null)
34const context = atom({ plugin: 'cache-meter', key: 'context' } as const, null)
35const threshold = atom(
36  { plugin: 'cache-meter', key: 'threshold' } as const,
37  DEFAULT_THRESHOLD,
38)
39
40const toFill = (live: SessionContextUsage): ContextFill | null =>
41  live.percent === undefined || live.tokens === undefined
42    ? null
43    : { percent: live.percent, tokens: live.tokens, window: live.window }
44
45const tokens = (count: number) => {
46  const thousands = Math.round(count / 1000)
47
48  if (count < 1000) {
49    return String(count)
50  }
51
52  return thousands < 1000
53    ? `${thousands}k`
54    : `${Number((count / 1_000_000).toFixed(1))}M`
55}
56
57const cacheTone = (percent: number): Color =>
58  percent < 40 ? 'warning' : 'text'
59
60const contextTone = (percent: number, limit: number): Color =>
61  percent >= limit ? 'error' : percent >= limit - 15 ? 'warning' : 'text'
62
63const meter = ({ Box, Text }: Tags, layout: Layout, one: Cell) => {
64  const filled =
65    one.percent === null
66      ? 0
67      : Math.max(
68          0,
69          Math.min(SEGMENTS, Math.round((one.percent / 100) * SEGMENTS)),
70        )
71
72  return (
73    <Box flexDirection="row" flexShrink={0}>
74      <Box
75        minWidth={LABEL_WIDTHS[one.column] ?? LABEL_WIDTHS[1]}
76        marginRight={1}
77        flexShrink={0}
78      >
79        <Text dimColor>{one.label}</Text>
80      </Box>
81      {layout.hasGauge && (
82        <Box minWidth={SEGMENTS} marginRight={1} flexShrink={0}>
83          <Text color={one.tone}>{'▮'.repeat(filled)}</Text>
84          <Text dimColor>{'▯'.repeat(SEGMENTS - filled)}</Text>
85        </Box>
86      )}
87      <Box
88        minWidth={PERCENT_WIDTH}
89        marginRight={1}
90        flexShrink={0}
91        justifyContent="flex-end"
92      >
93        <Text color={one.percent === null ? 'inactive' : one.tone}>
94          {one.percent === null ? '–' : `${Math.round(one.percent)}%`}
95        </Text>
96      </Box>
97      {layout.hasDetails && (
98        <Box minWidth={DETAIL_WIDTH} marginRight={1} flexShrink={0}>
99          <Text dimColor>{one.detail}</Text>
100        </Box>
101      )}
102    </Box>
103  )
104}
105
106const runCommand = ($: EngineInterface, command: string) => {
107  $.command
108    .run({ command })
109    .catch(() => $.ui.toast(`cache-meter: /${command} could not run`))
110}
111
112const cycleThreshold = async ($: EngineInterface) => {
113  const value = await update(
114    $,
115    threshold,
116    held => STEPS[(STEPS.indexOf(held) + 1) % STEPS.length] ?? DEFAULT_THRESHOLD,
117  )
118  await $.store.set(STORE_KEY, value)
119}
120
121const setThreshold = async ($: EngineInterface, value: number) => {
122  if (STEPS.includes(value)) {
123    await update($, threshold, () => value)
124    await $.store.set(STORE_KEY, value)
125  }
126}
127
128export const register: Register = (on, options) => {
129  const layout: Layout = {
130    hasGauge: options.display !== 'text',
131    hasDetails: options.details !== false,
132  }
133  // Re-armed only once the fill drops below the threshold, so a compaction that lands above it does not loop.
134  let isArmed = true
135
136  on('session.start', async ($, e, next) => {
137    const stored = await $.store.get(STORE_KEY)
138
139    if (typeof stored === 'number' && STEPS.includes(stored)) {
140      await update($, threshold, () => stored)
141    }
142
143    const usage = await $.session.usage()
144    await update($, context, () => toFill(usage.context))
145
146    return next(e)
147  })
148
149  on('session.end', async ($, e, next) => {
150    if (e.reason === 'clear') {
151      isArmed = true
152      await update($, cache, () => null)
153      await update($, context, () => null)
154    }
155
156    return next(e)
157  })
158
159  on('turn.complete', async ($, e, next) => {
160    if (e.agentId === undefined && e.usage !== undefined) {
161      const served = e.usage.cache_read_input_tokens
162      const total =
163        e.usage.input_tokens + served + e.usage.cache_creation_input_tokens
164
165      if (total > 0) {
166        await update($, cache, () => ({ read: served, total }))
167      }
168    }
169
170    return next(e)
171  })
172
173  on('session.measure', async ($, e, next) => {
174    if (e.changed.includes('context')) {
175      const fill = toFill(e.context)
176      const limit = await read($, threshold)
177      await update($, context, () => fill)
178
179      if (fill !== null && fill.percent < limit) {
180        isArmed = true
181      } else if (fill !== null && isArmed) {
182        isArmed = false
183        // A compaction is refused while the turn that raised this measurement still runs.
184        $.clock.after(250, () => runCommand($, 'compact'))
185      }
186    }
187
188    return next(e)
189  })
190
191  on('session.compact', async ($, e, next) => {
192    const compacted = await next(e)
193    const isLive = e.agentId === undefined && e.trigger !== 'precompute'
194
195    if (isLive && compacted.skip === undefined) {
196      const after = compacted.tokensAfter
197      await update($, context, held =>
198        held === null || after === undefined
199          ? null
200          : {
201              tokens: after,
202              window: held.window,
203              percent: Math.round((after / held.window) * 100),
204            },
205      )
206    }
207
208    return compacted
209  }).catch(($, e, next) => next(e))
210
211  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
212    const below = await next(e)
213
214    if (e.props.hasSurvey) {
215      return below
216    }
217
218    const { Box, Button, Text } = $.ui.resolve(e)
219    const hit = await read($, cache)
220    const fill = await read($, context)
221    const limit = await read($, threshold)
222    const hitPercent = hit === null ? null : (hit.read / hit.total) * 100
223    let picker
224
225    // The mobile surface draws no Select.
226    if (e.surface === 'mobile') {
227      picker = (
228        <Button
229          key="threshold"
230          label={`Auto ${limit}%`}
231          onPress={() => void cycleThreshold($)}
232        />
233      )
234    } else {
235      const { Select } = $.ui.resolve(e)
236      picker = (
237        <Select
238          key="threshold"
239          label="Auto"
240          options={OPTIONS}
241          value={String(limit)}
242          onSelect={value => void setThreshold($, Number(value))}
243        />
244      )
245    }
246
247    return (
248      <Box flexDirection="column">
249        <Box flexDirection="row" flexWrap="wrap">
250          {meter({ Box, Text }, layout, {
251            label: 'Cache',
252            column: 0,
253            percent: hitPercent,
254            tone: cacheTone(hitPercent ?? 100),
255            detail:
256              hit === null ? '' : `(${tokens(hit.read)}/${tokens(hit.total)})`,
257          })}
258          {meter({ Box, Text }, layout, {
259            label: 'Context',
260            column: 1,
261            percent: fill === null ? null : fill.percent,
262            tone: contextTone(fill === null ? 0 : fill.percent, limit),
263            detail:
264              fill === null
265                ? ''
266                : `(${tokens(fill.tokens)}/${tokens(fill.window)})`,
267          })}
268          <Box flexDirection="row" gap={1} flexShrink={0}>
269            {picker}
270            <Button
271              key="compact"
272              label="Compact"
273              onPress={() => runCommand($, 'compact')}
274            />
275            <Button
276              key="clear"
277              label="Clear"
278              onPress={() => runCommand($, 'clear')}
279            />
280          </Box>
281        </Box>
282        {below}
283      </Box>
284    )
285  })
286}
287

1export type CacheHit = { read: number; total: number }
2
3export type ContextFill = { percent: number; tokens: number; window: number }
4
5declare module 'claude-code' {
6  interface PluginState {
7    'cache-meter': {
8      cache: CacheHit | null
9      context: ContextFill | null
10      threshold: number
11    }
12  }
13}
14