Fix job records lost on concurrent background task launches

plugins/codex/scripts/lib/state.mjs

@@ -1,3 +1,4 @@
+import { spawnSync } from "node:child_process";
 import { createHash } from "node:crypto";
 import fs from "node:fs";
 import os from "node:os";
@@ -111,14 +112,237 @@ export function saveState(cwd, state) {
     removeFileIfExists(job.logFile);
   }
 
-  fs.writeFileSync(resolveStateFile(cwd), `${JSON.stringify(nextState, null, 2)}\n`, "utf8");
+  const stateFile = resolveStateFile(cwd);
+  const tmpFile = `${stateFile}.${process.pid}.tmp`;
+  fs.writeFileSync(tmpFile, `${JSON.stringify(nextState, null, 2)}\n`, "utf8");
+  fs.renameSync(tmpFile, stateFile); // atomic replace; readers never see a partial file
   return nextState;
 }
 
+// Blocking sleep for a sync context (no busy-wait) via Atomics.wait.
+function sleepSync(ms) {
+  Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, ms);
+}
+
+// A per-process, restart-stable token identifying a specific process *instance*
+// (not just its PID), so PID reuse can be told apart from the original owner.
+// Lock files live under a per-workspace OS temp dir (see resolveStateDir), i.e. a
+// single host, so the PID refers to a process on this machine.
+//
+// The source is chosen by platform and never mixed: two renderings of the same
+// live process must compare equal, so we must not stamp with one source and check
+// with another. Returns null when the process is gone or its start time can't be
+// read (callers treat null conservatively -- never as "different instance").
+//   - Linux: /proc/<pid>/stat field 22 is the process start time (clock ticks
+//     since boot); read directly, no subprocess.
+//   - else (macOS/BSD): `ps -o lstart` is the start timestamp, stable for the
+//     process lifetime. The env is pinned (TZ/locale) because lstart is rendered
+//     with strftime and would otherwise differ between a stamper and a checker
+//     running under different TZ/LC settings. spawnSync is only reached on the
+//     (rare) reclaim path, never on the uncontended fast path.
+function processStartToken(pid) {
+  if (!Number.isInteger(pid) || pid <= 0) return null;
+  if (process.platform === "linux") {
+    try {
+      const stat = fs.readFileSync(`/proc/${pid}/stat`, "utf8");
+      // comm (field 2) may contain spaces/parens; the numeric fields start after
+      // the last ')'. starttime is field 22 => index 19 of that remainder.
+      const rest = stat.slice(stat.lastIndexOf(")") + 2).split(/\s+/);
+      if (rest[19]) return `L:${rest[19]}`;
+    } catch {}
+    return null; // no cross-source fallback -- see note above
+  }
+  try {
+    const r = spawnSync("/bin/ps", ["-o", "lstart=", "-p", String(pid)], {
+      encoding: "utf8",
+      env: { ...process.env, TZ: "UTC0", LC_ALL: "C", LANG: "C" },
+    });
+    if (r.status === 0) {
+      const s = (r.stdout || "").trim();
+      if (s) return `P:${s}`;
+    }
+  } catch {}
+  return null;
+}
+
+// A process that has exited but not yet been reaped by its parent is a zombie:
+// process.kill(pid, 0) still succeeds and its start token is unchanged, so it
+// would otherwise look like a live owner forever. Detect it so its lock is
+// reclaimed instead of blocking every writer until the parent reaps it.
+function isZombie(pid) {
+  if (process.platform === "linux") {
+    try {
+      const stat = fs.readFileSync(`/proc/${pid}/stat`, "utf8");
+      // state (field 3) is the first token after the last ')'.
+      return stat.slice(stat.lastIndexOf(")") + 2).split(/\s+/)[0] === "Z";
+    } catch {
+      return false;
+    }
+  }
+  try {
+    const r = spawnSync("/bin/ps", ["-o", "state=", "-p", String(pid)], {
+      encoding: "utf8",
+      env: { ...process.env, LC_ALL: "C" },
+    });
+    return r.status === 0 && (r.stdout || "").trim().startsWith("Z");
+  } catch {
+    return false;
+  }
+}
+
+// The lock owner is identified as "<pid>.<startTokenHex>.<time>.<rand>". Compute
+// the current process's identity once per acquisition.
+function selfOwnerId() {
+  const tok = processStartToken(process.pid);
+  const tag = tok ? Buffer.from(tok).toString("hex") : "0";
+  return `${process.pid}.${tag}.${Date.now().toString(36)}.${Math.random().toString(36).slice(2, 8)}`;
+}
+
+// True only when the exact process instance that wrote `id` is gone -- never for
+// a live owner, no matter how long it has been holding the lock. This is what
+// lets reclaim run without any time-based expiry: a suspended-but-alive holder is
+// never reclaimed, and a dead owner whose PID was recycled is detected because
+// the recycled process reports a different start token.
+function isAbandoned(id) {
+  const parts = String(id).split(".");
+  const pid = Number.parseInt(parts[0], 10);
+  if (!Number.isInteger(pid) || pid <= 0) return true; // empty/garbled -> not a live owner
+  let alive;
+  try {
+    process.kill(pid, 0);
+    alive = true;
+  } catch (err) {
+    alive = err.code === "EPERM"; // exists but not ours (still alive); ESRCH => dead
+  }
+  if (!alive) return true; // owner process is gone
+  if (isZombie(pid)) return true; // exited but unreaped -> effectively gone
+  // PID is alive. Only declare it abandoned if we can PROVE it is a different
+  // instance. If the owner's stamp is unverifiable ("0"), or we can't read the
+  // current start token, treat the live PID as the same instance and do NOT
+  // reclaim -- otherwise a checker that *can* probe would delete a live owner's
+  // lock whose owner merely failed to self-probe at stamp time. (pid-death
+  // reclaim above still applies, so this never causes a permanent deadlock.)
+  if (!parts[1] || parts[1] === "0") return false;
+  const cur = processStartToken(pid);
+  if (cur === null) return false;
+  return Buffer.from(cur).toString("hex") !== parts[1]; // different instance => PID reuse => owner gone
+}
+
+// Publish a lock atomically: write the owner id into a unique temp file, then
+// hard-link it onto the fixed path. linkSync is atomic and fails EEXIST if the
+// path is already held, exactly like O_EXCL -- but unlike open()+write() the file
+// has its full content the instant it appears at the path, so a concurrent reader
+// can never observe an empty lock and mistake a just-created live lock for an
+// abandoned one. Returns true if claimed, false if already held.
+function claimLock(lockFile, ownerId) {
+  const tmp = `${lockFile}.tmp.${ownerId}`;
+  fs.writeFileSync(tmp, ownerId);
+  try {
+    fs.linkSync(tmp, lockFile);
+    return true;
+  } catch (err) {
+    if (err.code === "EEXIST") return false;
+    throw err;
+  } finally {
+    try { fs.unlinkSync(tmp); } catch {}
+  }
+}
+
+// Reclaim a lock only when its owner instance is gone, never a live one.
+//
+// isAbandoned is true only for a dead owner (or a PID recycled by a different
+// instance), so this never fires against the live owner. Removal is atomic via
+// capture-by-rename: renameSync of the fixed path has exactly-one-winner
+// semantics, so when several launches race to reclaim the same lock only one
+// captures it and the losers get ENOENT and fall back to re-contending. The
+// captor then confirms the file it captured is byte-for-byte the id it judged
+// abandoned (the id embeds time+rand, so this is unique); if instead it captured
+// a lock that had been recreated in the meantime -- i.e. possibly a live one --
+// it restores it rather than deleting it. This needs no second lock, so there is
+// no reclaim-lock to itself go stale and be cleaned up unsafely.
+//
+// Residual (fundamental to pure-fs locking): between the re-read and the rename
+// below the lock could be reclaimed by someone else and freshly re-claimed by a
+// live owner; renaming then captures that live lock, and if a third process
+// claims the momentarily-absent path the restore relink loses (EEXIST) and the
+// captured lock is dropped. The pre-rename re-read shrinks that window to ~2
+// adjacent syscalls with no subprocess in it; isAbandoned (which may spawn `ps`)
+// runs only *before* the re-read. Closing it completely needs an atomic
+// conditional-delete / OS advisory lock (flock), which Node's fs builtins do not
+// expose. Within that ~2-syscall window the worst case is a brief overlap (two
+// writers) or a losing contender that errors out at the 15s acquire deadline;
+// never a permanent deadlock, since a dead owner is always reclaimable next pass.
+function reclaimIfAbandoned(lockFile, selfId) {
+  let seen;
+  try {
+    seen = fs.readFileSync(lockFile, "utf8");
+  } catch {
+    return; // already gone
+  }
+  if (!isAbandoned(seen)) return; // live owner -> wait, don't touch
+  const tomb = `${lockFile}.rip.${selfId}`;
+  // Re-read immediately before capturing so isAbandoned's (possibly subprocess-
+  // backed) probe is not inside the capture window: only proceed if the lock is
+  // still the exact abandoned instance we judged.
+  try {
+    if (fs.readFileSync(lockFile, "utf8") !== seen) return;
+  } catch {
+    return; // vanished -- re-contend
+  }
+  try {
+    fs.renameSync(lockFile, tomb); // atomic: exactly one reclaimer captures the path
+  } catch {
+    return; // lost the race (ENOENT) -- another reclaimer took it; re-contend
+  }
+  try {
+    if (fs.readFileSync(tomb, "utf8") === seen) {
+      fs.unlinkSync(tomb); // captured the exact abandoned instance we judged -> drop it
+    } else {
+      // Captured a lock recreated after our read -> may be live; put it back.
+      try { fs.linkSync(tomb, lockFile); } catch {}
+      fs.unlinkSync(tomb);
+    }
+  } catch {
+    try { fs.unlinkSync(tomb); } catch {}
+  }
+}
+
+// Cross-process lock around the state.json read-modify-write. Without it,
+// concurrent `task --background` launches each read the same base state, add
+// only their own job, and clobber siblings on write (and saveState's prune
+// then deletes the "orphan" job files). Serializing the RMW fixes both.
+function withStateLock(cwd, fn) {
+  ensureStateDir(cwd);
+  const lockFile = path.join(resolveStateDir(cwd), "state.lock");
+  // Identifies this exact process instance (pid + start-time), so others can tell
+  // a live owner from a recycled PID; the full string is stamped into the lock so
+  // only the true owner removes it on release.
+  const ownerId = selfOwnerId();
+  const deadline = Date.now() + 15000;
+  for (;;) {
+    if (claimLock(lockFile, ownerId)) break;
+    reclaimIfAbandoned(lockFile, ownerId);
+    if (Date.now() > deadline) throw new Error("Timed out acquiring Codex state lock");
+    sleepSync(20 + Math.floor(Math.random() * 30)); // jittered backoff
+  }
+  try {
+    return fn();
+  } finally {
+    // Only remove the lock if we still own it: if we were ever reclaimed, the
+    // contents no longer match ownerId and we must not delete the lock another
+    // process now holds.
+    try {
+      if (fs.readFileSync(lockFile, "utf8") === ownerId) fs.unlinkSync(lockFile);
+    } catch {}
+  }
+}
+
 export function updateState(cwd, mutate) {
-  const state = loadState(cwd);
-  mutate(state);
-  return saveState(cwd, state);
+  return withStateLock(cwd, () => {
+    const state = loadState(cwd);
+    mutate(state);
+    return saveState(cwd, state);
+  });
 }
 
 export function generateJobId(prefix = "job") {

plugins/codex/scripts/session-lifecycle-hook.mjs

@@ -13,7 +13,7 @@ import {
   sendBrokerShutdown,
   teardownBrokerSession
 } from "./lib/broker-lifecycle.mjs";
-import { loadState, resolveStateFile, saveState } from "./lib/state.mjs";
+import { loadState, resolveStateFile, updateState } from "./lib/state.mjs";
 import { TRANSCRIPT_PATH_ENV } from "./lib/claude-session-transfer.mjs";
 import { resolveWorkspaceRoot } from "./lib/workspace.mjs";
 
@@ -50,28 +50,49 @@ function cleanupSessionJobs(cwd, sessionId) {
     return;
   }
 
-  const state = loadState(workspaceRoot);
-  const removedJobs = state.jobs.filter((job) => job.sessionId === sessionId);
-  if (removedJobs.length === 0) {
-    return;
-  }
-
-  for (const job of removedJobs) {
-    const stillRunning = job.status === "queued" || job.status === "running";
-    if (!stillRunning) {
-      continue;
+  // Drop this session's jobs through the locked read-modify-write so a concurrent
+  // upsertJob (task launch) can't be clobbered by a stale snapshot, and capture
+  // which running jobs to terminate. Process teardown runs in `finally`, after the
+  // lock is released, so a failed state write (e.g. the 15s lock-acquire timeout)
+  // can never leak this session's processes, and never aborts the rest of session
+  // shutdown (broker teardown) -- session cleanup is best-effort.
+  const isRunning = (job) => job.status === "queued" || job.status === "running";
+  const toTerminate = [];
+  try {
+    updateState(workspaceRoot, (state) => {
+      for (const job of state.jobs) {
+        if (job.sessionId === sessionId && isRunning(job)) {
+          toTerminate.push(job.pid ?? Number.NaN);
+        }
+      }
+      state.jobs = state.jobs.filter((job) => job.sessionId !== sessionId);
+    });
+  } catch {
+    // The locked update failed (e.g. lock-acquire timeout). Still identify this
+    // session's processes so we can tear them down -- via a best-effort unlocked
+    // read only. We deliberately do NOT write state here: an unlocked save is the
+    // very clobber this lock prevents; the stale records are removed on a later
+    // locked pass.
+    if (toTerminate.length === 0) {
+      try {
+        for (const job of loadState(workspaceRoot).jobs) {
+          if (job.sessionId === sessionId && isRunning(job)) {
+            toTerminate.push(job.pid ?? Number.NaN);
+          }
+        }
+      } catch {
+        // Nothing more we can do; fall through to whatever we collected.
+      }
     }
-    try {
-      terminateProcessTree(job.pid ?? Number.NaN);
-    } catch {
-      // Ignore teardown failures during session shutdown.
+  } finally {
+    for (const pid of toTerminate) {
+      try {
+        terminateProcessTree(pid);
+      } catch {
+        // Ignore teardown failures during session shutdown.
+      }
     }
   }
-
-  saveState(workspaceRoot, {
-    ...state,
-    jobs: state.jobs.filter((job) => job.sessionId !== sessionId)
-  });
 }
 
 function handleSessionStart(input) {