fix(broker): shut down the app-server broker after an idle timeout

plugins/codex/scripts/app-server-broker.mjs

@@ -11,6 +11,42 @@ import { parseBrokerEndpoint } from "./lib/broker-endpoint.mjs";
 
 const STREAMING_METHODS = new Set(["turn/start", "review/start", "thread/compact/start"]);
 
+// A broker outlives the client that spawned it: nothing in the protocol tells it the
+// client is gone for good, so without this it stays resident forever holding its socket
+// dir. Idle shutdown is decided BY THE BROKER because it is the only party that can see
+// whether it is serving anyone -- an external sweep cannot, and racing one is unsafe.
+const DEFAULT_IDLE_TIMEOUT_MS = 30 * 60 * 1000;
+const IDLE_TIMEOUT_ENV = "CODEX_COMPANION_BROKER_IDLE_MS";
+// setTimeout() overflows above 2^31-1 ms: it warns and then fires after 1ms, so an
+// over-large timeout would shut the broker down almost immediately -- the exact opposite
+// of what was asked for. Reject instead, so the mistake is visible at startup.
+const MAX_IDLE_TIMEOUT_MS = 2147483647;
+
+function resolveIdleTimeoutMs(rawOption, env = {}) {
+  const raw = rawOption ?? env[IDLE_TIMEOUT_ENV];
+  // Trim before the emptiness test: Number("  ") is 0, so a blank or whitespace-only
+  // value would otherwise DISABLE idle shutdown silently. Explicit "0" is the only
+  // way to turn it off; anything blank falls back to the default.
+  const text = raw === undefined || raw === null ? "" : String(raw).trim();
+  if (text === "") {
+    return DEFAULT_IDLE_TIMEOUT_MS;
+  }
+  const parsed = Number(text);
+  // Reject rather than silently falling back: a typo'd timeout that quietly became
+  // "never expire" would reintroduce the exact leak this exists to close.
+  if (!Number.isFinite(parsed) || parsed < 0) {
+    throw new Error(
+      `Invalid idle timeout ${JSON.stringify(text)}: expected a non-negative number of milliseconds.`
+    );
+  }
+  if (parsed > MAX_IDLE_TIMEOUT_MS) {
+    throw new Error(
+      `Invalid idle timeout ${JSON.stringify(text)}: must be at most ${MAX_IDLE_TIMEOUT_MS} ms (Node timer limit).`
+    );
+  }
+  return parsed;
+}
+
 function buildStreamThreadIds(method, params, result) {
   const threadIds = new Set();
   if (params?.threadId) {
@@ -48,11 +84,13 @@ function writePidFile(pidFile) {
 async function main() {
   const [subcommand, ...argv] = process.argv.slice(2);
   if (subcommand !== "serve") {
-    throw new Error("Usage: node scripts/app-server-broker.mjs serve --endpoint <value> [--cwd <path>] [--pid-file <path>]");
+    throw new Error(
+      "Usage: node scripts/app-server-broker.mjs serve --endpoint <value> [--cwd <path>] [--pid-file <path>] [--idle-timeout <ms>]"
+    );
   }
 
   const { options } = parseArgs(argv, {
-    valueOptions: ["cwd", "pid-file", "endpoint"]
+    valueOptions: ["cwd", "pid-file", "endpoint", "idle-timeout"]
   });
 
   if (!options.endpoint) {
@@ -63,13 +101,50 @@ async function main() {
   const endpoint = String(options.endpoint);
   const listenTarget = parseBrokerEndpoint(endpoint);
   const pidFile = options["pid-file"] ? path.resolve(options["pid-file"]) : null;
+  const idleTimeoutMs = resolveIdleTimeoutMs(options["idle-timeout"], process.env);
   writePidFile(pidFile);
 
   const appClient = await CodexAppServerClient.connect(cwd, { disableBroker: true });
   let activeRequestSocket = null;
   let activeStreamSocket = null;
   let activeStreamThreadIds = null;
   const sockets = new Set();
+  let idleTimer = null;
+
+  // Idle means nobody is connected AND nothing is in flight. Holding an open socket is
+  // enough to keep the broker alive, so a long streaming turn can never be cut short.
+  function isIdle() {
+    return sockets.size === 0 && activeRequestSocket === null && activeStreamSocket === null;
+  }
+
+  function disarmIdleTimer() {
+    if (idleTimer) {
+      clearTimeout(idleTimer);
+      idleTimer = null;
+    }
+  }
+
+  function armIdleTimer() {
+    disarmIdleTimer();
+    if (idleTimeoutMs <= 0 || !isIdle()) {
+      return;
+    }
+    idleTimer = setTimeout(() => {
+      idleTimer = null;
+      // Re-check at fire time: a client may have connected while the timer was pending.
+      if (!isIdle()) {
+        armIdleTimer();
+        return;
+      }
+      void shutdown(server).then(
+        () => process.exit(0),
+        () => process.exit(0)
+      );
+    }, idleTimeoutMs);
+    // The listening server keeps the event loop alive; the timer must not do so itself,
+    // or a broker with idle shutdown disabled could never exit cleanly.
+    idleTimer.unref();
+  }
 
   function clearSocketOwnership(socket) {
     if (activeRequestSocket === socket) {
@@ -100,11 +175,16 @@ async function main() {
   }
 
   async function shutdown(server) {
+    // Stop accepting FIRST. server.close() stops listening immediately and resolves once
+    // existing connections drain. Tearing down the app server first would leave the
+    // endpoint accepting throughout that await, so ensureBrokerSession's readiness probe
+    // could connect, judge a shutting-down broker "ready", and then lose the connection.
+    const closed = new Promise((resolve) => server.close(resolve));
     for (const socket of sockets) {
       socket.end();
     }
+    await closed;
     await appClient.close().catch(() => {});
-    await new Promise((resolve) => server.close(resolve));
     if (listenTarget.kind === "unix" && fs.existsSync(listenTarget.path)) {
       fs.unlinkSync(listenTarget.path);
     }
@@ -116,6 +196,7 @@ async function main() {
   appClient.setNotificationHandler(routeNotification);
 
   const server = net.createServer((socket) => {
+    disarmIdleTimer();
     sockets.add(socket);
     socket.setEncoding("utf8");
     let buffer = "";
@@ -225,11 +306,13 @@ async function main() {
     socket.on("close", () => {
       sockets.delete(socket);
       clearSocketOwnership(socket);
+      armIdleTimer();
     });
 
     socket.on("error", () => {
       sockets.delete(socket);
       clearSocketOwnership(socket);
+      armIdleTimer();
     });
   });
 
@@ -243,7 +326,9 @@ async function main() {
     process.exit(0);
   });
 
-  server.listen(listenTarget.path);
+  server.listen(listenTarget.path, () => {
+    armIdleTimer();
+  });
 }
 
 main().catch((error) => {

tests/broker-idle.test.mjs

@@ -0,0 +1,181 @@
+import assert from "node:assert/strict";
+import fs from "node:fs";
+import net from "node:net";
+import path from "node:path";
+import test from "node:test";
+import { spawn } from "node:child_process";
+import { fileURLToPath } from "node:url";
+
+import { makeTempDir } from "./helpers.mjs";
+import { buildEnv, installFakeCodex } from "./fake-codex-fixture.mjs";
+
+const ROOT = path.resolve(fileURLToPath(new URL("..", import.meta.url)));
+const BROKER = path.join(ROOT, "plugins", "codex", "scripts", "app-server-broker.mjs");
+
+function delay(ms) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+async function waitFor(predicate, { timeoutMs = 10000, intervalMs = 25 } = {}) {
+  const start = Date.now();
+  while (Date.now() - start < timeoutMs) {
+    if (await predicate()) {
+      return true;
+    }
+    await delay(intervalMs);
+  }
+  return false;
+}
+
+function startBroker({ idleTimeout } = {}) {
+  const binDir = makeTempDir();
+  installFakeCodex(binDir);
+  const sessionDir = makeTempDir("codex-broker-idle-");
+  const cwd = makeTempDir("codex-broker-cwd-");
+  const socketPath = path.join(sessionDir, "broker.sock");
+  const pidFile = path.join(sessionDir, "broker.pid");
+
+  const args = [BROKER, "serve", "--endpoint", `unix:${socketPath}`, "--cwd", cwd, "--pid-file", pidFile];
+  if (idleTimeout !== undefined) {
+    args.push("--idle-timeout", String(idleTimeout));
+  }
+
+  const child = spawn(process.execPath, args, { env: buildEnv(binDir), stdio: ["ignore", "pipe", "pipe"] });
+  let stderr = "";
+  child.stderr.on("data", (chunk) => {
+    stderr += chunk;
+  });
+
+  const exited = new Promise((resolve) => {
+    child.on("exit", (code, signal) => resolve({ code, signal }));
+  });
+
+  const alive = () => child.exitCode === null && child.signalCode === null;
+
+  // Bounded: a broker that never exits must turn into a RED test, not a hung run.
+  // Awaiting `exited` unbounded is what made this file hang against an unpatched broker.
+  const exitedWithin = (timeoutMs) =>
+    Promise.race([exited, delay(timeoutMs).then(() => null)]);
+
+  const dispose = () => {
+    if (alive()) {
+      child.kill("SIGKILL");
+    }
+  };
+
+  return {
+    child,
+    socketPath,
+    pidFile,
+    exited,
+    exitedWithin,
+    dispose,
+    stderr: () => stderr,
+    alive,
+    listening: () => waitFor(() => fs.existsSync(socketPath))
+  };
+}
+
+function connectTo(socketPath) {
+  return new Promise((resolve, reject) => {
+    const socket = net.createConnection({ path: socketPath });
+    socket.on("connect", () => resolve(socket));
+    socket.on("error", reject);
+  });
+}
+
+test("broker shuts itself down once it has been idle for the timeout", async (t) => {
+  const broker = startBroker({ idleTimeout: 400 });
+  t.after(() => broker.dispose());
+  assert.equal(await broker.listening(), true, `broker never listened: ${broker.stderr()}`);
+
+  const result = await broker.exitedWithin(8000);
+  assert.ok(result, `broker never exited after the idle timeout, stderr: ${broker.stderr()}`);
+  assert.equal(result.code, 0, `expected a clean idle exit, stderr: ${broker.stderr()}`);
+  // shutdown() must still clean up after itself on the idle path, or the next
+  // ensureBrokerSession would find a stale socket and a stale pidfile.
+  assert.equal(fs.existsSync(broker.socketPath), false, "idle shutdown left the socket behind");
+  assert.equal(fs.existsSync(broker.pidFile), false, "idle shutdown left the pidfile behind");
+});
+
+test("broker stays alive while a client is connected, then exits after it disconnects", async (t) => {
+  const broker = startBroker({ idleTimeout: 400 });
+  t.after(() => broker.dispose());
+  assert.equal(await broker.listening(), true, `broker never listened: ${broker.stderr()}`);
+
+  const socket = await connectTo(broker.socketPath);
+  // Well past the idle timeout: an open connection must hold the broker open, which is
+  // what stops a long streaming turn from being cut off mid-flight.
+  await delay(1600);
+  assert.equal(broker.alive(), true, "broker exited while a client was still connected");
+
+  socket.destroy();
+  const result = await broker.exitedWithin(8000);
+  assert.ok(result, `broker never exited after the client disconnected, stderr: ${broker.stderr()}`);
+  assert.equal(result.code, 0, `expected a clean idle exit after disconnect, stderr: ${broker.stderr()}`);
+});
+
+test("broker with --idle-timeout 0 never idles out", async (t) => {
+  const broker = startBroker({ idleTimeout: 0 });
+  t.after(() => broker.dispose());
+  assert.equal(await broker.listening(), true, `broker never listened: ${broker.stderr()}`);
+
+  await delay(1600);
+  assert.equal(broker.alive(), true, "idle shutdown ran even though it was disabled");
+
+  // SIGTERM shares shutdown() with the idle path, so this also covers the ordering there.
+  broker.child.kill("SIGTERM");
+  const result = await broker.exitedWithin(8000);
+  assert.ok(result, "broker did not exit on SIGTERM");
+  assert.equal(fs.existsSync(broker.socketPath), false, "SIGTERM shutdown left the socket behind");
+  assert.equal(fs.existsSync(broker.pidFile), false, "SIGTERM shutdown left the pidfile behind");
+});
+
+test("broker rejects a non-numeric --idle-timeout instead of silently never expiring", async (t) => {
+  const broker = startBroker({ idleTimeout: "not-a-number" });
+  t.after(() => broker.dispose());
+  const result = await broker.exitedWithin(8000);
+  assert.ok(result, "broker did not exit on an invalid idle timeout");
+  assert.equal(result.code, 1);
+  assert.match(broker.stderr(), /Invalid idle timeout/);
+});
+
+test("blank --idle-timeout falls back to the default instead of silently disabling", async (t) => {
+  // Number("  ") === 0, so a whitespace value must NOT be read as "never expire".
+  // Only an explicit 0 disables idle shutdown.
+  const broker = startBroker({ idleTimeout: "  " });
+  t.after(() => broker.dispose());
+  assert.equal(await broker.listening(), true, `broker never listened: ${broker.stderr()}`);
+
+  // The default is 30 minutes, so it must still be alive well past the short timeouts
+  // used elsewhere in this file.
+  await delay(1600);
+  assert.equal(broker.alive(), true, "a blank idle timeout disabled idle shutdown");
+
+  broker.child.kill("SIGTERM");
+  await broker.exitedWithin(8000);
+});
+
+test("broker rejects an idle timeout beyond Node's timer range", async (t) => {
+  // setTimeout() overflows above 2^31-1 and fires after 1ms, which would shut the broker
+  // down almost immediately. Asking for ~30 days must fail loudly, not silently invert.
+  const broker = startBroker({ idleTimeout: 30 * 24 * 60 * 60 * 1000 });
+  t.after(() => broker.dispose());
+
+  const result = await broker.exitedWithin(8000);
+  assert.ok(result, "broker did not exit on an out-of-range idle timeout");
+  assert.equal(result.code, 1);
+  assert.match(broker.stderr(), /must be at most 2147483647 ms/);
+});
+
+test("broker accepts an idle timeout exactly at the Node timer limit", async (t) => {
+  // Boundary: 2147483647 is valid, so the guard must not be off by one.
+  const broker = startBroker({ idleTimeout: 2147483647 });
+  t.after(() => broker.dispose());
+
+  assert.equal(await broker.listening(), true, `broker never listened: ${broker.stderr()}`);
+  assert.equal(broker.alive(), true, "broker rejected a timeout that is exactly at the limit");
+
+  broker.child.kill("SIGTERM");
+  await broker.exitedWithin(8000);
+});