mem-guard

1import type { EngineInterface, Register } from 'claude-code'
2
3import { checkRunnerRefusal, crowdedRefusal, invocations, isHeavy, isServe, parsePressure, refusal, runningRefusal, scoped } from './rules'
4import type { Invocation, Pressure, ScriptsByDir } from './rules'
5
6let cmdSlice = ''
7let home = ''
8
9// Stays empty while claude-cmd.slice is inactive, so it is looked up again on the next call.
10async function resolveCmdSlice($: EngineInterface): Promise<string> {
11  if (cmdSlice === '') {
12    const group = await $.process.run(['systemctl', '--user', 'show', 'claude-cmd.slice', '-p', 'ControlGroup', '--value']).catch(() => null)
13    const path = group?.exitCode === 0 ? group.stdout.trim() : ''
14    cmdSlice = path === '' ? '' : `/sys/fs/cgroup${path}`
15  }
16  return cmdSlice
17}
18
19async function readPressure($: EngineInterface): Promise<Pressure | null> {
20  const slice = await resolveCmdSlice($)
21  if (slice === '') {
22    return null
23  }
24  const read = (file: string) => $.fs.read(`${slice}/${file}`)
25  const [current, max, pressure] = await Promise.all([read('memory.current'), read('memory.max'), read('memory.pressure')]).catch(() => ['', '', ''])
26  return parsePressure(current ?? '', max ?? '', pressure ?? '')
27}
28
29async function scriptsIn($: EngineInterface, dir: string): Promise<Record<string, string>> {
30  const manifest = await $.fs.read(`${dir}/package.json`).catch(() => '{}')
31  try {
32    return (JSON.parse(manifest) as { scripts?: Record<string, string> }).scripts ?? {}
33  } catch {
34    return {}
35  }
36}
37
38async function scriptsOf($: EngineInterface, calls: readonly Invocation[]): Promise<ScriptsByDir> {
39  const dirs = [...new Set(calls.map(call => call.dir))]
40  const scripts = await Promise.all(dirs.map(dir => scriptsIn($, dir)))
41  return Object.fromEntries(dirs.map((dir, i) => [dir, scripts[i] ?? {}]))
42}
43
44export const register: Register = on => {
45  on('session.start', async ($, e, next) => {
46    await resolveCmdSlice($)
47    home = (await $.env.get('HOME')) ?? ''
48    return next(e)
49  })
50
51  on('tool.call', { tool: 'Bash' }, async ($, e, next) => {
52    if (e.agentId !== undefined) {
53      const agents = await $.agent.list().catch(() => [])
54      if (agents.find(agent => agent.id === e.agentId)?.type === 'check-runner') {
55        const writeRefusal = checkRunnerRefusal(e.command)
56        if (writeRefusal !== null) {
57          return { deny: `mem-guard: ${writeRefusal}` }
58        }
59      }
60    }
61    const calls = invocations(e.command, { cwd: await $.session.cwd(), home })
62    const refused = refusal(calls, await scriptsOf($, calls))
63    if (refused !== null) {
64      return { deny: `mem-guard: ${refused}` }
65    }
66    if (isHeavy(calls) || isServe(calls)) {
67      const pressure = await readPressure($)
68      const crowded = pressure === null ? null : crowdedRefusal(pressure)
69      if (crowded !== null) {
70        return { deny: `mem-guard: ${crowded}` }
71      }
72      const ps = await $.process.run(['ps', '-eo', 'cgroup:250=,args=', '--cols', '600']).catch(() => null)
73      const busyRefusal = ps?.exitCode === 0 ? runningRefusal(calls, ps.stdout) : null
74      if (busyRefusal !== null) {
75        return { deny: `mem-guard: ${busyRefusal}` }
76      }
77    }
78    return next({ ...e, command: scoped(e.command) })
79  })
80}
81

1export type Pressure = { usedBytes: number; maxBytes: number; psiAvg10: number }
2export type Shell = { cwd: string; home: string }
3export type ScriptsByDir = Record<string, Record<string, string>>
4export type Invocation = { line: string; tool: string; call: string; script: string | null; dir: string }
5
6const GIB = 1024 ** 3
7const FULL_SHARE = 0.75
8const PSI_LIMIT = 20
9
10const HEAVY_TOOLS = ['nx', 'jest', 'vitest', 'playwright', 'tsc', 'ngc']
11const NODE_TOOLS = ['node', 'npx', 'pnpm', 'npm', 'yarn', ...HEAVY_TOOLS]
12const PREFIX = /^(?:(?:sudo|time|nice|env|exec|command|then|do|else|if|while|until|\{|!)\s+|timeout\s+\S+\s+|\w+=\S*\s+)*/
13const PNPM_FLAGS = String.raw`(?:(?:--filter|-F|-C|--dir)\s+\S+\s+|-[\w-]+(?:=\S+)?\s+)*`
14const RUNNER = new RegExp(String.raw`^(?:npx\s+(?:-\S+\s+)*|npm\s+(?:exec|x)\s+(?:-\S+\s+)*|pnpm\s+${PNPM_FLAGS}(?:exec\s+|dlx\s+|run\s+)?|yarn\s+(?:run\s+)?|node\s+(?=\S*node_modules\/\.bin\/)|)(?:\S*node_modules\/\.bin\/)?`)
15const SCRIPT = new RegExp(String.raw`^(?:pnpm\s+${PNPM_FLAGS}(?:run\s+)?|npm\s+run\s+|yarn\s+(?:run\s+)?)([\w:-]+)`)
16const SHELL_C = /^(?:ba)?sh\s+(?:-\w+\s+)*-c\s+(?:'([^']*)'|"((?:[^"\\]|\\.)*)")/
17const CD = /^cd(?:\s+(\S+))?$/
18const NX_LIGHT = /^nx\s+(show|graph|reset|daemon|report|list|--version)\b/
19const NX_TEST = /^nx\s+(?:test\b|run\s+\S+:test(?::\S+)?(?=\s|$))|\s(?:-t|--targets?)[= ]\s*(?:\S*,)?test(?=[,\s]|$)/
20const PARALLEL_ONE = /--parallel[= ]1([^0-9]|$)/
21const WORKER_CAP = /\s(?:--runInBand|-i)(?=\s|$)|\s(?:--maxWorkers|--max-workers|-w)(?:=|\s+)(?![^\s%]*%)\S/
22const JEST_NO_RUN = /\s(?:--listTests|--showConfig|--clearCache|--version|-v|--help|-h)(?=\s|$)/
23const KILL_BY_NAME = /\s(?:-[a-zA-Z]*f[a-zA-Z]*|--full)(?=\s|$)/
24const KILL_VALUE_FLAG = /^(?:-[GgPstuU]|--(?:signal|euid|uid|group|pgroup|parent|session|terminal|ns|nslist))$/
25const WORD = /(?:'[^']*'|"[^"]*"|[^\s'"])+/g
26
27// Splits on control operators and command substitutions outside quotes; a heredoc body is no command.
28export function segments(command: string): string[] {
29  const parts: string[] = []
30  const heredocs: string[] = []
31  let part = ''
32  let quote: '' | "'" | '"' = ''
33  const opened: Array<'' | '"'> = []
34  const cut = () => {
35    parts.push(part.trim())
36    part = ''
37  }
38  for (let i = 0; i < command.length; i++) {
39    const c = command[i] ?? ''
40    const two = command.slice(i, i + 2)
41    if (quote === "'") {
42      quote = c === "'" ? '' : quote
43      part += c
44    } else if (c === '\\') {
45      part += two
46      i++
47    } else if (two === '$(') {
48      cut()
49      opened.push(quote)
50      quote = ''
51      i++
52    } else if (c === '`') {
53      cut()
54    } else if (quote === '"') {
55      quote = c === '"' ? '' : quote
56      part += c
57    } else if (c === "'" || c === '"') {
58      quote = c
59      part += c
60    } else if (two === '<<') {
61      const tag = /^<<-?\s*(['"]?)(\w+)\1/.exec(command.slice(i))?.[2]
62      if (tag !== undefined) {
63        heredocs.push(tag)
64      }
65      part += two
66      i++
67    } else if (c === '\n' && heredocs.length > 0) {
68      cut()
69      for (const tag of heredocs.splice(0)) {
70        const end = new RegExp(String.raw`\n\t*${tag}(?=\n|$)`).exec(command.slice(i))
71        i = end === null ? command.length : i + end.index + end[0].length - 1
72      }
73    } else if (two === '&&' || two === '||') {
74      cut()
75      i++
76    } else if (c === '(') {
77      cut()
78      opened.push('')
79    } else if (c === ')') {
80      cut()
81      quote = opened.pop() ?? ''
82    } else if (';|&\n'.includes(c)) {
83      cut()
84    } else {
85      part += c
86    }
87  }
88  cut()
89  return parts.filter(p => p.length > 0)
90}
91
92function resolvePath(dir: string, home: string, target = '~'): string {
93  const raw = target.replace(/^(['"])(.*)\1$/, '$2')
94  if (raw === '-') {
95    return dir
96  }
97  const path = raw.startsWith('~') ? home + raw.slice(1) : raw.startsWith('/') ? raw : `${dir}/${raw}`
98  const parts: string[] = []
99  for (const part of path.split('/')) {
100    if (part === '..') {
101      parts.pop()
102    } else if (part !== '' && part !== '.') {
103      parts.push(part)
104    }
105  }
106  return `/${parts.join('/')}`
107}
108
109// `FOO=1 timeout 60 pnpm exec jest x` → line `pnpm exec jest x`, tool `jest`, call `jest x`; a `bash -c '…'` is opened up.
110export function invocations(command: string, shell: Shell): Invocation[] {
111  let dir = shell.cwd
112  return segments(command).flatMap(part => {
113    const line = part.replace(PREFIX, '').replace(/[\s}]+$/, '')
114    const cd = CD.exec(line)
115    if (cd !== null) {
116      dir = resolvePath(dir, shell.home, cd[1])
117    }
118    const inner = SHELL_C.exec(line)
119    const nested = inner === null ? [] : invocations(inner[1] ?? (inner[2] ?? '').replace(/\\(.)/g, '$1'), { ...shell, cwd: dir })
120    const call = line.replace(RUNNER, '')
121    return [{ line, tool: call.split(/\s/)[0] ?? '', call, script: SCRIPT.exec(line)?.[1] ?? null, dir }, ...nested]
122  })
123}
124
125function liteSibling(script: string, scripts: Record<string, string>): string | null {
126  if (script.endsWith(':lite')) {
127    return null
128  }
129  return scripts[`${script}:lite`] === undefined ? null : script
130}
131
132function words(line: string): string[] {
133  return line.match(WORD) ?? []
134}
135
136// The first operand that is neither a flag nor a flag's value.
137function killPattern(line: string): string {
138  const args = words(line).slice(1)
139  for (let i = 0; i < args.length; i++) {
140    const arg = args[i] ?? ''
141    if (KILL_VALUE_FLAG.test(arg)) {
142      i++
143    } else if (!arg.startsWith('-')) {
144      return arg
145    }
146  }
147  return ''
148}
149
150function invocationRefusal({ line, tool, call, script }: Invocation, scripts: Record<string, string>): string | null {
151  if (tool === 'nx' && /^nx\s+(affected|run-many)\b/.test(call) && !PARALLEL_ONE.test(call)) {
152    return "nx affected/run-many must run with --parallel=1. Prefer the repo's :lite script (pnpm run lint:affected:lite, typecheck:lite, test:affected:lite)."
153  }
154  const base = script === null ? null : liteSibling(script, scripts)
155  if (base !== null) {
156    return `use 'pnpm run ${base}:lite' instead of '${script}'.`
157  }
158  if (script === 'ci:local') {
159    return 'ci:local runs everything in sequence (~45 min here). Push and let the pipeline be the gate.'
160  }
161  if (/^(?:\S*\/)?p(kill|grep)\s/.test(line) && KILL_BY_NAME.test(line) && !killPattern(line).includes('[')) {
162    return "pkill -f / pgrep -f with a plain pattern matches its own wrapper and kills it (exit 143). Stop it by port (fuser -k 4700/tcp), by saved PID, or bracket the pattern: pkill -f '[n]x serve'."
163  }
164  const isJest = (tool === 'jest' && !JEST_NO_RUN.test(call)) || (tool === 'nx' && NX_TEST.test(call))
165  if (isJest && !WORKER_CAP.test(call)) {
166    return 'jest runs need --maxWorkers=2 (or --runInBand): its pool otherwise sizes itself off the CPU count, per task.'
167  }
168  return null
169}
170
171export function refusal(calls: readonly Invocation[], scripts: ScriptsByDir): string | null {
172  for (const call of calls) {
173    const refused = invocationRefusal(call, scripts[call.dir] ?? {})
174    if (refused !== null) {
175      return refused
176    }
177  }
178  return null
179}
180
181export function isHeavy(calls: readonly Invocation[]): boolean {
182  return calls.some(
183    ({ line, tool, call, script }) =>
184      (HEAVY_TOOLS.includes(tool) && !NX_LIGHT.test(call)) ||
185      (script !== null && /^(test|build|lint|typecheck|e2e|serve)/.test(script) && !script.endsWith(':lite')) ||
186      /^docker(-compose|\s+compose)\b.*\sup\b/.test(line),
187  )
188}
189
190export function crowdedRefusal(pressure: Pressure): string | null {
191  const used = pressure.usedBytes / GIB
192  const max = pressure.maxBytes / GIB
193  if (pressure.usedBytes <= pressure.maxBytes * FULL_SHARE && pressure.psiAvg10 <= PSI_LIMIT) {
194    return null
195  }
196  return `claude-cmd.slice is at ${used.toFixed(1)}/${max.toFixed(0)}G, pressure ${pressure.psiAvg10.toFixed(0)}%. Wait for a running build or test to finish, or push and let CI run it.`
197}
198
199function runsNode({ line, tool }: Invocation): boolean {
200  const program = (line.split(/\s/)[0] ?? '').replace(/^.*\//, '')
201  return NODE_TOOLS.includes(program) || NODE_TOOLS.includes(tool)
202}
203
204// Leading `cd`s stay outside the wrapper so the shell's working directory still moves.
205export function scoped(command: string): string {
206  const cds = /^(?:\s*cd(?:\s+(?:'[^']*'|"[^"]*"|[^\s;&|'"]+))?\s*(?:&&|;))*/.exec(command)?.[0] ?? ''
207  const rest = command.slice(cds.length).trimStart()
208  if (/^systemd-run/.test(rest) || !invocations(rest, { cwd: '/', home: '' }).some(runsNode)) {
209    return command
210  }
211  const quoted = `'${rest.replace(/'/g, `'\\''`)}'`
212  const wrapped = `systemd-run --user --scope -q --slice=claude-cmd.slice --expand-environment=no -p MemoryMax=30% -p MemorySwapMax=4% -- bash -c ${quoted}`
213  return cds === '' ? wrapped : `${cds.trim()} ${wrapped}`
214}
215
216export function parsePressure(current: string, max: string, pressure: string): Pressure | null {
217  const usedBytes = Number(current.trim())
218  const maxBytes = Number(max.trim())
219  const psiAvg10 = Number(/^some avg10=([\d.]+)/m.exec(pressure)?.[1] ?? NaN)
220  return [usedBytes, maxBytes, psiAvg10].every(Number.isFinite) ? { usedBytes, maxBytes, psiAvg10 } : null
221}
222
223const MAX_HEAVY_AT_ONCE = 2
224// A test runner, a build or an nx task; the nx daemon and a dev server stay out, they live for hours.
225const HEAVY_PROCESS =
226  /\/jest(?:-cli)?\/bin\/|jest-worker\/build\/workers|\/vitest\/|playwright\/(?:cli\.js|lib\/common\/process)|\/typescript\/(?:bin|lib)\/tsc\b|\/compiler-cli\/bundles\/src\/bin\/ngc|nx\/dist\/bin\/run-executor\.js/
227const SERVE = /\bnx(?:\.js)?\s+serve\b|:serve(?::\S+)?(?:\s|$)/
228
229/** Each scope in claude-cmd.slice with its processes' args, from `ps -eo cgroup=,args=`. */
230function scopesOf(ps: string): string[][] {
231  const scopes = new Map<string, string[]>()
232  for (const row of ps.split('\n')) {
233    const match = /\/claude-cmd\.slice\/(run-[^\s/]+\.scope)\S*\s+(.*)$/.exec(row)
234    if (match) {
235      scopes.set(match[1], [...(scopes.get(match[1]) ?? []), match[2]])
236    }
237  }
238  return [...scopes.values()]
239}
240
241/** Heavy commands running in claude-cmd.slice, one per scope. */
242export function heavyScopes(ps: string): number {
243  return scopesOf(ps).filter(args => !args.some(arg => SERVE.test(arg)) && args.some(arg => HEAVY_PROCESS.test(arg))).length
244}
245
246/** Dev servers running in claude-cmd.slice, one per scope. */
247export function serveScopes(ps: string): number {
248  return scopesOf(ps).filter(args => args.some(arg => SERVE.test(arg))).length
249}
250
251export function isServe(calls: readonly Invocation[]): boolean {
252  return calls.some(({ call, script }) => SERVE.test(call) || (script !== null && /^serve\b/.test(script)))
253}
254
255// The API and one app.
256const MAX_SERVES_AT_ONCE = 2
257
258export function serveRefusal(running: number): string | null {
259  return running < MAX_SERVES_AT_ONCE
260    ? null
261    : `${running} dev servers are already running (at most ${MAX_SERVES_AT_ONCE}: the API and one app). Stop one by its port first (fuser -k 4700/tcp).`
262}
263
264export function concurrencyRefusal(running: number): string | null {
265  return running < MAX_HEAVY_AT_ONCE
266    ? null
267    : `${running} heavy commands are already running (at most ${MAX_HEAVY_AT_ONCE} at once). Wait for one to finish.`
268}
269
270/** A dev server start counts against the dev servers in `ps`, anything else against the heavy commands. */
271export function runningRefusal(calls: readonly Invocation[], ps: string): string | null {
272  return isServe(calls) ? serveRefusal(serveScopes(ps)) : concurrencyRefusal(heavyScopes(ps))
273}
274
275const RUNNER_WRITES: ReadonlyArray<readonly [RegExp, string]> = [
276  [/(?:^|[\s;&|(])git\s+(?:-[Cc]\s+\S+\s+|-\S+\s+)*(?:checkout|restore|reset|stash|clean|commit|push|rebase|revert|apply|cherry-pick|merge|switch|rm|mv|add)\b/, 'changes the working tree or git history'],
277  [/(?:^|[\s;&|(])sed\s+(?:\S+\s+)*?(?:-[a-zA-Z]*i|--in-place)\b/, 'edits a file in place'],
278  [/\s--(?:write|fix)(?:[=\s]|$)/, 'rewrites files'],
279  [/(?:^|[\s;&|(])(?:systemctl\s+(?:--user\s+)?(?:start|stop|restart|kill|reset-failed)|systemd-run\s.*--unit|docker(?:\s+compose)?\s+(?:start|stop|restart|kill|rm|down|up)|pkill|killall|kill|fuser\s+-k)\b/, 'starts or stops a server or process'],
280  [/(?:^|[^<>&\d])>>?(?!>)\s*(?!\/tmp\/|\/dev\/)[^\s&]/, 'writes a file outside /tmp'],
281]
282const FILE_WRITERS = ['tee', 'rm', 'touch', 'cp', 'mv']
283
284/** Whether a tee, rm, touch, cp or mv names a path outside /tmp and /dev; a cp's source is only read. */
285function writesOutsideTmp(command: string): boolean {
286  return segments(command).some(part => {
287    const [path = '', ...args] = words(part.replace(PREFIX, ''))
288    const program = path.replace(/^.*\//, '')
289    if (!FILE_WRITERS.includes(program)) {
290      return false
291    }
292    const operands = args.filter((arg, i) => !/^-|^\d*[<>]/.test(arg) && !/^\d*[<>]+&?$/.test(args[i - 1] ?? ''))
293    return operands.slice(program === 'cp' ? 1 : 0).some(arg => !/^['"]?\/(?:tmp|dev)\//.test(arg))
294  })
295}
296
297/** A check-runner runs checks and reports; anything that changes files, git or servers is refused. */
298export function checkRunnerRefusal(command: string): string | null {
299  const reason = RUNNER_WRITES.find(([pattern]) => pattern.test(command))?.[1] ?? (writesOutsideTmp(command) ? 'writes a file outside /tmp' : null)
300  return reason === null ? null : `a check-runner only runs checks: this command ${reason}. Report the failure instead of fixing it.`
301}
302
303