image-peek

1import type { Elements, EngineInterface, Register } from 'claude-code'
2
3type Size = { width: number; height: number }
4type Picture = { n: number; path: string; size: Size }
5
6// Pasting an image raises no prompt.edit, so the draft is read on a timer.
7const POLL_MS = 250
8// The shape drawn when the header cannot be read: a file over $.fs.read's size cap.
9const FALLBACK: Size = { width: 16, height: 10 }
10// Each tile adds a border around the picture and a label under it.
11const CHROME = { columns: 2, rows: 3 }
12
13export const imageTags = (text: string) => [...new Set([...text.matchAll(/\[Image #(\d+)\]/g)].map(m => Number(m[1])))]
14
15// Width and height from the PNG's IHDR chunk: bytes 16-23, big-endian.
16export const pngSize = (base64: string): Size | undefined => {
17  const bytes = Uint8Array.from(atob(base64.slice(0, 32)), c => c.charCodeAt(0))
18  if (bytes.length < 24 || String.fromCharCode(...bytes.slice(1, 4)) !== 'PNG') return undefined
19  const view = new DataView(bytes.buffer)
20  return { width: view.getUint32(16), height: view.getUint32(20) }
21}
22
23// The largest box of cells within columns × rows that keeps the picture's shape; a cell is about twice as tall as wide.
24export const fit = ({ width, height }: Size, columns: number, rows: number) => {
25  const clamp = (n: number) => Math.max(1, Math.min(255, Math.round(n)))
26  const across = Math.min(columns, (rows * 2 * width) / height)
27  return { columns: clamp(across), rows: clamp((across * height) / width / 2) }
28}
29
30let dir: string | undefined
31const pictures = new Map<number, Picture>()
32let draft: number[] = []
33let polling = false
34
35// Claude Code keeps each paste as <tmp>/<project>/<session>/images/<n>.png. The project folder is named
36// after the directory the session started in, which the session's cwd may since have left: find it by the id.
37async function imagesDir($: EngineInterface) {
38  if (dir !== undefined) return dir
39  const root = (await $.env.get('CLAUDE_CODE_TMPDIR')) ?? `/tmp/claude-${(await $.process.run(['id', '-u'])).stdout.trim()}`
40  const id = await $.session.id()
41  for (const project of await $.fs.list(root).catch(() => [])) {
42    if (await $.fs.exists(`${root}/${project.name}/${id}/images`)) return (dir = `${root}/${project.name}/${id}/images`)
43  }
44  return undefined
45}
46
47async function picture($: EngineInterface, n: number): Promise<Picture | undefined> {
48  if (pictures.has(n)) return pictures.get(n)
49  const folder = await imagesDir($)
50  const path = `${folder}/${n}.png`
51  if (folder === undefined || !(await $.fs.exists(path))) return undefined
52  const size = await $.fs.read(path, { as: 'bytes' }).then(b => pngSize(b.base64), () => FALLBACK)
53  if (size === undefined) return undefined
54  return pictures.set(n, { n, path, size }).get(n)
55}
56
57async function poll($: EngineInterface) {
58  if (polling) return
59  polling = true
60  try {
61    const tags = imageTags((await $.prompt.read()).text)
62    if (tags.join() === draft.join()) return
63    // A tag can show up before its file is written; it is looked for again on the next tick.
64    const ready = (await Promise.all(tags.map(n => picture($, n)))).filter(p => p !== undefined).map(p => p.n)
65    if (ready.join() === draft.join()) return
66    draft = ready
67    $.ui.invalidate('ui.render')
68  } catch {
69    // No prompt box (a dialog holds it, or the session is headless): nothing to show.
70  } finally {
71    polling = false
72  }
73}
74
75// One row of tiles sharing `columns`, each picture at most `rows` tall.
76function tiles({ Box, Image, Text }: Elements['terminal'], list: readonly Picture[], columns: number, rows: number) {
77  const share = Math.floor((columns - (list.length - 1)) / list.length) - CHROME.columns
78  return (
79    <Box flexDirection="row" columnGap={1}>
80      {list.map(p => (
81        <Box key={`image-${p.n}`} flexDirection="column" alignItems="center" borderStyle="round" borderDimColor>
82          {/* Tiles stretch to the row's tallest; a shorter picture sits in the middle, its label stays at the bottom. */}
83          <Box flexDirection="column" flexGrow={1} justifyContent="center">
84            <Image source={{ file: p.path, format: 'png' }} {...fit(p.size, Math.max(1, share), rows)} alt={`[Image #${p.n}]`} />
85          </Box>
86          <Text dimColor>#{p.n}</Text>
87        </Box>
88      ))}
89    </Box>
90  )
91}
92
93export const register: Register = on => {
94  on('session.start', async ($, e, next) => ($.clock.every(POLL_MS, () => poll($)), next(e)))
95
96  on('session.end', async ($, e, next) => (pictures.clear(), (draft = []), (dir = undefined), next(e)))
97
98  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
99    const below = await next(e)
100    const rows = Math.min(6, e.props.maxRows - CHROME.rows - 1)
101    if (e.surface !== 'terminal' || e.props.hasSurvey || draft.length === 0 || rows < 1) return below
102    const elements = $.ui.resolve(e as typeof e & { surface: 'terminal' })
103    const list = draft.map(n => pictures.get(n)).filter(p => p !== undefined)
104    return (
105      <elements.Box flexDirection="column">
106        {tiles(elements, list, e.props.bodyColumns, rows)}
107        {below}
108      </elements.Box>
109    )
110  })
111
112  // Pictures stay under the prompt that sent them, larger than in the draft.
113  on('ui.render', { component: 'UserMessage' }, async ($, e, next) => {
114    const below = await next(e)
115    if (e.surface !== 'terminal') return below
116    const list = (await Promise.all(imageTags(e.props.text).map(n => picture($, n)))).filter(p => p !== undefined)
117    if (list.length === 0) return below
118    const elements = $.ui.resolve(e as typeof e & { surface: 'terminal' })
119    const { columns = 80, rows = 24 } = e.viewport ?? {}
120    return (
121      <elements.Box flexDirection="column">
122        {below}
123        {tiles(elements, list, Math.min(120, columns - 4), Math.min(20, Math.floor(rows / 2)))}
124      </elements.Box>
125    )
126  })
127}
128