ctx-usage-dock

1/**
2 * ctx-usage-dock: the context window broken down by category, as /context
3 * breaks it down, in a pane docked beside the transcript and on the status line.
4 *
5 * Each main-loop turn, compaction and new session recounts with the free
6 * `summary` estimate; the pane's `Full count` button asks the token-count API
7 * for each category, as /context does.
8 *
9 * `/ctx-usage-dock off` closes the pane, clears the status line and stops the
10 * counting, in this session and the next ones, until `/ctx-usage-dock on`.
11 *
12 * Everything that takes `$` is a top-level declaration: the loader inventories
13 * what a hooks module reaches for through `$`.
14 */
15import { atom, read, update } from 'claude-code'
16import type { ContextBreakdownDetail, EngineInterface, RenderElement, Register, SessionContextBreakdown } from 'claude-code'
17
18import type { UsageCategory, UsageGroup, UsageItem, UsageSnapshot } from '../types'
19
20const PLUGIN = 'ctx-usage-dock'
21const PANE = 'ctx-usage'
22const TITLE = 'Context'
23/** Items listed under a group before `+N more`. */
24const TOP_ITEMS = 8
25/** Categories the status line names after the percentage. */
26const STATUS_CATEGORIES = 3
27/** A rules file sits in a `.claude/rules` folder; every other memory file is memory. */
28const RULES_PATH = /[\\/]\.claude[\\/]rules[\\/]/
29/** The `$.store` key that keeps the on/off switch across sessions; unset is on. */
30const ENABLED_KEY = 'enabled'
31
32// Change the tag when UsageSnapshot's fields change: it declines what an
33// earlier version stored, which outlives a reload.
34const snapshot = atom({ plugin: 'ctx-usage-dock', key: 'snapshot' } as const, null, { shape: 'snapshot-1' })
35const isCounting = atom({ plugin: 'ctx-usage-dock', key: 'isCounting' } as const, false)
36
37/** One count at a time: a turn that ends while a full count runs waits for the next. */
38let isRunning = false
39
40// ── Numbers ─────────────────────────────────────────────────────────────────
41
42function formatTokens(tokens: number): string {
43  if (tokens >= 1_000_000) return `${(tokens / 1_000_000).toFixed(1)}M`
44  if (tokens >= 1_000) return `${(tokens / 1_000).toFixed(1)}k`
45  return String(tokens)
46}
47
48function formatPercent(tokens: number, max: number): string {
49  return max > 0 ? `${((tokens / max) * 100).toFixed(1)}%` : ''
50}
51
52function formatTime(epochMs: number): string {
53  const at = new Date(epochMs)
54  return `${String(at.getHours()).padStart(2, '0')}:${String(at.getMinutes()).padStart(2, '0')}`
55}
56
57/** The last two segments of a path: `.claude/CLAUDE.md`. */
58function shortPath(path: string): string {
59  return path.split(/[\\/]/).filter(Boolean).slice(-2).join('/')
60}
61
62const byTokens = (a: { tokens: number }, b: { tokens: number }) => b.tokens - a.tokens
63const sumTokens = (items: readonly { tokens: number }[]) => items.reduce((sum, item) => sum + item.tokens, 0)
64
65// ── Breakdown to snapshot ───────────────────────────────────────────────────
66
67function group(title: string, items: UsageItem[], note?: string): UsageGroup {
68  return { title, tokens: sumTokens(items), note, items: [...items].sort(byTokens) }
69}
70
71/** One item per MCP server: the tokens of its loaded schemas, and how many are loaded. */
72function mcpServers(breakdown: SessionContextBreakdown): UsageItem[] {
73  const servers = new Map<string, { tokens: number; loaded: number; total: number }>()
74  for (const tool of breakdown.mcpTools) {
75    const server = servers.get(tool.serverName) ?? { tokens: 0, loaded: 0, total: 0 }
76    server.total += 1
77    if (tool.isLoaded) {
78      server.loaded += 1
79      server.tokens += tool.tokens
80    }
81    servers.set(tool.serverName, server)
82  }
83  return [...servers].map(([name, server]) => ({
84    label: name,
85    tokens: server.tokens,
86    note: `${server.loaded}/${server.total} loaded`,
87  }))
88}
89
90function memoryItems(breakdown: SessionContextBreakdown, isRules: boolean): UsageItem[] {
91  return breakdown.memoryFiles
92    .filter((file) => RULES_PATH.test(file.path) === isRules)
93    .map((file) => ({ label: shortPath(file.path), tokens: file.tokens, note: file.type }))
94}
95
96function toSnapshot(breakdown: SessionContextBreakdown, detail: ContextBreakdownDetail, countedAt: number): UsageSnapshot {
97  const { skills, slashCommands } = breakdown
98  const groups = [
99    skills &&
100      group(
101        'Skills',
102        skills.skillFrontmatter.map((skill) => ({ label: skill.name, tokens: skill.tokens, note: skill.pluginName })),
103        `${skills.includedSkills}/${skills.totalSkills} listed`,
104      ),
105    breakdown.mcpTools.length > 0 && group('MCP servers', mcpServers(breakdown)),
106    group('Memory files', memoryItems(breakdown, false)),
107    group('Rules', memoryItems(breakdown, true)),
108    breakdown.agents.length > 0 &&
109      group('Agents', breakdown.agents.map((agent) => ({ label: agent.agentType, tokens: agent.tokens, note: agent.source }))),
110    slashCommands && {
111      title: 'Slash commands',
112      tokens: slashCommands.tokens,
113      note: `${slashCommands.includedCommands}/${slashCommands.totalCommands} listed`,
114      items: [],
115    },
116  ].filter((one): one is UsageGroup => Boolean(one) && (one as UsageGroup).tokens > 0)
117
118  return {
119    detail,
120    countedAt,
121    model: breakdown.model,
122    totalTokens: breakdown.totalTokens,
123    maxTokens: breakdown.rawMaxTokens,
124    percentage: breakdown.percentage,
125    autoCompactAt: breakdown.isAutoCompactEnabled ? breakdown.autoCompactThreshold : undefined,
126    categories: breakdown.categories.map(({ name, tokens, color, kind }) => ({ name, tokens, color, kind })),
127    groups,
128  }
129}
130
131/** `ctx 42% 84.2k/200k │ Messages 30.1k · System tools 12.4k · MCP tools 8.0k` */
132function statusLine(snap: UsageSnapshot): string {
133  const head = `ctx ${snap.percentage}% ${formatTokens(snap.totalTokens)}/${formatTokens(snap.maxTokens)}`
134  const top = snap.categories
135    .filter((category) => category.kind === 'used' && category.tokens > 0)
136    .sort(byTokens)
137    .slice(0, STATUS_CATEGORIES)
138    .map((category) => `${category.name} ${formatTokens(category.tokens)}`)
139  return top.length > 0 ? `${head} │ ${top.join(' · ')}` : head
140}
141
142// ── On and off ──────────────────────────────────────────────────────────────
143
144async function isEnabled($: EngineInterface): Promise<boolean> {
145  return (await $.store.get(ENABLED_KEY)) !== false
146}
147
148async function turnOn($: EngineInterface): Promise<void> {
149  await $.store.set(ENABLED_KEY, true)
150  await $.ui.open({ id: PANE, title: TITLE })
151  await recount($, 'summary')
152}
153
154async function turnOff($: EngineInterface): Promise<void> {
155  await $.store.set(ENABLED_KEY, false)
156  await $.ui.close({ id: PANE })
157  $.ui.status(undefined)
158}
159
160/** The count a turn, a compaction or a new session asks for: none while turned off. */
161async function autoRecount($: EngineInterface): Promise<void> {
162  if (await isEnabled($)) await recount($, 'summary')
163}
164
165// ── Counting ────────────────────────────────────────────────────────────────
166
167async function recount($: EngineInterface, detail: ContextBreakdownDetail): Promise<void> {
168  if (isRunning) return
169  isRunning = true
170  try {
171    await update($, isCounting, () => true)
172    const { context } = await $.session.usage({ breakdown: detail })
173    if (!context.breakdown) return
174    const snap = toSnapshot(context.breakdown, detail, await $.clock.now())
175    await update($, snapshot, () => snap)
176    // Turned off while the count ran: the status line stays clear.
177    if (await isEnabled($)) $.ui.status(statusLine(snap))
178  } catch (error) {
179    $.ui.log(`${PLUGIN}: counting the context failed: ${String(error)}`, { to: 'debug' })
180  } finally {
181    isRunning = false
182    await update($, isCounting, () => false).catch(() => undefined)
183  }
184}
185
186// ── Drawing ─────────────────────────────────────────────────────────────────
187
188type Site = Parameters<EngineInterface['ui']['resolve']>[0]
189
190/** The window as one bar: each used category its own colour, the rest free. */
191function windowBar($: EngineInterface, site: Site, snap: UsageSnapshot, width: number): RenderElement {
192  const { Box, Text } = $.ui.resolve(site)
193  const max = Math.max(snap.maxTokens, snap.totalTokens, 1)
194  let used = 0
195  let drawn = 0
196  const segments = snap.categories
197    .filter((category) => category.kind === 'used' && category.tokens > 0)
198    .map((category) => {
199      used += category.tokens
200      const end = Math.round((used / max) * width)
201      const cells = end - drawn
202      drawn = end
203      return cells > 0 ? <Text color={category.color}>{'█'.repeat(cells)}</Text> : null
204    })
205  return (
206    <Box flexDirection="row">
207      {segments}
208      <Text dimColor>{'░'.repeat(Math.max(0, width - drawn))}</Text>
209    </Box>
210  )
211}
212
213function row($: EngineInterface, site: Site, label: string, value: string, color?: string): RenderElement {
214  const { Box, Text } = $.ui.resolve(site)
215  return (
216    <Box flexDirection="row" justifyContent="space-between" gap={1}>
217      <Text color={color} wrap="truncate-end">
218        {label}
219      </Text>
220      <Text>{value}</Text>
221    </Box>
222  )
223}
224
225function groupBlock($: EngineInterface, site: Site, usageGroup: UsageGroup): RenderElement {
226  const { Box, Text } = $.ui.resolve(site)
227  const shown = usageGroup.items.slice(0, TOP_ITEMS)
228  const hidden = usageGroup.items.length - shown.length
229  const title = usageGroup.note ? `${usageGroup.title} (${usageGroup.note})` : usageGroup.title
230  return (
231    <Box flexDirection="column" marginTop={1}>
232      <Box flexDirection="row" justifyContent="space-between" gap={1}>
233        <Text bold>{title}</Text>
234        <Text bold>{formatTokens(usageGroup.tokens)}</Text>
235      </Box>
236      {shown.map((item) =>
237        row($, site, `  ${item.label}${item.note ? ` · ${item.note}` : ''}`, formatTokens(item.tokens)),
238      )}
239      {hidden > 0 && <Text dimColor>{`  +${hidden} more`}</Text>}
240    </Box>
241  )
242}
243
244function snapshotView($: EngineInterface, site: Site, snap: UsageSnapshot, width: number): RenderElement {
245  const { Box, Text } = $.ui.resolve(site)
246  const compactAt = snap.autoCompactAt ? ` · compacts at ${formatTokens(snap.autoCompactAt)}` : ''
247  return (
248    <Box flexDirection="column">
249      <Text bold>{`${snap.percentage}% · ${formatTokens(snap.totalTokens)} / ${formatTokens(snap.maxTokens)}`}</Text>
250      <Text dimColor wrap="truncate-end">{`${snap.model} · ${snap.detail} · ${formatTime(snap.countedAt)}${compactAt}`}</Text>
251      {windowBar($, site, snap, width)}
252      <Box flexDirection="column" marginTop={1}>
253        {snap.categories.map((category) =>
254          row(
255            $,
256            site,
257            category.name,
258            `${formatTokens(category.tokens)} ${formatPercent(category.tokens, snap.maxTokens)}`.trim(),
259            category.kind === 'used' ? category.color : undefined,
260          ),
261        )}
262      </Box>
263      {snap.groups.map((usageGroup) => groupBlock($, site, usageGroup))}
264    </Box>
265  )
266}
267
268export const register: Register = (on) => {
269  on('session.start', async ($, e, next) => {
270    await $.command.register({
271      name: PLUGIN,
272      description: 'Show the context usage pane, or turn it off and on',
273      argumentHint: '[on|off]',
274    })
275    if (await isEnabled($)) void $.ui.open({ id: PANE, title: TITLE })
276    const started = await next(e)
277    await autoRecount($)
278    return started
279  })
280
281  on('command.run', { command: PLUGIN }, async ($, e) => {
282    const arg = e.args.trim().toLowerCase()
283    if (arg === 'off') {
284      await turnOff($)
285      return { text: 'ctx-usage-dock is off. Run /ctx-usage-dock on to turn it on again.' }
286    }
287    if (arg !== '' && arg !== 'on') return { text: 'Usage: /ctx-usage-dock [on|off]' }
288    await turnOn($)
289    return { text: 'ctx-usage-dock is on.' }
290  })
291
292  on('turn.complete', async ($, e, next) => {
293    const done = await next(e)
294    if (e.agentId === undefined) await autoRecount($)
295    return done
296  })
297
298  on('session.compact', async ($, e, next) => {
299    const compacted = await next(e)
300    if (e.agentId === undefined) await autoRecount($)
301    return compacted
302  })
303
304  // `/clear`, `/resume` and a fork switch sessions without another `session.start`.
305  on('classic.SessionStart', async ($, e, next) => {
306    const started = await next(e)
307    await autoRecount($)
308    return started
309  }).catch(($, e, next) => next(e))
310
311  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
312    const { Box, Text, Button } = $.ui.resolve(e)
313    const snap = await read($, snapshot)
314    const counting = await read($, isCounting)
315    const width = Math.max(10, e.props.bodyColumns)
316
317    return (
318      <Box flexDirection="column">
319        {snap ? snapshotView($, e, snap, width) : <Text dimColor>No count yet. Press Refresh.</Text>}
320        <Box flexDirection="row" gap={1} marginTop={1}>
321          <Button label="Refresh" hotkey="r" onPress={() => void recount($, 'summary')} />
322          <Button label="Full count" hotkey="f" onPress={() => void recount($, 'full')} />
323          <Button label="Turn off" hotkey="o" onPress={() => void turnOff($)} />
324          {counting && <Text dimColor>counting...</Text>}
325        </Box>
326      </Box>
327    )
328  })
329}
330

1/** How a count was made: `summary` estimates locally, `full` asks the token-count API. */
2export type CountDetail = 'summary' | 'full'
3
4/** What a /context row is: content in the window, free space, the compaction buffer, or on-demand tool schemas. */
5export type CategoryKind = 'used' | 'free' | 'buffer' | 'deferred'
6
7/** One line under a group: a skill, an MCP server, a memory file. */
8export type UsageItem = {
9  label: string
10  tokens: number
11  /** What follows the tokens: `3/20 loaded`. */
12  note?: string
13}
14
15/** A list /context prints beneath its grid: Skills, MCP tools, Memory files, Rules, Agents. */
16export type UsageGroup = {
17  title: string
18  tokens: number
19  /** What follows the title: `40/45 listed`. */
20  note?: string
21  items: UsageItem[]
22}
23
24/** One row /context prints beside its grid. */
25export type UsageCategory = {
26  name: string
27  tokens: number
28  /** The theme key /context draws the row in. */
29  color: string
30  kind: CategoryKind
31}
32
33/** The context window as the last count broke it down. */
34export type UsageSnapshot = {
35  detail: CountDetail
36  /** When the count finished, in epoch milliseconds. */
37  countedAt: number
38  model: string
39  totalTokens: number
40  maxTokens: number
41  percentage: number
42  autoCompactAt?: number
43  categories: UsageCategory[]
44  groups: UsageGroup[]
45}
46
47declare module 'claude-code' {
48  interface PluginState {
49    'ctx-usage-dock': {
50      /** The last count; null before the first one lands. */
51      snapshot: Shaped<UsageSnapshot | null>
52      /** True while a count runs. */
53      isCounting: boolean
54    }
55  }
56}
57