context-colour

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionContextUsage, SessionRateLimit } from 'claude-code'
3
4import type { Fill, Limit } from '../types'
5
6// The context window degrades well before it is full; a rate-limit window
7// only matters near its end.
8const CONTEXT_AT = { warn: 40, danger: 65 }
9const LIMIT_AT = { warn: 50, danger: 80 }
10
11// A fixed, short bar of one glyph: a width counted in cells wraps where the
12// surface's font is not the terminal's, and two glyphs do not sit on one line.
13const CELLS = 6
14const BAR = '━'
15// The painted bar's height, as a part of one row.
16const THICKNESS = 0.28
17
18// The part of a painted bar that is filled, in percent: a window in use always
19// shows a sliver. A whole number, as a Box's percentage width must be.
20export const shareOf = (percent: number): number =>
21  percent <= 0 ? 0 : Math.min(100, Math.max(3, Math.round(percent)))
22
23// The rate-limit windows drawn, in order, and what each is called.
24const WINDOWS = [
25  { kind: 'five_hour', title: '5h ' },
26  { kind: 'seven_day', title: 'Week ' },
27]
28
29// A window in use always shows at least one cell.
30export const usedOf = (percent: number): number =>
31  percent <= 0 ? 0 : Math.min(CELLS, Math.max(1, Math.round((CELLS * percent) / 100)))
32
33const fill = atom({ plugin: 'context-colour', key: 'fill' } as const, null)
34const limits = atom({ plugin: 'context-colour', key: 'limits' } as const, [])
35// The time the countdowns are told from: written each minute, so they redraw.
36const at = atom({ plugin: 'context-colour', key: 'at' } as const, 0)
37
38export const colourOf = (percent: number, at = CONTEXT_AT): string =>
39  percent >= at.danger ? 'error' : percent >= at.warn ? 'warning' : 'success'
40
41export const limitColourOf = (percent: number): string => colourOf(percent, LIMIT_AT)
42
43const short = (tokens: number): string =>
44  tokens >= 1_000_000
45    ? `${(tokens / 1_000_000).toFixed(1)}M`
46    : `${Math.round(tokens / 1000)}k`
47
48export const labelOf = (now: Fill): string =>
49  ` ${now.percent}% · ${short(now.tokens)}/${short(now.window)}`
50
51const measured = (context: SessionContextUsage): Fill | null =>
52  context.percent === undefined || context.tokens === undefined
53    ? null
54    : { percent: context.percent, tokens: context.tokens, window: context.window }
55
56const windowsOf = (rateLimits: SessionRateLimit[]): Limit[] =>
57  rateLimits.map(one => {
58    const resetsAt = one.resetsAt === undefined ? NaN : Date.parse(one.resetsAt)
59
60    return {
61      kind: one.kind,
62      percent: one.percentUsed,
63      ...(Number.isNaN(resetsAt) ? {} : { resetsAt }),
64    }
65  })
66
67const MINUTE = 60_000
68const HOUR = 60 * MINUTE
69const DAY = 24 * HOUR
70
71// How long until a window resets, in its two largest units.
72export const leftOf = (ms: number): string => {
73  if (ms <= 0) {
74    return 'now'
75  }
76
77  const minutes = Math.ceil(ms / MINUTE)
78  const days = Math.floor((minutes * MINUTE) / DAY)
79  const hours = Math.floor((minutes * MINUTE) / HOUR) % 24
80
81  return days > 0
82    ? `${days}d ${hours}h`
83    : hours > 0
84      ? `${hours}h ${minutes % 60}m`
85      : `${minutes}m`
86}
87
88export const limitLabelOf = (one: Limit, now: number): string =>
89  ` ${Math.round(one.percent)}%` +
90  (one.resetsAt === undefined ? '' : ` · resets ${leftOf(one.resetsAt - now)}`)
91
92// A label is its percentage, drawn in the meter's colour, and what follows the
93// first dot, drawn dim.
94const SPLIT = ' · '
95
96export const figureOf = (label: string): string => label.split(SPLIT)[0] ?? label
97
98export const detailOf = (label: string): string =>
99  label.includes(SPLIT) ? label.slice(label.indexOf(SPLIT)) : ''
100
101let ticking = false
102
103// Reads the clock into `at`, and from the first call on, once a minute. Local
104// work alone: no model call is made.
105const tell = async ($: EngineInterface): Promise<void> => {
106  const now = await $.clock.now()
107  await update($, at, () => now)
108
109  if (!ticking) {
110    ticking = true
111    $.clock.every(MINUTE, () => tell($))
112  }
113}
114
115export const register: Register = on => {
116  on('session.start', async ($, e, next) => {
117    const { context, rateLimits } = await $.session.usage()
118    await update($, fill, () => measured(context))
119    await update($, limits, () => windowsOf(rateLimits))
120    await tell($)
121
122    return next(e)
123  })
124
125  on('session.measure', async ($, e, next) => {
126    await update($, fill, () => measured(e.context))
127    await update($, limits, () => windowsOf(e.rateLimits))
128    await tell($)
129
130    return next(e)
131  })
132
133  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
134    const now = await read($, fill)
135    const windows = await read($, limits)
136    const below = await next(e)
137    const told = await read($, at)
138
139    const meters = [
140      ...(now === null
141        ? []
142        : [{ title: 'Context ', percent: now.percent, colour: colourOf(now.percent), label: labelOf(now) }]),
143      ...WINDOWS.flatMap(({ kind, title }) => {
144        const one = windows.find(each => each.kind === kind)
145
146        return one === undefined
147          ? []
148          : [{ title, percent: one.percent, colour: limitColourOf(one.percent), label: limitLabelOf(one, told) }]
149      }),
150    ]
151
152    if (e.props.hasSurvey || meters.length === 0) {
153      return below
154    }
155
156    const { Box, Text } = $.ui.resolve(e)
157    // The desktop paints boxes, so its bar is two of them at a part of a row's
158    // height: solid, centred, and as thick as THICKNESS says. A terminal has
159    // whole cells alone and keeps the line glyph.
160    const painted = e.surface === 'desktop'
161
162    return (
163      <Box flexDirection="column">
164        <Box flexWrap="wrap" columnGap={2} width="100%">
165          {meters.map(({ title, percent, colour, label }) => (
166            <Box flexWrap="nowrap" alignItems="center" flexGrow={painted ? 1 : 0}>
167              <Text color="text" wrap="truncate">{title}</Text>
168              {painted ? (
169                // The bars take what room the row has left, shared between them.
170                <Box alignItems="center" flexGrow={1} minWidth={CELLS}>
171                  <Box width={`${shareOf(percent)}%`} height={THICKNESS} backgroundColor={colour} />
172                  <Box width={`${100 - shareOf(percent)}%`} height={THICKNESS} backgroundColor="subtle" />
173                </Box>
174              ) : (
175                <Box>
176                  <Text color={colour} wrap="truncate">{BAR.repeat(usedOf(percent))}</Text>
177                  <Text dimColor wrap="truncate">{BAR.repeat(CELLS - usedOf(percent))}</Text>
178                </Box>
179              )}
180              <Text color={colour} wrap="truncate">{figureOf(label)}</Text>
181              {detailOf(label) === '' ? null : <Text dimColor wrap="truncate">{detailOf(label)}</Text>}
182            </Box>
183          ))}
184        </Box>
185        {below}
186      </Box>
187    )
188  })
189}
190

1export type Fill = { percent: number; tokens: number; window: number }
2
3// `resetsAt` is in `$.clock.now()`'s milliseconds.
4export type Limit = { kind: string; percent: number; resetsAt?: number }
5
6declare module 'claude-code' {
7  interface PluginState {
8    'context-colour': { fill: Fill | null; limits: Limit[]; at: number }
9  }
10}
11