usage-bars

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, ModelUsage, Register, Timer } from 'claude-code'
3
4import type { AgentContext, Band, Context, Limit } from '../types'
5
6const limits = atom({ plugin: 'usage-bars', key: 'limits' } as const, [] as Limit[])
7const context = atom({ plugin: 'usage-bars', key: 'context' } as const, null as Context | null)
8const agentContexts = atom(
9  { plugin: 'usage-bars', key: 'agentContexts' } as const,
10  {} as Record<string, AgentContext>,
11)
12const mainModel = atom({ plugin: 'usage-bars', key: 'mainModel' } as const, null as string | null)
13const band = atom({ plugin: 'usage-bars', key: 'band' } as const, null as Band | null)
14
15const LABELS: Record<string, string> = {
16  five_hour: 'Session',
17  seven_day: 'Weekly',
18  spend_limit: 'Spend',
19}
20const ORDER = ['five_hour', 'seven_day', 'spend_limit']
21const WEEKDAYS = ['Sun', 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat']
22// The band reports its column less a five-cell mark of its own.
23const BAND_MARK = 5
24// The engine draws the hint line inset from both edges of that column.
25const LINE_INSET = 4
26const COLUMN_GAP = 3
27const MIN_BAR = 8
28// A seven-eighths block leaves a sliver between stacked bars; a full block doesn't.
29const BAR = '▇'
30const LABEL_WIDTH = 7
31// "Sun 00:20", "142k/200k"
32const EXTRA_WIDTH = 9
33const PERCENT_WIDTH = 5
34const LIMITS_KEY = 'lastLimits'
35const FIXED_WIDTH = LABEL_WIDTH + 1 + PERCENT_WIDTH + 2 + EXTRA_WIDTH
36// The engine reports a window size for the main conversation only; a subagent
37// on another model is assumed to have the standard one.
38const DEFAULT_WINDOW = 200_000
39
40// `percent` is null while there is nothing to measure yet.
41type Cell = { key: string; label: string; percent: number | null; extra: string }
42
43const pad2 = (n: number) => String(n).padStart(2, '0')
44
45function formatReset(iso: string | undefined, now: number): string {
46  if (!iso) return ''
47  const at = new Date(iso)
48  if (at.getTime() <= now) return 'reset'
49  const time = `${pad2(at.getHours())}:${pad2(at.getMinutes())}`
50  const isToday = new Date(now).toDateString() === at.toDateString()
51  return `${isToday ? '' : `${WEEKDAYS[at.getDay()]} `}${time}`
52}
53
54function formatTokens(n: number): string {
55  return n >= 999_500 ? `${(n / 1_000_000).toFixed(1)}M` : `${Math.round(n / 1000)}k`
56}
57
58function barColor(percent: number) {
59  if (percent >= 90) return 'error'
60  if (percent >= 70) return 'warning'
61  return 'success'
62}
63
64const inputTokens = (u: ModelUsage) =>
65  u.input_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens
66
67function contextCell(tokens: number | undefined, window: number, percent?: number): Cell {
68  if (tokens === undefined) {
69    return { key: 'context', label: 'Context', percent: null, extra: `–/${formatTokens(window)}` }
70  }
71  return {
72    key: 'context',
73    label: 'Context',
74    percent: percent ?? Math.round((tokens / window) * 100),
75    extra: `${formatTokens(tokens)}/${formatTokens(window)}`,
76  }
77}
78
79// The engine only learns of a reset from the next API response, so a window
80// whose reset time has passed is shown as empty until then.
81function zeroPassed(list: Limit[], now: number): Limit[] {
82  return list.map(l =>
83    l.resetsAt && Date.parse(l.resetsAt) <= now ? { ...l, percentUsed: 0 } : l,
84  )
85}
86
87let resetTimer: Timer | undefined
88
89function scheduleReset($: EngineInterface, list: Limit[], now: number) {
90  resetTimer?.cancel()
91  resetTimer = undefined
92  const upcoming = list
93    .map(l => (l.resetsAt ? Date.parse(l.resetsAt) : NaN))
94    .filter(t => t > now)
95  if (upcoming.length === 0) return
96  resetTimer = $.clock.after(Math.min(...upcoming) - now + 1000, async () => {
97    const at = await $.clock.now()
98    const current = await update($, limits, cur => zeroPassed(cur, at))
99    scheduleReset($, current, at)
100  })
101}
102
103// An empty reading means the engine has none yet, not that the limits went
104// away, so the last figures stay up.
105async function storeLimits($: EngineInterface, next: Limit[]) {
106  if (next.length === 0) return
107  const now = await $.clock.now()
108  const stored = zeroPassed(next, now)
109  await update($, limits, () => stored)
110  scheduleReset($, stored, now)
111}
112
113// Limits are the account's, not the conversation's, so the last reading is
114// kept across sessions: a /clear or resume starts with no reading of its own.
115async function lastKnownLimits($: EngineInterface): Promise<Limit[]> {
116  const saved = await $.store.get(LIMITS_KEY)
117  return Array.isArray(saved) ? (saved as Limit[]) : []
118}
119
120async function hydrate($: EngineInterface) {
121  const usage = await $.session.usage()
122  if (usage.rateLimits.length > 0) {
123    await storeLimits($, usage.rateLimits)
124    await $.store.set(LIMITS_KEY, usage.rateLimits)
125  } else {
126    await storeLimits($, await lastKnownLimits($))
127  }
128  const { tokens, window, percent } = usage.context
129  await update($, context, () => ({ tokens, window, percent }))
130}
131
132let relaySeq = 0
133let isHydrating = false
134
135export const register: Register = on => {
136  on('session.start', async ($, e, next) => {
137    const result = await next(e)
138    await hydrate($)
139    return result
140  })
141
142  // A /clear or resume keeps the process but starts the new conversation with
143  // empty state and no session.start; redrawing lets the hint line notice.
144  on('session.end', ($, e, next) => {
145    if (e.reason === 'clear' || e.reason === 'resume') {
146      $.clock.after(100, () => $.ui.invalidate('ui.render'))
147    }
148    return next(e)
149  })
150
151  on('session.measure', async ($, e, next) => {
152    // Stored every time: a new reset time with under a point of movement is
153    // not reported in `changed`.
154    await storeLimits($, e.rateLimits)
155    if (e.rateLimits.length > 0) await $.store.set(LIMITS_KEY, e.rateLimits)
156    const { tokens, window, percent } = e.context
157    await update($, context, () => ({ tokens, window, percent }))
158    return next(e)
159  })
160
161  on('turn.step', async function* ($, e, next) {
162    const result = yield* next(e)
163    const usage = result.usage
164    if (!usage) return result
165    const tokens = inputTokens(usage)
166    const agentId = e.agentId
167    if (agentId === undefined) {
168      // session.measure only fires once the whole turn ends; this keeps the
169      // main bar moving during a long turn.
170      await update($, mainModel, () => usage.model)
171      await update($, context, cur =>
172        cur ? { tokens, window: cur.window, percent: Math.round((tokens / cur.window) * 100) } : cur,
173      )
174      return result
175    }
176    await update($, agentContexts, all => ({
177      ...all,
178      [agentId]: { tokens, model: usage.model },
179    }))
180    return result
181  })
182
183  // Only the band above the prompt is told which transcript is in view and how
184  // wide its column is, so it relays both to the hint line. A render may not
185  // write state, hence the timer; the sequence drops relays a newer one overtook.
186  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
187    const seen: Band = {
188      agentId: e.props.view.agentId ?? null,
189      columns: e.props.bodyColumns + BAND_MARK,
190    }
191    const stored = await read($, band)
192    if (stored?.agentId !== seen.agentId || stored.columns !== seen.columns) {
193      const seq = ++relaySeq
194      $.clock.after(0, () => {
195        if (seq === relaySeq) void update($, band, () => seen)
196      })
197    }
198    return next(e)
199  })
200
201  on('ui.render', { component: 'PromptHint' }, async ($, e, next) => {
202    const line = await next(e)
203    const relayed = await read($, band)
204    const columns = relayed?.columns ?? e.viewport?.columns
205    if (columns === undefined) return line
206    const now = await $.clock.now()
207
208    const right: Cell[] = [...(await read($, limits))]
209      .sort((a, b) => rank(a.kind) - rank(b.kind))
210      .map(l => ({
211        key: l.kind,
212        label: LABELS[l.kind] ?? l.kind,
213        percent: Math.round(l.percentUsed),
214        extra: formatReset(l.resetsAt, now),
215      }))
216
217    const main = await read($, context)
218    // Context is always set once a conversation is hydrated, so its absence
219    // marks one that started without session.start (/clear, resume).
220    if (main === null && !isHydrating) {
221      isHydrating = true
222      $.clock.after(0, () => {
223        void hydrate($).finally(() => {
224          isHydrating = false
225        })
226      })
227    }
228    const agentId = relayed?.agentId ?? null
229    let ctx: Cell | null = null
230    if (agentId === null) {
231      if (main) ctx = contextCell(main.tokens, main.window, main.percent)
232    } else {
233      const agent = (await read($, agentContexts))[agentId]
234      const sameModel = agent !== undefined && agent.model === (await read($, mainModel))
235      const window = main && sameModel ? main.window : DEFAULT_WINDOW
236      ctx = contextCell(agent?.tokens ?? 0, window)
237    }
238    const left = ctx ? [ctx] : []
239    if (left.length === 0 && right.length === 0) return line
240
241    const { Box, Text } = $.ui.resolve(e)
242    const room = columns - LINE_INSET
243    const halfWidth = Math.floor((room - COLUMN_GAP) / 2)
244    const isTwoColumn = halfWidth >= FIXED_WIDTH + MIN_BAR
245    const cellWidth = isTwoColumn ? halfWidth : room
246    // Every part but the bar is fixed-width so the bar only changes with the
247    // column; on a very narrow one the reset/token column goes first.
248    const showExtra = cellWidth >= FIXED_WIDTH + MIN_BAR
249    const fixed = showExtra ? FIXED_WIDTH : FIXED_WIDTH - 2 - EXTRA_WIDTH
250    const barWidth = Math.max(1, cellWidth - fixed)
251
252    const cell = (c: Cell | undefined) => {
253      if (!c) return null
254      const percent = c.percent ?? 0
255      const filled = Math.round((Math.min(percent, 100) / 100) * barWidth)
256      return (
257        <Box key={c.key}>
258          <Text dimColor>{c.label.padEnd(LABEL_WIDTH)} </Text>
259          <Text color={barColor(percent)}>{BAR.repeat(filled)}</Text>
260          <Text dimColor>{BAR.repeat(barWidth - filled)}</Text>
261          {c.percent === null ? (
262            <Text dimColor>{'–'.padStart(PERCENT_WIDTH)}</Text>
263          ) : (
264            <Text>{` ${String(c.percent).padStart(3)}%`}</Text>
265          )}
266          {showExtra ? <Text dimColor>{`  ${c.extra.padEnd(EXTRA_WIDTH)}`}</Text> : null}
267        </Box>
268      )
269    }
270
271    if (!isTwoColumn) {
272      return (
273        <Box flexDirection="column">
274          {line}
275          {[...left, ...right].map(cell)}
276        </Box>
277      )
278    }
279
280    const rows = Math.max(left.length, right.length)
281    return (
282      <Box flexDirection="column">
283        {line}
284        {Array.from({ length: rows }, (_, i) => (
285          <Box key={`row-${i}`}>
286            <Box width={cellWidth} flexShrink={0}>
287              {cell(left[i])}
288            </Box>
289            <Box marginLeft={COLUMN_GAP} flexShrink={0}>
290              {cell(right[i])}
291            </Box>
292          </Box>
293        ))}
294      </Box>
295    )
296  })
297}
298
299function rank(kind: string) {
300  const i = ORDER.indexOf(kind)
301  return i === -1 ? ORDER.length : i
302}
303

1export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
2export type Context = { tokens?: number; window: number; percent?: number }
3export type AgentContext = { tokens: number; model: string }
4export type Band = { agentId: string | null; columns: number }
5
6declare module 'claude-code' {
7  interface PluginState {
8    'usage-bars': {
9      limits: Limit[]
10      context: Context | null
11      agentContexts: Record<string, AgentContext>
12      mainModel: string | null
13      band: Band | null
14    }
15  }
16}
17