Reap leaked broker sessions across working directories

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

@@ -2,12 +2,15 @@
 
 import fs from "node:fs";
 import net from "node:net";
+import os from "node:os";
 import path from "node:path";
 import process from "node:process";
 
 import { parseArgs } from "./lib/args.mjs";
 import { BROKER_BUSY_RPC_CODE, CodexAppServerClient } from "./lib/app-server.mjs";
 import { parseBrokerEndpoint } from "./lib/broker-endpoint.mjs";
+import { clearBrokerSession, loadBrokerSession } from "./lib/broker-lifecycle.mjs";
+import { terminateProcessTree } from "./lib/process.mjs";
 
 const STREAMING_METHODS = new Set(["turn/start", "review/start", "thread/compact/start"]);
 
@@ -48,11 +51,12 @@ 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>] [--managed-session-dir]");
   }
 
   const { options } = parseArgs(argv, {
-    valueOptions: ["cwd", "pid-file", "endpoint"]
+    valueOptions: ["cwd", "pid-file", "endpoint"],
+    booleanOptions: ["managed-session-dir"]
   });
 
   if (!options.endpoint) {
@@ -63,14 +67,64 @@ async function main() {
   const endpoint = String(options.endpoint);
   const listenTarget = parseBrokerEndpoint(endpoint);
   const pidFile = options["pid-file"] ? path.resolve(options["pid-file"]) : null;
+  const managedSessionDir = options["managed-session-dir"] === true;
   writePidFile(pidFile);
 
   const appClient = await CodexAppServerClient.connect(cwd, { disableBroker: true });
   let activeRequestSocket = null;
   let activeStreamSocket = null;
   let activeStreamThreadIds = null;
+  let inFlightRequests = 0;
+  let shuttingDown = false;
+  let shutdownPromise = null;
+  let serverRef = null;
   const sockets = new Set();
 
+  // Bound each forwarded request so a hung app server cannot pin
+  // inFlightRequests forever, which would permanently disarm idle
+  // self-shutdown and make the broker unkillable by anything but a signal.
+  // Requests resolve at acceptance (streams ride notifications), so the
+  // default is generous; <= 0 disables the bound.
+  const requestTimeoutRaw = (process.env.CODEX_BROKER_REQUEST_TIMEOUT_MS ?? "").trim();
+  const requestTimeoutMs = /^-?\d+$/.test(requestTimeoutRaw)
+    ? Math.min(Number(requestTimeoutRaw), 2 ** 31 - 1)
+    : 10 * 60 * 1000;
+
+  // Forward a request to the app server while counting it as in-flight, so idle
+  // self-shutdown can never fire while real work is running, even if the calling
+  // client disconnected mid-request (which clears activeRequestSocket).
+  async function forwardAppRequest(method, params) {
+    inFlightRequests += 1;
+    try {
+      const request = appClient.request(method, params);
+      if (requestTimeoutMs <= 0) {
+        return await request;
+      }
+      return await Promise.race([
+        request,
+        new Promise((_, reject) => {
+          const timer = setTimeout(() => {
+            // The race cannot cancel the underlying request: releasing
+            // ownership while it might still execute would let a clientless
+            // turn keep running (and route its notifications to a later
+            // client). A request unanswered for this long means the app
+            // server is wedged, so terminate the whole broker instead —
+            // shutdown closes the app-server child with it, and callers
+            // respawn a fresh broker on demand.
+            reject(new Error(`Shared broker request ${method} timed out after ${requestTimeoutMs}ms; broker shutting down.`));
+            setTimeout(() => process.exit(1), 5000).unref(); // backstop if cleanup hangs
+            shutdown(serverRef).finally(() => process.exit(1));
+          }, requestTimeoutMs);
+          timer.unref();
+          request.finally(() => clearTimeout(timer)).catch(() => {});
+        })
+      ]);
+    } finally {
+      inFlightRequests -= 1;
+      scheduleIdleShutdown();
+    }
+  }
+
   function clearSocketOwnership(socket) {
     if (activeRequestSocket === socket) {
       activeRequestSocket = null;
@@ -95,27 +149,169 @@ async function main() {
         if (activeRequestSocket === target) {
           activeRequestSocket = null;
         }
+        // Releasing stream ownership can be the last activity on this broker;
+        // without rearming here an abandoned cwd would never become idle.
+        scheduleIdleShutdown();
       }
     }
   }
 
-  async function shutdown(server) {
+  // Every shutdown path shares one promise: a second caller (e.g. SIGTERM
+  // landing during an idle shutdown) awaits the same cleanup instead of
+  // returning early and letting its process.exit() abort the first caller's
+  // cleanup mid-flight.
+  function shutdown(server) {
+    if (!shutdownPromise) {
+      shuttingDown = true;
+      shutdownPromise = performShutdown(server);
+    }
+    return shutdownPromise;
+  }
+
+  async function performShutdown(server) {
+    // Whole-shutdown backstop: whatever below wedges, this process ends. The
+    // state record is cleared first, so a replacement broker is never blocked
+    // on this one finishing its cleanup.
+    setTimeout(() => process.exit(1), 15000).unref();
+    // Retire this broker's state record first, while its socket is still the
+    // live one for this cwd: no replacement broker can have been spawned yet,
+    // so the guarded clear cannot race a newer record, and clients probing
+    // from here on fall back to starting a fresh broker instead of connecting
+    // to a dying one. Guarded on the endpoint matching, so a record that was
+    // already replaced (e.g. this broker was deemed unresponsive) is kept.
+    try {
+      if (loadBrokerSession(cwd)?.endpoint === endpoint) {
+        clearBrokerSession(cwd);
+      }
+    } catch {
+      // Ignore unreadable or already-removed state records.
+    }
     for (const socket of sockets) {
       socket.end();
     }
-    await appClient.close().catch(() => {});
-    await new Promise((resolve) => server.close(resolve));
-    if (listenTarget.kind === "unix" && fs.existsSync(listenTarget.path)) {
-      fs.unlinkSync(listenTarget.path);
+    // Bound the app-server close: close() escalates only as far as SIGTERM,
+    // so a child that ignores it would wedge this shutdown indefinitely
+    // (record already cleared, session dir still present, both processes
+    // alive while the next command spawns a replacement). After the deadline,
+    // deliver a real force-kill; the child's exit settles the dangling close.
+    // On POSIX that must be SIGKILL to the pid itself: the child is not
+    // detached, so a process-group signal (kill(-pid)) hits ESRCH and
+    // terminateProcessTree delivers nothing there. On Windows the child is a
+    // cmd.exe wrapper, so the taskkill /T /F tree kill is the right tool.
+    await Promise.race([
+      appClient.close().catch(() => {}),
+      new Promise((resolve) => {
+        const timer = setTimeout(() => {
+          try {
+            if (appClient.proc && appClient.proc.exitCode === null) {
+              if (process.platform === "win32") {
+                terminateProcessTree(appClient.proc.pid);
+              } else {
+                appClient.proc.kill("SIGKILL");
+              }
+            }
+          } catch {
+            // Best effort; the whole-shutdown backstop above still applies.
+          }
+          resolve();
+        }, 5000);
+        timer.unref();
+      })
+    ]);
+    if (server) {
+      await new Promise((resolve) => server.close(resolve));
     }
-    if (pidFile && fs.existsSync(pidFile)) {
-      fs.unlinkSync(pidFile);
+    // Remove the broker's own session directory (socket, pid file, log) so a
+    // clean exit leaves nothing behind for the reaper to GC. Recursive removal
+    // requires the spawner to have declared it created the directory
+    // (--managed-session-dir, set by spawnBrokerProcess after mkdtemp) AND the
+    // plugin's own layout (cxc- prefix directly under the OS temp dir), so a
+    // manual invocation can never have a caller-selected directory deleted,
+    // even one named to look like ours; only the broker's own files are
+    // unlinked there.
+    const sessionDir = pidFile
+      ? path.dirname(pidFile)
+      : listenTarget.kind === "unix"
+        ? path.dirname(listenTarget.path)
+        : null;
+    try {
+      if (sessionDir && isManagedSessionDir(sessionDir)) {
+        fs.rmSync(sessionDir, { recursive: true, force: true });
+      } else {
+        if (pidFile && fs.existsSync(pidFile)) {
+          fs.unlinkSync(pidFile);
+        }
+        if (listenTarget.kind === "unix" && fs.existsSync(listenTarget.path)) {
+          fs.unlinkSync(listenTarget.path);
+        }
+      }
+    } catch {
+      // Ignore already-removed files or directories.
+    }
+  }
+
+  function isManagedSessionDir(dir) {
+    if (!managedSessionDir) {
+      return false;
+    }
+    try {
+      return (
+        path.basename(dir).startsWith("cxc-") &&
+        fs.realpathSync(path.dirname(dir)) === fs.realpathSync(os.tmpdir())
+      );
+    } catch {
+      return false;
     }
   }
 
   appClient.setNotificationHandler(routeNotification);
 
+  // Idle self-shutdown: a broker is spawned per working directory and is reused
+  // across sessions, so no external actor can safely decide it is done. Instead
+  // the broker exits itself once it has had no connections and no in-flight work
+  // for CODEX_BROKER_IDLE_MS (default 30 min; <= 0 disables). Callers respawn one
+  // on demand, so exiting when idle is safe and stops brokers from accumulating.
+  // Strict integer parsing: parseInt would truncate "30m" to 30ms and accept
+  // scientific notation, silently inverting an "effectively never" intent into
+  // near-instant shutdown. Malformed values fall back to the default, and the
+  // value is clamped below Node's 2^31-1 setTimeout ceiling (beyond it the
+  // timer fires after 1ms).
+  const idleRaw = (process.env.CODEX_BROKER_IDLE_MS ?? "").trim();
+  const idleTimeoutMs = /^-?\d+$/.test(idleRaw)
+    ? Math.min(Number(idleRaw), 2 ** 31 - 1)
+    : 30 * 60 * 1000;
+  let idleTimer = null;
+  function cancelIdleShutdown() {
+    if (idleTimer) {
+      clearTimeout(idleTimer);
+      idleTimer = null;
+    }
+  }
+  function isIdle() {
+    return sockets.size === 0 && inFlightRequests === 0 && !activeRequestSocket && !activeStreamSocket;
+  }
+  function scheduleIdleShutdown() {
+    cancelIdleShutdown();
+    if (idleTimeoutMs <= 0 || !isIdle()) {
+      return;
+    }
+    idleTimer = setTimeout(() => {
+      if (isIdle()) {
+        shutdown(server).finally(() => process.exit(0));
+      }
+    }, idleTimeoutMs);
+    idleTimer.unref(); // never keep the process alive solely to fire this timer
+  }
+
   const server = net.createServer((socket) => {
+    if (shuttingDown) {
+      // A connection that lands between shutdown starting and server.close()
+      // taking effect would otherwise hold the close (and the process exit)
+      // open until the client goes away on its own.
+      socket.destroy();
+      return;
+    }
+    cancelIdleShutdown();
     sockets.add(socket);
     socket.setEncoding("utf8");
     let buffer = "";
@@ -183,7 +379,7 @@ async function main() {
 
         if (allowInterruptDuringActiveStream) {
           try {
-            const result = await appClient.request(message.method, message.params ?? {});
+            const result = await forwardAppRequest(message.method, message.params ?? {});
             send(socket, { id: message.id, result });
           } catch (error) {
             send(socket, {
@@ -198,9 +394,13 @@ async function main() {
         activeRequestSocket = socket;
 
         try {
-          const result = await appClient.request(message.method, message.params ?? {});
+          const result = await forwardAppRequest(message.method, message.params ?? {});
           send(socket, { id: message.id, result });
-          if (isStreaming) {
+          // A socket that disconnected during the await already had its
+          // ownership cleared by the close handler; assigning it here would
+          // strand a dead socket in activeStreamSocket, busy-rejecting other
+          // clients and keeping isIdle() false with no event left to clear it.
+          if (isStreaming && !socket.destroyed) {
             activeStreamSocket = socket;
             activeStreamThreadIds = buildStreamThreadIds(message.method, message.params ?? {}, result);
           }
@@ -225,14 +425,30 @@ async function main() {
     socket.on("close", () => {
       sockets.delete(socket);
       clearSocketOwnership(socket);
+      scheduleIdleShutdown();
     });
 
     socket.on("error", () => {
       sockets.delete(socket);
       clearSocketOwnership(socket);
+      scheduleIdleShutdown();
     });
   });
 
+  serverRef = server;
+
+  // If the codex app-server child exits or its connection is lost, this broker
+  // can never serve another request, but its socket keeps accepting: endpoint
+  // probes pass, the reaper skips the live pid, and nothing external kills
+  // brokers anymore. Exit instead; callers respawn a fresh broker on demand.
+  Promise.resolve(appClient.exitPromise)
+    .catch(() => {})
+    .then(() => {
+      if (!shuttingDown) {
+        shutdown(server).finally(() => process.exit(1));
+      }
+    });
+
   process.on("SIGTERM", async () => {
     await shutdown(server);
     process.exit(0);
@@ -243,7 +459,9 @@ async function main() {
     process.exit(0);
   });
 
-  server.listen(listenTarget.path);
+  server.listen(listenTarget.path, () => {
+    scheduleIdleShutdown(); // exit if nobody ever connects
+  });
 }
 
 main().catch((error) => {

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

@@ -13,7 +13,7 @@ import process from "node:process";
 import { spawn } from "node:child_process";
 import readline from "node:readline";
 import { parseBrokerEndpoint } from "./broker-endpoint.mjs";
-import { ensureBrokerSession, loadBrokerSession } from "./broker-lifecycle.mjs";
+import { ensureBrokerSession, loadBrokerSession, waitForBrokerEndpoint } from "./broker-lifecycle.mjs";
 import { terminateProcessTree } from "./process.mjs";
 
 const PLUGIN_MANIFEST_URL = new URL("../../.claude-plugin/plugin.json", import.meta.url);
@@ -338,7 +338,14 @@ export class CodexAppServerClient {
     if (!options.disableBroker) {
       brokerEndpoint = options.brokerEndpoint ?? options.env?.[BROKER_ENDPOINT_ENV] ?? process.env[BROKER_ENDPOINT_ENV] ?? null;
       if (!brokerEndpoint && options.reuseExistingBroker) {
-        brokerEndpoint = loadBrokerSession(cwd)?.endpoint ?? null;
+        // Validate the recorded endpoint before reusing it: the record can be
+        // stale when the broker died without a clean exit. A stale record then
+        // behaves like no record, and the caller falls back to a direct app
+        // server instead of erroring against a dead socket.
+        const recordedEndpoint = loadBrokerSession(cwd)?.endpoint ?? null;
+        if (recordedEndpoint && (await waitForBrokerEndpoint(recordedEndpoint, 150).catch(() => false))) {
+          brokerEndpoint = recordedEndpoint;
+        }
       }
       if (!brokerEndpoint && !options.reuseExistingBroker) {
         const brokerSession = await ensureBrokerSession(cwd, { env: options.env });