firmware-budget

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Sample } from '../types'
5
6// Values a drawing reads, kept for the session and across a reload
7const samples = atom({ plugin: 'firmware-budget', key: 'samples' } as const, [] as Sample[])
8const fault = atom({ plugin: 'firmware-budget', key: 'fault' } as const, null as string | null)
9
10// What the manifest's userConfig resolves to, with defaults applied by the engine
11type Options = {
12  elf: string
13  flash_budget: number
14  ram_budget: number
15  ram_start: number
16  size_tool: string
17  build_pattern: string
18}
19
20const COMMIT = /\bgit\s+commit\b/
21
22// Parse SysV-format `size -A` output into flash and RAM use.
23//
24//   section          size        addr
25//   .text          143832   268435456
26//   .data            6968   536871184
27//   ...
28//
29// Flash is every section the linker placed below `ramStart` (the XIP image),
30// plus `.data`, whose initial image lives in flash and is copied to RAM at boot.
31// RAM is every section placed at or above `ramStart`, which on the Pico SDK
32// includes the vector table, .data, .bss and the SDK's minimum heap and stack
33// reservations. Debug sections have address 0 and count for neither.
34//
35// Berkeley `size` (no -A) is not used because the Pico SDK linker script marks
36// .data executable, and Berkeley size then folds it into the text column and
37// reports data 0, which under-reports RAM by the size of .data.
38export function parseSize(stdout: string, ramStart: number): { flash: number; ram: number } | null {
39  let flash = 0
40  let ram = 0
41  let data = 0
42  let rows = 0
43  for (const line of stdout.split('\n')) {
44    const cols = line.trim().split(/\s+/)
45    if (cols.length !== 3) continue
46    const name = cols[0] ?? ''
47    const size = Number(cols[1])
48    const addr = Number(cols[2])
49    if (!name.startsWith('.') || Number.isNaN(size) || Number.isNaN(addr)) continue
50    rows += 1
51    if (addr === 0) continue
52    if (addr >= ramStart) ram += size
53    else flash += size
54    if (name === '.data') data = size
55  }
56  if (rows === 0) return null
57  return { flash: flash + data, ram }
58}
59
60export function toSample(parsed: { flash: number; ram: number }, n: number, label: string): Sample {
61  return { n, flash: parsed.flash, ram: parsed.ram, label }
62}
63
64// A bar of `width` cells, filled in proportion to used / budget, capped at full
65export function bar(used: number, budget: number, width: number): string {
66  const fraction = budget > 0 ? Math.min(used / budget, 1) : 1
67  const filled = Math.round(fraction * width)
68  return '█'.repeat(filled) + '░'.repeat(width - filled)
69}
70
71export function percent(used: number, budget: number): string {
72  if (budget <= 0) return '?'
73  return ((used / budget) * 100).toFixed(1) + '%'
74}
75
76export function kb(bytes: number): string {
77  return (bytes / 1024).toFixed(1) + ' KB'
78}
79
80function overBudget(s: Sample, o: Options): string | null {
81  const reasons: string[] = []
82  if (s.flash > o.flash_budget) reasons.push(`flash ${s.flash} B is over the ${o.flash_budget} B budget by ${s.flash - o.flash_budget} B`)
83  if (s.ram > o.ram_budget) reasons.push(`RAM ${s.ram} B is over the ${o.ram_budget} B budget by ${s.ram - o.ram_budget} B`)
84  return reasons.length ? reasons.join('; ') : null
85}
86
87// Run the size tool on the ELF and record a sample. Returns what went wrong, or null.
88async function measure($: EngineInterface, o: Options, label: string): Promise<string | null> {
89  let ran: { exitCode: number; stdout: string; stderr: string }
90  try {
91    ran = await $.process.run([o.size_tool, '-A', o.elf], { timeoutMs: 15000 })
92  } catch (error) {
93    const why = `${o.size_tool} could not start: ${String(error)}`
94    await update($, fault, () => why)
95    return why
96  }
97  if (ran.exitCode !== 0) {
98    const why = `${o.size_tool} exited ${ran.exitCode}: ${ran.stderr.trim().split('\n')[0] ?? ''}`
99    await update($, fault, () => why)
100    return why
101  }
102  const parsed = parseSize(ran.stdout, o.ram_start)
103  if (parsed === null) {
104    const why = `${o.size_tool} -A output had no section rows`
105    await update($, fault, () => why)
106    return why
107  }
108  const history = await read($, samples)
109  const sample = toSample(parsed, history.length + 1, label)
110  await update($, samples, list => [...list, sample])
111  await update($, fault, () => null)
112  const over = overBudget(sample, o)
113  if (over !== null) $.ui.toast(`firmware-budget: ${over}`)
114  return null
115}
116
117function summary(history: Sample[], o: Options, lastFault: string | null): string {
118  if (history.length === 0) {
119    return lastFault ? `No measurement yet. Last attempt failed: ${lastFault}` : `No measurement yet. Build the firmware, or run /budget check.`
120  }
121  const rows = history.map((s, i) => {
122    const before = i > 0 ? history[i - 1] : undefined
123    const dFlash = before ? s.flash - before.flash : 0
124    const dRam = before ? s.ram - before.ram : 0
125    const sign = (d: number) => (d > 0 ? '+' : '') + d
126    return `${s.n}. flash ${s.flash} B (${sign(dFlash)})  ram ${s.ram} B (${sign(dRam)})  ${s.label}`
127  })
128  const last = history[history.length - 1]
129  if (last === undefined) return 'No measurement yet.'
130  const head = `flash ${last.flash} / ${o.flash_budget} B (${percent(last.flash, o.flash_budget)}), ram ${last.ram} / ${o.ram_budget} B (${percent(last.ram, o.ram_budget)})`
131  const over = overBudget(last, o)
132  return [head, over ? `OVER BUDGET: ${over}` : 'within budget', ...rows].join('\n')
133}
134
135// True when the command is a build: the pattern is tested against the command
136// with any leading path stripped from each word, so a full path to cmake or a
137// wrapper such as ninja-build still matches a pattern written for the bare name.
138export function isBuild(command: string, pattern: RegExp): boolean {
139  const bare = command.replace(/(^|\s)\S*\//g, '$1')
140  return pattern.test(bare)
141}
142
143export const register: Register = (on, options) => {
144  const o = options as unknown as Options
145  const build = new RegExp(o.build_pattern)
146
147  on('session.start', async ($, e, next) => {
148    const started = await next(e)
149    try {
150      await $.command.register({
151        name: 'budget',
152        description: 'Show flash and RAM use of the firmware ELF; "check" measures it now',
153        argumentHint: '[check]',
154      })
155    } catch {
156      // The name is taken; the band and the guard still work
157    }
158    return started
159  })
160
161  on('command.run', { command: 'budget' }, async ($, e) => {
162    if (e.args.trim() === 'check') {
163      const why = await measure($, o, '/budget check')
164      if (why !== null) return { text: why }
165    }
166    return { text: summary(await read($, samples), o, await read($, fault)) }
167  })
168
169  // One hook for Bash. Before the call: hold a commit while the last build is
170  // over budget, judged before next so a failure fails closed. After the call:
171  // if it was a build that succeeded, measure the ELF.
172  on('tool.call', { tool: 'Bash' }, async ($, e, next) => {
173    if (COMMIT.test(e.command)) {
174      const history = await read($, samples)
175      const last = history[history.length - 1]
176      const over = last === undefined ? null : overBudget(last, o)
177      if (over !== null) {
178        return {
179          deny: `firmware-budget refused this commit: ${over}. Reduce the firmware size, or ask the user to raise the budget in /config, then commit again.`,
180        }
181      }
182    }
183    const ran = await next(e)
184    if (!isBuild(e.command, build)) return ran
185    if (ran.deny !== undefined || ran.isError === true) return ran
186    const label = e.command.trim().replace(/\s+/g, ' ').slice(0, 40)
187    const why = await measure($, o, label)
188    if (why !== null) $.ui.log(why)
189    return ran
190  }).catch(($, e, next) =>
191    next.called ? next(e) : COMMIT.test(e.command) ? { deny: 'firmware-budget: its commit guard failed, so the commit was not run.' } : next(e),
192  )
193
194  // The band is shared with every other mod, so what the mods after this one
195  // draw is kept under the bars rather than replaced.
196  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
197    if (e.props.hasSurvey) return next(e)
198    const history = await read($, samples)
199    const lastFault = await read($, fault)
200    const { Box, Text } = $.ui.resolve(e)
201    const theirs = await next(e)
202    const withTheirs = (mine: typeof theirs) => Box({ flexDirection: 'column', children: [mine, theirs] })
203
204    if (history.length === 0) {
205      if (lastFault === null) return theirs
206      return withTheirs(Text({ dimColor: true, children: ['firmware-budget: ' + lastFault] }))
207    }
208
209    const last = history[history.length - 1]
210    if (last === undefined) return theirs
211    const prev = history.length > 1 ? history[history.length - 2] ?? null : null
212    const delta = (now: number, before: number | null) => (before === null ? '' : ' ' + (now - before >= 0 ? '+' : '') + (now - before) + ' B')
213    const line = (name: string, used: number, budget: number, before: number | null) => {
214      const over = used > budget
215      return Box({
216        flexDirection: 'row',
217        columnGap: 1,
218        children: [
219          Text({ bold: true, children: [name.padEnd(5)] }),
220          Text({ color: over ? 'red' : 'green', children: [bar(used, budget, 20)] }),
221          Text({ children: [`${percent(used, budget).padStart(6)}  ${kb(used)} / ${kb(budget)}`] }),
222          Text({ dimColor: true, children: [delta(used, before)] }),
223        ],
224      })
225    }
226    return withTheirs(
227      Box({
228        flexDirection: 'column',
229        children: [
230          line('flash', last.flash, o.flash_budget, prev ? prev.flash : null),
231          line('ram', last.ram, o.ram_budget, prev ? prev.ram : null),
232          Text({ dimColor: true, children: [`build ${last.n} · ${last.label}` + (lastFault ? ' · ' + lastFault : '')] }),
233        ],
234      }),
235    )
236  })
237}
238

1// One measurement of the ELF, taken after a build.
2export type Sample = {
3  // Build number within the session, from 1
4  n: number
5  // The XIP image plus the .data load image: what sits in flash
6  flash: number
7  // Every section the linker placed at a RAM address
8  ram: number
9  // The build command, trimmed to a short label
10  label: string
11}
12
13declare module 'claude-code' {
14  interface PluginState {
15    'firmware-budget': {
16      // Every sample this session, oldest first
17      samples: Sample[]
18      // The last size run that failed, or null
19      fault: string | null
20    }
21  }
22}
23