image-preview

1import type { EngineInterface, Register } from 'claude-code'
2
3type Thumb = { base64: string; columns: number; rows: number }
4
5const TOKEN = /\[Image #(\d+)\]/g
6const HAS_TOKEN = /\[Image #\d+\]/
7const SOURCE = /\[Image: source: ([^\]\n]+)\]/g
8const ALPHABET = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'
9
10// The card is at most this many cells. A cell is about twice as tall as it is wide.
11const MAX_COLUMNS = 40
12const MAX_ROWS = 14
13const BACKGROUND = '#373737'
14
15// Image number to the file the engine pasted it from.
16const paths = new Map<string, string>()
17const thumbs = new Map<string, Thumb | null>()
18const running = new Map<string, Promise<Thumb | null>>()
19const gaveUp = new Set<string>()
20let pending: string[] = []
21let hydrating: Promise<void> | null = null
22
23// The row that names a file is stored after the prompt row, so a first draw can come too early.
24const ATTEMPTS = 8
25const RETRY_MS = 150
26
27function hash(text: string): string {
28  let h = 0x811c9dc5
29
30  for (let n = 0; n < text.length; n++) {
31    h = Math.imul(h ^ text.charCodeAt(n), 0x01000193)
32  }
33
34  return (h >>> 0).toString(16) + text.length.toString(16)
35}
36
37// A prompt row names the numbers; the row after it names the files, in the same order.
38function note(text: string, into: Map<string, string>, waiting: string[]): string[] {
39  const numbers = [...new Set([...text.matchAll(TOKEN)].map(match => match[1]))]
40  const sources = [...text.matchAll(SOURCE)].map(match => match[1])
41  const names = numbers.length > 0 ? numbers : waiting
42
43  if (sources.length === 0) {
44    return names
45  }
46
47  sources.forEach((file, n) => {
48    if (names[n] !== undefined && !into.has(names[n])) {
49      into.set(names[n], file)
50    }
51  })
52
53  return []
54}
55
56async function hydrate($: EngineInterface) {
57  const found = new Map<string, string>()
58  let waiting: string[] = []
59
60  for (const message of await $.session.messages()) {
61    if (message.role === 'user') {
62      waiting = note(message.text, found, waiting)
63    }
64  }
65
66  for (const [number, file] of found) {
67    if (!paths.has(number)) {
68      paths.set(number, file)
69    }
70  }
71}
72
73// Width and height sit at bytes 16 to 23 of a PNG: the first 32 base64 characters hold them.
74function pngSize(base64: string): { width: number; height: number } | null {
75  const bytes: number[] = []
76  const head = base64.slice(0, 32)
77
78  for (let n = 0; n < head.length; n += 4) {
79    const a = ALPHABET.indexOf(head[n])
80    const b = ALPHABET.indexOf(head[n + 1])
81    const c = ALPHABET.indexOf(head[n + 2])
82    const d = ALPHABET.indexOf(head[n + 3])
83    bytes.push((a << 2) | (b >> 4), ((b & 15) << 4) | (c >> 2), ((c & 3) << 6) | d)
84  }
85
86  const read = (at: number) => ((bytes[at] << 24) | (bytes[at + 1] << 16) | (bytes[at + 2] << 8) | bytes[at + 3]) >>> 0
87  const width = read(16)
88  const height = read(20)
89
90  return width > 0 && height > 0 ? { width, height } : null
91}
92
93async function makeThumb($: EngineInterface, file: string): Promise<Thumb | null> {
94  const output = `/tmp/image-preview-${hash(file)}.png`
95  const run = await $.process.run(['sips', '-s', 'format', 'png', '-Z', '640', file, '--out', output], {
96    timeoutMs: 15_000,
97  })
98
99  if (run.exitCode !== 0) {
100    return null
101  }
102
103  const { base64 } = await $.fs.read(output, { as: 'bytes' })
104  const size = pngSize(base64)
105
106  if (!size) {
107    return null
108  }
109
110  const aspect = size.height / size.width
111  let columns = MAX_COLUMNS
112  let rows = Math.round((columns * aspect) / 2)
113
114  if (rows > MAX_ROWS) {
115    rows = MAX_ROWS
116    columns = Math.round((rows * 2) / aspect)
117  }
118
119  return { base64, columns: Math.max(1, columns), rows: Math.max(1, rows) }
120}
121
122async function thumbFor($: EngineInterface, number: string): Promise<Thumb | null> {
123  const file = paths.get(number)
124
125  if (!file) {
126    return null
127  }
128
129  if (thumbs.has(file)) {
130    return thumbs.get(file) ?? null
131  }
132
133  let job = running.get(file)
134
135  if (!job) {
136    job = makeThumb($, file).catch(() => null)
137    running.set(file, job)
138  }
139
140  const thumb = await job
141  thumbs.set(file, thumb)
142  running.delete(file)
143
144  return thumb
145}
146
147export const register: Register = on => {
148  // Drawing is pure, so the files are noted when a row is stored.
149  on('session.append', async ($, e, next) => {
150    const result = await next(e)
151
152    if (e.message.type === 'user') {
153      for (const block of e.message.content) {
154        if (block.type === 'text' && typeof block.text === 'string') {
155          pending = note(block.text, paths, pending)
156        }
157      }
158    }
159
160    return result
161  })
162
163  on('ui.render', { component: 'UserMessage' }, async ($, e, next) => {
164    if (e.surface !== 'terminal' || e.props.origin.kind !== 'composer' || !HAS_TOKEN.test(e.props.text)) {
165      return next(e)
166    }
167
168    const numbers = [...new Set([...e.props.text.matchAll(TOKEN)].map(match => match[1]))]
169
170    for (let attempt = 0; attempt < ATTEMPTS; attempt++) {
171      const missing = numbers.filter(number => !paths.has(number) && !gaveUp.has(number))
172
173      if (missing.length === 0) {
174        break
175      }
176
177      hydrating ??= hydrate($)
178        .catch(() => undefined)
179        .finally(() => {
180          hydrating = null
181        })
182      await hydrating
183
184      if (missing.every(number => paths.has(number))) {
185        break
186      }
187
188      if (attempt === ATTEMPTS - 1) {
189        missing.forEach(number => gaveUp.add(number))
190      } else {
191        await $.clock.sleep(RETRY_MS)
192      }
193    }
194
195    const drawn = new Map<string, Thumb | null>()
196    await Promise.all(numbers.map(async number => drawn.set(number, await thumbFor($, number))))
197
198    if ([...drawn.values()].every(thumb => thumb === null)) {
199      return next(e)
200    }
201
202    const { Box, Image, Text } = $.ui.resolve(e)
203
204    return (
205      <Box backgroundColor={BACKGROUND} flexDirection="row">
206        <Text dimColor>{'❯ '}</Text>
207        <Box flexDirection="column" flexGrow={1}>
208          {e.props.text.split('\n').map((line, i) => (
209            <Box key={`l${i}`} flexDirection="row" flexWrap="wrap">
210              {line.split(/(\[Image #\d+\])/).map((part, j) => {
211                if (part === '') {
212                  return null
213                }
214
215                const number = /^\[Image #(\d+)\]$/.exec(part)?.[1]
216                const thumb = number === undefined ? null : drawn.get(number)
217
218                if (!thumb) {
219                  return <Text key={`t${i}-${j}`}>{part}</Text>
220                }
221
222                return (
223                  <Box key={`k${i}-${j}`}>
224                    <Text color="blue" hover={{ underline: true }}>
225                      {part}
226                    </Text>
227                    <Box
228                      position="absolute"
229                      top={-(thumb.rows + 2)}
230                      left={0}
231                      display="none"
232                      hover={{ display: 'flex' }}
233                      borderStyle="round"
234                      borderColor="gray"
235                    >
236                      <Image
237                        key={`p${i}-${j}`}
238                        source={{ png: thumb.base64 }}
239                        columns={thumb.columns}
240                        rows={thumb.rows}
241                        alt={part}
242                      />
243                    </Box>
244                  </Box>
245                )
246              })}
247            </Box>
248          ))}
249        </Box>
250      </Box>
251    )
252  })
253}
254