mermaid-ascii

1import type { Register } from 'claude-code';
2
3const FENCE = /^ {0,3}(`{3,}|~{3,})(.*)$/;
4
5type Diagram = { start: number; end: number; source: string };
6
7// Tracks every fence, not just mermaid ones: a mermaid fence quoted inside
8// an outer fence is an example to copy, not a diagram to draw.
9export function findDiagrams(text: string): Diagram[] {
10  const diagrams: Diagram[] = [];
11  let open:
12    | { marker: string; start: number; body: number; isMermaid: boolean }
13    | undefined;
14  let at = 0;
15  for (const line of text.split('\n')) {
16    const lineEnd = at + line.length;
17    const fence = FENCE.exec(line);
18    if (fence) {
19      const [, marker = '', info = ''] = fence;
20      if (!open) {
21        if (!(marker[0] === '`' && info.includes('`'))) {
22          const isMermaid = info.trim().split(/\s/)[0] === 'mermaid';
23          open = { marker, start: at, body: lineEnd + 1, isMermaid };
24        }
25      } else if (
26        marker[0] === open.marker[0] &&
27        marker.length >= open.marker.length &&
28        info.trim() === ''
29      ) {
30        if (open.isMermaid) {
31          const source = text.slice(open.body, at);
32          diagrams.push({ start: open.start, end: lineEnd, source });
33        }
34        open = undefined;
35      }
36    }
37    at = lineEnd + 1;
38  }
39  return diagrams;
40}
41
42// The transcript indents a reply past its bullet, so the full width would wrap.
43const INDENT = 4;
44
45export const register: Register = (on) => {
46  // A streaming reply redraws on every flush; each diagram renders once. A
47  // diagram the tool rejects stays cached as undefined, but a run that could
48  // not finish (binary not installed yet, timeout) is dropped so a redraw
49  // retries it.
50  const cache = new Map<string, Promise<string | undefined>>();
51  let bullet: Promise<string> | undefined;
52
53  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
54    const { text } = e.props;
55    const columns = e.viewport ? e.viewport.columns - INDENT : undefined;
56
57    const toAscii = (source: string) => {
58      const key = `${columns}\0${source}`;
59      let ascii = cache.get(key);
60      if (!ascii) {
61        const width = columns ? ['--max-width', String(columns)] : [];
62        ascii = $.process
63          .run(['mermaid-ascii', '-f', '-', ...width], { stdin: source })
64          .then((r) => (r.exitCode === 0 ? r.stdout.trimEnd() : undefined));
65        ascii.catch(() => cache.delete(key));
66        cache.set(key, ascii);
67      }
68      return ascii;
69    };
70
71    const diagrams = findDiagrams(text);
72    const rendered = await Promise.all(diagrams.map((d) => toAscii(d.source)));
73    if (rendered.every((r) => r === undefined)) return next(e);
74
75    const { Box, Text } = $.ui.resolve(e);
76    const parts: Array<{ markdown: string } | { ascii: string }> = [];
77    let markdown = '';
78    let at = 0;
79    diagrams.forEach((d, i) => {
80      const ascii = rendered[i];
81      if (ascii === undefined) return;
82      markdown += text.slice(at, d.start);
83      if (markdown.trim()) parts.push({ markdown });
84      parts.push({ ascii });
85      markdown = '';
86      at = d.end;
87    });
88    markdown += text.slice(at);
89    if (markdown.trim()) parts.push({ markdown });
90
91    // Prose goes back through the engine so it keeps the reply's bullet and
92    // spacing. Only the first block draws the bullet, and the engine indents no
93    // block after it, so those take the bullet's width as padding.
94    const drawn = await Promise.all(
95      parts.map(async (p, i) => {
96        if ('markdown' in p) {
97          const isFirst = e.props.isFirstOfReply && i === 0;
98          const prose = await next({
99            ...e,
100            props: {
101              ...e.props,
102              text: p.markdown.replace(/^\s*\n/, '').trimEnd(),
103              isFirstOfReply: isFirst,
104            },
105          });
106          return i === 0 ? prose : <Box paddingLeft={2}>{prose}</Box>;
107        }
108        const diagram = (
109          <Box flexDirection="column">
110            {p.ascii.split('\n').map((line) => (
111              <Text wrap="truncate-end">
112                {runs(line).map((r) => (
113                  <Text {...PAINT[r.kind]}>{r.text}</Text>
114                ))}
115              </Text>
116            ))}
117          </Box>
118        );
119        if (!(e.props.isFirstOfReply && i === 0)) {
120          return (
121            <Box paddingLeft={2} marginTop={1}>
122              {diagram}
123            </Box>
124          );
125        }
126        // The engine draws no bullet for an empty block, so a reply that opens
127        // with a diagram draws its own, using the glyph the engine picks per OS.
128        bullet ??= $.process
129          .run(['uname'])
130          .then((r) => (r.stdout.trim() === 'Darwin' ? '⏺' : '●'));
131        return (
132          <Box flexDirection="row" marginTop={1}>
133            <Box width={2}>
134              <Text color="text">{await bullet}</Text>
135            </Box>
136            {diagram}
137          </Box>
138        );
139      }),
140    );
141    return <Box flexDirection="column">{drawn}</Box>;
142  }).catch((_$, e, next) => next(e));
143};
144
145type Kind = 'border' | 'arrow' | 'label';
146
147const PAINT: Record<Kind, { color?: string; bold?: boolean }> = {
148  border: { color: 'subtle' },
149  arrow: { color: 'claude' },
150  label: { bold: true },
151};
152
153function kindOf(ch: string): Kind {
154  if (/[►◄▲▼]/.test(ch)) return 'arrow';
155  return /[─-╿]/.test(ch) ? 'border' : 'label';
156}
157
158function runs(line: string): Array<{ kind: Kind; text: string }> {
159  const out: Array<{ kind: Kind; text: string }> = [];
160  for (const ch of line) {
161    const kind = ch === ' ' ? (out.at(-1)?.kind ?? 'label') : kindOf(ch);
162    const last = out.at(-1);
163    if (last?.kind === kind) last.text += ch;
164    else out.push({ kind, text: ch });
165  }
166  return out;
167}
168