@vellumai/assistant 0.11.7-dev.202608281503.cbafd6c → 0.11.7-dev.202608282003.8ce5421

package/src/util/worker-process.ts

@@ -19,8 +19,21 @@ import {
 import { dirname, join } from "node:path";
 import { fileURLToPath } from "node:url";
 
-import { getCurrentLogFilePath } from "./logger.js";
+import { getCurrentLogFilePath, getLogger } from "./logger.js";
+import { isProcessAlive } from "./process-liveness.js";
+import {
+  listProcessTable,
+  type ProcessTableRow,
+  readRawProcessCommand,
+} from "./process-table.js";
 import { workerMemoryEnv } from "./worker-memory.js";
+import {
+  classifyWorkerOwnership,
+  isDaemonCommand,
+  pid1OwnsWorkers,
+} from "./worker-ownership.js";
+
+const log = getLogger("worker-process");
 
 export interface WorkerProcessStatus {
   status: "running" | "not_running";
@@ -135,21 +148,40 @@ interface WorkerCommandRuntime {
   executableExists: (path: string) => boolean;
 }
 
-export function resolveWorkerCommand(
-  entry: URL,
-  packagedEntry: PackagedWorkerEntry | undefined,
-  runtime: WorkerCommandRuntime = {
+function defaultWorkerCommandRuntime(): WorkerCommandRuntime {
+  return {
     platform: process.platform,
     execPath: process.execPath,
     executableExists: existsSync,
-  },
-): string[] {
+  };
+}
+
+/**
+ * How a packaged Windows runtime spawns this worker, or null when the worker
+ * runs from source. Only packaged Windows runtimes ship the executable.
+ */
+function packagedWorkerSpawn(
+  packagedEntry: PackagedWorkerEntry | undefined,
+  runtime: WorkerCommandRuntime,
+): { executable: string; entry: PackagedWorkerEntry } | null {
   if (runtime.platform === "win32" && packagedEntry) {
     const executable = join(dirname(runtime.execPath), "vellum-worker.exe");
     if (runtime.executableExists(executable)) {
-      return [executable, packagedEntry];
+      return { executable, entry: packagedEntry };
     }
   }
+  return null;
+}
+
+export function resolveWorkerCommand(
+  entry: URL,
+  packagedEntry: PackagedWorkerEntry | undefined,
+  runtime: WorkerCommandRuntime = defaultWorkerCommandRuntime(),
+): string[] {
+  const packaged = packagedWorkerSpawn(packagedEntry, runtime);
+  if (packaged) {
+    return [packaged.executable, packaged.entry];
+  }
   return ["bun", "--smol", "run", fileURLToPath(entry)];
 }
 
@@ -194,10 +226,310 @@ async function waitForWorkerPidFile(
   return existsSync(pidPath) ? "ready" : "timeout";
 }
 
+/** Path separators differ by platform; compare command lines on one form. */
+function normalizeSeparators(value: string): string {
+  return value.replaceAll("\\", "/");
+}
+
+/** Whether a command line carries every fragment that marks this worker. */
+function matchesSignature(
+  command: string,
+  signature: readonly string[],
+): boolean {
+  const normalized = normalizeSeparators(command);
+  return signature.every((part) =>
+    normalized.includes(normalizeSeparators(part)),
+  );
+}
+
+/**
+ * Fragments of a command line that together mark a process as this worker,
+ * whichever install it came from. All must be present.
+ *
+ * This is a safety guard, not the ownership test. Ownership is decided by
+ * parentage; this only keeps us from signalling a process that happens to hold
+ * a recycled PID. Three trailing path segments are specific enough
+ * (`src/schedule/worker.ts`, `defaults/memory/worker.ts`) that an unrelated
+ * program running some other `worker.ts` does not match. Packaged workers all
+ * run one executable and are told apart by the subcommand, so the signature
+ * carries both: matching on the executable alone would let one packaged
+ * worker's slot reclaim another's process.
+ */
+export function workerKindSignature(
+  entry: URL,
+  packagedEntry: PackagedWorkerEntry | undefined,
+  runtime: WorkerCommandRuntime = defaultWorkerCommandRuntime(),
+): readonly string[] {
+  const packaged = packagedWorkerSpawn(packagedEntry, runtime);
+  if (packaged) {
+    // Matched as separate fragments rather than one adjacent string: a Windows
+    // command line may quote the executable path, so the two are not reliably
+    // neighbours.
+    return ["vellum-worker", packaged.entry];
+  }
+  return [
+    normalizeSeparators(fileURLToPath(entry))
+      .split("/")
+      .filter(Boolean)
+      .slice(-3)
+      .join("/"),
+  ];
+}
+
+/**
+ * What became of the orphan after an awaited gap, from liveness plus a fresh
+ * command-line read.
+ *
+ *   - `gone`: the PID is dead, or alive under a command line that is not this
+ *     worker (the orphan exited and the OS recycled its PID). The slot may be
+ *     released.
+ *   - `orphan`: still alive and still reads as this worker. Escalation may
+ *     proceed.
+ *   - `identity-unreadable`: alive, but the command line could not be read, so
+ *     "still the orphan" and "recycled PID" cannot be told apart. Uncertainty
+ *     must resolve to reuse: no signal, and the slot must NOT be released, or
+ *     a second worker would be spawned next to a live one.
+ */
+export type OrphanFate = "gone" | "orphan" | "identity-unreadable";
+
+export function classifyOrphanAfterWait(
+  alive: boolean,
+  command: string | null,
+  signature: readonly string[],
+): OrphanFate {
+  if (!alive) {
+    return "gone";
+  }
+  if (command == null) {
+    return "identity-unreadable";
+  }
+  return matchesSignature(command, signature) ? "orphan" : "gone";
+}
+
+/** Re-read the orphan's fate for `pid` from the live process table. */
+function orphanFate(pid: number, signature: readonly string[]): OrphanFate {
+  return classifyOrphanAfterWait(
+    isProcessAlive(pid),
+    readRawProcessCommand(pid),
+    signature,
+  );
+}
+
+/**
+ * How long an orphaned worker gets to exit after SIGTERM.
+ *
+ * Derived from the longest shutdown any worker can legitimately take, not
+ * picked. The memory worker reaps the ONNX embedding subprocess it owns before
+ * exiting and bounds that at `EMBEDDING_SHUTDOWN_BUDGET_MS + 1s`, so a shorter
+ * ceiling here would SIGKILL it mid-reap and strand the very subprocess that
+ * reap exists to collect. A test asserts this stays above that bound.
+ */
+export const ORPHAN_STOP_TIMEOUT_MS = 15_000;
+
+/** Poll interval while waiting for a stopped worker to exit. */
+const ORPHAN_STOP_POLL_INTERVAL_MS = 100;
+
+/**
+ * How long to wait for a SIGKILLed worker to actually disappear. Signal
+ * delivery and reaping are asynchronous, so a liveness probe taken straight
+ * after the kill can still see a process that is on its way out.
+ */
+const ORPHAN_KILL_CONFIRM_MS = 1_000;
+
+/** Poll until `pid` is gone or `timeoutMs` elapses. */
+async function waitForExit(pid: number, timeoutMs: number): Promise<void> {
+  const deadline = Date.now() + timeoutMs;
+  while (Date.now() < deadline && isProcessAlive(pid)) {
+    await Bun.sleep(ORPHAN_STOP_POLL_INTERVAL_MS);
+  }
+}
+
+/**
+ * Stop a worker orphaned by a previous owner and release its PID file.
+ * Reports whether it actually exited.
+ *
+ * Escalates to SIGKILL because the orphan runs a different generation's
+ * shutdown path, which this process cannot make assumptions about. Signalling
+ * can still fail outright, most plausibly EPERM for a process owned by another
+ * user, so the PID file is released only once the process is confirmed gone:
+ * dropping it while the worker still runs would let the caller spawn a second
+ * one against the same workspace. The unlink is also identity-checked, since a
+ * concurrent launcher may already have replaced the worker and deleting that
+ * successor's entry would strand it.
+ */
+async function stopOrphanedWorker(
+  pid: number,
+  pidPath: string,
+  workerLabel: string,
+  signature: readonly string[],
+): Promise<boolean> {
+  try {
+    process.kill(pid, "SIGTERM");
+  } catch {
+    // Exited between the probe and the signal.
+  }
+
+  await waitForExit(pid, ORPHAN_STOP_TIMEOUT_MS);
+  // Identity is re-proven after every awaited gap, not carried by the PID:
+  // the orphan may have exited during the wait and the OS may have handed its
+  // PID to a stranger. Only a positive re-match escalates, and only a proven
+  // "gone" releases the slot.
+  let fate = orphanFate(pid, signature);
+  if (fate === "orphan") {
+    try {
+      process.kill(pid, "SIGKILL");
+    } catch {
+      // Best-effort.
+    }
+    await waitForExit(pid, ORPHAN_KILL_CONFIRM_MS);
+    fate = orphanFate(pid, signature);
+  }
+
+  if (fate !== "gone") {
+    log.warn(
+      { pid, pidPath, fate },
+      fate === "orphan"
+        ? `${workerLabel} orphaned by a previous owner could not be stopped; reusing it rather than running a second one`
+        : `${workerLabel} is still live but its identity could not be re-read; reusing it rather than risking a duplicate`,
+    );
+    return false;
+  }
+
+  unlinkPidFileIfNames(pidPath, pid);
+  log.warn(
+    { pid, pidPath },
+    `${workerLabel} orphaned by a previous owner was stopped so this process can start its own`,
+  );
+  return true;
+}
+
+/**
+ * What to do about the process a worker's PID file currently names.
+ *
+ *   - `adopt`: reuse it. Either it is this process's own worker, or it belongs
+ *     to another live owner, or it is not recognisably one of our workers at
+ *     all and must never be signalled.
+ *   - `reclaim`: an orphan left by an owner that is gone. Stop it, then spawn.
+ *   - `spawn`: nothing is holding the slot.
+ */
+export type WorkerSlotDecision =
+  | { action: "adopt"; pid: number }
+  | { action: "reclaim"; pid: number }
+  | { action: "spawn" };
+
+/**
+ * The whole safety decision, with no IO, so every branch is testable.
+ *
+ * `reclaim` is the only outcome that signals a process, and it requires two
+ * independent things to agree: the command line marks the process as this
+ * worker, and parentage says no daemon owns it. A missing row, an unreadable
+ * command line, and a command line that does not match all fall back to
+ * `adopt`, so an uncertain answer never costs a process its life.
+ *
+ * `isOwnerAlive` is the caller's definition of a legitimate owner. These
+ * workers have exactly one, the daemon, so passing a plain liveness probe
+ * would let a recycled owner PID keep a stranded worker looking owned forever.
+ * The embed worker passes a liveness probe instead, because two processes may
+ * each legitimately run one of its workers.
+ */
+export function decideWorkerSlot(
+  status: WorkerProcessStatus,
+  row: Pick<ProcessTableRow, "pid" | "ppid" | "command"> | null,
+  signature: readonly string[],
+  selfPid: number,
+  isOwnerAlive: (pid: number) => boolean,
+  pid1OwnsWorkers: boolean,
+): WorkerSlotDecision {
+  if (status.status !== "running" || status.pid == null) {
+    return { action: "spawn" };
+  }
+  if (!row) {
+    return { action: "adopt", pid: status.pid };
+  }
+  if (!matchesSignature(row.command, signature)) {
+    return { action: "adopt", pid: status.pid };
+  }
+  const ownership = classifyWorkerOwnership(
+    row,
+    selfPid,
+    isOwnerAlive,
+    pid1OwnsWorkers,
+  );
+  return ownership === "orphan"
+    ? { action: "reclaim", pid: status.pid }
+    : { action: "adopt", pid: status.pid };
+}
+
+/**
+ * The slot decision for the worker at `pidPath`, made without yielding.
+ *
+ * Callers spawn inside their own synchronous prefix, so this stays sync: an
+ * `await` here would let a caller's next tick run before the child exists.
+ */
+function inspectWorkerSlot(
+  pidPath: string,
+  signature: readonly string[],
+): WorkerSlotDecision {
+  const status = probeWorkerPidFile(pidPath);
+  if (status.status !== "running" || status.pid == null) {
+    return { action: "spawn" };
+  }
+
+  // The worker's row, its owner's identity, and what PID 1 is all come from
+  // one snapshot. Reading the table is a subprocess on most platforms, so
+  // asking three separate times would be three of them per worker.
+  let table: readonly ProcessTableRow[];
+  try {
+    table = listProcessTable();
+  } catch {
+    table = [];
+  }
+  const commandOf = (pid: number): string | null =>
+    table.find((row) => row.pid === pid)?.command ?? null;
+
+  const row = table.find((entry) => entry.pid === status.pid) ?? null;
+  if (row && !matchesSignature(row.command, signature)) {
+    log.info(
+      { pid: row.pid, pidPath },
+      "Worker PID file names a live process this runtime does not recognise as its worker; reusing it rather than signalling it",
+    );
+  }
+
+  return decideWorkerSlot(
+    status,
+    row,
+    signature,
+    process.pid,
+    (pid) => isDaemonCommand(commandOf(pid)),
+    pid1OwnsWorkers(commandOf(1)),
+  );
+}
+
+/**
+ * Stop the orphan holding the slot, then report a PID to adopt instead of
+ * spawning: either the orphan itself when it could not be stopped, or a
+ * replacement a concurrent launcher brought up while we waited.
+ */
+async function reclaimWorkerSlot(
+  pid: number,
+  pidPath: string,
+  workerLabel: string,
+  signature: readonly string[],
+): Promise<number | null> {
+  if (!(await stopOrphanedWorker(pid, pidPath, workerLabel, signature))) {
+    // Still running and beyond our reach. One stale worker beats two live ones.
+    return pid;
+  }
+  const replacement = probeWorkerPidFile(pidPath);
+  return replacement.status === "running" && replacement.pid != null
+    ? replacement.pid
+    : null;
+}
+
 /**
  * Spawn a worker entry script as a background process and wait for it to
  * report readiness by writing its PID file. `opts.detached` (default `true`)
- * controls process parentage — see {@link SpawnWorkerProcessOptions}.
+ * controls process parentage, see {@link SpawnWorkerProcessOptions}.
  *
  * If a worker is already running (per the PID file), returns its PID with
  * `alreadyRunning: true` rather than spawning a second one. Throws
@@ -220,9 +552,21 @@ export async function spawnWorkerProcess(args: {
   const pidPollIntervalMs = opts.pidPollIntervalMs ?? PID_FILE_POLL_INTERVAL_MS;
   const detached = opts.detached ?? true;
 
-  const current = probeWorkerPidFile(args.pidPath);
-  if (current.status === "running" && current.pid != null) {
-    return { pid: current.pid, alreadyRunning: true };
+  const signature = workerKindSignature(args.entry, args.packagedEntry);
+  const decision = inspectWorkerSlot(args.pidPath, signature);
+  if (decision.action === "adopt") {
+    return { pid: decision.pid, alreadyRunning: true };
+  }
+  if (decision.action === "reclaim") {
+    const adopted = await reclaimWorkerSlot(
+      decision.pid,
+      args.pidPath,
+      args.workerLabel,
+      signature,
+    );
+    if (adopted != null) {
+      return { pid: adopted, alreadyRunning: true };
+    }
   }
 
   // Pipe the worker's stderr into the same daily log file the daemon writes
@@ -319,23 +663,28 @@ export function stopWorkerProcess(pidPath: string): WorkerProcessStatus {
 // ---------------------------------------------------------------------------
 
 /**
- * Remove the worker PID file only while it still names this process. A
- * successor worker overwrites the file with its own PID at startup, so an
- * exiting predecessor that unlinked unconditionally would delete the
- * successor's entry — the liveness supervisor would then respawn a
- * duplicate and orphan the successor.
+ * Remove `pidPath` only while it still names `pid`.
+ *
+ * A successor worker overwrites the file with its own PID at startup, so an
+ * unconditional unlink would delete the successor's entry: the liveness
+ * supervisor would then respawn a duplicate and orphan the successor.
  */
-export function cleanupWorkerPidFile(pidPath: string): void {
+export function unlinkPidFileIfNames(pidPath: string, pid: number): void {
   try {
     const raw = readFileSync(pidPath, "utf-8").trim();
-    if (parseInt(raw, 10) === process.pid) {
+    if (parseInt(raw, 10) === pid) {
       unlinkSync(pidPath);
     }
   } catch {
-    // best-effort — a missing or unreadable file needs no cleanup
+    // best-effort: a missing or unreadable file needs no cleanup
   }
 }
 
+/** Release this process's own PID file on the way out. */
+export function cleanupWorkerPidFile(pidPath: string): void {
+  unlinkPidFileIfNames(pidPath, process.pid);
+}
+
 /**
  * Watch the worker PID file and evict this process when the file stops
  * naming it. The PID file is a worker's sole tracking handle: stop