context-gauge

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Segment } from '../types'
5import { DEFAULT_THRESHOLD, bandView, evaluateAlert, formatTokens } from './gauge'
6
7const BAR_CELLS = 20
8
9const tokens = atom({ plugin: 'context-gauge', key: 'tokens' } as const, null)
10const isExpanded = atom({ plugin: 'context-gauge', key: 'isExpanded' } as const, false)
11const segments = atom({ plugin: 'context-gauge', key: 'segments' } as const, null)
12const isAlertArmed = atom({ plugin: 'context-gauge', key: 'isAlertArmed' } as const, true)
13
14async function refreshSegments($: EngineInterface) {
15  const { context } = await $.session.usage({ breakdown: 'summary' })
16  const rows: Segment[] = (context.breakdown?.categories ?? [])
17    .filter(category => category.kind === 'used' && category.tokens > 0)
18    .map(({ name, color, tokens }) => ({ name, color, tokens }))
19    .sort((a, b) => b.tokens - a.tokens)
20  await update($, segments, () => rows)
21}
22
23async function record(
24  $: EngineInterface,
25  { current, threshold }: { current: number | null; threshold: number },
26) {
27  await update($, tokens, () => current)
28
29  if (current !== null) {
30    const alert = evaluateAlert({
31      tokens: current,
32      threshold,
33      isArmed: await read($, isAlertArmed),
34    })
35    await update($, isAlertArmed, () => alert.isArmed)
36    if (alert.shouldToast) {
37      $.ui.toast(
38        `Context past ${formatTokens(threshold)} (${formatTokens(current)}). /compact at the next phase boundary.`,
39      )
40    }
41  }
42
43  await refreshSegments($)
44}
45
46export const register: Register = (on, options) => {
47  const threshold =
48    typeof options.threshold === 'number' && options.threshold > 0
49      ? options.threshold
50      : DEFAULT_THRESHOLD
51
52  on('session.start', async ($, e, next) => {
53    const result = await next(e)
54    const { context } = await $.session.usage()
55    await record($, { current: context.tokens ?? null, threshold })
56    return result
57  })
58
59  on('session.measure', async ($, e, next) => {
60    if (e.changed.includes('context')) {
61      await record($, { current: e.context.tokens ?? null, threshold })
62    }
63    return next(e)
64  })
65
66  on('session.compact', async ($, e, next) => {
67    const result = await next(e)
68    if (e.agentId === undefined && e.trigger !== 'precompute' && result.messages) {
69      await record($, { current: null, threshold })
70    }
71    return result
72  })
73
74  on('session.end', async ($, e, next) => {
75    if (e.reason === 'clear') {
76      await update($, isAlertArmed, () => true)
77      await record($, { current: null, threshold })
78    }
79    return next(e)
80  })
81
82  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
83    if (e.props.hasSurvey) {
84      return next(e)
85    }
86
87    const { Box, Button, Text } = $.ui.resolve(e)
88    const view = bandView({
89      tokens: await read($, tokens),
90      threshold,
91      segments: await read($, segments),
92      isExpanded: await read($, isExpanded),
93      surface: e.surface,
94    })
95
96    return (
97      <Box flexDirection="column">
98        <Box flexDirection="row" justifyContent="space-between">
99          <Box flexDirection="row" flexGrow={1} marginRight={2}>
100            <Text bold>Context</Text>
101            <Text color={view.count.color}>{`  ${view.count.text}  `}</Text>
102            <Box flexDirection="row" flexGrow={1} minWidth={BAR_CELLS} height={1} alignItems="center">
103              {view.fills.map(fill => (
104                <Box flexGrow={fill.grow} height={view.barHeight} backgroundColor={fill.color} />
105              ))}
106            </Box>
107          </Box>
108          <Button
109            key="toggle"
110            plain
111            label={view.toggleLabel}
112            onPress={() => update($, isExpanded, value => !value)}
113          />
114        </Box>
115        {view.details.state === 'measuring' && <Box marginTop={1}><Text dimColor>measuring…</Text></Box>}
116        {view.details.state === 'shown' && (
117          <Box flexDirection="row" flexWrap="wrap" columnGap={2} marginTop={1}>
118            {view.details.rows.map(row => (
119              <Box flexDirection="row" flexShrink={0} gap={1}>
120                <Text color={row.color}>■</Text>
121                <Text>{row.name}</Text>
122                <Text dimColor>{row.tokens}</Text>
123              </Box>
124            ))}
125          </Box>
126        )}
127      </Box>
128    )
129  })
130}
131

1import type { ThemeKey } from 'claude-code'
2
3import type { Segment } from '../types'
4
5export const DEFAULT_THRESHOLD = 150_000
6const DESKTOP_BAR_HEIGHT = '85%'
7const WARN_RATIO = 0.8
8
9type Zone = 'ok' | 'warn' | 'over'
10
11const ZONE_COLOR = {
12  ok: 'success',
13  warn: 'warning',
14  over: 'error',
15} as const satisfies Record<Zone, ThemeKey>
16
17const PENDING_COLOR: ThemeKey = 'subtle'
18
19export type Fill = { color?: string; grow: number }
20
21export type DetailsRow = { name: string; color: string; tokens: string }
22
23export type Details =
24  | { state: 'hidden' }
25  | { state: 'measuring' }
26  | { state: 'shown'; rows: DetailsRow[] }
27
28export type BandView = {
29  count: { text: string; color?: ThemeKey }
30  fills: Fill[]
31  barHeight: number | string
32  details: Details
33  toggleLabel: 'Show' | 'Hide'
34}
35
36type Level = { tokens: number; threshold: number }
37
38const zoneOf = ({ tokens, threshold }: Level): Zone => {
39  if (tokens > threshold) {
40    return 'over'
41  }
42
43  return tokens >= threshold * WARN_RATIO ? 'warn' : 'ok'
44}
45
46const barFills = ({
47  tokens,
48  threshold,
49  segments,
50}: Level & { segments: readonly Segment[] | null }): Fill[] => {
51  const filled = Math.min(1, Math.max(0, tokens / threshold)) * 100
52  const total = segments?.reduce((sum, segment) => sum + segment.tokens, 0) ?? 0
53  const fills: Fill[] =
54    segments?.length && total > 0
55      ? segments.map(segment => ({ color: segment.color, grow: (segment.tokens / total) * filled }))
56      : [{ color: PENDING_COLOR, grow: filled }]
57
58  return [...fills, { grow: 100 - filled }].filter(fill => fill.grow > 0)
59}
60
61const detailsOf = ({
62  segments,
63  isExpanded,
64}: {
65  segments: readonly Segment[] | null
66  isExpanded: boolean
67}): Details => {
68  if (!isExpanded) {
69    return { state: 'hidden' }
70  }
71  if (segments === null) {
72    return { state: 'measuring' }
73  }
74  return {
75    state: 'shown',
76    rows: segments.map(({ name, color, tokens }) => ({ name, color, tokens: formatTokens(tokens) })),
77  }
78}
79
80export const formatTokens = (tokens: number): string => {
81  if (tokens < 1000) {
82    return String(Math.round(tokens))
83  }
84  return `${Math.round(tokens / 1000)}k`
85}
86
87export const bandView = ({
88  tokens,
89  threshold,
90  segments,
91  isExpanded,
92  surface,
93}: {
94  tokens: number | null
95  threshold: number
96  segments: readonly Segment[] | null
97  isExpanded: boolean
98  surface: string
99}): BandView => {
100  const count = (tokens === null ? '—' : formatTokens(tokens)).padStart(4)
101
102  return {
103    count: {
104      text: `${count} of ${formatTokens(threshold)}`,
105      color: tokens === null ? undefined : ZONE_COLOR[zoneOf({ tokens, threshold })],
106    },
107    fills: barFills({ tokens: tokens ?? 0, threshold, segments }),
108    barHeight: surface === 'desktop' ? DESKTOP_BAR_HEIGHT : 1,
109    details: detailsOf({ segments, isExpanded }),
110    toggleLabel: isExpanded ? 'Hide' : 'Show',
111  }
112}
113
114export const evaluateAlert = ({
115  tokens,
116  threshold,
117  isArmed,
118}: Level & { isArmed: boolean }): { shouldToast: boolean; isArmed: boolean } => {
119  if (isArmed && tokens > threshold) {
120    return { shouldToast: true, isArmed: false }
121  }
122  if (!isArmed && tokens < threshold * WARN_RATIO) {
123    return { shouldToast: false, isArmed: true }
124  }
125  return { shouldToast: false, isArmed }
126}
127

1export type Segment = { name: string; color: string; tokens: number }
2
3declare module 'claude-code' {
4  interface PluginState {
5    'context-gauge': {
6      tokens: number | null
7      isExpanded: boolean
8      segments: Segment[] | null
9      isAlertArmed: boolean
10    }
11  }
12}
13