ssh-usage-band

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { Limit, Ram } from '../types'
5
6const limits = atom({ plugin: 'ssh-usage-band', key: 'limits' } as const, [])
7const now = atom({ plugin: 'ssh-usage-band', key: 'now' } as const, 0)
8const host = atom({ plugin: 'ssh-usage-band', key: 'host' } as const, '')
9const ram = atom({ plugin: 'ssh-usage-band', key: 'ram' } as const, null)
10// When the session last got a usage reading: each turn's measurement carries the
11// figures from that turn's last API response, so this is the session's last turn.
12const readAt = atom({ plugin: 'ssh-usage-band', key: 'readAt' } as const, 0)
13
14const LABELS: Record<string, string> = {
15  five_hour: '5h',
16  seven_day: '7d',
17  spend_limit: 'spend',
18}
19
20const BAR_WIDTH = 8
21const ORANGE = '#ff8700'
22
23export function bar(percent: number): string {
24  const filled = Math.round((Math.min(100, Math.max(0, percent)) / 100) * BAR_WIDTH)
25  return '█'.repeat(filled) + '░'.repeat(BAR_WIDTH - filled)
26}
27
28export function colorFor(percent: number): string {
29  if (percent >= 90) return 'red'
30  if (percent >= 80) return ORANGE
31  if (percent >= 70) return 'yellow'
32  return 'green'
33}
34
35function span(minutes: number): string {
36  const days = Math.floor(minutes / 1440)
37  const hours = Math.floor((minutes % 1440) / 60)
38  const mins = minutes % 60
39
40  if (days > 0) return `${days}d ${hours}h`
41  if (hours > 0) return `${hours}h ${mins}m`
42  return `${mins}m`
43}
44
45// A window whose reset time has passed: the session's reading of it predates the
46// reset, so its percentage is no longer true until the session's next turn.
47export function hasReset(resetsAt: string | undefined, nowMs: number): boolean {
48  if (!resetsAt || nowMs <= 0) return false
49  return Date.parse(resetsAt) <= nowMs
50}
51
52// A day or more out, whole days are enough and keep the row on one line.
53export function untilReset(resetsAt: string | undefined, nowMs: number): string | undefined {
54  if (!resetsAt) return undefined
55  const ms = Date.parse(resetsAt) - nowMs
56  if (Number.isNaN(ms) || ms <= 0) return undefined
57  const minutes = Math.ceil(ms / 60_000)
58  return minutes >= 1440 ? `↻ ${Math.floor(minutes / 1440)}d` : `↻ ${span(minutes)}`
59}
60
61// Readings younger than this are the session you're working in; no label.
62const AGE_SHOWN_FROM_MS = 5 * 60_000
63
64export function readingAge(readAtMs: number, nowMs: number): string | undefined {
65  if (readAtMs <= 0 || nowMs <= 0) return undefined
66  const ms = nowMs - readAtMs
67  if (ms < AGE_SHOWN_FROM_MS) return undefined
68  return `${span(Math.floor(ms / 60_000))} ago`
69}
70
71const RAM_POLL_MS = 10_000
72
73// Percent of MemTotal that MemAvailable is under. Under 30 % free is over 70 % used,
74// 20 % free is 80 % used, 10 % free is 90 % used: the same yellow, orange and red
75// steps as colorFor, so the RAM bar changes color where the usage bars do.
76const RAM_YELLOW_BELOW = 30
77const RAM_ORANGE_BELOW = 20
78const RAM_RED_BELOW = 10
79
80export function hostLabel(hostname: string, sshConnection: string | undefined): string {
81  // SSH_CONNECTION is "client_ip client_port server_ip server_port".
82  const serverIp = sshConnection?.trim().split(/\s+/)[2]
83  if (!serverIp) return hostname
84  return hostname ? `SSH ${hostname} (${serverIp})` : `SSH ${serverIp}`
85}
86
87export function parseMeminfo(text: string): Ram | null {
88  const kb = (field: string) => Number(new RegExp(`^${field}:\\s+(\\d+) kB$`, 'm').exec(text)?.[1])
89  const totalKb = kb('MemTotal')
90  const availableKb = kb('MemAvailable')
91  if (!(totalKb > 0) || Number.isNaN(availableKb)) return null
92  return { totalKb, availableKb }
93}
94
95export function ramColorFor(availablePercent: number): string {
96  if (availablePercent < RAM_RED_BELOW) return 'red'
97  if (availablePercent < RAM_ORANGE_BELOW) return ORANGE
98  if (availablePercent < RAM_YELLOW_BELOW) return 'yellow'
99  return 'green'
100}
101
102// The bar fills with RAM in use, like the usage bars; the color follows what's free.
103export function ramShown(r: Ram): { usedPercent: number; color: string } {
104  const availablePercent = (r.availableKb / r.totalKb) * 100
105  return {
106    usedPercent: 100 - availablePercent,
107    color: ramColorFor(availablePercent),
108  }
109}
110
111export const register: Register = on => {
112  on('session.start', async ($, e, next) => {
113    const hostname = await $.fs.read('/etc/hostname').then(
114      text => text.trim(),
115      () => '',
116    )
117    const ssh = await $.env.get('SSH_CONNECTION')
118    await update($, host, () => hostLabel(hostname, ssh))
119
120    // Write only when what the band draws changes, so polling doesn't redraw it.
121    let drawn = ''
122    const sampleRam = async () => {
123      let sample: Ram | null = null
124      try {
125        sample = parseMeminfo(await $.fs.read('/proc/meminfo'))
126      } catch {
127        sample = null
128      }
129      const shown = sample ? ramShown(sample) : null
130      const key = shown ? `${Math.round(shown.usedPercent)} ${shown.color}` : ''
131      if (key === drawn) return
132      drawn = key
133      await update($, ram, () => sample)
134    }
135    void sampleRam()
136    $.clock.every(RAM_POLL_MS, () => void sampleRam())
137
138    const usage = await $.session.usage()
139    await update($, limits, () => usage.rateLimits as Limit[])
140    await update($, now, () => 0)
141    if (usage.rateLimits.length > 0) {
142      const t = await $.clock.now()
143      await update($, readAt, () => t)
144    }
145
146    const tick = async () => {
147      const t = await $.clock.now()
148      await update($, now, () => t)
149    }
150    void tick()
151    $.clock.every(60_000, () => void tick())
152
153    return next(e)
154  })
155
156  on('session.measure', async ($, e, next) => {
157    // Every measurement carries the current figures, even when only the context
158    // moved, so it refreshes the reading's age; the bars redraw only on a change.
159    if (e.rateLimits.length > 0) {
160      const t = await $.clock.now()
161      if (e.changed.includes('rateLimits')) await update($, limits, () => e.rateLimits as Limit[])
162      await update($, readAt, () => t)
163      await update($, now, () => t)
164    }
165
166    return next(e)
167  })
168
169  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
170    const current = await read($, limits)
171    const hostText = await read($, host)
172    const mem = await read($, ram)
173    if (e.props.hasSurvey || (current.length === 0 && !hostText && !mem)) {
174      return next(e)
175    }
176
177    const nowMs = await read($, now)
178    const age = current.length > 0 ? readingAge(await read($, readAt), nowMs) : undefined
179    const { Box, Text } = $.ui.resolve(e)
180    const shown = mem ? ramShown(mem) : null
181
182    return (
183      <Box flexDirection="row" flexWrap="wrap" columnGap={2}>
184        {hostText ? <Text dimColor>{hostText}</Text> : null}
185        {shown ? (
186          <Box key="ram" flexDirection="row" gap={1}>
187            <Text dimColor>RAM</Text>
188            <Text color={shown.color}>{bar(shown.usedPercent)}</Text>
189            <Text bold color={shown.color}>
190              {`${Math.round(shown.usedPercent)}%`}
191            </Text>
192          </Box>
193        ) : null}
194        {current.map(limit => {
195          const label = LABELS[limit.kind] ?? limit.kind
196          if (hasReset(limit.resetsAt, nowMs)) {
197            return (
198              <Box key={limit.kind} flexDirection="row" gap={1}>
199                <Text dimColor>{label}</Text>
200                <Text dimColor>{bar(0)}</Text>
201                <Text dimColor>reset</Text>
202              </Box>
203            )
204          }
205          const reset = nowMs > 0 ? untilReset(limit.resetsAt, nowMs) : undefined
206          return (
207            <Box key={limit.kind} flexDirection="row" gap={1}>
208              <Text dimColor>{label}</Text>
209              <Text color={colorFor(limit.percentUsed)}>{bar(limit.percentUsed)}</Text>
210              <Text bold color={colorFor(limit.percentUsed)}>
211                {`${Math.round(limit.percentUsed)}%`}
212              </Text>
213              {reset ? <Text dimColor>{reset}</Text> : null}
214            </Box>
215          )
216        })}
217        {age ? <Text dimColor>{age}</Text> : null}
218      </Box>
219    )
220  })
221}
222

1export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
2export type Ram = { totalKb: number; availableKb: number }
3
4declare module 'claude-code' {
5  interface PluginState {
6    'ssh-usage-band': { limits: Limit[]; now: number; host: string; ram: Ram | null; readAt: number }
7  }
8}
9