mermaid-render

1import type { EngineInterface, Register } from 'claude-code'
2
3import { type Segment, splitMermaid } from './fences'
4
5type Mode = 'svg' | 'ascii'
6type Rendered = { ok: true; output: string } | { ok: false; error: string }
7type Drawn = Exclude<Segment, { kind: 'mermaid' }> | (Extract<Segment, { kind: 'mermaid' }> & { result: Rendered })
8
9// Room the terminal keeps beside a reply's text (bullet and margins).
10const TERMINAL_GUTTER = 4
11const RENDERER_FAILED = 'renderer failed'
12
13// Keyed by mode and source: a streaming reply redraws often, and each diagram renders once.
14const cache = new Map<string, Promise<Rendered>>()
15
16const renderSegments = async ($: EngineInterface, mode: Mode, segments: Segment[]): Promise<Drawn[]> => {
17  const keyOf = (source: string) => `${mode}\0${source}`
18  const missing = [
19    ...new Set(segments.flatMap(s => (s.kind === 'mermaid' && !cache.has(keyOf(s.source)) ? [s.source] : []))),
20  ]
21
22  if (missing.length > 0) {
23    const batch = runRenderer($, mode, missing)
24    missing.forEach((source, index) => {
25      const result = batch.then(results => results[index] ?? failed('renderer returned too few results'))
26      cache.set(keyOf(source), result)
27      // A renderer that could not run says nothing about the diagram, so the next draw tries again.
28      void result.then(r => {
29        if (!r.ok && r.error.startsWith(RENDERER_FAILED)) cache.delete(keyOf(source))
30      })
31    })
32  }
33
34  return Promise.all(
35    segments.map(async (segment): Promise<Drawn> => {
36      if (segment.kind !== 'mermaid') return segment
37      const result = (await cache.get(keyOf(segment.source))) ?? failed('diagram was not rendered')
38      return { ...segment, result }
39    }),
40  )
41}
42
43export const register: Register = on => {
44  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
45    const segments = splitMermaid(e.props.text)
46
47    if (!segments.some(s => s.kind === 'mermaid')) {
48      return next(e)
49    }
50
51    // The terminal has no vector element, so it gets ASCII art inside the reply's own markdown.
52    if (e.surface === 'terminal') {
53      const maxWidth = (e.viewport?.columns ?? Number.POSITIVE_INFINITY) - TERMINAL_GUTTER
54      const text = (await renderSegments($, 'ascii', segments))
55        .map(segment => {
56          if (segment.kind === 'markdown') return segment.text
57          const { result } = segment
58          if (!result.ok) return `${segment.fence}\n\n_mermaid-render: ${result.error}_`
59          const width = Math.max(...result.output.split('\n').map(line => line.length))
60          if (width > maxWidth) {
61            return `${segment.fence}\n\n_mermaid-render: diagram is ${width} columns wide, the terminal has ${maxWidth}_`
62          }
63          return `\`\`\`text\n${result.output}\n\`\`\``
64        })
65        .join('')
66
67      return next({ ...e, props: { ...e.props, text } })
68    }
69
70    const drawn = await renderSegments($, 'svg', segments)
71    const { Box, Markdown, Svg } = $.ui.resolve(e)
72
73    return (
74      <Box flexDirection="column">
75        {drawn.map(segment => {
76          if (segment.kind === 'markdown') {
77            return segment.text.trim() === '' ? null : <Markdown text={segment.text} />
78          }
79          const { result } = segment
80          if (!result.ok) {
81            return <Markdown text={`${segment.fence}\n\n_mermaid-render: ${result.error}_`} />
82          }
83          const header = segment.source.trim().split('\n')[0]
84          return <Svg source={result.output} alt={`Mermaid diagram: ${header}`} />
85        })}
86      </Box>
87    )
88  })
89}
90
91const failed = (reason: string): Rendered => ({ ok: false, error: `${RENDERER_FAILED}: ${reason}` })
92
93// One process per batch: the renderer needs Node and beautiful-mermaid, which the hooks module does not have.
94const runRenderer = async ($: EngineInterface, mode: Mode, sources: string[]): Promise<Rendered[]> => {
95  try {
96    const { exitCode, stdout, stderr } = await $.process.run(['node', 'render.mjs'], {
97      cwd: `${$.plugin.root}/renderer`,
98      stdin: JSON.stringify({ mode, sources }),
99    })
100    if (exitCode !== 0) {
101      return sources.map(() => failed(stderr.trim().split('\n')[0] || `exit code ${exitCode}`))
102    }
103    return JSON.parse(stdout) as Rendered[]
104  } catch (error) {
105    return sources.map(() => failed(error instanceof Error ? error.message : String(error)))
106  }
107}
108

1export type Segment = { kind: 'markdown'; text: string } | { kind: 'mermaid'; source: string; fence: string }
2
3// Closed fences only: while a reply streams, an unclosed block stays plain code until its closing fence arrives.
4const MERMAID_FENCE = /^ {0,3}(`{3,}|~{3,})[ \t]*mermaid[ \t]*\n([\s\S]*?)\n {0,3}\1[ \t]*$/gm
5
6export const splitMermaid = (text: string): Segment[] => {
7  const segments: Segment[] = []
8  let last = 0
9
10  for (const match of text.matchAll(MERMAID_FENCE)) {
11    const start = match.index ?? 0
12    if (start > last) segments.push({ kind: 'markdown', text: text.slice(last, start) })
13    segments.push({ kind: 'mermaid', source: match[2] ?? '', fence: match[0] })
14    last = start + match[0].length
15  }
16
17  if (last < text.length) segments.push({ kind: 'markdown', text: text.slice(last) })
18
19  return segments
20}
21