context-map

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { ContextMap, Segment } from '../types'
5import { SEQUENCE, VESUVIUS } from './palette'
6
7const map = atom({ plugin: 'context-map', key: 'map' } as const, null)
8const isExpanded = atom({ plugin: 'context-map', key: 'isExpanded' } as const, false)
9
10export function formatTokens(n: number): string {
11  return n >= 1000 ? `${Math.round(n / 1000)}k` : String(n)
12}
13
14// How many cells of the bar each segment gets: any segment with tokens gets at least one,
15// and the total never exceeds the width.
16export function allocateCells(segments: readonly Segment[], max: number, width: number): number[] {
17  const cells = segments.map(s => {
18    const t = s.kind === 'deferred' || max <= 0 ? 0 : s.tokens
19
20    return t > 0 ? Math.max(1, Math.round((t / max) * width)) : 0
21  })
22  let excess = cells.reduce((a, b) => a + b, 0) - width
23
24  // On overflow, take from the free space first, then from the largest shares.
25  segments.forEach((s, i) => {
26    const c = cells[i] ?? 0
27    if (excess > 0 && s.kind === 'free') {
28      const cut = Math.min(excess, c)
29      cells[i] = c - cut
30      excess -= cut
31    }
32  })
33  while (excess > 0) {
34    const largest = Math.max(...cells)
35    if (largest <= 1) break
36    const i = cells.indexOf(largest)
37    cells[i] = largest - 1
38    excess -= 1
39  }
40
41  return cells
42}
43
44export function colorOf(segments: readonly Segment[], i: number): string {
45  const s = segments[i]
46  if (s === undefined || s.kind !== 'used') return VESUVIUS.ash
47  const usedBefore = segments.slice(0, i).filter(x => x.kind === 'used').length
48
49  return SEQUENCE[usedBefore % SEQUENCE.length] ?? VESUVIUS.lava
50}
51
52export function alertColor(percent: number): string {
53  return percent >= 85 ? VESUVIUS.lava : percent >= 60 ? VESUVIUS.magma : VESUVIUS.tuff
54}
55
56async function refresh($: EngineInterface): Promise<void> {
57  const usage = await $.session.usage({ breakdown: 'summary' })
58  const b = usage.context.breakdown
59  if (b === undefined) return
60
61  const next: ContextMap = {
62    percent: b.percentage,
63    total: b.totalTokens,
64    max: b.rawMaxTokens,
65    // What is used first, then the compaction reserve, then the free space.
66    segments: [
67      ...b.categories.filter(c => c.kind === 'used'),
68      ...b.categories.filter(c => c.kind === 'buffer'),
69      ...b.categories.filter(c => c.kind === 'free'),
70    ].map(c => ({ name: c.name, tokens: c.tokens, kind: c.kind })),
71  }
72  await update($, map, () => next)
73}
74
75export const register: Register = on => {
76  on('session.start', async ($, e, next) => {
77    await refresh($).catch(() => undefined)
78
79    return next(e)
80  })
81
82  on('turn.complete', async ($, e, next) => {
83    const result = await next(e)
84    await refresh($).catch(() => undefined)
85
86    return result
87  })
88
89  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
90    const m = await read($, map)
91    if (e.props.hasSurvey || m === null || m.max <= 0) return next(e)
92
93    const { Box, Text, Button } = $.ui.resolve(e)
94    const isOpen = await read($, isExpanded)
95    const width = Math.max(10, Math.min(60, e.props.bodyColumns - 30))
96    const cells = allocateCells(m.segments, m.max, width)
97    const used = m.segments.filter(s => s.kind === 'used')
98
99    return (
100      <Box flexDirection="column">
101        <Box flexDirection="row" columnGap={1}>
102          <Text color={VESUVIUS.lava}>▲</Text>
103          <Text>
104            {m.segments.map((s, i) => (
105              <Text color={colorOf(m.segments, i)} dimColor={s.kind === 'free'}>
106                {(s.kind === 'used' ? '█' : s.kind === 'buffer' ? '▒' : '░').repeat(cells[i] ?? 0)}
107              </Text>
108            ))}
109          </Text>
110          <Text color={alertColor(m.percent)} bold={m.percent >= 85}>
111            {m.percent}%
112          </Text>
113          <Text color={VESUVIUS.ash}>
114            {formatTokens(m.total)}/{formatTokens(m.max)}
115          </Text>
116          <Button
117            key="details"
118            plain
119            label={isOpen ? 'close' : 'details'}
120            onPress={() => update($, isExpanded, v => !v)}
121          />
122        </Box>
123        {isOpen &&
124          used.map(s => (
125            <Text>
126              {'  '}
127              <Text color={colorOf(m.segments, m.segments.indexOf(s))}>■</Text> {s.name}{' '}
128              <Text color={VESUVIUS.ash}>
129                {formatTokens(s.tokens)} · {Math.round((s.tokens / m.max) * 100)}%
130              </Text>
131            </Text>
132          ))}
133      </Box>
134    )
135  })
136}
137

1// Vesuvius palette (placeholder): replace the hex values with the design system's.
2// Every mod in affreschi keeps its own copy, because a plugin cannot import another plugin's files.
3export const VESUVIUS = {
4  lava: '#E4572E',
5  magma: '#F28C28',
6  sulfur: '#F2C14E',
7  tuff: '#C9B79C',
8  ash: '#8A8D91',
9  gulf: '#2E86AB',
10  vineyard: '#7A9E4E',
11} as const
12
13// The order categories take their colors in, warmest first.
14export const SEQUENCE = [
15  VESUVIUS.lava,
16  VESUVIUS.magma,
17  VESUVIUS.sulfur,
18  VESUVIUS.tuff,
19  VESUVIUS.gulf,
20  VESUVIUS.vineyard,
21] as const
22

1export type Segment = { name: string; tokens: number; kind: 'used' | 'free' | 'buffer' | 'deferred' }
2
3export type ContextMap = {
4  percent: number
5  total: number
6  max: number
7  segments: Segment[]
8}
9
10declare module 'claude-code' {
11  interface PluginState {
12    'context-map': { map: ContextMap | null; isExpanded: boolean }
13  }
14}
15