context-bar

1// ABOUTME: context-bar mod: a stacked bar above the prompt showing the context window by /context category.
2// ABOUTME: Refreshes after each turn and compaction; /context-bar toggles it on and off.
3
4import { atom, read, update } from 'claude-code'
5import type { EngineInterface, Register } from 'claude-code'
6
7import { cellWidths, formatTokens, toSnapshot } from './bar'
8
9const COMMAND = 'context-bar'
10const snapshot = atom({ plugin: 'context-bar', key: 'snapshot' } as const, null)
11const isShown = atom({ plugin: 'context-bar', key: 'isShown' } as const, false)
12
13const refresh = async ($: EngineInterface) => {
14  const usage = await $.session.usage({ breakdown: 'summary' })
15  const breakdown = usage.context.breakdown
16  if (breakdown) await update($, snapshot, () => toSnapshot(breakdown))
17}
18
19export const register: Register = on => {
20  on('session.start', async ($, e, next) => {
21    await $.command.register({
22      name: COMMAND,
23      description: 'Toggle the context window bar above the prompt',
24    })
25    const result = await next(e)
26    await refresh($)
27
28    return result
29  })
30
31  on('turn.complete', async ($, e, next) => {
32    const result = await next(e)
33    await refresh($)
34
35    return result
36  })
37
38  on('session.compact', async ($, e, next) => {
39    const result = await next(e)
40    await refresh($)
41
42    return result
43  })
44
45  on('command.run', { command: COMMAND }, async $ => {
46    const shown = await update($, isShown, value => !value)
47    if (shown) await refresh($)
48
49    return { text: shown ? 'Context bar shown.' : 'Context bar hidden.' }
50  })
51
52  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
53    const snap = await read($, snapshot)
54    if (e.props.hasSurvey || !snap || !(await read($, isShown))) {
55      return next(e)
56    }
57
58    const { Box, Text } = $.ui.resolve(e)
59    const label = ` ${formatTokens(snap.usedTokens)}/${formatTokens(snap.maxTokens)} (${snap.percentage}%)`
60    // The band draws its [-] collapse marker at the right edge of the first row.
61    const widths = cellWidths(snap.segments, Math.max(10, e.props.bodyColumns - label.length - 5))
62    const used = snap.segments.filter(s => s.kind === 'used')
63
64    return (
65      <Box flexDirection="column">
66        <Box key="bar">
67          {snap.segments.map((s, i) => (
68            <Text color={s.color} dimColor={s.kind !== 'used'}>
69              {(s.kind === 'used' ? '█' : s.kind === 'buffer' ? '▒' : '░').repeat(widths[i] ?? 0)}
70            </Text>
71          ))}
72          <Text dimColor>{label}</Text>
73        </Box>
74        <Box key="legend" flexWrap="wrap" columnGap={2}>
75          {used.map(s => (
76            <Text>
77              <Text color={s.color}>■</Text>
78              <Text dimColor> {s.name} {formatTokens(s.tokens)}</Text>
79            </Text>
80          ))}
81        </Box>
82      </Box>
83    )
84  })
85}
86

1// ABOUTME: Pure helpers that turn a /context breakdown into bar segments and cell widths.
2// ABOUTME: Kept free of the engine so tests can exercise the arithmetic directly.
3
4import type { SessionContextBreakdown } from 'claude-code'
5
6import type { Segment, Snapshot } from '../types'
7
8// The four groups the bar draws, in order. System, Tools and Memory are the
9// same in every session and draw gray; Messages takes its color from how full
10// the window is. Colors are ANSI names, so the terminal's theme decides them.
11const GROUPS = [
12  { name: 'System', test: /system/i },
13  { name: 'Tools', test: /mcp|agent|skill|plugin|command/i },
14  { name: 'Memory', test: /memory/i },
15  { name: 'Messages', test: /message/i },
16] as const
17
18// The same thresholds as claude-statusline.sh's default CLAUDE_CONTEXT_GRADIENT.
19export const fillColor = (percentage: number): string =>
20  percentage <= 40 ? 'green' : percentage <= 70 ? 'yellow' : 'red'
21
22// Rows matching no group (new categories the engine adds) count as System.
23const groupOf = (name: string) =>
24  GROUPS.slice(1).find(g => g.test.test(name)) ?? GROUPS[0]
25
26export const toSnapshot = (breakdown: SessionContextBreakdown): Snapshot => {
27  const used = new Map<string, number>()
28  let free = 0
29  let buffer = 0
30  for (const row of breakdown.categories) {
31    if (row.kind === 'used') {
32      const group = groupOf(row.name).name
33      used.set(group, (used.get(group) ?? 0) + row.tokens)
34    } else if (row.kind === 'free') {
35      free += row.tokens
36    } else if (row.kind === 'buffer') {
37      buffer += row.tokens
38    }
39  }
40
41  const segments: Segment[] = GROUPS.filter(g => (used.get(g.name) ?? 0) > 0).map(g => ({
42    name: g.name,
43    color: g.name === 'Messages' ? fillColor(breakdown.percentage) : 'gray',
44    tokens: used.get(g.name) ?? 0,
45    kind: 'used',
46  }))
47  if (free > 0) segments.push({ name: 'Free space', color: 'gray', tokens: free, kind: 'free' })
48  if (buffer > 0) segments.push({ name: 'Autocompact buffer', color: 'gray', tokens: buffer, kind: 'buffer' })
49
50  return {
51    segments,
52    usedTokens: breakdown.totalTokens,
53    maxTokens: breakdown.rawMaxTokens,
54    percentage: breakdown.percentage,
55  }
56}
57
58// Largest-remainder split of `width` cells across segments by token share.
59// Every used segment gets at least one cell, taken from the widest segment,
60// so a small group like Memory stays visible on a 1M window.
61export const cellWidths = (segments: readonly Segment[], width: number): number[] => {
62  const total = segments.reduce((sum, s) => sum + s.tokens, 0)
63  if (total <= 0 || width <= 0) return segments.map(() => 0)
64
65  const exact = segments.map(s => (s.tokens / total) * width)
66  const cells = exact.map(Math.floor)
67  let left = width - cells.reduce((sum, n) => sum + n, 0)
68  const byRemainder = exact
69    .map((value, i) => ({ i, rest: value - Math.floor(value) }))
70    .sort((a, b) => b.rest - a.rest)
71  for (const { i } of byRemainder) {
72    if (left <= 0) break
73    cells[i] = (cells[i] ?? 0) + 1
74    left -= 1
75  }
76  segments.forEach((s, i) => {
77    if (s.kind !== 'used' || (cells[i] ?? 0) > 0) return
78    const widest = cells.indexOf(Math.max(...cells))
79    if ((cells[widest] ?? 0) <= 1) return
80    cells[widest] = (cells[widest] ?? 0) - 1
81    cells[i] = 1
82  })
83
84  return cells
85}
86
87export const formatTokens = (tokens: number): string =>
88  tokens >= 1_000_000
89    ? `${(tokens / 1_000_000).toFixed(1)}M`
90    : tokens >= 1000
91      ? `${Math.round(tokens / 1000)}k`
92      : `${tokens}`
93

1// ABOUTME: Type contract for the context-bar mod's session state.
2// ABOUTME: Declares the breakdown snapshot and the visibility flag under PluginState.
3
4export type Segment = { name: string; color: string; tokens: number; kind: 'used' | 'free' | 'buffer' }
5
6export type Snapshot = { segments: Segment[]; usedTokens: number; maxTokens: number; percentage: number }
7
8declare module 'claude-code' {
9  interface PluginState {
10    'context-bar': { snapshot: Snapshot | null; isShown: boolean }
11  }
12}
13