Fix test broker leaks, state races, and signal-masked command failures

plugins/codex/scripts/lib/broker-lifecycle.mjs

@@ -6,6 +6,7 @@ import process from "node:process";
 import { spawn } from "node:child_process";
 import { fileURLToPath } from "node:url";
 import { createBrokerEndpoint, parseBrokerEndpoint } from "./broker-endpoint.mjs";
+import { withLock } from "./locking.mjs";
 import { resolveStateDir } from "./state.mjs";
 
 export const PID_FILE_ENV = "CODEX_COMPANION_APP_SERVER_PID_FILE";
@@ -111,6 +112,16 @@ async function isBrokerEndpointReady(endpoint) {
 }
 
 export async function ensureBrokerSession(cwd, options = {}) {
+  const stateDir = resolveStateDir(cwd);
+  fs.mkdirSync(stateDir, { recursive: true });
+  // Serialize the check-then-create sequence per workspace so two concurrent
+  // callers cannot both miss the existing broker and spawn a duplicate.
+  return withLock(path.join(stateDir, ".broker.lock"), () => ensureBrokerSessionLocked(cwd, options), {
+    timeoutMs: options.lockTimeoutMs ?? 10000
+  });
+}
+
+async function ensureBrokerSessionLocked(cwd, options = {}) {
   const existing = loadBrokerSession(cwd);
   if (existing && (await isBrokerEndpointReady(existing.endpoint))) {
     return existing;

plugins/codex/scripts/lib/locking.mjs

@@ -0,0 +1,88 @@
+import fs from "node:fs";
+
+const DEFAULT_TIMEOUT_MS = 5000;
+const DEFAULT_STALE_MS = 30000;
+const RETRY_DELAY_MS = 25;
+
+function sleepSync(ms) {
+  Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, ms);
+}
+
+function sleep(ms) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+function tryAcquire(lockDir) {
+  try {
+    fs.mkdirSync(lockDir);
+    return true;
+  } catch (error) {
+    if (error?.code === "EEXIST") {
+      return false;
+    }
+    throw error;
+  }
+}
+
+function reclaimIfStale(lockDir, staleMs) {
+  try {
+    const age = Date.now() - fs.statSync(lockDir).mtimeMs;
+    if (age > staleMs) {
+      fs.rmdirSync(lockDir);
+    }
+  } catch {
+    // Lock released (or reclaimed by someone else) between stat and rmdir.
+  }
+}
+
+export function acquireLockSync(lockDir, options = {}) {
+  const timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS;
+  const staleMs = options.staleMs ?? DEFAULT_STALE_MS;
+  const deadline = Date.now() + timeoutMs;
+  while (!tryAcquire(lockDir)) {
+    reclaimIfStale(lockDir, staleMs);
+    if (Date.now() >= deadline) {
+      throw new Error(`Timed out waiting for lock: ${lockDir}`);
+    }
+    sleepSync(RETRY_DELAY_MS);
+  }
+}
+
+export async function acquireLock(lockDir, options = {}) {
+  const timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS;
+  const staleMs = options.staleMs ?? DEFAULT_STALE_MS;
+  const deadline = Date.now() + timeoutMs;
+  while (!tryAcquire(lockDir)) {
+    reclaimIfStale(lockDir, staleMs);
+    if (Date.now() >= deadline) {
+      throw new Error(`Timed out waiting for lock: ${lockDir}`);
+    }
+    await sleep(RETRY_DELAY_MS);
+  }
+}
+
+export function releaseLock(lockDir) {
+  try {
+    fs.rmdirSync(lockDir);
+  } catch {
+    // Already released or reclaimed as stale; nothing left to do.
+  }
+}
+
+export function withLockSync(lockDir, fn, options = {}) {
+  acquireLockSync(lockDir, options);
+  try {
+    return fn();
+  } finally {
+    releaseLock(lockDir);
+  }
+}
+
+export async function withLock(lockDir, fn, options = {}) {
+  await acquireLock(lockDir, options);
+  try {
+    return await fn();
+  } finally {
+    releaseLock(lockDir);
+  }
+}

plugins/codex/scripts/lib/process.mjs

@@ -16,7 +16,9 @@ export function runCommand(command, args = [], options = {}) {
   return {
     command,
     args,
-    status: result.status ?? 0,
+    // A null status with a signal means the process was killed; never report
+    // signal-terminated commands as exit 0.
+    status: result.status ?? (result.signal != null ? 1 : 0),
     signal: result.signal ?? null,
     stdout: result.stdout ?? "",
     stderr: result.stderr ?? "",

plugins/codex/scripts/lib/state.mjs

@@ -3,6 +3,7 @@ import fs from "node:fs";
 import os from "node:os";
 import path from "node:path";
 
+import { withLockSync } from "./locking.mjs";
 import { resolveWorkspaceRoot } from "./workspace.mjs";
 
 const STATE_VERSION = 1;
@@ -89,9 +90,29 @@ function removeFileIfExists(filePath) {
   }
 }
 
-export function saveState(cwd, state) {
+function resolveStateLockDir(cwd) {
+  return path.join(resolveStateDir(cwd), ".state.lock");
+}
+
+function writeStateFileAtomic(stateFile, nextState) {
+  const tmpFile = `${stateFile}.${process.pid}.${Math.random().toString(36).slice(2)}.tmp`;
+  try {
+    const fd = fs.openSync(tmpFile, "w");
+    try {
+      fs.writeFileSync(fd, `${JSON.stringify(nextState, null, 2)}\n`, "utf8");
+      fs.fsyncSync(fd);
+    } finally {
+      fs.closeSync(fd);
+    }
+    fs.renameSync(tmpFile, stateFile);
+  } catch (error) {
+    removeFileIfExists(tmpFile);
+    throw error;
+  }
+}
+
+function saveStateLocked(cwd, state) {
   const previousJobs = loadState(cwd).jobs;
-  ensureStateDir(cwd);
   const nextJobs = pruneJobs(state.jobs ?? []);
   const nextState = {
     version: STATE_VERSION,
@@ -111,14 +132,22 @@ export function saveState(cwd, state) {
     removeFileIfExists(job.logFile);
   }
 
-  fs.writeFileSync(resolveStateFile(cwd), `${JSON.stringify(nextState, null, 2)}\n`, "utf8");
+  writeStateFileAtomic(resolveStateFile(cwd), nextState);
   return nextState;
 }
 
+export function saveState(cwd, state) {
+  ensureStateDir(cwd);
+  return withLockSync(resolveStateLockDir(cwd), () => saveStateLocked(cwd, state));
+}
+
 export function updateState(cwd, mutate) {
-  const state = loadState(cwd);
-  mutate(state);
-  return saveState(cwd, state);
+  ensureStateDir(cwd);
+  return withLockSync(resolveStateLockDir(cwd), () => {
+    const state = loadState(cwd);
+    mutate(state);
+    return saveStateLocked(cwd, state);
+  });
 }
 
 export function generateJobId(prefix = "job") {

tests/broker-lifecycle.test.mjs

@@ -0,0 +1,39 @@
+import test from "node:test";
+import assert from "node:assert/strict";
+
+import { buildEnv, installFakeCodex } from "./fake-codex-fixture.mjs";
+import { makeTempDir } from "./helpers.mjs";
+import {
+  ensureBrokerSession,
+  loadBrokerSession,
+  sendBrokerShutdown
+} from "../plugins/codex/scripts/lib/broker-lifecycle.mjs";
+
+test("concurrent ensureBrokerSession calls share a single broker", async () => {
+  const workspace = makeTempDir();
+  const binDir = makeTempDir();
+  installFakeCodex(binDir);
+  const env = buildEnv(binDir);
+
+  let first = null;
+  try {
+    const [left, right] = await Promise.all([
+      ensureBrokerSession(workspace, { env }),
+      ensureBrokerSession(workspace, { env })
+    ]);
+    first = left ?? right;
+
+    assert.ok(left, "first ensureBrokerSession returned no session");
+    assert.ok(right, "second ensureBrokerSession returned no session");
+    assert.equal(left.endpoint, right.endpoint);
+    assert.equal(left.pid, right.pid);
+
+    const persisted = loadBrokerSession(workspace);
+    assert.ok(persisted);
+    assert.equal(persisted.endpoint, left.endpoint);
+  } finally {
+    if (first?.endpoint) {
+      await sendBrokerShutdown(first.endpoint);
+    }
+  }
+});

tests/helpers.mjs

@@ -3,11 +3,89 @@ import os from "node:os";
 import path from "node:path";
 import process from "node:process";
 import { spawnSync } from "node:child_process";
+import { after } from "node:test";
+
+import { loadBrokerSession, sendBrokerShutdown, teardownBrokerSession } from "../plugins/codex/scripts/lib/broker-lifecycle.mjs";
+import { terminateProcessTree } from "../plugins/codex/scripts/lib/process.mjs";
+
+const trackedTempDirs = [];
 
 export function makeTempDir(prefix = "codex-plugin-test-") {
-  return fs.mkdtempSync(path.join(os.tmpdir(), prefix));
+  const dir = fs.mkdtempSync(path.join(os.tmpdir(), prefix));
+  trackedTempDirs.push(dir);
+  return dir;
+}
+
+function sleep(ms) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+function pidAlive(pid) {
+  try {
+    process.kill(pid, 0);
+    return true;
+  } catch {
+    return false;
+  }
 }
 
+function readBrokerPid(session) {
+  if (Number.isFinite(session.pid)) {
+    return session.pid;
+  }
+  try {
+    const parsed = Number.parseInt(fs.readFileSync(session.pidFile, "utf8").trim(), 10);
+    return Number.isFinite(parsed) ? parsed : null;
+  } catch {
+    return null;
+  }
+}
+
+// Global teardown: any broker a test started (directly or lazily) and did not
+// tear down is shut down here, so the suite never leaves broker/app-server
+// processes behind — even when a test fails or returns early.
+after(async () => {
+  const leakedPids = [];
+  const leakedSessions = [];
+  for (const dir of trackedTempDirs) {
+    const session = loadBrokerSession(dir);
+    if (!session) {
+      continue;
+    }
+    leakedSessions.push(session);
+    if (session.endpoint) {
+      await sendBrokerShutdown(session.endpoint);
+    }
+    const pid = readBrokerPid(session);
+    if (Number.isFinite(pid)) {
+      leakedPids.push(pid);
+    }
+  }
+
+  const deadline = Date.now() + 3000;
+  for (const pid of leakedPids) {
+    while (pidAlive(pid) && Date.now() < deadline) {
+      await sleep(50);
+    }
+    if (pidAlive(pid)) {
+      terminateProcessTree(pid);
+      await sleep(200);
+    }
+    if (pidAlive(pid)) {
+      throw new Error(`Leaked broker process ${pid} survived suite teardown.`);
+    }
+  }
+
+  for (const session of leakedSessions) {
+    teardownBrokerSession({
+      endpoint: session.endpoint ?? null,
+      pidFile: session.pidFile ?? null,
+      logFile: session.logFile ?? null,
+      sessionDir: session.sessionDir ?? null
+    });
+  }
+});
+
 export function writeExecutable(filePath, source) {
   fs.writeFileSync(filePath, source, { encoding: "utf8", mode: 0o755 });
 }

tests/process.test.mjs

@@ -1,7 +1,24 @@
+import process from "node:process";
 import test from "node:test";
 import assert from "node:assert/strict";
 
-import { terminateProcessTree } from "../plugins/codex/scripts/lib/process.mjs";
+import { runCommand, runCommandChecked, terminateProcessTree } from "../plugins/codex/scripts/lib/process.mjs";
+
+const SELF_TERMINATING_SCRIPT = "process.kill(process.pid, 'SIGTERM'); setInterval(() => {}, 1000);";
+
+test("runCommand reports a signal-terminated process as a failure", { skip: process.platform === "win32" }, () => {
+  const result = runCommand(process.execPath, ["-e", SELF_TERMINATING_SCRIPT]);
+
+  assert.equal(result.signal, "SIGTERM");
+  assert.notEqual(result.status, 0);
+});
+
+test("runCommandChecked throws when the process dies from a signal", { skip: process.platform === "win32" }, () => {
+  assert.throws(
+    () => runCommandChecked(process.execPath, ["-e", SELF_TERMINATING_SCRIPT]),
+    /signal=SIGTERM/
+  );
+});
 
 test("terminateProcessTree uses taskkill on Windows", () => {
   let captured = null;

tests/runtime.test.mjs

@@ -2233,6 +2233,16 @@ test("status reports shared session runtime when a lazy broker is active", () =>
 
   assert.equal(result.status, 0, result.stderr);
   assert.match(result.stdout, /Session runtime: shared session/);
+
+  const cleanup = run("node", [SESSION_HOOK, "SessionEnd"], {
+    cwd: repo,
+    env: buildEnv(binDir),
+    input: JSON.stringify({
+      hook_event_name: "SessionEnd",
+      cwd: repo
+    })
+  });
+  assert.equal(cleanup.status, 0, cleanup.stderr);
 });
 
 test("setup and status honor --cwd when reading shared session runtime", () => {