limits-band

1import type { EngineInterface, Register } from 'claude-code'
2
3import { colorFor, formatShort, remainingPercent } from './format'
4
5type Limit = { kind: string; percentUsed: number; resetsAt?: string }
6type Sample = { limits: Limit[]; contextUsed?: number }
7
8// Limits belong to the account, not the chat. A live reading is saved to the
9// store so a chat with no reply yet can show the last one any session saw.
10// Two readings of one limit: a later reset time is a newer window; within the
11// same window, use only goes up, so the higher figure is the fresher one.
12function fresher(a: Limit, b: Limit): Limit {
13  const ta = a.resetsAt ? Date.parse(a.resetsAt) : 0
14  const tb = b.resetsAt ? Date.parse(b.resetsAt) : 0
15
16  if (ta !== tb) return ta > tb ? a : b
17
18  return a.percentUsed >= b.percentUsed ? a : b
19}
20
21function mergeLimits(live: Limit[], saved: Limit[]): Limit[] {
22  const merged = new Map<string, Limit>()
23
24  for (const l of [...saved, ...live]) {
25    const current = merged.get(l.kind)
26    merged.set(l.kind, current ? fresher(current, l) : l)
27  }
28
29  return [...merged.values()]
30}
31
32// Limits belong to the account, not the chat: merge this chat's last reading
33// with the one any session saved, and save the result back.
34async function sample($: EngineInterface): Promise<Sample> {
35  const usage = await $.session.usage()
36  const contextUsed = usage.context.percent
37  // 'limitsShared', not 'limits': sessions still running the first version
38  // keep overwriting 'limits' with their own stale reading.
39  const stored = await $.store.get('limitsShared')
40  const saved = Array.isArray(stored) ? (stored as Limit[]) : []
41  const limits = mergeLimits(usage.rateLimits, saved)
42
43  if (limits.length > 0) await $.store.set('limitsShared', limits)
44
45  return { limits, contextUsed }
46}
47
48export const register: Register = on => {
49  let limits: Limit[] = []
50  let contextUsed: number | undefined
51  let isTicking = false
52
53  // A resumed chat has no rate-limit reading until its first reply, so an
54  // empty sample must not wipe the last known one.
55  const apply = (s: Sample) => {
56    limits = s.limits
57    if (s.contextUsed !== undefined) contextUsed = s.contextUsed
58  }
59
60  on('session.start', async ($, e, next) => {
61    if (!isTicking) {
62      isTicking = true
63      $.clock.every(10_000, async () => {
64        apply(await sample($))
65        $.ui.invalidate('ui.render')
66      })
67    }
68    apply(await sample($))
69    $.ui.invalidate('ui.render')
70
71    return next(e)
72  })
73
74  on('prompt.submit', async ($, e, next) => {
75    $.clock.after(4_000, async () => {
76      apply(await sample($))
77      $.ui.invalidate('ui.render')
78    })
79
80    return next(e)
81  })
82
83  on('turn.complete', async ($, e, next) => {
84    apply(await sample($))
85    $.ui.invalidate('ui.render')
86    $.clock.after(3_000, async () => {
87      apply(await sample($))
88      $.ui.invalidate('ui.render')
89    })
90
91    return next(e)
92  })
93
94  on('ui.render', { component: 'AbovePrompt' }, ($, e, next) => {
95    const { Box, Text } = $.ui.resolve(e)
96    const hourly = limits.find(l => l.kind === 'five_hour')
97    const weekly = limits.find(l => l.kind === 'seven_day')
98
99    if (!hourly && !weekly && contextUsed === undefined) {
100      return (
101        <Box justifyContent="center" width="100%">
102          <Text dimColor>Limits: waiting for first reply…</Text>
103        </Box>
104      )
105    }
106
107    const cell = (label: string, left: number, reset = '') => (
108      <Box key={label} flexDirection="row">
109        <Text bold>{label} </Text>
110        <Text bold color={colorFor(left)}>
111          {left}%
112        </Text>
113        <Text dimColor>{reset ? ` (${reset})` : ''}</Text>
114      </Box>
115    )
116
117    const cells = []
118
119    // A saved reading whose window has since reset means a fresh, full window.
120    const windowCell = (label: string, l: Limit) => {
121      const msLeft = l.resetsAt ? Date.parse(l.resetsAt) - Date.now() : undefined
122
123      if (msLeft !== undefined && msLeft <= 0) return cell(label, 100)
124
125      return cell(label, remainingPercent(l.percentUsed), msLeft === undefined ? '' : formatShort(msLeft))
126    }
127
128    if (hourly) cells.push(windowCell('Hourly Left', hourly))
129    if (weekly) cells.push(windowCell('Weekly Left', weekly))
130
131    if (contextUsed !== undefined) {
132      cells.push(cell('Context Left', remainingPercent(contextUsed)))
133    }
134
135    const parts = cells.flatMap((c, i) =>
136      i === 0 ? [c] : [<Text key={`dot${i}`} dimColor>{'   ·   '}</Text>, c],
137    )
138
139    return (
140      <Box flexDirection="row" justifyContent="center" width="100%">
141        {parts}
142      </Box>
143    )
144  })
145}
146

1export const formatShort = (ms: number): string => {
2  const totalMin = Math.max(0, Math.floor(ms / 60000))
3  const days = Math.floor(totalMin / 1440)
4  const hours = Math.floor((totalMin % 1440) / 60)
5  const mins = totalMin % 60
6
7  if (days > 0) return `${days}d ${hours}h`
8  if (hours > 0) return `${hours}h ${mins}m`
9
10  return `${mins}m`
11}
12
13export const remainingPercent = (percentUsed: number): number =>
14  Math.min(100, Math.max(0, Math.round(100 - percentUsed)))
15
16export const colorFor = (remaining: number): string | undefined =>
17  remaining < 20 ? 'error' : undefined
18