tui-diagrams

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { DiagramMode, DiagramSettings, DiagramTheme } from '../types'
5
6// Draws ```mermaid, ```plantuml and ```widget fences in Claude's replies as
7// pictures (Claude Code's Image element: kitty, Ghostty, iTerm2, or a terminal
8// forced with CLAUDE_CODE_FORCE_TERMINAL_IMAGES=1) or as Unicode box art (any
9// terminal), in place, in the transcript. Rendering happens in render/render.mjs
10// on the host; the hook only ships source in and reads a JSON line out.
11
12const TEST_PANE = 'diagrams-test'
13const DEFAULTS: DiagramSettings = { mode: 'image', theme: 'dark' }
14
15// The model is NOT told about this mod from here. Where Claude Code's built-in
16// `sec-default` guard loads (machines with managed settings, and Team or
17// Enterprise logins), it skips every user-installed mod's prompt hooks, so a
18// `prompt.compose` hook would reach the model only some of the time. The
19// guidance (fences become pictures; which types render; the shape that fits;
20// when to draw one unprompted) lives in CLAUDE-snippet.md, pasted into the
21// person's ~/.claude/CLAUDE.md, which every session reads. Without it the
22// model tends to inspect the mod before writing a fence (one measured case:
23// almost three minutes).
24
25// Settings live in $.state so every drawing re-runs when they change, and in
26// $.store so they survive the session. `generation` is bumped when a render
27// lands, which redraws every drawing that read it while waiting.
28// `plugin` must equal the manifest's `name`: only a value's owner may write it.
29const settings = atom({ plugin: 'tui-diagrams', key: 'settings' } as const, DEFAULTS)
30const generation = atom({ plugin: 'tui-diagrams', key: 'generation' } as const, 0)
31
32type Kind = 'mermaid' | 'plantuml' | 'widget'
33type Block = { type: 'md'; text: string } | { type: 'diagram'; kind: Kind; source: string }
34type Rendered =
35  | { ok: true; png?: string; width?: number; height?: number; text?: string }
36  | { ok: false; error: string }
37type Slot = { status: 'pending' } | { status: 'done'; result: Rendered }
38
39// One fence: ```mermaid / ```plantuml / ```puml / ```widget, anything after the
40// language ignored, closed by a ``` on its own line. CRLF tolerated.
41//
42// `widget` is a whole HTML document drawn as a picture. It is deliberately NOT
43// spelled ```html: that fence has to keep meaning "here is some HTML source",
44// or every code example in a reply would turn into a screenshot.
45const FENCE = /```[ \t]*(mermaid|plantuml|puml|widget)\b[^\n]*\r?\n([\s\S]*?)\r?\n[ \t]*```/g
46
47function splitFences(text: string): Block[] {
48  const blocks: Block[] = []
49  let last = 0
50  for (const match of text.matchAll(FENCE)) {
51    const at = match.index ?? 0
52    if (at > last) blocks.push({ type: 'md', text: text.slice(last, at) })
53    const kind: Kind = match[1] === 'mermaid' ? 'mermaid' : match[1] === 'widget' ? 'widget' : 'plantuml'
54    let source = (match[2] ?? '').trim()
55    if (kind === 'plantuml' && !/@start\w+/.test(source)) source = `@startuml\n${source}\n@enduml`
56    blocks.push({ type: 'diagram', kind, source })
57    last = at + match[0].length
58  }
59  if (last < text.length) blocks.push({ type: 'md', text: text.slice(last) })
60  return blocks
61}
62
63// There is no "widget" language for the highlighter; a widget is HTML.
64function codeLanguage(kind: Kind): string {
65  return kind === 'widget' ? 'html' : kind
66}
67
68// FNV-1a, enough to key a render cache for one session.
69function hashOf(text: string): string {
70  let h = 0x811c9dc5
71  for (let i = 0; i < text.length; i++) {
72    h ^= text.charCodeAt(i)
73    h = Math.imul(h, 0x01000193) >>> 0
74  }
75  return h.toString(16).padStart(8, '0')
76}
77
78// A terminal cell is about twice as tall as it is wide; the picture is scaled
79// to fill the cell box, so the box must keep the picture's aspect ratio.
80const CELL_ASPECT = 0.5
81const MAX_ROWS = 40
82
83function cellBox(width: number, height: number, maxColumns: number): { columns: number; rows: number } {
84  const aspect = width / Math.max(1, height)
85  // Never upscale past ~8 px per cell, or a small PlantUML picture goes blurry.
86  let columns = Math.max(8, Math.min(maxColumns, Math.ceil(width / 8), 255))
87  let rows = Math.round((columns * CELL_ASPECT) / aspect)
88  if (rows > MAX_ROWS) {
89    rows = MAX_ROWS
90    columns = Math.round((rows * aspect) / CELL_ASPECT)
91  }
92  return { columns: Math.max(1, Math.min(255, columns)), rows: Math.max(1, Math.min(255, rows)) }
93}
94
95// Reads stored settings back, ignoring fields an older version may have written.
96function settingsFrom(value: unknown): DiagramSettings | undefined {
97  if (typeof value !== 'object' || value === null) return undefined
98  const v = value as Record<string, unknown>
99  if (!(v.mode === 'image' || v.mode === 'text') || !(v.theme === 'dark' || v.theme === 'light')) return undefined
100  return { mode: v.mode, theme: v.theme }
101}
102
103// Renders done or in flight this session, by kind + theme + source. A module
104// variable, so a hot reload starts it empty; that costs one re-render per
105// diagram on screen, nothing else.
106const renders = new Map<string, Slot>()
107
108async function runRenderer($: EngineInterface, kind: Kind, source: string, theme: DiagramTheme): Promise<Rendered> {
109  try {
110    // A widget is HTML laid out at a CSS width, so 1600 would be a wall of
111    // whitespace; and it has no box-art form, so png only.
112    const argv = kind === 'widget'
113      ? ['--kind', 'widget', '--format', 'png', '--width', '760', '--theme', theme]
114      : ['--kind', kind, '--format', 'both', '--width', '1600', '--theme', theme]
115    const run = await $.process.run(
116      ['node', `${$.plugin.root}/render/render.mjs`, ...argv],
117      { stdin: source, timeoutMs: 90_000 },
118    )
119    const line = run.stdout.trim().split('\n').pop() ?? ''
120    try {
121      return JSON.parse(line) as Rendered
122    } catch {
123      return { ok: false, error: (run.stderr || run.stdout || `renderer exited ${run.exitCode}`).trim().slice(0, 500) }
124    }
125  } catch (error) {
126    return { ok: false, error: error instanceof Error ? error.message : String(error) }
127  }
128}
129
130// Answers the render when it is done. Otherwise starts it, once, on a timer:
131// a timer runs as a dispatch of its own, so a redraw that abandons the drawing
132// which asked for it (the pane taking focus, the reply still streaming) does
133// not abort the renderer. When it lands, `generation` is bumped and every
134// drawing that read it is drawn again.
135function renderOf($: EngineInterface, kind: Kind, source: string, theme: DiagramTheme): Rendered | undefined {
136  const key = `${kind}|${theme}|${hashOf(source)}`
137  const slot = renders.get(key)
138  if (slot?.status === 'done') return slot.result
139  if (slot) return undefined
140  renders.set(key, { status: 'pending' })
141  $.clock.after(0, () => {
142    void (async () => {
143      const result = await runRenderer($, kind, source, theme)
144      // An abort is the engine's doing, not the diagram's: forget it so the
145      // next drawing asks again. A renderer error stays, or every redraw
146      // would run it again.
147      if (!result.ok && /abort/i.test(result.error)) renders.delete(key)
148      else renders.set(key, { status: 'done', result })
149      await update($, generation, n => (n ?? 0) + 1)
150    })()
151  })
152  return undefined
153}
154
155export const register: Register = on => {
156  on('session.start', async ($, e, next) => {
157    await $.command.register({
158      name: 'diagram',
159      description: 'Diagrams in the transcript: image or text mode, theme, and a terminal test',
160      argumentHint: '[image|text|dark|light|test]',
161    })
162    const stored = settingsFrom(await $.store.get('settings'))
163    if (stored) await update($, settings, () => stored)
164    return next(e)
165  })
166
167  on('command.run', { command: 'diagram' }, async ($, e) => {
168    const arg = e.args.trim().toLowerCase()
169    const current = await read($, settings)
170
171    if (arg === 'test') {
172      const opened = await $.ui.open({ id: TEST_PANE, title: 'Diagram test', focus: true })
173      return {
174        text:
175          `Diagram test pane opened (placed=${opened.isPlaced}, fullscreen=${e.presentation.isFullscreen}, ` +
176          `columns=${e.presentation.columns}). Dim alt text in place of a picture means Claude Code decided ` +
177          `this terminal does not draw images; blank rows mean it sent the picture and the terminal lost it ` +
178          `(on Windows: ConPTY). render/kitty-probe.py tests the terminal without Claude Code in the loop.`,
179      }
180    }
181
182    let next: DiagramSettings | undefined
183    if (arg === 'image' || arg === 'text') next = { ...current, mode: arg as DiagramMode }
184    if (arg === 'dark' || arg === 'light') next = { ...current, theme: arg as DiagramTheme }
185    if (next) {
186      await $.store.set('settings', next)
187      await update($, settings, () => next)
188      return { text: `Diagrams: mode=${next.mode}, theme=${next.theme} (saved).` }
189    }
190
191    return {
192      text:
193        `Diagrams: mode=${current.mode}, theme=${current.theme}\n` +
194        `/diagram image   pictures where the terminal draws them (kitty, Ghostty, iTerm2, or CLAUDE_CODE_FORCE_TERMINAL_IMAGES=1); elsewhere the box art shows as alt text\n` +
195        `/diagram text    Unicode box art in every terminal\n` +
196        `/diagram dark|light   picture theme\n` +
197        `/diagram test    open a pane with a test picture\n` +
198        `The model learns to use fences from CLAUDE-snippet.md in the mod folder, pasted into ~/.claude/CLAUDE.md.`,
199    }
200  })
201
202  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
203    if (e.surface !== 'terminal' || e.props.isSummary) return next(e)
204    const blocks = splitFences(e.props.text)
205    if (!blocks.some(b => b.type === 'diagram')) return next(e)
206
207    await read($, generation)
208    const { mode, theme } = await read($, settings)
209    const { Box, Text, Markdown, Code, Image } = $.ui.resolve(e)
210    const gutter = 2
211    const maxColumns = Math.max(20, (e.viewport?.columns ?? 80) - gutter - 1)
212
213    const parts: ReturnType<typeof Box>[] = []
214    for (const block of blocks) {
215      if (block.type === 'md') {
216        if (block.text.trim()) parts.push(<Box><Markdown text={block.text.replace(/^\n+|\n+$/g, '')} /></Box>)
217        continue
218      }
219      const r = renderOf($, block.kind, block.source, theme)
220      if (!r) {
221        parts.push(
222          <Box flexDirection="column">
223            <Code language={codeLanguage(block.kind)} source={block.source} />
224            <Text dimColor>rendering {block.kind} diagram…</Text>
225          </Box>,
226        )
227        continue
228      }
229      if (!r.ok) {
230        parts.push(
231          <Box flexDirection="column">
232            <Code language={codeLanguage(block.kind)} source={block.source} />
233            <Text color="red">{block.kind} diagram did not render: {r.error}</Text>
234          </Box>,
235        )
236        continue
237      }
238      if (mode === 'text' || !r.png || !r.width || !r.height) {
239        parts.push(
240          r.text
241            ? <Box paddingY={1}><Text>{r.text}</Text></Box>
242            : <Box flexDirection="column"><Code language={codeLanguage(block.kind)} source={block.source} /><Text dimColor>({block.kind}: no text rendering for this diagram type; /diagram image draws it)</Text></Box>,
243        )
244        continue
245      }
246      const box = cellBox(r.width, r.height, maxColumns)
247      const alt = r.text ?? `[${block.kind} diagram: this terminal does not draw images; /diagram text shows it as box art]`
248      parts.push(
249        <Box paddingY={1}>
250          <Image source={{ png: r.png }} columns={box.columns} rows={box.rows} alt={alt} />
251        </Box>,
252      )
253    }
254
255    return (
256      <Box flexDirection="row">
257        <Box width={gutter}><Text>{e.props.isFirstOfReply ? '●' : ''}</Text></Box>
258        <Box flexDirection="column" flexGrow={1}>{parts}</Box>
259      </Box>
260    )
261  })
262
263  // `/diagram test`: a pane with a plain PNG and a rendered Mermaid sample, so
264  // "does this terminal draw pictures" is answered without a reply involved.
265  on('ui.render', { component: 'Pane', requestId: TEST_PANE }, async ($, e) => {
266    const facts = `viewport ${e.viewport?.columns ?? '?'}x${e.viewport?.rows ?? '?'} fullscreen=${String(e.viewport?.isFullscreen)} body=${e.props.bodyColumns}`
267    if (e.surface !== 'terminal') {
268      const { Box, Text } = $.ui.resolve(e)
269      return <Box flexDirection="column"><Text>{facts}</Text><Text dimColor>Image is a terminal-only element.</Text></Box>
270    }
271    const { Box, Text, Image } = $.ui.resolve(e)
272    await read($, generation)
273    const { theme } = await read($, settings)
274    const sample = 'flowchart LR\n  A[Mermaid source] --> B[render.mjs]\n  B --> C{Terminal draws images?}\n  C -->|yes| D[PNG inline]\n  C -->|no| E[Unicode box art]'
275    const r = renderOf($, 'mermaid', sample, theme)
276    const { base64 } = await $.fs.read(`${$.plugin.root}/assets/test.png`, { as: 'bytes' })
277    const width = Math.max(20, e.props.bodyColumns - 2)
278    let sampleRow: ReturnType<typeof Box> | ReturnType<typeof Text> | ReturnType<typeof Image>
279    if (!r) sampleRow = <Text dimColor>rendering the sample…</Text>
280    else if (!r.ok) sampleRow = <Text color="red">render failed: {r.error}</Text>
281    else if (!r.png || !r.width || !r.height) sampleRow = <Text color="red">render failed: no picture</Text>
282    else {
283      const box = cellBox(r.width, r.height, width)
284      sampleRow = <Image source={{ png: r.png }} columns={box.columns} rows={box.rows} alt={r.text ?? '[mermaid diagram]'} />
285    }
286    return (
287      <Box flexDirection="column">
288        <Text dimColor>{facts}</Text>
289        <Text>1. A plain PNG (gradient, green border, white diagonal):</Text>
290        <Image source={{ png: base64 }} columns={Math.min(40, width)} rows={10} alt="[no picture: this terminal did not draw the PNG]" />
291        <Text>2. A rendered Mermaid flowchart; where pictures are not drawn, its box art shows instead:</Text>
292        {sampleRow}
293        <Text dimColor>Esc closes this pane.</Text>
294      </Box>
295    )
296  })
297}
298

1export type DiagramMode = 'image' | 'text'
2export type DiagramTheme = 'dark' | 'light'
3export type DiagramSettings = { mode: DiagramMode; theme: DiagramTheme }
4
5declare module 'claude-code' {
6  interface PluginState {
7    'tui-diagrams': { settings: DiagramSettings; generation: number }
8  }
9}
10