fix(codex): coalesce prompt probes and cancel abandoned probes

src/codex/prompt-text-probe.ts

@@ -100,41 +100,236 @@ function resolveCodexBinary(): string | null {
 /** 8 MiB is far above any real prompt and far below anything that hurts the server. */
 const MAX_PROBE_OUTPUT_BYTES = 8 * 1024 * 1024;
 
-function runProbe(binary: string, cwd: string, timeoutMs: number): Promise<string | null> {
-  return new Promise(resolve => {
+interface ProbeCommand {
+  binary: string;
+  args: string[];
+  cwd: string;
+  timeoutMs: number;
+  promptRevision: string | null;
+}
+
+interface PromptProbeFlight {
+  key: string;
+  controller: AbortController;
+  result: Promise<string | null>;
+  closed: Promise<void>;
+  waiters: number;
+  joinable: boolean;
+  resultSettled: boolean;
+  settled: boolean;
+}
+
+interface PromptProbeExecution {
+  result: Promise<string | null>;
+  closed: Promise<void>;
+}
+
+type SharedPromptProbeOutcome =
+  | { kind: "output"; raw: string }
+  | { kind: "failed" }
+  | { kind: "busy" };
+
+let activePromptProbe: PromptProbeFlight | null = null;
+let probeCommandForTests: { binary: string; args: string[] } | null = null;
+let probeSpawnAttemptsForTests = 0;
+let probeCloseBarrierForTests: Promise<void> | null = null;
+
+function commandKey(command: ProbeCommand): string {
+  return JSON.stringify([
+    command.binary,
+    command.args,
+    command.cwd,
+    command.timeoutMs,
+    command.promptRevision,
+  ]);
+}
+
+function completedExecution(value: string | null): PromptProbeExecution {
+  return { result: Promise.resolve(value), closed: Promise.resolve() };
+}
+
+function runProbe(
+  command: ProbeCommand,
+  signal: AbortSignal,
+  onStopping: () => void,
+): PromptProbeExecution {
+  if (signal.aborted) return completedExecution(null);
+  let resolveResult!: (value: string | null) => void;
+  let resolveClosed!: () => void;
+  const result = new Promise<string | null>(resolve => { resolveResult = resolve; });
+  const closed = new Promise<void>(resolve => { resolveClosed = resolve; });
+  let resultSettled = false;
+  let closeSettled = false;
+
+  const finishResult = (value: string | null) => {
+    if (resultSettled) return;
+    resultSettled = true;
+    resolveResult(value);
+  };
+  const finishClosed = () => {
+    if (closeSettled) return;
+    closeSettled = true;
+    resolveClosed();
+  };
+
+  try {
     // A probe must never hang OR balloon the management API: it is bounded in
     // time AND in bytes, and every failure degrades to "unavailable" rather than
     // an error page.
-    const child = spawn(binary, ["debug", "prompt-input"], {
-      cwd,
-      stdio: ["ignore", "pipe", "ignore"],
-    });
+    let child: ReturnType<typeof spawn>;
+    try {
+      if (probeCommandForTests) probeSpawnAttemptsForTests += 1;
+      child = spawn(command.binary, command.args, {
+        cwd: command.cwd,
+        stdio: ["ignore", "pipe", "ignore"],
+      });
+    } catch {
+      finishResult(null);
+      finishClosed();
+      return { result, closed };
+    }
     const chunks: Buffer[] = [];
     let size = 0;
     let settled = false;
-    // One settlement path: a timeout that resolved before `close` used to leave
-    // the child streaming into a buffer nobody would ever read.
-    const settle = (value: string | null) => {
+    let stopping = false;
+    let timer: ReturnType<typeof setTimeout> | undefined;
+
+    const finish = (value: string | null) => {
       if (settled) return;
       settled = true;
-      clearTimeout(timer);
+      if (timer) clearTimeout(timer);
+      signal.removeEventListener("abort", onAbort);
+      finishResult(value);
+      finishClosed();
+    };
+
+    // Keep the flight admitted until `close`: kill() only requests termination
+    // and does not prove the exact child has released its process and stdio.
+    const terminate = () => {
+      if (settled || stopping) return;
+      stopping = true;
+      onStopping();
+      if (timer) clearTimeout(timer);
+      signal.removeEventListener("abort", onAbort);
       child.stdout?.destroy();
-      if (child.exitCode === null && child.signalCode === null) child.kill("SIGKILL");
-      resolve(value);
+      // The caller is bounded even if OS termination later fails. Admission is
+      // retained separately by `closed`, and later probes fail soft while this
+      // exact child remains unproven terminal.
+      finishResult(null);
+      if (child.exitCode !== null || child.signalCode !== null) return;
+      try {
+        child.kill("SIGKILL");
+      } catch {
+        // Exact child state is ambiguous. Keep admission non-joinable until its
+        // own `close` proves terminal instead of targeting a reusable numeric PID.
+      }
     };
-    const timer = setTimeout(() => settle(null), timeoutMs);
+    const onAbort = () => terminate();
+
+    timer = setTimeout(terminate, command.timeoutMs);
+    signal.addEventListener("abort", onAbort, { once: true });
     child.stdout?.on("data", (chunk: Buffer) => {
       size += chunk.length;
-      if (size > MAX_PROBE_OUTPUT_BYTES) { settle(null); return; }
+      if (size > MAX_PROBE_OUTPUT_BYTES) { terminate(); return; }
       chunks.push(chunk);
     });
-    child.on("error", () => settle(null));
+    child.on("error", () => {
+      // No PID means spawn itself failed, so there is no live child to drain.
+      if (child.pid === undefined) {
+        finish(null);
+      }
+      else terminate();
+    });
     child.on("close", code => {
       // Decode once, at the end: `String(chunk)` per chunk corrupts any UTF-8
       // character that straddles a chunk boundary.
-      settle(code === 0 ? Buffer.concat(chunks).toString("utf8") : null);
+      const recordClose = () => {
+        finish(!stopping && code === 0 ? Buffer.concat(chunks).toString("utf8") : null);
+      };
+      const barrier = probeCloseBarrierForTests;
+      if (barrier) void barrier.then(recordClose, recordClose);
+      else recordClose();
     });
+    // Close the race between the pre-spawn check and listener registration.
+    if (signal.aborted) terminate();
+  } catch {
+    finishResult(null);
+    finishClosed();
+  }
+  return { result, closed };
+}
+
+function startPromptProbeFlight(command: ProbeCommand): PromptProbeFlight {
+  const controller = new AbortController();
+  const flight: PromptProbeFlight = {
+    key: commandKey(command),
+    controller,
+    result: Promise.resolve(null),
+    closed: Promise.resolve(),
+    waiters: 0,
+    joinable: true,
+    resultSettled: false,
+    settled: false,
+  };
+  const execution = runProbe(command, controller.signal, () => {
+    flight.joinable = false;
   });
+  flight.result = execution.result
+    .catch(() => null)
+    .finally(() => {
+      flight.resultSettled = true;
+    });
+  flight.closed = execution.closed
+    .finally(() => {
+      flight.settled = true;
+      if (activePromptProbe === flight) activePromptProbe = null;
+    });
+  activePromptProbe = flight;
+  return flight;
+}
+
+async function runSharedPromptProbe(
+  command: ProbeCommand,
+  signal?: AbortSignal,
+): Promise<SharedPromptProbeOutcome> {
+  const key = commandKey(command);
+  if (signal?.aborted) return { kind: "failed" };
+  const active = activePromptProbe;
+  if (!active) {
+    const raw = await waitForPromptProbeFlight(startPromptProbeFlight(command), signal);
+    return raw === null ? { kind: "failed" } : { kind: "output", raw };
+  }
+  if (active.key === key && active.joinable && !active.controller.signal.aborted) {
+    const raw = await waitForPromptProbeFlight(active, signal);
+    return raw === null ? { kind: "failed" } : { kind: "output", raw };
+  }
+  // A different or terminating flight still owns the sole process slot. Never
+  // wait unboundedly for an unproven close and never launch beside it.
+  return { kind: "busy" };
+}
+
+async function waitForPromptProbeFlight(flight: PromptProbeFlight, signal?: AbortSignal): Promise<string | null> {
+  if (signal?.aborted) {
+    if (flight.waiters === 0 && !flight.settled) flight.controller.abort();
+    return null;
+  }
+  flight.waiters += 1;
+  let onAbort: (() => void) | undefined;
+  try {
+    if (!signal) return await flight.result;
+    const aborted = new Promise<null>(resolve => {
+      onAbort = () => resolve(null);
+      signal.addEventListener("abort", onAbort, { once: true });
+      if (signal.aborted) onAbort();
+    });
+    return await Promise.race([flight.result, aborted]);
+  } finally {
+    if (onAbort) signal?.removeEventListener("abort", onAbort);
+    flight.waiters = Math.max(0, flight.waiters - 1);
+    if (flight.waiters === 0 && !flight.resultSettled) {
+      flight.controller.abort(new DOMException("All prompt probe callers cancelled", "AbortError"));
+    }
+  }
 }
 
 /** Pull every `<tag>...</tag>` section out of the rendered developer messages. */
@@ -182,20 +377,44 @@ export const extractSectionsForTests = extractSections;
  * `cwd` matters: AGENTS.md and environment context are directory-dependent, so a
  * probe from the wrong place would describe a prompt the user never sees.
  */
-export async function probePromptText(timeoutMs = 15_000): Promise<PromptTextProbe> {
+export async function probePromptText(
+  timeoutMs = 15_000,
+  signal?: AbortSignal,
+  promptRevision: string | null = null,
+): Promise<PromptTextProbe> {
   // The probe runs in CODEX_HOME, never in a caller-supplied directory. A `cwd`
   // parameter let an authenticated request read any readable folder's AGENTS.md,
   // and it also described a prompt that depends on where Codex happened to run.
   // The global home is the one context this page can honestly report on.
   const codexHome = resolveCodexHomeDir();
-  const binary = resolveCodexBinary();
+  if (signal?.aborted) {
+    return { ok: false, codexHome, layers: {}, detail: "prompt probe cancelled" };
+  }
+  const binary = probeCommandForTests?.binary ?? resolveCodexBinary();
   if (!binary) {
     return { ok: false, codexHome, layers: {}, detail: "codex binary not found" };
   }
-  const raw = await runProbe(binary, codexHome, timeoutMs);
-  if (raw === null) {
-    return { ok: false, codexHome, layers: {}, detail: "codex debug prompt-input failed" };
+  const command: ProbeCommand = {
+    binary,
+    args: probeCommandForTests?.args ?? ["debug", "prompt-input"],
+    cwd: codexHome,
+    timeoutMs,
+    promptRevision,
+  };
+  const outcome = await runSharedPromptProbe(command, signal);
+  if (outcome.kind !== "output") {
+    return {
+      ok: false,
+      codexHome,
+      layers: {},
+      detail: signal?.aborted
+        ? "prompt probe cancelled"
+        : outcome.kind === "busy"
+          ? "another prompt probe is still finishing; retry shortly"
+          : "codex debug prompt-input failed",
+    };
   }
+  const raw = outcome.raw;
   const sections = extractSections(raw);
   if (sections.size === 0) {
     // Zero sections from a zero-exit probe means the output did not parse, which
@@ -236,3 +455,35 @@ export async function probePromptText(timeoutMs = 15_000): Promise<PromptTextPro
   }
   return { ok: true, codexHome, layers };
 }
+
+/** Test-only command seam; production always resolves the installed Codex binary. */
+export function setPromptTextProbeCommandForTests(command: { binary: string; args: string[] } | null): void {
+  probeCommandForTests = command ? { binary: command.binary, args: [...command.args] } : null;
+}
+
+/** Test-only process-start counter for proving admission without timing guesses. */
+export function promptTextProbeSpawnAttemptsForTests(): number {
+  return probeSpawnAttemptsForTests;
+}
+
+/** Test-only close barrier for proving admission is not released at process exit. */
+export function setPromptTextProbeCloseBarrierForTests(barrier: Promise<void> | null): void {
+  probeCloseBarrierForTests = barrier;
+}
+
+/** Test-only fail-closed drain so one failed lifecycle case cannot poison another. */
+export async function resetPromptTextProbeForTests(): Promise<void> {
+  const active = activePromptProbe;
+  if (active && !active.settled) {
+    active.controller.abort(new DOMException("Prompt probe test reset", "AbortError"));
+    const drained = await Promise.race([
+      active.closed.then(() => true),
+      Bun.sleep(2_000).then(() => false),
+    ]);
+    if (!drained) throw new Error("prompt probe child did not close during test reset");
+  }
+  if (activePromptProbe === active) activePromptProbe = null;
+  probeCommandForTests = null;
+  probeSpawnAttemptsForTests = 0;
+  probeCloseBarrierForTests = null;
+}

src/server/management/codex-prompt-routes.ts

@@ -320,7 +320,11 @@ export async function handleCodexPromptRoutes(ctx: ManagementContext): Promise<R
     // A `cwd` parameter would have let any authenticated request read an arbitrary
     // folder's AGENTS.md through this endpoint.
     const { probePromptText } = await import("../../codex/prompt-text-probe");
-    return jsonResponse(await probePromptText(), 200, req, ctx.config);
+    // A write can complete while an older probe is still running. Include the
+    // exact prompt-file revision in single-flight admission so a post-write read
+    // fails soft instead of joining and returning the pre-write result.
+    const promptRevision = readPromptLayers(paths(ctx)).revision;
+    return jsonResponse(await probePromptText(15_000, req.signal, promptRevision), 200, req, ctx.config);
   }
 
   if (url.pathname === "/api/codex-prompt/toggle" && req.method === "PUT") {