usage-band-for-claude

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Limit, Segment, Snapshot } from '../types'
5
6const snapshot = atom({ plugin: 'usage-band-for-claude', key: 'snapshot' } as const, null)
7// The effort the last main-loop request actually used, and whether the
8// current prompt opted into ultracode with the keyword.
9const stepEffort = atom({ plugin: 'usage-band-for-claude', key: 'stepEffort' } as const, null)
10const ultracodeTurn = atom({ plugin: 'usage-band-for-claude', key: 'ultracodeTurn' } as const, false)
11
12const EFFORT_COLORS: Record<string, string> = {
13  low: 'inactive',
14  medium: 'suggestion',
15  high: 'success',
16  xhigh: 'warning',
17  max: 'error',
18}
19
20const LIMIT_NAMES: Record<string, string> = {
21  five_hour: '5h',
22  seven_day: '7d',
23  spend_limit: 'spend',
24}
25
26const GAP: Segment = { text: '   ' }
27
28export function formatTokens(n: number): string {
29  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`
30  if (n >= 1_000) return `${Math.round(n / 1_000)}k`
31  return String(n)
32}
33
34export function formatSpan(ms: number): string {
35  const minutes = Math.max(0, Math.round(ms / 60_000))
36  const days = Math.floor(minutes / 1440)
37  const hours = Math.floor((minutes % 1440) / 60)
38  const mins = minutes % 60
39  if (days > 0) return `${days}d${hours}h`
40  if (hours > 0) return `${hours}h${String(mins).padStart(2, '0')}m`
41  return `${mins}m`
42}
43
44// Green while there's room, yellow when getting close, red near the limit.
45export function levelColor(percent: number): string {
46  if (percent >= 85) return 'error'
47  if (percent >= 60) return 'warning'
48  return 'success'
49}
50
51export function bar(percent: number, width: number, full: string, empty: string): Segment[] {
52  const clamped = Math.min(100, Math.max(0, percent))
53  const filled = Math.round((clamped / 100) * width)
54  return [
55    { text: full.repeat(filled), color: levelColor(percent) },
56    { text: empty.repeat(width - filled), color: 'inactive' },
57  ]
58}
59
60export function limitSegments(limit: Limit, now: number, width: number): Segment[] {
61  const name = LIMIT_NAMES[limit.kind] ?? limit.kind
62  const segments: Segment[] = [
63    { text: `${name} `, dim: true },
64    ...bar(limit.percentUsed, width, '▰', '▱'),
65    { text: ` ${Math.round(limit.percentUsed)}%`, color: levelColor(limit.percentUsed), bold: true },
66  ]
67  const resetsAt = limit.resetsAt ? Date.parse(limit.resetsAt) : NaN
68  if (!Number.isNaN(resetsAt)) {
69    segments.push({ text: ` ↻ ${formatSpan(resetsAt - now)}`, color: 'suggestion' })
70  }
71  return segments
72}
73
74export function lineSegments(s: Snapshot, columns: number): Segment[] {
75  // Narrow terminals get shorter bars so everything still fits on one row.
76  const isNarrow = columns < 110
77  const ctxWidth = isNarrow ? 10 : 20
78  const limitWidth = isNarrow ? 5 : 10
79
80  const percent = s.contextPercent ?? 0
81  const used = s.contextTokens ?? 0
82  const segments: Segment[] = [
83    { text: `◆ ${s.model.replace(/^claude-/, '')}`, color: 'claude', bold: true },
84  ]
85  if (s.effort !== undefined) {
86    segments.push(
87      { text: ' ' },
88      { text: `⚡${s.effort}`, color: EFFORT_COLORS[s.effort] ?? 'suggestion', bold: true },
89    )
90  }
91  if (s.ultracode === 'off') {
92    segments.push({ text: ' ' }, { text: 'ultracode off', dim: true })
93  } else {
94    const label = s.ultracode === 'turn' ? 'ULTRACODE (this prompt)' : 'ULTRACODE'
95    segments.push({ text: ' ' }, { text: ` ${label} `, color: 'inverseText', backgroundColor: 'planMode', bold: true })
96  }
97  segments.push(
98    GAP,
99    { text: 'ctx ', dim: true },
100    ...bar(percent, ctxWidth, '█', '░'),
101    { text: ` ${percent}%`, color: levelColor(percent), bold: true },
102    { text: ` ${formatTokens(used)}/${formatTokens(s.contextWindow)}`, dim: true },
103  )
104
105  for (const limit of s.limits) {
106    segments.push(GAP, ...limitSegments(limit, s.now, limitWidth))
107  }
108
109  segments.push(GAP, { text: `⏱ ${formatSpan(s.now - s.startedAt)}`, dim: true })
110  return segments.filter(seg => seg.text !== '')
111}
112
113async function refresh($: EngineInterface): Promise<void> {
114  const [usage, model, now, settings, lastEffort, isUltracodeTurn] = await Promise.all([
115    $.session.usage(),
116    $.session.model(),
117    $.clock.now(),
118    $.settings.read(),
119    read($, stepEffort),
120    read($, ultracodeTurn),
121  ])
122  const configuredEffort = typeof settings.effortLevel === 'string' ? settings.effortLevel : undefined
123  const next: Snapshot = {
124    model,
125    effort: lastEffort ?? configuredEffort,
126    ultracode: settings.ultracode === true ? 'on' : isUltracodeTurn ? 'turn' : 'off',
127    contextTokens: usage.context.tokens,
128    contextWindow: usage.context.window,
129    contextPercent: usage.context.percent,
130    limits: usage.rateLimits.map(l => ({ kind: l.kind, percentUsed: l.percentUsed, resetsAt: l.resetsAt })),
131    startedAt: usage.startedAt,
132    now,
133  }
134  await update($, snapshot, () => next)
135}
136
137export const register: Register = on => {
138  let tick: { cancel: () => void } | undefined
139
140  on('session.start', async ($, e, next) => {
141    const result = await next(e)
142    await refresh($)
143    tick?.cancel()
144    // Keep the reset countdowns current while idle.
145    tick = $.clock.every(60_000, () => refresh($))
146    return result
147  })
148
149  on('prompt.submit', async ($, e, next) => {
150    // Never let the band get in the way of sending a prompt.
151    await update($, ultracodeTurn, () => /\bultracode\b/i.test(e.text))
152      .then(() => refresh($))
153      .catch(() => undefined)
154    return next(e)
155  })
156
157  on('turn.step', async function* ($, e, next) {
158    if (e.agentId === undefined && e.effort !== undefined) {
159      const effort = String(e.effort)
160      await read($, stepEffort)
161        .then(async last => {
162          if (last === effort) return
163          await update($, stepEffort, () => effort)
164          await refresh($)
165        })
166        .catch(() => undefined)
167    }
168    return yield* next(e)
169  })
170
171  on('session.measure', async ($, e, next) => {
172    await refresh($)
173    return next(e)
174  })
175
176  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
177    const s = await read($, snapshot)
178    if (e.props.hasSurvey || s === null) {
179      return next(e)
180    }
181
182    const { Text } = $.ui.resolve(e)
183
184    return (
185      <Text wrap="truncate-end">
186        {lineSegments(s, e.props.bodyColumns).map(seg => (
187          <Text color={seg.color} backgroundColor={seg.backgroundColor} dimColor={seg.dim} weight={seg.bold ? 'bold' : undefined}>
188            {seg.text}
189          </Text>
190        ))}
191      </Text>
192    )
193  })
194}
195

1export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
2
3export type Segment = { text: string; color?: string; backgroundColor?: string; dim?: boolean; bold?: boolean }
4
5export type Snapshot = {
6  model: string
7  effort?: string
8  ultracode: 'on' | 'turn' | 'off'
9  contextTokens?: number
10  contextWindow: number
11  contextPercent?: number
12  limits: Limit[]
13  startedAt: number
14  now: number
15}
16
17declare module 'claude-code' {
18  interface PluginState {
19    'usage-band-for-claude': { snapshot: Snapshot | null; stepEffort: string | null; ultracodeTurn: boolean }
20  }
21}
22