mermaid-render

1import type { EngineInterface, Register, RenderElement, RenderInput, RenderSurface, RenderViewport } from 'claude-code'
2import { parseSegments, type Segment } from './segments'
3
4type Format = 'svg' | 'ascii'
5type Rendered = { ok: true; output: string } | { ok: false; reason: string }
6type Diagram = { kind: 'diagram'; output: string; alt: string }
7type Block = { kind: 'markdown'; text: string } | Diagram
8
9const FORMAT_BY_SURFACE = {
10  desktop: 'svg',
11  vscode: 'svg',
12  mobile: 'svg',
13  terminal: 'ascii',
14} as const satisfies Record<RenderSurface, Format>
15
16const widestLine = (text: string) => Math.max(...text.split('\n').map(line => [...line].length))
17
18const FITS: Record<Format, (output: string, viewport: RenderViewport | undefined) => boolean> = {
19  svg: () => true,
20  ascii: (output, viewport) => viewport === undefined || widestLine(output) <= viewport.columns,
21}
22
23// The renderer runs out of process: its bundle (beautiful-mermaid with elkjs) is 3.4 MB,
24// and the hooks module may not import a file over 1 MiB.
25const RENDERER = 'renderer/render.mjs'
26const RUNTIMES = ['node', '/opt/homebrew/bin/node', '/usr/local/bin/node', 'bun', '/opt/homebrew/bin/bun']
27const RENDER_TIMEOUT_MS = 20_000
28const CACHE_LIMIT = 256
29
30const cache = new Map<string, Promise<Rendered>>()
31let runtime: Promise<string | undefined> | undefined
32let hasLoggedNoRuntime = false
33
34function findRuntime($: EngineInterface): Promise<string | undefined> {
35  runtime ??= (async () => {
36    for (const candidate of RUNTIMES) {
37      try {
38        await $.process.run([candidate, '--version'], { timeoutMs: 5_000 })
39        return candidate
40      } catch {}
41    }
42    // A run can also reject because its dispatch was abandoned, so a miss is retried on the next draw.
43    runtime = undefined
44    if (!hasLoggedNoRuntime) {
45      hasLoggedNoRuntime = true
46      $.ui.log(`no node or bun could start (tried ${RUNTIMES.join(', ')}); mermaid fences stay code`, { to: 'debug' })
47    }
48    return undefined
49  })()
50  return runtime
51}
52
53function parseRendered(stdout: string, count: number): Rendered[] {
54  const parsed: unknown = JSON.parse(stdout)
55  const isRendered = (value: unknown): value is Rendered =>
56    typeof value === 'object' &&
57    value !== null &&
58    (('ok' in value && value.ok === true && 'output' in value && typeof value.output === 'string') ||
59      ('ok' in value && value.ok === false && 'reason' in value && typeof value.reason === 'string'))
60  if (!Array.isArray(parsed) || parsed.length !== count || !parsed.every(isRendered)) {
61    throw new Error(`renderer answered ${stdout.slice(0, 80)}`)
62  }
63  return parsed
64}
65
66async function runRenderer($: EngineInterface, format: Format, sources: string[]): Promise<Rendered[]> {
67  const argv0 = await findRuntime($)
68  if (argv0 === undefined) throw new Error('no runtime')
69  const { exitCode, stdout, stderr } = await $.process.run([argv0, `${$.plugin.root}/${RENDERER}`], {
70    stdin: JSON.stringify({ format, sources }),
71    timeoutMs: RENDER_TIMEOUT_MS,
72  })
73  if (exitCode !== 0) throw new Error(`renderer exited ${exitCode}: ${stderr.slice(0, 200)}`)
74  return parseRendered(stdout, sources.length)
75}
76
77function remember(key: string, rendered: Promise<Rendered>): void {
78  cache.set(key, rendered)
79  for (const oldest of cache.keys()) {
80    if (cache.size <= CACHE_LIMIT) break
81    cache.delete(oldest)
82  }
83}
84
85function renderAll($: EngineInterface, format: Format, sources: string[]): Promise<Rendered[]> {
86  const keyOf = (source: string) => `${format}\0${source}`
87  const missing = [...new Set(sources)].filter(source => !cache.has(keyOf(source)))
88  if (missing.length > 0) {
89    const batch = runRenderer($, format, missing).then(
90      results => ({ isTransportFailure: false, results }),
91      (error: unknown) => ({
92        isTransportFailure: true,
93        results: missing.map((): Rendered => ({ ok: false, reason: String(error) })),
94      }),
95    )
96    missing.forEach((source, index) => {
97      const key = keyOf(source)
98      const rendered = batch.then(({ isTransportFailure, results }) => {
99        if (isTransportFailure && cache.get(key) === rendered) cache.delete(key)
100        return results[index] ?? { ok: false as const, reason: 'missing result' }
101      })
102      remember(key, rendered)
103    })
104  }
105  return Promise.all(
106    sources.map(source => cache.get(keyOf(source)) ?? Promise.resolve({ ok: false as const, reason: 'evicted' })),
107  )
108}
109
110function toBlocks(segments: Segment[], rendered: Rendered[]): Block[] {
111  const blocks: Block[] = []
112  let diagramIndex = 0
113  const pushMarkdown = (text: string) => {
114    const last = blocks.at(-1)
115    if (last?.kind !== 'markdown') blocks.push({ kind: 'markdown', text })
116    else last.text += (last.text.endsWith('\n') || text.startsWith('\n') ? '' : '\n') + text
117  }
118  for (const segment of segments) {
119    if (segment.kind === 'markdown') {
120      pushMarkdown(segment.text)
121      continue
122    }
123    const result = rendered[diagramIndex++]
124    if (result?.ok === true) {
125      const alt = segment.source.trim().split('\n')[0]?.trim() || 'mermaid diagram'
126      blocks.push({ kind: 'diagram', output: result.output, alt })
127    } else {
128      pushMarkdown(segment.fence)
129    }
130  }
131  return blocks
132}
133
134export const register: Register = on => {
135  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
136    if (e.props.isSummary === true) return next(e)
137    const segments = parseSegments(e.props.text)
138    const sources = segments.flatMap(segment => (segment.kind === 'mermaid' ? [segment.source] : []))
139    if (sources.length === 0) return next(e)
140
141    const format = FORMAT_BY_SURFACE[e.surface]
142    const fitting = (result: Rendered): Rendered =>
143      result.ok && !FITS[format](result.output, e.viewport) ? { ok: false, reason: 'too wide' } : result
144    const blocks = toBlocks(segments, (await renderAll($, format, sources)).map(fitting))
145    if (blocks.every(block => block.kind === 'markdown')) return next(e)
146
147    return drawBlocks($, e, blocks)
148  }).catch(($, e, next) => next(e))
149}
150
151function diagramDrawer($: EngineInterface, e: RenderInput<'AssistantMessage'>): (diagram: Diagram) => RenderElement {
152  if (e.surface === 'terminal') {
153    const { Code } = $.ui.resolve(e)
154    return ({ output }) => <Code source={output} wrap="truncate-end" />
155  }
156  const { Svg } = $.ui.resolve(e)
157  return ({ output, alt }) => <Svg source={output} alt={alt} />
158}
159
160function drawBlocks($: EngineInterface, e: RenderInput<'AssistantMessage'>, blocks: Block[]): RenderElement {
161  const { Box, Markdown } = $.ui.resolve(e)
162  const drawDiagram = diagramDrawer($, e)
163  const body = (
164    <Box flexDirection="column" gap={1}>
165      {blocks.map(block => (block.kind === 'markdown' ? <Markdown text={block.text} /> : drawDiagram(block)))}
166    </Box>
167  )
168  if (e.surface !== 'terminal') return body
169  // The terminal's reply bullet and gutter belong to the component this tree replaces.
170  const { Text } = $.ui.resolve(e)
171  return (
172    <Box flexDirection="row">
173      <Text>{e.props.isFirstOfReply ? '⏺ ' : '  '}</Text>
174      {body}
175    </Box>
176  )
177}
178

1export type Segment =
2  | { kind: 'markdown'; text: string }
3  | { kind: 'mermaid'; source: string; fence: string }
4
5const OPENER = /^( {0,3})(`{3,}|~{3,})(.*)$/
6
7type Fence = { char: string; length: number; indent: number; isMermaid: boolean; start: number; body: string[] }
8
9function openFence(line: string, start: number): Fence | undefined {
10  const match = OPENER.exec(line)
11  if (match === null) return undefined
12  const [, indent = '', marker = '', info = ''] = match
13  if (marker[0] === '`' && info.includes('`')) return undefined
14  const language = info.trim().split(/\s+/)[0]?.toLowerCase()
15  return { char: marker[0] ?? '`', length: marker.length, indent: indent.length, isMermaid: language === 'mermaid', start, body: [] }
16}
17
18function closes(fence: Fence, line: string): boolean {
19  const match = /^ {0,3}(`+|~+)[ \t]*$/.exec(line)
20  const marker = match?.[1]
21  return marker !== undefined && marker[0] === fence.char && marker.length >= fence.length
22}
23
24function dedent(line: string, indent: number): string {
25  let cut = 0
26  while (cut < indent && line[cut] === ' ') cut++
27  return line.slice(cut)
28}
29
30export function parseSegments(text: string): Segment[] {
31  const lines = text.split(/(?<=\n)/)
32  const segments: Segment[] = []
33  let fence: Fence | undefined
34  let offset = 0
35  let markdownStart = 0
36
37  for (const raw of lines) {
38    const line = raw.replace(/\r?\n$/, '')
39    if (fence === undefined) {
40      fence = openFence(line, offset)
41    } else if (closes(fence, line)) {
42      if (fence.isMermaid) {
43        segments.push({ kind: 'markdown', text: text.slice(markdownStart, fence.start) })
44        segments.push({
45          kind: 'mermaid',
46          source: fence.body.join('\n'),
47          fence: text.slice(fence.start, offset + line.length),
48        })
49        markdownStart = offset + line.length
50      }
51      fence = undefined
52    } else {
53      fence.body.push(dedent(line, fence.indent))
54    }
55    offset += raw.length
56  }
57  segments.push({ kind: 'markdown', text: text.slice(markdownStart) })
58
59  return segments.filter(segment => segment.kind === 'mermaid' || segment.text.trim() !== '')
60}
61