load

1import type { EngineInterface, On } from "claude-code";
2import { type BrowserCall, type Census, browserCalls, judge } from "./browser";
3import { type Machine, judgeBuilds, loadPerCore, planBuilds, runningBuilds } from "./builds";
4import { segments } from "./command";
5
6const MOD = "load";
7const MAIN = "main";
8
9const ENDED = new Set(["completed", "failed", "killed"]);
10const CONTINUING = new Set(["clear", "resume"]);
11
12// A background subagent stopped with TaskStop ends with no event of its own, so
13// its sessions are looked for on a timer while one is open.
14const SWEEP_MS = 30_000;
15
16interface State {
17  /** Session to the agent that launched it. */
18  owners: Map<string, string>;
19  pending: Set<string>;
20  /** Builds allowed to start whose call has not returned. */
21  building: string[];
22  sweep: { cancel: () => void } | undefined;
23}
24
25async function liveSessions($: EngineInterface): Promise<Set<string>> {
26  const { exitCode, stdout, stderr } = await $.process.run(
27    ["agent-browser", "session", "list", "--json"],
28    { timeoutMs: 5000 },
29  );
30  if (exitCode !== 0) throw new Error(`agent-browser session list exited ${exitCode}: ${stderr}`);
31  const parsed: unknown = JSON.parse(stdout);
32  const sessions =
33    typeof parsed === "object" &&
34    parsed !== null &&
35    "data" in parsed &&
36    typeof parsed.data === "object" &&
37    parsed.data !== null &&
38    "sessions" in parsed.data
39      ? parsed.data.sessions
40      : undefined;
41  if (!Array.isArray(sessions)) throw new Error(`unexpected session list: ${stdout}`);
42  return new Set(sessions.filter((s) => typeof s === "string"));
43}
44
45async function machine($: EngineInterface): Promise<Machine> {
46  const [ps, sysctl] = await Promise.all([
47    $.process.run(["ps", "-axo", "command="], { timeoutMs: 5000 }),
48    $.process.run(["sysctl", "-n", "vm.loadavg", "hw.ncpu"], { timeoutMs: 5000 }),
49  ]);
50  return { running: runningBuilds(ps.stdout), loadPerCore: loadPerCore(sysctl.stdout) };
51}
52
53function census(state: State, owner: string, isSubagent: boolean, live: Set<string>): Census {
54  const mine = new Set<string>();
55  for (const [session, by] of state.owners) {
56    if (by === owner && (live.has(session) || state.pending.has(session))) mine.add(session);
57  }
58  // The main loop may close every session.
59  const others = isSubagent ? new Set([...live].filter((s) => !mine.has(s))) : new Set<string>();
60  return { live, pending: state.pending, mine, others };
61}
62
63async function close(
64  $: EngineInterface,
65  state: State,
66  session: string,
67  trigger: string,
68  timeoutMs: number,
69): Promise<void> {
70  const owner = state.owners.get(session);
71  const agent = owner === MAIN ? null : (owner ?? null);
72  state.owners.delete(session);
73  const { exitCode, stderr } = await $.process
74    .run(["agent-browser", "--session", session, "close"], { timeoutMs })
75    .catch((error: unknown) => ({ exitCode: -1, stderr: String(error) }));
76  void $.modEvents.emit({
77    mod: MOD,
78    event: "browser.close",
79    ok: exitCode === 0,
80    detail: { session, agent, trigger, exitCode, stderr: stderr.trim() },
81  });
82}
83
84async function sweep($: EngineInterface, state: State): Promise<void> {
85  const owned = [...state.owners].filter(([, owner]) => owner !== MAIN);
86  if (owned.length === 0) {
87    state.sweep?.cancel();
88    state.sweep = undefined;
89    return;
90  }
91  state.sweep ??= $.clock.every(SWEEP_MS, () => {
92    sweep($, state).catch((error: unknown) => {
93      failed($, "sweep", "threw", String(error));
94    });
95  });
96  const agents = await $.agent.list();
97  // A workflow's agents are absent from the list, so only a listed agent that
98  // ended has its sessions closed. The rest wait for the idle timeout.
99  const ended = new Set(agents.filter((a) => ENDED.has(a.status)).map((a) => a.id));
100  await Promise.all(
101    owned
102      .filter(([session, owner]) => ended.has(owner) && !state.pending.has(session))
103      .map(([session]) => close($, state, session, "agent-end", 5000)),
104  );
105}
106
107function record(state: State, calls: BrowserCall[], owner: string, live: Set<string>): string[] {
108  const launched: string[] = [];
109  for (const call of calls) {
110    if (call.effect !== "launch" || live.has(call.session) || state.owners.has(call.session)) {
111      continue;
112    }
113    state.owners.set(call.session, owner);
114    state.pending.add(call.session);
115    launched.push(call.session);
116  }
117  return launched;
118}
119
120function forget(state: State, calls: BrowserCall[]): void {
121  for (const call of calls) {
122    if (call.effect === "close") state.owners.delete(call.session);
123    if (call.effect === "close-all") state.owners.clear();
124  }
125}
126
127function failed($: EngineInterface, hook: string, kind: string, message?: string): void {
128  $.modEvents
129    .emit({
130      mod: MOD,
131      event: "error",
132      ok: false,
133      detail: { hook, kind, message: message ?? null },
134    })
135    .catch(() => undefined);
136}
137
138export function register(on: On): void {
139  const state: State = { owners: new Map(), pending: new Set(), building: [], sweep: undefined };
140
141  on("session.start", ($, e, next) => {
142    void $.modEvents.emit({ mod: MOD, event: "session.start" });
143    return next(e);
144  });
145
146  on("tool.call", { tool: "Bash" }, async ($, e, next) => {
147    const parts = segments(e.command);
148    const calls = browserCalls(parts);
149    const plan = planBuilds(parts);
150    if (calls.length === 0 && plan.kinds.length === 0) return next(e);
151
152    const owner = e.agentId ?? MAIN;
153    const agent = e.agentId ?? null;
154    let live = new Set<string>();
155    if (calls.length > 0) {
156      live = await liveSessions($);
157      const verdict = judge(calls, census(state, owner, e.agentId !== undefined, live));
158      if (verdict.deny !== undefined) {
159        void $.modEvents.emit({
160          mod: MOD,
161          event: "browser.deny",
162          detail: {
163            reason: verdict.reason,
164            sessions: calls.map((c) => c.session),
165            live: live.size,
166            agent,
167          },
168        });
169        return { deny: verdict.deny };
170      }
171    }
172
173    if (plan.kinds.length > 0) {
174      const probed = await machine($);
175      const now = { ...probed, running: [...probed.running, ...state.building] };
176      const verdict = judgeBuilds(plan, now);
177      void $.modEvents.emit({
178        mod: MOD,
179        event: verdict.deny === undefined ? "build.start" : "build.deny",
180        detail: {
181          reason: verdict.reason ?? null,
182          kinds: plan.kinds,
183          running: now.running,
184          loadPerCore: now.loadPerCore ?? null,
185          agent,
186        },
187      });
188      if (verdict.deny !== undefined) return { deny: verdict.deny };
189    }
190
191    state.building.push(...plan.kinds);
192    const launched = record(state, calls, owner, live);
193    for (const session of launched) {
194      void $.modEvents.emit({
195        mod: MOD,
196        event: "browser.launch",
197        detail: { session, live: live.size, agent },
198      });
199    }
200    try {
201      return await next(e);
202    } finally {
203      for (const session of launched) state.pending.delete(session);
204      for (const kind of plan.kinds) state.building.splice(state.building.indexOf(kind), 1);
205      forget(state, calls);
206      if (launched.length > 0) await sweep($, state);
207    }
208    // Fails open: refusing on a broken probe would block every matching call.
209  }).catch(($, e, next) => {
210    failed($, "tool.call", next.error.kind, next.error.message);
211    return next(e);
212  });
213
214  on("turn.complete", async ($, e, next) => {
215    await sweep($, state);
216    return next(e);
217  }).catch(($, e, next) => {
218    failed($, "turn.complete", next.error.kind, next.error.message);
219    return next(e);
220  });
221
222  on("session.end", async ($, e, next) => {
223    if (CONTINUING.has(e.reason)) return next(e);
224    state.sweep?.cancel();
225    state.sweep = undefined;
226    // session.end shares 1.5s across hooks. The idle timeout catches stragglers.
227    await Promise.all([...state.owners.keys()].map((s) => close($, state, s, "session-end", 1000)));
228    return next(e);
229  }).catch(($, e, next) => {
230    failed($, "session.end", next.error.kind, next.error.message);
231    return next(e);
232  });
233}
234

1import { type Segment, splitPrefix } from "./command";
2
3export type Effect = "launch" | "close" | "close-all" | "none";
4
5export interface BrowserCall {
6  session: string;
7  effect: Effect;
8}
9
10const RUNNERS = new Set(["npx", "bunx", "pnpx"]);
11
12// Options whose value is a separate word.
13const VALUED = new Set([
14  "--session",
15  "--namespace",
16  "--profile",
17  "--session-name",
18  "--state",
19  "--headers",
20  "--executable-path",
21  "--extension",
22  "--init-script",
23  "--enable",
24  "--args",
25  "--user-agent",
26  "--proxy",
27  "--proxy-bypass",
28  "--ca-cert",
29  "--cdp",
30  "-p",
31  "--provider",
32  "--device",
33  "--screenshot-dir",
34  "--screenshot-quality",
35  "--screenshot-format",
36  "--input-mode",
37  "--idle-timeout",
38  "--restore-save",
39  "--model",
40]);
41
42const INERT = new Set([
43  "session",
44  "skills",
45  "help",
46  "install",
47  "upgrade",
48  "doctor",
49  "dashboard",
50  "profiles",
51  "auth",
52  "plugin",
53  "mcp",
54]);
55
56const INERT_FLAGS = new Set(["--help", "-h", "--version", "-V"]);
57
58function isAgentBrowser(word: string | undefined): boolean {
59  return word !== undefined && (word === "agent-browser" || word.endsWith("/agent-browser"));
60}
61
62function parse(argv: string[], envSession: string | undefined): BrowserCall {
63  let session = envSession;
64  let subcommand: string | undefined;
65  const rest: string[] = [];
66  for (let i = 0; i < argv.length; i++) {
67    const word = argv[i] ?? "";
68    if (word.startsWith("--session=")) {
69      session = word.slice("--session=".length);
70    } else if (word === "--session") {
71      session = argv[i + 1];
72      i++;
73    } else if (VALUED.has(word)) {
74      i++;
75    } else if (subcommand === undefined && !word.startsWith("-")) {
76      subcommand = word;
77    } else {
78      rest.push(word);
79    }
80  }
81  const resolved = session === undefined || session === "" ? "default" : session;
82  if (subcommand === "close") {
83    return { session: resolved, effect: rest.includes("--all") ? "close-all" : "close" };
84  }
85  const isInert =
86    subcommand === undefined || INERT.has(subcommand) || rest.some((word) => INERT_FLAGS.has(word));
87  return { session: resolved, effect: isInert ? "none" : "launch" };
88}
89
90export function browserCalls(segments: Segment[]): BrowserCall[] {
91  const calls: BrowserCall[] = [];
92  for (const { words } of segments) {
93    const { env, argv } = splitPrefix(words);
94    let start = 0;
95    if (RUNNERS.has(argv[0] ?? "")) {
96      start = 1;
97      while ((argv[start] ?? "").startsWith("-")) start++;
98    }
99    if (!isAgentBrowser(argv[start])) continue;
100    calls.push(parse(argv.slice(start + 1), env.get("AGENT_BROWSER_SESSION")));
101  }
102  return calls;
103}
104
105export interface Verdict {
106  deny?: string;
107  reason?: "owned" | "cap" | "close-all" | "dynamic";
108}
109
110export interface Census {
111  live: Set<string>;
112  /** Allowed to launch, call not yet returned. */
113  pending: Set<string>;
114  mine: Set<string>;
115  others: Set<string>;
116}
117
118export const MAX_DAEMONS = 3;
119
120const list = (sessions: Iterable<string>) => [...sessions].map((s) => `\`${s}\``).join(", ");
121
122export function judge(calls: BrowserCall[], census: Census): Verdict {
123  if (calls.some((call) => call.effect === "close-all") && census.others.size > 0) {
124    return {
125      reason: "close-all",
126      deny: `\`agent-browser close --all\` would close browser sessions other agents are using (${list(census.others)}). Close only your own: \`agent-browser --session <name> close\`.`,
127    };
128  }
129  const dynamic = calls.find((call) => call.effect === "launch" && call.session.includes("$"));
130  if (dynamic !== undefined) {
131    return {
132      reason: "dynamic",
133      deny: `Browser session \`${dynamic.session}\` is named by a shell expansion, so each run can start another browser. Name one session for this agent literally and navigate it from item to item.`,
134    };
135  }
136  const launching = new Set(
137    calls
138      .filter((call) => call.effect === "launch" && !census.live.has(call.session))
139      .map((call) => call.session),
140  );
141  const mine = new Set(census.mine);
142  const pending = new Set(census.pending);
143  for (const session of launching) {
144    if (pending.has(session)) continue;
145    const held = [...mine].filter((owned) => owned !== session);
146    if (held.length > 0) {
147      return {
148        reason: "owned",
149        deny: `You already have browser session ${list(held)} open. Reuse it for this item with \`--session ${held[0]}\` (navigate it to the next page), or run \`agent-browser --session ${held[0]} close\` before starting \`${session}\`. Keep one browser session per agent.`,
150      };
151    }
152    const running = new Set([...census.live, ...pending]);
153    if (running.size >= MAX_DAEMONS) {
154      return {
155        reason: "cap",
156        deny: `${running.size} browser sessions are already running on this machine (${list(running)}), the limit is ${MAX_DAEMONS}. Reuse one you own, close one you are done with, or finish this item without a browser and retry once another agent finishes.`,
157      };
158    }
159    mine.add(session);
160    pending.add(session);
161  }
162  return {};
163}
164

1import { type Segment, splitPrefix } from "./command";
2
3const HEAVY: { kind: string; pattern: RegExp }[] = [
4  { kind: "ladle", pattern: /\bladle(?:\.js)? build\b/ },
5  { kind: "vite", pattern: /\bvite(?:\.js)? build\b/ },
6  { kind: "storybook", pattern: /\bstorybook(?:\.js)? build\b|\bbuild-storybook\b/ },
7  {
8    kind: "pytest-xdist",
9    pattern: /\bpytest\b.* (?:-n ?|--numprocesses[ =])(?:auto|logical|\d+)\b/,
10  },
11];
12
13const PARALLEL_RUNNERS = /^(?:parallel|xargs .*-P ?(?:0|[2-9]|\d{2,})\b)/;
14
15// Programs that start a build. Any other program only mentions one.
16const LAUNCHERS =
17  /^(?:npx|bunx|pnpx|pnpm|npm|yarn|bun|node|deno|uv|poetry|xargs|parallel|ladle|vite|storybook|build-storybook|pytest|python[\d.]*)$/;
18
19const basename = (word: string | undefined) => word?.split("/").at(-1) ?? "";
20
21function kindOf(line: string): string | undefined {
22  return HEAVY.find(({ pattern }) => pattern.test(line))?.kind;
23}
24
25export interface BuildPlan {
26  kinds: string[];
27  isParallel: boolean;
28}
29
30const LOOPS = new Set(["for", "while", "until"]);
31
32export function planBuilds(segments: Segment[]): BuildPlan {
33  const kinds: string[] = [];
34  let backgrounded = 0;
35  let fanned = false;
36  const isLoop = segments.some(({ words }) => LOOPS.has(words[0] ?? ""));
37  for (const { words, isBackground } of segments) {
38    const { argv } = splitPrefix(words);
39    if (!LAUNCHERS.test(basename(argv[0]))) continue;
40    const line = argv.join(" ");
41    const kind = kindOf(line);
42    if (kind === undefined) continue;
43    kinds.push(kind);
44    if (isBackground) backgrounded++;
45    if (PARALLEL_RUNNERS.test(line)) fanned = true;
46  }
47  // A backgrounded build inside a loop starts one build per iteration at once.
48  const isParallel = fanned || (backgrounded > 0 && (kinds.length > 1 || isLoop));
49  return { kinds, isParallel };
50}
51
52// Wrappers (`npm exec`, `bun x`, the Bash tool's `zsh -c`) repeat a build's words.
53const WORKER = /^(?:node|bun|deno|ladle|vite|storybook|build-storybook|pytest|python[\d.]*)$/;
54const RUNNER_STEPS = new Set(["x", "exec", "dlx", "run", "npx", "npx-cli.js", "npm-cli.js"]);
55
56export function runningBuilds(ps: string): string[] {
57  const kinds: string[] = [];
58  for (const line of ps.split("\n")) {
59    const [program, step] = line.trim().split(/\s+/).map(basename);
60    if (!WORKER.test(program ?? "") || RUNNER_STEPS.has(step ?? "")) continue;
61    const kind = kindOf(line);
62    if (kind !== undefined) kinds.push(kind);
63  }
64  return kinds;
65}
66
67/** Parses `sysctl -n vm.loadavg hw.ncpu`. */
68export function loadPerCore(sysctl: string): number | undefined {
69  const match = /\{\s*([\d.]+)[^}]*\}\s+(\d+)/.exec(sysctl);
70  if (match === null) return undefined;
71  const cores = Number(match[2]);
72  return cores > 0 ? Number(match[1]) / cores : undefined;
73}
74
75export const MAX_BUILDS = 2;
76export const MAX_LOAD_PER_CORE = 3;
77
78export interface Machine {
79  running: string[];
80  loadPerCore: number | undefined;
81}
82
83export interface BuildVerdict {
84  deny?: string;
85  reason?: "parallel" | "running" | "load";
86}
87
88export function judgeBuilds(plan: BuildPlan, machine: Machine): BuildVerdict {
89  if (plan.kinds.length === 0) return {};
90  if (plan.isParallel) {
91    return {
92      reason: "parallel",
93      deny: `This command starts ${plan.kinds.length} heavy builds side by side (${plan.kinds.join(", ")}). Run them one at a time, each to completion, in separate commands or a sequential loop without \`&\`.`,
94    };
95  }
96  if (machine.running.length >= MAX_BUILDS) {
97    return {
98      reason: "running",
99      deny: `${machine.running.length} heavy builds are already running on this machine (${machine.running.join(", ")}), the limit is ${MAX_BUILDS}. Wait for one to finish before starting this one.`,
100    };
101  }
102  if (machine.loadPerCore !== undefined && machine.loadPerCore > MAX_LOAD_PER_CORE) {
103    return {
104      reason: "load",
105      deny: `The machine is saturated (load ${machine.loadPerCore.toFixed(1)}x its core count). Wait for running work to finish before starting a ${plan.kinds.join(", ")} build.`,
106    };
107  }
108  return {};
109}
110

1export interface Segment {
2  words: string[];
3  isBackground: boolean;
4}
5
6const SEPARATORS = new Set([";", "|", "\n", "(", ")", "`"]);
7
8/** Splits a command line into simple commands. `2>&1` and `&>` stay inside one. */
9export function segments(command: string): Segment[] {
10  const out: Segment[] = [];
11  let words: string[] = [];
12  let word: string | undefined;
13
14  const endWord = () => {
15    if (word !== undefined) words.push(word);
16    word = undefined;
17  };
18  const endSegment = (isBackground: boolean) => {
19    endWord();
20    if (words.length > 0) out.push({ words, isBackground });
21    words = [];
22  };
23
24  for (let i = 0; i < command.length; i++) {
25    const char = command.charAt(i);
26    if (char === "'" || char === '"') {
27      const end = command.indexOf(char, i + 1);
28      const stop = end === -1 ? command.length : end;
29      word = (word ?? "") + command.slice(i + 1, stop);
30      i = stop;
31    } else if (char === "\\") {
32      word = (word ?? "") + command.charAt(i + 1);
33      i++;
34    } else if (char === "&" && (command[i - 1] === ">" || command[i + 1] === ">")) {
35      word = (word ?? "") + char;
36    } else if (char === "&" && command[i + 1] === "&") {
37      endSegment(false);
38      i++;
39    } else if (char === "&") {
40      endSegment(true);
41    } else if (SEPARATORS.has(char)) {
42      endSegment(false);
43    } else if (/\s/.test(char)) {
44      endWord();
45    } else {
46      word = (word ?? "") + char;
47    }
48  }
49  endSegment(false);
50  return out;
51}
52
53const DURATION = /^\d+(?:\.\d+)?[smhd]?$/;
54
55const ASSIGNMENT = /^([A-Za-z_][A-Za-z0-9_]*)=(.*)$/s;
56
57const PREFIXES = new Set([
58  "do",
59  "then",
60  "else",
61  "{",
62  "!",
63  "time",
64  "nohup",
65  "exec",
66  "command",
67  "env",
68  "nice",
69]);
70
71export function splitPrefix(words: string[]): { env: Map<string, string>; argv: string[] } {
72  const env = new Map<string, string>();
73  let i = 0;
74  for (; i < words.length; i++) {
75    const word = words[i] ?? "";
76    if (PREFIXES.has(word)) continue;
77    if (word === "timeout") {
78      while (i + 1 < words.length && !DURATION.test(words[i + 1] ?? "")) i++;
79      i++;
80      continue;
81    }
82    const match = ASSIGNMENT.exec(word);
83    if (match === null) break;
84    env.set(match[1] ?? "", match[2] ?? "");
85  }
86  return { env, argv: words.slice(i) };
87}
88