diff --git a/plugins/codex/scripts/app-server-broker.mjs b/plugins/codex/scripts/app-server-broker.mjs index 1954274f..722ea462 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 72b30a76..29d66642 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 ef763819..c757b9ed 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 fead00cc..8225cf68 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 778571e6..2c894563 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 b3fc1146..d0f10813 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 8f276835..d0c1eb03 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);