context-bar

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface as Engine, Register } from 'claude-code'
3
4import type { Segment, Snapshot } from '../types'
5
6const isOn = atom({ plugin: 'context-bar', key: 'isOn' } as const, true)
7const snapshot = atom({ plugin: 'context-bar', key: 'snapshot' } as const, null)
8
9const PALETTE = [
10  'rgb(136,136,136)',
11  'rgb(177,185,249)',
12  'rgb(22,163,74)',
13  'rgb(215,119,87)',
14  'rgb(255,193,7)',
15  'rgb(130,125,189)',
16  'rgb(80,200,220)',
17  'rgb(230,100,160)',
18]
19
20const compact = (n: number) => (n >= 1000 ? `${(n / 1000).toFixed(1)}k` : `${n}`)
21
22// Cells per segment: proportional to tokens, at least one for any non-empty segment,
23// summing to exactly `width` so the bar fills one line and never wraps.
24const cellsFor = (segments: Segment[], width: number, max: number) => {
25  const sum = segments.reduce((a, s) => a + Math.max(0, s.tokens), 0) || max || 1
26  const raw = segments.map(s => (s.tokens > 0 ? (s.tokens / sum) * width : 0))
27  const cells = raw.map((r, i) => ((segments[i]?.tokens ?? 0) > 0 ? Math.max(1, Math.floor(r)) : 0))
28  let diff = width - cells.reduce((a, b) => a + b, 0)
29
30  // Short: hand spare cells to the largest remainders.
31  const byRemainder = raw.map((r, i) => i).filter(i => (raw[i] ?? 0) > 0)
32  byRemainder.sort((a, b) => ((raw[b] ?? 0) % 1) - ((raw[a] ?? 0) % 1))
33  for (let k = 0; diff > 0 && byRemainder.length > 0; k++) {
34    const i = byRemainder[k % byRemainder.length] ?? 0
35    cells[i] = (cells[i] ?? 0) + 1
36    diff--
37  }
38
39  // Over (the one-cell minimums): take from the widest segments, and once every
40  // segment is down to one cell, drop the smallest ones from the bar.
41  while (diff < 0) {
42    let widest = -1
43    cells.forEach((c, i) => {
44      if (c > 1 && (widest < 0 || c > (cells[widest] ?? 0))) widest = i
45    })
46    if (widest < 0) {
47      let smallest = -1
48      cells.forEach((c, i) => {
49        const t = segments[i]?.tokens ?? 0
50        if (c > 0 && (smallest < 0 || t < (segments[smallest]?.tokens ?? 0))) smallest = i
51      })
52      if (smallest < 0) break
53      widest = smallest
54    }
55    cells[widest] = (cells[widest] ?? 0) - 1
56    diff++
57  }
58  return cells
59}
60
61async function refresh($: Engine) {
62  const { context } = await $.session.usage({ breakdown: 'summary' })
63  const b = context.breakdown
64  if (!b) return
65  const segments: Segment[] = b.categories
66    .filter(c => c.kind !== 'deferred')
67    .map(c => ({ name: c.name, tokens: c.tokens, kind: c.kind as Segment['kind'] }))
68  const snap: Snapshot = { segments, total: b.totalTokens, max: b.rawMaxTokens, percent: b.percentage }
69  await update($, snapshot, () => snap)
70}
71
72export const register: Register = on => {
73  on('session.start', async ($, e, next) => {
74    await $.command.register({ name: 'context-bar', description: 'Toggle the context-window bar above the prompt' })
75    await refresh($)
76    return next(e)
77  })
78
79  on('command.run', { command: 'context-bar' }, async $ => {
80    const now = !(await read($, isOn))
81    await update($, isOn, () => now)
82    if (now) await refresh($)
83    return { text: `Context bar ${now ? 'on' : 'off'}.` }
84  })
85
86  on('prompt.submit', async ($, e, next) => {
87    const ran = await next(e)
88    void refresh($)
89    return ran
90  }).catch(($, e, next) => next(e))
91
92  on('turn.complete', async ($, e, next) => {
93    await refresh($)
94    return next(e)
95  })
96
97  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
98    const snap = await read($, snapshot)
99    if (!(await read($, isOn)) || snap === null || e.props.hasSurvey) return next(e)
100
101    const { Box, Text } = $.ui.resolve(e)
102    // The band's own width, which excludes the engine's `[-]` at the right end.
103    const width = Math.max(10, e.props.bodyColumns)
104    // The usage figure sits at the bar's right end; the bar takes the rest of the line.
105    const usage = ` ${compact(snap.total)}/${compact(snap.max)} (${snap.percent}%)`
106    const cells = cellsFor(snap.segments, Math.max(1, width - usage.length), snap.max)
107    let used = 0
108
109    const colors = snap.segments.map(s => {
110      if (s.kind !== 'used') return undefined
111      return PALETTE[used++ % PALETTE.length]
112    })
113
114    const legend = snap.segments
115      .map((s, i) => ({ s, color: colors[i] }))
116      .filter(x => x.s.kind === 'used' && x.s.tokens > 0)
117
118    return (
119      <Box flexDirection="column">
120        <Box flexDirection="row" width={width}>
121          <Text wrap="truncate-end">
122            {snap.segments.map((s, i) => {
123              const n = cells[i] ?? 0
124              return n === 0 ? null : s.kind === 'used' ? (
125                <Text color={colors[i]}>{'█'.repeat(n)}</Text>
126              ) : (
127                <Text dimColor>{(s.kind === 'buffer' ? '▒' : '░').repeat(n)}</Text>
128              )
129            })}
130          </Text>
131          <Box flexShrink={0}>
132            <Text dimColor>{usage}</Text>
133          </Box>
134        </Box>
135        <Box flexDirection="row" flexWrap="wrap" columnGap={2} width={width}>
136          {legend.map(x => (
137            <Box key={x.s.name} flexShrink={0}>
138              <Text color={x.color}>
139                ■ <Text dimColor>{x.s.name} {compact(x.s.tokens)}</Text>
140              </Text>
141            </Box>
142          ))}
143        </Box>
144      </Box>
145    )
146  })
147}
148

1export type Segment = { name: string; tokens: number; kind: 'used' | 'free' | 'buffer' }
2export type Snapshot = { segments: Segment[]; total: number; max: number; percent: number }
3
4declare module 'claude-code' {
5  interface PluginState {
6    'context-bar': { isOn: boolean; snapshot: Snapshot | null }
7  }
8}
9