context-band

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionMessage } from 'claude-code'
3
4import type { Band, Contribution } from '../types'
5
6const STEP_TOKENS = 150_000
7const CEILING_TOKENS = 250_000
8const CACHE_TTL_MS = 60 * 60 * 1000
9const CACHE_WARN_MS = 10 * 60 * 1000
10const CACHE_RELEVANT_TOKENS = 30_000
11const STREAK_SHOWN = 3
12const SPARK_TURNS = 12
13const TOP_TURNS = 5
14const TOP_SHOWN = 2
15const TOGGLE_GAP = 4
16const LIST_ROWS = 10
17const METER_CELLS = 10
18const GROUP = ' │ '
19const LABEL_CELLS = 18
20const FIVE_HOUR = 'five_hour'
21const HINT_INDENT = 4
22const CHARS_PER_TOKEN = 4
23const TICK_MS = 15_000
24const GRAPHIFY_LOCK = '/tmp/graphify-autosync.lock'
25const HOOKS_LABEL = 'hook reminders'
26const TEXT_LABEL = 'text'
27
28const EMPTY: Band = {
29  startedAt: 0,
30  readings: [],
31  window: 0,
32  lastTurnAt: null,
33  streak: 0,
34  turn: 0,
35  contributions: [],
36  cost: null,
37  fiveHour: null,
38}
39
40const band = atom({ plugin: 'context-band', key: 'band' } as const, EMPTY)
41const now = atom({ plugin: 'context-band', key: 'now' } as const, 0)
42const isRebuilding = atom({ plugin: 'context-band', key: 'isRebuilding' } as const, false)
43const isExpanded = atom({ plugin: 'context-band', key: 'isExpanded' } as const, false)
44
45const complete = (stored: Band): Band => ({ ...EMPTY, ...stored })
46
47export const kilo = (tokens: number): string => {
48  if (tokens >= 1_000_000) return `${Number((tokens / 1_000_000).toFixed(1))}M`
49  if (tokens >= 1000) return `${(tokens / 1000).toFixed(1)}k`
50  return `${tokens}`
51}
52
53export const tokenColor = (tokens: number): string | undefined => {
54  if (tokens >= CEILING_TOKENS) return 'red'
55  if (tokens >= STEP_TOKENS) return 'magenta'
56  if (tokens >= STEP_TOKENS - 50_000) return 'yellow'
57  return undefined
58}
59
60export const stepLabel = (tokens: number): string => {
61  if (tokens >= CEILING_TOKENS) return '250k ceiling passed: /clear with a handoff note'
62  if (tokens >= STEP_TOKENS) return '150k step reached: /clear at the end of the phase (ceiling 250k)'
63  return `/clear advised in ${kilo(STEP_TOKENS - tokens)} (step 150k, ceiling 250k)`
64}
65
66export type CacheState = { label: string; color?: string; detail?: string } | null
67
68export const cacheState = (lastTurnAt: number | null, at: number, tokens: number): CacheState => {
69  if (lastTurnAt === null || tokens < CACHE_RELEVANT_TOKENS) return null
70  const left = CACHE_TTL_MS - (at - lastTurnAt)
71  if (left <= 0) {
72    return { label: 'Cache expired', color: 'red', detail: `next turn rereads ${kilo(tokens)} at full price` }
73  }
74  const minutes = Math.ceil(left / 60_000)
75  return {
76    label: `Cache (${minutes}m)`,
77    color: left < CACHE_WARN_MS ? 'yellow' : undefined,
78  }
79}
80
81export const nextStreak = (streak: number, toolsInStep: number): number => {
82  if (toolsInStep === 1) return streak + 1
83  if (toolsInStep > 1) return 0
84  return streak
85}
86
87export const estimate = (chars: number): number => Math.round(chars / CHARS_PER_TOKEN)
88
89const basename = (path: unknown): string =>
90  typeof path === 'string' ? (path.split('/').pop() ?? path) : ''
91
92const clip = (text: unknown, length: number): string => {
93  if (typeof text !== 'string') return ''
94  const line = text.replace(/\s+/g, ' ').trim()
95  return line.length > length ? `${line.slice(0, length - 1)}…` : line
96}
97
98const host = (url: unknown): string =>
99  typeof url === 'string' ? (url.replace(/^https?:\/\//, '').split('/')[0] ?? '') : ''
100
101export const callLabel = (tool: string, input: Record<string, unknown>): string => {
102  switch (tool) {
103    case 'Read':
104    case 'Edit':
105    case 'Write':
106      return `${tool} ${basename(input.file_path)}`
107    case 'NotebookEdit':
108      return `${tool} ${basename(input.notebook_path)}`
109    case 'Bash':
110      return `Bash ${clip(input.description ?? input.command, 20)}`
111    case 'Grep':
112    case 'Glob':
113      return `${tool} ${clip(input.pattern, 24)}`
114    case 'Agent':
115      return `Agent ${clip(input.description ?? input.subagent_type, 28)}`
116    case 'Skill':
117      return `Skill ${clip(input.skill, 24)}`
118    case 'WebFetch':
119      return `WebFetch ${clip(host(input.url), 24)}`
120    default:
121      return tool
122  }
123}
124
125export const topContributions = (
126  contributions: Contribution[],
127  sinceTurn: number,
128  count: number,
129): { label: string; tokens: number }[] => {
130  const totals = new Map<string, number>()
131  for (const one of contributions) {
132    if (one.turn < sinceTurn) continue
133    totals.set(one.label, (totals.get(one.label) ?? 0) + one.tokens)
134  }
135  return [...totals.entries()]
136    .map(([label, tokens]) => ({ label, tokens }))
137    .filter(one => one.tokens > 0)
138    .sort((a, b) => b.tokens - a.tokens)
139    .slice(0, count)
140}
141
142export const replay = (
143  messages: readonly SessionMessage[],
144): { turn: number; streak: number; contributions: Contribution[] } => {
145  let turn = 0
146  let streak = 0
147  let contributions: Contribution[] = []
148  let toolsInStep = 0
149  for (const message of messages) {
150    if (message.role === 'user') {
151      streak = nextStreak(streak, toolsInStep)
152      toolsInStep = 0
153    }
154    const isPrompt = message.role === 'user' && message.text !== '' && (message.toolResults ?? []).length === 0
155    if (isPrompt) turn += 1
156    if (turn === 0) continue
157    if (message.text !== '') contributions.push({ turn, label: TEXT_LABEL, tokens: estimate(message.text.length) })
158    if (message.role !== 'assistant') continue
159    toolsInStep += message.toolUses.length
160    for (const use of message.toolUses) {
161      contributions.push({
162        turn,
163        label: callLabel(use.tool, use.input),
164        tokens: estimate(JSON.stringify(use.input).length + (use.text ?? '').length),
165      })
166    }
167  }
168  streak = nextStreak(streak, toolsInStep)
169  contributions = contributions.filter(one => one.turn > turn - TOP_TURNS)
170  return { turn, streak, contributions }
171}
172
173export const calibrate = (
174  calls: { label: string; tokens: number }[],
175  delta: number | null,
176): { calls: { label: string; tokens: number }[]; text: number } => {
177  const measured = calls.reduce((sum, one) => sum + one.tokens, 0)
178  if (delta === null || delta <= 0) return { calls, text: 0 }
179  if (measured <= delta) return { calls, text: delta - measured }
180  const ratio = delta / measured
181  return { calls: calls.map(one => ({ ...one, tokens: Math.round(one.tokens * ratio) })), text: 0 }
182}
183
184export const meter = (percent: number): string => {
185  const filled = Math.min(METER_CELLS, Math.max(0, Math.floor((percent * METER_CELLS) / 100)))
186  return '█'.repeat(filled) + '░'.repeat(METER_CELLS - filled)
187}
188
189export const meterColor = (percent: number): string => {
190  if (percent >= 85) return 'red'
191  if (percent >= 60) return 'yellow'
192  return 'green'
193}
194
195export const span = (ms: number): string => {
196  const minutes = Math.floor(Math.max(0, ms) / 60_000)
197  if (minutes < 1) return `${Math.floor(Math.max(0, ms) / 1000)}s`
198  if (minutes < 60) return `${minutes}m`
199  return `${Math.floor(minutes / 60)}h${minutes % 60}m`
200}
201
202export const dollars = (usd: number): string => `$${usd.toFixed(2)}`
203
204async function recordStep($: EngineInterface, tools: number): Promise<void> {
205  await update($, band, current => {
206    const full = complete(current)
207    return { ...full, streak: nextStreak(full.streak, tools) }
208  })
209}
210
211export const register: Register = on => {
212  let toolsInStep = 0
213  let pending: { label: string; tokens: number }[] = []
214
215  on('session.start', async ($, e, next) => {
216    const tick = async () => {
217      const at = await $.clock.now()
218      await update($, now, () => at)
219      const isLocked = await $.fs.exists(GRAPHIFY_LOCK)
220      await update($, isRebuilding, () => isLocked)
221    }
222    $.clock.every(TICK_MS, () => void tick())
223    await tick()
224
225    const result = await next(e)
226    if (complete(await read($, band)).turn === 0) {
227      const usage = await $.session.usage()
228      const replayed = replay(await $.session.messages())
229      const tokens = usage.context.tokens
230      await update($, band, () => ({
231        ...EMPTY,
232        ...replayed,
233        startedAt: usage.startedAt,
234        window: usage.context.window,
235        readings: tokens === undefined || replayed.turn === 0 ? [] : [tokens],
236        cost: usage.cost?.usd ?? null,
237      }))
238    }
239
240    return result
241  })
242
243  on('tool.call', async ($, e, next) => {
244    if (e.agentId !== undefined) return next(e)
245    toolsInStep += 1
246    const result = await next(e)
247    const input = e as unknown as Record<string, unknown>
248    const output = typeof result.deny === 'string' ? result.deny : (result.text ?? '')
249    pending.push({
250      label: callLabel(e.tool, input),
251      tokens: estimate(JSON.stringify(input).length + output.length),
252    })
253    const reminders = (result.context ?? []).reduce((sum, one) => sum + one.length, 0)
254    if (reminders > 0) pending.push({ label: HOOKS_LABEL, tokens: estimate(reminders) })
255
256    return result
257  })
258
259  on('turn.step', async function* ($, e, next) {
260    if (e.agentId === undefined && e.index > 0) {
261      await recordStep($, toolsInStep)
262      toolsInStep = 0
263    }
264
265    return yield* next(e)
266  })
267
268  on('turn.complete', async ($, e, next) => {
269    const result = await next(e)
270    if (e.agentId !== undefined) return result
271    await recordStep($, toolsInStep)
272    toolsInStep = 0
273    const calls = pending
274    pending = []
275    const usage = await $.session.usage()
276    const at = await $.clock.now()
277    const tokens = usage.context.tokens
278    await update($, band, stored => {
279      const current = complete(stored)
280      const base = current.startedAt === usage.startedAt ? current : { ...EMPTY, startedAt: usage.startedAt }
281      const cost = usage.cost?.usd ?? base.cost
282      const limit = usage.rateLimits.find(one => one.kind === FIVE_HOUR)
283      const resetsAt = limit?.resetsAt === undefined ? null : Date.parse(limit.resetsAt)
284      const fiveHour = limit === undefined ? base.fiveHour : { percent: limit.percentUsed, resetsAt }
285      if (tokens === undefined) return { ...base, cost, fiveHour, lastTurnAt: at }
286      const turn = base.turn + 1
287      const previous = base.readings[base.readings.length - 1]
288      const calibrated = calibrate(calls, previous === undefined ? null : tokens - previous)
289      const added: Contribution[] = [
290        ...calibrated.calls.map(one => ({ turn, ...one })),
291        ...(calibrated.text > 0 ? [{ turn, label: TEXT_LABEL, tokens: calibrated.text }] : []),
292      ]
293      return {
294        ...base,
295        turn,
296        cost,
297        fiveHour,
298        window: usage.context.window,
299        readings: [...base.readings, tokens].slice(-SPARK_TURNS),
300        contributions: [...base.contributions, ...added].filter(one => one.turn > turn - TOP_TURNS),
301        lastTurnAt: at,
302      }
303    })
304    await update($, now, () => at)
305
306    return result
307  })
308
309  on('ui.render', { component: 'PromptHint' }, async ($, e, next) => {
310    const state = complete(await read($, band))
311    if (state.readings.length === 0) return next(e)
312
313    const hint = await next(e)
314    const columns = e.viewport === undefined ? undefined : Math.max(0, e.viewport.columns - HINT_INDENT)
315
316    const { Box, Button, Text } = $.ui.resolve(e)
317    const at = await read($, now)
318    const tokens = state.readings[state.readings.length - 1] ?? 0
319    const previous = state.readings[state.readings.length - 2]
320    const delta = previous === undefined ? 0 : tokens - previous
321    const percent = state.window > 0 ? Math.round((tokens * 100) / state.window) : null
322    const color = tokenColor(tokens)
323    const cache = e.props.isWorking ? null : cacheState(state.lastTurnAt, at, tokens)
324    const rebuilding = await read($, isRebuilding)
325    const grouping = state.streak >= STREAK_SHOWN
326    const limit = state.fiveHour
327    const resetIn = limit?.resetsAt == null || limit.resetsAt <= at ? null : span(limit.resetsAt - at)
328    const turns = Math.min(TOP_TURNS, state.turn)
329    const all = topContributions(state.contributions, state.turn - TOP_TURNS + 1, Number.MAX_SAFE_INTEGER)
330    const hidden = Math.max(0, all.length - TOP_SHOWN)
331    const expanded = await read($, isExpanded)
332    const shown = expanded ? all.slice(0, LIST_ROWS) : all.slice(0, TOP_SHOWN)
333    const showsTop = all.length > 0
334    const toggle = (
335      <Button
336        key="top-toggle"
337        label={expanded ? '[Collapse]' : hidden > 0 ? `[Show all (+${hidden})]` : '[Show all]'}
338        plain
339        dimColor
340        onPress={() => update($, isExpanded, value => !value)}
341      />
342    )
343    const header = `Top ${turns === 1 ? 'last turn' : `${turns} turns`}`
344
345    return (
346      <Box flexDirection="column">
347        <Box flexDirection="column" width={columns}>
348          <Box height={1}>
349            <Box flexGrow={1} flexShrink={1}>
350              <Text wrap="truncate-end">
351                <Text dimColor>Context </Text>
352                {percent !== null && <Text color={meterColor(percent)}>{meter(percent)} </Text>}
353                <Text color={color} bold={color !== undefined}>{kilo(tokens)}</Text>
354                {state.window > 0 && <Text dimColor>/{kilo(state.window)}</Text>}
355                {delta !== 0 && <Text dimColor> ({delta > 0 ? '+' : '−'}{kilo(Math.abs(delta))} this turn)</Text>}
356                {[
357                  cache !== null && <Text key="cache" color={cache.color} dimColor={cache.color === undefined}>{cache.label}</Text>,
358                  grouping && <Text key="grouping" color="yellow">{state.streak} turns</Text>,
359                  rebuilding && <Text key="graph" dimColor>graph rebuilding</Text>,
360                ].filter(Boolean).map((one, index) => (
361                  <Text key={`advice-${index}`}>
362                    <Text dimColor>{index === 0 ? GROUP : ' · '}</Text>
363                    {one}
364                  </Text>
365                ))}
366                {limit !== null && <Text dimColor>{GROUP}Lim 5h </Text>}
367                {limit !== null && <Text color={meterColor(limit.percent)}>{meter(limit.percent)}</Text>}
368                {limit !== null && <Text dimColor> {Math.round(limit.percent)}%{resetIn === null ? '' : ` ${resetIn}`}</Text>}
369                <Text dimColor>{GROUP}Session {span(at - state.startedAt)} · {state.turn} turn{state.turn > 1 ? 's' : ''}</Text>
370                {state.cost !== null && <Text dimColor> · {dollars(state.cost)}</Text>}
371                {showsTop && !expanded && <Text dimColor>{GROUP}{header} : </Text>}
372                {showsTop && !expanded && shown.map((one, index) => (
373                  <Text key={`top-${index}`}>
374                    {index > 0 && <Text dimColor> · </Text>}
375                    <Text>{one.label} </Text>
376                    <Text bold>{kilo(one.tokens)}</Text>
377                  </Text>
378                ))}
379              </Text>
380            </Box>
381            <Box flexShrink={0} paddingLeft={TOGGLE_GAP}>{toggle}</Box>
382          </Box>
383          {expanded && (
384            <Box flexDirection="column" paddingLeft={2}>
385              <Box>
386                <Box width={LABEL_CELLS} flexShrink={0}><Text dimColor>Step</Text></Box>
387                <Text color={color} dimColor={color === undefined} wrap="truncate-end">{stepLabel(tokens)}</Text>
388              </Box>
389              {cache?.detail !== undefined && (
390                <Box>
391                  <Box width={LABEL_CELLS} flexShrink={0}><Text dimColor>Cache</Text></Box>
392                  <Text color={cache.color} wrap="truncate-end">{cache.detail}</Text>
393                </Box>
394              )}
395              {grouping && (
396                <Box>
397                  <Box width={LABEL_CELLS} flexShrink={0}><Text dimColor>Batch</Text></Box>
398                  <Text color="yellow" wrap="truncate-end">
399                    {`${state.streak} turns with a single tool call: send independent calls together (each turn resends the whole context)`}
400                  </Text>
401                </Box>
402              )}
403              {showsTop && (
404                <Box>
405                  <Box width={LABEL_CELLS} flexShrink={0}><Text dimColor>{header}</Text></Box>
406                  <Box flexDirection="column">
407                    {shown.map((one, index) => (
408                      <Text key={`top-${index}`} wrap="truncate-end">
409                        <Text dimColor>{index + 1}. </Text>
410                        <Text>{one.label} </Text>
411                        <Text bold>{kilo(one.tokens)}</Text>
412                      </Text>
413                    ))}
414                  </Box>
415                </Box>
416              )}
417            </Box>
418          )}
419        </Box>
420        {hint}
421      </Box>
422    )
423  })
424}
425

1export type Contribution = { turn: number; label: string; tokens: number }
2
3export type Band = {
4  startedAt: number
5  readings: number[]
6  window: number
7  lastTurnAt: number | null
8  streak: number
9  turn: number
10  contributions: Contribution[]
11  cost: number | null
12  fiveHour: { percent: number; resetsAt: number | null } | null
13}
14
15declare module 'claude-code' {
16  interface PluginState {
17    'context-band': { band: Band; now: number; isRebuilding: boolean; isExpanded: boolean }
18  }
19}
20