usage-limits

1import { atom, read, update } from 'claude-code'
2import type {
3  EngineInterface,
4  Register,
5  RenderChildren,
6  SessionAppendMessage,
7  SessionContextUsage,
8  SessionRateLimit,
9  SessionUsage,
10} from 'claude-code'
11
12import type { ContextFill, LimitWindow, Limits } from '../types'
13
14const TICK_MS = 30_000
15const MINUTE_MS = 60_000
16const HOUR_MS = 60 * MINUTE_MS
17const DAY_MS = 24 * HOUR_MS
18const WARNING_AT = 70
19const ERROR_AT = 90
20const BAR_WIDTH = 8
21// The context is judged on the percent as it is shown: green up to 40, yellow
22// 41 to 60, dark red from 61. The limits keep WARNING_AT and ERROR_AT.
23const CONTEXT_GREEN_UNTIL = 40
24const CONTEXT_YELLOW_UNTIL = 60
25const DARK_RED = '#b91c1c'
26// The editor's theme leaves its `warning` colour unset, and a drawing cannot
27// use a theme colour at all, so the editor gets its colours given outright.
28const EDITOR_GREEN = '#25E712'
29const EDITOR_YELLOW = '#ca8a04'
30const EDITOR_RED = '#C1121C'
31const RING_TRACK = 'rgba(140, 140, 140, 0.35)'
32const RING_SIZE = 16
33const RING_RADIUS = 6
34const RING_STROKE = 2.5
35// Below this many columns the rings and the text do not fit one line, and the
36// editor draws the readings as text alone. The editor counts a column as the
37// width of a box-drawing character in its own font, some 13 to 16 px, not as a
38// monospace cell. 0 means the width is not known yet.
39const EDITOR_RINGS_FROM = 32
40const SEPARATOR = ' · '
41const COMPACT_ICON = '♻️'
42const CLOSE_ICON = '×'
43const COMMAND = 'limits'
44const COMPACTION_SETTLE_MS = 1_500
45
46type Segment = {
47  key: string
48  label: string
49  percent: number
50  color: string
51  detail: string
52}
53
54const limits = atom(
55  { plugin: 'usage-limits', key: 'limits' } as const,
56  { session: null, week: null },
57)
58const context = atom({ plugin: 'usage-limits', key: 'context' } as const, null)
59const compacting = atom({ plugin: 'usage-limits', key: 'isCompacting' } as const, false)
60const hidden = atom({ plugin: 'usage-limits', key: 'isHidden' } as const, false)
61const tick = atom({ plugin: 'usage-limits', key: 'tick' } as const, 0)
62
63const toContextFill = (usage: SessionContextUsage): ContextFill | null => {
64  const { percent, tokens, window } = usage
65
66  if (percent !== undefined) {
67    return { percent, tokens: tokens ?? null, window }
68  }
69  if (tokens === undefined || window <= 0) {
70    return null
71  }
72
73  return { percent: (tokens / window) * 100, tokens, window }
74}
75
76const toWindow = (rate: SessionRateLimit | undefined): LimitWindow | null =>
77  rate === undefined
78    ? null
79    : { percentUsed: rate.percentUsed, resetsAt: rate.resetsAt ?? null }
80
81const toLimits = (rates: readonly SessionRateLimit[]): Limits => ({
82  session: toWindow(rates.find(rate => rate.kind === 'five_hour')),
83  week: toWindow(rates.find(rate => rate.kind === 'seven_day')),
84})
85
86const formatTokens = (count: number): string => {
87  if (count >= 1_000_000) return `${Number((count / 1_000_000).toFixed(1))}M`
88  if (count >= 1_000) return `${Math.round(count / 1_000)}k`
89  return String(count)
90}
91
92const formatRemaining = (ms: number): string => {
93  const days = Math.floor(ms / DAY_MS)
94  const hours = Math.floor((ms % DAY_MS) / HOUR_MS)
95  const minutes = Math.floor((ms % HOUR_MS) / MINUTE_MS)
96
97  if (days > 0) return `${days}d ${hours}h`
98  if (hours > 0) return `${hours}h ${minutes}m`
99  return `${Math.max(minutes, 1)}m`
100}
101
102const describeContext = (fill: ContextFill): string =>
103  fill.tokens === null
104    ? `window ${formatTokens(fill.window)}`
105    : `${formatTokens(fill.tokens)}/${formatTokens(fill.window)}`
106
107const hasReset = (window: LimitWindow, now: number): boolean =>
108  window.resetsAt !== null && Date.parse(window.resetsAt) <= now
109
110const describeReset = (window: LimitWindow, now: number): string => {
111  if (window.resetsAt === null) return ''
112
113  return hasReset(window, now)
114    ? 'just reset'
115    : formatRemaining(Date.parse(window.resetsAt) - now)
116}
117
118// The engine reports a limit only with a response, so a window that has reset
119// since still carries the reading from before. It starts again from nothing.
120const usedNow = (window: LimitWindow, now: number): number =>
121  hasReset(window, now) ? 0 : window.percentUsed
122
123const drawBar = (percent: number): string => {
124  const filled = Math.round((Math.min(percent, 100) * BAR_WIDTH) / 100)
125
126  return '█'.repeat(filled) + '░'.repeat(BAR_WIDTH - filled)
127}
128
129const pickColor = (percent: number): 'success' | 'warning' | 'error' => {
130  if (percent >= ERROR_AT) return 'error'
131  if (percent >= WARNING_AT) return 'warning'
132  return 'success'
133}
134
135const pickContextColor = (percent: number): string => {
136  const shown = Math.round(percent)
137
138  if (shown <= CONTEXT_GREEN_UNTIL) return 'success'
139  if (shown <= CONTEXT_YELLOW_UNTIL) return 'warning'
140  return DARK_RED
141}
142
143const editorColor = (color: string): string => {
144  if (color === 'success') return EDITOR_GREEN
145  if (color === 'warning') return EDITOR_YELLOW
146  if (color === 'error' || color === DARK_RED) return EDITOR_RED
147  return color
148}
149
150// A ring that fills clockwise from the top, as a drawing: the editor's font
151// does not keep the blocks of a text bar to one width.
152const drawRing = (percent: number, color: string): string => {
153  const circumference = 2 * Math.PI * RING_RADIUS
154  const filled = (Math.min(Math.max(percent, 0), 100) * circumference) / 100
155  const centre = RING_SIZE / 2
156  const circle = (stroke: string, extra = '') =>
157    `<circle cx="${centre}" cy="${centre}" r="${RING_RADIUS}" fill="none" stroke="${stroke}" stroke-width="${RING_STROKE}"${extra}/>`
158  const arc =
159    filled > 0
160      ? circle(
161          color,
162          ` stroke-dasharray="${filled.toFixed(2)} ${circumference.toFixed(2)}" stroke-linecap="round" transform="rotate(-90 ${centre} ${centre})"`,
163        )
164      : ''
165
166  return `<svg xmlns="http://www.w3.org/2000/svg" width="${RING_SIZE}" height="${RING_SIZE}" viewBox="0 0 ${RING_SIZE} ${RING_SIZE}">${circle(RING_TRACK)}${arc}</svg>`
167}
168
169const toSegments = (
170  fill: ContextFill | null,
171  { session, week }: Limits,
172  now: number,
173): Segment[] => {
174  const segments: Segment[] = []
175
176  if (fill !== null) {
177    segments.push({
178      key: 'context',
179      label: 'Context',
180      percent: fill.percent,
181      color: pickContextColor(fill.percent),
182      detail: describeContext(fill),
183    })
184  }
185  if (session !== null) {
186    const percent = usedNow(session, now)
187
188    segments.push({
189      key: 'session',
190      label: 'Session',
191      percent,
192      color: pickColor(percent),
193      detail: describeReset(session, now),
194    })
195  }
196  if (week !== null) {
197    const percent = usedNow(week, now)
198
199    segments.push({
200      key: 'week',
201      label: 'Week',
202      percent,
203      color: pickColor(percent),
204      detail: describeReset(week, now),
205    })
206  }
207
208  return segments
209}
210
211// What the context is worth once a compaction is done. The engine reports its
212// own figure only with the next response, so until then its estimate of the
213// live transcript is the best reading there is.
214// The figure the engine reports, or else its estimate of the live transcript.
215const fillFrom = (usage: SessionUsage): ContextFill | null =>
216  toContextFill(usage.context) ??
217  toContextFill({ window: usage.context.window, tokens: usage.context.breakdown?.totalTokens })
218
219const refreshContext = async ($: EngineInterface) => {
220  try {
221    const fill = fillFrom(await $.session.usage({ breakdown: 'summary' }))
222
223    if (fill !== null) {
224      await update($, context, () => fill)
225    }
226  } catch {
227    // The bar is a convenience: the next response reports the real figure.
228  }
229}
230
231// Straight after a compaction the estimate has not counted the new
232// conversation in full yet, and it settles within a second or two.
233const refreshSoon = ($: EngineInterface) => {
234  $.clock.after(COMPACTION_SETTLE_MS, () => {
235    void refreshContext($)
236  })
237}
238
239// A compaction that nothing announced still shows: the engine has no figure
240// for the context until the next response, while the bar still has the old one.
241const catchUpWithCompaction = async ($: EngineInterface) => {
242  if ((await read($, context)) === null) return
243
244  try {
245    const { context: usage } = await $.session.usage()
246
247    if (usage.tokens === undefined) {
248      await refreshContext($)
249    }
250  } catch {
251    // The next tick tries again.
252  }
253}
254
255// What a finished /compact leaves in the conversation: the command and its
256// output, as the engine wraps them.
257const COMPACT_ROW = /<command-name>\/?compact<\/command-name>|<local-command-stdout>\s*Compacted/
258
259const isCompactRow = ({ content }: SessionAppendMessage): boolean =>
260  content.some(block => block.type === 'text' && COMPACT_ROW.test(String(block.text)))
261
262// Runs the real `/compact`, so it shows in the chat as if it were typed. It is
263// queued until the session is idle, so a press during a turn is not lost.
264const compactNow = async ($: EngineInterface) => {
265  await update($, compacting, () => true)
266
267  try {
268    await $.command.run({ command: 'compact' })
269    refreshSoon($)
270  } catch {
271    $.ui.toast('Compact could not be started')
272  } finally {
273    await update($, compacting, () => false)
274  }
275}
276
277// A refused name must not stop the timer and the readings that follow.
278const registerCommand = async ($: EngineInterface) => {
279  try {
280    await $.command.register({
281      name: COMMAND,
282      description: 'Show or hide the row with the context and the limits',
283      immediate: true,
284    })
285  } catch {
286    $.ui.toast(`/${COMMAND} could not be registered`)
287  }
288}
289
290export const register: Register = on => {
291  on('session.start', async ($, e, next) => {
292    $.clock.every(TICK_MS, () => {
293      void catchUpWithCompaction($)
294
295      return update($, tick, n => n + 1)
296    })
297    await update($, compacting, () => false)
298    await registerCommand($)
299
300    const usage = await $.session.usage()
301    await update($, limits, () => toLimits(usage.rateLimits))
302    await update($, context, () => toContextFill(usage.context))
303
304    return next(e)
305  })
306
307  // The same switch the close button throws, so `/limits` brings a closed band
308  // back and hides an open one.
309  on('command.run', { command: COMMAND }, async $ => {
310    const isNowHidden = !(await read($, hidden))
311
312    await update($, hidden, () => isNowHidden)
313
314    return {
315      text: isNowHidden
316        ? 'The limits row is hidden. /limits brings it back.'
317        : 'The limits row is shown.',
318    }
319  }).catch(() => ({ text: 'Could not toggle the limits row.' }))
320
321  // Every compaction leaves the context smaller. A subagent's own does not
322  // count. Not every surface announces one here: the desktop's /compact does
323  // not, and the rows it leaves in the conversation stand in for it below.
324  on('session.compact', async ($, e, next) => {
325    const result = await next(e)
326
327    if (e.trigger !== 'precompute' && e.agentId === undefined && result.skip === undefined) {
328      refreshSoon($)
329    }
330
331    return result
332  }).catch(($, e, next) => next(e))
333
334  // The command's rows come once it is done, so a typed /compact is caught as
335  // well as the button's.
336  on('session.append', { door: 'command' }, async ($, e, next) => {
337    const result = await next(e)
338
339    if (e.agentId === undefined && isCompactRow(e.message)) {
340      refreshSoon($)
341    }
342
343    return result
344  }).catch(($, e, next) => next(e))
345
346  on('session.measure', async ($, e, next) => {
347    await update($, limits, () => toLimits(e.rateLimits))
348    // Just after a compaction the engine has no context figure until the next
349    // response, and the last estimate is kept rather than blanked.
350    await update($, context, fill => toContextFill(e.context) ?? fill)
351
352    return next(e)
353  })
354
355  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
356    await read($, tick)
357
358    const now = await $.clock.now()
359    const isBusy = await read($, compacting)
360    const isHidden = await read($, hidden)
361    const segments = toSegments(await read($, context), await read($, limits), now)
362
363    if (e.props.hasSurvey || isHidden || segments.length === 0) {
364      return next(e)
365    }
366
367    const { Box, Button, Text } = $.ui.resolve(e)
368    const hasContext = segments.some(segment => segment.key === 'context')
369
370    // In the editor's panel the band is one small line in the caption size, with
371    // no labels. The readings are rings beside the text when the band is wide
372    // enough, else the text alone.
373    const isEditor = e.surface === 'vscode'
374    const Svg = e.surface === 'vscode' ? $.ui.resolve(e).Svg : null
375    const columns = e.props.bodyColumns
376    const hasRings = Svg !== null && (columns === 0 || columns >= EDITOR_RINGS_FROM)
377    // The editor wraps a Text unless it says otherwise, and the parts inside the
378    // one that truncates wrap on their own, so each says it.
379    const textProps = isEditor
380      ? { size: 'caption' as const, wrap: 'truncate-end' as const }
381      : {}
382    const tint = (color: string) => (isEditor ? editorColor(color) : color)
383
384    const drawSegment = ({ key, label, percent, color, detail }: Segment) => (
385      <Text key={key} {...textProps}>
386        {isEditor ? null : `${label} `}
387        {isEditor ? null : (
388          <Text color={color} mono>
389            {`${drawBar(percent)} `}
390          </Text>
391        )}
392        <Text color={tint(color)} mono={!isEditor} {...textProps}>
393          {`${Math.round(percent)}%`}
394        </Text>
395        {detail === '' ? null : (
396          <Text dimColor {...textProps}>
397            {` ${detail}`}
398          </Text>
399        )}
400      </Text>
401    )
402
403    const drawRingSegment = ({ key, label, percent, color, detail }: Segment) =>
404      Svg === null ? null : (
405        <Box key={key} flexDirection="row" alignItems="center" columnGap="sm">
406          <Svg
407            key="ring"
408            source={drawRing(percent, tint(color))}
409            alt={`${label} ${Math.round(percent)}%`}
410            width={RING_SIZE}
411            height={RING_SIZE}
412          />
413          <Text wrap="truncate-end" {...textProps}>
414            <Text color={tint(color)} {...textProps}>
415              {`${Math.round(percent)}%`}
416            </Text>
417            {detail === '' ? null : (
418              <Text dimColor {...textProps}>
419                {` ${detail}`}
420              </Text>
421            )}
422          </Text>
423        </Box>
424      )
425
426    const separator = (index: number) => (
427      <Text key={`separator-${index}`} dimColor {...textProps}>
428        {SEPARATOR}
429      </Text>
430    )
431
432    const joined = (draw: (segment: Segment) => RenderChildren) =>
433      segments.flatMap((segment, index) =>
434        index === 0 ? [draw(segment)] : [separator(index), draw(segment)],
435      )
436
437    // The icon stands for the context beside it, which stays plain coloured
438    // text: the desktop draws a Button's children without their colours.
439    // Pressed during a turn it only toasts, as the compaction is refused.
440    const drawCompact = () =>
441      isBusy ? (
442        <Text dimColor>⏳</Text>
443      ) : (
444        <Button
445          key="compact"
446          label={COMPACT_ICON}
447          dimColor
448          onPress={() => {
449            void compactNow($)
450          }}
451        />
452      )
453
454    // The outer row spans the band: the compact button at its left edge, the
455    // close button at its right, and the readings centred in the room between.
456    return (
457      <Box
458        flexDirection="row"
459        alignItems="center"
460        columnGap="sm"
461        width="100%"
462        {...(isEditor ? { flexWrap: 'nowrap' as const } : {})}
463      >
464        <Box flexDirection="row" alignItems="center">
465          {hasContext ? drawCompact() : null}
466        </Box>
467        <Box
468          flexDirection="row"
469          alignItems="center"
470          justifyContent="center"
471          flexGrow={1}
472          flexShrink={1}
473          {...(isEditor ? { minWidth: 0 } : {})}
474        >
475          {hasRings ? (
476            <Box flexDirection="row" alignItems="center" columnGap="sm" flexWrap="nowrap" overflow="hidden">
477              {joined(drawRingSegment)}
478            </Box>
479          ) : (
480            // One Text, so the segments keep a single row on every surface.
481            <Text wrap="truncate-end" {...textProps}>
482              {joined(drawSegment)}
483            </Text>
484          )}
485        </Box>
486        {/* The desktop draws a dismiss Button as a close control. It hides the
487            band in this chat only: the state is the session's, so a new chat
488            shows it again. */}
489        <Button
490          key="close"
491          role="dismiss"
492          label={CLOSE_ICON}
493          dimColor
494          onPress={() => {
495            void update($, hidden, () => true)
496          }}
497        />
498      </Box>
499    )
500  })
501}
502

1export type LimitWindow = {
2  /** 0 to 100, as the rate-limit notices read it. */
3  percentUsed: number
4  /** ISO 8601 timestamp of the reset, when the API reported one. */
5  resetsAt: string | null
6}
7
8export type Limits = {
9  /** The five-hour session window. */
10  session: LimitWindow | null
11  /** The seven-day weekly window. */
12  week: LimitWindow | null
13}
14
15export type ContextFill = {
16  /** How much of the context window the last response filled, 0 to 100. */
17  percent: number
18  /** Input tokens the last response was answered over, when reported. */
19  tokens: number | null
20  /** The model's context window, in tokens. */
21  window: number
22}
23
24declare module 'claude-code' {
25  interface PluginState {
26    'usage-limits': {
27      limits: Limits
28      /** The live context window of this session; null before a reading. */
29      context: ContextFill | null
30      /** True while a compaction started from the context button runs. */
31      isCompacting: boolean
32      /** True once the person closed the band in this chat. */
33      isHidden: boolean
34      /** Bumped by a timer so the countdown redraws between API responses. */
35      tick: number
36    }
37  }
38}
39