mcp-squeeze

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Memory, Seen, Stat, Total } from '../types'
5import { groupOf, MCP_TOOL, OTHER, packOf, switches } from './packs'
6import { hash, imageNote, imageTokens, isRepeatIn, pngSize, squeezeBlock } from './squeeze'
7
8// Each agent's blocks, images and tokens returned this conversation, by hash: `whole` the last copy
9// went whole, so the next is left out; `pointed` it was left out, so the next goes whole. By agent,
10// '' the main one.
11const seenAtom = atom({ plugin: 'mcp-squeeze', key: 'seen' } as const, {} as Record<string, Seen>)
12const statsAtom = atom({ plugin: 'mcp-squeeze', key: 'stats' } as const, {} as Record<string, Stat>)
13// What the packs hold of each agent's conversation (a page's last snapshot), by agent as `seen` is.
14const memoryAtom = atom({ plugin: 'mcp-squeeze', key: 'memory' } as const, {} as Record<string, Memory>)
15
16const OFF_KEY = 'off'
17// Kept true while a pack, or the tools no pack takes, is off by itself: `off:figma`, `off:other`.
18const packOffKey = (pack: string) => `off:${pack}`
19// What it has kept out since it was first on, across sessions (Total).
20const TOTAL_KEY = 'total'
21// Characters a token, roughly, over code and text in either language: for the estimate alone, as
22// no API counts the tokens of a text without a model call that would cost about what it saves.
23const CHARS_PER_TOKEN = 3.5
24
25// The stores of the total one after another, so two calls answering at once both count.
26let totalWrites: Promise<unknown> = Promise.resolve()
27
28type Block = { type: string; text?: string; [field: string]: unknown }
29
30// The text blocks of an MCP tool's record, however it came: blocks, `{ content }`, or a string.
31const blocksOf = (result: unknown): { blocks: Block[]; wrap: (blocks: Block[]) => unknown } | undefined => {
32  if (typeof result === 'string') return { blocks: [{ type: 'text', text: result }], wrap: b => b.map(x => x.text ?? '').join('') }
33  if (Array.isArray(result) && result.every(x => typeof x === 'object' && x !== null && typeof x.type === 'string'))
34    return { blocks: result as Block[], wrap: b => b }
35  if (typeof result === 'object' && result !== null && Array.isArray((result as { content?: unknown }).content)) {
36    const inner = blocksOf((result as { content: unknown }).content)
37    return inner && { blocks: inner.blocks, wrap: b => ({ ...result, content: b }) }
38  }
39  return undefined
40}
41
42const textLength = (blocks: readonly Block[]) => blocks.reduce((n, b) => n + (b.type === 'text' ? (b.text ?? '').length : 0), 0)
43
44// An image block's data and kind, as MCP sends it (`data`, `mimeType`) or as the API takes it
45// (`source.data`, `source.media_type`).
46const imageOf = (b: Block) => {
47  const source = (typeof b.source === 'object' && b.source !== null ? b.source : {}) as { data?: unknown; media_type?: unknown }
48  const data = typeof b.data === 'string' ? b.data : typeof source.data === 'string' ? source.data : undefined
49  const mimeType = typeof b.mimeType === 'string' ? b.mimeType : typeof source.media_type === 'string' ? source.media_type : 'image'
50  return data === undefined ? undefined : { data, mimeType }
51}
52
53// The blocks squeezed, each text as squeezeBlock tells it and an image returned whole before left
54// out; with the images left out and the tokens they would have cost.
55async function squeeze($: EngineInterface, tool: string, agent: string, blocks: readonly Block[]) {
56  const seen = { ...((await read($, seenAtom))[agent] ?? {}) }
57  const memory = { ...((await read($, memoryAtom))[agent] ?? {}) }
58  let images = 0
59  let imageCost = 0
60  const out = blocks.map((b): Block => {
61    if (b.type === 'image') {
62      // Only an image whose size it reads (a PNG, as Figma sends) is left out, so what it saves
63      // is counted; any other passes as it came.
64      const image = imageOf(b)
65      const size = image && pngSize(image.data)
66      if (!image || !size || !isRepeatIn(seen, `image:${hash(image.data)}`)) return b
67      images++
68      imageCost += imageTokens(size)
69      return { type: 'text', text: imageNote(image.mimeType, size) }
70    }
71    if (b.type !== 'text' || typeof b.text !== 'string') return b
72    const { text } = squeezeBlock(tool, b.text, seen, memory)
73    return text === b.text ? b : { ...b, text }
74  })
75  await update($, seenAtom, all => ({ ...all, [agent]: seen }))
76  await update($, memoryAtom, all => ({ ...all, [agent]: memory }))
77  return { blocks: out, images, imageCost }
78}
79
80const kilo = (n: number) => (n < 1000 ? String(n) : `${(n / 1000).toFixed(1)}k`)
81const percent = (s: Stat) => (s.before === 0 ? 0 : Math.round(100 - (100 * s.after) / s.before))
82// The text's change, signed: `−23%` smaller, `+40%` larger (a repeated image's note is text where
83// the image was none).
84const change = (s: Stat) => (percent(s) < 0 ? `+${-percent(s)}%` : `−${percent(s)}%`)
85const tokens = (chars: number) => Math.round(chars / CHARS_PER_TOKEN)
86
87// What one call adds to a tool's stat; an earlier version's stat counts no images.
88const added = (was: Stat | undefined, call: Omit<Stat, 'calls'>): Stat => {
89  const s = was ?? { calls: 0, before: 0, after: 0 }
90  return {
91    calls: s.calls + 1,
92    before: s.before + call.before,
93    after: s.after + call.after,
94    images: (s.images ?? 0) + (call.images ?? 0),
95    imageTokens: (s.imageTokens ?? 0) + (call.imageTokens ?? 0),
96  }
97}
98
99const sum = (stats: Readonly<Record<string, Stat>> | readonly Stat[]) =>
100  Object.values(stats).reduce(
101    (s, r) => ({
102      calls: s.calls + r.calls,
103      before: s.before + r.before,
104      after: s.after + r.after,
105      images: (s.images ?? 0) + (r.images ?? 0),
106      imageTokens: (s.imageTokens ?? 0) + (r.imageTokens ?? 0),
107    }),
108    { calls: 0, before: 0, after: 0, images: 0, imageTokens: 0 } as Stat,
109  )
110
111// A stat's key, `<pack>:<tool>`, or `<server>:<tool>` for a tool no pack takes; a key an earlier
112// version kept is a tool's name alone, of the Figma pack, the only one it had.
113const keyOf = (pack: string, tool: string) => `${pack}:${tool}`
114const splitKey = (key: string) => {
115  const at = key.indexOf(':')
116  return at < 0 ? { pack: 'figma', tool: key } : { pack: key.slice(0, at), tool: key.slice(at + 1) }
117}
118
119// The stats each under its `<pack>:<tool>` key, an earlier version's added to its own.
120const keyed = (stats: Readonly<Record<string, Stat>>) => {
121  const out: Record<string, Stat> = {}
122  for (const [key, s] of Object.entries(stats)) {
123    const { pack, tool } = splitKey(key)
124    const k = keyOf(pack, tool)
125    out[k] = out[k] ? sum([out[k], s]) : s
126  }
127  return out
128}
129
130// How a row is named: by its tool alone while one pack or server has answered, by it and the tool
131// past that.
132const namer = (keys: readonly string[]) => {
133  const isMany = new Set(keys.map(k => splitKey(k).pack)).size > 1
134  return (key: string) => {
135    const { pack, tool } = splitKey(key)
136    return isMany ? `${pack} ${tool}` : tool
137  }
138}
139
140// The tokens a stat kept out: its text's, estimated, and its images'.
141const saved = (s: Stat) => tokens(s.before - s.after) + (s.imageTokens ?? 0)
142
143// The total as kept, or none yet; a value of another shape (an earlier version's) counts as none.
144async function totalIn($: EngineInterface): Promise<Total | undefined> {
145  const t = (await $.store.get(TOTAL_KEY).catch(() => undefined)) as Partial<Total> | undefined
146  return t && typeof t.since === 'number' && typeof t.tools === 'object' && t.tools !== null ? { since: t.since, tools: keyed(t.tools) } : undefined
147}
148
149function addToTotal($: EngineInterface, key: string, call: Omit<Stat, 'calls'>) {
150  totalWrites = totalWrites
151    .then(async () => {
152      const t = (await totalIn($)) ?? { since: await $.clock.now(), tools: {} }
153      await $.store.set(TOTAL_KEY, { since: t.since, tools: { ...t.tools, [key]: added(t.tools[key], call) } })
154    })
155    .catch(() => {})
156  return totalWrites
157}
158
159// One line of what was kept out: the calls, the tokens saved, the text before and after, the
160// images left out.
161const gainLine = (label: string, s: Stat, width: number) =>
162  `${label.padEnd(width)}  ${String(s.calls).padStart(4)} ${s.calls === 1 ? 'call ' : 'calls'}  saved ≈ ${kilo(saved(s)).padStart(6)} tokens  text ${kilo(tokens(s.before)).padStart(6)} → ${kilo(tokens(s.after)).padStart(6)} (${change(s)})${s.images ? `  images left out: ${s.images} (≈ ${kilo(s.imageTokens ?? 0)})` : ''}`
163
164function gain(conversation: Readonly<Record<string, Stat>>, total: Total | undefined) {
165  const since = total ? new Date(total.since).toISOString().slice(0, 10) : ''
166  const name = namer(Object.keys(total?.tools ?? {}))
167  const tools = Object.entries(total?.tools ?? {})
168    .sort((a, b) => saved(b[1]) - saved(a[1]))
169    .map(([k, s]) => [name(k), s] as const)
170  const width = Math.max(`all time, since ${since}`.length, ...tools.map(([t]) => t.length + 2))
171  return [
172    `Tokens the MCP tools' results did not bring into the context (text estimated at ${CHARS_PER_TOKEN} characters a token, an image at a token per 750 pixels):`,
173    '',
174    gainLine('this conversation', sum(conversation), width),
175    ...(total ? [gainLine(`all time, since ${since}`, sum(total.tools), width), '', 'By tool, all time:', ...tools.map(([t, s]) => gainLine(`  ${t}`, s, width))] : []),
176  ].join('\n')
177}
178
179const FOOT = `/squeeze off turns it off, /squeeze on back on; /squeeze <pack> off one pack alone, /squeeze ${OTHER} off the tools no pack takes; /squeeze-gain says the tokens saved, all time too.`
180const USAGE = `Usage: /squeeze, /squeeze on, /squeeze off, /squeeze <pack> on, /squeeze <pack> off (${switches.join(', ')})`
181
182// Each pack, and the tools no pack takes, on or off by itself: "Packs: figma on, other on."
183const packsLine = (offPacks: ReadonlySet<string>) => `Packs: ${switches.map(p => `${p} ${offPacks.has(p) ? 'off' : 'on'}`).join(', ')}.`
184
185function report(stats: Readonly<Record<string, Stat>>, isOn: boolean, offPacks: ReadonlySet<string>) {
186  const name = namer(Object.keys(stats))
187  const rows = Object.entries(stats)
188    .sort((a, b) => b[1].before - a[1].before)
189    .map(([k, s]) => [name(k), s] as const)
190  const all = sum(stats)
191  const head = `MCP squeeze is ${isOn ? 'on' : 'off'}.`
192  if (all.calls === 0) return `${head} No MCP tool has answered in this conversation yet.\n${packsLine(offPacks)}\n${FOOT}`
193  const width = Math.max(...rows.map(([t]) => t.length))
194  const images = all.images ? `, ${all.images} repeated ${all.images === 1 ? 'image' : 'images'} left out` : ''
195  return [
196    `${head} This conversation: ${all.calls} MCP ${all.calls === 1 ? 'call' : 'calls'}, ${kilo(all.before)} → ${kilo(all.after)} characters (${change(all)})${images}; about ${kilo(saved(all))} tokens kept out of the context.`,
197    ...rows.map(([t, r]) => `  ${t.padEnd(width)}  ${String(r.calls).padStart(3)} × ${kilo(r.before).padStart(6)} → ${kilo(r.after).padStart(6)}  ${change(r)}`),
198    packsLine(offPacks),
199    FOOT,
200  ].join('\n')
201}
202
203const isOnIn = async ($: EngineInterface) => (await $.store.get(OFF_KEY).catch(() => undefined)) !== true
204const isPackOnIn = async ($: EngineInterface, pack: string) => (await $.store.get(packOffKey(pack)).catch(() => undefined)) !== true
205
206async function offPacksIn($: EngineInterface) {
207  const off = new Set<string>()
208  for (const p of switches) if (!(await isPackOnIn($, p))) off.add(p)
209  return off
210}
211
212export const register: Register = on => {
213  on('session.start', async ($, e, next) => {
214    await $.command.register({ name: 'squeeze', description: 'What squeezing the MCP tools has kept out of the context; on or off, all or one pack', argumentHint: '[on|off|<pack> on|<pack> off]' })
215    await $.command.register({ name: 'squeeze-gain', description: 'Tokens squeezing the MCP tools has saved: this conversation, all time, by tool' })
216    return next(e)
217  })
218
219  on('command.run', { command: 'squeeze-gain' }, async $ => ({ text: gain(keyed(await read($, statsAtom)), await totalIn($)) }))
220
221  // A compact starts the compacting agent's context over: a block left out now would point at what
222  // the summary took the place of. Once the compact stood, not for one skipped, nor for a summary
223  // worked out ahead of time (`precompute`), which leaves the context as it was.
224  on('session.compact', async ($, e, next) => {
225    const ran = await next(e)
226    if (e.trigger !== 'precompute' && ran.skip === undefined) {
227      await update($, seenAtom, all => ({ ...all, [e.agentId ?? '']: {} }))
228      await update($, memoryAtom, all => ({ ...all, [e.agentId ?? '']: {} }))
229    }
230    return ran
231  }).catch(($, e, next) => next(e))
232
233  // A /clear starts a new conversation: what was returned and what was kept out start over.
234  on('session.end', async ($, e, next) => {
235    const ran = await next(e)
236    if (e.reason === 'clear') {
237      await update($, seenAtom, () => ({}))
238      await update($, memoryAtom, () => ({}))
239      await update($, statsAtom, () => ({}))
240    }
241    return ran
242  }).catch(($, e, next) => next(e))
243
244  // An MCP tool's result is squeezed after it ran, by its pack if one takes it; an error, a shape it
245  // does not know, or a pack turned off passes as it came, as does a tool of no MCP server, which the
246  // hook never sees.
247  on('tool.call', { tool: MCP_TOOL }, async ($, e, next) => {
248    const ran = await next(e)
249    if (ran.deny !== undefined || ran.isError || !(await isOnIn($)) || !(await isPackOnIn($, packOf(e.tool)?.id ?? OTHER))) return ran
250    const found = blocksOf(ran.result)
251    if (!found) return ran
252    const { blocks, images, imageCost } = await squeeze($, e.tool, e.agentId ?? '', found.blocks)
253    // An image left out leaves a note in its place, counted with the text.
254    const isChanged = blocks.some((b, k) => b !== found.blocks[k])
255    const before = textLength(found.blocks)
256    const after = isChanged ? textLength(blocks) : before
257    const call = { before, after, images, imageTokens: imageCost }
258    const key = keyOf(groupOf(e.tool), e.tool.slice(e.tool.lastIndexOf('__') + 2))
259    await update($, statsAtom, was => {
260      const s = keyed(was)
261      return { ...s, [key]: added(s[key], call) }
262    })
263    await addToTotal($, key, call)
264    if (!isChanged) return ran
265    return { result: found.wrap(blocks), ...(ran.context ? { context: ran.context } : {}) }
266  }).catch(($, e, next) => next(e))
267
268  on('command.run', { command: 'squeeze' }, async ($, e) => {
269    const args = e.args.trim().toLowerCase().split(/\s+/).filter(Boolean)
270    const [first, second] = args
271    const isSwitch = (word: string | undefined) => word === 'on' || word === 'off'
272    if (args.length === 1 && isSwitch(first)) await $.store.set(OFF_KEY, first === 'off')
273    else if (args.length === 2 && switches.includes(first!) && isSwitch(second)) await $.store.set(packOffKey(first!), second === 'off')
274    else if (args.length !== 0) return { text: USAGE }
275    return { text: report(keyed(await read($, statsAtom)), await isOnIn($), await offPacksIn($)) }
276  })
277}
278

1// The packs: each one squeezes the tools of the MCP servers it knows, by what each tool returns. A
2// tool no pack takes has only what the mod does to every MCP tool (repeats left out).
3import type { Memory } from '../types'
4import { chrome } from './chrome.ts'
5import { figma } from './figma.ts'
6
7export type Pack = {
8  // How the pack is named in the stats and in `/squeeze <id> off`.
9  id: string
10  // The names of the servers whose tools the pack takes, whatever a config names them, matched
11  // whole and in any case, as they read in a tool's name: `claude_ai_Figma` in
12  // `mcp__claude_ai_Figma__get_metadata`.
13  server: RegExp
14  // A line the pack starts a kind of block with, passed over where a repeat quotes how one starts.
15  legend?: string
16  // The lines that say a state the server keeps and says again after each call, each pattern's
17  // first group the label and its second the value: where the value is the one the last such line
18  // said, the line says so in its place.
19  states?: readonly RegExp[]
20  // One text block of the tool's result, squeezed; the caller keeps the original where it is no
21  // shorter.
22  squeezeText: (tool: string, text: string) => string
23  // The block squeezed, as the pack tells it by what it holds of this agent's conversation:
24  // `memory`, its keys under its id, kept up for the next (a page's last snapshot); `isAgain` for
25  // the copy the model called for again, after one left out. `isWhole` where the pack gives it
26  // whole as the model called for it so, which no repeat then leaves out.
27  tell?: (tool: string, text: string, memory: Memory, isAgain: boolean) => { text: string; isWhole?: boolean }
28}
29
30export const packs: readonly Pack[] = [figma, chrome]
31
32const isServer = (pack: Pack, server: string) => new RegExp(`^(?:${pack.server.source})$`, 'i').test(server)
33
34// The pack that takes an MCP tool, by its server, the part of the name before the last `__`, so
35// another server's `figma_export` is no Figma tool; none for a tool of no MCP server.
36export const packOf = (tool: string) => {
37  if (!tool.startsWith('mcp__') || tool.lastIndexOf('__') < 5) return undefined
38  const server = tool.slice(5, tool.lastIndexOf('__'))
39  return packs.find(p => isServer(p, server))
40}
41
42// Every MCP server's tools, for a hook's matcher, so the hook never sees any other: the server the
43// part of the name before the last `__`.
44export const MCP_TOOL = /^mcp__.+__(?:(?!__).)+$/
45
46// The tools no pack takes, as `/squeeze other off` names them.
47export const OTHER = 'other'
48
49// What `/squeeze <name> on|off` turns: each pack, and the tools no pack takes.
50export const switches = [...packs.map(p => p.id), OTHER]
51
52// The tool's group in the stats: its pack, or else its server.
53export const groupOf = (tool: string) => packOf(tool)?.id ?? tool.slice(5, tool.lastIndexOf('__'))
54

1// What the mod does to the text an MCP tool returns, whichever pack squeezes it, and the notes it
2// leaves where a repeat was left out. What each pack does to its own tools is in its file.
3import type { Memory, Seen } from '../types'
4import { packOf, packs } from './packs.ts'
5
6export { collapseInstances, hoistClasses, isCode, outline, OUTLINE_LEGEND, roundNoise } from './figma.ts'
7
8// FNV-1a over a text: which blocks were returned before.
9export const hash = (text: string) => {
10  let h = 0x811c9dc5
11  for (const ch of text) {
12    h ^= ch.codePointAt(0)!
13    h = Math.imul(h, 0x01000193) >>> 0
14  }
15  return h.toString(16).padStart(8, '0')
16}
17
18// One text block of an MCP tool's result, squeezed by its server's pack; a tool no pack takes as it
19// came. Never longer than it was. Its line breaks read as `\n`, however it came (`\r\n` from a
20// server on Windows); a block with nothing else to squeeze comes as it came, line breaks and all.
21export const squeezeText = (tool: string, text: string) => {
22  const lf = text.replaceAll('\r\n', '\n')
23  const out = packOf(tool)?.squeezeText(tool, lf) ?? lf
24  return out.length < lf.length ? out : text
25}
26
27// An opaque token written out long: a JWT, or a run of base64 (or base64url).
28const TOKEN = /\beyJ[\w-]+\.[\w-]+\.[\w-]+|(?<![\w+/=-])[A-Za-z0-9+/_-]+={0,2}(?![\w+/=-])/g
29// A token this long or longer is told by a note where it came whole before.
30export const TOKEN_MIN = 160
31
32// Long, with capitals, small letters and digits, as random bytes written out are, and no path:
33// base64 has a slash in 64 characters, a path one every few.
34const isOpaque = (t: string) => t.length >= TOKEN_MIN && /[A-Z]/.test(t) && /[a-z]/.test(t) && /\d/.test(t) && (t.match(/\//g)?.length ?? 0) * 16 < t.length
35
36// The opaque tokens of a text, as they are written.
37export const tokensOf = (text: string) => [...text.matchAll(TOKEN)].map(m => m[0]).filter(isOpaque)
38
39export const tokenKey = (token: string) => `token:${hash(token)}`
40
41export const tokenNote = (token: string) => `[the same ${token.length}-character token as earlier in this conversation, ${token.slice(0, 12)}…${token.slice(-6)}]`
42
43// A text with each opaque token `isRepeat` says to leave out told by a note: the token stays in the
44// context where it came whole before (a bearer token in every request of a page, a data URI).
45export const shortenTokens = (text: string, isRepeat: (key: string) => boolean) =>
46  text.replace(TOKEN, t => (isOpaque(t) && isRepeat(tokenKey(t)) ? tokenNote(t) : t))
47
48// A block this long or longer is left out where it was returned whole before.
49export const REPEAT_MIN = 300
50
51// Whether this copy is left out, by what `seen` holds: so when the last copy went whole, and the next
52// goes whole again; kept in `seen` for the next.
53export const isRepeatIn = (seen: Seen, key: string) => {
54  const was = seen[key] === 'whole'
55  seen[key] = was ? 'pointed' : 'whole'
56  return was
57}
58
59// What a state line says in place of a value the last such line said.
60export const STATE_SAME = '(as before)'
61
62// A text with each state line whose value is the one the last line of its label said told so, a
63// line with another value kept as the state now; with `isAgain`, each line kept, as the state now.
64const tellStates = (text: string, states: readonly RegExp[], pack: string, seen: Seen, isAgain: boolean) =>
65  text
66    .split('\n')
67    .map(line => {
68      const m = states.map(re => re.exec(line)).find(Boolean)
69      if (!m) return line
70      const [, label = '', value = ''] = m
71      const prefix = `state:${pack}:${hash(label)}:`
72      if (!isAgain && seen[prefix + hash(value)] === 'whole' && STATE_SAME.length < value.length) return `${label}${STATE_SAME}`
73      for (const key of Object.keys(seen)) if (key.startsWith(prefix)) delete seen[key]
74      seen[prefix + hash(value)] = 'whole'
75      return line
76    })
77    .join('\n')
78
79// One text block as the model is to read it, by what this agent's conversation holds (`seen` and
80// the packs' `memory`, kept up for the next): squeezed by its pack and told as its pack tells it,
81// left out where the same came whole before, or else each state line that says what the last did
82// told so and each token that came whole before told by a note. The copy after one left out is the
83// one the model called for again, its earlier copy out of view: it comes whole, states and tokens
84// and all. With the text its pack alone made of it.
85export const squeezeBlock = (tool: string, text: string, seen: Seen, memory: Memory) => {
86  const packed = squeezeText(tool, text)
87  const key = hash(packed)
88  const isLong = packed.length >= REPEAT_MIN
89  const pack = packOf(tool)
90  // Told whatever comes of it, so what the pack holds is what the model reads now.
91  const tale = pack?.tell?.(tool, packed, memory, isLong && seen[key] === 'pointed')
92  const told = tale?.text ?? packed
93  const isAgain = (isLong && seen[key] === 'pointed') || tale?.isWhole === true
94  if (isAgain && isLong) seen[key] = 'whole'
95  else if (isLong && isRepeatIn(seen, key)) {
96    const note = repeatNote(packed)
97    return { packed, text: told.length < note.length ? told : note }
98  }
99  const stated = pack?.states ? tellStates(told, pack.states, pack.id, seen, isAgain) : told
100  if (!isAgain) return { packed, text: shortenTokens(stated, k => isRepeatIn(seen, k)) }
101  for (const t of tokensOf(stated)) seen[tokenKey(t)] = 'whole'
102  return { packed, text: stated }
103}
104
105const BASE64 = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'
106
107// A PNG's width and height from its header (the first 24 bytes, 32 base64 characters); undefined
108// for any other image.
109export const pngSize = (base64: string) => {
110  const bytes: number[] = []
111  for (let i = 0; i + 4 <= 32 && i + 4 <= base64.length; i += 4) {
112    const n = [...base64.slice(i, i + 4)].reduce((v, ch) => (v << 6) | Math.max(0, BASE64.indexOf(ch)), 0)
113    bytes.push((n >> 16) & 255, (n >> 8) & 255, n & 255)
114  }
115  const isPng = bytes.length >= 24 && bytes[1] === 0x50 && bytes[2] === 0x4e && bytes[3] === 0x47
116  const u32 = (at: number) => ((bytes[at]! << 24) | (bytes[at + 1]! << 16) | (bytes[at + 2]! << 8) | bytes[at + 3]!) >>> 0
117  return isPng ? { width: u32(16), height: u32(20) } : undefined
118}
119
120// The tokens an image of this size costs the model, by Anthropic's rule of thumb: about one per 750
121// pixels, after it is scaled to fit 1568 pixels a side and 1.15 megapixels.
122export const imageTokens = (size: { width: number; height: number }) => {
123  const { width: w, height: h } = size
124  if (w <= 0 || h <= 0) return 0
125  const scale = Math.min(1, 1568 / Math.max(w, h), Math.sqrt(1_150_000 / (w * h)))
126  return Math.round((w * scale * (h * scale)) / 750)
127}
128
129export const imageNote = (mimeType: string, size: { width: number; height: number }) =>
130  `[The same image (${mimeType}, ${size.width}×${size.height}) this tool returned earlier in this conversation: left out to save context. If it is no longer in view, call the tool again and the next copy comes whole.]\n`
131
132// Quoting how the block starts, past a pack's legend, which every block of its kind starts with.
133const pastLegend = (text: string) => packs.reduce((t, p) => (p.legend ? t.replace(`${p.legend}\n`, '') : t), text)
134
135export const repeatNote = (text: string) =>
136  `[The same ${text.length} characters this tool returned earlier in this conversation, starting "${pastLegend(text).slice(0, 60).replace(/\s+/g, ' ').trim()}…": left out to save context. If that copy is no longer in view, call the tool again and the next copy comes whole.]\n`
137

1// The Chrome DevTools pack: what the mod does to the text the Chrome DevTools server's tools return
2// (`chrome-devtools`, `chrome_devtools`, `chrome-devtools-mcp`). It keeps what the model acts on
3// (every uid, role, name, value, URL) and drops only what spells it long.
4import type { Memory } from '../types'
5import type { Pack } from './packs.ts'
6
7// A run of percent-escapes of bytes past ASCII, as UTF-8 writes a character past it.
8const ESCAPED = /(?:%[89A-Fa-f][0-9A-Fa-f])+/g
9// A character better left escaped: one that does not show, or shows as a space.
10const UNSEEN = /[\p{C}\p{Z}]/u
11
12// A URL with each character past ASCII that is percent-escaped written as itself: the same URL to a
13// browser, one character where it was six. An escaped ASCII character (%20, %2F, %26) stays, as
14// unescaped it may mean something else; so does a run that is no UTF-8, or spells a character that
15// does not show.
16export const unescapeUrl = (url: string) =>
17  url.replace(ESCAPED, run => {
18    let text: string
19    try {
20      text = decodeURIComponent(run)
21    } catch {
22      return run
23    }
24    return UNSEEN.test(text) ? run : text
25  })
26
27// The page's snapshot, from its heading on: `take_snapshot`'s, and the one after a click, a fill, a
28// wait.
29const SNAPSHOT = '## Latest page snapshot\n'
30const URL_ATTR = / url="([^"]*)"/g
31
32// A result with each `url="…"` of its snapshot unescaped; whatever comes before the snapshot (a URL
33// navigated to) is left as it is, and so is a result with none.
34export const unescapeSnapshot = (text: string) => {
35  const at = text.startsWith(SNAPSHOT) ? 0 : text.indexOf(`\n${SNAPSHOT}`)
36  if (at < 0) return text
37  return text.slice(0, at) + text.slice(at).replace(URL_ATTR, (all, url: string) => ` url="${unescapeUrl(url)}"`)
38}
39
40// A snapshot's line: its indent, then its node's uid, which the node keeps from one snapshot to the
41// next while it is on the page.
42const NODE = /^(\s*)uid=(\S+)(?:\s|$)/
43const HEADING = '## Latest page snapshot'
44
45type Mark = ' ' | '+' | '-' | '~'
46type Entry = { mark: Mark; line: string; keep: boolean }
47
48const indentOf = (line: string) => line.length - line.trimStart().length
49const uidOf = (line: string) => NODE.exec(line)![2]!
50
51// The longest run of `positions` that only rises, as their indexes: the nodes of a snapshot in the
52// order they had in the last one.
53const rising = (positions: readonly number[]) => {
54  const tails: number[] = []
55  const before: number[] = []
56  positions.forEach((p, k) => {
57    if (p < 0) return
58    let lo = 0
59    let hi = tails.length
60    while (lo < hi) {
61      const mid = (lo + hi) >> 1
62      if (positions[tails[mid]!]! < p) lo = mid + 1
63      else hi = mid
64    }
65    before[k] = lo > 0 ? tails[lo - 1]! : -1
66    tails[lo] = k
67  })
68  const out = new Set<number>()
69  for (let k = tails.at(-1) ?? -1; k >= 0; k = before[k]!) out.add(k)
70  return out
71}
72
73// The snapshot `now` told by its changes since `last`, both its lines: each line marked new (+),
74// gone (-, a node gone with all in it told once) or changed (~), with the lines above it it is in,
75// and each run of lines as they were told by their count. Undefined where a node's uid is on two
76// lines, which no snapshot has.
77export const snapshotChanges = (last: readonly string[], now: readonly string[]) => {
78  const lastAt = new Map(last.map((line, k) => [uidOf(line), k]))
79  if (lastAt.size !== last.length || new Set(now.map(uidOf)).size !== now.length) return undefined
80  const positions = now.map(line => lastAt.get(uidOf(line)) ?? -1)
81  const kept = rising(positions)
82  const entries: Entry[] = []
83  let i = 0
84  // The last snapshot's lines from `i` up to `to`, the next one still on the page, as gone: each
85  // told, but for a line in one told before it, which it counts. A line in it that is still on the
86  // page is past `to`, and no line past `to` is counted, so each line is told once.
87  const gone = (to: number) => {
88    while (i < to) {
89      const line = last[i]!
90      let count = 0
91      while (i + 1 + count < to && indentOf(last[i + 1 + count]!) > indentOf(line)) count++
92      entries.push({ mark: '-', line: count ? `${line} (and the ${count} ${count === 1 ? 'line' : 'lines'} in it)` : line, keep: true })
93      i += 1 + count
94    }
95  }
96  now.forEach((line, k) => {
97    if (!kept.has(k)) {
98      entries.push({ mark: '+', line, keep: true })
99      return
100    }
101    gone(positions[k]!)
102    i = positions[k]! + 1
103    entries.push(line === last[positions[k]!] ? { mark: ' ', line, keep: false } : { mark: '~', line, keep: true })
104  })
105  gone(last.length)
106  // Each change keeps the lines it is in, the nearest above it of each lesser indent. A line gone is
107  // in no line of the snapshot now, so it closes none.
108  const open: Entry[] = []
109  for (const entry of entries) {
110    const indent = indentOf(entry.line)
111    if (entry.mark === '-') {
112      for (const parent of open) if (indentOf(parent.line) < indent) parent.keep = true
113      continue
114    }
115    while (open.length && indentOf(open.at(-1)!.line) >= indent) open.pop()
116    if (entry.mark !== ' ') for (const parent of open) parent.keep = true
117    open.push(entry)
118  }
119  if (entries.every(e => e.mark === ' ')) return []
120  const out: string[] = []
121  for (let k = 0; k < entries.length; ) {
122    const entry = entries[k]!
123    if (entry.keep) {
124      out.push(`${entry.mark} ${entry.line}`)
125      k++
126      continue
127    }
128    // A run as deep as its first line or deeper, so the count stays in what that line is in.
129    let end = k
130    while (end < entries.length && !entries[end]!.keep && indentOf(entries[end]!.line) >= indentOf(entry.line)) end++
131    out.push(`  ${' '.repeat(indentOf(entry.line))}… ${end - k} ${end - k === 1 ? 'line' : 'lines'} as before`)
132    k = end
133  }
134  return out
135}
136
137// A snapshot told by its changes is used where it comes to this part of the whole, or less.
138const CHANGES_AT_MOST = 0.6
139
140// A result with its snapshot told by its changes since the last one this agent read, as the pack
141// remembers it: where there is one, the changes come to little enough, and this is no copy called
142// for again nor a take_snapshot called next after a snapshot told so, which come whole. Either way
143// the snapshot is remembered as the one the model reads now.
144const tellSnapshot = (tool: string, text: string, memory: Memory, isAgain: boolean): { text: string; isWhole?: boolean } => {
145  const lines = text.split('\n')
146  const at = lines.indexOf(HEADING)
147  let end = at + 1
148  while (at >= 0 && end < lines.length && NODE.test(lines[end]!)) end++
149  const wasTold = memory['chrome:told'] === 'yes'
150  // Any call between a snapshot told so and a take_snapshot makes the take_snapshot no call for it.
151  memory['chrome:told'] = 'no'
152  if (at < 0 || end === at + 1) return { text }
153  const now = lines.slice(at + 1, end)
154  const last = memory['chrome:snapshot']
155  memory['chrome:snapshot'] = now.join('\n')
156  if (wasTold && tool.endsWith('__take_snapshot')) return { text, isWhole: true }
157  if (last === undefined || isAgain) return { text }
158  const changes = snapshotChanges(last.split('\n'), now)
159  if (!changes) return { text }
160  const head =
161    changes.length === 0
162      ? [`${HEADING}: the same as the last one, its ${now.length} lines as they were. A take_snapshot called next gives it whole.`]
163      : [`${HEADING}, told by its changes since the last one: + a line new, - gone, ~ changed, each with the lines it is in; "…" lines as they were. A take_snapshot called next gives it whole.`, ...changes]
164  const whole = now.join('\n').length
165  if (head.join('\n').length > whole * CHANGES_AT_MOST) return { text }
166  memory['chrome:told'] = 'yes'
167  return { text: [...lines.slice(0, at), ...head, ...lines.slice(end)].join('\n') }
168}
169
170export const chrome: Pack = {
171  id: 'chrome',
172  server: /.*chrome[-_]?dev[-_]?tools.*/,
173  // What it emulates (`Emulating viewport: {…}`), said again after each call while it lasts.
174  states: [/^(Emulating[^:\n]*: )(.+)$/],
175  squeezeText: (tool, text) => unescapeSnapshot(text),
176  tell: tellSnapshot,
177}
178

1// The Figma pack: what the mod does to the text the Figma servers' tools return (`figma`,
2// `claude_ai_Figma`). Each step keeps what the model acts on (node ids, names, positions, sizes,
3// every class) and drops only what repeats it or spells it long.
4import type { Pack } from './packs.ts'
5
6// A number written with seven or more places, as a float32 value read as a double is
7// (1.1299999952316284 for 1.13).
8const LONG_NUMBER = /(?<![\w.])-?\d+\.\d{7,}(?![\w.])/g
9
10// The shortest writing of a float32 value that reads back as the same float32; a number that is no
11// float32 value (Math.PI, 0.00001234) is left as it is, so nothing it says is lost.
12const shortFloat32 = (written: string) => {
13  const v = Number(written)
14  if (Math.fround(v) !== v) return written
15  for (let digits = 1; digits <= 9; digits++) {
16    const short = String(Number(v.toPrecision(digits)))
17    if (Math.fround(Number(short)) === v) return short.length < written.length ? short : written
18  }
19  return written
20}
21
22export const roundNoise = (text: string) => text.replace(LONG_NUMBER, shortFloat32)
23
24const OPEN = /^(\s*)<([\w-]+)((?:\s+[\w-]+="[^"]*")*)\s*\/?>$/
25const ATTR = /([\w-]+)="([^"]*)"/g
26const CLOSING = /^\s*(?:<\/[\w-]+>)+/
27
28export const OUTLINE_LEGEND = '[Figma metadata as an outline: type id "name" x,y width×height, each child one space in from its parent; a node whose name is its type (Frame, Group…) goes without it]'
29
30// `get_metadata`'s XML as an outline, one node a line, the closing tags dropped since the indent
31// says where a node ends; a line that is no tag (the note after it) is kept as it is. A line may end
32// `\r\n`, as a server on Windows may send it.
33export const outline = (text: string) => {
34  const out: string[] = []
35  let isTree = false
36  for (const raw of text.split(/\r?\n/)) {
37    const line = raw.replace(CLOSING, '')
38    if (line !== raw) isTree = true
39    if (line.trim() === '') {
40      if (line === raw) out.push(line)
41      continue
42    }
43    const m = OPEN.exec(line)
44    if (!m) {
45      out.push(line)
46      continue
47    }
48    isTree = true
49    const [, indent = '', tag = '', attrs = ''] = m
50    const { id, name, x, y, width, height, ...rest } = Object.fromEntries([...attrs.matchAll(ATTR)].map(a => [a[1]!, a[2]!]))
51    const parts = [tag]
52    if (id !== undefined) parts.push(id)
53    if (name !== undefined && name.toLowerCase() !== tag.toLowerCase()) parts.push(`"${name}"`)
54    if (x !== undefined && y !== undefined) parts.push(`${x},${y}`)
55    else if (x !== undefined || y !== undefined) parts.push(...(x !== undefined ? [`x=${x}`] : []), ...(y !== undefined ? [`y=${y}`] : []))
56    if (width !== undefined && height !== undefined) parts.push(`${width}×${height}`)
57    else if (width !== undefined || height !== undefined) parts.push(...(width !== undefined ? [`width=${width}`] : []), ...(height !== undefined ? [`height=${height}`] : []))
58    for (const [k, v] of Object.entries(rest)) parts.push(`${k}=${v}`)
59    out.push(`${' '.repeat(Math.floor(indent.length / 2))}${parts.join(' ')}`)
60  }
61  // Ended by a line break: the server sends its note after the tree as a block of its own, which
62  // the model reads joined on, so the last node keeps a line to itself.
63  return isTree ? `${OUTLINE_LEGEND}\n${out.join('\n').replace(/\n*$/, '\n')}` : text
64}
65
66// A `data-name` that says only what kind of node it is.
67const GENERIC_NAME = / data-name="(?:Frame|Group|Rectangle|Vector|Ellipse|Line|Container)"/g
68
69const CLASS = /className="([^"]*)"/g
70// Where the constants go: right before the component, at the start of a line, so whatever leads
71// the code (imports over several lines, a 'use client') stays as it was. They are read only when a
72// component renders, so one declared above it may use them too.
73const COMPONENT = [/^export default function /m, /^(?:export )?(?:default )?function \w/m]
74
75// `get_design_context`'s code with each className used more than once, and long enough to be worth
76// it, declared once as a constant before the component; the code means the same, the value kept
77// exactly (written as a JSON string, so a `\/` in a class stays one). Code with no component to put
78// them before is left as it is.
79export const hoistClasses = (code: string, minLength = 24) => {
80  const at = COMPONENT.map(re => code.search(re)).find(i => i >= 0) ?? -1
81  if (at < 0) return code
82  const counts = new Map<string, number>()
83  for (const m of code.matchAll(CLASS)) counts.set(m[1]!, (counts.get(m[1]!) ?? 0) + 1)
84  const hoisted = [...counts].filter(([c, n]) => n > 1 && c.length >= minLength).map(([c]) => c)
85  if (hoisted.length === 0) return code
86  const prefix = /\bcn\d+\b/.test(code) ? 'cls' : 'cn'
87  const names = new Map(hoisted.map((c, i) => [c, `${prefix}${i + 1}`]))
88  const decls = hoisted.map(c => `const ${names.get(c)} = ${JSON.stringify(c)}`).join('\n')
89  const use = (part: string) => part.replace(CLASS, (all, c: string) => (names.has(c) ? `className={${names.get(c)}}` : all))
90  return `${use(code.slice(0, at))}${decls}\n\n${use(code.slice(at))}`
91}
92
93export const isCode = (text: string) => text.includes('className="') && /\bexport default function\b|\bfunction \w+\(/.test(text)
94
95const NODE_ID = /data-node-id="([^"]+)"/
96const DATA_NAME = /data-name="([^"]*)"/
97const SRC = /src=\{([^}]*)\}/g
98const BARE_ID = /^\d+:\d+$/
99
100const indentOf = (line: string) => line.length - line.trimStart().length
101const isTag = (line: string) => line.trimStart().startsWith('<')
102// A tag whose element goes on over the lines after it: not one that closes an element, closes
103// itself, or whose element closes on its own line (`<p data-node-id="1:3">Pro</p>`).
104const CLOSED_ON_LINE = /<\/[\w.-]*>\s*$/
105const isOpening = (line: string) => isTag(line) && !line.trimStart().startsWith('</') && !line.trimEnd().endsWith('/>') && line.trimEnd().endsWith('>') && !CLOSED_ON_LINE.test(line)
106
107// The line that closes the element `lines[at]` opens: the next closing tag at its indent, as the
108// code Figma writes puts each tag on a line of its own; -1 where there is none.
109const closingOf = (lines: readonly string[], at: number) => {
110  const indent = indentOf(lines[at]!)
111  for (let i = at + 1; i < lines.length; i++) if (indentOf(lines[i]!) === indent && lines[i]!.trimStart().startsWith('</')) return i
112  return -1
113}
114
115type Slot = { id: string; value: string }
116
117// An element's inside as a shape to compare: its lines from its own indent, each node id less
118// `prefix`, each text and each `src={…}` a slot; and the slots, by the node they are in.
119const shapeOf = (lines: readonly string[], indent: number, prefix: string) => {
120  const slots: Slot[] = []
121  const open: string[] = []
122  const shape = lines.map(raw => {
123    const line = raw.slice(indent).split(`data-node-id="${prefix}`).join('data-node-id="')
124    if (!isTag(line)) {
125      slots.push({ id: open.at(-1) ?? '', value: line.trim() })
126      return `${' '.repeat(indentOf(line))}§`
127    }
128    if (line.trimStart().startsWith('</')) {
129      open.pop()
130      return line
131    }
132    const id = NODE_ID.exec(line)?.[1] ?? open.at(-1) ?? ''
133    for (const m of line.matchAll(SRC)) slots.push({ id, value: `src={${m[1]}}` })
134    if (isOpening(line)) open.push(id)
135    return line.replace(SRC, 'src={§}')
136  })
137  return { shape: shape.join('\n'), slots }
138}
139
140type Template = { name: string; label: string; prefix: string; shape: string; slots: Slot[] }
141
142// `get_design_context`'s code with each instance whose inside is a component's already in the code
143// (the function Figma wrote for it, or an earlier instance written out whole) told in one line in
144// place of that inside: what it is the same as, how its node ids read, and the texts and images that
145// differ, by node. Nothing it says is lost; an inside that differs in anything else stays whole.
146export const collapseInstances = (code: string) => {
147  const lines = code.split(/\r?\n/)
148  const templates: Template[] = []
149  // The components Figma wrote as functions: the element each returns, its node ids as they are.
150  // Only the first tag in the function's own body counts, and only with a node id of its own.
151  for (let i = 0; i < lines.length; i++) {
152    const fn = /^function (\w+)\(/.exec(lines[i]!)
153    if (!fn) continue
154    const body = lines.findIndex((l, k) => k > i && l.startsWith('}'))
155    const root = lines.findIndex((l, k) => k > i && (body < 0 || k < body) && isTag(l))
156    const end = root >= 0 && isOpening(lines[root]!) && NODE_ID.test(lines[root]!) ? closingOf(lines, root) : -1
157    if (end < 0 || (body >= 0 && end > body)) continue
158    const id = NODE_ID.exec(lines[root]!)![1]!
159    const { shape, slots } = shapeOf(lines.slice(root + 1, end), indentOf(lines[root]!), '')
160    templates.push({ name: DATA_NAME.exec(lines[root]!)?.[1] ?? '', label: `<${fn[1]}> (${id})`, prefix: '', shape, slots })
161  }
162  const out: string[] = []
163  for (let i = 0; i < lines.length; i++) {
164    const line = lines[i]!
165    const id = NODE_ID.exec(line)?.[1]
166    const end = id && BARE_ID.test(id) && isOpening(line) ? closingOf(lines, i) : -1
167    const inside = end < 0 ? [] : lines.slice(i + 1, end)
168    const prefix = `I${id};`
169    const isInstance = inside.some(l => NODE_ID.test(l)) && inside.every(l => !NODE_ID.test(l) || NODE_ID.exec(l)![1]!.startsWith(prefix))
170    if (!isInstance) {
171      out.push(line)
172      continue
173    }
174    const indent = indentOf(line)
175    const { shape, slots } = shapeOf(inside, indent, prefix)
176    const name = DATA_NAME.exec(line)?.[1] ?? ''
177    const same = templates.find(t => t.name === name && t.shape === shape)
178    if (!same) {
179      templates.push({ name, label: `${id} above`, prefix, shape, slots })
180      out.push(line)
181      continue
182    }
183    // A slot is named by the component's own id of its node, as the line says how to read it here;
184    // its value a JSON string in which the slash of any star-slash is written as the JSON escape
185    // for a slash (backslash, u002f): the same string to JSON.parse, and a text holding a
186    // star-slash no longer ends the comment.
187    const differs = slots.filter((s, k) => s.value !== same.slots[k]?.value).map(s => `${s.id} ${JSON.stringify(s.value).replaceAll('*/', '*\\u002f')}`)
188    const ids = same.prefix ? `"${prefix}" + its own ids (there after "${same.prefix}")` : `"${prefix}" + its own ids`
189    const note = `${' '.repeat(indent + 2)}{/* inside as ${same.label}, node ids ${ids}${differs.length ? `; differs at its own ids: ${differs.join(', ')}` : ''} */}`
190    if (note.length >= inside.join('\n').length) {
191      out.push(line)
192      continue
193    }
194    out.push(line, note, lines[end]!)
195    i = end
196  }
197  return out.join('\n')
198}
199
200// One text block of a Figma tool's result, squeezed by what the tool is.
201export const figma: Pack = {
202  id: 'figma',
203  // `figma`, `claude_ai_Figma`, `figma-dev-mode`, a plugin's `plugin_figma_figma`.
204  server: /.*figma.*/,
205  legend: OUTLINE_LEGEND,
206  squeezeText: (tool, text) => {
207    const out = roundNoise(text)
208    if (tool.endsWith('__get_metadata')) return outline(out)
209    if (tool.endsWith('__get_design_context') && isCode(out)) return hoistClasses(collapseInstances(out).replace(GENERIC_NAME, ''))
210    return out
211  },
212}
213

1// What became of one tool's results this conversation: the calls, their text before and after in
2// characters, and the repeated images left out with the tokens they would have cost (absent in a
3// stat an earlier version kept).
4export type Stat = { calls: number; before: number; after: number; images?: number; imageTokens?: number }
5
6// What became of each tool's results since the mod was first on, across sessions, kept in the
7// store: `since` in milliseconds.
8export type Total = { since: number; tools: Record<string, Stat> }
9
10// One agent's blocks returned this conversation, by hash: `whole` its last copy went whole, so the
11// next one is left out; `pointed` it was left out, so the next goes whole.
12export type Seen = Record<string, 'whole' | 'pointed'>
13
14// What the packs hold of one agent's conversation, each under keys of its id (`chrome:snapshot`, the
15// last page snapshot as the model knows it).
16export type Memory = Record<string, string>
17
18declare module 'claude-code' {
19  interface PluginState {
20    'mcp-squeeze': {
21      seen: Record<string, Seen>
22      memory: Record<string, Memory>
23      stats: Record<string, Stat>
24    }
25  }
26}
27