latex-render

1import type { Register, EngineInterface } from 'claude-code'
2
3// Display math in an assistant reply ($$...$$ or \[...\]) is typeset with
4// tectonic, rasterised with pdftocairo, and drawn inline as an Image over the
5// kitty graphics protocol (Ghostty, kitty). Inline $...$ is left as text.
6//
7// Inside tmux the engine's image upload does not reach the terminal unless
8// Claude Code runs under scripts/cc-tmux-bridge.py (which sets
9// CC_TMUX_BRIDGE=1); without the bridge, formulas are shown as LaTeX source.
10
11const DPI = 400 // rasterisation; above the screen's pixels per point, but every formula is held decoded by the terminal
12const PT_PER_ROW = 8 // points of typeset height per terminal row; lower = bigger
13const CELL_ASPECT = 2.1 // cell height / cell width, used when no measured size is available
14const TEXT_COLOR = 'white' // the terminal is dark; change for a light theme
15
16type Piece = { kind: 'md'; text: string } | { kind: 'tex'; src: string }
17
18// Fenced blocks and inline code spans are skipped: a `$$` quoted in prose
19// must not pair with a real delimiter.
20const FENCE = /(```[\s\S]*?```|~~~[\s\S]*?~~~|`[^`\n]*`)/
21const DISPLAY = /\$\$([\s\S]+?)\$\$|\\\[([\s\S]+?)\\\]/g
22
23function split(text: string): Piece[] {
24  const out: Piece[] = []
25  for (const chunk of text.split(FENCE)) {
26    if (!chunk) continue
27    if (chunk.startsWith('```') || chunk.startsWith('~~~') || chunk.startsWith('`')) {
28      out.push({ kind: 'md', text: chunk })
29      continue
30    }
31    let last = 0
32    for (const m of chunk.matchAll(DISPLAY)) {
33      const src = (m[1] ?? m[2] ?? '').trim()
34      if (m.index! > last) out.push({ kind: 'md', text: chunk.slice(last, m.index) })
35      out.push({ kind: 'tex', src })
36      last = m.index! + m[0].length
37    }
38    if (last < chunk.length) out.push({ kind: 'md', text: chunk.slice(last) })
39  }
40  return out
41}
42
43const hashes = new Map<string, Promise<string>>() // formula source -> cache key
44function sha(text: string): Promise<string> {
45  let p = hashes.get(text)
46  if (!p) {
47    p = crypto.subtle.digest('SHA-256', new TextEncoder().encode(text)).then(buf =>
48      [...new Uint8Array(buf)].map(b => b.toString(16).padStart(2, '0')).join('').slice(0, 16),
49    )
50    hashes.set(text, p)
51  }
52  return p
53}
54
55function pngSize(base64: string): { width: number; height: number } | null {
56  const b = Uint8Array.fromBase64(base64)
57  if (b.length < 24 || b[12] !== 0x49 || b[13] !== 0x48 || b[14] !== 0x44 || b[15] !== 0x52) return null
58  const v = new DataView(b.buffer, b.byteOffset)
59  return { width: v.getUint32(16), height: v.getUint32(20) }
60}
61
62type Rendered = { png: string; width: number; height: number }
63const done = new Map<string, Rendered | null>() // settled renders, by key
64const pending = new Map<string, string>() // key -> source, waiting for the next batch
65let batchTimer: ReturnType<typeof setTimeout> | null = null
66let batchRunning = false
67let home: Promise<string> | null = null
68
69function cacheDir($: EngineInterface): Promise<string> {
70  if (!home) home = $.env.get('HOME').then(h => `${h ?? '/tmp'}/.cache/claude-latex`)
71  return home
72}
73
74// Answers at once: the settled picture, or `undefined` while it renders. The
75// formulas of one redraw are compiled together (one tectonic run), and the
76// batch settling asks the engine to draw this plugin's sites again, once.
77function render($: EngineInterface, key: string, src: string): Rendered | null | undefined {
78  if (done.has(key)) return done.get(key)
79  if (!pending.has(key)) pending.set(key, src)
80  // Every formula of a redraw is registered within the same hook call, so the
81  // window only has to outlast the current tick.
82  if (!batchTimer && !batchRunning) batchTimer = setTimeout(() => void runBatch($), 10)
83  return undefined
84}
85
86async function runBatch($: EngineInterface): Promise<void> {
87  batchTimer = null
88  if (batchRunning || pending.size === 0) return
89  batchRunning = true
90  const items = [...pending].map(([key, src]) => ({ key, src }))
91  pending.clear()
92  try {
93    const dir = await cacheDir($)
94    const fresh: typeof items = []
95    for (const it of items) (await $.fs.exists(`${dir}/${it.key}.png`)) ? void 0 : fresh.push(it)
96    if (fresh.length > 0) {
97      try {
98        await compile($, dir, fresh)
99      } catch (err) {
100        // One bad formula fails the whole compile: retry them one at a time.
101        $.ui.log(`latex-render: batch failed (${String(err).slice(0, 200)}); retrying singly`, { to: 'debug' })
102        for (const it of fresh) {
103          await compile($, dir, [it]).catch(e => $.ui.log(`latex-render: ${it.key}: ${String(e).slice(0, 200)}`, { to: 'debug' }))
104        }
105      }
106    }
107    for (const it of items) done.set(it.key, await load($, `${dir}/${it.key}.png`))
108  } finally {
109    batchRunning = false
110    $.ui.invalidate('ui.render')
111    if (pending.size > 0) batchTimer = setTimeout(() => void runBatch($), 0)
112  }
113}
114
115// One tectonic run for several formulas: each display becomes its own tightly
116// cropped page (the preview package), rasterised to <key>.png.
117async function compile($: EngineInterface, dir: string, items: { key: string; src: string }[]): Promise<void> {
118  const id = items.length === 1 ? items[0]!.key : `batch-${items.map(i => i.key.slice(0, 4)).join('')}`
119  const tex = `${dir}/${id}.tex`
120  await $.fs.write(
121    tex,
122    [
123      '\\documentclass{article}',
124      '\\usepackage{amsmath,amssymb,amsfonts,xcolor}',
125      // Each formula is its own tight preview box around displaystyle math. (The
126      // package's displaymath extraction keeps the display's whole line, which
127      // leaves the formula centred in a canvas of transparent space.)
128      '\\usepackage[active,tightpage]{preview}',
129      '\\setlength\\PreviewBorder{3pt}',
130      '\\begin{document}',
131      ...items.map(i => `\\begin{preview}\\color{${TEXT_COLOR}}\\(\\displaystyle ${i.src} \\)\\end{preview}`),
132      '\\end{document}',
133      '',
134    ].join('\n'),
135  )
136  // Even with --only-cached, tectonic contacts the bundle server on every run
137  // and waits on it (seconds on a slow link, with almost no CPU used). Pointing
138  // its proxy at a closed local port makes that attempt fail at once, so the
139  // compile runs from the cache in well under a second. A package not cached
140  // yet makes this run fail, and the fallback compiles with the network open.
141  const offline = { https_proxy: 'http://127.0.0.1:9', http_proxy: 'http://127.0.0.1:9', HTTPS_PROXY: 'http://127.0.0.1:9', HTTP_PROXY: 'http://127.0.0.1:9' }
142  // tectonic reruns TeX when the .aux changes, and a first pass always writes
143  // one. Writing the two lines it would produce skips that second pass.
144  await $.fs.write(`${dir}/${id}.aux`, `\\relax \n\\gdef \\@abspage@last{${items.length}}\n`)
145  let tect = await $.process.run(['tectonic', '-o', dir, '--chatter', 'minimal', '--only-cached', tex], { env: offline, timeoutMs: 60000 })
146  if (tect.exitCode !== 0) {
147    tect = await $.process.run(['tectonic', '-o', dir, '--chatter', 'minimal', tex], { timeoutMs: 120000 })
148  }
149  if (tect.exitCode !== 0) throw new Error(`tectonic: ${tect.stderr.slice(0, 300)}`)
150  // One pdftocairo run for every page (it names them <prefix>-<n>.png, n
151  // zero-padded to the digits of the page count), then each page is moved to
152  // its formula's cache name; one process instead of one per formula.
153  const pdf = `${dir}/${id}.pdf`
154  const prefix = `${dir}/${id}`
155  const width = String(items.length).length
156  // Ghostty draws cairo's PNG byte stream at a fraction of its size (the same
157  // pixels re-saved by any other encoder draw right), so each page is re-saved
158  // by sips on the way to its cache name; alpha survives.
159  const moves = items
160    // A page wider than 4096 px is downsampled first (never enlarged): the
161    // engine refuses a larger Image, and with it the whole message's tree.
162    .map((it, n) => {
163      const page = `${prefix}-${String(n + 1).padStart(width, '0')}.png`
164      return `w=$(sips -g pixelWidth "${page}" | awk '/pixelWidth/{print $2}'); if [ "$w" -gt 4096 ]; then sips --resampleWidth 4096 "${page}" >/dev/null; fi; sips -s format png "${page}" --out "${dir}/${it.key}.png" >/dev/null`
165    })
166    .join(' && ')
167  const cairo = await $.process.run(
168    ['sh', '-c', `pdftocairo -png -transp -r ${DPI} "${pdf}" "${prefix}" && ${moves} && rm -f "${prefix}"-*.png`],
169    { timeoutMs: 30000 },
170  )
171  if (cairo.exitCode !== 0) throw new Error(`pdftocairo: ${cairo.stderr.slice(0, 300)}`)
172}
173
174async function load($: EngineInterface, png: string): Promise<Rendered | null> {
175  try {
176    const { base64 } = await $.fs.read(png, { as: 'bytes' })
177    const size = pngSize(base64)
178    return size ? { png: base64, ...size } : null
179  } catch {
180    return null
181  }
182}
183
184// The terminal's cell size in pixels, as cc-tmux-bridge measures it and leaves
185// in ~/.cache/cc-tmux-bridge.cellsize ("<width> <height>"); CELL_ASPECT else.
186let cellAspect: Promise<number> | null = null
187function measuredCellAspect($: EngineInterface): Promise<number> {
188  if (!cellAspect) {
189    cellAspect = (async () => {
190      try {
191        const home = (await $.env.get('HOME')) ?? ''
192        const [w, h] = (await $.fs.read(`${home}/.cache/cc-tmux-bridge.cellsize`)).trim().split(/\s+/).map(Number)
193        if (w! > 0 && h! > 0) return h! / w!
194      } catch {}
195      return CELL_ASPECT
196    })()
197  }
198  return cellAspect
199}
200
201function cells(r: Rendered, maxColumns: number, aspectRatio: number): { columns: number; rows: number } {
202  const heightPt = (r.height * 72) / DPI
203  const aspect = (r.width / r.height) * aspectRatio // columns per row
204  let rows = Math.max(1, Math.round(heightPt / PT_PER_ROW))
205  // Err a little wide: the bridge trims the box to the picture's exact shape
206  // (it can only shrink it), and a box too narrow would shrink the picture.
207  let columns = Math.max(1, Math.ceil(rows * aspect * 1.04))
208  if (columns > maxColumns) {
209    columns = maxColumns
210    rows = Math.max(1, Math.round(columns / aspect))
211  }
212  return { columns: Math.min(columns, 255), rows: Math.min(rows, 255) }
213}
214
215// The model avoids LaTeX in a terminal unless told it will be drawn, and then
216// this mod has nothing to render. One system-prompt section says how to write
217// math here.
218const GUIDANCE =
219  'This terminal typesets display math. Write any standalone equation or formula as display LaTeX ' +
220  'between `$$` and `$$` on its own lines (amsmath and amssymb are available); it is drawn as a typeset image. ' +
221  'Inline `$...$` is not rendered, so keep short inline expressions as plain Unicode text (x², ∑, √, π). ' +
222  'Never put display math inside a code block or backticks, where it stays as source.'
223
224export const register: Register = on => {
225  on('prompt.compose', async ($, e, next) => {
226    const composed = await next(e)
227    // Only the terminal draws the pictures.
228    if (!e.surfaces.includes('terminal')) return composed
229    return { sections: [...composed.sections, { id: 'latex-render:math', text: GUIDANCE, scope: 'session' }] }
230  })
231
232  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
233    if (e.surface !== 'terminal') return next(e)
234    const pieces = split(e.props.text)
235    if (!pieces.some(p => p.kind === 'tex')) return next(e)
236
237    const { Box, Text, Markdown, Image } = $.ui.resolve(e)
238    const maxColumns = Math.min(255, Math.max(10, (e.viewport?.columns ?? 80) - 4))
239    const bridged = await $.env.get('CC_TMUX_BRIDGE').then(Boolean, () => false)
240    const inTmux = !bridged && (await $.env.get('TMUX').then(Boolean, () => false))
241    const aspectRatio = await measuredCellAspect($)
242
243    const keys = await Promise.all(pieces.map(p => (p.kind === 'tex' ? sha(p.src) : '')))
244    const nodes = pieces.map((p, i) => {
245      if (p.kind === 'md') {
246        const text = p.text.replace(/^\n+|\n+$/g, '')
247        return text ? <Markdown text={text.slice(0, 10000)} /> : null
248      }
249      if (inTmux) return <Markdown text={'```latex\n' + p.src + '\n```'} />
250      const r = render($, keys[i]!, p.src)
251      if (r === undefined) return <Markdown text={'$$ ' + p.src + ' $$'} dimColor />
252      if (r === null) return <Markdown text={'```latex\n' + p.src + '\n```'} />
253      const { columns, rows } = cells(r, maxColumns, aspectRatio)
254      return (
255        <Box marginLeft={2} marginTop={1} marginBottom={1}>
256          <Image source={{ png: r.png }} columns={columns} rows={rows} alt={p.src} />
257        </Box>
258      )
259    })
260
261    return (
262      <Box flexDirection="row">
263        <Text>{e.props.isFirstOfReply ? '● ' : '  '}</Text>
264        <Box flexDirection="column" flexGrow={1}>
265          {nodes}
266        </Box>
267      </Box>
268    )
269  })
270}
271