omnirush 0.8.6 → 0.9.1

package/assets/extensions/omnirush/memory-lib.ts

@@ -0,0 +1,516 @@
+// memory-lib — memory-aware admission of sub-agents (spawn_agents).
+//
+// Every sub-agent is a separate agent process (its own runtime and engine,
+// ~150-250 MB each), and spawn_agents has no count cap: 20 or 40 children at
+// once used to start together whatever the machine had left, and machines
+// ran out of memory. Admission keeps the "no fixed cap" rule but starts a
+// child only while the memory allows it:
+//
+//   - the system keeps at least MIN_FREE of its RAM available after the
+//     child's expected footprint (default 20% of RAM; inside a cgroup with a
+//     memory limit — a container, a systemd scope — the limit counts as RAM
+//     and what it has left as available);
+//   - this agent's whole process tree (the top agent, every sub-agent at
+//     every nesting layer, their tools) stays under BUDGET (default 50% of
+//     RAM);
+//   - critically low (available under CRITICAL_FREE, default 10% of RAM):
+//     nothing new starts, not even a first child, and the user is warned;
+//     running children go on. Otherwise one child may always run, so a
+//     batch makes progress on a small machine.
+//
+// The rest wait in the queue ("queued: waiting for memory") and start as
+// memory frees up (a child finishing, or the next check, every second).
+// Thresholds: OMNIRUSH_AGENT_MIN_FREE, OMNIRUSH_AGENT_MEMORY_BUDGET,
+// OMNIRUSH_AGENT_CRITICAL_FREE and OMNIRUSH_AGENT_CHILD_ESTIMATE ("20%",
+// "1.5G", "800M", a plain number is MiB), OMNIRUSH_AGENT_MEMORY=off turns
+// admission off; the same keys (minFree, budget, criticalFree,
+// childEstimate, enabled) under "subagentMemory" in <omnirush
+// dir>/settings.json. The environment wins.
+//
+// Measuring: Linux reads /proc/meminfo (MemAvailable), the cgroup's
+// memory.max / memory.current / memory.stat, and /proc/<pid>/statm of the
+// process tree; macOS asks vm_stat and ps; elsewhere os.freemem() and an
+// estimate per running child stand in for the tree.
+
+import { execFileSync } from "node:child_process";
+import { readdirSync, readFileSync } from "node:fs";
+import os from "node:os";
+import path from "node:path";
+
+const MIB = 1024 * 1024;
+
+/** Environment variable the top agent hands down: the pid whose process tree the budget covers. */
+export const ENV_MEMORY_ROOT = "OMNIRUSH_AGENT_MEMORY_ROOT";
+
+export interface MemoryPolicy {
+  enabled: boolean;
+  /** Keep at least this much of RAM available (bytes, resolved). */
+  minFreeBytes: number;
+  /** The whole process tree stays under this (bytes, resolved). */
+  budgetBytes: number;
+  /** Below this nothing new starts and the user is warned (bytes, resolved). */
+  criticalFreeBytes: number;
+  /** Footprint expected of a child before it is measured (bytes). */
+  childEstimateBytes: number;
+  /** Where each value came from ("env", "settings", "default"), for /agents memory. */
+  sources: Record<"minFree" | "budget" | "criticalFree" | "childEstimate" | "enabled", "env" | "settings" | "default">;
+}
+
+export const DEFAULT_MIN_FREE = "20%";
+export const DEFAULT_BUDGET = "50%";
+export const DEFAULT_CRITICAL_FREE = "10%";
+export const DEFAULT_CHILD_ESTIMATE = "256M";
+
+/**
+ * "20%" of `total`, "1.5G", "800M", "512K", or a plain number of MiB; null
+ * when unreadable. Percentages are clamped to 0-100.
+ */
+export function parseAmount(raw: unknown, total: number): number | null {
+  if (typeof raw === "number" && Number.isFinite(raw) && raw >= 0) return Math.round(raw * MIB);
+  if (typeof raw !== "string") return null;
+  const text = raw.trim().toLowerCase();
+  if (!text) return null;
+  const match = /^(\d+(?:\.\d+)?)\s*(%|b|k|kb|kib|m|mb|mib|g|gb|gib|t|tb|tib)?$/.exec(text);
+  if (!match) return null;
+  const value = Number(match[1]);
+  const unit = match[2] ?? "m";
+  if (unit === "%") return Math.round((Math.min(100, value) / 100) * total);
+  const scale = unit === "b" ? 1 : unit.startsWith("k") ? 1024 : unit.startsWith("m") ? MIB : unit.startsWith("g") ? 1024 * MIB : 1024 * 1024 * MIB;
+  return Math.round(value * scale);
+}
+
+/** The "subagentMemory" block of <dir>/settings.json (empty when absent). */
+export function readMemorySettings(dir: string): Record<string, unknown> {
+  try {
+    const parsed = JSON.parse(readFileSync(path.join(dir, "settings.json"), "utf8"));
+    const block = parsed?.subagentMemory;
+    return block && typeof block === "object" && !Array.isArray(block) ? block : {};
+  } catch {
+    return {};
+  }
+}
+
+/** Resolve the thresholds for a machine with `totalBytes` of RAM (env over settings over defaults). */
+export function resolveMemoryPolicy(totalBytes: number, env: NodeJS.ProcessEnv = process.env, settings: Record<string, unknown> = {}): MemoryPolicy {
+  const sources = {} as MemoryPolicy["sources"];
+  const pick = (key: keyof MemoryPolicy["sources"], envName: string, fallback: string): number => {
+    const fromEnv = parseAmount(env[envName], totalBytes);
+    if (fromEnv !== null) {
+      sources[key] = "env";
+      return fromEnv;
+    }
+    const fromSettings = parseAmount(settings[key], totalBytes);
+    if (fromSettings !== null) {
+      sources[key] = "settings";
+      return fromSettings;
+    }
+    sources[key] = "default";
+    return parseAmount(fallback, totalBytes) ?? 0;
+  };
+  const rawEnabled = String(env.OMNIRUSH_AGENT_MEMORY ?? "").trim().toLowerCase();
+  let enabled = true;
+  if (rawEnabled) {
+    sources.enabled = "env";
+    enabled = !["off", "0", "false", "no", "disabled"].includes(rawEnabled);
+  } else if (typeof settings.enabled === "boolean") {
+    sources.enabled = "settings";
+    enabled = settings.enabled;
+  } else {
+    sources.enabled = "default";
+  }
+  const minFreeBytes = pick("minFree", "OMNIRUSH_AGENT_MIN_FREE", DEFAULT_MIN_FREE);
+  const budgetBytes = pick("budget", "OMNIRUSH_AGENT_MEMORY_BUDGET", DEFAULT_BUDGET);
+  const criticalFreeBytes = Math.min(pick("criticalFree", "OMNIRUSH_AGENT_CRITICAL_FREE", DEFAULT_CRITICAL_FREE), minFreeBytes);
+  const childEstimateBytes = Math.max(32 * MIB, pick("childEstimate", "OMNIRUSH_AGENT_CHILD_ESTIMATE", DEFAULT_CHILD_ESTIMATE));
+  return { enabled, minFreeBytes, budgetBytes, criticalFreeBytes, childEstimateBytes, sources };
+}
+
+// --- measuring ---------------------------------------------------------------
+
+export interface SystemMemory {
+  /** RAM, or the cgroup's memory limit when that is smaller. */
+  totalBytes: number;
+  /** What can still be used before the system (or the cgroup) runs out. */
+  availableBytes: number;
+  source: string;
+}
+
+function readText(file: string): string | null {
+  try {
+    return readFileSync(file, "utf8");
+  } catch {
+    return null;
+  }
+}
+
+/** MemTotal / MemAvailable from /proc/meminfo (bytes), or null. */
+export function parseMeminfo(text: string | null): { total: number; available: number } | null {
+  if (!text) return null;
+  const field = (name: string) => {
+    const match = new RegExp(`^${name}:\\s+(\\d+)\\s*kB`, "m").exec(text);
+    return match ? Number(match[1]) * 1024 : null;
+  };
+  const total = field("MemTotal");
+  let available = field("MemAvailable");
+  if (available === null) {
+    const free = field("MemFree");
+    const cached = field("Cached");
+    const buffers = field("Buffers");
+    available = free === null ? null : free + (cached ?? 0) + (buffers ?? 0);
+  }
+  return total !== null && available !== null ? { total, available } : null;
+}
+
+/**
+ * The tightest cgroup v2 memory limit on the way from this process's cgroup
+ * to the root: {limit, used}, `used` without the page cache the kernel can
+ * drop (file pages minus shmem). Null outside a limited cgroup.
+ */
+export function readCgroupLimit(root = "/sys/fs/cgroup", selfCgroup = readText("/proc/self/cgroup")): { limit: number; used: number; available: number } | null {
+  if (!selfCgroup) return null;
+  const line = selfCgroup.split("\n").find((candidate) => candidate.startsWith("0::"));
+  if (!line) return null;
+  let dir = line.slice(3).trim() || "/";
+  let best: { limit: number; used: number; available: number } | null = null;
+  for (let guard = 0; guard < 32; guard++) {
+    const base = path.join(root, dir);
+    const max = readText(path.join(base, "memory.max"))?.trim();
+    if (max && max !== "max" && /^\d+$/.test(max)) {
+      const limit = Number(max);
+      const current = Number(readText(path.join(base, "memory.current"))?.trim() ?? NaN);
+      const stat = readText(path.join(base, "memory.stat")) ?? "";
+      const statField = (name: string) => Number(new RegExp(`^${name} (\\d+)`, "m").exec(stat)?.[1] ?? 0);
+      const reclaimable = Math.max(0, statField("file") - statField("shmem"));
+      const used = Number.isFinite(current) ? Math.max(0, current - reclaimable) : 0;
+      const available = Math.max(0, limit - used);
+      if (!best || available < best.available) best = { limit, used, available };
+    }
+    if (dir === "/" || dir === "") break;
+    dir = path.dirname(dir);
+  }
+  return best;
+}
+
+/** macOS: available ≈ free + inactive + speculative + purgeable pages (vm_stat). */
+export function parseVmStat(text: string): number | null {
+  const pageSize = Number(/page size of (\d+) bytes/.exec(text)?.[1] ?? 4096);
+  const pages = (name: string) => Number(new RegExp(`^Pages ${name}:\\s+(\\d+)`, "m").exec(text)?.[1] ?? NaN);
+  const free = pages("free");
+  if (!Number.isFinite(free)) return null;
+  const extra = ["inactive", "speculative", "purgeable"].map(pages).filter(Number.isFinite).reduce((a, b) => a + b, 0);
+  return (free + extra) * pageSize;
+}
+
+/** RAM and what is available of it right now (the cgroup's limit when tighter). */
+export function readSystemMemory(): SystemMemory {
+  let total = os.totalmem();
+  let available = os.freemem();
+  let source = "os";
+  if (process.platform === "linux") {
+    const info = parseMeminfo(readText("/proc/meminfo"));
+    if (info) {
+      total = info.total;
+      available = info.available;
+      source = "meminfo";
+    }
+    const cgroup = readCgroupLimit();
+    if (cgroup && cgroup.limit < total) {
+      total = cgroup.limit;
+      source = "cgroup";
+    }
+    if (cgroup && cgroup.available < available) {
+      available = cgroup.available;
+      source = "cgroup";
+    }
+  } else if (process.platform === "darwin") {
+    try {
+      const parsed = parseVmStat(String(execFileSync("vm_stat", { timeout: 2_000, maxBuffer: 64 * 1024 })));
+      if (parsed !== null) {
+        available = Math.min(total, parsed);
+        source = "vm_stat";
+      }
+    } catch {
+      /* os.freemem() stands */
+    }
+  }
+  return { totalBytes: total, availableBytes: Math.max(0, available), source };
+}
+
+export interface ProcessSample {
+  pid: number;
+  ppid: number;
+  /** Private resident memory (bytes): resident minus file-backed shared pages. */
+  privateBytes: number;
+  /** File-backed shared resident pages (the runtime binary, mapped files). */
+  sharedBytes: number;
+}
+
+/** Every process's parent and memory (Linux /proc, macOS ps); null where neither works. */
+export function listProcesses(): ProcessSample[] | null {
+  if (process.platform === "linux") {
+    const page = 4096;
+    const out: ProcessSample[] = [];
+    let names: string[];
+    try {
+      names = readdirSync("/proc");
+    } catch {
+      return null;
+    }
+    for (const name of names) {
+      if (!/^\d+$/.test(name)) continue;
+      const stat = readText(`/proc/${name}/stat`);
+      const statm = readText(`/proc/${name}/statm`);
+      if (!stat || !statm) continue;
+      // Fields after the ")" of the command name: state ppid ...
+      const after = stat.slice(stat.lastIndexOf(")") + 2).split(" ");
+      const ppid = Number(after[1]);
+      const [, resident, shared] = statm.trim().split(/\s+/).map(Number);
+      if (!Number.isFinite(ppid) || !Number.isFinite(resident)) continue;
+      out.push({ pid: Number(name), ppid, privateBytes: Math.max(0, resident - (shared || 0)) * page, sharedBytes: (shared || 0) * page });
+    }
+    return out;
+  }
+  if (process.platform === "darwin") {
+    try {
+      const text = String(execFileSync("ps", ["-A", "-o", "pid=,ppid=,rss="], { timeout: 3_000, maxBuffer: 8 * MIB }));
+      return text
+        .split("\n")
+        .map((line) => line.trim().split(/\s+/).map(Number))
+        .filter((parts) => parts.length === 3 && parts.every(Number.isFinite))
+        .map(([pid, ppid, rss]) => ({ pid, ppid, privateBytes: rss * 1024, sharedBytes: 0 }));
+    } catch {
+      return null;
+    }
+  }
+  return null;
+}
+
+/**
+ * Memory of `rootPid` and all its descendants: private memory summed, the
+ * shared file pages counted once (the largest; every agent process maps the
+ * same runtime binary). `byPid` has each member's private + shared bytes.
+ */
+export function treeMemory(processes: readonly ProcessSample[], rootPid: number): { bytes: number; count: number; byPid: Map<number, number> } {
+  const children = new Map<number, ProcessSample[]>();
+  const self = processes.find((sample) => sample.pid === rootPid);
+  for (const sample of processes) {
+    const list = children.get(sample.ppid);
+    if (list) list.push(sample);
+    else children.set(sample.ppid, [sample]);
+  }
+  const byPid = new Map<number, number>();
+  let privateTotal = 0;
+  let sharedMax = 0;
+  const stack: ProcessSample[] = self ? [self] : (children.get(rootPid) ?? []).slice();
+  const seen = new Set<number>();
+  while (stack.length > 0) {
+    const sample = stack.pop()!;
+    if (seen.has(sample.pid)) continue;
+    seen.add(sample.pid);
+    privateTotal += sample.privateBytes;
+    sharedMax = Math.max(sharedMax, sample.sharedBytes);
+    byPid.set(sample.pid, sample.privateBytes + sample.sharedBytes);
+    for (const child of children.get(sample.pid) ?? []) stack.push(child);
+  }
+  return { bytes: privateTotal + sharedMax, count: seen.size, byPid };
+}
+
+// --- the admission decision ---------------------------------------------------
+
+export interface AdmissionInput {
+  policy: MemoryPolicy;
+  system: SystemMemory;
+  /** The process tree's memory now (null when it cannot be measured). */
+  treeBytes: number | null;
+  /** Children of this process running now. */
+  running: number;
+  /**
+   * Memory the running children have not taken yet but will (young children
+   * still starting up): counted as used already.
+   */
+  reservedBytes: number;
+  /** What the next child is expected to take. */
+  estimateBytes: number;
+}
+
+export type AdmissionVerdict =
+  | { ok: true }
+  | { ok: false; reason: "critical" | "free" | "budget"; detail: string };
+
+function mb(bytes: number): string {
+  return bytes >= 10 * 1024 * MIB ? `${(bytes / (1024 * MIB)).toFixed(1)} GB` : `${Math.round(bytes / MIB)} MB`;
+}
+
+/** Whether one more child may start now. */
+export function admit(input: AdmissionInput): AdmissionVerdict {
+  const { policy, system } = input;
+  if (!policy.enabled) return { ok: true };
+  const available = Math.max(0, system.availableBytes - input.reservedBytes);
+  if (available < policy.criticalFreeBytes) {
+    return {
+      ok: false,
+      reason: "critical",
+      detail: `memory critically low: ${mb(available)} free of ${mb(system.totalBytes)} (the floor is ${mb(policy.criticalFreeBytes)})`,
+    };
+  }
+  const after = available - input.estimateBytes;
+  if (input.running === 0) {
+    // Progress: one child may run whenever the machine is not critically low.
+    if (after >= policy.criticalFreeBytes) return { ok: true };
+    return {
+      ok: false,
+      reason: "free",
+      detail: `${mb(available)} free, a sub-agent needs about ${mb(input.estimateBytes)} and ${mb(policy.criticalFreeBytes)} must stay free`,
+    };
+  }
+  if (after < policy.minFreeBytes) {
+    return {
+      ok: false,
+      reason: "free",
+      detail: `${mb(available)} free of ${mb(system.totalBytes)}; ${mb(policy.minFreeBytes)} stays free, a sub-agent needs about ${mb(input.estimateBytes)}`,
+    };
+  }
+  if (input.treeBytes !== null && input.treeBytes + input.reservedBytes + input.estimateBytes > policy.budgetBytes) {
+    return {
+      ok: false,
+      reason: "budget",
+      detail: `agents use ${mb(input.treeBytes + input.reservedBytes)} of their ${mb(policy.budgetBytes)} budget; a sub-agent needs about ${mb(input.estimateBytes)}`,
+    };
+  }
+  return { ok: true };
+}
+
+/** A child this process started: its pid and when, for the reservation. */
+interface Tracked {
+  pid: number | null;
+  startedAt: number;
+}
+
+/** How long a new child counts at least its estimate (it is still loading). */
+export const CHILD_WARMUP_MS = 15_000;
+
+/**
+ * The admission state of one agent process: measures (at most every
+ * `sampleMs`), keeps track of its own running children, and answers
+ * `check()`. Injectable readers keep it testable.
+ */
+export class MemoryAdmission {
+  readonly policy: MemoryPolicy;
+  private readonly rootPid: number;
+  private readonly now: () => number;
+  private readonly readSystem: () => SystemMemory;
+  private readonly readProcesses: () => ProcessSample[] | null;
+  private readonly sampleMs: number;
+  private readonly tracked = new Map<object, Tracked>();
+  private sample: { at: number; system: SystemMemory; tree: ReturnType<typeof treeMemory> | null } | null = null;
+  /** Largest footprint measured of a child of this process (bytes). */
+  private observedChildBytes = 0;
+
+  constructor(options: {
+    policy?: MemoryPolicy;
+    env?: NodeJS.ProcessEnv;
+    settings?: Record<string, unknown>;
+    rootPid?: number;
+    now?: () => number;
+    readSystem?: () => SystemMemory;
+    readProcesses?: () => ProcessSample[] | null;
+    sampleMs?: number;
+  } = {}) {
+    const env = options.env ?? process.env;
+    this.now = options.now ?? Date.now;
+    this.readSystem = options.readSystem ?? readSystemMemory;
+    this.readProcesses = options.readProcesses ?? listProcesses;
+    this.sampleMs = options.sampleMs ?? 1_000;
+    const rootFromEnv = Number(env[ENV_MEMORY_ROOT]);
+    this.rootPid = options.rootPid ?? (Number.isInteger(rootFromEnv) && rootFromEnv > 0 ? rootFromEnv : process.pid);
+    this.policy = options.policy ?? resolveMemoryPolicy(this.readSystem().totalBytes, env, options.settings ?? {});
+  }
+
+  /** The environment a child gets: the tree root it counts against. */
+  childEnv(): Record<string, string> {
+    return { [ENV_MEMORY_ROOT]: String(this.rootPid) };
+  }
+
+  private measure(force = false) {
+    const at = this.now();
+    if (!force && this.sample && at - this.sample.at < this.sampleMs) return this.sample;
+    const system = this.readSystem();
+    const processes = this.readProcesses();
+    const tree = processes ? treeMemory(processes, this.rootPid) : null;
+    if (tree) {
+      for (const entry of this.tracked.values()) {
+        const bytes = entry.pid !== null ? tree.byPid.get(entry.pid) : undefined;
+        if (bytes !== undefined && at - entry.startedAt >= CHILD_WARMUP_MS) this.observedChildBytes = Math.max(this.observedChildBytes, bytes);
+      }
+    }
+    this.sample = { at, system, tree };
+    return this.sample;
+  }
+
+  /** What the next child is expected to take: the estimate, or more once children were measured bigger. */
+  estimate(): number {
+    return Math.max(this.policy.childEstimateBytes, Math.min(this.observedChildBytes, 4 * this.policy.childEstimateBytes));
+  }
+
+  /** Memory young children will still take (estimate minus what they use now). */
+  private reserved(tree: ReturnType<typeof treeMemory> | null): number {
+    const at = this.now();
+    let total = 0;
+    for (const entry of this.tracked.values()) {
+      if (at - entry.startedAt >= CHILD_WARMUP_MS) continue;
+      const used = entry.pid !== null ? tree?.byPid.get(entry.pid) ?? 0 : 0;
+      total += Math.max(0, this.estimate() - used);
+    }
+    return total;
+  }
+
+  /** May one more child start now? */
+  check(): AdmissionVerdict & { system: SystemMemory; treeBytes: number | null } {
+    const sample = this.measure();
+    const treeBytes = sample.tree ? sample.tree.bytes : null;
+    // Without a process list (Windows), the children's estimates stand in for the tree.
+    const tree = treeBytes ?? (this.tracked.size > 0 ? this.tracked.size * this.estimate() : null);
+    const verdict = admit({
+      policy: this.policy,
+      system: sample.system,
+      treeBytes: tree,
+      running: this.tracked.size,
+      reservedBytes: this.reserved(sample.tree),
+      estimateBytes: this.estimate(),
+    });
+    return { ...verdict, system: sample.system, treeBytes: tree };
+  }
+
+  /** A child started (its pid once known): it counts from now on. */
+  started(key: object, pid: number | null = null): void {
+    // No new measurement for it: until the next sample (and while it warms
+    // up) the reservation counts its estimate.
+    this.tracked.set(key, { pid, startedAt: this.now() });
+  }
+
+  /** The child's pid arrived after it was started. */
+  spawned(key: object, pid: number | null | undefined): void {
+    const entry = this.tracked.get(key);
+    if (entry && typeof pid === "number") entry.pid = pid;
+  }
+
+  /** A child ended. */
+  finished(key: object): void {
+    this.tracked.delete(key);
+    if (this.sample) this.sample = { ...this.sample, at: -Infinity };
+  }
+
+  /** Children of this process that count now. */
+  get running(): number {
+    return this.tracked.size;
+  }
+
+  /** One line for /agents and status: RAM, available, the tree, the thresholds. */
+  describe(): string {
+    const sample = this.measure(true);
+    const tree = sample.tree ? mb(sample.tree.bytes) : "unknown";
+    const p = this.policy;
+    if (!p.enabled) return `memory admission off (${mb(sample.system.availableBytes)} free of ${mb(sample.system.totalBytes)})`;
+    return `memory: ${mb(sample.system.availableBytes)} free of ${mb(sample.system.totalBytes)} (${sample.system.source}); agents ${tree} of ${mb(p.budgetBytes)} budget; keeps ${mb(p.minFreeBytes)} free, stops at ${mb(p.criticalFreeBytes)}; ~${mb(this.estimate())} per sub-agent`;
+  }
+}