convo-tree-2

1import { atom, read, update } from 'claude-code'
2import type { Hook, Register, SessionMessage } from 'claude-code'
3
4import type { Bridge, Tree, View } from '../types'
5import {
6  ROOT,
7  addBranches,
8  branchContext,
9  branchOf,
10  compact,
11  emptyTree,
12  findBranch,
13  finishTurn,
14  isReturning,
15  labelOf,
16  listText,
17  pathTo,
18  patchBranch,
19  ranges,
20  restore,
21  rows,
22  startTurn,
23  strip,
24  switchTo,
25  tokensOf,
26  turnsOf,
27} from './tree'
28import type { Row } from './tree'
29
30type Api = Parameters<Hook<'session.start'>>[0]
31
32const PANE = 'convo-tree-2'
33// How often each surface asked to draw the pane since load, for /tree to report
34const renders: Record<string, number> = {}
35const TITLE = 'Conversation'
36const COLORS = ['claude', 'suggestion', 'success', 'warning', 'planMode', 'autoAccept', 'ide', 'merged'] as const
37
38const tree = atom({ plugin: 'convo-tree-2', key: 'tree' } as const, emptyTree())
39const view = atom({ plugin: 'convo-tree-2', key: 'view' } as const, { openId: null, openTurnId: null } as View)
40const bridge = atom({ plugin: 'convo-tree-2', key: 'bridge' } as const, null as Bridge | null)
41const pending = atom({ plugin: 'convo-tree-2', key: 'pending' } as const, null as string | null)
42
43const snapshotOf = async ($: Api, id: string) =>
44  (await $.state.get({ plugin: 'convo-tree-2', key: 'snapshot', id })).value ?? null
45
46const saveSnapshot = ($: Api, id: string, messages: SessionMessage[]) =>
47  $.state.set({ plugin: 'convo-tree-2', key: 'snapshot', id }, messages)
48
49// The last turn on the main thread has no reply yet
50const isRunning = (t: Tree) => t.turns.at(-1)?.reply === null
51
52// Hands the tree to the browser view's server, when one is up
53const publish = async ($: Api) => {
54  const at = await read($, bridge)
55  if (at === null) return
56  await $.http
57    .fetch(`http://127.0.0.1:${at.port}/tree`, {
58      method: 'POST',
59      headers: { 'x-token': at.token, 'content-type': 'application/json' },
60      body: JSON.stringify(await read($, tree)),
61    })
62    .catch(() => undefined)
63}
64
65// Every tree write goes through here, so the browser view sees it too
66const change = async ($: Api, fn: (t: Tree) => Tree) => {
67  await update($, tree, fn)
68  await publish($)
69}
70
71// Creation order, as the browser page colors them
72const colorOf = (t: Tree, id: string) =>
73  id === ROOT ? 'text' : COLORS[(t.branches.findIndex(b => b.id === id) - 1) % COLORS.length]
74
75const cut = (text: string, room: number) => (text.length <= room ? text : `${text.slice(0, Math.max(0, room - 1))}…`)
76
77const oneLine = (text: string) => text.replace(/\s+/g, ' ').trim()
78
79const haiku = async ($: Api, prompt: string, maxTokens: number) => {
80  const reply = await $.model.complete({ model: 'haiku', prompt, maxTokens, timeoutMs: 20000 })
81  return reply.isAnswered ? reply.text.trim() : undefined
82}
83
84const labelBranch = async ($: Api, id: string, prompt: string) => {
85  const label = await haiku(
86    $,
87    `Give a 2-5 word title for a conversation thread that starts with this message. Answer with the title alone, no quotes.\n\n<message>\n${prompt.slice(0, 1500)}\n</message>`,
88    20,
89  )
90  if (label !== undefined) await change($, t => patchBranch(t, id, { label: cut(oneLine(label), 40) }))
91}
92
93const finishedOf = (t: Tree, id: string) => turnsOf(t, id).filter(x => x.reply !== null)
94
95const summarizeBranch = async ($: Api, id: string) => {
96  const all = finishedOf(await read($, tree), id)
97  const turns = all.slice(-6)
98  if (turns.length === 0) return
99  const body = turns
100    .map(t => `User: ${t.prompt.slice(0, 600)}\nAssistant: ${(t.reply ?? '').slice(0, 600)}`)
101    .join('\n\n')
102  const summary = await haiku(
103    $,
104    `Summarize where this conversation thread stands in one sentence of at most 18 words: what is being explored and what was decided or is open. Answer with the sentence alone.\n\n<thread>\n${body}\n</thread>`,
105    50,
106  )
107  if (summary !== undefined) {
108    await change($, t => patchBranch(t, id, { summary: oneLine(summary), summarizedTurns: all.length }))
109  }
110}
111
112// Titles every untitled branch that has turns and refreshes every stale summary,
113// so work a reload dropped is picked up by the next turn or load
114const catchUp = async ($: Api) => {
115  for (const { id } of (await read($, tree)).branches) {
116    const t = await read($, tree)
117    const branch = branchOf(t, id)
118    const first = turnsOf(t, id).find(x => x.prompt !== '')
119    if (branch === undefined || first === undefined) continue
120    if (branch.label === null) await labelBranch($, id, first.prompt)
121    if (branch.summarizedTurns !== finishedOf(t, id).length) await summarizeBranch($, id)
122  }
123}
124
125let isBusy = false
126
127// Runs catchUp outside the hook's budget, one at a time; a call skipped meanwhile is caught up next turn
128const settle = ($: Api) =>
129  $.clock.after(0, () => {
130    if (isBusy) return
131    isBusy = true
132    void catchUp($)
133      .catch(() => undefined)
134      .finally(() => {
135        isBusy = false
136      })
137  })
138
139const MID_TURN = 'Branch once this turn finishes: its transcript is still being written.'
140
141// Forks the live transcript: the parent and the new branch both start from it, and the session stays on it.
142// Undefined while a turn runs, whose rest would land on the fork
143const newBranch = async ($: Api, label: string | null) => {
144  if (isRunning(await read($, tree))) return undefined
145  const live = strip(await $.session.messages())
146  await saveSnapshot($, (await read($, tree)).activeId, live)
147  let id = ''
148  await change($, t => {
149    const made = addBranches(t, [label])
150    id = made.branches.at(-1)?.id ?? ROOT
151    return switchTo(made, id)
152  })
153  await saveSnapshot($, id, live)
154  await update($, view, () => ({ openId: id, openTurnId: null }))
155  return id
156}
157
158const splitReply = async ($: Api) => {
159  const t = await read($, tree)
160  if (isRunning(t)) {
161    $.ui.toast(MID_TURN)
162    return
163  }
164  const last = finishedOf(t, t.activeId).at(-1)
165  if (last === undefined) {
166    $.ui.toast('Nothing to split: this branch has no reply yet.')
167    return
168  }
169  const answer = await haiku(
170    $,
171    `The assistant reply below may offer several distinct options, ideas or directions to explore. List each as a 2-5 word title, one per line, at most 6, no numbering. If it offers no distinct options, answer NONE.\n\n<reply>\n${last.reply ?? ''}\n</reply>`,
172    120,
173  )
174  const labels = (answer ?? '')
175    .split('\n')
176    .map(line => oneLine(line.replace(/^[-*\d.)\s]+/, '')))
177    .filter(line => line !== '' && line !== 'NONE')
178    .slice(0, 6)
179    .map(line => cut(line, 40))
180  if (labels.length === 0) {
181    $.ui.toast('No distinct options found in the last reply.')
182    return
183  }
184  const live = strip(await $.session.messages())
185  let ids: string[] = []
186  await change($, cur => {
187    const made = addBranches(cur, labels)
188    ids = made.branches.slice(cur.branches.length).map(b => b.id)
189    return made
190  })
191  await Promise.all(ids.map(id => saveSnapshot($, id, live)))
192  $.ui.toast(`Split into ${labels.length} branches.`)
193}
194
195const AGENT = 'convo-tree-2:branch'
196const BRANCH_LINE = /^\s*Branch:\s*(.+?)\s*$/m
197
198const AGENT_PROMPT = `You are a subagent that sees one branch of the user's conversation with Claude, not the whole conversation. Its turns (user prompts and final assistant replies; tool calls omitted) open your task message, under "# Conversation branch". Other branches of the conversation are deliberately left out: work from this branch alone, and say so when the task needs something it does not hold. Finish with a concise answer for the main conversation.`
199
200// The branch's turns, then the task: what a branch subagent is handed
201const withBranch = (t: Tree, id: string, task: string) => `${branchContext(t, id) ?? ''}\n\n---\n\nTask: ${task}`
202
203const recordAgent =
204  (agentId: string, branchId: string, task: string) =>
205  (cur: Tree): Tree => ({
206    ...cur,
207    agents: [...cur.agents, { agentId, branchId, task, isDone: false }],
208  })
209
210const startAgent = async ($: Api, ref: string, task: string) => {
211  const t = await read($, tree)
212  const branch = findBranch(t, ref)
213  if (branch === undefined) return `No branch "${ref}". Branches:\n${listText(t)}`
214  if (task.trim() === '') return `Give the subagent a task: /branch-agent ${branch.id} <task>`
215  const spawned = await $.agent.spawn({
216    subagentType: AGENT,
217    prompt: withBranch(t, branch.id, task),
218    description: `Branch ${branch.id}: ${labelOf(branch)}`,
219  })
220  if (spawned.deny !== undefined) return `Subagent refused: ${spawned.deny}`
221  if (spawned.agentId !== undefined) await change($, recordAgent(spawned.agentId, branch.id, task))
222  return `Started a subagent on branch ${branch.id} "${labelOf(branch)}" with ${turnsOf(t, branch.id).length} of its turns.`
223}
224
225// The branch a switch is installing, and the prompt that started the turn installing it, if one did
226let switching: { id: string; prompt: string | null } | null = null
227
228// Saves the live transcript as the active branch's and puts the asked-for branch's in its place,
229// through a compaction this module answers itself: no summary is made, no model call.
230// The engine runs a plugin's own compaction hook only when another of its hooks asks, and not a
231// command or prompt hook (they hold the turn): so a press, an edit of the prompt or a turn's start.
232const install = async ($: Api, prompt: string | null = null) => {
233  const id = await read($, pending)
234  // A turn's start means the last one is over, even one that ended with no reply
235  if (id === null || switching !== null || (prompt === null && isRunning(await read($, tree)))) return
236  switching = { id, prompt }
237  try {
238    const done = await $.session.compact({ instructions: `convo-tree-2: switch to branch ${id}` })
239    $.ui.toast(done.skip === undefined ? `Switched the session to branch ${id}.` : `Could not switch: ${done.skip}`)
240  } catch {
241    // busy: the next press, edit or turn tries again
242  } finally {
243    switching = null
244  }
245}
246
247// Asks for a switch; the next press, edit of the prompt or turn installs it, before Claude reads anything
248const swapTo = async ($: Api, id: string) => {
249  const t = await read($, tree)
250  const branch = branchOf(t, id)
251  if (branch === undefined) return `No branch "${id}".`
252  if (id === t.activeId) return `Already on branch ${id}.`
253  await update($, pending, () => id)
254  return `Your next prompt goes to branch ${id} "${labelOf(branch)}".`
255}
256
257const showBranch = async ($: Api, id: string) => {
258  const said = await swapTo($, id)
259  $.ui.toast(said)
260}
261
262// The browser view: browser/server.ts, run with bun, serves the page and prints its clicks
263const BUNS = ['bun', '/opt/homebrew/bin/bun', '/usr/local/bin/bun']
264
265const onAction = async ($: Api, action: Record<string, unknown>) => {
266  if (action.action === 'switch' && typeof action.id === 'string') await showBranch($, action.id)
267  if (action.action === 'new' && (await newBranch($, null)) === undefined) $.ui.toast(MID_TURN)
268  if (action.action === 'split') await splitReply($)
269}
270
271// `READY {"port","token"}` once the server listens, `ACTION {...}` per click on the page
272const onLine = async ($: Api, line: string) => {
273  const [, kind, json = ''] = /^(READY|ACTION) (.+)$/.exec(line) ?? []
274  let body: Record<string, unknown>
275  try {
276    body = JSON.parse(json) ?? {}
277  } catch {
278    return
279  }
280  if (kind === 'READY' && typeof body.port === 'number' && typeof body.token === 'string') {
281    const at = { port: body.port, token: body.token }
282    await update($, bridge, () => at)
283    await publish($)
284  }
285  if (kind === 'ACTION') await onAction($, body)
286}
287
288// Runs one server until it exits; false when no bun could start it
289const runServer = async ($: Api) => {
290  const saved = await read($, bridge)
291  const script = `${$.plugin.root}/browser/server.ts`
292  for (const bun of BUNS) {
293    try {
294      const child = $.process.spawn({ argv: [bun, script, String(saved?.port ?? 0), saved?.token ?? ''] })
295      let rest = ''
296      for await (const piece of child) {
297        if (piece.stream !== 'stdout') continue
298        const lines = (rest + piece.text).split('\n')
299        rest = lines.pop() ?? ''
300        for (const line of lines) await onLine($, line)
301      }
302      return true
303    } catch {
304      // not at this path: try the next
305    }
306  }
307  return false
308}
309
310// For the module's life: a server that exits starts again on the same port, five times at most.
311// A reload kills the child and drops this loop's timer, and the next load starts its own.
312const serve = async ($: Api) => {
313  for (let restarts = 0; restarts < 5; restarts++) {
314    if (!(await runServer($))) break
315    await $.clock.sleep(1000)
316  }
317  await update($, bridge, () => null)
318}
319
320const browserUrl = (at: Bridge) => `http://127.0.0.1:${at.port}/?t=${at.token}`
321
322const openBrowser = async ($: Api) => {
323  const at = await read($, bridge)
324  if (at === null) return 'The browser view is not running: it needs bun (https://bun.sh).'
325  await $.process.run(['open', browserUrl(at)]).catch(() => undefined)
326  return `Opened ${browserUrl(at)}`
327}
328
329export const register: Register = on => {
330  on('session.start', async ($, e, next) => {
331    await Promise.all([
332      $.command.register({ name: 'tree', description: 'Open the conversation tree pane' }),
333      $.command.register({ name: 'branches', description: 'List the conversation branches' }),
334      $.command.register({ name: 'tree-web', description: 'Open the conversation tree in your browser' }),
335      $.command.register({
336        name: 'branch',
337        description: 'Fork the session into a new branch off the active one; your next prompts go there',
338        argumentHint: '[label]',
339      }),
340      $.command.register({
341        name: 'branch-to',
342        description: 'Switch the session to another branch: Claude sees only that branch from then on',
343        argumentHint: '<branch>',
344      }),
345      $.command.register({
346        name: 'branch-agent',
347        description: 'Run a subagent that sees only one branch of the conversation',
348        argumentHint: '<branch> <task>',
349      }),
350      $.agent.register({
351        name: 'branch',
352        description:
353          'Runs a task with only one branch of this conversation as context, the branches the user made in the convo-tree-2 pane (A, B, … or "root"). Use when the user asks for an agent that sees just one branch. Start the prompt with a line "Branch: <id or label>", then the task; the branch\'s turns are added for you, so do not restate them.',
354        prompt: AGENT_PROMPT,
355      }),
356    ])
357    void $.ui.open({ id: PANE, title: TITLE })
358    void serve($).catch(() => undefined) // an unload cancels its sleep
359    settle($)
360
361    return next(e)
362  })
363
364  on('command.run', { command: 'tree' }, async $ => {
365    const opened = await $.ui.open({ id: PANE, title: TITLE, focus: true })
366    const surfaces = (await $.session.surfaces()).join(', ') || 'none'
367    if (opened.isPlaced) {
368      const drawn =
369        Object.entries(renders)
370          .map(([s, n]) => `${s} ${n}×`)
371          .join(', ') || 'not yet'
372      return { text: `Conversation tree opened (surfaces: ${surfaces}; drawn on: ${drawn}).` }
373    }
374    // Where no surface places panes, say so and show the forks as text
375    return {
376      text: `The tree pane is open but not drawn (${opened.reason}; surfaces: ${surfaces}). Branches:\n${listText(await read($, tree))}`,
377    }
378  })
379
380  on('command.run', { command: 'branches' }, async $ => ({ text: listText(await read($, tree)) }))
381
382  on('command.run', { command: 'tree-web' }, async $ => ({ text: await openBrowser($) }))
383
384  on('command.run', { command: 'branch' }, async ($, e) => {
385    const label = oneLine(e.args)
386    const id = await newBranch($, label === '' ? null : cut(label, 40))
387    if (id === undefined) return { text: MID_TURN }
388    return { text: `On new branch ${id}, forked from here. Your next prompt continues there.` }
389  }).catch(() => ({ text: 'convo-tree-2: could not start a branch.' }))
390
391  on('command.run', { command: 'branch-to' }, async ($, e) => {
392    const branch = findBranch(await read($, tree), e.args)
393    if (branch === undefined)
394      return { text: `No branch "${e.args.trim()}". Branches:\n${listText(await read($, tree))}` }
395    return { text: await swapTo($, branch.id) }
396  })
397
398  // The switch: the transcript handed in is the branch being left; the one handed back is the branch's
399  on('session.compact', async ($, e, next) => {
400    const at = switching
401    if (e.agentId !== undefined || at === null) return next(e)
402    await update($, pending, () => null)
403    const target = await snapshotOf($, at.id)
404    if (target === null) return { skip: `branch ${at.id} has no saved transcript` }
405    // The engine hands its transcript in; read it when a host does not
406    const all = Array.isArray(e.messages) ? e.messages : await $.session.messages()
407    // At a turn's start the transcript may already end on its prompt: that prompt moves with the switch
408    const last = all.at(-1)
409    const isPrompt =
410      at.prompt !== null &&
411      last?.role === 'user' &&
412      last.toolResults === undefined &&
413      last.text.includes(at.prompt.trim())
414    const live = isPrompt ? all.slice(0, -1) : all
415    await saveSnapshot($, (await read($, tree)).activeId, strip(live))
416    await change($, cur => switchTo(cur, at.id))
417    await update($, view, () => ({ openId: at.id, openTurnId: null }))
418    const restored = restore(live, target)
419    return { messages: isPrompt && last !== undefined ? [...restored, last] : restored }
420  }).catch(($, e, next) =>
421    next.called || switching === null ? next(e) : { skip: 'convo-tree-2: could not switch branches' },
422  )
423
424  // A switch asked for from the pane, the browser or /branch-to goes in at the first of these
425  on('ui.press', async ($, e, next) => {
426    const done = await next(e)
427    await install($)
428    return done
429  }).catch(($, e, next) => (next.called ? undefined : next(e)))
430
431  on('prompt.edit', async ($, e, next) => {
432    await install($)
433    return next(e)
434  }).catch(($, e, next) => next(e))
435
436  on('command.run', { command: 'branch-agent' }, async ($, e) => {
437    const [ref = '', ...rest] = e.args.trim().split(/\s+/)
438    return { text: await startAgent($, ref, rest.join(' ')) }
439  })
440
441  // Claude spawns convo-tree-2:branch with "Branch: A"; the branch's turns are spliced in here
442  on('agent.spawn', async ($, e, next) => {
443    if (e.subagentType !== AGENT) return next(e)
444    const t = await read($, tree)
445    const ref = BRANCH_LINE.exec(e.prompt)?.[1]
446    const branch = ref === undefined ? undefined : findBranch(t, ref)
447    if (branch === undefined) {
448      return {
449        deny: `Start the prompt with a line "Branch: <id or label>" naming one of these branches:\n${listText(t)}`,
450      }
451    }
452    const task = e.prompt.replace(BRANCH_LINE, '').trim()
453    const spawned = await next({ ...e, prompt: withBranch(t, branch.id, task) })
454    if (spawned.agentId !== undefined) await change($, recordAgent(spawned.agentId, branch.id, task))
455    return spawned
456  }).catch(($, e, next) => (next.called ? next(e) : { deny: 'convo-tree-2: could not attach the branch.' }))
457
458  // Back on a branch after talking elsewhere: the transcript is, or is about to be, that branch's; say so once
459  on('prompt.submit', async ($, e, next) => {
460    const t = await read($, tree)
461    const asked = await read($, pending)
462    const branch = branchOf(t, asked ?? t.activeId)
463    if (branch === undefined || (asked === null && !isReturning(t))) return next(e)
464    const note = `[convo-tree-2] The user switched the session to branch ${branch.id} "${labelOf(branch)}"; the conversation above is that branch's.${branch.summary === '' ? '' : ` Where it stood: ${branch.summary}`}`
465    return next({ ...e, context: [...(e.context ?? []), note] })
466  }).catch(($, e, next) => next(e))
467
468  on('turn.start', async ($, e, next) => {
469    await install($, e.text)
470    await change($, t => startTurn(t, e.turnId, e.text))
471    return next(e)
472  })
473
474  on('turn.complete', async ($, e, next) => {
475    const done = await next(e)
476
477    if (e.agentId !== undefined) {
478      const mine = (await read($, tree)).agents.find(a => a.agentId === e.agentId)
479      if (mine !== undefined) {
480        await change($, cur => ({
481          ...cur,
482          agents: cur.agents.map(a => (a.agentId === e.agentId ? { ...a, isDone: true } : a)),
483        }))
484        $.ui.toast(`Branch ${mine.branchId} subagent finished.`)
485      }
486      return done
487    }
488
489    const u = e.usage
490    const tokens = u === undefined ? 0 : u.input_tokens + u.cache_creation_input_tokens + u.output_tokens
491    await change($, t => finishTurn(t, e.turnId, e.answer, tokens))
492    settle($)
493
494    return done
495  })
496
497  // While a switch waits to go in, a band says where the next prompt goes
498  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
499    const id = await read($, pending)
500    const branch = id === null ? undefined : branchOf(await read($, tree), id)
501    if (branch === undefined) return next(e)
502    const { Text } = $.ui.resolve(e)
503    return (
504      <Text dimColor>
505        ⑂ Next prompt goes to branch {branch.id} · {labelOf(branch)}
506      </Text>
507    )
508  })
509
510  // One bordered card per branch, indented by depth; opening a card shows its last turns and actions
511  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
512    renders[e.surface] = (renders[e.surface] ?? 0) + 1
513    const { Box, Text, Button, Markdown } = $.ui.resolve(e)
514    const t = await read($, tree)
515    const v = await read($, view)
516    const width = Math.max(28, e.props.bodyColumns)
517    const path = pathTo(t, t.activeId)
518    const trail = path.map(b => (b.id === ROOT ? labelOf(b) : `${b.id} · ${labelOf(b)}`))
519
520    const toggleCard = (id: string) => () =>
521      update($, view, cur => ({ openId: cur.openId === id ? null : id, openTurnId: null }))
522    const toggleTurn = (id: string) => () =>
523      update($, view, cur => ({ ...cur, openTurnId: cur.openTurnId === id ? null : id }))
524
525    // Dotted guides beside a card: ancestors' lines carry on, its own meets the card's title row
526    const guides = (prefix: string, height: number) => {
527      const chunks = prefix.match(/.{3}/g) ?? []
528      const isLast = chunks.pop()?.startsWith('└') ?? true
529      const lead = chunks.map(c => (c.startsWith('│') ? '┊  ' : '   ')).join('')
530      const own = (i: number) => (i < 2 ? '┊  ' : i === 2 ? (isLast ? '╰┈┈' : '├┈┈') : isLast ? '   ' : '┊  ')
531      return (
532        <Box flexDirection="column">
533          {Array.from({ length: height }, (_, i) => (
534            <Text dimColor>{lead + own(i)}</Text>
535          ))}
536        </Box>
537      )
538    }
539
540    // Lines a card takes, to draw its guides as tall: exact when closed, an estimate when open
541    const heightOf = (id: string, room: number) => {
542      if (id !== v.openId) return 4
543      const turns = turnsOf(t, id)
544      const shown = turns.slice(-3)
545      const reply = shown.find(x => x.id === v.openTurnId)?.reply ?? ''
546      const replyLines = cut(reply, 1500)
547        .split('\n')
548        .reduce((sum, line) => sum + Math.max(1, Math.ceil(line.length / Math.max(8, room - 2))), 0)
549      const agents = Math.min(3, t.agents.filter(a => a.branchId === id).length)
550      return 4 + 1 + (turns.length > 0 ? 1 : 0) + shown.length + (reply === '' ? 0 : replyLines) + agents + 2
551    }
552
553    const card = ({ branch, depth, prefix }: Row) => {
554      const id = branch.id
555      const turns = turnsOf(t, id)
556      const isActive = id === t.activeId
557      const isOpen = id === v.openId
558      const room = width - depth * 3 - 4
559      const count = turns.length === 0 ? 'empty' : `${turns.length} turn${turns.length === 1 ? '' : 's'}`
560      const tag = id === ROOT ? '' : `${id} `
561      const summary =
562        branch.summary !== ''
563          ? branch.summary
564          : turns.length === 0
565            ? 'No turns yet.'
566            : oneLine(turns.at(-1)?.prompt ?? '')
567      const border = isActive ? { borderColor: colorOf(t, id) } : { borderDimColor: true }
568      return (
569        <Box key={`row-${id}`}>
570          {depth > 0 && guides(prefix, 1 + heightOf(id, room))}
571          <Box
572            key={`card-${id}`}
573            flexDirection="column"
574            flexGrow={1}
575            marginTop={depth > 0 ? 1 : 0}
576            borderStyle="round"
577            paddingX={1}
578            {...border}
579          >
580            <Box justifyContent="space-between">
581              <Button key={`branch-${id}`} plain onPress={toggleCard(id)}>
582                <Text color={colorOf(t, id)}>
583                  {isActive ? '●' : turns.length === 0 ? '◌' : '○'} {tag}
584                </Text>
585                <Text bold={isActive}>{cut(labelOf(branch), Math.max(4, room - tag.length - count.length - 4))}</Text>
586              </Button>
587              <Text dimColor>{count}</Text>
588            </Box>
589            <Text dimColor>{cut(oneLine(summary), Math.max(8, room))}</Text>
590            {isOpen && (
591              <Box flexDirection="column" marginTop={1}>
592                {turns.length > 0 && (
593                  <Text dimColor>
594                    {ranges(turns.map(x => x.n))} · {compact(tokensOf(t, id))} tokens
595                  </Text>
596                )}
597                {turns.slice(-3).map(turn => (
598                  <Box key={`turn-${turn.id}`} flexDirection="column">
599                    <Button
600                      key={`open-${turn.id}`}
601                      plain
602                      dimColor={v.openTurnId !== turn.id}
603                      label={`${v.openTurnId === turn.id ? '▾' : '▸'} T${turn.n}  ${cut(oneLine(turn.prompt), room - 6 - String(turn.n).length)}`}
604                      onPress={toggleTurn(turn.id)}
605                    />
606                    {v.openTurnId === turn.id && (
607                      <Markdown key={`reply-${turn.id}`} text={cut(turn.reply ?? '_Still running…_', 1500)} />
608                    )}
609                  </Box>
610                ))}
611                {t.agents
612                  .filter(a => a.branchId === id)
613                  .slice(-3)
614                  .map(a => (
615                    <Text key={`agent-${a.agentId}`} dimColor>
616                      ⧉ {cut(oneLine(a.task), room - 12)} {a.isDone ? '· done' : '· running'}
617                    </Text>
618                  ))}
619                <Box marginTop={1} flexWrap="wrap" columnGap={1}>
620                  {isActive ? (
621                    <Button key={`split-${id}`} label="⑂ Split reply" onPress={() => void splitReply($)} />
622                  ) : (
623                    <Button
624                      key={`switch-${id}`}
625                      variant="primary"
626                      label="Switch here"
627                      onPress={() => showBranch($, id)}
628                    />
629                  )}
630                  <Button
631                    key={`agent-${id}`}
632                    label="⧉ Agent"
633                    onPress={() => void $.prompt.fill({ text: `/branch-agent ${id} `, mode: 'replace' })}
634                  />
635                </Box>
636              </Box>
637            )}
638          </Box>
639        </Box>
640      )
641    }
642
643    return (
644      <Box flexDirection="column">
645        <Text dimColor>{cut(trail.join(' › '), width)}</Text>
646        {rows(t).map(card)}
647        <Box marginTop={1} flexWrap="wrap" columnGap={1}>
648          <Button
649            key="new"
650            label="+ Branch"
651            onPress={async () => {
652              if ((await newBranch($, null)) === undefined) $.ui.toast(MID_TURN)
653            }}
654          />
655          <Button key="web" label="↗ Map" onPress={() => void openBrowser($)} />
656        </Box>
657      </Box>
658    )
659  })
660}
661

1import type { SessionMessage } from 'claude-code'
2
3import type { Branch, Tree, Turn } from '../types'
4
5export const ROOT = 'root'
6const MAX_PROMPT = 2000
7const MAX_REPLY = 4000
8const MAX_CONTEXT = 60000
9
10export const emptyTree = (): Tree => ({
11  branches: [{ id: ROOT, parentId: null, label: null, summary: '', summarizedTurns: 0 }],
12  turns: [],
13  agents: [],
14  activeId: ROOT,
15})
16
17// 0 → A, 25 → Z, 26 → AA
18const keyOf = (n: number): string => (n >= 26 ? keyOf(Math.floor(n / 26) - 1) : '') + String.fromCharCode(65 + (n % 26))
19
20export const labelOf = (b: Branch) => b.label ?? (b.id === ROOT ? 'Root' : 'New branch')
21
22export const branchOf = (tree: Tree, id: string) => tree.branches.find(b => b.id === id)
23
24export const findBranch = (tree: Tree, ref: string) => {
25  const want = ref.trim().toLowerCase()
26  return (
27    tree.branches.find(b => b.id.toLowerCase() === want) ??
28    tree.branches.find(b => labelOf(b).toLowerCase() === want) ??
29    tree.branches.find(b => labelOf(b).toLowerCase().includes(want))
30  )
31}
32
33const childrenOf = (tree: Tree, id: string) => tree.branches.filter(b => b.parentId === id)
34
35export const turnsOf = (tree: Tree, id: string) => tree.turns.filter(t => t.branchId === id)
36
37export const tokensOf = (tree: Tree, id: string) => turnsOf(tree, id).reduce((sum, t) => sum + t.tokens, 0)
38
39export const pathTo = (tree: Tree, id: string): Branch[] => {
40  const branch = branchOf(tree, id)
41  if (branch === undefined) return []
42  return branch.parentId === null ? [branch] : [...pathTo(tree, branch.parentId), branch]
43}
44
45// Children of the active branch; ids follow on from the last, as branches are never removed
46export const addBranches = (tree: Tree, labels: readonly (string | null)[]): Tree => ({
47  ...tree,
48  branches: [
49    ...tree.branches,
50    ...labels.map(
51      (label, i): Branch => ({
52        id: keyOf(tree.branches.length - 1 + i),
53        parentId: tree.activeId,
54        label,
55        summary: '',
56        summarizedTurns: 0,
57      }),
58    ),
59  ],
60})
61
62export const switchTo = (tree: Tree, id: string): Tree =>
63  branchOf(tree, id) === undefined ? tree : { ...tree, activeId: id }
64
65// What the host wraps around a typed prompt (reminders, slash-command echoes), not the person's words
66// A tag with no closing tag runs to the end
67const HOST_TAGS =
68  /<(system-reminder|local-command-caveat|local-command-stdout|command-name|command-message|command-args)>[\s\S]*?(<\/\1>|$)/g
69
70const cleanPrompt = (text: string) => text.replace(HOST_TAGS, '').trim()
71
72// A turn belongs to the branch that is active when it starts
73export const startTurn = (tree: Tree, id: string, prompt: string): Tree => ({
74  ...tree,
75  turns: [
76    ...tree.turns,
77    {
78      id,
79      branchId: tree.activeId,
80      n: tree.turns.length + 1,
81      prompt: cleanPrompt(prompt).slice(0, MAX_PROMPT),
82      reply: null,
83      tokens: 0,
84    },
85  ],
86})
87
88export const finishTurn = (tree: Tree, id: string, reply: string, tokens: number): Tree => ({
89  ...tree,
90  turns: tree.turns.map(t => (t.id === id ? { ...t, reply: reply.slice(0, MAX_REPLY), tokens } : t)),
91})
92
93export const patchBranch = (tree: Tree, id: string, patch: Partial<Branch>): Tree => ({
94  ...tree,
95  branches: tree.branches.map(b => (b.id === id ? { ...b, ...patch } : b)),
96})
97
98// The person is on a branch with turns, and the last turn was on another one
99export const isReturning = (tree: Tree) => {
100  const last = tree.turns.at(-1)
101  return last !== undefined && last.branchId !== tree.activeId && turnsOf(tree, tree.activeId).length > 0
102}
103
104// "T2–3, T6" from turn numbers
105export const ranges = (ns: readonly number[]): string => {
106  const runs: [number, number][] = []
107  for (const n of ns) {
108    const last = runs.at(-1)
109    if (last !== undefined && n === last[1] + 1) last[1] = n
110    else runs.push([n, n])
111  }
112  return runs.map(([a, b]) => (a === b ? `T${a}` : `T${a}–${b}`)).join(', ')
113}
114
115export const compact = (n: number) => (n >= 1000 ? `${(n / 1000).toFixed(n >= 10000 ? 0 : 1)}k` : `${n}`)
116
117export type Row = { branch: Branch; prefix: string; depth: number }
118
119// Depth-first rows with box-drawing prefixes
120export const rows = (tree: Tree): Row[] => {
121  const walk = (id: string, lead: string, isLast: boolean, depth: number): Row[] => {
122    const branch = branchOf(tree, id)
123    if (branch === undefined) return []
124    const prefix = depth === 0 ? '' : lead + (isLast ? '└─ ' : '├─ ')
125    const childLead = depth === 0 ? '' : lead + (isLast ? '   ' : '│  ')
126    const kids = childrenOf(tree, id)
127    return [
128      { branch, prefix, depth },
129      ...kids.flatMap((kid, i) => walk(kid.id, childLead, i === kids.length - 1, depth + 1)),
130    ]
131  }
132  return walk(ROOT, '', true, 0)
133}
134
135const nameOf = (b: Branch) => (b.id === ROOT ? `root "${labelOf(b)}"` : `${b.id} "${labelOf(b)}"`)
136
137const turnText = (t: Turn) =>
138  `### T${t.n} · user\n${t.prompt}\n\n### T${t.n} · assistant\n${t.reply ?? '(still running)'}`
139
140// Root → branch: each ancestor's turns from before the branch began, then the branch's own
141export const branchContext = (tree: Tree, id: string): string | undefined => {
142  const path = pathTo(tree, id)
143  const branch = path.at(-1)
144  if (branch === undefined) return undefined
145  const own = turnsOf(tree, id)
146  const start = own[0]?.n ?? Infinity
147  const sections = path.slice(0, -1).flatMap(a => {
148    const lead = turnsOf(tree, a.id).filter(t => t.n < start)
149    return lead.length === 0 ? [] : [`## Lead-in: ${nameOf(a)}\n\n${lead.map(turnText).join('\n\n')}`]
150  })
151  sections.push(
152    `## Branch ${nameOf(branch)}\n\n${own.length === 0 ? '(no turns yet)' : own.map(turnText).join('\n\n')}`,
153  )
154  const head = [
155    `# Conversation branch ${nameOf(branch)}`,
156    `Path: ${path.map(nameOf).join(' › ')}`,
157    branch.summary === '' ? '' : `Summary: ${branch.summary}`,
158    'Only the user prompts and final assistant replies are kept; tool calls and outputs are omitted.',
159  ].filter(Boolean)
160  const text = [...head, ...sections].join('\n\n')
161  return text.length > MAX_CONTEXT ? `…(earlier text cut)\n${text.slice(-MAX_CONTEXT)}` : text
162}
163
164export const listText = (tree: Tree): string =>
165  rows(tree)
166    .map(({ branch, prefix }) => {
167      const turns = turnsOf(tree, branch.id).map(t => t.n)
168      const mark = branch.id === tree.activeId ? ' (active)' : ''
169      const where = turns.length === 0 ? 'no turns' : ranges(turns)
170      const summary = branch.summary === '' ? '' : ` — ${branch.summary}`
171      return `${prefix}${nameOf(branch)}${mark}: ${where}${summary}`
172    })
173    .join('\n')
174
175// A message as a snapshot keeps it: the engine's handle is only good for the compaction that stamped it
176export const strip = (messages: readonly SessionMessage[]): SessionMessage[] => messages.map(({ handle: _, ...m }) => m)
177
178const idsOf = (m: SessionMessage) =>
179  [...m.toolUses.map(u => u.tool_use_id), '|', ...(m.toolResults ?? []).map(r => r.tool_use_id)].join(',')
180
181const isSame = (a: SessionMessage, b: SessionMessage) => a.role === b.role && a.text === b.text && idsOf(a) === idsOf(b)
182
183// The branch's transcript, with the opening it shares with the live one taken from the live one,
184// so those messages stand as the engine has them and only the rest is rebuilt
185export const restore = (live: readonly SessionMessage[], target: SessionMessage[]): SessionMessage[] => {
186  const shared = target.findIndex((m, i) => live[i] === undefined || !isSame(live[i], m))
187  const i = shared === -1 ? target.length : shared
188  return [...live.slice(0, i), ...target.slice(i)]
189}
190

1// A prompt and its reply, on the branch that was active when it was sent
2export type Turn = {
3  id: string
4  branchId: string
5  n: number
6  prompt: string
7  reply: string | null // null while the turn runs
8  tokens: number
9}
10
11export type Branch = {
12  id: string
13  parentId: string | null
14  label: string | null // null until Haiku titles it
15  summary: string
16  summarizedTurns: number // how many finished turns the summary covers
17}
18
19export type BranchAgent = {
20  agentId: string
21  branchId: string
22  task: string
23  isDone: boolean
24}
25
26// Branch ids, turn counts and token totals are all derived from these
27export type Tree = {
28  branches: Branch[]
29  turns: Turn[]
30  agents: BranchAgent[]
31  activeId: string
32}
33
34// Which card in the pane is opened up, and which of its turns shows its reply
35export type View = { openId: string | null; openTurnId: string | null }
36
37// Where the browser view's server listens; kept so a reload asks for the same port
38export type Bridge = { port: number; token: string }
39
40declare module 'claude-code' {
41  interface PluginState {
42    // pending: the branch a switch was asked for and not yet installed
43    // snapshot: a branch's whole transcript, as it stood when the session last left it
44    'convo-tree-2': {
45      tree: Tree
46      view: View
47      bridge: Bridge | null
48      pending: string | null
49      snapshot: StateFamily<SessionMessage[] | null>
50    }
51  }
52}
53