diagram

1import type { Register } from 'claude-code'
2
3const TOOL = 'mcp__diagram__diagram'
4
5type Tone = 'start' | 'step' | 'decision' | 'end' | 'warn' | 'good'
6type Node = { id: string; label: string; note?: string; tone?: Tone }
7type Edge = { from: string; to: string; label?: string }
8type Chart = { title: string; nodes: Node[]; edges: Edge[] }
9
10const TONE: Record<Tone, { border: string; text: string; style: string }> = {
11  start: { border: 'cyan', text: 'cyan', style: 'round' },
12  step: { border: 'blue', text: 'white', style: 'round' },
13  decision: { border: 'yellow', text: 'yellow', style: 'double' },
14  end: { border: 'green', text: 'green', style: 'bold' },
15  warn: { border: 'red', text: 'red', style: 'round' },
16  good: { border: 'green', text: 'green', style: 'round' },
17}
18
19// Layered layout: each node's rank is its longest path from a source
20const layers = (c: Chart): Node[][] => {
21  const rank = new Map<string, number>(c.nodes.map(n => [n.id, 0]))
22  for (let pass = 0; pass < c.nodes.length; pass++) {
23    let changed = false
24    for (const e of c.edges) {
25      const r = (rank.get(e.from) ?? 0) + 1
26      if (rank.has(e.to) && r > (rank.get(e.to) ?? 0) && r < c.nodes.length) {
27        rank.set(e.to, r)
28        changed = true
29      }
30    }
31    if (!changed) break
32  }
33  const out: Node[][] = []
34  for (const n of c.nodes) {
35    const r = rank.get(n.id) ?? 0
36    ;(out[r] ??= []).push(n)
37  }
38  return out.filter(Boolean)
39}
40
41const parse = (out: unknown): Chart | null => {
42  const text = typeof out === 'string' ? out
43    : Array.isArray(out) ? out.map(b => (b && typeof b === 'object' && 'text' in b ? String((b as { text: unknown }).text) : '')).join('')
44    : ''
45  try {
46    const c = JSON.parse(text) as Chart
47    return Array.isArray(c.nodes) && c.nodes.length > 0 ? c : null
48  } catch {
49    return null
50  }
51}
52
53export const register: Register = on => {
54  on('session.start', async ($, e, next) => {
55    await $.tool.register({
56      name: 'diagram',
57      description:
58        'Draw a flowchart inline in the chat (boxes and arrows, laid out top to bottom). Use to explain a process, a decision, ' +
59        'a framework or an architecture visually, like a mermaid chart. Keep labels under 28 characters, notes under 50, at most 14 nodes. ' +
60        'tone per node: start | step | decision | end | warn | good. Edges may carry a short label (e.g. "yes", "no").',
61      inputSchema: {
62        type: 'object',
63        properties: {
64          title: { type: 'string' },
65          nodes: {
66            type: 'array',
67            items: {
68              type: 'object',
69              properties: {
70                id: { type: 'string' }, label: { type: 'string' }, note: { type: 'string' },
71                tone: { type: 'string', enum: ['start', 'step', 'decision', 'end', 'warn', 'good'] },
72              },
73              required: ['id', 'label'],
74            },
75          },
76          edges: {
77            type: 'array',
78            items: {
79              type: 'object',
80              properties: { from: { type: 'string' }, to: { type: 'string' }, label: { type: 'string' } },
81              required: ['from', 'to'],
82            },
83          },
84        },
85        required: ['title', 'nodes', 'edges'],
86      },
87    })
88    return next(e)
89  })
90
91  on('tool.call', { tool: TOOL }, async ($, e) => {
92    const input = e as unknown as Partial<Chart>
93    const nodes = (input.nodes ?? []).slice(0, 14).map(n => ({ ...n, id: String(n.id), label: String(n.label) }))
94    if (nodes.length === 0) return { deny: 'diagram: give at least one node.' }
95    const ids = new Set(nodes.map(n => n.id))
96    const edges = (input.edges ?? []).filter(x => ids.has(String(x.from)) && ids.has(String(x.to)))
97    return { result: JSON.stringify({ title: String(input.title ?? 'Diagram'), nodes, edges }) }
98  })
99
100  on('ui.render', { component: 'ToolResult' }, async ($, e, next) => {
101    if (e.props.tool !== TOOL || e.props.isErrored) return next(e)
102    const c = parse(e.props.output)
103    if (!c) return next(e)
104    const { Box, Text } = $.ui.resolve(e)
105    const rows = layers(c)
106    const width = Math.max(40, (e.viewport?.columns ?? 100) - 8)
107    const labelOf = new Map(c.nodes.map(n => [n.id, n.label]))
108
109    return (
110      <Box flexDirection="column" alignItems="center" width={width}>
111        <Text bold>{c.title}</Text>
112        {rows.map((layer, i) => {
113          const ids = new Set(layer.map(n => n.id))
114          const out = c.edges.filter(x => ids.has(x.from))
115          return (
116            <Box key={`layer-${i}`} flexDirection="column" alignItems="center">
117              <Box flexDirection="row" gap={2} justifyContent="center" flexWrap="wrap">
118                {layer.map(n => {
119                  const t = TONE[n.tone ?? 'step'] ?? TONE.step
120                  return (
121                    <Box key={`n-${n.id}`} borderStyle={t.style} borderColor={t.border} paddingX={1} flexDirection="column" alignItems="center">
122                      <Text bold color={t.text}>{n.tone === 'decision' ? `◆ ${n.label} ?` : n.label}</Text>
123                      {n.note ? <Text dimColor>{n.note}</Text> : null}
124                    </Box>
125                  )
126                })}
127              </Box>
128              {i < rows.length - 1 && out.length > 0 && (
129                <Box flexDirection="column" alignItems="center">
130                  <Text color="gray">│</Text>
131                  {out.some(x => x.label) || layer.length > 1 || out.length > 1
132                    ? out.map((x, k) => (
133                      <Text key={`e-${i}-${k}`} dimColor>
134                        {labelOf.get(x.from)}<Text color="magenta"> ──{x.label ? ` ${x.label} ` : ''}──▶ </Text>{labelOf.get(x.to)}
135                      </Text>
136                    ))
137                    : null}
138                  <Text color="gray">▼</Text>
139                </Box>
140              )}
141            </Box>
142          )
143        })}
144      </Box>
145    )
146  })
147}
148