context-tree

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { CtxNode, CtxView } from '../types'
5import { PANE, VIEW, VIEW_LIMIT, allIds, languageOf, nvimCommand, paneRows, splitPaneId, toggled, usageText } from './pane'
6import { CHUNK_AWK, CHUNK_CHARS, HEAD_AWK, chunkLines, startDrafts, turnDrafts, withoutPrompt } from './transcript'
7import {
8  START,
9  addToTurn,
10  annotateUsage,
11  attachmentDraft,
12  fillStart,
13  firstLine,
14  isStartFilled,
15  mentionNode,
16  prependTurns,
17  startNode,
18  toolDraft,
19  turnCount,
20  turnDraft,
21  withStart,
22} from './tree'
23import type { Draft } from './tree'
24
25const tree = atom({ plugin: 'context-tree', key: 'tree' } as const, [])
26const expanded = atom({ plugin: 'context-tree', key: 'expanded' } as const, [START])
27const viewing = atom({ plugin: 'context-tree', key: 'viewing' } as const, null)
28const pending = atom({ plugin: 'context-tree', key: 'pending' } as const, [])
29const usage = atom({ plugin: 'context-tree', key: 'usage' } as const, null)
30
31type $ = EngineInterface
32
33/** Writes text to a capture file of this session and answers its path. */
34const capture = async ($: $, name: string, text: string) => {
35  const tmp = ((await $.env.get('TMPDIR')) ?? '/tmp/').replace(/\/?$/, '/')
36  const dir = `${tmp}claude-context-tree/${await $.session.id()}`
37  const path = `${dir}/${Date.now()}-${crypto.randomUUID().slice(0, 6)}-${name}`
38  await $.fs.write(path, text)
39  return path
40}
41
42/** Writes the capture files of a draft and its children, and answers the nodes. */
43const materialize = async ($: $, draft: Draft): Promise<CtxNode> => {
44  const { capture: text, children, ...n } = draft
45  const kids = await Promise.all(children.map(c => materialize($, c)))
46  if (!text) return { ...n, children: kids }
47  return { ...n, path: await capture($, text.file, text.text), isCapture: true, children: kids }
48}
49
50/** Reads the context usage, and the token counts of the instruction files. */
51const refreshUsage = async ($: $, event: string) => {
52  const now = await $.session.usage({ breakdown: 'summary' })
53  const { percent, tokens: used, window, breakdown } = now.context
54  await update($, usage, () => ({ percent, tokens: used, window }))
55  await update($, tree, list => annotateUsage(list, breakdown?.memoryFiles ?? [], event))
56}
57
58const transcriptPath = async ($: $) => {
59  const config = (await $.env.get('CLAUDE_CONFIG_DIR')) ?? `${await $.env.get('HOME')}/.claude`
60  const slug = (await $.session.cwd()).replace(/[^a-zA-Z0-9]/g, '-')
61  return `${config}/projects/${slug}/${await $.session.id()}.jsonl`
62}
63
64let loading: Promise<void> | undefined
65// True after a /resume in this process: session.start does not fire again, so the next turn rebuilds the tree.
66let isResumed = false
67
68/** Fills the session start node from the head of the transcript, once per conversation. */
69const loadStart = ($: $, event: string) => {
70  loading ??= readStart($, event).finally(() => {
71    loading = undefined
72  })
73  return loading
74}
75
76const readStart = async ($: $, event: string) => {
77  await update($, tree, list => withStart(list, event))
78  if (isStartFilled(await read($, tree))) return
79  const head = await $.process.run(['awk', HEAD_AWK, await transcriptPath($)]).catch(() => undefined)
80  if (!head || head.exitCode !== 0) return
81  const drafts = startDrafts(head.stdout, event)
82  if (!drafts) return
83  const children = await Promise.all(drafts.map(d => materialize($, d)))
84  await update($, tree, list => fillStart(list, children))
85  await refreshUsage($, event)
86}
87
88/** Reads every line of the transcript, in chunks under the read limit of the engine. */
89const readTranscript = async ($: $) => {
90  const path = await transcriptPath($)
91  const lines: string[] = []
92  for (let from = 0; ; ) {
93    const chunk = await $.process.run(['awk', '-v', `from=${from}`, '-v', `max=${CHUNK_CHARS}`, CHUNK_AWK, path])
94    if (chunk.exitCode !== 0) return undefined
95    if (!chunk.stdout) return lines
96    const { lines: got, consumed } = chunkLines(chunk.stdout, chunk.isStdoutTruncated)
97    if (consumed === 0) return lines
98    lines.push(...got)
99    from += consumed
100  }
101}
102
103/**
104 * Adds the turns of a resumed conversation from its transcript, before the turns of this process.
105 * `prompt` is the prompt of a turn that starts now: the file holds it, and its own hooks add it.
106 */
107const rebuildTurns = async ($: $, event: string, prompt?: string) => {
108  if (turnCount(await read($, tree)) > 0) return
109  const lines = await readTranscript($).catch(() => undefined)
110  if (!lines) return
111  const drafts = turnDrafts(lines)
112  const turns = await Promise.all((prompt === undefined ? drafts : withoutPrompt(drafts, prompt)).map(d => materialize($, d)))
113  if (turns.length > 0) await update($, tree, list => prependTurns(list, turns, event))
114}
115
116/** Adds a draft to the current turn. */
117const addDraft = async ($: $, draft: Draft, event: string, agentId?: string) => {
118  const n = await materialize($, draft)
119  await update($, tree, list => addToTurn(list, n, event, agentId))
120}
121
122/** Opens the file in nvim, in a new herdr pane to the right of this one. */
123const openInHerdr = async ($: $, view: CtxView) => {
124  const pane = await $.env.get('HERDR_PANE_ID')
125  if (!pane) {
126    $.ui.toast('context-tree: HERDR_PANE_ID is not set; run Claude Code inside herdr')
127    return
128  }
129  const bin = (await $.env.get('HERDR_BIN_PATH')) ?? 'herdr'
130  const split = await $.process.run([bin, 'pane', 'split', pane, '--direction', 'right', '--focus'])
131  const id = split.exitCode === 0 ? splitPaneId(split.stdout) : undefined
132  if (!id) {
133    $.ui.toast(`context-tree: herdr pane split failed: ${firstLine(split.stderr || split.stdout)}`)
134    return
135  }
136  await $.process.run([bin, 'pane', 'run', id, nvimCommand(view)])
137}
138
139const openViewer = async ($: $, view: CtxView) => {
140  await update($, viewing, () => view)
141  await $.ui.open({ id: VIEW, title: view.label.slice(0, 60), focus: true, closeOnEscape: true, holdToasts: true })
142}
143
144export const register: Register = on => {
145  on('session.start', async ($, e, next) => {
146    await $.command.register({ name: 'context-tree', description: 'Show the context window as a tree in a pane' })
147    await update($, tree, list => withStart(list, 'session.start'))
148    void loadStart($, 'session.start').then(() => refreshUsage($, 'session.start')).catch(() => undefined)
149    void rebuildTurns($, 'session.start').catch(() => undefined)
150    void $.ui.open({ id: PANE, title: 'Context' })
151    return next(e)
152  })
153
154  on('command.run', { command: 'context-tree' }, async $ => {
155    await $.ui.open({ id: PANE, title: 'Context' })
156    return { text: 'Context tree opened.' }
157  })
158
159  on('session.end', async ($, e, next) => {
160    // A /resume in this process loads another conversation: the next turn rebuilds it from its transcript.
161    if (e.reason === 'clear' || e.reason === 'resume') {
162      await update($, tree, () => [startNode('session.end')])
163      await update($, expanded, () => [START])
164      await update($, usage, () => null)
165      isResumed = e.reason === 'resume'
166    }
167    return next(e)
168  })
169
170  on('prompt.attachment', async ($, e, next) => {
171    const out = await next(e)
172    if (out.text === null) return out
173    const hookEvent = e.origin.kind === 'hook' ? e.origin.event : undefined
174    const input = { type: e.type, text: out.text ?? e.text, agentId: e.agentId, hookEvent }
175    const draft = attachmentDraft(turnCount(await read($, tree)) <= 1, input, 'prompt.attachment')
176    if (draft) await addDraft($, draft, 'prompt.attachment', e.agentId)
177    return out
178  })
179
180  on('prompt.mention', async ($, e, next) => {
181    const out = await next(e)
182    if (!e.agentId) await update($, pending, list => [...list, mentionNode(e.mention, e.path, e.offset)])
183    return out
184  }).catch(($, e, next) => next(e))
185
186  on('turn.start', async ($, e, next) => {
187    const out = await next(e)
188    if (isResumed) {
189      isResumed = false
190      await rebuildTurns($, 'turn.start', e.text).catch(() => undefined)
191    }
192    const mentions = await read($, pending)
193    await update($, pending, () => [])
194    const turn = await materialize($, turnDraft(turnCount(await read($, tree)) + 1, e.turnId, 'turn.start', e.text, mentions))
195    await update($, tree, list => [...withStart(list, 'turn.start'), turn])
196    return out
197  })
198
199  on('turn.complete', async ($, e, next) => {
200    const out = await next(e)
201    void loadStart($, 'turn.complete').then(() => refreshUsage($, 'turn.complete')).catch(() => undefined)
202    return out
203  })
204
205  on('tool.call', async ($, e, next) => {
206    const ran = await next(e)
207    void loadStart($, 'tool.call').catch(() => undefined)
208    await addDraft($, toolDraft(e as unknown as Record<string, unknown>, ran, 'tool.call'), 'tool.call', e.agentId)
209    return ran
210  }).catch(($, e, next) => next(e))
211
212  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
213    const { Box, Text, Button } = $.ui.resolve(e)
214    const list = await read($, tree)
215    const open = new Set(await read($, expanded))
216    const now = await read($, usage)
217
218    const rows = paneRows(list, open, e.props.bodyColumns).map(row => (
219      <Box key={`row:${row.id}`} flexDirection="column">
220        <Box flexDirection="row">
221          <Text>{row.indent}</Text>
222          <Button
223            key={`n:${row.id}`}
224            plain
225            dimColor={!row.hasKids && !row.view}
226            onPress={() =>
227              row.hasKids
228                ? update($, expanded, ids => toggled(ids, row.id))
229                : row.view
230                  ? openViewer($, row.view)
231                  : undefined
232            }
233          >
234            {row.label}
235          </Button>
236          {row.view && (
237            <Button key={`o:${row.id}`} plain dimColor onPress={() => row.view && openInHerdr($, row.view)}>
238              {' ↗'}
239            </Button>
240          )}
241        </Box>
242        {row.hint && (
243          <Text dimColor wrap="truncate-end">
244            {`${row.indent}    ${row.hint}`}
245          </Text>
246        )}
247      </Box>
248    ))
249
250    return (
251      <Box flexDirection="column">
252        <Box flexDirection="row" gap={1}>
253          <Text bold>Context</Text>
254          <Text dimColor>{usageText(now)}</Text>
255        </Box>
256        <Box flexDirection="row" gap={1}>
257          <Button key="expand" plain dimColor onPress={() => update($, expanded, () => allIds(list))}>
258            expand all
259          </Button>
260          <Button key="collapse" plain dimColor onPress={() => update($, expanded, () => [])}>
261            collapse all
262          </Button>
263        </Box>
264        {rows}
265      </Box>
266    )
267  })
268
269  on('ui.render', { component: 'Pane', requestId: VIEW }, async ($, e) => {
270    const { Box, Text, Button, Code } = $.ui.resolve(e)
271    const view = await read($, viewing)
272    if (!view) return <Text dimColor>Nothing selected.</Text>
273    let source: string
274    try {
275      source = await $.fs.read(view.path)
276    } catch (err) {
277      source = `Cannot read ${view.path}: ${String(err)}`
278    }
279    const isCut = source.length > VIEW_LIMIT
280    return (
281      <Box flexDirection="column">
282        <Box flexDirection="row" gap={1}>
283          <Button key="nvim" hotkey="o" autoFocus onPress={() => openInHerdr($, view)}>
284            open in nvim
285          </Button>
286          <Button key="close" hotkey="q" role="dismiss" onPress={() => $.ui.close({ id: VIEW })}>
287            close
288          </Button>
289        </Box>
290        <Text dimColor wrap="truncate-start">{view.path}</Text>
291        {isCut && <Text color="warning">{`Shows the first ${VIEW_LIMIT} characters; open in nvim for the rest.`}</Text>}
292        <Code source={source.slice(0, VIEW_LIMIT)} language={languageOf(view.path)} startLine={1} />
293      </Box>
294    )
295  })
296}
297

1import type { CtxNode, CtxUsage, CtxView } from '../types'
2import { tokens } from './tree'
3
4// What the panes show, as plain values. register.tsx turns them into elements,
5// because only a hook holds the element table of the surface.
6
7export const PANE = 'context-tree'
8export const VIEW = 'context-view'
9export const VIEW_LIMIT = 100_000
10
11/** One line of the tree pane. */
12export type PaneRow = {
13  id: string
14  indent: string
15  label: string
16  /** The dim line under the label: the hint, and the event when it differs from the event of the parent. */
17  hint: string
18  hasKids: boolean
19  /** The file a press opens; nothing for a node without one. */
20  view?: CtxView
21}
22
23/** The lines of the tree pane: each node, and the children of each open node. */
24export const paneRows = (list: CtxNode[], open: Set<string>, width: number): PaneRow[] => {
25  const rows: PaneRow[] = []
26  const walk = (nodes: CtxNode[], depth: number, parentEvent?: string) => {
27    for (const n of nodes) {
28      const hasKids = n.children.length > 0
29      const isOpen = open.has(n.id)
30      const mark = hasKids ? (isOpen ? '▾ ' : '▸ ') : '  '
31      const indent = '  '.repeat(depth)
32      const room = Math.max(8, width - indent.length - 4)
33      const text = `${mark}${n.label}`
34      const label = text.length > room ? `${text.slice(0, room - 1)}…` : text
35      // A child shows its event only when it differs from the event of its parent.
36      const event = n.event !== parentEvent ? n.event : undefined
37      const hint = [n.hint, event && `on: ${event}`].filter(Boolean).join(' · ')
38      const view = n.path ? { label: n.label, path: n.path, line: n.line, isCapture: n.isCapture } : undefined
39      rows.push({ id: n.id, indent, label, hint, hasKids, view })
40      if (isOpen) walk(n.children, depth + 1, n.event)
41    }
42  }
43  walk(list, 0)
44  return rows
45}
46
47/** The ids of every node that has children. */
48export const allIds = (list: CtxNode[]): string[] =>
49  list.flatMap(n => (n.children.length ? [n.id, ...allIds(n.children)] : []))
50
51/** The ids with `id` added, or removed when they hold it. */
52export const toggled = (ids: string[], id: string) => (ids.includes(id) ? ids.filter(x => x !== id) : [...ids, id])
53
54/** The fill of the context window, for the head of the tree pane. */
55export const usageText = (now: CtxUsage | null) =>
56  now ? `${now.percent ?? '?'}% · ${now.tokens ? tokens(now.tokens) : '?'} of ${tokens(now.window)}` : 'no usage yet'
57
58const LANGS: Record<string, string> = {
59  ts: 'typescript', tsx: 'tsx', js: 'javascript', json: 'json', md: 'markdown', sh: 'bash',
60  py: 'python', go: 'go', rs: 'rust', lua: 'lua', toml: 'toml', yaml: 'yaml', yml: 'yaml',
61}
62
63/** The language the viewer highlights a file in, by its extension. */
64export const languageOf = (path: string) => LANGS[path.split('.').pop() ?? '']
65
66const shellQuote = (s: string) => `'${s.replace(/'/g, `'\\''`)}'`
67
68/** The command a herdr pane runs to open the file in nvim; a capture opens read-only. */
69export const nvimCommand = (view: CtxView) => {
70  const args = [view.isCapture ? '-R' : '', view.line ? `+${view.line}` : '', shellQuote(view.path)]
71  return `exec nvim ${args.filter(Boolean).join(' ')}`
72}
73
74/** The id of the pane that `herdr pane split` printed, or nothing. */
75export const splitPaneId = (stdout: string): string | undefined => {
76  try {
77    const id: unknown = JSON.parse(stdout)?.result?.pane?.pane_id
78    return typeof id === 'string' ? id : undefined
79  } catch {
80    return undefined
81  }
82}
83

1import {
2  HIDDEN_TYPES,
3  START_ATTACHMENTS,
4  START_BLOCKS,
5  START_FILES,
6  START_OWN_GROUP,
7  START_SYSTEM,
8  START_TOOLS,
9  attachmentDraft,
10  captureDraft,
11  estimate,
12  firstLine,
13  node,
14  toolDraft,
15  turnDraft,
16} from './tree'
17import type { Draft } from './tree'
18
19// Reads the session transcript (the jsonl file) into drafts of the tree.
20// register.tsx runs the reads; the functions here only parse what they print.
21
22/** The event the tree shows on the nodes it rebuilt from the transcript. */
23const EVENT = 'transcript'
24
25// Prints the transcript up to the first prompt snapshot that lists the tools:
26// the rows of the session start and the first prompt.
27export const HEAD_AWK = '{ print } /"prompt_snapshot"/ && /"tools":\\[/ { exit } NR >= 400 { exit }'
28
29// Prints the lines after line `from` until they hold about `max` characters.
30// A read of the engine stops at 4 MiB, so register.tsx reads the file in chunks.
31export const CHUNK_AWK = 'NR > from { print; n += length($0) + 1; if (n >= max) exit }'
32export const CHUNK_CHARS = 3 * 1024 * 1024
33
34/**
35 * The complete lines of a chunk, and the number of lines it moves past.
36 * A cut chunk ends in a part of a line over the read limit: the read skips that line.
37 */
38export const chunkLines = (stdout: string, isCut: boolean) => {
39  const lines = stdout.split('\n')
40  // A chunk that is not cut ends with a newline; a cut chunk ends with the part of a line.
41  lines.pop()
42  return { lines, consumed: lines.length + (isCut ? 1 : 0) }
43}
44
45type Attachment = Record<string, unknown> & { type: string }
46type Block = { type?: string; text?: unknown; id?: string; name?: string; input?: unknown; tool_use_id?: string; content?: unknown; is_error?: boolean }
47type Row = {
48  type?: string
49  uuid?: string
50  promptId?: string
51  isMeta?: boolean
52  isSidechain?: boolean
53  toolUseResult?: unknown
54  message?: { content?: string | Block[] }
55  attachment?: Attachment
56  rendered?: { content?: unknown }[]
57}
58
59const parseRows = (lines: string[]): Row[] =>
60  lines.flatMap(line => {
61    try {
62      return [JSON.parse(line) as Row]
63    } catch {
64      return []
65    }
66  })
67
68/** The text of a message content: a string, or the text of its blocks. */
69const contentText = (content: unknown) =>
70  typeof content === 'string'
71    ? content
72    : Array.isArray(content)
73      ? content.map(b => (typeof b?.text === 'string' ? b.text : '')).join('\n')
74      : ''
75
76/** The text the model read for a transcript row, as the engine rendered it. */
77const renderedText = (row: Row) => (row.rendered ?? []).map(r => contentText(r.content)).join('\n')
78
79/** A short label for a system prompt section: its first heading or line. */
80const sectionLabel = (text: string) => firstLine(text.replace(/^\s*#+\s*/, ''), 50) || '(empty)'
81
82/** A group node; `texts` are what its children hold, for the token estimate. */
83const groupOf = (id: string, label: string, event: string, children: Draft[], texts: string[]): Draft => ({
84  id,
85  label,
86  hint: `${children.length} items · ${estimate(texts.join(''))}`,
87  event,
88  children,
89})
90
91/** The groups of the session start node, from the head of the transcript; nothing before the first snapshot. */
92export const startDrafts = (text: string, event: string): Draft[] | undefined => {
93  let snapshot: Attachment | undefined
94  let tools: unknown[] = []
95  const atts: Row[] = []
96  for (const row of parseRows(text.split('\n'))) {
97    const a = row.attachment
98    if (row.type !== 'attachment' || !a) continue
99    if (a.type === 'prompt_snapshot') {
100      snapshot ??= a
101      if (Array.isArray(a.tools)) {
102        tools = a.tools
103        break
104      }
105    } else if (!snapshot) {
106      // Rows after the first snapshot belong to the first turn.
107      atts.push(row)
108    }
109  }
110  if (!snapshot) return undefined
111
112  const instructions = atts.find(r => r.attachment?.type === 'instructions')?.attachment
113  const fileList = Array.isArray(instructions?.files) ? (instructions.files as { path: string; type: string; content: string }[]) : []
114  const files = fileList.map(f => node(f.path, { path: f.path, hint: `${f.type} · ${estimate(f.content)}`, event }))
115
116  const context = atts.find(r => r.attachment?.type === 'session_context')?.attachment?.context
117  const blockTexts = Object.entries((context ?? {}) as Record<string, unknown>).map(([name, value]) => ({
118    name,
119    text: typeof value === 'string' ? value : JSON.stringify(value, null, 2),
120  }))
121  const blocks = blockTexts.map(b => captureDraft(event, b.name, `${b.name}.md`, b.text, estimate(b.text)))
122
123  const prefix = typeof snapshot.cliPrefix === 'string' ? [snapshot.cliPrefix] : []
124  const sections = [...prefix, ...(Array.isArray(snapshot.systemPrompt) ? (snapshot.systemPrompt as string[]) : [])]
125  const system = sections.map((text, i) => captureDraft(event, sectionLabel(text), `system-${i}.md`, text, estimate(text)))
126
127  const toolTexts = tools.map(t => ({
128    name: String((t as { name?: unknown }).name ?? 'tool'),
129    text: JSON.stringify(t, null, 2),
130  }))
131  const schemas = toolTexts.map(t =>
132    captureDraft(event, t.name, `tool-${t.name.replace(/[^\w-]/g, '_')}.json`, t.text, estimate(t.text)),
133  )
134
135  const attTexts = atts
136    .filter(r => r.attachment && !START_OWN_GROUP.has(r.attachment.type) && !HIDDEN_TYPES.has(r.attachment.type))
137    .map(r => ({ type: r.attachment!.type, text: renderedText(r) }))
138    .filter(r => r.text)
139  const attachments = attTexts.map(r => captureDraft(event, r.type, `${r.type}.md`, r.text, estimate(r.text)))
140
141  return [
142    groupOf(START_FILES, 'Instruction files', event, files, fileList.map(f => f.content)),
143    groupOf(START_BLOCKS, 'Context blocks', event, blocks, blockTexts.map(b => b.text)),
144    groupOf(START_SYSTEM, 'System prompt', event, system, sections),
145    groupOf(START_TOOLS, 'Tool schemas', event, schemas, toolTexts.map(t => t.text)),
146    groupOf(START_ATTACHMENTS, 'Attachments', event, attachments, attTexts.map(r => r.text)),
147  ]
148}
149
150/** The turns without a last turn that holds `prompt`: the turn that starts now and has its row in the file. */
151export const withoutPrompt = (turns: Draft[], prompt: string) =>
152  turns.at(-1)?.children[0]?.capture?.text === prompt ? turns.slice(0, -1) : turns
153
154/** True for a prompt that a person or a command sent, not a tool result or a row the engine added. */
155const isPrompt = (row: Row) => {
156  if (row.type !== 'user' || row.isMeta || row.toolUseResult !== undefined) return false
157  const content = row.message?.content
158  return typeof content === 'string' || (Array.isArray(content) && !content.some(b => b.type === 'tool_result'))
159}
160
161/**
162 * The turns of the main thread: each prompt, then its tool calls and attachments.
163 * The subagents write their rows to other files, so their calls are not in the turns.
164 */
165export const turnDrafts = (lines: string[]): Draft[] => {
166  const turns: Draft[] = []
167  const calls = new Map<string, { name: string; input: Record<string, unknown> }>()
168  for (const row of parseRows(lines)) {
169    if (row.isSidechain) continue
170    const turn = turns[turns.length - 1]
171    if (isPrompt(row)) {
172      const text = contentText(row.message?.content)
173      if (text.trim()) turns.push(turnDraft(turns.length + 1, row.promptId ?? row.uuid ?? crypto.randomUUID(), EVENT, text))
174      continue
175    }
176    // The rows before the first prompt belong to the session start node.
177    if (!turn) continue
178    const content = row.message?.content
179    if (row.type === 'assistant' && Array.isArray(content)) {
180      for (const b of content) {
181        if (b.type === 'tool_use' && b.id && b.name) calls.set(b.id, { name: b.name, input: (b.input ?? {}) as Record<string, unknown> })
182      }
183    } else if (row.type === 'user' && Array.isArray(content)) {
184      for (const b of content) {
185        const call = b.type === 'tool_result' && b.tool_use_id ? calls.get(b.tool_use_id) : undefined
186        if (!call) continue
187        const ran = { text: contentText(b.content), isError: b.is_error === true }
188        turn.children.push(toolDraft({ ...call.input, tool: call.name }, ran, EVENT))
189      }
190    } else if (row.type === 'attachment' && row.attachment && !START_OWN_GROUP.has(row.attachment.type)) {
191      const text = renderedText(row)
192      const hookEvent = typeof row.attachment.hookEvent === 'string' ? row.attachment.hookEvent : undefined
193      const draft = text && attachmentDraft(turns.length === 1, { type: row.attachment.type, text, hookEvent }, EVENT)
194      if (draft) turn.children.push(draft)
195    }
196  }
197  return turns
198}
199

1import type { CtxNode } from '../types'
2
3// The functions here are pure: they take the tree and answer a new one, or
4// answer drafts of nodes. A module cannot pass `$` across an import, so
5// register.tsx does the reads, the writes and the captures.
6
7export const START = 'start'
8export const START_FILES = 'start:files'
9export const START_BLOCKS = 'start:blocks'
10export const START_SYSTEM = 'start:system'
11export const START_ATTACHMENTS = 'start:attachments'
12export const START_TOOLS = 'start:tools'
13
14// Attachments that describe the session, not the prompt that carried them.
15// The session start node reads them from the transcript, so the turn skips them.
16export const START_TYPES = new Set([
17  'skill_listing',
18  'deferred_tools_delta',
19  'mcp_instructions',
20  'mcp_instructions_delta',
21  'agent_listing_delta',
22  'environment',
23  'model',
24  'auto_mode',
25  'date',
26  'remote_session_change',
27])
28// Attachments the tree never shows, at the session start or in a turn.
29export const HIDDEN_TYPES = new Set(['total_tokens_reminder'])
30// Attachment rows the session start node shows in a group of their own.
31export const START_OWN_GROUP = new Set(['instructions', 'session_context', 'prompt_snapshot'])
32
33/** A node before its capture file exists: `capture` holds the text to write. */
34export type Draft = Omit<CtxNode, 'children'> & { capture?: { file: string; text: string }; children: Draft[] }
35
36export const node = (label: string, extra: Partial<CtxNode> = {}): CtxNode => ({
37  id: crypto.randomUUID(),
38  label,
39  children: [],
40  ...extra,
41})
42
43const group = (id: string, label: string, event: string): CtxNode => ({ id, label, event, children: [] })
44
45export const startNode = (event: string): CtxNode => ({
46  id: START,
47  label: 'Session start',
48  event,
49  children: [
50    group(START_FILES, 'Instruction files', event),
51    group(START_BLOCKS, 'Context blocks', event),
52    group(START_SYSTEM, 'System prompt', event),
53    group(START_TOOLS, 'Tool schemas', event),
54    group(START_ATTACHMENTS, 'Attachments', event),
55  ],
56})
57
58/** Applies `fn` to the node with `id`, wherever it sits. */
59const mapNode = (list: CtxNode[], id: string, fn: (n: CtxNode) => CtxNode): CtxNode[] =>
60  list.map(n => (n.id === id ? fn(n) : { ...n, children: mapNode(n.children, id, fn) }))
61
62const findNode = (list: CtxNode[], id: string): CtxNode | undefined => {
63  for (const n of list) {
64    if (n.id === id) return n
65    const hit = findNode(n.children, id)
66    if (hit) return hit
67  }
68  return undefined
69}
70
71/** A capture node: the viewer opens the file that register.tsx writes for it. */
72export const captureDraft = (event: string, label: string, file: string, text: string, hint?: string): Draft => ({
73  ...node(label, { hint, event }),
74  capture: { file, text },
75})
76
77export const tokens = (n: number) => (n >= 1000 ? `${(n / 1000).toFixed(1)}k tok` : `${n} tok`)
78
79export const firstLine = (text: string, max = 80) => {
80  const line = text.trim().split('\n')[0] ?? ''
81  return line.length > max ? `${line.slice(0, max - 1)}…` : line
82}
83
84export const estimate = (text: string) => `~${tokens(Math.ceil(text.length / 4))}`
85
86/** The tree with the session start node first. */
87export const withStart = (list: CtxNode[], event: string): CtxNode[] =>
88  list.some(n => n.id === START) ? list : [startNode(event), ...list]
89
90/** Adds a node under the current turn, or under a subagent group inside it. */
91export const addToTurn = (list: CtxNode[], child: CtxNode, event: string, agentId?: string): CtxNode[] => {
92  const full = withStart(list, event)
93  // The last turn, or the session start node before the first turn.
94  const root = full[full.length - 1]?.id ?? START
95  if (!agentId) return mapNode(full, root, n => ({ ...n, children: [...n.children, child] }))
96  const agentGroup = `${root}:agent:${agentId}`
97  return mapNode(full, root, n => {
98    const has = n.children.some(c => c.id === agentGroup)
99    const children = has ? n.children : [...n.children, group(agentGroup, `subagent ${agentId.slice(0, 8)}`, event)]
100    return { ...n, children: mapNode(children, agentGroup, g => ({ ...g, children: [...g.children, child] })) }
101  })
102}
103
104/** The number of turns in the tree. */
105export const turnCount = (list: CtxNode[]) => list.filter(n => n.id !== START).length
106
107/** A turn with its prompt and the mentions the prompt held. */
108export const turnDraft = (number: number, id: string, event: string, text?: string, children: Draft[] = []): Draft => {
109  const title = text ? firstLine(text, 60) : '(continuation)'
110  const prompt = text ? [captureDraft(event, 'prompt', 'prompt.txt', text)] : []
111  return { id, label: `Turn ${number}: ${title}`, event, children: [...prompt, ...children] }
112}
113
114export const mentionNode = (mention: string, path: string, offset?: number): CtxNode =>
115  node(`@${mention}`, { path, line: offset, event: 'prompt.mention' })
116
117export type AttachmentInput = { type: string; text: string; agentId?: string; hookEvent?: string }
118
119/** The node for an attachment, or nothing when the session start node holds it or the tree hides it. */
120export const attachmentDraft = (isFirstTurn: boolean, a: AttachmentInput, event: string): Draft | undefined => {
121  const label = a.hookEvent ? `${a.type} (${a.hookEvent} hook)` : a.type
122  const isStart = !a.agentId && (a.hookEvent === 'SessionStart' || (isFirstTurn && START_TYPES.has(a.type)))
123  if (isStart || HIDDEN_TYPES.has(a.type)) return undefined
124  return captureDraft(event, label, `${a.type}.txt`, a.text, estimate(a.text))
125}
126
127export type ToolRun = { deny?: string; text?: string; result?: unknown; isError?: boolean }
128
129/** The node for a tool call: its input, and its output or the reason for the deny. */
130export const toolDraft = (e: Record<string, unknown>, ran: ToolRun, event: string): Draft => {
131  const tool = String(e.tool)
132  if (ran.deny !== undefined) return node(`${tool} (denied)`, { hint: firstLine(ran.deny, 60), event })
133  const out = ran.text ?? JSON.stringify(ran.result, null, 2) ?? ''
134  const str = (k: string) => (typeof e[k] === 'string' ? (e[k] as string) : undefined)
135  const num = (k: string) => (typeof e[k] === 'number' ? (e[k] as number) : undefined)
136  const err = ran.isError === true ? 'error' : undefined
137  const filePath = str('file_path') ?? str('notebook_path')
138  if (filePath && ['Read', 'Write', 'Edit', 'MultiEdit', 'NotebookEdit'].includes(tool)) {
139    const offset = num('offset')
140    const label = `${tool} ${filePath}${offset ? `:${offset}` : ''}`
141    return node(label, { path: filePath, line: offset, hint: err, event })
142  }
143  const input = { ...e }
144  for (const k of ['tool', 'tool_use_id', 'agentId', 'consent']) delete input[k]
145  const output = captureDraft(event, 'output', 'output.txt', out, tokens(Math.ceil(out.length / 4)))
146  if (tool === 'Bash') {
147    const command = str('command') ?? ''
148    return {
149      ...node(`$ ${firstLine(command)}`, { hint: err ?? str('description'), event }),
150      children: [captureDraft(event, 'command', 'command.sh', command), output],
151    }
152  }
153  const summary = str('pattern') ?? str('url') ?? str('skill') ?? str('description') ?? str('query') ?? str('path') ?? ''
154  return {
155    ...node(`${tool} ${firstLine(summary, 60)}`.trim(), { hint: err, event }),
156    children: [captureDraft(event, 'input', 'input.json', JSON.stringify(input, null, 2)), output],
157  }
158}
159
160type MemoryFile = { path: string; type: string; tokens: number }
161
162/** Puts the token counts of the context breakdown on the instruction files. */
163export const annotateUsage = (list: CtxNode[], memoryFiles: MemoryFile[], event: string): CtxNode[] => {
164  const byPath = new Map(memoryFiles.map(f => [f.path, f]))
165  if (byPath.size === 0) return list
166  if (findNode(list, START_FILES)?.children.length === 0) {
167    // The transcript has not filled the session start node yet: list the files the breakdown knows.
168    const files = [...byPath.values()].map(f => node(f.path, { path: f.path, hint: `${f.type} · ${tokens(f.tokens)}`, event }))
169    return mapNode(list, START_FILES, g => ({ ...g, children: files }))
170  }
171  const annotate = (nodes: CtxNode[]): CtxNode[] =>
172    nodes.map(n => {
173      const f = n.path ? byPath.get(n.path) : undefined
174      const hint = f ? `${n.hint?.split(' · ')[0] ?? f.type} · ${tokens(f.tokens)}` : n.hint
175      return { ...n, hint, children: annotate(n.children) }
176    })
177  return mapNode(list, START_FILES, g => ({ ...g, children: annotate(g.children) }))
178}
179
180/** True when the transcript filled the session start node. */
181export const isStartFilled = (list: CtxNode[]) => Boolean(findNode(list, START)?.hint)
182
183/** The tree with `children` in the session start node. */
184export const fillStart = (list: CtxNode[], children: CtxNode[]) =>
185  mapNode(list, START, n => ({ ...n, hint: 'from the transcript', children }))
186
187/** The tree with `turns` after the session start node and before the turns it holds. */
188export const prependTurns = (list: CtxNode[], turns: CtxNode[], event: string) => {
189  const [start, ...rest] = withStart(list, event)
190  return [start!, ...turns, ...rest]
191}
192

1export type CtxNode = {
2  id: string
3  label: string
4  /** Dim text after the label: tokens, a tier, an error mark. */
5  hint?: string
6  /** The file the viewer and nvim open. */
7  path?: string
8  line?: number
9  /** True when `path` is a capture the mod wrote, not a file of the project. */
10  isCapture?: boolean
11  /** The hook event that added the node; the pane shows it as an `on:` child row. */
12  event?: string
13  children: CtxNode[]
14}
15
16export type CtxView = { label: string; path: string; line?: number; isCapture?: boolean }
17
18export type CtxUsage = { percent?: number; tokens?: number; window: number }
19
20declare module 'claude-code' {
21  interface PluginState {
22    'context-tree': {
23      tree: CtxNode[]
24      expanded: string[]
25      viewing: CtxView | null
26      pending: CtxNode[]
27      usage: CtxUsage | null
28    }
29  }
30}
31