xray-mod

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Row, Shift, Snapshot } from '../types'
5
6const PANE = 'xray-mod'
7const TITLE = 'Context'
8const LOG_SIZE = 12
9// Ignore estimate jitter below this many tokens when logging a shift.
10const NOISE = 100
11
12const snapshot = atom({ plugin: 'xray-mod', key: 'snapshot' } as const, null)
13const log = atom({ plugin: 'xray-mod', key: 'log' } as const, [])
14
15export const fmt = (n: number): string =>
16  Math.abs(n) >= 1000 ? `${(n / 1000).toFixed(1)}k` : String(n)
17
18const signed = (n: number): string => (n > 0 ? `+${fmt(n)}` : `-${fmt(-n)}`)
19
20// Names the call in the shift log: which skill, which search, which file.
21export const describe = (tool: string, input: Readonly<Record<string, unknown>>): string => {
22  const arg =
23    tool === 'Skill' ? input.skill
24    : tool === 'ToolSearch' ? input.query
25    : tool === 'Agent' ? input.subagent_type
26    : tool === 'Read' ? String(input.file_path ?? '').split(/[\\/]/).pop()
27    : undefined
28
29  return typeof arg === 'string' && arg !== '' ? `${tool}(${arg})` : tool
30}
31
32export const diffRows = (
33  rows: Omit<Row, 'delta'>[],
34  previous: Snapshot | null,
35): Row[] =>
36  rows.map(row => {
37    const before = previous?.rows.find(one => one.name === row.name)
38
39    return { ...row, delta: before ? row.tokens - before.tokens : 0 }
40  })
41
42export const summarize = (rows: Row[]): string =>
43  rows
44    .filter(row => (row.kind === 'used' || row.kind === 'deferred') && Math.abs(row.delta) >= NOISE)
45    .map(row => `${row.name} ${signed(row.delta)}`)
46    .join(', ')
47
48const pushShift = ($: EngineInterface, shift: Shift) =>
49  update($, log, list => [...list, shift].slice(-LOG_SIZE))
50
51const measure = async (
52  $: EngineInterface,
53  detail: 'summary' | 'full',
54  cause: string,
55  isForced = false,
56) => {
57  const usage = await $.session.usage({ breakdown: detail })
58  const b = usage.context.breakdown
59  if (!b) return
60
61  const previous = await read($, snapshot)
62  const rows = diffRows(
63    b.categories.map(c => ({ name: c.name, kind: c.kind, color: c.color, tokens: c.tokens })),
64    previous,
65  )
66
67  const servers = new Map<string, { loaded: number; deferred: number; tokens: number }>()
68  for (const tool of b.mcpTools) {
69    const s = servers.get(tool.serverName) ?? { loaded: 0, deferred: 0, tokens: 0 }
70    if (tool.isLoaded) s.loaded += 1
71    else s.deferred += 1
72    s.tokens += tool.tokens
73    servers.set(tool.serverName, s)
74  }
75
76  const at = await $.clock.now()
77  const next: Snapshot = {
78    at,
79    detail,
80    model: b.model,
81    total: b.totalTokens,
82    max: b.rawMaxTokens,
83    percent: b.percentage,
84    autoCompactAt: b.autoCompactThreshold,
85    rows,
86    memory: [...b.memoryFiles].sort((x, y) => y.tokens - x.tokens),
87    mcp: [...servers].map(([server, s]) => ({ server, ...s })).sort((x, y) => y.tokens - x.tokens),
88    skills: b.skills && {
89      included: b.skills.includedSkills,
90      total: b.skills.totalSkills,
91      tokens: b.skills.tokens,
92      top: [...b.skills.skillFrontmatter]
93        .sort((x, y) => y.tokens - x.tokens)
94        .slice(0, 5)
95        .map(s => ({ name: s.name, tokens: s.tokens })),
96    },
97    agents: b.agents.map(a => ({ name: a.agentType, tokens: a.tokens })),
98  }
99  await update($, snapshot, () => next)
100
101  const changes = summarize(rows)
102  if (changes !== '' || isForced) {
103    await pushShift($, { at, cause, changes: changes || 'no change' })
104  }
105}
106
107// Measurements run one at a time, in order, so each diff is against the last.
108let queue: Promise<void> = Promise.resolve()
109// Main-thread tool calls since the last measurement, named for the shift log.
110const causes: string[] = []
111
112const enqueue = (
113  $: EngineInterface,
114  detail: 'summary' | 'full',
115  cause: string,
116  isForced = false,
117) => {
118  queue = queue.then(() => measure($, detail, cause, isForced)).catch(() => {})
119}
120
121const openPane = async ($: EngineInterface) => {
122  await $.ui.open({ id: PANE, title: TITLE })
123  enqueue($, 'full', 'opened')
124}
125
126// What both /xray-mod and /xrm do: open the pane (which starts an exact count).
127// The engine prefixes command output with the plugin name, so the text omits it.
128const openCommand = async ($: EngineInterface) => {
129  await openPane($)
130
131  return { text: 'Opened; exact count running.' }
132}
133
134export const register: Register = on => {
135  on('session.start', async ($, e, next) => {
136    await $.command.register({
137      name: 'xray-mod',
138      description: 'Open the context x-ray: every token, traced to its source',
139    })
140    await $.command.register({ name: 'xrm', description: 'Same as /xray-mod' })
141
142    return next(e)
143  })
144
145  // /xray-mod and its short alias /xrm share one handler (openCommand, above).
146  on('command.run', { command: 'xray-mod' }, openCommand)
147  on('command.run', { command: 'xrm' }, openCommand)
148
149  // Note which main-thread tools ran; the next model step measures what they
150  // added. A skill's body or a searched tool's schema enters the window after
151  // the tool result, so measuring at the result would miss it.
152  on('tool.call', async ($, e, next) => {
153    const ran = await next(e)
154    if (!e.agentId) {
155      // A tool's arguments sit on the event itself (e.skill, e.query, ...).
156      causes.push(describe(e.tool, e as unknown as Record<string, unknown>))
157    }
158
159    return ran
160  }).catch(($, e, next) => next(e)) // an observer never blocks a tool
161
162  // turn.step streams, so its hook is an async generator forwarding the chunks.
163  on('turn.step', async function* ($, e, next) {
164    if (e.agentId) {
165      // Subagents run in windows of their own: log their steps, don't measure them.
166      await pushShift($, {
167        at: await $.clock.now(),
168        cause: `subagent step ${e.index}`,
169        changes: `${e.messageCount} msgs, own window`,
170        agentId: e.agentId,
171      })
172    } else if (e.index > 0) {
173      // Main thread, mid-turn: everything since the last request is in the window now.
174      const cause = causes.splice(0).join(', ') || 'between requests'
175      enqueue($, 'summary', cause)
176    }
177
178    return yield* next(e)
179  })
180
181  // After each main turn, when the context fill moved: the exact count.
182  // Rate-limit and cost ticks fire this too; a full count there would be waste.
183  on('session.measure', ($, e, next) => {
184    if (e.changed.includes('context')) {
185      enqueue($, 'full', causes.splice(0).join(', ') || 'turn end')
186    }
187
188    return next(e)
189  })
190
191  on('session.compact', async ($, e, next) => {
192    const done = await next(e)
193    if (!e.agentId) enqueue($, 'summary', `compaction (${e.trigger})`, true)
194
195    return done
196  }).catch(($, e, next) => next(e))
197
198  // The opener: a live percentage pill beside the footer's mode labels.
199  on('ui.render', { component: 'SessionMode' }, async ($, e, next) => {
200    const { Box, Button } = $.ui.resolve(e)
201    const snap = await read($, snapshot)
202    const own = await next(e)
203    const label = snap ? `◧ ${snap.percent}%` : '◧ context'
204
205    return (
206      <Box>
207        {own}
208        <Button key="xray-mod" label={label} onPress={() => openPane($)} />
209      </Box>
210    )
211  })
212
213  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
214    const { Box, Text } = $.ui.resolve(e)
215    const snap = await read($, snapshot)
216    const shifts = await read($, log)
217
218    if (!snap) return <Text dimColor>Measuring…</Text>
219
220    const used = snap.rows.filter(row => row.kind !== 'deferred')
221    const deferred = snap.rows.filter(row => row.kind === 'deferred')
222
223    return (
224      <Box flexDirection="column">
225        <Text bold>
226          {fmt(snap.total)} / {fmt(snap.max)} ({snap.percent}%) · {snap.model}
227        </Text>
228        <Text dimColor>
229          {snap.detail === 'full' ? 'exact count' : 'estimate'}
230          {snap.autoCompactAt ? ` · autocompact at ${fmt(snap.autoCompactAt)}` : ' · autocompact off'}
231        </Text>
232
233        <Text> </Text>
234        {used.map(row => (
235          <Text key={row.name}>
236            <Text color={row.color}>■ </Text>
237            {row.name.padEnd(22)} {fmt(row.tokens).padStart(7)}{' '}
238            {Math.abs(row.delta) >= NOISE && <Text color={row.delta > 0 ? 'warning' : 'success'}>{signed(row.delta)}</Text>}
239          </Text>
240        ))}
241        {deferred.map(row => (
242          <Text key={row.name} dimColor>
243            □ {row.name.padEnd(22)} {fmt(row.tokens).padStart(7)} (outside window)
244          </Text>
245        ))}
246
247        {snap.memory.length > 0 && <Text bold>{'\n'}Memory files</Text>}
248        {snap.memory.slice(0, 6).map(m => (
249          <Text key={m.path} dimColor>
250            {fmt(m.tokens).padStart(6)} {m.type.padEnd(8)} {m.path}
251          </Text>
252        ))}
253
254        {snap.skills && (
255          <Text bold>
256            {'\n'}Skills {snap.skills.included}/{snap.skills.total} listed · {fmt(snap.skills.tokens)}
257          </Text>
258        )}
259        {snap.skills?.top.map(s => (
260          <Text key={s.name} dimColor>
261            {fmt(s.tokens).padStart(6)} {s.name}
262          </Text>
263        ))}
264
265        {snap.mcp.length > 0 && <Text bold>{'\n'}MCP servers</Text>}
266        {snap.mcp.slice(0, 6).map(s => (
267          <Text key={s.server} dimColor>
268            {fmt(s.tokens).padStart(6)} {s.server} ({s.loaded} loaded, {s.deferred} deferred)
269          </Text>
270        ))}
271
272        <Text bold>{'\n'}Shifts</Text>
273        {shifts.length === 0 && <Text dimColor>None yet.</Text>}
274        {[...shifts].reverse().map((s, i) => (
275          <Text key={`${s.at}-${i}`} dimColor={!!s.agentId}>
276            {new Date(s.at).toTimeString().slice(0, 5)} {s.cause} → {s.changes}
277          </Text>
278        ))}
279      </Box>
280    )
281  })
282}
283

1export type Row = {
2  name: string
3  kind: 'used' | 'free' | 'buffer' | 'deferred'
4  color: string
5  tokens: number
6  delta: number
7}
8
9export type Snapshot = {
10  at: number
11  detail: 'summary' | 'full'
12  model: string
13  total: number
14  max: number
15  percent: number
16  autoCompactAt?: number
17  rows: Row[]
18  memory: { path: string; type: string; tokens: number }[]
19  mcp: { server: string; loaded: number; deferred: number; tokens: number }[]
20  skills?: { included: number; total: number; tokens: number; top: { name: string; tokens: number }[] }
21  agents: { name: string; tokens: number }[]
22}
23
24export type Shift = {
25  at: number
26  cause: string
27  changes: string
28  agentId?: string
29}
30
31declare module 'claude-code' {
32  interface PluginState {
33    'xray-mod': { snapshot: Snapshot | null; log: Shift[] }
34  }
35}
36