evidence-saver

1import type { Register } from 'claude-code'
2
3import { findImages } from './images'
4
5const MAX_B64 = 8_000_000
6const KEY = 'evidence-dir'
7const NO = "Don't save screenshots"
8const LATER = 'Ask me next time'
9
10/**
11 * Evidence Saver: persists screenshots from tool results immediately, into a folder the
12 * user chooses. No default location: the first screenshot asks, the answer is remembered
13 * (`$.store`), `/evidence-dir` shows or changes it.
14 * Base64 goes to a temp text file via $.fs.write (no argv limit), python3 decodes it.
15 */
16/** Absolute folder path from user input (`~/` expanded with home), or null. */
17const resolvePath = (home: string, raw: string): string | null => {
18  const p = raw.trim().replace(/^["']|["']$/g, '')
19
20  if (p.startsWith('~/')) return `${home}${p.slice(1)}`
21
22  return p.startsWith('/') ? p.replace(/\/+$/, '') : null
23}
24
25export const register: Register = on => {
26  let counter = 0
27  let declined = false
28
29  on('session.start', async ($, e, next) => {
30    // a name already taken must not break the mod
31    await $.command
32      .register({
33        name: 'evidence-dir',
34        description: 'Show, set (/evidence-dir /absolute/path) or reset (/evidence-dir reset) where screenshots are saved',
35      })
36      .catch(() => undefined)
37
38    return next(e)
39  })
40
41  on('command.run', { command: 'evidence-dir' }, async ($, e) => {
42    const arg = String(e.args ?? '').trim()
43
44    if (arg === 'reset') {
45      await $.store.delete(KEY)
46      declined = false
47
48      return { text: 'Evidence Saver: folder cleared. It will ask again at the next screenshot.' }
49    }
50
51    if (arg) {
52      const dir = resolvePath(String((await $.env.get('HOME')) ?? ''), arg)
53
54      if (!dir) return { text: 'Evidence Saver: give an absolute path (starting with / or ~/).' }
55
56      await $.store.set(KEY, dir)
57      declined = false
58
59      return { text: `Evidence Saver: screenshots will be saved to ${dir}` }
60    }
61
62    const dir = String((await $.store.get(KEY)) ?? '')
63
64    return { text: dir ? `Evidence Saver folder: ${dir}` : 'Evidence Saver: no folder set yet. Set one with /evidence-dir /absolute/path.' }
65  })
66
67  on('tool.call', async ($, e, next) => {
68    const ran = await next(e)
69    const saved: string[] = []
70
71    try {
72      const images = findImages(ran)
73
74      if (images.length && !declined) {
75        let dir = String((await $.store.get(KEY)) ?? '')
76
77        if (!dir) {
78          const answer = await $.ui
79            .ask(
80              'Evidence Saver: a screenshot was just returned. Which folder should screenshots be saved to? Type the absolute folder path under "Other", or pick an option.',
81              [NO, LATER],
82            )
83            .catch(() => LATER)
84
85          if (answer === NO) declined = true
86          else if (answer !== LATER) {
87            const picked = resolvePath(String((await $.env.get('HOME')) ?? ''), answer)
88
89            if (picked) {
90              dir = picked
91              await $.store.set(KEY, picked)
92            } else {
93              $.ui.log('Evidence Saver: that was not an absolute path (start with / or ~/); not saved.')
94            }
95          }
96        }
97
98        if (dir) {
99          const iso = new Date(await $.clock.now()).toISOString()
100
101          for (const img of images) {
102            if (img.b64.length > MAX_B64) {
103              $.ui.log(`Evidence Saver: image too large (${img.b64.length} b64 chars), not saved`)
104              continue
105            }
106
107            const base = `${dir}/${iso.slice(0, 10)}/${iso.slice(11, 19).replaceAll(':', '')}-${String(++counter).padStart(3, '0')}-${String(e.tool).replace(/\W+/g, '_')}`
108            const file = `${base}.${img.ext}`
109
110            await $.fs.write(`${base}.b64`, img.b64)
111            await $.process.run([
112              'python3',
113              '-c',
114              'import sys,base64,os;open(sys.argv[2],"wb").write(base64.b64decode(open(sys.argv[1]).read()));os.remove(sys.argv[1])',
115              `${base}.b64`,
116              file,
117            ])
118
119            saved.push(file)
120            $.ui.log(`Evidence Saver: saved ${file}`)
121          }
122        }
123      }
124    } catch {
125      // saving evidence must never break the tool call
126    }
127
128    // Logs are not drawn in VS Code; `context` reaches the model, which relays it.
129    return saved.length && ran.deny === undefined && ran.isError !== true
130      ? { ...ran, context: [...(ran.context ?? []), `evidence-saver: screenshot saved to ${saved.join(', ')}. Tell the user the path.`] }
131      : ran
132  })
133}
134

1export interface Img {
2  ext: string
3  b64: string
4}
5
6const EXT: Record<string, string> = { 'image/png': 'png', 'image/jpeg': 'jpg', 'image/webp': 'webp', 'image/gif': 'gif' }
7
8/** Find base64 image blocks: {type:'image', source:{data,media_type}} or {data, mimeType}. */
9export function findImages(v: unknown, out: Img[] = []): Img[] {
10  if (Array.isArray(v)) v.forEach(x => findImages(x, out))
11  else if (v && typeof v === 'object') {
12    const o = v as Record<string, any>
13    const src = o.source && typeof o.source === 'object' ? o.source : o
14    const data = src.data
15    const mime = src.media_type ?? src.mimeType ?? src.mime_type
16
17    if (typeof data === 'string' && data.length > 100 && typeof mime === 'string' && EXT[mime]) {
18      out.push({ ext: EXT[mime]!, b64: data })
19    } else Object.values(o).forEach(x => findImages(x, out))
20  }
21
22  return out
23}
24