pretty-markdown

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { FoldId } from '../types'
5import { asHunk, parse, spanWidth, wrapSpans } from './parse'
6import type { CalloutType, Line, Part, Span } from './parse'
7
8const CHIP_BG = '#2f3440'
9const CHIP_FG = '#e5c07b'
10const FOLD_OVER = 20
11const FOLD_TO = 12
12const NUMBER_OVER = 10
13
14const CALLOUT: Record<CalloutType, { color: string; title: string }> = {
15  note: { color: 'suggestion', title: 'ℹ Note' },
16  tip: { color: 'success', title: '✓ Tip' },
17  important: { color: 'merged', title: '◆ Important' },
18  warning: { color: 'warning', title: '⚠ Warning' },
19  caution: { color: 'error', title: '✕ Caution' },
20}
21
22// Code blocks the person unfolded, by message id and block index.
23const expanded = atom({ plugin: 'pretty-markdown', key: 'expanded' } as const, [] as readonly FoldId[])
24
25const widthOf = (list: Span[]) => list.reduce((n, s) => n + spanWidth(s), 0)
26
27type Indexed = Part & { i: number }
28
29// Pairs each heading with the part under it, so the two draw without a gap.
30function sections(parts: Part[]): [Indexed | undefined, Indexed | undefined][] {
31  const out: [Indexed | undefined, Indexed | undefined][] = []
32  const list = parts.map((p, i) => ({ ...p, i }) as Indexed)
33  for (let k = 0; k < list.length; k++) {
34    const p = list[k]!
35    if (p.kind !== 'heading') out.push([undefined, p])
36    else if (list[k + 1] && list[k + 1]!.kind !== 'heading') out.push([p, list[++k]])
37    else out.push([p, undefined])
38  }
39  return out
40}
41
42export const register: Register = on => {
43  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
44    const parts = parse(e.props.text)
45    const isPlain = parts.every(p => p.kind === 'markdown')
46    if (isPlain) return next(e)
47
48    const { Box, Text, Code, Markdown, Link, Button } = $.ui.resolve(e)
49    const open = await read($, expanded)
50    const room = (e.viewport?.columns ?? 100) - 4
51
52    const spans = (list: Span[], at: string) =>
53      list.map((s, k) =>
54        s.kind === 'code' ? (
55          <Text key={`${at}.${k}`} backgroundColor={CHIP_BG} color={CHIP_FG}>{` ${s.text} `}</Text>
56        ) : s.kind === 'link' ? (
57          <Link key={`${at}.${k}`} href={s.href} label={s.label} />
58        ) : (
59          <Text key={`${at}.${k}`} weight={s.bold ? 'bold' : undefined} italic={s.italic}>{s.text}</Text>
60        ),
61      )
62
63    const lines = (list: Line[], at: string) =>
64      list.map((line, j) => (
65        <Text key={`${at}.${j}`} weight={line.isHeading ? 'bold' : undefined}>
66          {line.prefix}
67          {spans(line.spans, `${at}.${j}`)}
68        </Text>
69      ))
70
71    const draw = (part: Indexed) => {
72      const i = part.i
73
74      if (part.kind === 'markdown') return <Markdown key={`m${i}`} text={part.text} />
75
76      if (part.kind === 'heading') {
77        return (
78          <Box key={`h${i}`} flexDirection="column">
79            <Text color="claude" weight="bold">{spans(part.spans, `h${i}`)}</Text>
80            <Text dimColor>{'─'.repeat(Math.min(room, widthOf(part.spans)))}</Text>
81          </Box>
82        )
83      }
84
85      if (part.kind === 'table') {
86        const widths = part.header.map((h, k) =>
87          Math.max(widthOf(h), ...part.rows.map(r => widthOf(r[k] ?? []))))
88        // Too wide: take the overflow from the widest columns, whose cells then
89        // wrap; the engine's own table is the fallback only when that can't fit.
90        const MIN_COL = 12
91        let over = widths.reduce((a, b) => a + b, 0) + 3 * (widths.length - 1) + 4 - room
92        while (over > 0) {
93          const k = widths.indexOf(Math.max(...widths))
94          if (widths[k]! <= MIN_COL) break
95          const cut = Math.min(over, widths[k]! - MIN_COL)
96          widths[k] = widths[k]! - cut
97          over -= cut
98        }
99        if (over > 0) return <Markdown key={`t${i}`} text={part.source} />
100        // Drawn line by line, so every wrapped line carries its column rules.
101        const rule = (l: string, m: string, r: string, at: string) => (
102          <Text key={at} color="subtle">{`${l}─${widths.map(w => '─'.repeat(w)).join(`─${m}─`)}─${r}`}</Text>
103        )
104        const row = (cells: Span[][], at: string, isHead: boolean) => {
105          const wrapped = cells.map((cell, k) => wrapSpans(cell, widths[k]!))
106          const height = Math.max(...wrapped.map(w => w.length))
107          return Array.from({ length: height }, (_, j) => (
108            <Text key={`${at}.${j}`} color={isHead ? 'claude' : undefined} weight={isHead ? 'bold' : undefined}>
109              {wrapped.map((lines, k) => {
110                const line = lines[j] ?? []
111                const room = widths[k]! - line.reduce((n, s) => n + spanWidth(s), 0)
112                const align = part.align[k] ?? 'left'
113                const before = align === 'right' ? room : align === 'center' ? Math.floor(room / 2) : 0
114                return (
115                  <Text key={`${at}.${j}.${k}`}>
116                    <Text color="subtle">{k === 0 ? '│ ' : ' │ '}</Text>
117                    {' '.repeat(before)}
118                    {spans(line, `${at}.${j}.${k}`)}
119                    {' '.repeat(room - before)}
120                  </Text>
121                )
122              })}
123              <Text color="subtle">{' │'}</Text>
124            </Text>
125          ))
126        }
127        return (
128          <Box key={`t${i}`} flexDirection="column" flexShrink={0}>
129            {rule('╭', '┬', '╮', `t${i}.top`)}
130            {row(part.header, `t${i}.h`, true)}
131            {part.rows.flatMap((r, j) => [rule('├', '┼', '┤', `t${i}.r${j}`), ...row(r, `t${i}.${j}`, false)])}
132            {rule('╰', '┴', '╯', `t${i}.end`)}
133          </Box>
134        )
135      }
136
137      if (part.kind === 'next') {
138        return (
139          <Box key={`n${i}`} borderStyle="round" borderColor="claude" paddingX={1}>
140            <Text>
141              <Text color="claude" weight="bold">{'→ Next  '}</Text>
142              {spans(part.spans, `n${i}`)}
143            </Text>
144          </Box>
145        )
146      }
147
148      if (part.kind === 'callout') {
149        const style = CALLOUT[part.type]
150        return (
151          <Box key={`a${i}`} flexDirection="column" borderStyle="round" borderColor={style.color} paddingX={1}>
152            <Text color={style.color} weight="bold">{style.title}</Text>
153            {lines(part.lines, `a${i}`)}
154          </Box>
155        )
156      }
157
158      if (part.kind === 'code') {
159        const isDiff = part.language === 'diff' || part.language === 'patch'
160        const all = part.source.split('\n')
161        const foldId = `${e.requestId}:${i}`
162        const isFolded = all.length > FOLD_OVER && !open.includes(foldId)
163        const shown = isFolded ? all.slice(0, FOLD_TO).join('\n') : part.source
164        return (
165          <Box key={`c${i}`} flexDirection="column" borderStyle="round" borderColor="subtle" paddingX={1}>
166            <Box flexDirection="row" justifyContent="space-between">
167              <Text dimColor italic>{part.language}</Text>
168              <Button
169                key={`copy${i}`}
170                label="⧉ copy"
171                plain
172                dimColor
173                onPress={async press => {
174                  const { isCopied } = await $.ui.copy({ text: part.source, surface: press.surface })
175                  $.ui.toast(isCopied ? 'Copied to clipboard' : 'Copy failed')
176                }}
177              />
178            </Box>
179            {isDiff ? (
180              <Code source={asHunk(part.source)} format="diff" />
181            ) : (
182              <Code
183                source={shown}
184                language={part.language || undefined}
185                startLine={all.length > NUMBER_OVER ? 1 : undefined}
186              />
187            )}
188            {all.length > FOLD_OVER && !isDiff ? (
189              <Button
190                key={`fold${i}`}
191                label={isFolded ? `▾ show all ${all.length} lines` : '▴ fold'}
192                plain
193                dimColor
194                onPress={() =>
195                  update($, expanded, ids =>
196                    ids.includes(foldId) ? ids.filter(id => id !== foldId) : [...ids, foldId],
197                  )
198                }
199              />
200            ) : null}
201          </Box>
202        )
203      }
204
205      return (
206        <Box key={`l${i}`} flexDirection="column">
207          {lines(part.lines, `l${i}`)}
208        </Box>
209      )
210    }
211
212    // A tree of our own replaces the engine's row, bullet and indent included.
213    return (
214      <Box flexDirection="row">
215      <Box width={2} flexShrink={0}>
216        <Text>{e.props.isFirstOfReply ? '●' : ' '}</Text>
217      </Box>
218      <Box flexDirection="column" gap={1} flexGrow={1} flexShrink={1}>
219        {sections(parts).map(([head, body]) =>
220          head === undefined ? draw(body!) : (
221            <Box key={`s${head.i}`} flexDirection="column">
222              {draw(head)}
223              {body ? draw(body) : null}
224            </Box>
225          ),
226        )}
227      </Box>
228      </Box>
229    )
230  })
231}
232

1export type Part =
2  | { kind: 'code'; language: string; source: string }
3  | { kind: 'markdown'; text: string }
4  | { kind: 'lines'; lines: Line[] }
5  | { kind: 'callout'; type: CalloutType; lines: Line[] }
6  | { kind: 'next'; spans: Span[] }
7  | { kind: 'heading'; spans: Span[] }
8  | { kind: 'table'; header: Span[][]; rows: Span[][][]; align: Align[]; source: string }
9
10export type Align = 'left' | 'center' | 'right'
11
12export type CalloutType = 'note' | 'tip' | 'important' | 'warning' | 'caution'
13
14export type Line = { prefix: string; isHeading: boolean; spans: Span[] }
15
16export type Span =
17  | { kind: 'text'; text: string; bold?: boolean; italic?: boolean }
18  | { kind: 'code'; text: string }
19  | { kind: 'link'; label: string; href: string }
20
21const FENCE = /^ {0,3}(`{3,}|~{3,})\s*([\w+#.-]*)/
22
23// Splits a reply into code fences and prose; prose blocks holding inline code
24// become our own lines, every other block stays with the engine's Markdown.
25export function parse(text: string): Part[] {
26  const parts: Part[] = []
27  const prose: string[] = []
28  const lines = text.split('\n')
29
30  const flushProse = () => {
31    if (prose.length) parts.push(...parseProse(prose.join('\n')))
32    prose.length = 0
33  }
34
35  for (let i = 0; i < lines.length; i++) {
36    const open = FENCE.exec(lines[i]!)
37    if (!open) {
38      prose.push(lines[i]!)
39      continue
40    }
41    const marker = open[1]!
42    const body: string[] = []
43    let j = i + 1
44    while (j < lines.length && !lines[j]!.trimStart().startsWith(marker)) body.push(lines[j++]!)
45    flushProse()
46    parts.push({ kind: 'code', language: open[2] ?? '', source: body.join('\n') })
47    i = j
48  }
49  flushProse()
50  return merge(parts)
51}
52
53const HEADING = /^\s*(?:#{1,6}\s+(.+?)\s*#*|\*\*([^*\n]+)\*\*)\s*$/
54const TABLE_ROW = /^\s*\|.*\|\s*$/
55const TABLE_RULE = /^\s*\|?(\s*:?-{2,}:?\s*\|)+\s*(:?-{2,}:?\s*)?\|?\s*$/
56
57const cells = (row: string) => row.trim().replace(/^\||\|$/g, '').split(/(?<!\\)\|/).map(c => c.trim())
58
59function parseTable(blockLines: string[]): Part | null {
60  if (blockLines.length < 2 || !blockLines.every(l => TABLE_ROW.test(l)) || !TABLE_RULE.test(blockLines[1]!)) return null
61  const align = cells(blockLines[1]!).map((c): Align =>
62    c.startsWith(':') && c.endsWith(':') ? 'center' : c.endsWith(':') ? 'right' : 'left')
63  const header = cells(blockLines[0]!)
64  const rows = blockLines.slice(2).map(r => {
65    const row = cells(r)
66    return header.map((_, k) => parseInline(row[k] ?? ''))
67  })
68  return { kind: 'table', header: header.map(h => parseInline(h)), rows, align, source: blockLines.join('\n') }
69}
70
71function parseProse(text: string): Part[] {
72  return text.split(/\n\s*\n/).map(b => b.replace(/^\n+|\n+$/g, '')).filter(b => b.trim()).flatMap(parseBlock)
73}
74
75// A heading line opens a block: it becomes its own part, and the lines under it
76// are read again as a block of their own.
77function parseBlock(block: string): Part[] {
78  const blockLines = block.split('\n')
79  const heading = HEADING.exec(blockLines[0]!)
80  if (heading) {
81    const rest = blockLines.slice(1).join('\n')
82    return [{ kind: 'heading', spans: parseInline(heading[1] ?? heading[2]!) }, ...(rest.trim() ? parseBlock(rest) : [])]
83  }
84  const table = parseTable(blockLines)
85  if (table) return [table]
86  return [parseOther(block, blockLines)]
87}
88
89function parseOther(block: string, blockLines: string[]): Part {
90    const callout = /^\s*>\s*\[!(NOTE|TIP|IMPORTANT|WARNING|CAUTION)\]\s*$/i.exec(blockLines[0]!)
91    if (callout && blockLines.every(l => /^\s*>/.test(l))) {
92      const body = blockLines.slice(1).map(l => l.replace(/^\s*>\s?/, ''))
93      return { kind: 'callout', type: callout[1]!.toLowerCase() as CalloutType, lines: body.map(parseLine) }
94    }
95
96    const nextLine = /^\s*(?:\*\*)?Next:(?:\*\*)?\s*(.*)$/s.exec(block)
97    if (nextLine) return { kind: 'next', spans: parseInline(nextLine[1]!.replace(/\s*\n\s*/g, ' ')) }
98
99    // Tables, quotes and indented blocks are the engine's to draw.
100    const isSimple = block.includes('`') &&
101      blockLines.every(l => !/^\s*(\||>)/.test(l) && !/^ {4,}\S/.test(l) && !/^\s*([-*_])\s*\1\s*\1/.test(l))
102    return isSimple
103      ? { kind: 'lines', lines: blockLines.map(parseLine) }
104      : { kind: 'markdown', text: block }
105}
106
107// Joins neighbouring blocks of one kind so spacing stays even.
108function merge(parts: Part[]): Part[] {
109  const out: Part[] = []
110  for (const p of parts) {
111    const last = out[out.length - 1]
112    if (p.kind === 'markdown' && last?.kind === 'markdown') last.text += '\n\n' + p.text
113    else out.push(p)
114  }
115  return out
116}
117
118function parseLine(line: string): Line {
119  const heading = /^\s*#{1,6}\s+(.*)$/.exec(line)
120  if (heading) return { prefix: '', isHeading: true, spans: parseInline(heading[1]!) }
121  const item = /^(\s*)([-*+]|\d+[.)])\s+(.*)$/.exec(line)
122  if (item) {
123    const marker = /\d/.test(item[2]!) ? item[2]! : '•'
124    return { prefix: `${item[1]}${marker} `, isHeading: false, spans: parseInline(item[3]!) }
125  }
126  return { prefix: '', isHeading: false, spans: parseInline(line) }
127}
128
129const INLINE = /(`+)([^`]|[^`].*?[^`])\1(?!`)|\*\*(.+?)\*\*|\[([^\]]+)\]\((https?:[^)\s]+|file:[^)\s]+)\)|(?<![\w*])\*(?!\s)(.+?)(?<!\s)\*(?![\w*])|(?<!\w)_(?!\s)(.+?)(?<!\s)_(?!\w)/g
130
131export function parseInline(text: string): Span[] {
132  const spans: Span[] = []
133  let at = 0
134  for (const m of text.matchAll(INLINE)) {
135    if (m.index! > at) spans.push({ kind: 'text', text: text.slice(at, m.index) })
136    if (m[1]) spans.push({ kind: 'code', text: m[2]!.trim() || m[2]! })
137    else if (m[3]) spans.push({ kind: 'text', text: m[3], bold: true })
138    else if (m[4]) spans.push({ kind: 'link', label: m[4], href: m[5]! })
139    else spans.push({ kind: 'text', text: m[6] ?? m[7] ?? '', italic: true })
140    at = m.index! + m[0].length
141  }
142  if (at < text.length) spans.push({ kind: 'text', text: text.slice(at) })
143  return spans
144}
145
146// shortcut: counts one cell per character, so wide glyphs (CJK, emoji) misalign
147// a column; measure display width if tables with them show up.
148export const spanWidth = (s: Span) => (s.kind === 'code' ? s.text.length + 2 : s.kind === 'link' ? s.label.length : s.text.length)
149
150// Breaks spans into lines at most `width` cells wide, at spaces where it can and
151// mid-word where a word is longer than a line. A link too long for a line loses
152// its target and wraps as text.
153export function wrapSpans(spans: Span[], width: number): Span[][] {
154  const words: { span: Span; isSpaced: boolean }[] = []
155  let isSpaced = false
156  for (const s of spans) {
157    if (s.kind !== 'text') {
158      words.push({ span: s, isSpaced })
159      isSpaced = false
160      continue
161    }
162    for (const bit of s.text.split(/(\s+)/)) {
163      if (/^\s+$/.test(bit)) isSpaced = true
164      else if (bit) {
165        words.push({ span: { ...s, text: bit }, isSpaced })
166        isSpaced = false
167      }
168    }
169  }
170
171  const chunks = (s: Span): Span[] => {
172    if (spanWidth(s) <= width) return [s]
173    const plain = s.kind === 'link' ? { kind: 'text' as const, text: s.label } : s
174    const step = Math.max(1, plain.kind === 'code' ? width - 2 : width)
175    const out: Span[] = []
176    for (let at = 0; at < plain.text.length; at += step) out.push({ ...plain, text: plain.text.slice(at, at + step) })
177    return out
178  }
179
180  const lines: Span[][] = [[]]
181  let used = 0
182  for (const word of words) {
183    chunks(word.span).forEach((piece, n) => {
184      const w = spanWidth(piece)
185      const gap = n === 0 && word.isSpaced && used > 0 ? 1 : 0
186      if (used > 0 && used + gap + w > width) {
187        lines.push([])
188        used = 0
189      } else if (gap) {
190        lines[lines.length - 1]!.push({ kind: 'text', text: ' ' })
191        used += 1
192      }
193      lines[lines.length - 1]!.push(piece)
194      used += w
195    })
196  }
197  return lines
198}
199
200// A diff fence written by hand often has no `@@` header, which Code's diff
201// format needs: add one counted from the lines, and mark bare lines as context.
202export function asHunk(source: string): string {
203  if (/^@@ /m.test(source)) return source
204  const lines = source.split('\n').filter(l => !/^(---|\+\+\+) /.test(l))
205    .map(l => (/^[+\- ]/.test(l) ? l : ` ${l}`))
206  const before = lines.filter(l => !l.startsWith('+')).length
207  const after = lines.filter(l => !l.startsWith('-')).length
208  return [`@@ -1,${before} +1,${after} @@`, ...lines].join('\n')
209}
210

1/** A code block's address: the message id, a colon, the block's index. */
2export type FoldId = string
3
4declare module 'claude-code' {
5  interface PluginState {
6    'pretty-markdown': { expanded: readonly FoldId[] }
7  }
8}
9