context-tracker

1import { atom, read, update } from 'claude-code'
2import type { Register, EngineInterface } from 'claude-code'
3
4import type { Breakdown, Sample, SampleKind, SkillEvent, View } from '../types'
5
6
7const PANE = 'context-tracker'
8const MAX_SAMPLES = 600
9const MAX_SKILL_EVENTS = 300
10
11const samples = atom({ plugin: 'context-tracker', key: 'samples' } as const, [] as Sample[])
12const breakdown = atom({ plugin: 'context-tracker', key: 'breakdown' } as const, null as Breakdown | null)
13const skillEvents = atom({ plugin: 'context-tracker', key: 'skillEvents' } as const, [] as SkillEvent[])
14const view = atom({ plugin: 'context-tracker', key: 'view' } as const, 'categories' as View)
15
16// ---------------------------------------------------------------- palette
17
18const CATEGORY_COLORS: Record<string, string> = {
19  'System prompt': '#7aa2f7',
20  'System tools': '#bb9af7',
21  'MCP tools': '#ff9e64',
22  'Custom agents': '#e0af68',
23  'Memory files': '#9ece6a',
24  'Skills': '#2ac3de',
25  'Messages': '#f7768e',
26  'Free space': '#3b4261',
27  'Autocompact buffer': '#565f89',
28}
29const KIND_COLORS = { used: '#c0caf5', free: '#3b4261', buffer: '#565f89', deferred: '#414868' }
30const SPIKE_RED = '#f7768e'
31const SPIKE_YELLOW = '#e0af68'
32const CALM_GREEN = '#9ece6a'
33const COMPACT_BLUE = '#7dcfff'
34
35const colorOf = (name: string, kind: keyof typeof KIND_COLORS) =>
36  CATEGORY_COLORS[name] ?? KIND_COLORS[kind]
37
38const fillColor = (percent: number) =>
39  percent >= 80 ? SPIKE_RED : percent >= 50 ? SPIKE_YELLOW : CALM_GREEN
40
41/** Red for a big jump, yellow for a notable one, dim otherwise. */
42const deltaColor = (delta: number, window: number) => {
43  if (delta >= Math.max(10_000, window * 0.05)) return SPIKE_RED
44  if (delta >= 3_000) return SPIKE_YELLOW
45  if (delta < 0) return COMPACT_BLUE
46  return undefined
47}
48
49// ---------------------------------------------------------------- formatting
50
51const fmtK = (n: number) => {
52  const abs = Math.abs(n)
53  if (abs < 1_000) return `${Math.round(n)}`
54  if (abs < 1_000_000) return `${(n / 1_000).toFixed(1)}k`
55  return `${(n / 1_000_000).toFixed(2)}M`
56}
57const fmtDelta = (n: number) => (n >= 0 ? `+${fmtK(n)}` : `-${fmtK(-n)}`)
58const pad2 = (n: number) => String(n).padStart(2, '0')
59const fmtTime = (t: number) => {
60  const d = new Date(t)
61  return `${pad2(d.getHours())}:${pad2(d.getMinutes())}:${pad2(d.getSeconds())}`
62}
63const fmtMs = (ms: number) => (ms >= 60_000 ? `${(ms / 60_000).toFixed(1)}m` : `${Math.round(ms / 1000)}s`)
64const padEnd = (s: string, n: number) => (s.length >= n ? s.slice(0, n) : s + ' '.repeat(n - s.length))
65const padStart = (s: string, n: number) => (s.length >= n ? s : ' '.repeat(n - s.length) + s)
66
67const bar = (ratio: number, width: number) => {
68  const r = Math.max(0, Math.min(1, ratio))
69  const filled = Math.round(r * width)
70  return '█'.repeat(filled) + '░'.repeat(Math.max(0, width - filled))
71}
72
73const SPARK = '▁▂▃▄▅▆▇█'
74const sparkline = (values: number[], max: number) =>
75  values
76    .map(v => SPARK[Math.min(SPARK.length - 1, Math.max(0, Math.round((v / Math.max(1, max)) * (SPARK.length - 1))))])
77    .join('')
78
79const toolsLabel = (tools: Record<string, number>, limit = 3) =>
80  Object.entries(tools)
81    .sort((a, b) => b[1] - a[1])
82    .slice(0, limit)
83    .map(([name, n]) => (n > 1 ? `${name}×${n}` : name))
84    .join(' ')
85
86// ---------------------------------------------------------------- recording
87
88type Recorder = {
89  /** The current main-thread turn's tool calls, by name. */
90  tools: Record<string, number>
91  /** Skills invoked in the current main-thread turn. */
92  skills: string[]
93  /** Main-thread turns seen since load. */
94  turn: number
95  /** Last context size recorded, to compute deltas without a read. */
96  lastTokens: number
97  window: number
98}
99
100async function addSample(
101  $: EngineInterface,
102  rec: Recorder,
103  input: { tokens: number; percent: number; kind: SampleKind; durationMs?: number; fromTokens?: number },
104) {
105  const t = await $.clock.now()
106  let isNew = false
107  await update($, samples, list => {
108    const last = list[list.length - 1]
109    // A measurement right after a turn repeats the turn's own figure: keep one point.
110    if (last && input.kind === 'measure' && last.tokens === input.tokens) return list
111    isNew = true
112    const sample: Sample = {
113      t,
114      tokens: input.tokens,
115      percent: input.percent,
116      kind: input.kind,
117      turn: rec.turn,
118      tools: input.kind === 'turn' ? { ...rec.tools } : {},
119      skills: input.kind === 'turn' ? [...rec.skills] : [],
120      durationMs: input.durationMs,
121      fromTokens: input.fromTokens,
122    }
123    return [...list, sample].slice(-MAX_SAMPLES)
124  })
125  if (isNew) {
126    await update($, skillEvents, list =>
127      list.map(ev => (ev.tokensAfter === null ? { ...ev, tokensAfter: input.tokens } : ev)),
128    )
129  }
130  rec.lastTokens = input.tokens
131}
132
133async function refreshBreakdown($: EngineInterface) {
134  try {
135    const usage = await $.session.usage({ breakdown: 'summary', columns: 80 })
136    const b = usage.context.breakdown
137    if (!b) return
138    const snapshot: Breakdown = {
139      at: await $.clock.now(),
140      total: b.totalTokens,
141      max: b.rawMaxTokens,
142      percent: b.percentage,
143      categories: b.categories.map(c => ({ name: c.name, tokens: c.tokens, kind: c.kind })),
144    }
145    await update($, breakdown, () => snapshot)
146  } catch {
147    // No session bound (a -p run, a test without the op answered): keep what we had.
148  }
149}
150
151// ---------------------------------------------------------------- register
152
153export const register: Register = on => {
154  const rec: Recorder = { tools: {}, skills: [], turn: 0, lastTokens: 0, window: 0 }
155
156  on('session.start', async ($, e, next) => {
157    await $.command.register({
158      name: 'ctx',
159      description: 'Open the context tracker pane (categories, timeline, skills)',
160    })
161    void $.ui.open({ id: PANE, title: 'Context' })
162    void refreshBreakdown($)
163    return next(e)
164  })
165
166  on('command.run', { command: 'ctx' }, async $ => {
167    await $.ui.open({ id: PANE, title: 'Context', focus: true })
168    await refreshBreakdown($)
169    return { text: 'Context tracker pane opened. Hotkeys 1/2/3 switch views, r refreshes.' }
170  })
171
172  on('prompt.submit', ($, e, next) => {
173    rec.tools = {}
174    rec.skills = []
175    return next(e)
176  })
177
178  on('tool.call', ($, e, next) => {
179    if (!e.agentId) rec.tools[e.tool] = (rec.tools[e.tool] ?? 0) + 1
180    return next(e)
181  })
182
183  on('skill.prompt', async ($, e, next) => {
184    rec.skills.push(e.skill)
185    const ev: SkillEvent = {
186      t: await $.clock.now(),
187      skill: e.skill,
188      turn: rec.turn + 1,
189      promptTokens: Math.round(e.text.length / 4),
190      tokensBefore: rec.lastTokens,
191      tokensAfter: null,
192    }
193    await update($, skillEvents, list => [...list, ev].slice(-MAX_SKILL_EVENTS))
194    return next(e)
195  })
196
197  on('turn.complete', async ($, e, next) => {
198    if (!e.agentId) {
199      rec.turn += 1
200      const u = e.usage
201      if (u) {
202        const tokens = u.input_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens
203        const percent = rec.window > 0 ? Math.round((tokens / rec.window) * 100) : 0
204        await addSample($, rec, { tokens, percent, kind: 'turn', durationMs: e.durationMs })
205      }
206      rec.tools = {}
207      rec.skills = []
208    }
209    return next(e)
210  })
211
212  on('session.measure', async ($, e, next) => {
213    rec.window = e.context.window
214    if (e.changed.includes('context') && e.context.tokens !== undefined) {
215      const tokens = e.context.tokens
216      const percent = e.context.percent ?? Math.round((tokens / e.context.window) * 100)
217      const delta = tokens - rec.lastTokens
218      await addSample($, rec, { tokens, percent, kind: 'measure' })
219      await refreshBreakdown($)
220      const tag = delta === 0 || rec.lastTokens === tokens ? '' : ` ${fmtDelta(delta)}`
221      $.ui.status(`ctx ${percent}%${tag}`)
222    }
223    return next(e)
224  })
225
226  on('session.compact', async ($, e, next) => {
227    const result = await next(e)
228    if (!e.agentId && e.trigger !== 'precompute' && 'messages' in result && result.messages) {
229      const from = result.tokensBefore ?? rec.lastTokens
230      const to = result.tokensAfter ?? 0
231      const percent = rec.window > 0 ? Math.round((to / rec.window) * 100) : 0
232      await addSample($, rec, { tokens: to, percent, kind: 'compact', fromTokens: from })
233      await refreshBreakdown($)
234    }
235    return result
236  })
237
238  // ---------------------------------------------------------------- drawing
239
240  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
241    const { Box, Text, Button } = $.ui.resolve(e)
242    const width = Math.max(40, e.props.bodyColumns ?? e.viewport?.columns ?? 80)
243    const [list, snap, events, current] = await Promise.all([
244      read($, samples),
245      read($, breakdown),
246      read($, skillEvents),
247      read($, view),
248    ])
249    const last = list[list.length - 1]
250    const prev = list[list.length - 2]
251    const window = snap?.max ?? rec.window
252    const tokensNow = last?.tokens ?? snap?.total ?? 0
253    const pctNow = last?.percent ?? snap?.percent ?? 0
254    const lastDelta = last && prev ? last.tokens - prev.tokens : 0
255
256    const tab = (key: View, hotkey: string, label: string) => (
257      <Button
258        key={`view-${key}`}
259        hotkey={hotkey}
260        plain
261        label={label}
262        dimColor={current !== key}
263        onPress={() => update($, view, () => key)}
264      />
265    )
266
267    const header = (
268      <Box flexDirection="column">
269        <Box>
270          <Text bold>Context </Text>
271          <Text color={fillColor(pctNow)} bold>{fmtK(tokensNow)}</Text>
272          <Text dimColor> / {fmtK(window)} ({pctNow}%)</Text>
273          {last && prev && (
274            <Text color={deltaColor(lastDelta, window)}> last {fmtDelta(lastDelta)}</Text>
275          )}
276          <Text dimColor> · {list.length} pts · {events.length} skill runs</Text>
277        </Box>
278        <Box gap={2}>
279          {tab('categories', '1', 'Categories')}
280          {tab('timeline', '2', 'Timeline')}
281          {tab('skills', '3', 'Skills')}
282          <Button key="refresh" hotkey="r" plain dimColor label="Refresh" onPress={() => refreshBreakdown($)} />
283        </Box>
284      </Box>
285    )
286
287    let body
288    if (current === 'categories') body = drawCategories()
289    else if (current === 'timeline') body = drawTimeline()
290    else body = drawSkills()
291
292    return (
293      <Box flexDirection="column">
294        {header}
295        <Text> </Text>
296        {body}
297      </Box>
298    )
299
300    // ------------------------------------------------ categories view
301    function drawCategories() {
302      if (!snap) return <Text dimColor>No breakdown yet. Press r after the first response.</Text>
303      const inWindow = snap.categories.filter(c => c.kind !== 'deferred')
304      const deferred = snap.categories.filter(c => c.kind === 'deferred')
305      const stackWidth = width - 2
306      const stack = inWindow
307        .map(c => ({ c, cells: Math.round((c.tokens / Math.max(1, snap.max)) * stackWidth) }))
308        .filter(x => x.cells > 0)
309      const nameW = Math.min(20, Math.max(...inWindow.map(c => c.name.length), 8))
310      const barW = Math.max(10, width - nameW - 20)
311      return (
312        <Box flexDirection="column">
313          <Box>
314            {stack.map(({ c, cells }) => (
315              <Text color={colorOf(c.name, c.kind)}>{'█'.repeat(cells)}</Text>
316            ))}
317          </Box>
318          <Text> </Text>
319          {inWindow.map(c => {
320            const share = c.tokens / Math.max(1, snap.max)
321            return (
322              <Box>
323                <Text color={colorOf(c.name, c.kind)}>■ </Text>
324                <Text dimColor={c.kind !== 'used'}>{padEnd(c.name, nameW)} </Text>
325                <Text>{padStart(fmtK(c.tokens), 7)} </Text>
326                <Text dimColor>{padStart(`${Math.round(share * 100)}%`, 4)} </Text>
327                <Text color={colorOf(c.name, c.kind)}>{bar(share, barW)}</Text>
328              </Box>
329            )
330          })}
331          {deferred.length > 0 && (
332            <Box flexDirection="column" marginTop={1}>
333              {deferred.map(c => (
334                <Text dimColor>
335                  ◌ {padEnd(c.name, nameW)} {padStart(fmtK(c.tokens), 7)} (outside the window, loads on demand)
336                </Text>
337              ))}
338            </Box>
339          )}
340          <Text dimColor>
341            {'\n'}measured {fmtTime(snap.at)} · {fmtK(snap.total)} of {fmtK(snap.max)} ({snap.percent}%)
342          </Text>
343        </Box>
344      )
345    }
346
347    // ------------------------------------------------ timeline view
348    function drawTimeline() {
349      if (list.length === 0) return <Text dimColor>No measurements yet. Points appear after each turn.</Text>
350      const maxTokens = Math.max(window, ...list.map(s => s.tokens))
351      const sparkWidth = width - 2
352      const sparkPoints = list.slice(-sparkWidth)
353      const barW = Math.max(8, Math.min(24, Math.floor(width * 0.25)))
354      const rowsRoom = Math.max(10, (e.viewport?.rows ?? 40) - 8)
355      const shown = list.slice(-rowsRoom)
356      const firstIndex = list.length - shown.length
357      const labelW = Math.max(10, width - 8 - 1 - barW - 1 - 7 - 1 - 7 - 1 - 4 - 2)
358      return (
359        <Box flexDirection="column">
360          <Text color={fillColor(pctNow)}>{sparkline(sparkPoints.map(s => s.tokens), maxTokens)}</Text>
361          <Text dimColor>
362            {padEnd('time', 8)} {padEnd('fill', barW)} {padStart('tokens', 7)} {padStart('delta', 7)} {padEnd('turn', 4)}  what happened
363          </Text>
364          {shown.map((s, i) => {
365            const before = list[firstIndex + i - 1]
366            const delta = before ? s.tokens - before.tokens : 0
367            const dc = deltaColor(delta, window)
368            const isSpike = dc === SPIKE_RED
369            const parts: string[] = []
370            if (s.kind === 'compact') parts.push(`compacted from ${fmtK(s.fromTokens ?? 0)}`)
371            if (s.skills.length) parts.push(`/${s.skills.join(' /')}`)
372            if (Object.keys(s.tools).length) parts.push(toolsLabel(s.tools))
373            if (s.durationMs !== undefined) parts.push(fmtMs(s.durationMs))
374            const label = parts.join(' · ')
375            return (
376              <Box>
377                <Text dimColor>{fmtTime(s.t)} </Text>
378                <Text color={fillColor(s.percent)}>{bar(s.tokens / Math.max(1, maxTokens), barW)} </Text>
379                <Text bold={isSpike}>{padStart(fmtK(s.tokens), 7)} </Text>
380                <Text color={dc} dimColor={dc === undefined} bold={isSpike}>
381                  {padStart(before ? fmtDelta(delta) : '', 7)}{' '}
382                </Text>
383                <Text dimColor>{padEnd(s.kind === 'compact' ? 'cmp' : `T${s.turn}`, 4)}  </Text>
384                <Text color={s.kind === 'compact' ? COMPACT_BLUE : isSpike ? SPIKE_RED : undefined} wrap="truncate-end">
385                  {isSpike ? '▲ ' : ''}{label.slice(0, labelW)}
386                </Text>
387              </Box>
388            )
389          })}
390        </Box>
391      )
392    }
393
394    // ------------------------------------------------ skills view
395    function drawSkills() {
396      if (events.length === 0) return <Text dimColor>No skill invoked yet.</Text>
397      type Agg = { skill: string; runs: number; prompt: number; growth: number; max: number; last: number }
398      const byName = new Map<string, Agg>()
399      for (const ev of events) {
400        const growth = ev.tokensAfter === null ? 0 : ev.tokensAfter - ev.tokensBefore
401        const a = byName.get(ev.skill) ?? { skill: ev.skill, runs: 0, prompt: 0, growth: 0, max: 0, last: 0 }
402        a.runs += 1
403        a.prompt += ev.promptTokens
404        a.growth += growth
405        a.max = Math.max(a.max, growth)
406        a.last = Math.max(a.last, ev.t)
407        byName.set(ev.skill, a)
408      }
409      const aggs = [...byName.values()].sort((a, b) => b.growth - a.growth)
410      const nameW = Math.min(24, Math.max(6, ...aggs.map(a => a.skill.length)))
411      const maxGrowth = Math.max(1, ...aggs.map(a => a.growth))
412      const barW = Math.max(8, Math.min(20, width - nameW - 40))
413      const recent = events.slice(-Math.max(5, (e.viewport?.rows ?? 40) - aggs.length - 12)).reverse()
414      return (
415        <Box flexDirection="column">
416          <Text dimColor>
417            {padEnd('skill', nameW)} {padStart('runs', 4)} {padStart('prompt≈', 8)} {padStart('growth', 8)} {padStart('max', 7)}  share of growth
418          </Text>
419          {aggs.map(a => (
420            <Box>
421              <Text bold>{padEnd(a.skill, nameW)} </Text>
422              <Text>{padStart(String(a.runs), 4)} </Text>
423              <Text dimColor>{padStart(fmtK(a.prompt), 8)} </Text>
424              <Text color={deltaColor(a.max, window)}>{padStart(fmtDelta(a.growth), 8)} </Text>
425              <Text color={deltaColor(a.max, window)}>{padStart(fmtDelta(a.max), 7)}  </Text>
426              <Text color={deltaColor(a.max, window) ?? CALM_GREEN}>{bar(a.growth / maxGrowth, barW)}</Text>
427            </Box>
428          ))}
429          <Text> </Text>
430          <Text dimColor>recent invocations (turn growth = context after the turn minus before the skill ran)</Text>
431          {recent.map(ev => {
432            const growth = ev.tokensAfter === null ? null : ev.tokensAfter - ev.tokensBefore
433            const dc = growth === null ? undefined : deltaColor(growth, window)
434            return (
435              <Box>
436                <Text dimColor>{fmtTime(ev.t)} </Text>
437                <Text dimColor>T{String(ev.turn).padEnd(3)} </Text>
438                <Text bold>{padEnd('/' + ev.skill, nameW + 1)} </Text>
439                <Text dimColor>prompt≈{padStart(fmtK(ev.promptTokens), 6)} </Text>
440                <Text color={dc} dimColor={dc === undefined}>
441                  {growth === null ? 'pending' : `turn ${fmtDelta(growth)}`}
442                </Text>
443                <Text dimColor> ({fmtK(ev.tokensBefore)} → {ev.tokensAfter === null ? '…' : fmtK(ev.tokensAfter)})</Text>
444              </Box>
445            )
446          })}
447        </Box>
448      )
449    }
450  })
451}
452

1export type SampleKind = 'turn' | 'measure' | 'compact'
2
3/** One point on the context timeline. */
4export type Sample = {
5  /** When it was taken, ms since the epoch. */
6  t: number
7  /** Context tokens the next request re-sends (input + cache read + cache write). */
8  tokens: number
9  /** `tokens` over the model's window, whole percent. */
10  percent: number
11  kind: SampleKind
12  /** Main-thread turn number the sample belongs to (0 before the first). */
13  turn: number
14  /** Tool calls made in that turn, by tool name. */
15  tools: Record<string, number>
16  /** Skills invoked in that turn. */
17  skills: string[]
18  /** Wall-clock length of the turn, when known. */
19  durationMs?: number
20  /** For a compaction: the size it came down from. */
21  fromTokens?: number
22}
23
24export type Category = {
25  name: string
26  tokens: number
27  kind: 'used' | 'free' | 'buffer' | 'deferred'
28}
29
30export type Breakdown = {
31  at: number
32  total: number
33  max: number
34  percent: number
35  categories: Category[]
36}
37
38/** One skill invocation and the growth attributed to it. */
39export type SkillEvent = {
40  t: number
41  skill: string
42  turn: number
43  /** Estimated tokens of the skill's own prompt text. */
44  promptTokens: number
45  /** Context size at the last sample before the skill ran. */
46  tokensBefore: number
47  /** Context size at the first sample after it; null until measured. */
48  tokensAfter: number | null
49}
50
51export type View = 'categories' | 'timeline' | 'skills'
52
53declare module 'claude-code' {
54  interface PluginState {
55    'context-tracker': {
56      samples: Sample[]
57      breakdown: Breakdown | null
58      skillEvents: SkillEvent[]
59      view: View
60    }
61  }
62}
63