usage-limits

1import { atom, read, update } from 'claude-code'
2import type { Color, Elements, EngineInterface, Register } from 'claude-code'
3
4import type { Reading, UsageWindow } from '../types'
5
6const LABELS: Record<string, string> = {
7  five_hour: '5h',
8  seven_day: 'Week',
9  spend_limit: 'Spend',
10}
11const ORDER = ['five_hour', 'seven_day', 'spend_limit']
12const MINUTE = 60_000
13const HOUR = 60 * MINUTE
14const DAY = 24 * HOUR
15const STORE_KEY = 'reading'
16
17// Shared with cache-meter 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 windows = atom({ plugin: 'usage-limits', key: 'windows' } as const, [])
34const now = atom({ plugin: 'usage-limits', key: 'now' } as const, 0)
35const staleSince = atom(
36  { plugin: 'usage-limits', key: 'staleSince' } as const,
37  null,
38)
39
40const labelOf = (kind: string) => LABELS[kind] ?? kind.replaceAll('_', ' ')
41
42const rankOf = (kind: string) => {
43  const rank = ORDER.indexOf(kind)
44
45  return rank === -1 ? ORDER.length : rank
46}
47
48const toneOf = (percent: number): Color =>
49  percent >= 90 ? 'error' : percent >= 75 ? 'warning' : 'text'
50
51const span = (duration: number) => {
52  const days = Math.floor(duration / DAY)
53  const hours = Math.floor((duration % DAY) / HOUR)
54  const minutes = Math.max(1, Math.floor((duration % HOUR) / MINUTE))
55
56  if (days > 0) {
57    return `${days}d ${hours}h`
58  }
59
60  return hours > 0 ? `${hours}h ${minutes}m` : `${minutes}m`
61}
62
63const untilReset = (resetsAt: string | undefined, at: number) => {
64  const left = resetsAt === undefined ? Number.NaN : Date.parse(resetsAt) - at
65
66  if (Number.isNaN(left)) {
67    return ''
68  }
69
70  return left <= 0 ? '(now)' : `(${span(left)})`
71}
72
73const isOpen = (one: UsageWindow, at: number) =>
74  one.resetsAt === undefined || Date.parse(one.resetsAt) > at
75
76const meter = ({ Box, Text }: Tags, layout: Layout, one: Cell) => {
77  const filled =
78    one.percent === null
79      ? 0
80      : Math.max(
81          0,
82          Math.min(SEGMENTS, Math.round((one.percent / 100) * SEGMENTS)),
83        )
84
85  return (
86    <Box flexDirection="row" flexShrink={0}>
87      <Box
88        minWidth={LABEL_WIDTHS[one.column] ?? LABEL_WIDTHS[1]}
89        marginRight={1}
90        flexShrink={0}
91      >
92        <Text dimColor>{one.label}</Text>
93      </Box>
94      {layout.hasGauge && (
95        <Box minWidth={SEGMENTS} marginRight={1} flexShrink={0}>
96          <Text color={one.tone}>{'▮'.repeat(filled)}</Text>
97          <Text dimColor>{'▯'.repeat(SEGMENTS - filled)}</Text>
98        </Box>
99      )}
100      <Box
101        minWidth={PERCENT_WIDTH}
102        marginRight={1}
103        flexShrink={0}
104        justifyContent="flex-end"
105      >
106        <Text color={one.percent === null ? 'inactive' : one.tone}>
107          {one.percent === null ? '–' : `${Math.round(one.percent)}%`}
108        </Text>
109      </Box>
110      {layout.hasDetails && (
111        <Box minWidth={DETAIL_WIDTH} marginRight={1} flexShrink={0}>
112          <Text dimColor>{one.detail}</Text>
113        </Box>
114      )}
115    </Box>
116  )
117}
118
119const refresh = async ($: EngineInterface, list?: readonly UsageWindow[]) => {
120  const at = await $.clock.now()
121  await update($, now, () => at)
122
123  if (list === undefined) {
124    return
125  }
126
127  const sorted = [...list].sort((a, b) => rankOf(a.kind) - rankOf(b.kind))
128  await update($, windows, () => sorted)
129  await update($, staleSince, () => null)
130
131  if (sorted.length > 0) {
132    await $.store.set(STORE_KEY, { windows: sorted, at })
133  }
134}
135
136const restore = async ($: EngineInterface) => {
137  const at = await $.clock.now()
138  const stored = (await $.store.get(STORE_KEY)) as Reading | undefined
139  await update($, now, () => at)
140
141  if (stored !== undefined) {
142    await update($, windows, () => stored.windows.filter(one => isOpen(one, at)))
143    await update($, staleSince, () => stored.at)
144  }
145}
146
147export const register: Register = (on, options) => {
148  const layout: Layout = {
149    hasGauge: options.display !== 'text',
150    hasDetails: options.details !== false,
151  }
152
153  on('session.start', async ($, e, next) => {
154    const usage = await $.session.usage()
155
156    // The engine has no reading before the first response of a session.
157    if (usage.rateLimits.length > 0) {
158      await refresh($, usage.rateLimits)
159    } else {
160      await restore($)
161    }
162
163    $.clock.every(MINUTE, () => void refresh($))
164
165    return next(e)
166  })
167
168  on('session.measure', async ($, e, next) => {
169    await refresh($, e.rateLimits)
170
171    return next(e)
172  })
173
174  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
175    const below = await next(e)
176
177    if (e.props.hasSurvey) {
178      return below
179    }
180
181    const { Box, Text } = $.ui.resolve(e)
182    const list = await read($, windows)
183    const at = await read($, now)
184    const since = await read($, staleSince)
185
186    return (
187      <Box flexDirection="column">
188        {list.length === 0 ? (
189          <Text dimColor>Usage limits: no reading yet</Text>
190        ) : (
191          <Box flexDirection="row" flexWrap="wrap">
192            {list.map((one, column) =>
193              meter({ Box, Text }, layout, {
194                label: labelOf(one.kind),
195                column,
196                percent: one.percentUsed,
197                tone: toneOf(one.percentUsed),
198                detail: untilReset(one.resetsAt, at),
199              }),
200            )}
201            {since !== null && (
202              <Text dimColor>{`as of ${span(at - since)} ago`}</Text>
203            )}
204          </Box>
205        )}
206        {below}
207      </Box>
208    )
209  })
210}
211

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