usage-band

1/* @jsx h */
2import type { EngineInterface, Register, SessionMeasureInput } from 'claude-code'
3
4// The limit windows and this session's context fill on one line above the prompt.
5//
6// session.measure pushes what the last API response already reported, after each turn and whenever a
7// limit window moves a whole point, so it costs no request and no tokens and the render hook, which
8// runs on every keystroke, never reads anything. Before the first response there is no reading at all.
9//
10// The nudge exists because long windows are where the money goes: /usage says most of the spend
11// happens above 150k context, and a dim band is easy to stop seeing, so each new tenth of the
12// window past the threshold also toasts once.
13
14const NUDGE_AT = 55
15const SHOWN_KEY = 'usage-band:shown'
16const SGD_KEY = 'usage-band:sgd'
17let SGD_RATE = 1.3
18
19type Reading = {
20  fiveHour?: number
21  fiveHourResetsAt?: number
22  sevenDay?: number
23  ctxPercent?: number
24  ctxTokens?: number
25  usd?: number
26}
27
28type Delta = { in: number; out: number; cache: number }
29
30let reading: Reading | undefined
31let delta: Delta | undefined
32let account: string | undefined
33let model: string | undefined
34let project: string | undefined
35let shown = true
36let sgdShown = true
37let nudgedAt = 0
38
39// Which login the figures belong to. $.session.usage() does not carry it, and the same terminal
40// switches accounts mid-session with /login, so it is re-read each turn from the file the CLI
41// writes the logged-in account into.
42const readAccount = async ($: EngineInterface): Promise<void> => {
43  const home = await $.env.get('HOME').catch(() => undefined)
44  if (!home) return
45  const raw = await $.fs.read(`${home}/.claude.json`).catch(() => undefined)
46  if (typeof raw !== 'string') return
47  const found = accountIn(raw)
48  if (found) account = found.email
49}
50
51// The CLI rewrites this file, so a half-written one is unreadable rather than a logout: undefined
52// keeps the account as it was, a parsed file without an email is a logout.
53export const accountIn = (raw: string): { email: string | undefined } | undefined => {
54  let parsed: unknown
55  try {
56    parsed = JSON.parse(raw)
57  } catch {
58    return undefined
59  }
60  const email = (parsed as { oauthAccount?: { emailAddress?: unknown } } | null)?.oauthAccount?.emailAddress
61  return { email: typeof email === 'string' ? email : undefined }
62}
63
64// The model can change mid-session with /model and the folder is fixed, but both are one cheap
65// in-process read, so they go with the account each turn rather than keeping their own state.
66const readContext = async ($: EngineInterface): Promise<void> => {
67  model = await $.session.model().catch(() => undefined)
68  const cwd = await $.session.cwd().catch(() => undefined)
69  project = cwd ? cwd.split('/').filter(Boolean).at(-1) : undefined
70}
71
72type Usage = Pick<SessionMeasureInput, 'context' | 'rateLimits' | 'cost'>
73
74const take = (u: Usage): void => {
75  const window = (kind: string) => u.rateLimits.find((l) => l.kind === kind)
76  const resetsAt = Date.parse(window('five_hour')?.resetsAt ?? '')
77  reading = {
78    fiveHour: window('five_hour')?.percentUsed,
79    fiveHourResetsAt: Number.isFinite(resetsAt) ? resetsAt : undefined,
80    sevenDay: window('seven_day')?.percentUsed,
81    ctxPercent: u.context.percent,
82    ctxTokens: u.context.tokens,
83    usd: u.cost?.usd,
84  }
85}
86
87const read = async ($: EngineInterface): Promise<void> => {
88  const u = await $.session.usage().catch(() => undefined)
89  if (u) take(u)
90}
91
92const nudge = ($: EngineInterface): void => {
93  const percent = reading?.ctxPercent
94  if (percent === undefined) return
95  // A clear or a compact drops the fill, and the nudge arms itself again from there.
96  if (percent < nudgedAt) nudgedAt = 0
97  if (percent >= NUDGE_AT && percent >= nudgedAt + 10) {
98    nudgedAt = Math.floor(percent / 10) * 10
99    $.ui.toast(`context ${Math.round(percent)}% · /clear if the next thing is a new task, /compact to keep going`, {
100      timeoutMs: 8000,
101    })
102  }
103}
104
105type Part = { text: string; dim?: boolean; bold?: boolean; color?: string }
106type Segment = { parts: Part[]; drop: number }
107
108const family = (id: string) => /opus|sonnet|fable|haiku/.exec(id)?.[0] ?? id.replace(/^claude-/, '').replace(/\[1m\]$/, '')
109
110const GAP = '  '
111const width = (list: Segment[]) => list.reduce((n, s) => n + s.parts.reduce((m, p) => m + p.text.length, 0), 0) + GAP.length * Math.max(0, list.length - 1)
112
113// The whole line as one string, measured against the band's columns and shed by priority until it
114// fits: the turn delta, project, account, then the model to its family word, then cost, then the
115// nudge. The windows and the context fill are never shed. The "usage" label goes under 60 columns.
116export const layout = (state: {
117  account?: string
118  project?: string
119  model?: string
120  reading: Reading
121  delta?: Delta
122  now: number
123  columns: number
124}): Part[] => {
125  const { fiveHour, fiveHourResetsAt, sevenDay, ctxPercent, ctxTokens, usd } = state.reading
126  const left = fiveHourResetsAt !== undefined ? countdown(fiveHourResetsAt - state.now) : undefined
127  const segments: Segment[] = []
128  if (state.columns >= 60) segments.push({ parts: [{ text: 'usage', dim: true }], drop: 0 })
129  if (state.account !== undefined) segments.push({ parts: [{ text: state.account, color: 'cyan' }], drop: 3 })
130  if (state.project !== undefined) segments.push({ parts: [{ text: state.project, color: 'magenta' }], drop: 2 })
131  if (state.model !== undefined) segments.push({ parts: [{ text: state.model, dim: true }], drop: 4 })
132  if (fiveHour !== undefined) {
133    segments.push({
134      parts: [{ text: '5h ', dim: true }, { text: `${Math.round(fiveHour)}%`, bold: true, color: heat(fiveHour) }, ...(left ? [{ text: ` ↻${left}`, dim: true }] : [])],
135      drop: 0,
136    })
137  }
138  if (sevenDay !== undefined) {
139    segments.push({ parts: [{ text: '7d ', dim: true }, { text: `${Math.round(sevenDay)}%`, bold: true, color: heat(sevenDay) }], drop: 0 })
140  }
141  if (ctxPercent !== undefined) {
142    segments.push({
143      parts: [
144        { text: 'ctx ', dim: true },
145        { text: `${Math.round(ctxPercent)}%`, bold: true, color: heat(ctxPercent) },
146        ...(ctxTokens !== undefined ? [{ text: ` ${tokens(ctxTokens)}`, dim: true }] : []),
147      ],
148      drop: 0,
149    })
150  }
151  if (usd !== undefined) segments.push({ parts: [{ text: `$${usd.toFixed(2)}`, dim: true }], drop: 5 })
152  if (state.delta) {
153    const d = state.delta
154    segments.push({ parts: [{ text: `last ${tokens(d.in)} in · ${tokens(d.out)} out · ${tokens(d.cache)} cache`, dim: true }], drop: 1 })
155  }
156  if (ctxPercent !== undefined && ctxPercent >= NUDGE_AT) segments.push({ parts: [{ text: '· /clear on a new task', color: 'yellow' }], drop: 6 })
157
158  let kept = segments
159  let shortened = false
160  while (width(kept) > state.columns) {
161    const droppable = kept.filter((s) => s.drop > 0)
162    if (!droppable.length) break
163    const gone = droppable.reduce((a, b) => (a.drop < b.drop ? a : b))
164    if (gone.drop === 4 && !shortened && state.model !== undefined) {
165      shortened = true
166      kept = kept.map((s) => (s === gone ? { parts: [{ text: family(state.model ?? ''), dim: true }], drop: 4 } : s))
167      continue
168    }
169    kept = kept.filter((s) => s !== gone)
170  }
171  return kept.flatMap((s, i) => (i ? [{ text: GAP, dim: true }, ...s.parts] : s.parts))
172}
173
174export const countdown = (ms: number) => {
175  const m = Math.max(0, Math.ceil(ms / 60000))
176  return m >= 60 ? `${Math.floor(m / 60)}h${String(m % 60).padStart(2, '0')}m` : `${m}m`
177}
178
179const heat = (percent: number) => (percent >= 75 ? 'red' : percent >= 50 ? 'yellow' : 'green')
180
181const tokens = (n: number) => (n >= 1000 ? `${Math.round(n / 1000)}k` : `${n}`)
182
183// The first thing anyone does when a mod misbehaves is try to turn it off. `CLAUDE_MODS_DISABLE=all`,
184// or a comma list naming this mod, makes every hook here a pass-through and registers no command.
185const MOD = 'usage-band'
186let disabled = false
187const readDisabled = async ($: EngineInterface): Promise<boolean> => {
188  const raw = (await $.env.get('CLAUDE_MODS_DISABLE').catch(() => undefined)) ?? ''
189  disabled = raw
190    .split(',')
191    .map(v => v.trim())
192    .some(v => v === 'all' || v === MOD)
193  return disabled
194}
195
196export const register: Register = (on, options) => {
197  const rate = Number(options.sgdRate)
198  if (Number.isFinite(rate) && rate > 0) SGD_RATE = rate
199
200  on('session.start', async ($, e, next) => {
201    const r = await next(e)
202    if (await readDisabled($)) return r
203    if ((await $.store.get(SHOWN_KEY).catch(() => undefined)) === false) shown = false
204    if ((await $.store.get(SGD_KEY).catch(() => undefined)) === false) sgdShown = false
205    await readAccount($)
206    await readContext($)
207    await $.command
208      .register({
209        name: 'usage-band',
210        description: 'The usage line above the prompt: limit windows, context fill, cost (usage-band)',
211        argumentHint: '[on | off | sgd on | sgd off]',
212        immediate: true,
213      })
214      .catch((err) => $.ui.log(`usage-band: /usage-band not registered: ${err}`))
215    return r
216  })
217
218  on('command.run', { command: 'usage-band' }, async ($, e) => {
219    const arg = e.args.trim().toLowerCase()
220    if (arg === 'sgd on' || arg === 'sgd off') {
221      sgdShown = arg === 'sgd on'
222      await $.store.set(SGD_KEY, sgdShown).catch(() => undefined)
223      $.ui.invalidate('ui.render')
224      return { text: sgdShown ? `S$ footer label on (rate ${SGD_RATE}, set sgdRate in /config)` : 'S$ footer label off' }
225    }
226    if (arg === 'off') {
227      shown = false
228      await $.store.set(SHOWN_KEY, false).catch(() => undefined)
229      $.ui.invalidate('ui.render')
230      return { text: 'usage band off' }
231    }
232    if (arg === 'on' || arg === '') {
233      shown = true
234      await $.store.set(SHOWN_KEY, true).catch(() => undefined)
235      await read($)
236      $.ui.invalidate('ui.render')
237      return { text: 'usage band on' }
238    }
239    return { text: `usage-band: no such argument "${arg}" — use on, off, sgd on or sgd off` }
240  })
241
242  on('session.measure', async ($, e, next) => {
243    if (disabled) return next(e)
244    const r = await next(e)
245    if (!shown) return r
246    take(e)
247    nudge($)
248    $.ui.invalidate('ui.render')
249    return r
250  })
251
252  // The last turn's own token counts come with the turn, not the measurement; the account, model and
253  // folder are re-read here because a turn is when any of them can have changed.
254  on('turn.complete', async ($, e, next) => {
255    if (disabled) return next(e)
256    const r = await next(e)
257    if (!shown || e.agentId) return r
258    if (e.usage) {
259      delta = {
260        in: e.usage.input_tokens,
261        out: e.usage.output_tokens,
262        cache: e.usage.cache_read_input_tokens + e.usage.cache_creation_input_tokens,
263      }
264    }
265    await readAccount($)
266    await readContext($)
267    $.ui.invalidate('ui.render')
268    return r
269  })
270
271  // The session's cost in dollars that mean something here, as a footer mode label beside `focus`.
272  on('ui.render', { component: 'SessionMode' }, async ($, e, next) => {
273    if (disabled || !shown || !sgdShown || reading?.usd === undefined) return next(e)
274    const label = `S$${(reading.usd * SGD_RATE).toFixed(2)}`
275    return next({ ...e, props: { ...e.props, modes: [...e.props.modes, label] } })
276  })
277
278  // A render hook that throws unmounts the module and takes the whole mod with it, so a bad frame
279  // falls back to the band as it was rather than costing the session its usage line.
280  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
281    if (disabled) return next(e)
282    if (!shown || e.props.hasSurvey || e.surface !== 'terminal') return next(e)
283    try {
284      const { Box, Text } = await $.ui.resolve(e)
285      const rest = await next(e)
286
287      if (!reading) {
288        return (
289          <Box flexDirection="column">
290            <Text dimColor wrap="truncate">{account ? `usage · ${account} · no reading yet` : 'usage · no reading yet'}</Text>
291            {rest}
292          </Box>
293        )
294      }
295
296      const now = await $.clock.now()
297      const parts = layout({ account, project, model, reading, delta, now, columns: e.props.bodyColumns || 80 })
298      return (
299        <Box flexDirection="column">
300          <Text wrap="truncate">
301            {parts.map((p, i) => (
302              <Text key={String(i)} dimColor={p.dim} bold={p.bold} color={p.color}>
303                {p.text}
304              </Text>
305            ))}
306          </Text>
307          {rest}
308        </Box>
309      )
310    } catch (err) {
311      $.ui.log(`usage-band: render failed: ${err}`)
312      return next(e)
313    }
314  })
315}
316