@shogo-ai/worker 1.0.10

package/src/lib/resource-limits.ts

@@ -0,0 +1,340 @@
+// SPDX-License-Identifier: MIT
+// Copyright (C) 2026 Shogo Technologies, Inc.
+/**
+ * Cross-platform resource limits for host-spawned agent-runtime process groups.
+ *
+ * A single project runtime is a tree: `bun run agent-runtime` -> vite build
+ * --watch, `bun run server.tsx`, tsserver, pyright, expo, etc. Left unbounded
+ * its RSS climbs until the OS OOM-kills it (macOS jetsam SIGKILL is the
+ * canonical case), taking the whole app's memory with it. We cap each tree,
+ * using the strongest primitive each platform offers:
+ *
+ *   - Linux:   cgroup v2. Prefer a transient `systemd-run --scope` (no root,
+ *              hard `memory.max`/`cpu.max`/`pids.max`), fall back to writing a
+ *              cgroup directly, fall back to the RSS watchdog.
+ *   - Windows: a Job Object with ProcessMemoryLimit + KILL_ON_JOB_CLOSE via an
+ *              optional native helper; fall back to the RSS watchdog.
+ *   - macOS:   no cgroups and RLIMIT_AS is unreliable for JIT/mmap-heavy Bun,
+ *              so an RSS watchdog samples the process group and asks the
+ *              manager to restart on breach.
+ *
+ * The JS-heap env cap (see {@link applyHeapEnvCap}) is a secondary hint layered
+ * on top everywhere — OS-level limits are the real enforcement.
+ */
+import { spawnSync } from 'node:child_process';
+import { existsSync, mkdirSync, writeFileSync } from 'node:fs';
+import { join } from 'node:path';
+
+export interface RuntimeResourceLimits {
+  /** Per-project RAM ceiling in MB. Always > 0 when limits are active. */
+  memoryMB: number;
+  /** CPU ceiling as percent of ONE core (100 = one full core). 0 = no cap. */
+  cpuPercent: number;
+  /** Max tasks/pids in the group (cgroup pids.max). */
+  tasksMax: number;
+}
+
+type Logger = Pick<Console, 'log' | 'warn' | 'error'>;
+
+/** cgroup CPU period in microseconds (100ms), the systemd/cgroup default. */
+const CPU_PERIOD_US = 100_000;
+
+/**
+ * Resolve the active resource limits from the environment. Returns `null` when
+ * no ceiling is configured (`RUNTIME_MEMORY_MB` unset or <= 0), which disables
+ * all enforcement for backwards compatibility.
+ */
+export function resolveResourceLimits(env: NodeJS.ProcessEnv): RuntimeResourceLimits | null {
+  const mb = parseInt(env.RUNTIME_MEMORY_MB || '', 10);
+  if (!Number.isFinite(mb) || mb <= 0) return null;
+
+  const cpuRaw = parseInt(env.RUNTIME_CPU_PERCENT || '0', 10);
+  const cpuPercent = Number.isFinite(cpuRaw) && cpuRaw > 0 ? cpuRaw : 0;
+
+  const tasksRaw = parseInt(env.RUNTIME_TASKS_MAX || '', 10);
+  const tasksMax = Number.isFinite(tasksRaw) && tasksRaw > 0 ? tasksRaw : 4096;
+
+  return { memoryMB: mb, cpuPercent, tasksMax };
+}
+
+/**
+ * Append a JS-heap ceiling to `NODE_OPTIONS` (mutates `env` in place).
+ *
+ * This bounds the V8/JSC old-space of Node children (tsserver, vite, esbuild)
+ * and modern Bun, which honours `--max-old-space-size`. It does NOT bound
+ * native RSS, so it is only a hint — the OS-level limit is the real cap. We
+ * size it below the RSS ceiling (~75%) so the heap limit trips (recoverable
+ * JS OOM) before the native RSS watchdog / cgroup kill does.
+ */
+export function applyHeapEnvCap(env: NodeJS.ProcessEnv, memoryMB: number): void {
+  const heapMB = Math.max(256, Math.floor(memoryMB * 0.75));
+  const flag = `--max-old-space-size=${heapMB}`;
+  const existing = env.NODE_OPTIONS?.trim();
+  if (existing) {
+    if (/--max-old-space-size=/.test(existing)) return; // caller already set one
+    env.NODE_OPTIONS = `${existing} ${flag}`;
+  } else {
+    env.NODE_OPTIONS = flag;
+  }
+}
+
+// ─── Linux: cgroup v2 ────────────────────────────────────────────────
+
+/** True if `systemd-run` is usable for a rootless `--user --scope` unit. */
+export function hasSystemdRun(): boolean {
+  if (process.platform !== 'linux') return false;
+  try {
+    const r = spawnSync('systemd-run', ['--version'], { stdio: 'ignore', timeout: 3000 });
+    return r.status === 0;
+  } catch {
+    return false;
+  }
+}
+
+/**
+ * Wrap a spawn command in a transient rootless cgroup v2 scope via
+ * `systemd-run --user --scope`. The child (and everything it forks) runs inside
+ * a scope unit with hard `MemoryMax`/`CPUQuota`/`TasksMax`. Returns the original
+ * command untouched (`wrapped: false`) on non-Linux or when systemd-run is
+ * absent — the caller then falls back to {@link attachDirectCgroup} or the
+ * watchdog.
+ */
+export function wrapSpawnForCgroup(opts: {
+  command: string;
+  args: string[];
+  limits: RuntimeResourceLimits;
+  scopeName: string;
+  logger?: Logger;
+}): { command: string; args: string[]; wrapped: boolean } {
+  const { command, args, limits, scopeName, logger } = opts;
+  if (process.platform !== 'linux' || !hasSystemdRun()) {
+    return { command, args, wrapped: false };
+  }
+
+  const props = [
+    '-p', `MemoryMax=${limits.memoryMB}M`,
+    // MemoryHigh throttles + reclaims before the hard OOM at MemoryMax, so a
+    // spiky build is slowed rather than instantly killed.
+    '-p', `MemoryHigh=${Math.max(128, Math.floor(limits.memoryMB * 0.9))}M`,
+    '-p', `TasksMax=${limits.tasksMax}`,
+  ];
+  if (limits.cpuPercent > 0) props.push('-p', `CPUQuota=${limits.cpuPercent}%`);
+
+  const wrappedArgs = [
+    '--user',
+    '--scope',
+    '--quiet',
+    '--collect', // GC the transient unit as soon as it exits
+    `--unit=${scopeName}`,
+    ...props,
+    '--',
+    command,
+    ...args,
+  ];
+  logger?.log(`[resource-limits] wrapping spawn in systemd scope ${scopeName} (MemoryMax=${limits.memoryMB}M)`);
+  return { command: 'systemd-run', args: wrappedArgs, wrapped: true };
+}
+
+/**
+ * Fallback for Linux hosts without systemd: write a cgroup v2 directly and move
+ * `pid` (plus its future children) into it. Best-effort — returns `false` when
+ * cgroup v2 is unavailable or not writable (the common non-root, non-delegated
+ * case), so the caller can start the watchdog instead.
+ */
+export function attachDirectCgroup(opts: {
+  pid: number;
+  limits: RuntimeResourceLimits;
+  logger?: Logger;
+}): boolean {
+  const { pid, limits, logger } = opts;
+  if (process.platform !== 'linux') return false;
+
+  const root = '/sys/fs/cgroup';
+  // cgroup v2 unified hierarchy exposes cgroup.controllers at the root.
+  if (!existsSync(join(root, 'cgroup.controllers'))) return false;
+
+  const dir = join(root, 'shogo', `rt-${pid}`);
+  try {
+    mkdirSync(dir, { recursive: true });
+    writeFileSync(join(dir, 'memory.max'), String(limits.memoryMB * 1024 * 1024));
+    if (limits.cpuPercent > 0) {
+      const quota = Math.floor((limits.cpuPercent / 100) * CPU_PERIOD_US);
+      writeFileSync(join(dir, 'cpu.max'), `${quota} ${CPU_PERIOD_US}`);
+    }
+    try { writeFileSync(join(dir, 'pids.max'), String(limits.tasksMax)); } catch { /* controller may be absent */ }
+    // Moving the pid in also captures children forked after this point.
+    writeFileSync(join(dir, 'cgroup.procs'), String(pid));
+    logger?.log(`[resource-limits] attached pid ${pid} to cgroup ${dir} (memory.max=${limits.memoryMB}M)`);
+    return true;
+  } catch (err: any) {
+    logger?.warn(`[resource-limits] direct cgroup attach failed for pid ${pid}: ${err?.message ?? err}`);
+    return false;
+  }
+}
+
+// ─── Windows: Job Object (optional native helper) ────────────────────
+
+/**
+ * Assign `pid` to a Windows Job Object with a hard process-memory limit and
+ * KILL_ON_JOB_CLOSE, using an optional native helper. Node has no built-in Job
+ * Object API, so we load a helper lazily:
+ *
+ *   - a module path in `SHOGO_WINJOB_HELPER`, or
+ *   - a `@shogo/winjob` package if installed,
+ *
+ * exposing `limitProcess(pid, { memoryBytes, killOnClose }): boolean`.
+ *
+ * Returns `false` when no helper is available so the caller falls back to the
+ * RSS watchdog (which also enforces the ceiling on Windows, just via sampling).
+ */
+export function tryAttachJobObject(opts: {
+  pid: number;
+  limits: RuntimeResourceLimits;
+  logger?: Logger;
+}): boolean {
+  const { pid, limits, logger } = opts;
+  if (process.platform !== 'win32') return false;
+
+  const helperPath = process.env.SHOGO_WINJOB_HELPER || '@shogo/winjob';
+  try {
+    // eslint-disable-next-line @typescript-eslint/no-var-requires
+    const helper = require(helperPath) as {
+      limitProcess?: (pid: number, o: { memoryBytes: number; killOnClose: boolean }) => boolean;
+    };
+    if (typeof helper.limitProcess !== 'function') return false;
+    const ok = helper.limitProcess(pid, {
+      memoryBytes: limits.memoryMB * 1024 * 1024,
+      killOnClose: true,
+    });
+    if (ok) {
+      logger?.log(`[resource-limits] attached pid ${pid} to Job Object (${limits.memoryMB}M)`);
+    }
+    return ok;
+  } catch {
+    // Helper not installed — expected on stock installs. Watchdog covers us.
+    return false;
+  }
+}
+
+// ─── RSS watchdog (macOS always; Windows / Linux fallback) ───────────
+
+/**
+ * Sum the resident-set size (KB) of the process group led by `pid`.
+ *
+ *   - POSIX: `ps -o rss= -g <pgid>` (pid is the group leader because the
+ *     runtime is spawned `detached`), summed across the group.
+ *   - Windows: a PowerShell pass that sums WorkingSet64 of `pid` and all of its
+ *     descendants (walking Win32_Process.ParentProcessId).
+ *
+ * Returns `null` when sampling fails (process gone, tool missing) so the
+ * watchdog treats it as "no reading" rather than "0".
+ */
+export function samplePgidRssKb(pid: number, platform: NodeJS.Platform = process.platform): number | null {
+  try {
+    if (platform === 'win32') {
+      const script =
+        `$ErrorActionPreference='SilentlyContinue';` +
+        `$root=${pid};` +
+        `$all=Get-CimInstance Win32_Process;` +
+        `$seen=@{};$stack=New-Object System.Collections.Stack;$stack.Push($root);$sum=0;` +
+        `while($stack.Count -gt 0){$id=$stack.Pop();if($seen[$id]){continue};$seen[$id]=$true;` +
+        `foreach($p in $all){if($p.ProcessId -eq $id){$sum+=$p.WorkingSetSize};` +
+        `if($p.ParentProcessId -eq $id){$stack.Push($p.ProcessId)}}};` +
+        `[math]::Floor($sum/1024)`;
+      const r = spawnSync('powershell.exe', ['-NoProfile', '-NonInteractive', '-Command', script], {
+        encoding: 'utf-8',
+        timeout: 10_000,
+        windowsHide: true,
+      });
+      if (r.status !== 0 || !r.stdout) return null;
+      const kb = parseInt(r.stdout.trim(), 10);
+      return Number.isFinite(kb) ? kb : null;
+    }
+
+    // POSIX: sum RSS (KB) of every process in the group.
+    const r = spawnSync('ps', ['-o', 'rss=', '-g', String(pid)], {
+      encoding: 'utf-8',
+      timeout: 5_000,
+    });
+    if (r.status !== 0 || !r.stdout) return null;
+    let total = 0;
+    let sawRow = false;
+    for (const line of r.stdout.split('\n')) {
+      const v = parseInt(line.trim(), 10);
+      if (Number.isFinite(v)) { total += v; sawRow = true; }
+    }
+    return sawRow ? total : null;
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * Periodically samples a process group's RSS and invokes `onBreach` once the
+ * ceiling has been exceeded for {@link RssWatchdogOptions.strikes} consecutive
+ * samples (so a single transient spike doesn't trigger a restart). Used where
+ * no kernel-level cap is available.
+ */
+export interface RssWatchdogOptions {
+  pid: number;
+  ceilingMB: number;
+  platform?: NodeJS.Platform;
+  intervalMs?: number;
+  /** Consecutive over-ceiling samples required before onBreach fires. */
+  strikes?: number;
+  logger?: Logger;
+  /** Called with the observed RSS (MB) when the ceiling is breached. */
+  onBreach: (rssMB: number) => void;
+}
+
+export class RssWatchdog {
+  private timer: ReturnType<typeof setInterval> | null = null;
+  private consecutive = 0;
+  private fired = false;
+  private readonly pid: number;
+  private readonly ceilingMB: number;
+  private readonly platform: NodeJS.Platform;
+  private readonly intervalMs: number;
+  private readonly strikes: number;
+  private readonly logger?: Logger;
+  private readonly onBreach: (rssMB: number) => void;
+
+  constructor(opts: RssWatchdogOptions) {
+    this.pid = opts.pid;
+    this.ceilingMB = opts.ceilingMB;
+    this.platform = opts.platform ?? process.platform;
+    this.intervalMs = opts.intervalMs ?? parseInt(process.env.RUNTIME_RSS_WATCHDOG_MS || '15000', 10);
+    this.strikes = opts.strikes ?? parseInt(process.env.RUNTIME_RSS_WATCHDOG_STRIKES || '2', 10);
+    this.logger = opts.logger;
+    this.onBreach = opts.onBreach;
+  }
+
+  start(): void {
+    if (this.timer) return;
+    this.timer = setInterval(() => this.sample(), Math.max(2000, this.intervalMs));
+    try { this.timer.unref?.(); } catch { /* unref is best-effort */ }
+  }
+
+  stop(): void {
+    if (this.timer) { clearInterval(this.timer); this.timer = null; }
+  }
+
+  private sample(): void {
+    if (this.fired) return;
+    const kb = samplePgidRssKb(this.pid, this.platform);
+    if (kb == null) { this.consecutive = 0; return; }
+    const rssMB = Math.floor(kb / 1024);
+    if (rssMB <= this.ceilingMB) { this.consecutive = 0; return; }
+
+    this.consecutive += 1;
+    this.logger?.warn(
+      `[resource-limits] pid ${this.pid} RSS ${rssMB}MB over ceiling ${this.ceilingMB}MB ` +
+        `(strike ${this.consecutive}/${this.strikes})`,
+    );
+    if (this.consecutive >= this.strikes) {
+      this.fired = true;
+      this.stop();
+      this.onBreach(rssMB);
+    }
+  }
+}