reply-polish

1import type { Register, RenderElement } from 'claude-code'
2
3import { displayWidth, layout, polish, reflow } from './format'
4import type { Segment } from './format'
5
6// Layout rules, in the order of the four design principles.
7//
8// Contrast:    headings are bold, the first two levels in the accent colour; table borders are dim,
9//              the table header takes the accent colour; body text stays plain.
10// Repetition:  one accent colour, one heading style per level, one table style, one reply marker.
11// Alignment:   one left edge for everything: text, headings and tables start at the same cell, and a
12//              table is never wider than the text column. The column itself is centred on the screen.
13// Proximity:   a heading sits right above what it introduces and a blank line apart from what came
14//              before; other blocks are one blank line apart.
15const ACCENT = 'cyan'
16// Cells across the text column: about 40 Chinese characters or 80 Latin ones, a comfortable measure.
17const MEASURE = 80
18// The column takes at most this share of the window.
19const FILL = 0.72
20// A nested list item moves in by this many cells.
21const INDENT = 2
22// One Markdown block is at most this long, or the whole tree is refused.
23const MAX_BLOCK = 10000
24// The reply's own marker (the bullet) takes this many cells at the left.
25const MARKER = 2
26
27const text = (props: Record<string, unknown>, content: string): RenderElement =>
28  ({ type: 'Text', props: { wrap: 'truncate-end', ...props }, children: [content] }) as RenderElement
29
30const box = (props: Record<string, unknown>, children: RenderElement[]): RenderElement =>
31  ({ type: 'Box', props, children }) as RenderElement
32
33export const register: Register = on => {
34  // A rewrite changes the drawing only; the stored reply, and ctrl+o, keep the original.
35  on('ui.render', { component: 'AssistantMessage' }, ($, e, next) => {
36    const columns = e.viewport?.columns
37    // At most MEASURE cells, and a share of a narrow window, so the margins on both sides stay visible.
38    const measure = columns === undefined ? undefined : Math.max(30, Math.min(MEASURE, Math.floor(columns * FILL)))
39    const segments = measure === undefined ? [] : layout(e.props.text, measure)
40    const isDrawable = segments.every(segment => segment.kind !== 'text' || segment.text.length <= MAX_BLOCK) &&
41      segments.every(segment => segment.kind !== 'list' || segment.items.every(item => item.text.length <= MAX_BLOCK))
42
43    if (columns === undefined || measure === undefined || !isDrawable) {
44      const rewritten = polish(e.props.text)
45      return next(rewritten === e.props.text ? e : { ...e, props: { ...e.props, text: rewritten } })
46    }
47
48    const left = Math.max(0, Math.floor((columns - MARKER - measure) / 2))
49
50    const draw = (segment: Segment): RenderElement => {
51      if (segment.kind === 'text') {
52        return { type: 'Markdown', props: { text: reflow(segment.text, measure) } } as RenderElement
53      }
54
55      if (segment.kind === 'list') {
56        // Markers share one column so the item texts line up; each text hangs after its marker.
57        const markerWidth = Math.max(...segment.items.map(item => displayWidth(item.marker))) + 1
58
59        return box(
60          { flexDirection: 'column' },
61          segment.items.map(item => {
62            const indent = INDENT * item.depth
63            const itemWidth = measure - indent - markerWidth
64
65            return box({ flexDirection: 'row', marginLeft: indent }, [
66              text({ color: ACCENT }, item.marker.padEnd(markerWidth)),
67              box({ width: itemWidth }, [
68                { type: 'Markdown', props: { text: reflow(item.text, itemWidth) } } as RenderElement,
69              ]),
70            ])
71          }),
72        )
73      }
74
75      if (segment.kind === 'heading') {
76        const isTop = segment.level <= 2
77        const title = text({ bold: true, ...(isTop ? { color: ACCENT } : {}) }, segment.text)
78        const rule = text({ dimColor: true }, '─'.repeat(Math.min(displayWidth(segment.text), measure)))
79
80        return box({ flexDirection: 'column' }, segment.level === 1 ? [title, rule] : [title])
81      }
82
83      return box(
84        { flexDirection: 'column' },
85        segment.cells.map((row, k) => {
86          if (row === null) return text({ dimColor: true }, segment.lines[k] ?? '')
87
88          const isHeader = k === 1
89          const sources = segment.sources[k] ?? []
90          // A body cell is drawn as markdown in a box of the column's width, so code spans and bold keep
91          // their style; the header is plain text in the accent colour.
92          const draw = (cell: string, c: number): RenderElement => {
93            const source = sources[c] ?? ''
94            if (isHeader) return text({ bold: true, color: ACCENT }, cell)
95            if (source === '') return text({}, cell)
96            return box({ width: segment.widths[c] ?? 0 }, [
97              { type: 'Markdown', props: { text: source } } as RenderElement,
98            ])
99          }
100          const pieces = row.flatMap((cell, c) => [text({ dimColor: true }, c === 0 ? '│ ' : ' │ '), draw(cell, c)])
101
102          return box({ flexDirection: 'row' }, [...pieces, text({ dimColor: true }, ' │')])
103        }),
104      )
105    }
106
107    const children = segments.map((segment, i) => {
108      const isFirst = i === 0 && e.props.isFirstOfReply
109      const isTight = (i === 0 && !isFirst) || (i > 0 && segments[i - 1]?.kind === 'heading')
110      // The marker of a reply is kept: it sits in the gutter just left of the text column.
111      const hasMarker = isFirst && left >= MARKER
112      const block = box({ width: measure }, [draw(segment)])
113
114      // A reply starts one blank line below whatever came before it, a hook's notice for one.
115      return box(
116        { flexDirection: 'row', marginLeft: hasMarker ? left - MARKER : left, marginTop: isTight ? 0 : 1 },
117        hasMarker ? [text({}, '● '), block] : [block],
118      )
119    })
120
121    return box({ flexDirection: 'column' }, children)
122  })
123}
124

1// Pure text rewriting for an assistant reply. The result is still markdown: the engine's own
2// renderer draws it, so themes, links and code highlighting are kept.
3
4export const MAX_CODE_LINES = 30
5export const KEEP_CODE_LINES = 12
6// Without a measured width, a table stays a table when its widest row fits and it has few columns.
7export const WIDE_TABLE = 72
8export const MAX_COLUMNS = 3
9// A table segment carries what `drawTable` returns, line by line (see there).
10export type ListItem = { depth: number; marker: string; text: string }
11
12export type Segment =
13  | { kind: 'text'; text: string }
14  | { kind: 'heading'; level: number; text: string }
15  | { kind: 'list'; items: ListItem[] }
16  | ({ kind: 'table' } & Drawn)
17
18const OPEN_FENCE = /^\s*(`{3,}|~{3,})/
19const LIST_ITEM = /^(\s*)([-*+]|\d{1,3}[.)])\s+(.*)$/
20const BULLETS = ['•', '◦', '▪']
21// Characters that may not start a line (closing punctuation) or end one (opening punctuation).
22const NO_START = '，。、；：！？）】》」』”’…%,.;:!?)]}'
23const NO_END = '（【《「『“‘([{'
24const HEADING = /^(#{1,6})\s+(.*?)\s*#*\s*$/
25const SEPARATOR = /^\s*\|?\s*:?-+:?\s*(\|\s*:?-+:?\s*)*\|?\s*$/
26
27const PIPE = '\u0000'
28
29const cells = (line: string): string[] =>
30  line
31    .trim()
32    .replace(/\\\|/g, PIPE)
33    .replace(/^\|/, '')
34    .replace(/\|$/, '')
35    .split('|')
36    .map(cell => cell.split(PIPE).join('|').trim())
37
38const closesFence = (line: string, marker: string): boolean => {
39  const text = line.trim()
40  const mark = text[0] ?? ''
41  return text.length >= marker.length && mark === marker[0] && text === mark.repeat(text.length)
42}
43
44const isTableStart = (lines: string[], i: number): boolean => {
45  const line = lines[i] ?? ''
46  const next = lines[i + 1]
47  return (
48    next !== undefined &&
49    line.includes('|') &&
50    next.includes('|') &&
51    SEPARATOR.test(next) &&
52    cells(line).length >= 2 &&
53    cells(line).length === cells(next).length
54  )
55}
56
57const isWide = (cp: number): boolean =>
58  (cp >= 0x1100 && cp <= 0x115f) ||
59  (cp >= 0x2e80 && cp <= 0xa4cf) ||
60  (cp >= 0xac00 && cp <= 0xd7a3) ||
61  (cp >= 0xf900 && cp <= 0xfaff) ||
62  (cp >= 0xfe30 && cp <= 0xfe6f) ||
63  (cp >= 0xff00 && cp <= 0xff60) ||
64  (cp >= 0xffe0 && cp <= 0xffe6) ||
65  (cp >= 0x1f300 && cp <= 0x1faff)
66
67// Terminal cells a string takes: CJK and emoji are two cells wide.
68export const displayWidth = (text: string): number => {
69  let width = 0
70  for (const char of text) width += isWide(char.codePointAt(0) ?? 0) ? 2 : 1
71  return width
72}
73
74// A cell or a heading is drawn as plain text, so inline markdown is reduced to what it shows.
75const plain = (cell: string): string =>
76  cell
77    .replace(/\*\*(.+?)\*\*/g, '$1')
78    .replace(/`([^`]+)`/g, '$1')
79    .replace(/\[([^\]]+)\]\([^)]*\)/g, '$1')
80
81// A line cut to `width` cells. A Latin word or number stays with the Chinese character after it
82// ("20个"), closing punctuation never starts a line and opening punctuation never ends one.
83export const wrapLine = (line: string, width: number): string[] => {
84  if (displayWidth(line) <= width) return [line]
85
86  const pieces: string[] = []
87  let run = ''
88  for (const char of line) {
89    if (char === ' ') {
90      if (run !== '') pieces.push(run)
91      run = ''
92      pieces.push(' ')
93    } else if (isWide(char.codePointAt(0) ?? 0)) {
94      pieces.push(run + char)
95      run = ''
96    } else {
97      run += char
98    }
99  }
100  if (run !== '') pieces.push(run)
101
102  const units: string[] = []
103  for (const piece of pieces) {
104    const prev = units[units.length - 1]
105    const isGlued =
106      prev !== undefined &&
107      prev !== ' ' &&
108      piece !== ' ' &&
109      (NO_START.includes(piece.charAt(0)) || NO_END.includes(prev.charAt(prev.length - 1)))
110    if (isGlued) units[units.length - 1] = prev + piece
111    else units.push(piece)
112  }
113
114  const out: string[] = []
115  let current = ''
116  let used = 0
117  // Backticks are never shown; `**` is hidden as bold markup, but shown inside a code span.
118  let isInCode = false
119  for (const unit of units) {
120    if (unit === ' ') {
121      if (current !== '') {
122        current += ' '
123        used += 1
124      }
125      continue
126    }
127    const ticks = (unit.match(/`/g) ?? []).length
128    const shown = isInCode || ticks > 0 ? unit.replace(/`/g, '') : unit.replace(/\*\*/g, '')
129    if (ticks % 2 === 1) isInCode = !isInCode
130    const size = displayWidth(shown)
131    if (used + size > width && current.trim() !== '') {
132      out.push(current.trimEnd())
133      current = ''
134      used = 0
135    }
136    current += unit
137    used += size
138  }
139  if (current.trim() !== '') out.push(current.trimEnd())
140
141  return out
142}
143
144// Paragraph lines cut to `width`; code, quotes, tables and indented lines are left as written.
145export const reflow = (text: string, width: number): string => {
146  let isFenced = false
147
148  return text
149    .split('\n')
150    .flatMap(raw => {
151      if (OPEN_FENCE.test(raw)) {
152        isFenced = !isFenced
153        return [raw]
154      }
155      // One to three leading spaces mean nothing in markdown; four or more make a code block.
156      const line = isFenced || !/^ {1,3}\S/.test(raw) ? raw : raw.trimStart()
157      return isFenced || /^(\s|>|\||<)/.test(line) ? [line] : wrapLine(line, width)
158    })
159    .join('\n')
160}
161
162// A table with this many body rows or more gets a rule between its rows, so a long table stays readable.
163export const ROW_RULES_FROM = 7
164
165type Drawn = {
166  lines: string[]
167  // Per line: the padded plain cells of a content line, null for a rule.
168  cells: (string[] | null)[]
169  // Per line: the cells as written in markdown (code spans, bold), null for a rule.
170  sources: (string[] | null)[]
171  widths: number[]
172  width: number
173}
174
175// The table as box-drawing lines, columns sized to their widest cell. Line 1 is the header.
176export const drawTable = (header: string[], rows: string[][]): Drawn => {
177  const raw = [header, ...rows.map(row => header.map((_, k) => row[k] ?? ''))]
178  const grid = raw.map(row => row.map(plain))
179  const widths = header.map((_, k) => Math.max(...grid.map(row => displayWidth(row[k] ?? ''))))
180  const pad = (text: string, width: number) => text + ' '.repeat(width - displayWidth(text))
181  const rule = (left: string, mid: string, right: string) =>
182    left + widths.map(width => '─'.repeat(width + 2)).join(mid) + right
183  const padded = grid.map(row => row.map((cell, k) => pad(cell, widths[k] ?? 0)))
184  const [top, middle, bottom] = [rule('┌', '┬', '┐'), rule('├', '┼', '┤'), rule('└', '┴', '┘')]
185  const hasRowRules = rows.length >= ROW_RULES_FROM
186
187  type Entry = { rule: string } | { row: number }
188  const entries: Entry[] = [{ rule: top }, { row: 0 }, { rule: middle }]
189  for (let r = 1; r < grid.length; r += 1) {
190    entries.push({ row: r })
191    if (hasRowRules && r < grid.length - 1) entries.push({ rule: middle })
192  }
193  entries.push({ rule: bottom })
194
195  return {
196    lines: entries.map(entry => ('rule' in entry ? entry.rule : '│ ' + (padded[entry.row] ?? []).join(' │ ') + ' │')),
197    cells: entries.map(entry => ('rule' in entry ? null : (padded[entry.row] ?? []))),
198    sources: entries.map(entry => ('rule' in entry ? null : (raw[entry.row] ?? []))),
199    widths,
200    width: displayWidth(top),
201  }
202}
203
204// Two columns: one bullet per row, "- **first** — second".
205// More columns: a bullet per row with one sub-bullet per cell, "  - head: cell".
206// Returns null when the table is small enough to stay as it is, unless `force` is set.
207const tableItems = (header: string[], rows: string[][]): ListItem[] =>
208  rows.flatMap((row): ListItem[] => {
209    const [first = '', ...rest] = header.map((_, k) => row[k] ?? '')
210    const title = first === '' ? '' : first.includes('**') ? first : `**${first}**`
211
212    if (rest.length === 1) {
213      const only = rest[0] ?? ''
214      return [{ depth: 0, marker: '•', text: `${title}${title !== '' && only !== '' ? ' — ' : ''}${only}` }]
215    }
216
217    const parts = rest.flatMap((cell, k): ListItem[] =>
218      cell === '' ? [] : [{ depth: 1, marker: '◦', text: `${header[k + 1] ?? ''}: ${cell}` }],
219    )
220    return [{ depth: 0, marker: '•', text: title }, ...parts]
221  })
222
223const tableToList = (header: string[], rows: string[][], rawLines: string[], force = false): string[] | null => {
224  const widest = Math.max(...rawLines.map(line => line.length))
225  if (!force && widest <= WIDE_TABLE && header.length <= MAX_COLUMNS) return null
226
227  return tableItems(header, rows).map(item => `${'  '.repeat(item.depth)}- ${item.text}`)
228}
229
230// The reply as text blocks, headings and tables. With `width` (the cells a table may take, the text
231// column's) a table that fits is returned as a table to draw and one that does not becomes a list,
232// and headings are blocks of their own; without it, only a wide table becomes a list, headings
233// become bold lines and the rest stays markdown for the engine.
234export const layout = (text: string, width?: number): Segment[] => {
235  const lines = text.split('\n')
236  const out: string[] = []
237  const segments: Segment[] = []
238  const lastIsBlank = () => out.length > 0 && (out[out.length - 1] ?? '').trim() === ''
239  const flush = () => {
240    const block = out.join('\n').replace(/^\n+|\n+$/g, '')
241    if (block !== '') segments.push({ kind: 'text', text: block })
242    out.length = 0
243  }
244  let i = 0
245
246  while (i < lines.length) {
247    const line = lines[i] ?? ''
248    const fence = OPEN_FENCE.exec(line)
249
250    if (fence) {
251      const marker = fence[1] ?? '```'
252      let end = i + 1
253      while (end < lines.length && !closesFence(lines[end] ?? '', marker)) end += 1
254      const isClosed = end < lines.length
255      const body = lines.slice(i + 1, end)
256
257      out.push(line)
258      if (isClosed && body.length > MAX_CODE_LINES) {
259        out.push(...body.slice(0, KEEP_CODE_LINES))
260        out.push(`… ${body.length - KEEP_CODE_LINES} more lines hidden (ctrl+o shows the whole reply)`)
261      } else {
262        out.push(...body)
263      }
264      if (isClosed) out.push(lines[end] ?? '')
265      i = isClosed ? end + 1 : lines.length
266      continue
267    }
268
269    if (isTableStart(lines, i)) {
270      let end = i + 2
271      while (end < lines.length && (lines[end] ?? '').trim() !== '' && (lines[end] ?? '').includes('|')) end += 1
272      const raw = lines.slice(i, end)
273      const header = cells(line)
274      const rows = lines.slice(i + 2, end).map(cells)
275      const drawn = width === undefined ? null : drawTable(header, rows)
276
277      if (drawn !== null && width !== undefined && drawn.width <= width) {
278        flush()
279        segments.push({ kind: 'table', ...drawn })
280      } else if (width !== undefined) {
281        flush()
282        segments.push({ kind: 'list', items: tableItems(header, rows) })
283      } else {
284        out.push(...(tableToList(header, rows, raw) ?? raw))
285      }
286      i = end
287      continue
288    }
289
290    const heading = HEADING.exec(line)
291    if (heading) {
292      const level = (heading[1] ?? '#').length
293      const title = (heading[2] ?? '').replace(/^\*\*(.+)\*\*$/, '$1')
294
295      if (width !== undefined) {
296        flush()
297        segments.push({ kind: 'heading', level, text: plain(title) })
298      } else {
299        if (level <= 2 && out.length > 0 && !lastIsBlank()) out.push('')
300        out.push(`**${title}**`)
301      }
302      i += 1
303      continue
304    }
305
306    if (width !== undefined && LIST_ITEM.test(line)) {
307      flush()
308      const items: ListItem[] = []
309      const indents: number[] = []
310      let end = i
311
312      while (end < lines.length) {
313        const current = lines[end] ?? ''
314        const item = LIST_ITEM.exec(current)
315
316        if (item) {
317          const indent = (item[1] ?? '').length
318          while (indents.length > 0 && indent < (indents[indents.length - 1] ?? 0)) indents.pop()
319          if (indents.length === 0 || indent > (indents[indents.length - 1] ?? 0)) indents.push(indent)
320          const depth = Math.min(indents.length - 1, BULLETS.length - 1)
321          const mark = item[2] ?? '-'
322          items.push({
323            depth,
324            marker: /^\d/.test(mark) ? mark : (BULLETS[depth] ?? '•'),
325            text: (item[3] ?? '').trim(),
326          })
327        } else if (items.length > 0 && /^\s{2,}\S/.test(current) && !OPEN_FENCE.test(current)) {
328          const last = items[items.length - 1]
329          if (last !== undefined) last.text += ' ' + current.trim()
330        } else {
331          break
332        }
333        end += 1
334      }
335
336      segments.push({ kind: 'list', items })
337      i = end
338      continue
339    }
340
341    if (line.trim() === '') {
342      if (!lastIsBlank()) out.push('')
343    } else {
344      out.push(line)
345    }
346    i += 1
347  }
348
349  flush()
350  return segments.length > 0 ? segments : [{ kind: 'text', text: '' }]
351}
352
353// The reply as one markdown string, for a surface that cannot be measured.
354export const polish = (text: string): string =>
355  layout(text)
356    .map(segment => {
357      if (segment.kind === 'text') return segment.text
358      if (segment.kind === 'heading') return `**${segment.text}**`
359      if (segment.kind === 'list') return segment.items.map(item => `${item.marker} ${item.text}`).join('\n')
360      return segment.lines.join('\n')
361    })
362    .join('\n')
363