From 8884609a91d4537cd93149027d1b099358e840f1 Mon Sep 17 00:00:00 2001 From: Bryan Lee <38807139+liby@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:23:15 +0800 Subject: [PATCH] fix: preserve shared broker across session exits SessionEnd treated the workspace broker as owned by the exiting Claude session, so it could terminate another session's active Codex task and leave the caller waiting forever after the transport disappeared. Let the broker reject shutdown while requests, streams, or other clients are active, then exit after five minutes of complete idleness so crashed sessions cannot leave it alive indefinitely. Bound shutdown RPCs to one second and only tear down resources after an accepted response. Race task completion against transport exit, and fall back to a direct client only for marked broker connection failures. --- plugins/codex/scripts/app-server-broker.mjs | 110 ++++- plugins/codex/scripts/lib/app-server.mjs | 39 +- .../codex/scripts/lib/broker-lifecycle.mjs | 119 ++++- plugins/codex/scripts/lib/codex.mjs | 26 +- .../codex/scripts/session-lifecycle-hook.mjs | 40 +- tests/broker-endpoint.test.mjs | 74 ++++ tests/runtime.test.mjs | 413 +++++++++++++++++- 7 files changed, 755 insertions(+), 66 deletions(-) diff --git a/plugins/codex/scripts/app-server-broker.mjs b/plugins/codex/scripts/app-server-broker.mjs index 1954274fe..722ea4628 100644 --- a/plugins/codex/scripts/app-server-broker.mjs +++ b/plugins/codex/scripts/app-server-broker.mjs @@ -10,6 +10,8 @@ import { BROKER_BUSY_RPC_CODE, CodexAppServerClient } from "./lib/app-server.mjs import { parseBrokerEndpoint } from "./lib/broker-endpoint.mjs"; const STREAMING_METHODS = new Set(["turn/start", "review/start", "thread/compact/start"]); +const BROKER_IDLE_TIMEOUT_ENV = "CODEX_COMPANION_BROKER_IDLE_TIMEOUT_MS"; +const DEFAULT_BROKER_IDLE_TIMEOUT_MS = 5 * 60 * 1000; function buildStreamThreadIds(method, params, result) { const threadIds = new Set(); @@ -45,6 +47,14 @@ function writePidFile(pidFile) { fs.writeFileSync(pidFile, `${process.pid}\n`, "utf8"); } +function resolveIdleTimeoutMs(value) { + const parsed = Number(value); + if (!Number.isFinite(parsed) || parsed <= 0) { + return DEFAULT_BROKER_IDLE_TIMEOUT_MS; + } + return Math.max(1, Math.floor(parsed)); +} + async function main() { const [subcommand, ...argv] = process.argv.slice(2); if (subcommand !== "serve") { @@ -70,6 +80,28 @@ async function main() { let activeStreamSocket = null; let activeStreamThreadIds = null; const sockets = new Set(); + const idleTimeoutMs = resolveIdleTimeoutMs(process.env[BROKER_IDLE_TIMEOUT_ENV]); + let idleTimer = null; + let shutdownPromise = null; + + function cancelIdleShutdown() { + if (idleTimer) { + clearTimeout(idleTimer); + idleTimer = null; + } + } + + function isFullyIdle() { + return sockets.size === 0 && !activeRequestSocket && !activeStreamSocket; + } + + function isBusyForShutdown(requestingSocket) { + return Boolean( + activeRequestSocket || + activeStreamSocket || + [...sockets].some((socket) => socket !== requestingSocket) + ); + } function clearSocketOwnership(socket) { if (activeRequestSocket === socket) { @@ -99,23 +131,57 @@ async function main() { } } - async function shutdown(server) { - 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); + function shutdown(server) { + if (!shutdownPromise) { + shutdownPromise = (async () => { + cancelIdleShutdown(); + const serverClosed = new Promise((resolve) => server.close(resolve)); + if (listenTarget.kind === "unix") { + fs.rmSync(listenTarget.path, { force: true }); + } + for (const socket of sockets) { + socket.end(); + } + await appClient.close().catch(() => {}); + await serverClosed; + if (pidFile) { + fs.rmSync(pidFile, { force: true }); + } + })(); } - if (pidFile && fs.existsSync(pidFile)) { - fs.unlinkSync(pidFile); + return shutdownPromise; + } + + function scheduleIdleShutdown(server) { + cancelIdleShutdown(); + if (!isFullyIdle() || shutdownPromise) { + return; } + + idleTimer = setTimeout(() => { + idleTimer = null; + if (!isFullyIdle() || shutdownPromise) { + return; + } + shutdown(server).then( + () => process.exit(0), + (error) => { + process.stderr.write(`${error instanceof Error ? error.message : String(error)}\n`); + process.exit(1); + } + ); + }, idleTimeoutMs); + idleTimer.unref?.(); } appClient.setNotificationHandler(routeNotification); const server = net.createServer((socket) => { + if (shutdownPromise) { + socket.destroy(); + return; + } + cancelIdleShutdown(); sockets.add(socket); socket.setEncoding("utf8"); let buffer = ""; @@ -158,6 +224,13 @@ async function main() { } if (message.id !== undefined && message.method === "broker/shutdown") { + if (isBusyForShutdown(socket)) { + send(socket, { + id: message.id, + error: buildJsonRpcError(BROKER_BUSY_RPC_CODE, "Shared Codex broker is busy.") + }); + continue; + } send(socket, { id: message.id, result: {} }); await shutdown(server); process.exit(0); @@ -222,15 +295,15 @@ async function main() { } }); - socket.on("close", () => { - sockets.delete(socket); + const releaseSocket = () => { + const removed = sockets.delete(socket); clearSocketOwnership(socket); - }); - - socket.on("error", () => { - sockets.delete(socket); - clearSocketOwnership(socket); - }); + if (removed) { + scheduleIdleShutdown(server); + } + }; + socket.on("close", releaseSocket); + socket.on("error", releaseSocket); }); process.on("SIGTERM", async () => { @@ -243,6 +316,7 @@ async function main() { process.exit(0); }); + server.on("listening", () => scheduleIdleShutdown(server)); server.listen(listenTarget.path); } diff --git a/plugins/codex/scripts/lib/app-server.mjs b/plugins/codex/scripts/lib/app-server.mjs index 72b30a764..29d66642d 100644 --- a/plugins/codex/scripts/lib/app-server.mjs +++ b/plugins/codex/scripts/lib/app-server.mjs @@ -1,5 +1,5 @@ /** - * @typedef {Error & { data?: unknown, rpcCode?: number }} ProtocolError + * @typedef {Error & { code?: string, codexTransport?: "broker", data?: unknown, rpcCode?: number }} ProtocolError * @typedef {import("./app-server-protocol").AppServerMethod} AppServerMethod * @typedef {import("./app-server-protocol").AppServerNotification} AppServerNotification * @typedef {import("./app-server-protocol").AppServerNotificationHandler} AppServerNotificationHandler @@ -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, isBrokerSessionReady, loadBrokerSession } from "./broker-lifecycle.mjs"; import { terminateProcessTree } from "./process.mjs"; const PLUGIN_MANIFEST_URL = new URL("../../.claude-plugin/plugin.json", import.meta.url); @@ -54,6 +54,22 @@ function createProtocolError(message, data) { return error; } +function createBrokerConnectionClosedError() { + const error = /** @type {ProtocolError} */ (new Error("codex app-server connection closed.")); + error.code = "ECONNRESET"; + return error; +} + +function markBrokerTransportError(error) { + const marked = /** @type {ProtocolError} */ (error instanceof Error ? error : new Error(String(error))); + marked.codexTransport = "broker"; + return marked; +} + +export function isBrokerTransportError(error) { + return error instanceof Error && /** @type {ProtocolError} */ (error).codexTransport === "broker"; +} + class AppServerClientBase { constructor(cwd, options = {}) { this.cwd = cwd; @@ -175,6 +191,11 @@ class AppServerClientBase { this.resolveExit(undefined); } + async waitForExit() { + await this.exitPromise; + throw this.exitError ?? new Error("codex app-server connection closed."); + } + sendMessage(_message) { throw new Error("sendMessage must be implemented by subclasses."); } @@ -298,7 +319,7 @@ class BrokerCodexAppServerClient extends AppServerClientBase { this.handleExit(error); }); this.socket.on("close", () => { - this.handleExit(this.exitError); + this.handleExit(this.exitError ?? createBrokerConnectionClosedError()); }); }); @@ -338,7 +359,8 @@ 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; + const brokerSession = loadBrokerSession(cwd); + brokerEndpoint = (await isBrokerSessionReady(brokerSession)) ? brokerSession.endpoint : null; } if (!brokerEndpoint && !options.reuseExistingBroker) { const brokerSession = await ensureBrokerSession(cwd, { env: options.env }); @@ -348,7 +370,14 @@ export class CodexAppServerClient { const client = brokerEndpoint ? new BrokerCodexAppServerClient(cwd, { ...options, brokerEndpoint }) : new SpawnedCodexAppServerClient(cwd, options); - await client.initialize(); + try { + await client.initialize(); + } catch (error) { + if (client.transport === "broker") { + throw markBrokerTransportError(error); + } + throw error; + } return client; } } diff --git a/plugins/codex/scripts/lib/broker-lifecycle.mjs b/plugins/codex/scripts/lib/broker-lifecycle.mjs index ef763819c..c757b9ed3 100644 --- a/plugins/codex/scripts/lib/broker-lifecycle.mjs +++ b/plugins/codex/scripts/lib/broker-lifecycle.mjs @@ -10,7 +10,9 @@ import { resolveStateDir } from "./state.mjs"; export const PID_FILE_ENV = "CODEX_COMPANION_APP_SERVER_PID_FILE"; export const LOG_FILE_ENV = "CODEX_COMPANION_APP_SERVER_LOG_FILE"; +export const BROKER_SHUTDOWN_ACCEPTED = "accepted"; const BROKER_STATE_FILE = "broker.json"; +const BROKER_SHUTDOWN_TIMEOUT_MS = 1000; export function createBrokerSessionDir(prefix = "cxc-") { return fs.mkdtempSync(path.join(os.tmpdir(), prefix)); @@ -40,19 +42,68 @@ export async function waitForBrokerEndpoint(endpoint, timeoutMs = 2000) { return false; } -export async function sendBrokerShutdown(endpoint) { - await new Promise((resolve) => { - const socket = connectToEndpoint(endpoint); +export async function sendBrokerShutdown(endpoint, timeoutMs = BROKER_SHUTDOWN_TIMEOUT_MS) { + return new Promise((resolve) => { + let socket; + let timer = null; + let settled = false; + let buffer = ""; + + const finish = (outcome) => { + if (settled) { + return; + } + settled = true; + if (timer) { + clearTimeout(timer); + } + if (socket && !socket.destroyed) { + socket.destroy(); + } + resolve(outcome); + }; + + try { + socket = connectToEndpoint(endpoint); + } catch { + finish("rejected"); + return; + } + + const boundedTimeoutMs = Number.isFinite(timeoutMs) ? Math.max(1, timeoutMs) : BROKER_SHUTDOWN_TIMEOUT_MS; + timer = setTimeout(() => finish("timeout"), boundedTimeoutMs); socket.setEncoding("utf8"); socket.on("connect", () => { socket.write(`${JSON.stringify({ id: 1, method: "broker/shutdown", params: {} })}\n`); }); - socket.on("data", () => { - socket.end(); - resolve(); + socket.on("data", (chunk) => { + buffer += chunk; + let newlineIndex = buffer.indexOf("\n"); + while (newlineIndex !== -1) { + const line = buffer.slice(0, newlineIndex); + buffer = buffer.slice(newlineIndex + 1); + newlineIndex = buffer.indexOf("\n"); + if (!line.trim()) { + continue; + } + + let response; + try { + response = JSON.parse(line); + } catch { + finish("rejected"); + return; + } + if (response.id !== 1) { + continue; + } + const hasResult = Object.prototype.hasOwnProperty.call(response, "result"); + finish(response.error || !hasResult ? "rejected" : BROKER_SHUTDOWN_ACCEPTED); + return; + } }); - socket.on("error", resolve); - socket.on("close", resolve); + socket.on("error", () => finish("rejected")); + socket.on("close", () => finish("rejected")); }); } @@ -93,26 +144,53 @@ export function saveBrokerSession(cwd, session) { } export function clearBrokerSession(cwd) { - const stateFile = resolveBrokerStateFile(cwd); - if (fs.existsSync(stateFile)) { - fs.unlinkSync(stateFile); + fs.rmSync(resolveBrokerStateFile(cwd), { force: true }); +} + +export function isBrokerSessionActive(session) { + if (!session?.endpoint) { + return false; + } + if (!Number.isFinite(session.pid)) { + return true; + } + + try { + process.kill(session.pid, 0); + } catch { + return false; + } + + if (session.pidFile && !fs.existsSync(session.pidFile)) { + return false; + } + + try { + const target = parseBrokerEndpoint(session.endpoint); + return target.kind !== "unix" || fs.existsSync(target.path); + } catch { + return false; } } -async function isBrokerEndpointReady(endpoint) { +async function isBrokerEndpointReady(endpoint, timeoutMs = 150) { if (!endpoint) { return false; } try { - return await waitForBrokerEndpoint(endpoint, 150); + return await waitForBrokerEndpoint(endpoint, timeoutMs); } catch { return false; } } +export async function isBrokerSessionReady(session, timeoutMs = 150) { + return isBrokerSessionActive(session) && (await isBrokerEndpointReady(session.endpoint, timeoutMs)); +} + export async function ensureBrokerSession(cwd, options = {}) { const existing = loadBrokerSession(cwd); - if (existing && (await isBrokerEndpointReady(existing.endpoint))) { + if (existing && (await isBrokerSessionReady(existing))) { return existing; } @@ -179,19 +257,18 @@ export function teardownBrokerSession({ endpoint = null, pidFile, logFile, sessi } } - if (pidFile && fs.existsSync(pidFile)) { - fs.unlinkSync(pidFile); + if (pidFile) { + fs.rmSync(pidFile, { force: true }); } - - if (logFile && fs.existsSync(logFile)) { - fs.unlinkSync(logFile); + if (logFile) { + fs.rmSync(logFile, { force: true }); } if (endpoint) { try { const target = parseBrokerEndpoint(endpoint); - if (target.kind === "unix" && fs.existsSync(target.path)) { - fs.unlinkSync(target.path); + if (target.kind === "unix") { + fs.rmSync(target.path, { force: true }); } } catch { // Ignore malformed or already-removed broker endpoints during teardown. diff --git a/plugins/codex/scripts/lib/codex.mjs b/plugins/codex/scripts/lib/codex.mjs index fead00cc4..8225cf68b 100644 --- a/plugins/codex/scripts/lib/codex.mjs +++ b/plugins/codex/scripts/lib/codex.mjs @@ -40,8 +40,13 @@ import os from "node:os"; import path from "node:path"; import { readJsonFile } from "./fs.mjs"; -import { BROKER_BUSY_RPC_CODE, BROKER_ENDPOINT_ENV, CodexAppServerClient } from "./app-server.mjs"; -import { loadBrokerSession } from "./broker-lifecycle.mjs"; +import { + BROKER_BUSY_RPC_CODE, + BROKER_ENDPOINT_ENV, + CodexAppServerClient, + isBrokerTransportError +} from "./app-server.mjs"; +import { isBrokerSessionActive, loadBrokerSession } from "./broker-lifecycle.mjs"; import { binaryAvailable } from "./process.mjs"; const SERVICE_NAME = "claude_code_codex_plugin"; @@ -603,7 +608,7 @@ async function captureTurn(client, threadId, startRequest, options = {}) { completeTurn(state, response.turn); } - return await state.completion; + return await Promise.race([state.completion, client.waitForExit()]); } finally { clearCompletionTimer(state); client.setNotificationHandler(previousHandler ?? null); @@ -618,10 +623,16 @@ async function withAppServer(cwd, fn) { await client.close(); return result; } catch (error) { - const brokerRequested = client?.transport === "broker" || Boolean(process.env[BROKER_ENDPOINT_ENV]); + const brokerTransportError = isBrokerTransportError(error); + const brokerConnectionFailed = + brokerTransportError && + (error?.code === "ENOENT" || + error?.code === "ECONNREFUSED" || + error?.code === "ECONNRESET" || + error?.code === "EPIPE"); const shouldRetryDirect = (client?.transport === "broker" && error?.rpcCode === BROKER_BUSY_RPC_CODE) || - (brokerRequested && (error?.code === "ENOENT" || error?.code === "ECONNREFUSED")); + brokerConnectionFailed; if (client) { await client.close().catch(() => {}); @@ -904,7 +915,10 @@ export function getCodexAvailability(cwd) { } export function getSessionRuntimeStatus(env = process.env, cwd = process.cwd()) { - const endpoint = env?.[BROKER_ENDPOINT_ENV] ?? loadBrokerSession(cwd)?.endpoint ?? null; + const brokerSession = loadBrokerSession(cwd); + const endpoint = + env?.[BROKER_ENDPOINT_ENV] ?? + (isBrokerSessionActive(brokerSession) ? brokerSession.endpoint : null); if (endpoint) { return { mode: "shared", diff --git a/plugins/codex/scripts/session-lifecycle-hook.mjs b/plugins/codex/scripts/session-lifecycle-hook.mjs index 778571e6c..2c894563e 100644 --- a/plugins/codex/scripts/session-lifecycle-hook.mjs +++ b/plugins/codex/scripts/session-lifecycle-hook.mjs @@ -6,7 +6,7 @@ import process from "node:process"; import { terminateProcessTree } from "./lib/process.mjs"; import { BROKER_ENDPOINT_ENV } from "./lib/app-server.mjs"; import { - clearBrokerSession, + BROKER_SHUTDOWN_ACCEPTED, LOG_FILE_ENV, loadBrokerSession, PID_FILE_ENV, @@ -74,6 +74,20 @@ function cleanupSessionJobs(cwd, sessionId) { }); } +function hasActiveJobFromAnotherSession(cwd, sessionId) { + const workspaceRoot = resolveWorkspaceRoot(cwd); + const stateFile = resolveStateFile(workspaceRoot); + if (!fs.existsSync(stateFile)) { + return false; + } + + return loadState(workspaceRoot).jobs.some( + (job) => + (job.status === "queued" || job.status === "running") && + (!sessionId || !job.sessionId || job.sessionId !== sessionId) + ); +} + function handleSessionStart(input) { appendEnvVar(SESSION_ID_ENV, input.session_id); appendEnvVar(TRANSCRIPT_PATH_ENV, input.transcript_path); @@ -82,6 +96,7 @@ function handleSessionStart(input) { async function handleSessionEnd(input) { const cwd = input.cwd || process.cwd(); + const sessionId = input.session_id || process.env[SESSION_ID_ENV] || null; const brokerSession = loadBrokerSession(cwd) ?? (process.env[BROKER_ENDPOINT_ENV] @@ -92,25 +107,24 @@ async function handleSessionEnd(input) { } : null); const brokerEndpoint = brokerSession?.endpoint ?? null; - const pidFile = brokerSession?.pidFile ?? null; - const logFile = brokerSession?.logFile ?? null; - const sessionDir = brokerSession?.sessionDir ?? null; - const pid = brokerSession?.pid ?? null; - if (brokerEndpoint) { - await sendBrokerShutdown(brokerEndpoint); + cleanupSessionJobs(cwd, sessionId); + if (!brokerEndpoint || hasActiveJobFromAnotherSession(cwd, sessionId)) { + return; + } + + const shutdownOutcome = await sendBrokerShutdown(brokerEndpoint); + if (shutdownOutcome !== BROKER_SHUTDOWN_ACCEPTED) { + return; } - cleanupSessionJobs(cwd, input.session_id || process.env[SESSION_ID_ENV]); + // A new broker may replace broker.json after the old listener closes. + // Remove only resources captured from the broker that accepted shutdown. teardownBrokerSession({ + ...brokerSession, endpoint: brokerEndpoint, - pidFile, - logFile, - sessionDir, - pid, killProcess: terminateProcessTree }); - clearBrokerSession(cwd); } async function main() { diff --git a/tests/broker-endpoint.test.mjs b/tests/broker-endpoint.test.mjs index b3fc11467..d0f108132 100644 --- a/tests/broker-endpoint.test.mjs +++ b/tests/broker-endpoint.test.mjs @@ -1,7 +1,47 @@ +import fs from "node:fs"; +import net from "node:net"; import test from "node:test"; import assert from "node:assert/strict"; import { createBrokerEndpoint, parseBrokerEndpoint } from "../plugins/codex/scripts/lib/broker-endpoint.mjs"; +import { sendBrokerShutdown } from "../plugins/codex/scripts/lib/broker-lifecycle.mjs"; +import { makeTempDir } from "./helpers.mjs"; + +async function listenForBrokerResponse(t, response = undefined) { + const sessionDir = makeTempDir("codex-broker-response-"); + const endpoint = createBrokerEndpoint(sessionDir); + const target = parseBrokerEndpoint(endpoint); + const sockets = new Set(); + const server = net.createServer((socket) => { + sockets.add(socket); + socket.on("data", () => { + if (response !== undefined) { + socket.write(`${JSON.stringify(response)}\n`); + } + }); + socket.on("close", () => sockets.delete(socket)); + }); + + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(target.path, () => { + server.off("error", reject); + resolve(); + }); + }); + + t.after(async () => { + for (const socket of sockets) { + socket.destroy(); + } + await new Promise((resolve) => server.close(resolve)); + if (target.kind === "unix") { + fs.rmSync(target.path, { force: true }); + } + }); + + return endpoint; +} test("createBrokerEndpoint uses Unix sockets on non-Windows platforms", () => { const endpoint = createBrokerEndpoint("/tmp/cxc-12345", "darwin"); @@ -20,3 +60,37 @@ test("createBrokerEndpoint uses named pipes on Windows", () => { path: "\\\\.\\pipe\\cxc-12345-codex-app-server" }); }); + +test("sendBrokerShutdown accepts an explicit successful response", async (t) => { + const endpoint = await listenForBrokerResponse(t, { id: 1, result: {} }); + + assert.equal(await sendBrokerShutdown(endpoint, 100), "accepted"); +}); + +test("sendBrokerShutdown reports an explicit broker rejection", async (t) => { + const endpoint = await listenForBrokerResponse(t, { + id: 1, + error: { + code: -32001, + message: "Shared Codex broker is busy." + } + }); + + assert.equal(await sendBrokerShutdown(endpoint, 100), "rejected"); +}); + +test("sendBrokerShutdown rejects a response without a result", async (t) => { + const endpoint = await listenForBrokerResponse(t, { id: 1 }); + + assert.equal(await sendBrokerShutdown(endpoint, 100), "rejected"); +}); + +test("sendBrokerShutdown stops waiting when the broker does not respond", async (t) => { + const endpoint = await listenForBrokerResponse(t); + + const startedAt = Date.now(); + const outcome = await sendBrokerShutdown(endpoint, 50); + + assert.equal(outcome, "timeout"); + assert.ok(Date.now() - startedAt < 1000); +}); diff --git a/tests/runtime.test.mjs b/tests/runtime.test.mjs index 8f276835b..d0c1eb037 100644 --- a/tests/runtime.test.mjs +++ b/tests/runtime.test.mjs @@ -1,4 +1,5 @@ import fs from "node:fs"; +import net from "node:net"; import path from "node:path"; import test from "node:test"; import assert from "node:assert/strict"; @@ -7,7 +8,15 @@ import { fileURLToPath } from "node:url"; import { buildEnv, installFakeCodex } from "./fake-codex-fixture.mjs"; import { initGitRepo, makeTempDir, run } from "./helpers.mjs"; -import { loadBrokerSession, saveBrokerSession } from "../plugins/codex/scripts/lib/broker-lifecycle.mjs"; +import { BROKER_ENDPOINT_ENV, CodexAppServerClient } from "../plugins/codex/scripts/lib/app-server.mjs"; +import { + isBrokerSessionActive, + loadBrokerSession, + saveBrokerSession, + sendBrokerShutdown +} from "../plugins/codex/scripts/lib/broker-lifecycle.mjs"; +import { createBrokerEndpoint, parseBrokerEndpoint } from "../plugins/codex/scripts/lib/broker-endpoint.mjs"; +import { terminateProcessTree } from "../plugins/codex/scripts/lib/process.mjs"; import { resolveStateDir } from "../plugins/codex/scripts/lib/state.mjs"; const ROOT = path.resolve(path.dirname(fileURLToPath(import.meta.url)), ".."); @@ -1801,6 +1810,396 @@ test("cancel sends turn interrupt to the shared app-server before killing a brok assert.equal(cleanup.status, 0, cleanup.stderr); }); +test("ending another session preserves an active brokered task", async (t) => { + const repo = makeTempDir(); + const binDir = makeTempDir(); + installFakeCodex(binDir, "interruptible-slow-task"); + initGitRepo(repo); + fs.writeFileSync(path.join(repo, "README.md"), "hello\n"); + run("git", ["add", "README.md"], { cwd: repo }); + run("git", ["commit", "-m", "init"], { cwd: repo }); + + const ownerEnv = { + ...buildEnv(binDir), + CODEX_COMPANION_SESSION_ID: "sess-owner" + }; + const launched = run("node", [SCRIPT, "task", "--background", "--json", "investigate the flaky worker timeout"], { + cwd: repo, + env: ownerEnv + }); + + assert.equal(launched.status, 0, launched.stderr); + const { jobId } = JSON.parse(launched.stdout); + const stateDir = resolveStateDir(repo); + const runningJob = await waitFor(() => { + const state = JSON.parse(fs.readFileSync(path.join(stateDir, "state.json"), "utf8")); + const job = state.jobs.find((candidate) => candidate.id === jobId); + return job?.status === "running" && job.threadId && job.turnId ? job : null; + }, { timeoutMs: 15000 }); + const brokerBefore = loadBrokerSession(repo); + assert.ok(brokerBefore?.pid); + + t.after(() => { + for (const pid of [runningJob.pid, brokerBefore.pid]) { + try { + terminateProcessTree(pid); + } catch {} + } + }); + + const otherSessionEnd = run("node", [SESSION_HOOK, "SessionEnd"], { + cwd: repo, + env: { + ...buildEnv(binDir), + CODEX_COMPANION_SESSION_ID: "sess-other" + }, + input: JSON.stringify({ + hook_event_name: "SessionEnd", + session_id: "sess-other", + cwd: repo + }) + }); + + assert.equal(otherSessionEnd.status, 0, otherSessionEnd.stderr); + const brokerAfter = loadBrokerSession(repo); + assert.equal(brokerAfter?.pid, brokerBefore.pid); + + const stateAfter = JSON.parse(fs.readFileSync(path.join(stateDir, "state.json"), "utf8")); + assert.equal(stateAfter.jobs.find((job) => job.id === jobId)?.status, "running"); +}); + +test("session end preserves a broker that rejects shutdown while another client is connected", async (t) => { + const repo = makeTempDir(); + const binDir = makeTempDir(); + installFakeCodex(binDir); + initGitRepo(repo); + fs.writeFileSync(path.join(repo, "README.md"), "hello\n"); + run("git", ["add", "README.md"], { cwd: repo }); + run("git", ["commit", "-m", "init"], { cwd: repo }); + + const env = { + ...buildEnv(binDir), + CODEX_COMPANION_SESSION_ID: "" + }; + const taskResult = run("node", [SCRIPT, "task", "exercise the shared broker"], { + cwd: repo, + env + }); + assert.equal(taskResult.status, 0, taskResult.stderr); + + const brokerSession = loadBrokerSession(repo); + assert.ok(brokerSession?.pid); + const client = net.createConnection({ path: parseBrokerEndpoint(brokerSession.endpoint).path }); + await new Promise((resolve, reject) => { + client.once("connect", resolve); + client.once("error", reject); + }); + + t.after(() => { + client.destroy(); + try { + terminateProcessTree(brokerSession.pid); + } catch {} + }); + + const busySessionEnd = run("node", [SESSION_HOOK, "SessionEnd"], { + cwd: repo, + env, + input: JSON.stringify({ + hook_event_name: "SessionEnd", + session_id: "sess-exiting", + cwd: repo + }) + }); + + assert.equal(busySessionEnd.status, 0, busySessionEnd.stderr); + assert.equal(loadBrokerSession(repo)?.pid, brokerSession.pid); + process.kill(brokerSession.pid, 0); + + const clientClosed = new Promise((resolve) => client.once("close", resolve)); + client.end(); + await clientClosed; + + const idleSessionEnd = run("node", [SESSION_HOOK, "SessionEnd"], { + cwd: repo, + env, + input: JSON.stringify({ + hook_event_name: "SessionEnd", + session_id: "sess-exiting", + cwd: repo + }) + }); + + assert.equal(idleSessionEnd.status, 0, idleSessionEnd.stderr); + assert.equal(isBrokerSessionActive(brokerSession), false); +}); + +test("accepted shutdown stops accepting new broker clients immediately", async (t) => { + const repo = makeTempDir(); + const binDir = makeTempDir(); + installFakeCodex(binDir, "interruptible-slow-task"); + initGitRepo(repo); + fs.writeFileSync(path.join(repo, "README.md"), "hello\n"); + run("git", ["add", "README.md"], { cwd: repo }); + run("git", ["commit", "-m", "init"], { cwd: repo }); + + const env = buildEnv(binDir); + const launched = run("node", [SCRIPT, "task", "--background", "--json", "leave an orphaned slow turn"], { + cwd: repo, + env + }); + assert.equal(launched.status, 0, launched.stderr); + + const { jobId } = JSON.parse(launched.stdout); + const stateDir = resolveStateDir(repo); + const runningJob = await waitFor(() => { + const state = JSON.parse(fs.readFileSync(path.join(stateDir, "state.json"), "utf8")); + const job = state.jobs.find((candidate) => candidate.id === jobId); + return job?.status === "running" && job.threadId && job.turnId ? job : null; + }, { timeoutMs: 15000 }); + const brokerSession = loadBrokerSession(repo); + assert.ok(brokerSession?.pid); + + let lateClient = null; + t.after(() => { + lateClient?.destroy(); + for (const pid of [runningJob.pid, brokerSession.pid]) { + try { + terminateProcessTree(pid); + } catch {} + } + }); + + terminateProcessTree(runningJob.pid); + await waitFor(() => { + try { + process.kill(runningJob.pid, 0); + return false; + } catch (error) { + return error?.code === "ESRCH"; + } + }); + + await waitFor(async () => { + const outcome = await sendBrokerShutdown(brokerSession.endpoint, 200); + return outcome === "accepted"; + }); + + lateClient = net.createConnection({ path: parseBrokerEndpoint(brokerSession.endpoint).path }); + const protocolAccepted = await new Promise((resolve) => { + let settled = false; + let buffer = ""; + const finish = (accepted) => { + if (settled) { + return; + } + settled = true; + clearTimeout(timer); + resolve(accepted); + }; + const timer = setTimeout(() => { + lateClient.destroy(); + finish(false); + }, 1000); + lateClient.once("connect", () => { + lateClient.write(`${JSON.stringify({ id: 99, method: "initialize", params: {} })}\n`); + }); + lateClient.on("data", (chunk) => { + buffer += chunk; + if (buffer.includes("\n")) { + const response = JSON.parse(buffer.slice(0, buffer.indexOf("\n"))); + finish(response.id === 99 && Boolean(response.result)); + } + }); + lateClient.once("error", () => { + finish(false); + }); + lateClient.once("close", () => { + finish(false); + }); + }); + + assert.equal(protocolAccepted, false); +}); + +test("task falls back to direct when the broker exits between readiness and initialize", async (t) => { + const repo = makeTempDir(); + const binDir = makeTempDir(); + const brokerDir = makeTempDir("codex-one-shot-broker-"); + const endpoint = createBrokerEndpoint(brokerDir); + const endpointPath = parseBrokerEndpoint(endpoint).path; + const readyFile = path.join(brokerDir, "ready"); + const pidFile = path.join(brokerDir, "broker.pid"); + const logFile = path.join(brokerDir, "broker.log"); + installFakeCodex(binDir); + initGitRepo(repo); + fs.writeFileSync(path.join(repo, "README.md"), "hello\n"); + run("git", ["add", "README.md"], { cwd: repo }); + run("git", ["commit", "-m", "init"], { cwd: repo }); + + const oneShotBroker = spawn( + process.execPath, + [ + "-e", + ` + const fs = require("node:fs"); + const net = require("node:net"); + const [endpointPath, readyFile] = process.argv.slice(1); + const server = net.createServer((socket) => { + server.close(() => process.exit(0)); + if (process.platform !== "win32") { + fs.rmSync(endpointPath, { force: true }); + } + socket.end(); + }); + server.listen(endpointPath, () => fs.writeFileSync(readyFile, "ready\\n", "utf8")); + `, + endpointPath, + readyFile + ], + { + cwd: repo, + stdio: "ignore" + } + ); + + t.after(() => { + try { + terminateProcessTree(oneShotBroker.pid); + } catch {} + }); + + await waitFor(() => fs.existsSync(readyFile)); + fs.writeFileSync(pidFile, `${oneShotBroker.pid}\n`, "utf8"); + fs.writeFileSync(logFile, "", "utf8"); + saveBrokerSession(repo, { + endpoint, + pidFile, + logFile, + sessionDir: brokerDir, + pid: oneShotBroker.pid + }); + + const result = run("node", [SCRIPT, "task", "exercise broker connection fallback"], { + cwd: repo, + env: buildEnv(binDir) + }); + + assert.equal(result.status, 0, result.stderr); + assert.match(result.stdout, /Handled the requested task/); +}); + +test("broker exits after remaining fully idle for the configured timeout", async (t) => { + const repo = makeTempDir(); + const binDir = makeTempDir(); + installFakeCodex(binDir); + initGitRepo(repo); + fs.writeFileSync(path.join(repo, "README.md"), "hello\n"); + run("git", ["add", "README.md"], { cwd: repo }); + run("git", ["commit", "-m", "init"], { cwd: repo }); + + const env = { + ...buildEnv(binDir), + CODEX_COMPANION_BROKER_IDLE_TIMEOUT_MS: "100" + }; + delete env[BROKER_ENDPOINT_ENV]; + const taskResult = run("node", [SCRIPT, "task", "exercise broker idle cleanup"], { + cwd: repo, + env + }); + assert.equal(taskResult.status, 0, taskResult.stderr); + + const brokerSession = loadBrokerSession(repo); + assert.ok(brokerSession?.pid); + t.after(() => { + try { + terminateProcessTree(brokerSession.pid); + } catch {} + }); + + await waitFor(() => { + try { + process.kill(brokerSession.pid, 0); + return false; + } catch (error) { + return error?.code === "ESRCH"; + } + }); + + assert.equal(fs.existsSync(brokerSession.pidFile), false); + const target = parseBrokerEndpoint(brokerSession.endpoint); + if (target.kind === "unix") { + assert.equal(fs.existsSync(target.path), false); + } + + const statusResult = run("node", [SCRIPT, "status", "--json"], { + cwd: repo, + env + }); + assert.equal(statusResult.status, 0, statusResult.stderr); + + const reusedClient = await CodexAppServerClient.connect(repo, { + env, + reuseExistingBroker: true + }); + t.after(async () => { + await reusedClient.close().catch(() => {}); + }); + + assert.equal(loadBrokerSession(repo)?.pid, brokerSession.pid); + assert.equal(JSON.parse(statusResult.stdout).sessionRuntime.mode, "direct"); + assert.equal(reusedClient.transport, "direct"); +}); + +test("broker disconnect marks a background task failed", async (t) => { + const repo = makeTempDir(); + const binDir = makeTempDir(); + installFakeCodex(binDir, "interruptible-slow-task"); + initGitRepo(repo); + fs.writeFileSync(path.join(repo, "README.md"), "hello\n"); + run("git", ["add", "README.md"], { cwd: repo }); + run("git", ["commit", "-m", "init"], { cwd: repo }); + + const env = { + ...buildEnv(binDir), + CODEX_COMPANION_SESSION_ID: "sess-owner" + }; + const launched = run("node", [SCRIPT, "task", "--background", "--json", "investigate the flaky worker timeout"], { + cwd: repo, + env + }); + + assert.equal(launched.status, 0, launched.stderr); + const { jobId } = JSON.parse(launched.stdout); + const stateDir = resolveStateDir(repo); + const runningJob = await waitFor(() => { + const state = JSON.parse(fs.readFileSync(path.join(stateDir, "state.json"), "utf8")); + const job = state.jobs.find((candidate) => candidate.id === jobId); + return job?.status === "running" && job.threadId && job.turnId ? job : null; + }, { timeoutMs: 15000 }); + const brokerSession = loadBrokerSession(repo); + assert.ok(brokerSession?.pid); + + t.after(() => { + for (const pid of [runningJob.pid, brokerSession.pid]) { + try { + terminateProcessTree(pid); + } catch {} + } + }); + + terminateProcessTree(brokerSession.pid); + + const failedJob = await waitFor(() => { + const state = JSON.parse(fs.readFileSync(path.join(stateDir, "state.json"), "utf8")); + const job = state.jobs.find((candidate) => candidate.id === jobId); + return job?.status === "failed" ? job : null; + }, { timeoutMs: 5000 }); + + assert.match(failedJob.errorMessage, /connection closed/i); + assert.equal(failedJob.pid, null); +}); + test("session end fully cleans up jobs for the ending session", async (t) => { const repo = makeTempDir(); initGitRepo(repo); @@ -2238,19 +2637,27 @@ test("status reports shared session runtime when a lazy broker is active", () => test("setup and status honor --cwd when reading shared session runtime", () => { const targetWorkspace = makeTempDir(); const invocationWorkspace = makeTempDir(); + const binDir = makeTempDir(); + fs.symlinkSync(process.execPath, path.join(binDir, "node")); + const env = { + ...process.env, + PATH: binDir + }; saveBrokerSession(targetWorkspace, { endpoint: "unix:/tmp/fake-broker.sock" }); const status = run("node", [SCRIPT, "status", "--cwd", targetWorkspace], { - cwd: invocationWorkspace + cwd: invocationWorkspace, + env }); assert.equal(status.status, 0, status.stderr); assert.match(status.stdout, /Session runtime: shared session/); const setup = run("node", [SCRIPT, "setup", "--cwd", targetWorkspace, "--json"], { - cwd: invocationWorkspace + cwd: invocationWorkspace, + env }); assert.equal(setup.status, 0, setup.stderr); const payload = JSON.parse(setup.stdout);