memory-bar

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionContextUsage, SessionRateLimit } from 'claude-code'
3
4import type { Limit, Memory } from '../types'
5
6const WIDTH = 20
7const LABEL = 9
8const ACCENT = '#D97757'
9const AMBER = '#E2A84B'
10const CRIMSON = '#E5484D'
11const memory = atom({ plugin: 'memory-bar', key: 'memory' } as const, null)
12const limits = atom({ plugin: 'memory-bar', key: 'limits' } as const, [])
13
14const toMemory = (context: SessionContextUsage): Memory | null =>
15  context.tokens === undefined
16    ? null
17    : {
18        tokens: context.tokens,
19        window: context.window,
20        percent: context.percent ?? Math.round((context.tokens / context.window) * 100),
21      }
22
23const toLimits = (rateLimits: readonly SessionRateLimit[]): Limit[] =>
24  rateLimits.map(limit => ({ kind: limit.kind, percent: limit.percentUsed, resetsAt: limit.resetsAt }))
25
26const NAMES: Record<string, string> = { five_hour: '5-hour', seven_day: 'weekly', spend_limit: 'spend' }
27
28const compact = (n: number) =>
29  n >= 1_000_000 ? `${(n / 1_000_000).toFixed(1)}M` : n >= 1000 ? `${Math.round(n / 1000)}k` : `${n}`
30
31// "2h 14m", "4d 3h", "12m"
32const until = (resetsAt: string | undefined, now: number) => {
33  const at = resetsAt === undefined ? NaN : Date.parse(resetsAt)
34  if (Number.isNaN(at)) return ''
35  const minutes = Math.max(0, Math.round((at - now) / 60_000))
36  const days = Math.floor(minutes / 1440)
37  const hours = Math.floor((minutes % 1440) / 60)
38  const text = days > 0 ? `${days}d ${hours}h` : hours > 0 ? `${hours}h ${minutes % 60}m` : `${minutes}m`
39  return `resets in ${text}`
40}
41
42// The icon fills like a gauge: ○ ◔ ◑ ◕ ●
43const icon = (percent: number) => '○◔◑◕●'[Math.max(0, Math.min(4, Math.floor(percent / 20)))]
44const toneOf = (percent: number) => (percent >= 85 ? CRIMSON : percent >= 60 ? AMBER : ACCENT)
45
46export const bar = (percent: number) => {
47  const filled = Math.max(0, Math.min(WIDTH, Math.round((percent / 100) * WIDTH)))
48  return { filled: '━'.repeat(filled), empty: '─'.repeat(WIDTH - filled) }
49}
50
51// Rate limits are the account's, but a session only learns them from its first reply. The latest
52// reading goes in the cross-session store, so every session shows it from the start and picks up
53// readings other sessions made.
54const READING = 'rateLimits'
55type Reading = { at: number; limits: Limit[] }
56
57async function saveReading($: EngineInterface, list: Limit[]) {
58  if (list.length > 0) {
59    const reading: Reading = { at: await $.clock.now(), limits: list }
60    await $.store.set(READING, reading)
61  }
62}
63
64async function loadReading($: EngineInterface) {
65  const reading = (await $.store.get(READING)) as Reading | undefined
66  if (reading !== undefined && reading.limits.length > 0) {
67    await update($, limits, () => reading.limits)
68  }
69}
70
71async function tick($: EngineInterface) {
72  await loadReading($)
73  // Keep the reset countdowns current while the session sits idle.
74  $.ui.invalidate('ui.render')
75}
76
77export const register: Register = on => {
78  on('session.start', async ($, e, next) => {
79    const result = await next(e)
80    const usage = await $.session.usage()
81    await update($, memory, () => toMemory(usage.context))
82    const live = toLimits(usage.rateLimits)
83    if (live.length > 0) {
84      await update($, limits, () => live)
85      await saveReading($, live)
86    } else {
87      await loadReading($)
88    }
89    $.clock.every(60_000, () => tick($))
90    return result
91  })
92
93  on('session.measure', async ($, e, next) => {
94    if (e.changed.includes('context')) {
95      await update($, memory, () => toMemory(e.context))
96    }
97    if (e.changed.includes('rateLimits')) {
98      const live = toLimits(e.rateLimits)
99      await update($, limits, () => live)
100      await saveReading($, live)
101    }
102    return next(e)
103  })
104
105  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
106    // Draw beneath whatever other bands (task-steps) drew, so they stack.
107    const above = await next(e)
108    if (e.props.hasSurvey) {
109      return above
110    }
111
112    const { Box, Text } = $.ui.resolve(e)
113    const value = await read($, memory)
114    const windows = await read($, limits)
115    const now = await $.clock.now()
116
117    const row = (label: string, percent: number, detail: string, hint?: string) => {
118      const { filled, empty } = bar(percent)
119      const tone = toneOf(percent)
120      return (
121        <Box>
122          <Text color={tone}>{`${icon(percent)} `}</Text>
123          <Text dimColor>{label.padEnd(LABEL)}</Text>
124          <Text color={tone}>{filled}</Text>
125          <Text dimColor>{empty}</Text>
126          <Text color={tone} bold>{`${Math.round(percent)}%`.padStart(6)}</Text>
127          <Text dimColor>{detail === '' ? '' : `  ${detail}`}</Text>
128          {hint === undefined ? null : <Text dimColor>{`  · ${hint}`}</Text>}
129        </Box>
130      )
131    }
132
133    return (
134      <Box flexDirection="column">
135        {above}
136        {value === null ? (
137          <Box>
138            <Text color={ACCENT}>{'○ '}</Text>
139            <Text dimColor>{`${'context'.padEnd(LABEL)}${'─'.repeat(WIDTH)}  waiting for first reply`}</Text>
140          </Box>
141        ) : (
142          row(
143            'context',
144            value.percent,
145            `${compact(value.tokens)}/${compact(value.window)}`,
146            value.percent >= 85 ? '/compact' : undefined,
147          )
148        )}
149        {windows.map(limit => {
150          // A saved reading may outlive its window: past its reset, the window starts over empty.
151          const isReset = limit.resetsAt !== undefined && Date.parse(limit.resetsAt) <= now
152          return isReset
153            ? row(NAMES[limit.kind] ?? limit.kind, 0, 'just reset')
154            : row(NAMES[limit.kind] ?? limit.kind, limit.percent, until(limit.resetsAt, now))
155        })}
156      </Box>
157    )
158  })
159}
160

1export type Memory = { tokens: number; window: number; percent: number }
2export type Limit = { kind: string; percent: number; resetsAt?: string }
3
4declare module 'claude-code' {
5  interface PluginState {
6    'memory-bar': { memory: Memory | null; limits: Limit[] }
7  }
8}
9