mermaid-inline

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4// The source of the newest diagram drawn; the band above the prompt acts on it.
5const latest = atom({ plugin: 'mermaid-inline', key: 'latest' } as const, null)
6
7type Diagram = { base64: string; width: number; height: number }
8type Part = { kind: 'text'; text: string } | { kind: 'diagram'; source: string }
9
10const cache = new Map<string, Diagram | null>()
11const running = new Map<string, Promise<Diagram | null>>()
12
13const FENCE = /```mermaid[^\S\n]*\n([\s\S]*?)\n[^\S\n]*```[^\S\n]*(?:\n|$)/g
14const ALPHABET = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'
15
16function hash(text: string): string {
17  let h = 0x811c9dc5
18
19  for (let n = 0; n < text.length; n++) {
20    h = Math.imul(h ^ text.charCodeAt(n), 0x01000193)
21  }
22
23  return (h >>> 0).toString(16) + text.length.toString(16)
24}
25
26function split(text: string): Part[] {
27  const parts: Part[] = []
28  let last = 0
29
30  for (const match of text.matchAll(FENCE)) {
31    const start = match.index ?? 0
32
33    if (start > last) {
34      parts.push({ kind: 'text', text: text.slice(last, start) })
35    }
36
37    parts.push({ kind: 'diagram', source: match[1] })
38    last = start + match[0].length
39  }
40
41  if (last < text.length) {
42    parts.push({ kind: 'text', text: text.slice(last) })
43  }
44
45  return parts
46}
47
48// Width and height sit at bytes 16 to 23 of a PNG: the first 32 base64 characters hold them.
49function pngSize(base64: string): { width: number; height: number } | null {
50  const bytes: number[] = []
51  const head = base64.slice(0, 32)
52
53  for (let n = 0; n < head.length; n += 4) {
54    const a = ALPHABET.indexOf(head[n])
55    const b = ALPHABET.indexOf(head[n + 1])
56    const c = ALPHABET.indexOf(head[n + 2])
57    const d = ALPHABET.indexOf(head[n + 3])
58    bytes.push((a << 2) | (b >> 4), ((b & 15) << 4) | (c >> 2), ((c & 3) << 6) | d)
59  }
60
61  const read = (at: number) => ((bytes[at] << 24) | (bytes[at + 1] << 16) | (bytes[at + 2] << 8) | bytes[at + 3]) >>> 0
62  const width = read(16)
63  const height = read(20)
64
65  return width > 0 && height > 0 ? { width, height } : null
66}
67
68async function render($: EngineInterface, source: string): Promise<Diagram | null> {
69  const key = hash(source)
70  const input = `/tmp/mermaid-inline-${key}.mmd`
71  const output = `/tmp/mermaid-inline-${key}.png`
72
73  await $.fs.write(input, source)
74
75  const run = await $.process.run(
76    ['mmdc', '-i', input, '-o', output, '-t', 'dark', '-b', 'transparent', '-s', String(RENDER_SCALE), '-q'],
77    { timeoutMs: 60_000 },
78  )
79
80  if (run.exitCode !== 0) {
81    return null
82  }
83
84  const { base64 } = await $.fs.read(output, { as: 'bytes' })
85  const size = pngSize(base64)
86
87  return size ? { base64, ...size } : null
88}
89
90async function diagramFor($: EngineInterface, source: string): Promise<Diagram | null> {
91  const key = hash(source)
92
93  if (cache.has(key)) {
94    return cache.get(key) ?? null
95  }
96
97  let job = running.get(key)
98
99  if (!job) {
100    job = render($, source).catch(() => null)
101    running.set(key, job)
102  }
103
104  const diagram = await job
105  cache.set(key, diagram)
106  running.delete(key)
107
108  return diagram
109}
110
111// The PNG is RENDER_SCALE times its CSS size. A cell is about 9 by 18 CSS pixels; ZOOM draws it larger.
112const RENDER_SCALE = 3
113const ZOOM = 1.5
114
115// One scale fits both limits, so a tall diagram shrinks and keeps its aspect ratio.
116function cells(diagram: Diagram, maxColumns: number, maxRows: number): { columns: number; rows: number } {
117  const naturalColumns = (diagram.width / RENDER_SCALE / 9) * ZOOM
118  const naturalRows = (diagram.height / RENDER_SCALE / 18) * ZOOM
119  const scale = Math.min(1, maxColumns / naturalColumns, maxRows / naturalRows)
120  const columns = Math.max(1, Math.min(255, Math.round(naturalColumns * scale)))
121  const rows = Math.max(1, Math.min(255, Math.round(naturalRows * scale)))
122
123  return { columns, rows }
124}
125
126async function copyCode($: EngineInterface, source: string, surface: string) {
127  const { isCopied } = await $.ui.copy({ text: source, surface: surface as never })
128
129  $.ui.toast(isCopied ? 'Mermaid code copied' : 'Could not copy the mermaid code')
130}
131
132// The saved file is a fresh render on a white background, so it reads in any viewer. It opens at once.
133async function savePng($: EngineInterface, source: string) {
134  const key = hash(source)
135  const home = await $.env.get('HOME')
136  const input = `/tmp/mermaid-inline-${key}.mmd`
137  const output = `${home}/Downloads/mermaid-${key}.png`
138
139  await $.fs.write(input, source)
140
141  const run = await $.process.run(
142    ['mmdc', '-i', input, '-o', output, '-t', 'default', '-b', 'white', '-s', '3', '-q'],
143    { timeoutMs: 60_000 },
144  )
145
146  if (run.exitCode !== 0) {
147    $.ui.toast('Could not save the PNG')
148
149    return
150  }
151
152  $.ui.toast(`Saved ${output}`)
153  await $.process.run(['open', output], { timeoutMs: 10_000 })
154}
155
156export const register: Register = on => {
157  // Drawing is pure, so the newest diagram is noted when the reply is stored.
158  on('session.append', { door: 'response' }, async ($, e, next) => {
159    const result = await next(e)
160
161    for (const block of e.message.content) {
162      if (block.type !== 'text') {
163        continue
164      }
165
166      const sources = [...block.text.matchAll(FENCE)].map(match => match[1])
167      const newest = sources[sources.length - 1]
168
169      if (newest !== undefined) {
170        await update($, latest, () => newest)
171      }
172    }
173
174    return result
175  })
176
177  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
178    if (e.surface !== 'terminal' || !e.props.text.includes('```mermaid')) {
179      return next(e)
180    }
181
182    const parts = split(e.props.text)
183
184    if (!parts.some(part => part.kind === 'diagram')) {
185      return next(e)
186    }
187
188    const drawn = await Promise.all(
189      parts.map(async part => (part.kind === 'diagram' ? await diagramFor($, part.source) : null)),
190    )
191
192    if (drawn.every((diagram, n) => parts[n].kind === 'text' || diagram === null)) {
193      return next(e)
194    }
195
196    const { Box, Image, Markdown } = $.ui.resolve(e)
197    const maxColumns = Math.max(20, Math.min(200, (e.viewport?.columns ?? 100) - 4))
198    const maxRows = Math.max(20, Math.min(80, (e.viewport?.rows ?? 50) - 6))
199
200    return (
201      <Box flexDirection="column">
202        {parts.map((part, n) => {
203          const diagram = drawn[n]
204
205          if (part.kind === 'text') {
206            return <Markdown key={`t${n}`} text={part.text} />
207          }
208
209          if (!diagram) {
210            return <Markdown key={`d${n}`} text={'```mermaid\n' + part.source + '\n```'} />
211          }
212
213          const size = cells(diagram, maxColumns, maxRows)
214
215          return (
216            <Image
217              key={`d${n}`}
218              source={{ png: diagram.base64 }}
219              columns={size.columns}
220              rows={size.rows}
221              alt="mermaid diagram"
222            />
223          )
224        })}
225      </Box>
226    )
227  })
228
229  // The transcript draws no Buttons, so the actions sit in the band above the prompt.
230// Another mod may draw there too, so this hook stacks its row over what `next` returns.
231  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
232    const source = await read($, latest)
233
234    if (e.props.hasSurvey || source === null) {
235      return next(e)
236    }
237
238    const below = await next(e)
239    const { Box, Button, Text } = $.ui.resolve(e)
240
241    return (
242      <Box flexDirection="column">
243        <Box flexDirection="row" gap={1}>
244          <Text dimColor>Last mermaid diagram:</Text>
245          <Button key="copy" label="Copy code" onPress={press => copyCode($, source, press.surface)} />
246          <Button key="save" label="Save PNG" onPress={() => savePng($, source)} />
247          <Button key="hide" label="Hide" onPress={() => update($, latest, () => null)} />
248        </Box>
249        {below ?? null}
250      </Box>
251    )
252  })
253}
254

1declare module 'claude-code' {
2  interface PluginState {
3    'mermaid-inline': { latest: string | null }
4  }
5}
6