prompt-hud

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { Cat, Limit, Pair, Snap, Tokens } from '../types'
5
6const limits = atom({ plugin: 'prompt-hud', key: 'limits' } as const, [])
7const tokens = atom({ plugin: 'prompt-hud', key: 'tokens' } as const, null)
8const error = atom({ plugin: 'prompt-hud', key: 'error' } as const, null)
9const snap = atom({ plugin: 'prompt-hud', key: 'snap' } as const, null)
10
11const DAY = 864e5
12const TTL = 60_000
13
14// Read-only aggregate over the transcripts. Streamed messages repeat one id with a
15// growing count, so the last row per (session, id) wins. `all` counts rows since $s
16// (the weekly window start); `chat` counts every row of session $sid. Input is
17// uncached + cache-write tokens; cache reads are not counted.
18export const JQ = `reduce (inputs | select(contains("\\"output_tokens\\"")) | (try fromjson catch empty)
19  | select(.type == "assistant" and .message.usage != null)
20  | (input_filename | ltrimstr($b) | split("/")) as $p
21  | {s: (if ($p | length) > 2 then $p[1] else ($p[1] | sub("\\\\.jsonl$"; "")) end), m: (.message.id // .uuid),
22     w: ((.timestamp // "") >= $s),
23     i: ((.message.usage.input_tokens // 0) + (.message.usage.cache_creation_input_tokens // 0)),
24     o: (.message.usage.output_tokens // 0)}
25) as $r ({}; .[$r.s + "|" + $r.m] = $r)
26| reduce .[] as $r ({all: {i: 0, o: 0}, chat: {i: 0, o: 0}};
27    (if $r.w then .all.i += $r.i | .all.o += $r.o else . end)
28    | (if $r.s == $sid then .chat.i += $r.i | .chat.o += $r.o else . end))`
29
30export const fmt = (n: number): string =>
31  n >= 1e9 ? `${(n / 1e9).toFixed(1)}B`
32  : n >= 1e6 ? `${(n / 1e6).toFixed(1)}M`
33  : n >= 1e3 ? `${Math.round(n / 1e3)}k`
34  : `${n}`
35
36export const until = (iso: string | undefined, now: number): string => {
37  if (!iso) return ''
38  const ms = Date.parse(iso) - now
39  if (!(ms > 0)) return 'now'
40  const h = Math.floor(ms / 36e5)
41  const m = Math.floor((ms % 36e5) / 6e4)
42  return h >= 24 ? `${Math.floor(h / 24)}d ${h % 24}h` : h > 0 ? `${h}h ${m}m` : `${m}m`
43}
44
45// Split `width` cells between two segments of a total; rounding never overfills.
46export const cells = (a: number, b: number, total: number, width: number): [number, number] => {
47  if (total <= 0) return [0, 0]
48  const x = Math.round((a / total) * width)
49  const y = Math.min(width - x, Math.round((b / total) * width))
50  return [x, y]
51}
52
53export const pad = (s: string, n: number): string => s + ' '.repeat(Math.max(0, n - s.length))
54
55export const summary = (l5: Limit | undefined, l7: Limit | undefined, chat: Pair, all: Pair): string =>
56  `5h ${l5 ? Math.round(l5.percentUsed) + '%' : '—'} · wk ${l7 ? Math.round(l7.percentUsed) + '%' : '—'}` +
57  ` · chat ↑${fmt(chat.i)} ↓${fmt(chat.o)} · all ↑${fmt(all.i)} ↓${fmt(all.o)}`
58
59// colours are theme keys, so light and dark both read
60const tone = (p: number) => (p >= 90 ? 'error' : p >= 70 ? 'warning' : 'success')
61
62
63let busy = false
64let lastScan = 0
65
66
67async function scan($: any, resetsAt: string) {
68  const home = await $.env.get('HOME')
69  const sid = await $.session.id()
70  const base = `${home}/.claude/projects/`
71  const start = Date.parse(resetsAt) - 7 * DAY // the weekly window began 7 days before its reset
72  const now = await $.clock.now()
73  const mins = Math.ceil((now - start) / 6e4) + 1
74  const found = await $.process.run(['find', base, '-type', 'f', '-name', '*.jsonl', '-mmin', `-${mins}`])
75  const files = found.stdout.split('\n').filter(Boolean)
76  let out = { all: { i: 0, o: 0 }, chat: { i: 0, o: 0 } }
77  if (files.length) {
78    const r = await $.process.run(
79      ['jq', '-R', '-n', '-c', '--arg', 's', new Date(start).toISOString(), '--arg', 'b', base, '--arg', 'sid', sid, JQ, ...files],
80      { timeoutMs: 120_000 },
81    )
82    if (r.exitCode !== 0) throw new Error(`jq exit ${r.exitCode}`)
83    out = JSON.parse(r.stdout || '{}')
84  }
85  await update($, tokens, () => ({ chat: out.chat, all: out.all, at: now }))
86  await update($, error, () => null)
87}
88
89// Never awaited by a hook: a slow scan only delays its own numbers.
90async function refreshUsage($: any) {
91  if (busy) return
92  busy = true
93  try {
94    const u = await $.session.usage()
95    if (u.rateLimits.length) await update($, limits, () => u.rateLimits)
96    const wk = u.rateLimits.find((l: Limit) => l.kind === 'seven_day')
97    if (wk?.resetsAt && (await $.clock.now()) - lastScan >= TTL) {
98      lastScan = await $.clock.now()
99      await scan($, wk.resetsAt)
100    }
101  } catch (err) {
102    await update($, error, () => String((err as Error)?.message ?? err))
103  } finally {
104    busy = false
105  }
106}
107
108
109const NAMES: [RegExp, string][] = [
110  [/system prompt/i, 'system prompt'],
111  [/mcp/i, 'mcp tools'],
112  [/tools/i, 'tools'],
113  [/agents/i, 'agents'],
114  [/memory/i, 'memory files'],
115  [/skills/i, 'skills'],
116  [/messages/i, 'messages'],
117]
118
119export const catLabel = (c: Pick<Cat, 'name' | 'kind'>): string =>
120  c.kind === 'free' ? 'free' : c.kind === 'buffer' ? 'compact buffer' : (NAMES.find(([re]) => re.test(c.name))?.[1] ?? c.name.toLowerCase())
121
122// The engine's own theme colour per row, as /context draws it; free space is always muted.
123export const hue = (c: Pick<Cat, 'color' | 'kind'>): string => (c.kind === 'free' ? 'subtle' : c.color || 'subtle')
124
125export const fmtK = (n: number): string =>
126  n >= 1e6 ? `${+(n / 1e6).toFixed(1)}M` : n >= 1e3 ? `${+(n / 1e3).toFixed(n < 1e4 ? 1 : 0)}k` : `${n}`
127
128export const share = (n: number, of: number): string => {
129  const p = (n / of) * 100
130  return n > 0 && p < 1 ? '<1%' : `${Math.round(p)}%`
131}
132
133export type Seg = { color: string; pct: number; isMark?: boolean }
134
135// Percent-width segments in category order; the compaction threshold becomes a 1-cell marker
136// splitting the category it falls in. Widths are percentages, so the bar fills any container.
137export const segments = (cats: Pick<Cat, 'tokens' | 'color' | 'kind'>[], threshold: number | null): Seg[] => {
138  const total = cats.reduce((a, c) => a + c.tokens, 0)
139  if (total <= 0) return []
140  const out: Seg[] = []
141  let at = 0
142  let marked = threshold === null
143  for (const c of cats) {
144    const from = at
145    const to = at + c.tokens
146    at = to
147    const color = hue(c)
148    if (!marked && threshold! >= from && threshold! < to) {
149      if (threshold! > from) out.push({ color, pct: ((threshold! - from) / total) * 100 })
150      out.push({ color: 'warning', pct: 0, isMark: true })
151      if (to > threshold!) out.push({ color, pct: ((to - threshold!) / total) * 100 })
152      marked = true
153    } else if (c.tokens > 0) out.push({ color, pct: (c.tokens / total) * 100 })
154  }
155  return out
156}
157
158async function refreshContext($: any) {
159  try {
160    const u = await $.session.usage({ breakdown: 'summary' })
161    const b = u.context.breakdown
162    if (!b) {
163      // no breakdown yet (fresh session): fall back to the status line's figures
164      const used = u.context.tokens ?? 0
165      await update($, snap, () => ({ cats: [], used, window: u.context.window, threshold: null, percent: u.context.percent ?? 0, err: null }))
166      return
167    }
168    await update($, snap, () => ({
169      cats: b.categories
170        .filter((c: Cat) => c.kind !== 'deferred' && (c.tokens > 0 || c.kind === 'free'))
171        .map((c: Cat) => ({ name: c.name, tokens: c.tokens, kind: c.kind, color: c.color })),
172      used: b.totalTokens,
173      window: b.rawMaxTokens,
174      threshold: b.isAutoCompactEnabled ? (b.autoCompactThreshold ?? null) : null,
175      percent: b.percentage,
176      err: null,
177    }))
178  } catch (err) {
179    await update($, snap, s => ({ cats: [], used: 0, window: 0, threshold: null, percent: 0, ...s, err: String((err as Error)?.message ?? err) }))
180  }
181}
182
183
184export function contextCard({ Box, Text }: any, s: Snap) {
185  // the bar is percent-width empty Boxes with a background: it fills the card, whatever the font does
186  const cats = s.cats
187  const segs = segments(cats, s.threshold === null || !s.window ? null : s.threshold)
188  if (!cats.length) {
189    const p = s.window ? Math.min(100, (s.used / s.window) * 100) : 0
190    segs.push({ color: 'claude', pct: p }, { color: 'subtle', pct: 100 - p })
191  }
192
193  const tone = s.percent >= 85 ? 'error' : s.percent >= 60 ? 'warning' : 'success'
194  const head = s.threshold === null ? 'auto-compact off' : `compacts at ${fmtK(s.threshold)}`
195
196  const mine = (
197    <Box flexDirection="column" paddingX={1} borderStyle="round" borderColor="claude" borderDimColor width="100%">
198      <Box justifyContent="space-between">
199        <Box>
200          <Text bold color="claude">✻ </Text>
201          <Text bold>context</Text>
202        </Box>
203        <Box columnGap={1}>
204          <Text bold>{fmtK(s.used)}</Text>
205          <Text dimColor>of {fmtK(s.window)} · {head}</Text>
206          <Text bold inverse color={tone}> {Math.round(s.percent)}% </Text>
207        </Box>
208      </Box>
209      <Box width="100%" height={1}>
210        {segs.map((r, i) => (
211          r.isMark
212            ? <Box key={String(i)} width={1} height={1} flexShrink={0} backgroundColor={r.color} />
213            : <Box key={String(i)} width={0} flexGrow={Math.max(1, Math.round(r.pct * 100))} minWidth={1} height={1} backgroundColor={r.color} />
214        ))}
215      </Box>
216      <Box columnGap={2} flexWrap="wrap">
217        {cats.map(c => (
218          <Box key={c.name} columnGap={1}>
219            <Text color={hue(c)}>●</Text>
220            <Text>{catLabel(c)}</Text>
221            <Text bold>{fmtK(c.tokens)}</Text>
222            <Text dimColor>{share(c.tokens, s.window)}</Text>
223          </Box>
224        ))}
225        {s.err ? <Text color="error">breakdown unavailable: {s.err}</Text> : cats.length ? <Text dimColor italic>est.</Text> : null}
226      </Box>
227    </Box>
228  )
229  return mine
230}
231
232
233export type ViewData = { limits: Limit[]; t: Tokens | null; err: string | null; now: number; cols: number }
234
235// One card: token numbers on the left, the 5h and week bars on the right, row for row.
236export function usageCard({ Box, Text }: any, d: ViewData, hasAbove: boolean) {
237  const l5 = d.limits.find(l => l.kind === 'five_hour')
238  const l7 = d.limits.find(l => l.kind === 'seven_day')
239  const { t, err, now } = d
240  const LABEL = 6, UP = 9, DOWN = 9, PCT = 5, RESET = 13
241  const stacked = d.cols < 90
242
243  const cell = (w: number, child: any) => <Box width={w} flexShrink={0}>{child}</Box>
244  const label = (name: string) => cell(LABEL, <Text bold color="claude">{name}</Text>)
245
246  const tokenRow = (name: string, p?: Pair) => (
247    <Box>
248      {label(name)}
249      {cell(UP, <Text color="suggestion">{p ? '↑' + fmt(p.i) : '…'}</Text>)}
250      {cell(DOWN, <Text color="claude">{p ? '↓' + fmt(p.o) : ''}</Text>)}
251    </Box>
252  )
253
254  // the bar takes the width left over; its segments are whole-number flexGrow weights (percent widths are refused)
255  const gaugeRow = (name: string, l?: Limit) => {
256    if (!l) return <Box>{label(name)}<Text dimColor>—</Text></Box>
257    const p = Math.min(100, Math.max(0, l.percentUsed))
258    const fill = Math.round(p * 100)
259    return (
260      <Box columnGap={1}>
261        {label(name)}
262        <Box flexGrow={1} height={1}>
263          {fill > 0 ? <Box width={0} flexGrow={fill} height={1} backgroundColor={tone(p)} /> : null}
264          {fill < 10000 ? <Box width={0} flexGrow={10000 - fill} height={1} backgroundColor="subtle" /> : null}
265        </Box>
266        {cell(PCT, <Text bold color={tone(p)}>{Math.round(l.percentUsed) + '%'}</Text>)}
267        {cell(RESET, <Text dimColor>{'↻ ' + until(l.resetsAt, now)}</Text>)}
268      </Box>
269    )
270  }
271
272  const mine = (
273    <Box flexDirection="column" paddingX={1} borderStyle="round" borderColor="claude" borderDimColor width="100%" marginTop={hasAbove ? 1 : 0}>
274      <Box>
275        <Text bold color="claude">✻ </Text>
276        <Text bold>usage</Text>
277        <Text dimColor>{'  tokens (chat, all projects this week) · plan limits'}</Text>
278      </Box>
279      <Box flexDirection={stacked ? 'column' : 'row'} columnGap={3}>
280        <Box flexDirection="column" flexShrink={0}>
281          {tokenRow('chat', t?.chat)}
282          {tokenRow('all', t?.all)}
283        </Box>
284        <Box flexDirection="column" flexGrow={1}>
285          {gaugeRow('5h', l5)}
286          {gaugeRow('week', l7)}
287        </Box>
288      </Box>
289      {t ? null : (
290        <Text dimColor>
291          {err ? `tokens unavailable: ${err}` : l7?.resetsAt ? 'scanning transcripts…' : 'tokens need a plan with a weekly reset'}
292        </Text>
293      )}
294    </Box>
295  )
296  return mine
297}
298
299// One tree, one fixed order: the context card, a row of space, the usage card. Nothing the engine
300// hands back from next(e) goes in it: that placeholder is not allowed under a sized Box.
301export function hud(el: any, d: ViewData, s: Snap | null) {
302  const { Box } = el
303  return (
304    <Box flexDirection="column">
305      {s === null ? null : contextCard(el, s)}
306      {usageCard(el, d, s !== null)}
307    </Box>
308  )
309}
310
311export const register: Register = on => {
312  on('session.start', async ($, e, next) => {
313    void refreshUsage($)
314    $.clock.every(TTL, () => void refreshUsage($))
315    await refreshContext($)
316    return next(e)
317  })
318
319  on('session.measure', async ($, e, next) => {
320    if (e.rateLimits.length) await update($, limits, () => e.rateLimits)
321    if (e.changed.includes('context')) await refreshContext($)
322    return next(e)
323  })
324
325  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
326    if (e.props.hasSurvey) return next(e)
327    const data = {
328      limits: await read($, limits),
329      t: await read($, tokens),
330      err: await read($, error),
331      now: await $.clock.now(),
332      cols: e.props.bodyColumns,
333    }
334    return hud($.ui.resolve(e), data, await read($, snap))
335  })
336}
337