heatmap

1import type { Register } from 'claude-code'
2
3const TOOL = 'mcp__heatmap__heatmap'
4
5type RGB = [number, number, number]
6type Map = {
7  title: string
8  rowLabels: string[]
9  colLabels: string[]
10  values: number[][]
11  palette: 'heat' | 'green' | 'diverging'
12  unit: string
13}
14
15const clamp = (v: number) => Math.max(0, Math.min(255, Math.round(v)))
16const pack = (c: RGB) => (clamp(c[0]) << 16) | (clamp(c[1]) << 8) | clamp(c[2])
17const DEFAULT = 0x01000000
18const ramp = (stops: RGB[], k: number): RGB => {
19  const t = Math.max(0, Math.min(0.9999, k)) * (stops.length - 1)
20  const i = Math.floor(t)
21  const a = stops[i]!
22  const b = stops[i + 1]!
23  const f = t - i
24  return [a[0] + (b[0] - a[0]) * f, a[1] + (b[1] - a[1]) * f, a[2] + (b[2] - a[2]) * f]
25}
26const PALETTES: Record<Map['palette'], RGB[]> = {
27  heat: [[18, 22, 48], [40, 70, 160], [30, 170, 170], [240, 210, 70], [235, 80, 50]],
28  green: [[22, 27, 34], [14, 68, 41], [0, 109, 50], [38, 166, 65], [57, 211, 83]],
29  diverging: [[50, 90, 200], [120, 170, 230], [235, 235, 235], [240, 150, 100], [200, 50, 50]],
30}
31const lum = (c: RGB) => 0.299 * c[0] + 0.587 * c[1] + 0.114 * c[2]
32
33const fmt = (v: number) =>
34  Math.abs(v) >= 1000 ? `${(v / 1000).toFixed(v >= 10000 ? 0 : 1)}k` : Number.isInteger(v) ? String(v) : v.toFixed(1)
35
36// Printable width-1 characters only: anything else becomes '?'
37const safe = (s: string) => s.replace(/[^\x20-\x7e]/g, '?')
38
39const toBase64 = (words: Uint32Array) =>
40  (new Uint8Array(words.buffer) as unknown as { toBase64: () => string }).toBase64()
41
42const draw = (m: Map, width: number) => {
43  const nr = m.values.length
44  const nc = Math.max(...m.values.map(r => r.length), 1)
45  const labelW = Math.min(16, Math.max(...m.rowLabels.map(l => l.length), 1)) + 1
46  const cw = Math.max(2, Math.min(7, Math.floor((width - labelW) / nc)))
47  const cols = Math.min(512, labelW + cw * nc)
48  const rows = Math.min(256, nr + 1)
49
50  const flat = m.values.flat().filter(v => Number.isFinite(v))
51  const lo = Math.min(...flat)
52  const hi = Math.max(...flat)
53  const span = m.palette === 'diverging' ? Math.max(Math.abs(lo), Math.abs(hi)) || 1 : hi - lo || 1
54  const kOf = (v: number) => (m.palette === 'diverging' ? (v / span + 1) / 2 : (v - lo) / span)
55
56  const words = new Uint32Array(cols * rows * 3)
57  const put = (r: number, c: number, ch: string, fg: number, bg: number) => {
58    if (c >= cols || r >= rows) return
59    const i = (r * cols + c) * 3
60    words[i] = ch.charCodeAt(0)
61    words[i + 1] = fg
62    words[i + 2] = bg
63  }
64  for (let r = 0; r < rows; r++) for (let c = 0; c < cols; c++) put(r, c, ' ', DEFAULT, DEFAULT)
65
66  // header row: column labels, truncated to the cell width
67  for (let j = 0; j < nc; j++) {
68    const label = safe(m.colLabels[j] ?? '').slice(0, cw - (cw > 2 ? 1 : 0))
69    for (let k = 0; k < label.length; k++) put(0, labelW + j * cw + k, label[k]!, 0x9aa4b2, DEFAULT)
70  }
71  for (let i = 0; i < nr; i++) {
72    const label = safe(m.rowLabels[i] ?? '').slice(0, labelW - 1)
73    for (let k = 0; k < label.length; k++) put(i + 1, k, label[k]!, 0xc8cdd5, DEFAULT)
74    for (let j = 0; j < nc; j++) {
75      const v = m.values[i]?.[j]
76      const has = typeof v === 'number' && Number.isFinite(v)
77      const col: RGB = has ? ramp(PALETTES[m.palette], kOf(v)) : [40, 40, 40]
78      const bg = pack(col)
79      const fg = lum(col) > 140 ? 0x101010 : 0xf0f0f0
80      const text = has ? fmt(v) : ''
81      const showText = text.length <= cw - 1
82      for (let k = 0; k < cw - (cw > 2 ? 1 : 0); k++) {
83        const ch = showText ? (text[k - Math.floor((cw - 1 - text.length) / 2)] ?? ' ') : ' '
84        put(i + 1, labelW + j * cw + k, ch, fg, bg)
85      }
86    }
87  }
88  return { cols, rows, cells: toBase64(words), lo, hi }
89}
90
91const legend = (m: Map, lo: number, hi: number, n = 24) => {
92  const words = new Uint32Array(n * 3)
93  for (let k = 0; k < n; k++) {
94    words[k * 3] = 0x2588
95    words[k * 3 + 1] = pack(ramp(PALETTES[m.palette], k / (n - 1)))
96    words[k * 3 + 2] = DEFAULT
97  }
98  return { n, cells: toBase64(words), lo, hi }
99}
100
101// The result travels as JSON text (a string, or text content blocks)
102const parseOutput = (out: unknown): Map | null => {
103  const text = typeof out === 'string' ? out
104    : Array.isArray(out) ? out.map(b => (b && typeof b === 'object' && 'text' in b ? String((b as { text: unknown }).text) : '')).join('')
105    : out && typeof out === 'object' ? JSON.stringify(out) : ''
106  try {
107    const m = JSON.parse(text) as Map
108    return Array.isArray(m.values) ? m : null
109  } catch {
110    return null
111  }
112}
113
114export const register: Register = on => {
115  on('session.start', async ($, e, next) => {
116    await $.tool.register({
117      name: 'heatmap',
118      description:
119        'Draw a truecolor heatmap inline in the chat for the user. Use whenever a grid of numbers ' +
120        '(activity by day x hour, scores by item x criterion, a correlation or comparison matrix) is clearer as colour. ' +
121        'values is rows x columns. palette: heat (default), green (activity/streaks), diverging (signed data, centred on 0). ' +
122        'Keep it under 30 rows and 40 columns; labels short.',
123      inputSchema: {
124        type: 'object',
125        properties: {
126          title: { type: 'string' },
127          rowLabels: { type: 'array', items: { type: 'string' } },
128          colLabels: { type: 'array', items: { type: 'string' } },
129          values: { type: 'array', items: { type: 'array', items: { type: 'number' } } },
130          palette: { type: 'string', enum: ['heat', 'green', 'diverging'] },
131          unit: { type: 'string', description: 'unit shown in the legend, e.g. "msgs" or "%"' },
132        },
133        required: ['title', 'rowLabels', 'colLabels', 'values'],
134      },
135    })
136    return next(e)
137  })
138
139  on('tool.call', { tool: TOOL }, async ($, e) => {
140    const input = e as unknown as Partial<Map>
141    const values = Array.isArray(input.values) ? input.values.slice(0, 30).map(r => (Array.isArray(r) ? r.slice(0, 40).map(Number) : [])) : []
142    if (values.length === 0) return { deny: 'heatmap: values must be a non-empty rows x columns array of numbers.' }
143    const map: Map = {
144      title: String(input.title ?? 'Heatmap'),
145      rowLabels: (input.rowLabels ?? []).map(String),
146      colLabels: (input.colLabels ?? []).map(String),
147      values,
148      palette: input.palette === 'green' || input.palette === 'diverging' ? input.palette : 'heat',
149      unit: String(input.unit ?? ''),
150    }
151    const flat = values.flat().filter(Number.isFinite)
152    return {
153      result: JSON.stringify({
154        ...map,
155        summary: `Drew "${map.title}" (${values.length}x${Math.max(...values.map(r => r.length))}) for the user, range ${Math.min(...flat)} to ${Math.max(...flat)}.`,
156      }),
157    }
158  })
159
160  on('ui.render', { component: 'ToolResult' }, async ($, e, next) => {
161    if (e.props.tool !== TOOL || e.props.isErrored || e.surface !== 'terminal') return next(e)
162    const m = parseOutput(e.props.output)
163    if (!m) return next(e)
164    const { Box, Text, Raster } = $.ui.resolve(e)
165    const width = Math.max(30, (e.viewport?.columns ?? 100) - 6)
166    const g = draw(m, width)
167    const lg = legend(m, g.lo, g.hi)
168    const unit = m.unit ? ` ${m.unit}` : ''
169    return (
170      <Box flexDirection="column">
171        <Text bold>{m.title}</Text>
172        <Raster key={`hm-${e.props.tool_use_id}`} columns={g.cols} rows={g.rows} cells={g.cells} />
173        <Box gap={1}>
174          <Text dimColor>{fmt(g.lo)}{unit}</Text>
175          <Raster key={`lg-${e.props.tool_use_id}`} columns={lg.n} rows={1} cells={lg.cells} />
176          <Text dimColor>{fmt(g.hi)}{unit}</Text>
177        </Box>
178      </Box>
179    )
180  })
181}
182