image-diet

1import type { EngineInterface, McpContentBlock, Register } from 'claude-code'
2
3const DEFAULT_MAX_EDGE = 1280
4const FULL_SIZE_MARKER = '.full.'
5// These tools take click coordinates in the pixels of the screenshot they
6// returned, so a shrunk screenshot would move every click.
7const DEFAULT_SKIP_TOOLS = 'computer|cua|browser_batch'
8const EXT: Record<string, string> = { 'image/png': 'png', 'image/jpeg': 'jpg' }
9
10// Prints "<oldW> <oldH> <newW> <newH>" then the resized base64 on one line.
11// Exit 3 means the image already fits. sips ships with macOS; ImageMagick
12// covers Linux.
13const SHRINK_SH = `set -eu
14max=$1 ext=$2
15d=$(mktemp -d); trap 'rm -rf "$d"' EXIT
16in="$d/in.$ext" out="$d/out.$ext"
17base64 -d > "$in"
18dims() {
19  if command -v sips >/dev/null 2>&1; then
20    sips -g pixelWidth -g pixelHeight "$1" | awk '/pixel/{printf "%s ", $2}'
21  else
22    magick identify -format '%w %h' "$1[0]"
23  fi
24}
25set -- $(dims "$in")
26[ "$1" -gt "$max" ] || [ "$2" -gt "$max" ] || exit 3
27if command -v sips >/dev/null 2>&1; then
28  sips -Z "$max" "$in" --out "$out" >/dev/null
29else
30  magick "$in[0]" -resize "\${max}x\${max}>" "$out"
31fi
32echo "$1 $2 $(dims "$out")"
33base64 < "$out" | tr -d '\\n'
34`
35
36type Shrunk = { base64: string; from: string; to: string; width: number; height: number }
37type ReadImage = {
38  type: 'image'
39  file: { base64: string; type: string; dimensions?: Record<string, number> }
40}
41type McpResult = { content: McpContentBlock[] }
42
43const isReadImage = (r: unknown): r is ReadImage =>
44  (r as ReadImage | undefined)?.type === 'image' && typeof (r as ReadImage).file?.base64 === 'string'
45const isMcpResult = (r: unknown): r is McpResult => Array.isArray((r as McpResult | undefined)?.content)
46
47async function shrink($: EngineInterface, base64: string, mime: string | undefined): Promise<Shrunk | undefined> {
48  const ext = mime && EXT[mime]
49  if (!ext) return undefined
50  const maxEdge = Number(await $.env.get('IMAGE_DIET_MAX_EDGE')) || DEFAULT_MAX_EDGE
51  const ran = await $.process.run(['sh', '-c', SHRINK_SH, 'sh', String(maxEdge), ext], {
52    stdin: base64,
53    timeoutMs: 10_000,
54  })
55  if (ran.exitCode !== 0 || ran.isStdoutTruncated) return undefined
56  const newline = ran.stdout.indexOf('\n')
57  const [ow, oh, w = 0, h = 0] = ran.stdout.slice(0, newline).trim().split(/\s+/).map(Number)
58  const out = ran.stdout.slice(newline + 1).trim()
59  if (!out || !(w > 0 && h > 0)) return undefined
60  return { base64: out, from: `${ow}x${oh}`, to: `${w}x${h}`, width: w, height: h }
61}
62
63const note = (what: string, s: Shrunk) =>
64  `image-diet: shrank ${what} from ${s.from} to ${s.to} to save context. ` +
65  `For full resolution, copy the file to a name containing "${FULL_SIZE_MARKER}" and Read that copy.`
66
67export const register: Register = on => {
68  on('tool.call', { tool: 'Read' }, async ($, e, next) => {
69    const ran = await next(e)
70    try {
71      if (!('result' in ran) || e.file_path.includes(FULL_SIZE_MARKER)) return ran
72      const result = ran.result
73      if (!isReadImage(result)) return ran
74      const s = await shrink($, result.file.base64, result.file.type)
75      if (!s) return ran
76      const dims = result.file.dimensions
77      return {
78        result: {
79          ...result,
80          file: {
81            ...result.file,
82            base64: s.base64,
83            dimensions: dims && { ...dims, displayWidth: s.width, displayHeight: s.height },
84          },
85        },
86        context: [note(e.file_path, s)],
87      }
88    } catch {
89      return ran
90    }
91  }).catch(($, e, next) => next(e))
92
93  on('tool.call', async ($, e, next) => {
94    const ran = await next(e)
95    if (!e.tool.startsWith('mcp__') || !('result' in ran) || !isMcpResult(ran.result)) return ran
96    const result = ran.result
97    try {
98      const skip = new RegExp((await $.env.get('IMAGE_DIET_SKIP_TOOLS')) || DEFAULT_SKIP_TOOLS, 'i')
99      if (skip.test(e.tool)) return ran
100      const notes: string[] = []
101      const content = await Promise.all(
102        result.content.map(async block => {
103          if (block.type !== 'image' || typeof block.data !== 'string') return block
104          const s = await shrink($, block.data, block.mimeType)
105          if (!s) return block
106          notes.push(note(`a ${e.tool} image`, s))
107          return { ...block, data: s.base64 }
108        }),
109      )
110      if (!notes.length) return ran
111      return { result: { ...result, content }, context: notes }
112    } catch {
113      return ran
114    }
115  }).catch(($, e, next) => next(e))
116}
117