usage-band

1import { atom, read, update } from 'claude-code'
2import type {
3  EngineInterface,
4  PromptOrigin,
5  Register,
6  RenderElement,
7  Timer,
8} from 'claude-code'
9
10import type { Limit } from '../types'
11
12const limits = atom({ plugin: 'usage-band', key: 'limits' } as const, [])
13const isHidden = atom({ plugin: 'usage-band', key: 'isHidden' } as const, false)
14const isLocked = atom({ plugin: 'usage-band', key: 'isLocked' } as const, false)
15const draft = atom({ plugin: 'usage-band', key: 'draft' } as const, '')
16
17const MINUTE = 60_000
18
19// One mark per window: an icon for the kind of window, then a dial. The dial's
20// outer ring is the share of the limit used, the wedge in its centre the share
21// of the window's time still left.
22const MARK_WIDTH = 41
23const MARK_HEIGHT = 22
24const RING = 59.06 // the circumference of the outer ring, r = 9.4
25const WEDGE = 18.85 // the circumference the wedge is drawn on, r = 3
26// The Svg is drawn as an image, so it cannot read the app's theme: these are
27// mid-tones that hold on both the light and the dark ground. The percentage
28// takes the same colour as its ring.
29const LOW = '#3A82DA'
30const WARNING = '#C98414'
31const CRITICAL = '#D2544A'
32const QUIET = '#87867F'
33const TRACK = '#8C8C8C'
34
35const HOURGLASS =
36  '<path d="M3.4 1.6h6.2M3.4 11.4h6.2M4 1.6v1.6l2.5 3.3L9 3.2V1.6M4 11.4V9.8l2.5-3.3L9 9.8v1.6"/>'
37const CALENDAR =
38  '<rect x="1.8" y="2.8" width="9.4" height="8.4" rx="1.6"/><path d="M1.8 5.7h9.4M4.4 1.5v2.4M8.6 1.5v2.4"/>'
39
40// The padlock is drawn like the window icons, on its own: open and quiet while
41// the session is free, shut and in the band's blue while it is locked.
42const LOCK_WIDTH = 13
43const PADLOCK_BODY = '<rect x="2.6" y="5.9" width="7.8" height="5.5" rx="1.4"/>'
44const SHACKLE_OPEN = 'M4.3 5.9V3.3a2.2 2.2 0 0 1 4.4 0'
45const SHACKLE_SHUT = 'M4.3 5.9V4.2a2.2 2.2 0 0 1 4.4 0v1.7'
46const KEYHOLE = 'M6.5 8.2v.9'
47// In the band the padlock is the lock's whole control. A Button's face is text
48// and an Svg takes no press, so a Button with a blank label is laid over the
49// padlock and takes the press for it. The label is an en space: it draws
50// nothing, and keeps the Button about as wide as it is tall.
51const BLANK = '\u2002'
52// How far the Box that holds the Button reaches past the padlock on each side,
53// in columns. The desktop draws a Button no wider than the Box it is in, and
54// this Button is wider than the padlock.
55const REACH = 2
56
57const SESSION = {
58  kind: 'five_hour',
59  label: 'Current session',
60  icon: HOURGLASS,
61  span: 5 * 60 * MINUTE,
62}
63const WEEK = {
64  kind: 'seven_day',
65  label: 'Weekly',
66  icon: CALENDAR,
67  span: 7 * 24 * 60 * MINUTE,
68}
69
70// The lock is one pane. `focus` asks the surface for the keyboard. Under
71// `closeOnEscape` Escape is the person's close of the pane, which the ui.close
72// hook refuses while locked: the key ends there and never stops the turn. Left
73// out, Escape would hand the keyboard back to the prompt box instead.
74const PANE = 'usage-band-lock'
75const OPEN = {
76  id: PANE,
77  title: 'Session lock',
78  focus: true,
79  closeOnEscape: true,
80} as const
81// How often a locked session checks that the pane still has the keyboard.
82const WATCH = 250
83
84// A prompt the person typed somewhere the lock cannot close: Enter at a
85// terminal, a message over Remote Control, a ping from Slack. The desktop app
86// is not on the list. It submits what is typed in its box as `sdk`, but it
87// sends the same way what goes out on a click or of its own accord, such as an
88// artifact comment sent to Claude, so `sdk` cannot tell typing apart. It need
89// not: the box takes no typing while locked (see prompt.edit), so nothing typed
90// there can reach the gate. Every other prompt passes a locked session: a
91// background task or agent reporting back, a schedule, another session, a
92// plugin, and whatever the desktop app sends.
93const PERSONS: readonly PromptOrigin['kind'][] = [
94  'composer',
95  'bridge',
96  'slack-ping',
97]
98
99const LOCK_ON = 'Session lock: on'
100const LOCK_OFF = 'Session lock: off'
101const LOCKED = 'Session locked.'
102const UNLOCKED = 'Session unlocked.'
103const NEEDS_DESKTOP = 'The session lock needs the Claude desktop app.'
104const USAGE = 'Usage: /usage-band, /usage-band lock or /usage-band unlock.'
105const REFUSAL =
106  'The session is locked, so this prompt was not sent. Unlock it first: press Unlock, or run /usage-band unlock.'
107const STAYS = 'the session is locked'
108const NOT_RETURNED =
109  'usage-band: your draft could not be put back. Run /usage-band unlock to try again.'
110
111const windowOf = (list: readonly Limit[], kind: string) =>
112  list.find(one => one.kind === kind) ??
113  list.find(one => one.kind.startsWith(kind))
114
115// Past its reset time a window is empty again: it is shown at 0% until the next
116// reply reports the new one.
117const hasReset = (limit: Limit, now: number) =>
118  limit.resetsAt !== undefined && Date.parse(limit.resetsAt) <= now
119
120const remaining = (resetsAt: string | undefined, now: number) => {
121  const left = resetsAt === undefined ? NaN : Date.parse(resetsAt) - now
122
123  if (Number.isNaN(left)) {
124    return ''
125  }
126
127  const minutes = Math.ceil(left / MINUTE)
128  const days = Math.floor(minutes / 1440)
129  const hours = Math.floor((minutes % 1440) / 60)
130
131  if (days > 0) {
132    return `resets in ${days}d ${hours}h`
133  }
134
135  if (hours > 0) {
136    return `resets in ${hours}h ${minutes % 60}m`
137  }
138
139  return `resets in ${minutes}m`
140}
141
142const share = (value: number) =>
143  Number.isNaN(value) ? 0 : Math.min(1, Math.max(0, value))
144
145const toneOf = (percent: number) =>
146  percent >= 90 ? CRITICAL : percent >= 70 ? WARNING : LOW
147
148const mark = (icon: string, used: number, left: number, colour: string) =>
149  [
150    `<svg xmlns="http://www.w3.org/2000/svg" width="${MARK_WIDTH}" height="${MARK_HEIGHT}" viewBox="0 0 ${MARK_WIDTH} ${MARK_HEIGHT}">`,
151    `<g transform="translate(0 4.5)" fill="none" stroke="${QUIET}" stroke-width="1.3" stroke-linecap="round" stroke-linejoin="round">${icon}</g>`,
152    `<circle cx="30" cy="11" r="9.4" fill="none" stroke="${TRACK}" stroke-opacity="0.3" stroke-width="2.6"/>`,
153    `<circle cx="30" cy="11" r="9.4" fill="none" stroke="${colour}" stroke-width="2.6" stroke-dasharray="${(RING * used).toFixed(2)} ${RING}" transform="rotate(-90 30 11)"/>`,
154    `<circle cx="30" cy="11" r="3" fill="none" stroke="${TRACK}" stroke-opacity="0.3" stroke-width="6"/>`,
155    `<circle cx="30" cy="11" r="3" fill="none" stroke="${QUIET}" stroke-width="6" stroke-dasharray="${(WEDGE * left).toFixed(2)} ${WEDGE}" transform="rotate(-90 30 11)"/>`,
156    '</svg>',
157  ].join('')
158
159const padlock = (isShut: boolean) =>
160  [
161    `<svg xmlns="http://www.w3.org/2000/svg" width="${LOCK_WIDTH}" height="${MARK_HEIGHT}" viewBox="0 0 ${LOCK_WIDTH} ${MARK_HEIGHT}">`,
162    `<g transform="translate(0 4.5)" fill="none" stroke="${isShut ? LOW : QUIET}" stroke-width="1.3" stroke-linecap="round" stroke-linejoin="round">`,
163    PADLOCK_BODY,
164    `<path d="${isShut ? SHACKLE_SHUT : SHACKLE_OPEN}${KEYHOLE}"/>`,
165    '</g>',
166    '</svg>',
167  ].join('')
168
169// The engine's own band is an empty Box; another mod's band has content.
170const hasContent = (drawn: RenderElement) =>
171  drawn.type !== 'Box' || (drawn.children?.length ?? 0) > 0
172
173const sameLimits = (a: readonly Limit[], b: readonly Limit[]) =>
174  a.length === b.length &&
175  a.every((one, i) => {
176    const other = b[i]
177
178    return (
179      other !== undefined &&
180      one.kind === other.kind &&
181      one.percentUsed === other.percentUsed &&
182      one.resetsAt === other.resetsAt
183    )
184  })
185
186// Claude Code stops reporting a window once its reset time has passed, so a
187// reading replaces the stored window of its kind and never removes one. A write
188// redraws the band, so only a list that differs is written.
189const keep = async ($: EngineInterface, reported: readonly Limit[]) => {
190  const stored = await read($, limits)
191  const next = [
192    ...stored.map(old => reported.find(one => one.kind === old.kind) ?? old),
193    ...reported.filter(one => !stored.some(old => old.kind === one.kind)),
194  ]
195
196  if (!sameLimits(stored, next)) {
197    await update($, limits, () => next)
198  }
199}
200
201const refresh = async ($: EngineInterface) => {
202  const { rateLimits } = await $.session.usage()
203  await keep($, rateLimits)
204}
205
206const isPersons = (origin: PromptOrigin) => PERSONS.includes(origin.kind)
207
208// Empties the prompt box and answers what was in it. Over text there the
209// surface refuses the pane the keyboard, and Escape in the box stops the turn.
210const emptied = async ($: EngineInterface) => {
211  const { text } = await $.prompt.read()
212
213  if (text.trim() === '') {
214    return ''
215  }
216
217  const { isFilled } = await $.prompt.fill({ text: '', mode: 'replace' })
218
219  return isFilled ? text : ''
220}
221
222// The timer and the two flags live as long as the module: a reload starts them
223// over, and the session's start puts the lock back.
224let watch: Timer | undefined
225let isHolding = false
226let hadKeyboard = false
227
228// Gives the pane the keyboard while the session is locked: opens it when it is
229// gone, and asks for the keyboard again when the person has clicked away.
230const hold = async ($: EngineInterface) => {
231  if (isHolding || !(await read($, isLocked))) {
232    return
233  }
234
235  isHolding = true
236
237  try {
238    const pane = (await $.ui.panes()).find(one => one.id === PANE)
239
240    if (pane?.isFocused) {
241      return
242    }
243
244    // What has landed in the box while locked was not meant for it: it goes.
245    await emptied($)
246    await $.ui.open(OPEN)
247  } finally {
248    isHolding = false
249  }
250}
251
252const disarm = () => {
253  watch?.cancel()
254  watch = undefined
255}
256
257const arm = ($: EngineInterface) => {
258  disarm()
259  watch = $.clock.every(WATCH, () => void hold($))
260}
261
262const engage = async ($: EngineInterface) => {
263  if (!(await read($, isLocked))) {
264    // The draft in the box as the lock goes on is the person's: it is set
265    // aside, and comes back on unlock.
266    const text = await emptied($)
267
268    if (text !== '') {
269      await update($, draft, () => text)
270    }
271
272    await update($, isLocked, () => true)
273  }
274
275  arm($)
276  await hold($)
277}
278
279const release = async ($: EngineInterface) => {
280  await update($, isLocked, () => false)
281  disarm()
282  await $.ui.close({ id: PANE })
283
284  const text = await read($, draft)
285
286  if (text === '') {
287    return
288  }
289
290  const { isFilled } = await $.prompt.fill({ text, mode: 'append' })
291
292  if (isFilled) {
293    await update($, draft, () => '')
294  } else {
295    $.ui.toast(NOT_RETURNED)
296  }
297}
298
299// A reload keeps $.state and drops the timer; /clear, /resume and /branch
300// reset $.state and may leave the pane up.
301const settle = async ($: EngineInterface) => {
302  if (await read($, isLocked)) {
303    arm($)
304    await hold($)
305
306    return
307  }
308
309  disarm()
310
311  if ((await $.ui.panes()).some(one => one.id === PANE)) {
312    await $.ui.close({ id: PANE })
313  }
314}
315
316export const register: Register = on => {
317  on('session.start', async ($, e, next) => {
318    await $.command.register({
319      name: 'usage-band',
320      description: 'Show or hide the usage band, or lock and unlock the session',
321      argumentHint: '[lock|unlock]',
322      // A session is locked, and unlocked, while it works.
323      immediate: true,
324    })
325    await refresh($)
326    $.clock.every(MINUTE, () => $.ui.invalidate('ui.render'))
327    await settle($)
328
329    return next(e)
330  })
331
332  // /clear, /resume and /branch reset $.state without a session.start.
333  on(
334    'classic.SessionStart',
335    { source: ['clear', 'resume', 'fork'] },
336    async ($, e, next) => {
337      await refresh($)
338      await settle($)
339
340      return next(e)
341    },
342  )
343
344  on('session.measure', async ($, e, next) => {
345    // Every measurement carries the exact figures, whatever `changed` names.
346    await keep($, e.rateLimits)
347
348    return next(e)
349  })
350
351  on('command.run', { command: 'usage-band' }, async ($, e) => {
352    const wanted = e.args.trim()
353
354    if (wanted === 'lock') {
355      if (!(await $.session.surfaces()).includes('desktop')) {
356        return { text: NEEDS_DESKTOP }
357      }
358
359      await engage($)
360
361      return { text: LOCKED }
362    }
363
364    if (wanted === 'unlock') {
365      await release($)
366
367      return { text: UNLOCKED }
368    }
369
370    if (wanted !== '') {
371      return { text: USAGE }
372    }
373
374    const hidden = await update($, isHidden, value => !value)
375
376    return { text: hidden ? 'Usage band hidden.' : 'Usage band shown.' }
377  })
378
379  // The prompt box takes no typing while locked: every edit is answered with an
380  // empty box, and the pane asks for the keyboard back.
381  on('prompt.edit', async ($, e, next) => {
382    if (!(await read($, isLocked))) {
383      return next(e)
384    }
385
386    void hold($)
387
388    return { text: '', cursor: 0 }
389  }).catch((_, e, next) => next(e))
390
391  // A prompt typed somewhere other than the desktop's box is refused while
392  // locked, and not kept. A dead gate must never swallow a prompt: it fails open.
393  on('prompt.submit', async ($, e, next) => {
394    if (!(await read($, isLocked)) || !isPersons(e.origin)) {
395      return next(e)
396    }
397
398    // The prompt box has just emptied, so the pane can have the keyboard back.
399    await $.ui.open(OPEN)
400
401    return { drop: REFUSAL }
402  }).catch((_, e, next) => next(e))
403
404  // While locked the pane closes for nobody. Escape and the close mark are the
405  // person's close, and refusing it is what keeps Escape from stopping the turn.
406  on('ui.close', { id: PANE }, async ($, e, next) =>
407    (await read($, isLocked)) ? { deny: STAYS } : next(e),
408  ).catch((_, e, next) => next(e))
409
410  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e, next) => {
411    if (!(await read($, isLocked))) {
412      return next(e)
413    }
414
415    // The keyboard has just left the pane. A drawing may not write, so the
416    // asking waits for the draw to end.
417    if (hadKeyboard && !e.props.isFocused) {
418      $.clock.after(0, () => void hold($))
419    }
420
421    hadKeyboard = e.props.isFocused
422
423    const { Box, Button, Text } = $.ui.resolve(e)
424    let sign: RenderElement | null = null
425
426    // The padlock is an Svg, which the terminal does not draw.
427    if (e.surface === 'desktop') {
428      const { Svg } = $.ui.resolve(e)
429      sign = (
430        <Svg
431          source={padlock(true)}
432          alt={LOCK_ON}
433          width={LOCK_WIDTH}
434          height={MARK_HEIGHT}
435        />
436      )
437    }
438
439    return (
440      <Box flexDirection="column" rowGap={1}>
441        <Box alignItems="center" columnGap={1}>
442          {sign}
443          <Text bold>Session locked</Text>
444        </Box>
445        <Box>
446          <Button key="unlock" variant="primary" onPress={() => release($)}>
447            Unlock
448          </Button>
449        </Box>
450      </Box>
451    )
452  })
453
454  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
455    // The band is one tree for every mod: each adds its own line and passes the
456    // lines beneath through as they came. Ours goes under what another mod drew;
457    // we never open that up or set it in our row.
458    const beneath = await next(e)
459
460    if (e.surface !== 'desktop' || e.props.hasSurvey) {
461      return beneath
462    }
463
464    const now = await $.clock.now()
465    const list = await read($, limits)
466    const session = windowOf(list, SESSION.kind)
467    const week = windowOf(list, WEEK.kind)
468    const isQuiet = (!session && !week) || (await read($, isHidden))
469
470    if (isQuiet) {
471      return beneath
472    }
473
474    const isShut = await read($, isLocked)
475    const { Box, Button, Svg, Text } = $.ui.resolve(e)
476
477    const meter = (one: typeof SESSION, limit: Limit) => {
478      const isOver = hasReset(limit, now)
479      const used = isOver ? 0 : limit.percentUsed
480      const percent = Math.round(used)
481      const reset = isOver ? 'reset' : remaining(limit.resetsAt, now)
482      const left =
483        limit.resetsAt === undefined
484          ? 0
485          : share((Date.parse(limit.resetsAt) - now) / one.span)
486      const colour = toneOf(percent)
487      const alt =
488        reset === ''
489          ? `${one.label}: ${percent}% used`
490          : `${one.label}: ${percent}% used, ${reset}`
491
492      return (
493        <Box alignItems="center" columnGap={1}>
494          <Svg
495            source={mark(one.icon, share(used / 100), left, colour)}
496            alt={alt}
497            width={MARK_WIDTH}
498            height={MARK_HEIGHT}
499          />
500          <Text bold color={colour}>
501            {`${percent}%`}
502          </Text>
503        </Box>
504      )
505    }
506
507    const ownLine = (
508      <Box alignItems="center" columnGap={3} flexWrap="wrap">
509        {session && meter(SESSION, session)}
510        {week && meter(WEEK, week)}
511        <Box>
512          <Svg
513            source={padlock(isShut)}
514            alt={isShut ? LOCK_ON : LOCK_OFF}
515            width={LOCK_WIDTH}
516            height={MARK_HEIGHT}
517          />
518          <Box
519            position="absolute"
520            top={0}
521            bottom={0}
522            left={-REACH}
523            right={-REACH}
524            alignItems="center"
525            justifyContent="center"
526          >
527            <Button
528              key="lock"
529              label={BLANK}
530              plain
531              onPress={async () =>
532                (await read($, isLocked)) ? release($) : engage($)
533              }
534            />
535          </Box>
536        </Box>
537      </Box>
538    )
539
540    if (!hasContent(beneath)) {
541      return ownLine
542    }
543
544    return (
545      <Box flexDirection="column" rowGap="md">
546        {beneath}
547        {ownLine}
548      </Box>
549    )
550  })
551}
552

1export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
2
3declare module 'claude-code' {
4  interface PluginState {
5    'usage-band': {
6      limits: Limit[]
7      isHidden: boolean
8      isLocked: boolean
9      draft: string
10    }
11  }
12}
13