mermaid-dianliu

1import type { EngineInterface, Register, RenderElement } from 'claude-code'
2
3import { chunkMarkdown, splitMermaid } from './split'
4
5// 把助理回覆裡已閉合的 ```mermaid 區塊畫成套好主題的 SVG。只改畫面，
6// 對話紀錄不動。轉圖由 node（找不到就 bun）執行 renderer/render.mjs：
7// 排版引擎 ELK 超過 hook 模組可載入的大小，所以不能直接 import。
8// 畫不出來的圖（不支援的圖種、語法錯誤、SVG 過大）維持原本的程式碼區塊。
9
10type Entry =
11  | { status: 'pending' }
12  | { status: 'done'; svg: string }
13  // isFinal：腳本明確回報畫不出來，不必再試；否則是執行環境的問題，之後重試。
14  | { status: 'failed'; isFinal: boolean; attempts: number }
15
16// Svg 元件的原始碼上限。
17const MAX_SVG = 131072
18const MAX_ATTEMPTS = 3
19const RUNTIMES = ['node', 'bun']
20
21// 以 mermaid 原始碼為 key：重畫（捲動、縮放視窗、串流中的每次更新）不會重轉。
22const cache = new Map<string, Entry>()
23
24type Outcome = { svg: string } | { reason: string; isFinal: boolean }
25
26async function render($: EngineInterface, source: string): Promise<Outcome> {
27  const script = `${$.plugin.root}/renderer/render.mjs`
28  let reason = '找不到 node 或 bun'
29  for (const bin of RUNTIMES) {
30    try {
31      const ran = await $.process.run([bin, script], { stdin: source, timeoutMs: 20000 })
32      // 腳本一律以 JSON 回報；啟動不了或讀不到 JSON 就換下一個執行環境。
33      const result = JSON.parse(ran.stdout) as { ok: boolean; svg?: string; reason?: string }
34      if (result.ok && typeof result.svg === 'string') {
35        if (result.svg.length > MAX_SVG) return { reason: 'SVG 超過 Svg 元件上限', isFinal: true }
36        return { svg: result.svg }
37      }
38      return { reason: result.reason ?? '轉圖失敗', isFinal: true }
39    } catch (error) {
40      reason = error instanceof Error ? error.message : String(error)
41    }
42  }
43  return { reason, isFinal: false }
44}
45
46/** 查快取；沒有就在背景開始轉圖，轉完要求重畫。 */
47function lookup($: EngineInterface, source: string): Entry {
48  const known = cache.get(source)
49  if (known && !(known.status === 'failed' && !known.isFinal && known.attempts < MAX_ATTEMPTS)) return known
50  const attempts = known?.status === 'failed' ? known.attempts + 1 : 1
51  if (cache.size > 200) cache.clear()
52  cache.set(source, { status: 'pending' })
53  void render($, source).then(outcome => {
54    if ('svg' in outcome) {
55      cache.set(source, { status: 'done', svg: outcome.svg })
56    } else {
57      cache.set(source, { status: 'failed', isFinal: outcome.isFinal, attempts })
58      $.ui.log(`mermaid-dianliu：這張圖維持程式碼區塊（${outcome.reason}）`, { to: 'debug' })
59    }
60    // 轉圖失敗又還能重試時不主動重畫，等下一次自然重畫再試，避免空轉。
61    if ('svg' in outcome || outcome.isFinal) $.ui.invalidate('ui.render')
62  })
63  return { status: 'pending' }
64}
65
66export const register: Register = on => {
67  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
68    // 這個 mod 只處理桌面版；其他介面照引擎原樣畫。
69    if (e.surface !== 'desktop') return next(e)
70
71    const segments = splitMermaid(e.props.text)
72    if (!segments.some(s => s.kind === 'mermaid')) return next(e)
73
74    const entries = segments.map(s => (s.kind === 'mermaid' ? lookup($, s.source) : undefined))
75    // 全部轉完才一起換上，避免同一則回覆一張一張跳動。
76    if (entries.some(entry => entry?.status === 'pending')) return next(e)
77    if (!entries.some(entry => entry?.status === 'done')) return next(e)
78
79    const { Box, Markdown, Svg } = $.ui.resolve(e)
80    const children: RenderElement[] = []
81    for (const [i, segment] of segments.entries()) {
82      const entry = entries[i]
83      if (segment.kind === 'mermaid' && entry?.status === 'done') {
84        const firstLine = segment.source.trim().split('\n')[0] ?? ''
85        children.push(
86          <Box key={`diagram-${i}`} marginY={1}>
87            <Svg source={entry.svg} alt={`Mermaid 圖：${firstLine.trim()}`} />
88          </Box>,
89        )
90        continue
91      }
92      // 一般文字，或畫不出來而保留程式碼區塊的圖。
93      const text = segment.kind === 'mermaid' ? segment.raw : segment.text
94      if (text.trim() === '') continue
95      if (children.length === 0) {
96        // 開頭的文字交回引擎自己畫，保留它原本的樣子。
97        children.push(await next({ ...e, props: { ...e.props, text } }))
98        continue
99      }
100      const chunks = chunkMarkdown(text)
101      // 有一段長到 Markdown 元件放不下：整則回覆退回引擎原樣。
102      if (chunks === undefined) return next(e)
103      chunks.forEach((chunk, j) => children.push(<Markdown key={`md-${i}-${j}`} text={chunk} />))
104    }
105
106    return <Box flexDirection="column">{children}</Box>
107  })
108}
109

1// 把一段回覆文字切成「一般 markdown」與「mermaid 區塊」。純函式，不碰引擎。
2
3export type Segment = { kind: 'markdown'; text: string } | { kind: 'mermaid'; source: string; raw: string }
4
5// Markdown 元件單段文字的上限。
6export const MAX_MARKDOWN = 10000
7
8// 開頭的 fence：縮排、fence 本身、語言名。
9const OPEN = /^([ \t]*)(`{3,}|~{3,})[ \t]*([^\s`]*)[^`]*$/
10
11function isClose(line: string, fence: string): boolean {
12  const trimmed = line.trim()
13  if (trimmed.length < fence.length) return false
14  // 收尾的 fence：同一種字元、長度不短於開頭、後面沒有別的字。
15  return trimmed.split('').every(c => c === fence[0])
16}
17
18/**
19 * 逐行掃，只認「已閉合」的 mermaid fence：還在串流中的區塊維持原樣，
20 * 包在其他 fence 裡（例如示範用的 ````markdown）的 mermaid 也不動。
21 */
22export function splitMermaid(text: string): Segment[] {
23  const lines = text.split('\n')
24  const segments: Segment[] = []
25  let plain: string[] = []
26  const flush = () => {
27    if (plain.length > 0) segments.push({ kind: 'markdown', text: plain.join('\n') })
28    plain = []
29  }
30
31  let i = 0
32  while (i < lines.length) {
33    const line = lines[i] ?? ''
34    const open = line.match(OPEN)
35    if (!open) {
36      plain.push(line)
37      i += 1
38      continue
39    }
40    const indent = open[1] ?? ''
41    const fence = open[2] ?? ''
42    let end = i + 1
43    while (end < lines.length && !isClose(lines[end] ?? '', fence)) end += 1
44    const closed = end < lines.length
45    const block = lines.slice(i, closed ? end + 1 : lines.length)
46    const body = lines
47      .slice(i + 1, end)
48      .map(l => (l.startsWith(indent) ? l.slice(indent.length) : l))
49      .join('\n')
50    if (closed && (open[3] ?? '').toLowerCase() === 'mermaid' && body.trim() !== '') {
51      flush()
52      segments.push({ kind: 'mermaid', source: body, raw: block.join('\n') })
53    } else {
54      plain.push(...block)
55    }
56    i += block.length
57  }
58  flush()
59  return segments
60}
61
62/**
63 * 把過長的 markdown 在段落空行處切開，每段不超過 max；
64 * 空行在 fence 裡不算段落邊界。單一段落就超過上限時回傳 undefined。
65 */
66export function chunkMarkdown(text: string, max: number = MAX_MARKDOWN): string[] | undefined {
67  if (text.length <= max) return [text]
68  const blocks: string[] = []
69  let current: string[] = []
70  let fence: string | undefined
71  for (const line of text.split('\n')) {
72    if (fence === undefined) {
73      fence = line.match(OPEN)?.[2]
74    } else if (isClose(line, fence)) {
75      fence = undefined
76    }
77    if (fence === undefined && line.trim() === '' && current.length > 0) {
78      blocks.push(current.join('\n'))
79      current = []
80    } else {
81      current.push(line)
82    }
83  }
84  if (current.length > 0) blocks.push(current.join('\n'))
85
86  const chunks: string[] = []
87  let chunk = ''
88  for (const block of blocks) {
89    if (block.length > max) return undefined
90    if (chunk !== '' && chunk.length + 2 + block.length > max) {
91      chunks.push(chunk)
92      chunk = ''
93    }
94    chunk = chunk === '' ? block : `${chunk}\n\n${block}`
95  }
96  if (chunk !== '') chunks.push(chunk)
97  return chunks
98}
99