heavy-job-guard

1import type { EngineInterface, Register } from 'claude-code'
2
3type Job = { pid: number; elapsed: string; command: string }
4
5const SHOWN_JOBS = 3
6const SHOWN_CHARS = 120
7
8/** Claude's Bash tool wraps each command as `zsh -c source … && eval '<command>' …`. */
9function unwrap(command: string): string {
10  const evaled = /\beval '((?:[^']|'\\'')*)'/.exec(command)
11
12  return (evaled?.[1] ?? command).replaceAll(`'\\''`, `'`)
13}
14
15function clip(text: string): string {
16  return text.length > SHOWN_CHARS ? `${text.slice(0, SHOWN_CHARS - 1)}…` : text
17}
18
19/** Seconds as `ps` prints an elapsed time: `mm:ss`, `hh:mm:ss` or `d-hh:mm:ss`. */
20export function etime(seconds: number): string {
21  const s = Math.max(0, Math.floor(seconds))
22  const two = (n: number) => String(n).padStart(2, '0')
23  const days = Math.floor(s / 86400)
24  const hours = Math.floor((s % 86400) / 3600)
25  const clock = `${two(Math.floor((s % 3600) / 60))}:${two(s % 60)}`
26
27  if (days > 0) {
28    return `${days}-${two(hours)}:${clock}`
29  }
30
31  return hours > 0 ? `${two(hours)}:${clock}` : clock
32}
33
34/**
35 * The machine's processes as `pid etime command` lines: `ps` on macOS and Linux, and
36 * Windows' process table through PowerShell where `ps` is missing or is not the BSD/procps one.
37 */
38async function processList($: EngineInterface): Promise<string | undefined> {
39  const ps = await $.process.run(['ps', '-axo', 'pid=,etime=,command=']).catch(() => undefined)
40
41  if (ps !== undefined && ps.exitCode === 0) {
42    return ps.stdout
43  }
44
45  const script =
46    'Get-CimInstance Win32_Process | Where-Object { $_.CommandLine } | ForEach-Object { ' +
47    "'{0} {1} {2}' -f $_.ProcessId, [int]((Get-Date) - $_.CreationDate).TotalSeconds, $_.CommandLine }"
48  const win = await $.process
49    .run(['powershell', '-NoProfile', '-NonInteractive', '-Command', script], { timeoutMs: 20000 })
50    .catch(() => undefined)
51
52  if (win === undefined || win.exitCode !== 0) {
53    $.ui.toast(`${$.plugin.name}: cannot list processes on this machine, so heavy jobs are not checked`)
54
55    return undefined
56  }
57
58  return win.stdout
59    .split(/\r?\n/)
60    .map(line => line.replace(/^(\d+) (\d+) /, (_, pid: string, secs: string) => `${pid} ${etime(Number(secs))} `))
61    .join('\n')
62}
63
64/** Every process on the machine whose command line matches `pattern` and not `ignore`, parents first. */
65export async function runningJobs(
66  $: EngineInterface,
67  pattern: RegExp,
68  ignore: RegExp,
69): Promise<Job[]> {
70  const listing = await processList($)
71
72  if (listing === undefined) {
73    return []
74  }
75
76  const jobs: Job[] = []
77
78  for (const line of listing.split('\n')) {
79    const match = /^\s*(\d+)\s+(\S+)\s+(.*)$/.exec(line)
80    const command = match?.[3] ?? ''
81
82    if (match === null || !pattern.test(command) || ignore.test(command)) {
83      continue
84    }
85
86    jobs.push({ pid: Number(match[1]), elapsed: match[2] ?? '', command: unwrap(command) })
87  }
88
89  return jobs.sort((a, b) => a.pid - b.pid)
90}
91
92export function describeJobs(jobs: readonly Job[]): string {
93  const shown = jobs
94    .slice(0, SHOWN_JOBS)
95    .map(job => `  pid ${job.pid} (running ${job.elapsed}): ${clip(job.command)}`)
96  const more = jobs.length > SHOWN_JOBS ? [`  …and ${jobs.length - SHOWN_JOBS} more`] : []
97
98  return [...shown, ...more].join('\n')
99}
100
101export const register: Register = (on, options) => {
102  const pattern = new RegExp(String(options.pattern))
103  const ignore = new RegExp(String(options.ignore))
104
105  on('tool.call', { tool: 'Bash' }, async ($, e, next) => {
106    if (!pattern.test(e.command)) {
107      return next(e)
108    }
109
110    if (options.subagents === 'deny' && e.agentId !== undefined) {
111      return {
112        deny:
113          `${$.plugin.name}: subagents don't run heavy jobs here (typecheck, tests, builds). ` +
114          'Finish your edits and report back; the main agent runs this one.',
115      }
116    }
117
118    const jobs = await runningJobs($, pattern, ignore)
119
120    if (jobs.length === 0) {
121      return next(e)
122    }
123
124    if (options.mode === 'warn') {
125      $.ui.toast(`${$.plugin.name}: starting a heavy job while ${jobs.length} already run`)
126
127      return next(e)
128    }
129
130    return {
131      deny: [
132        `${$.plugin.name}: a heavy job is already running on this machine, and two at once ` +
133          'compete for its memory and CPU:',
134        describeJobs(jobs),
135        'Wait for it to finish (watch it with Monitor, or come back to this later), then run ' +
136          'the command again. Do not reword the command to get past this check.',
137      ].join('\n'),
138    }
139  }).catch(($, e, next) => next(e))
140}
141