context-band

1import { atom, read, update } from 'claude-code'
2import type {
3  EngineInterface,
4  PluginOptions,
5  Register,
6  SessionContextUsage,
7  SessionCost,
8  SessionRateLimit,
9} from 'claude-code'
10
11import type { Figures } from '../types'
12
13const figures = atom({ plugin: 'context-band', key: 'figures' } as const, null)
14
15const LABELS: Record<string, string> = { five_hour: 'Session', seven_day: 'Weekly' }
16const USAGE_URL = 'https://api.anthropic.com/api/oauth/usage'
17const MINUTE = 60_000
18const MODELS_EVERY = 5 * MINUTE
19const MODELS_RETRY = MINUTE
20const MODELS_SOON = 10_000
21
22// Every part of the line, in the order it is drawn. Each is a boolean field of
23// the manifest's userConfig: a row in /config, and a screen at install time.
24const PARTS = [
25  'model', 'folder', 'branch',
26  'context', 'tokens', 'session', 'weekly', 'resets', 'models',
27  'cost', 'turns', 'duration', 'clock',
28] as const
29type Part = (typeof PARTS)[number]
30type Show = Record<Part, boolean>
31
32function showing(options: PluginOptions): Show {
33  const show = {} as Show
34  for (const p of PARTS) show[p] = options[p] === true
35  return show
36}
37
38// The line's colors: the base one, and the two a figure turns to at and
39// above the warn and high percentages.
40type Style = { color: string; warnAt: number; warnColor: string; highAt: number; highColor: string }
41
42function styling(options: PluginOptions): Style {
43  return {
44    color: String(options.color),
45    warnAt: Number(options.warnAt),
46    warnColor: String(options.warnColor),
47    highAt: Number(options.highAt),
48    highColor: String(options.highColor),
49  }
50}
51
52function tone(percent: number | null, style: Style): string {
53  if (percent === null) return style.color
54  if (percent >= style.highAt) return style.highColor
55  if (percent >= style.warnAt) return style.warnColor
56  return style.color
57}
58
59type UsageBody = {
60  limits?: { percent: number; resets_at?: string | null; scope: { model: { display_name: string } | null } | null }[]
61}
62
63// Per-model weekly windows come only from the OAuth usage endpoint, as the
64// `limits[]` entries with a model scope; the response headers that feed
65// `rateLimits` carry five_hour and seven_day alone. Null when the request
66// failed (a 429 among them): the caller keeps the reading it has.
67async function modelWindows($: EngineInterface): Promise<Figures['models'] | null> {
68  try {
69    const auth = await $.session.authorize()
70    if (auth === null) return []
71    const r = await $.http.fetch(USAGE_URL, { auth: auth.handle })
72    if (!r.ok) return null
73    const body = JSON.parse(r.text) as UsageBody
74    return (body.limits ?? []).flatMap(l =>
75      l.scope?.model ? [{ name: l.scope.model.display_name, percentUsed: l.percent, resetsAt: l.resets_at ?? null }] : [],
76    )
77  } catch {
78    return null
79  }
80}
81
82// The endpoint is asked again MODELS_EVERY after a reading, MODELS_RETRY
83// after a failed request; in between the last reading stands. A failed
84// request is also tried once more MODELS_SOON later, off the clock, so a
85// reload that collides with the engine's own request (a 429) shows the
86// entry within seconds instead of after the next turn.
87let modelsDueAt = -Infinity
88let models: Figures['models'] = []
89let soonRetry = false
90
91async function modelWindowsCached($: EngineInterface, now: number): Promise<Figures['models']> {
92  if (now >= modelsDueAt) {
93    const fresh = await modelWindows($)
94    if (fresh !== null) models = fresh
95    modelsDueAt = now + (fresh === null ? MODELS_RETRY : MODELS_EVERY)
96    if (fresh === null && !soonRetry) {
97      soonRetry = true
98      $.clock.after(MODELS_SOON, async () => {
99        soonRetry = false
100        const again = await modelWindows($)
101        if (again === null) return
102        models = again
103        modelsDueAt = (await $.clock.now()) + MODELS_EVERY
104        await update($, figures, f => (f === null ? f : { ...f, models: again }))
105      })
106    }
107  }
108  return models
109}
110
111function parentOf(dir: string): string | null {
112  const parent = dir.replace(/[\\/][^\\/]*$/, '')
113  return parent === dir || parent === '' ? null : parent
114}
115
116function basename(dir: string): string {
117  const m = /[^\\/]+$/.exec(dir)
118  return m ? m[0] : dir
119}
120
121// The branch of the repository `cwd` is in: `.git/HEAD` of the nearest `.git`
122// upward, following a worktree's `gitdir:` file. Null outside a repository.
123async function gitBranch($: EngineInterface, cwd: string): Promise<string | null> {
124  try {
125    for (let dir: string | null = cwd; dir !== null; dir = parentOf(dir)) {
126      const dotGit = `${dir}/.git`
127      if (!(await $.fs.exists(dotGit))) continue
128      let gitDir = dotGit
129      if ((await $.fs.stat(dotGit)).kind === 'file') {
130        const m = /^gitdir:\s*(.+)$/m.exec(await $.fs.read(dotGit))
131        if (!m) return null
132        const target = m[1]!.trim()
133        gitDir = /^([a-zA-Z]:)?[\\/]/.test(target) ? target : `${dir}/${target}`
134      }
135      const head = (await $.fs.read(`${gitDir}/HEAD`)).trim()
136      const ref = /^ref:\s*refs\/heads\/(.+)$/.exec(head)
137      return ref ? ref[1]! : head.slice(0, 7)
138    }
139  } catch {
140    // unreadable repository: no branch
141  }
142  return null
143}
144
145type Reading = { context: SessionContextUsage; rateLimits: SessionRateLimit[]; cost?: SessionCost }
146
147async function measure($: EngineInterface, show: Show, u: Reading) {
148  const cwd = show.folder || show.branch ? await $.session.cwd() : null
149  const now = await $.clock.now()
150  const next: Figures = {
151    now,
152    startedAt: show.duration ? (await $.session.usage()).startedAt : null,
153    context: {
154      percent: u.context.percent ?? null,
155      tokens: u.context.tokens ?? null,
156      window: u.context.window ?? null,
157    },
158    limits: u.rateLimits.map(r => ({ kind: r.kind, percentUsed: r.percentUsed, resetsAt: r.resetsAt ?? null })),
159    models: show.models ? await modelWindowsCached($, now) : [],
160    costUsd: u.cost?.usd ?? null,
161    turns: show.turns ? await $.session.turns() : null,
162    model: show.model ? await $.session.model() : null,
163    folder: cwd !== null && show.folder ? basename(cwd) : null,
164    branch: cwd !== null && show.branch ? await gitBranch($, cwd) : null,
165  }
166  await update($, figures, () => next)
167}
168
169function span(ms: number): string {
170  const mins = Math.max(0, Math.round(ms / MINUTE))
171  if (mins < 60) return `${mins}m`
172  const h = Math.floor(mins / 60)
173  const m = mins % 60
174  if (h < 24) return m ? `${h}hr ${m}m` : `${h}hr`
175  const d = Math.floor(h / 24)
176  const hh = h % 24
177  return hh ? `${d}d ${hh}hr` : `${d}d`
178}
179
180function until(resetsAt: string | null, now: number): string | null {
181  if (resetsAt === null) return null
182  const ms = Date.parse(resetsAt) - now
183  return Number.isFinite(ms) && ms > 0 ? span(ms) : null
184}
185
186function tokens(n: number): string {
187  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1).replace(/\.0$/, '')}M`
188  if (n >= 1_000) return `${Math.round(n / 1_000)}k`
189  return `${n}`
190}
191
192function clock(now: number): string {
193  const d = new Date(now)
194  return `${String(d.getHours()).padStart(2, '0')}:${String(d.getMinutes()).padStart(2, '0')}`
195}
196
197type Piece = { text: string; percent: number | null }
198
199function line(f: Figures, show: Show): Piece[] {
200  const parts: Piece[] = []
201  const plain = (text: string) => parts.push({ text, percent: null })
202  if (show.model && f.model !== null) plain(f.model)
203  if (show.folder && f.folder !== null) plain(show.branch && f.branch !== null ? `${f.folder} (${f.branch})` : f.folder)
204  else if (show.branch && f.branch !== null) plain(f.branch)
205  if (show.context) {
206    const pct = f.context.percent === null ? '--' : `${f.context.percent}%`
207    const size = show.tokens && f.context.tokens !== null && f.context.window !== null
208      ? ` (${tokens(f.context.tokens)}/${tokens(f.context.window)})`
209      : ''
210    parts.push({ text: `Context ${pct}${size}`, percent: f.context.percent })
211  }
212  for (const l of f.limits) {
213    if ((l.kind === 'five_hour' && !show.session) || (l.kind === 'seven_day' && !show.weekly)) continue
214    const reset = show.resets ? until(l.resetsAt, f.now) : null
215    parts.push({ text: `${LABELS[l.kind] ?? l.kind} ${l.percentUsed}%${reset === null ? '' : ` (${reset})`}`, percent: l.percentUsed })
216  }
217  if (show.models) {
218    for (const m of f.models) {
219      const reset = show.resets ? until(m.resetsAt, f.now) : null
220      parts.push({ text: `${m.name} ${m.percentUsed}%${reset === null ? '' : ` (${reset})`}`, percent: m.percentUsed })
221    }
222  }
223  if (show.cost && f.costUsd !== null) plain(`$${f.costUsd.toFixed(2)}`)
224  if (show.turns && f.turns !== null) plain(`${f.turns} ${f.turns === 1 ? 'turn' : 'turns'}`)
225  if (show.duration && f.startedAt !== null) plain(`Elapsed ${span(f.now - f.startedAt)}`)
226  if (show.clock) plain(clock(f.now))
227  return parts
228}
229
230export const register: Register = (on, options) => {
231  const show = showing(options)
232  const style = styling(options)
233
234  on('session.start', async ($, e, next) => {
235    await measure($, show, await $.session.usage())
236    if (show.clock || show.duration) {
237      $.clock.every(MINUTE, async () => {
238        const now = await $.clock.now()
239        await update($, figures, f => (f === null ? f : { ...f, now }))
240      })
241    }
242    return next(e)
243  })
244
245  on('session.measure', async ($, e, next) => {
246    await measure($, show, e)
247    return next(e)
248  })
249
250  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
251    const f = await read($, figures)
252    if (f === null || e.props.hasSurvey) return next(e)
253    const parts = line(f, show)
254    if (parts.length === 0) return next(e)
255
256    const { Box, Text } = $.ui.resolve(e)
257    return (
258      <Box>
259        <Text color={style.color}>
260          {parts.map((p, i) => [i === 0 ? '' : ' · ', <Text color={tone(p.percent, style)}>{p.text}</Text>])}
261        </Text>
262      </Box>
263    )
264  })
265}
266

1export type Figures = {
2  /** When the figures were read, in epoch milliseconds; moved on by the minute tick. */
3  now: number
4  /** When the session started, in epoch milliseconds; null unless the duration is shown. */
5  startedAt: number | null
6  context: { percent: number | null; tokens: number | null; window: number | null }
7  limits: { kind: string; percentUsed: number; resetsAt: string | null }[]
8  models: { name: string; percentUsed: number; resetsAt: string | null }[]
9  costUsd: number | null
10  turns: number | null
11  model: string | null
12  folder: string | null
13  branch: string | null
14}
15
16declare module 'claude-code' {
17  interface PluginState {
18    'context-band': { figures: Figures | null }
19  }
20}
21