diagram-render

1import type { EngineInterface, Register } from 'claude-code'
2import { blockHash, cells, failureLine, mermaidBlocks, pngSize, type Size } from './diagram.ts'
3
4const ENABLED_KEY = 'enabled'
5
6const USAGE = 'expects nothing (the status), on or off'
7
8const INSTALL_HINT = 'mmdc is not installed, so mermaid blocks stay text: npm i -g @mermaid-js/mermaid-cli'
9
10/** mmdc starts a headless browser; a large diagram takes seconds. */
11const RENDER_TIMEOUT_MS = 60_000
12
13/**
14 * Blocks drawn in a reply wait in `wanted` until the turn ends; `ready` holds the rendered pictures and
15 * `failed` the blocks mmdc refused, so neither is tried again. `hasMmdc` is undefined until checked.
16 */
17type State = {
18  enabled: boolean
19  wanted: Map<string, string>
20  ready: Map<string, { png: string; size: Size }>
21  failed: Set<string>
22  busy: boolean
23  hasMmdc?: boolean
24  dir?: string
25}
26
27/**
28 * Reads the on/off setting from the store, which every window shares, so a change made in another
29 * window applies here at the next hook that acts on it. A changed setting redraws the replies, as
30 * `on` and `off` do.
31 */
32async function readSettings($: EngineInterface, state: State): Promise<void> {
33  const was = state.enabled
34  state.enabled = (await $.store.get(ENABLED_KEY)) !== false
35  if (state.enabled !== was) $.ui.invalidate('ui.render')
36}
37
38function errorText(err: unknown): string {
39  return err instanceof Error ? err.message : String(err)
40}
41
42/** Whether mmdc runs; the answer is kept, and a missing mmdc is said once. */
43async function mmdcReady($: EngineInterface, state: State): Promise<boolean> {
44  if (state.hasMmdc !== undefined) return state.hasMmdc
45  const found = await $.process.run(['mmdc', '--version'], { timeoutMs: 20_000 }).then(r => r.exitCode === 0, () => false)
46  state.hasMmdc = found
47  if (!found) $.ui.log(INSTALL_HINT)
48  return found
49}
50
51async function workDir($: EngineInterface, state: State): Promise<string> {
52  state.dir ??= `${((await $.env.get('TMPDIR')) ?? '/tmp').replace(/\/+$/, '')}/diagram-render`
53  return state.dir
54}
55
56/** Renders one block to `<hash>.png` and answers its size, or throws with mmdc's reason. */
57async function renderOne($: EngineInterface, state: State, hash: string, source: string): Promise<{ png: string; size: Size }> {
58  const dir = await workDir($, state)
59  const input = `${dir}/${hash}.mmd`
60  const png = `${dir}/${hash}.png`
61  await $.fs.write(input, `${source}\n`)
62  const r = await $.process.run(['mmdc', '-i', input, '-o', png, '-b', 'transparent', '-t', 'dark', '-q'], { timeoutMs: RENDER_TIMEOUT_MS })
63  if (r.exitCode !== 0) throw new Error(failureLine(`${r.stderr}\n${r.stdout}`))
64  const size = pngSize((await $.fs.read(png, { as: 'bytes' })).base64)
65  if (size === undefined) throw new Error(`${png} is not a PNG`)
66  return { png, size }
67}
68
69/** Renders the waiting blocks one at a time, in the background; a failed block is logged and not tried again. */
70async function drain($: EngineInterface, state: State): Promise<void> {
71  if (state.busy || state.wanted.size === 0) return
72  state.busy = true
73  try {
74    if (!(await mmdcReady($, state))) return state.wanted.clear()
75    for (const [hash, source] of state.wanted) {
76      state.wanted.delete(hash)
77      await renderOne($, state, hash, source).then(
78        picture => { state.ready.set(hash, picture); $.ui.invalidate('ui.render') },
79        (err: unknown) => { state.failed.add(hash); $.ui.log(`a diagram was not rendered: ${errorText(err)}`) },
80      )
81    }
82  } finally {
83    state.busy = false
84  }
85}
86
87/** Queues the blocks of a drawn reply that have no picture yet. */
88function want(state: State, blocks: readonly string[]): void {
89  if (state.hasMmdc === false) return
90  for (const source of blocks) {
91    const hash = blockHash(source)
92    if (!state.ready.has(hash) && !state.failed.has(hash)) state.wanted.set(hash, source)
93  }
94}
95
96async function runCommand($: EngineInterface, state: State, args: string): Promise<string> {
97  const word = args.trim()
98  if (word === 'on' || word === 'off') {
99    await $.store.set(ENABLED_KEY, word === 'on')
100    state.enabled = word === 'on'
101    $.ui.invalidate('ui.render')
102    return word === 'on' ? 'on: mermaid blocks render under their reply after each turn' : 'off: mermaid blocks stay text'
103  }
104  if (word !== '') return USAGE
105  await readSettings($, state)
106  const mmdc = state.hasMmdc === undefined ? 'mmdc not checked yet' : state.hasMmdc ? 'mmdc found' : 'mmdc missing'
107  return `${state.enabled ? 'on' : 'off'}; ${mmdc}; ${state.ready.size} rendered, ${state.failed.size} failed this session`
108}
109
110export const register: Register = on => {
111  const state: State = { enabled: true, wanted: new Map(), ready: new Map(), failed: new Set(), busy: false }
112
113  on('session.start', async ($, e, next) => {
114    const r = await next(e)
115    await $.command.register({ name: 'diagram-render', description: 'Mermaid blocks as pictures: status, on, off (diagram-render)', argumentHint: '[on | off]' })
116    await readSettings($, state)
117    return r
118  })
119
120  // The engine prints the plugin name in front of command text and log lines, so the texts do not repeat it.
121  on('command.run', { command: 'diagram-render' }, async ($, e) => ({ text: await runCommand($, state, String(e.args ?? '')) }))
122
123  on('turn.complete', async ($, e, next) => {
124    const r = await next(e)
125    if (e.agentId !== undefined) return r
126    await readSettings($, state)
127    if (!state.enabled) return r
128    want(state, mermaidBlocks(e.answer))
129    // The renders run on a timer: mmdc takes seconds, the next prompt must not wait for it, and
130    // since 2.1.288 a process call still in flight when the turn's dispatch closes is aborted. A
131    // timer outlives the dispatch, so the picture still arrives after the turn's end.
132    $.clock.after(0, () => void drain($, state))
133    return r
134  })
135
136  // A render runs at every redraw and scroll, so it reads the setting the last turn's end read, not the store.
137  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
138    const blocks = state.enabled && e.surface === 'terminal' ? mermaidBlocks(e.props.text) : []
139    if (blocks.length === 0 || e.surface !== 'terminal') return next(e)
140    want(state, blocks)
141    const pictures = blocks.map(blockHash).map(h => state.ready.get(h)).filter(p => p !== undefined)
142    const drawn = await next(e)
143    if (pictures.length === 0) return drawn
144    const { Box, Image } = $.ui.resolve(e)
145    const maxColumns = Math.min(100, Math.max(10, (e.viewport?.columns ?? 84) - 4))
146    return (
147      <Box flexDirection="column">
148        {drawn}
149        {pictures.map((p, i) => {
150          const box = cells(p.size, maxColumns)
151          return <Image key={`diagram:${i}`} source={{ file: p.png, format: 'png' }} columns={box.columns} rows={box.rows} alt={`mermaid diagram ${i + 1}`} />
152        })}
153      </Box>
154    )
155  })
156}
157

1/** The mermaid blocks of a reply, their file names, and the box of cells a rendered picture is drawn in. */
2
3/** The source of each closed ```mermaid block in `text`, in order; a block still streaming has no closing fence yet. */
4export function mermaidBlocks(text: string): string[] {
5  return [...text.matchAll(/^[ \t]*```mermaid[^\n]*\n([\s\S]*?)^[ \t]*```[ \t]*$/gm)].map(m => (m[1] ?? '').trim()).filter(s => s !== '')
6}
7
8/** A file name for a block, FNV-1a 32 of its source, so the same block renders once. */
9export function blockHash(source: string): string {
10  let h = 0x811c9dc5
11  for (const ch of source) h = Math.imul(h ^ (ch.codePointAt(0) ?? 0), 0x01000193) >>> 0
12  return h.toString(16).padStart(8, '0')
13}
14
15const B64 = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'
16
17/** The first bytes of a base64 string; a module has no Node Buffer. */
18export function headBytes(base64: string, count: number): number[] {
19  const bytes: number[] = []
20  let bits = 0
21  let value = 0
22  for (const ch of base64) {
23    const n = B64.indexOf(ch)
24    if (n < 0 || bytes.length >= count) break
25    value = (value << 6) | n
26    bits += 6
27    if (bits >= 8) {
28      bits -= 8
29      bytes.push((value >> bits) & 0xff)
30    }
31  }
32  return bytes
33}
34
35export type Size = { width: number; height: number }
36
37const u32 = (b: number[], at: number): number => (((b[at] ?? 0) << 24) >>> 0) + ((b[at + 1] ?? 0) << 16) + ((b[at + 2] ?? 0) << 8) + (b[at + 3] ?? 0)
38
39/** The size in a PNG's IHDR chunk, or undefined for bytes that are not a PNG. */
40export function pngSize(base64: string): Size | undefined {
41  const b = headBytes(base64, 24)
42  const signature = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]
43  if (!signature.every((v, i) => b[i] === v) || String.fromCharCode(...b.slice(12, 16)) !== 'IHDR') return undefined
44  const size = { width: u32(b, 16), height: u32(b, 20) }
45  return size.width > 0 && size.height > 0 ? size : undefined
46}
47
48/** The tallest diagram, in rows. */
49export const MAX_ROWS = 30
50
51/** Pixels per cell across; a cell is about twice as tall as it is wide. */
52const PX_PER_COLUMN = 8
53
54/** The box of cells that keeps the picture's shape, at most `maxColumns` wide and MAX_ROWS tall. */
55export function cells(size: Size, maxColumns: number): { columns: number; rows: number } {
56  const fit = Math.max(1, Math.min(maxColumns, Math.round(size.width / PX_PER_COLUMN)))
57  const rows = Math.max(1, Math.round((fit * size.height) / size.width / 2))
58  if (rows <= MAX_ROWS) return { columns: fit, rows }
59  return { columns: Math.max(1, Math.round((MAX_ROWS * 2 * size.width) / size.height)), rows: MAX_ROWS }
60}
61
62/** The line of an mmdc failure worth showing: the first that names an error, else the first. */
63export function failureLine(output: string): string {
64  const lines = output.split('\n').map(l => l.trim()).filter(l => l !== '')
65  return (lines.find(l => /error/i.test(l)) ?? lines[0] ?? 'no output').slice(0, 200)
66}
67