diff --git a/plugins/codex/scripts/lib/app-server.mjs b/plugins/codex/scripts/lib/app-server.mjs index 72b30a76..2e2b43be 100644 --- a/plugins/codex/scripts/lib/app-server.mjs +++ b/plugins/codex/scripts/lib/app-server.mjs @@ -13,7 +13,12 @@ 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 { + assertBrokerSessionActive, + ensureBrokerSession, + resolveSessionId, + reuseBrokerSession +} from "./broker-lifecycle.mjs"; import { terminateProcessTree } from "./process.mjs"; const PLUGIN_MANIFEST_URL = new URL("../../.claude-plugin/plugin.json", import.meta.url); @@ -334,14 +339,31 @@ class BrokerCodexAppServerClient extends AppServerClientBase { export class CodexAppServerClient { static async connect(cwd, options = {}) { + const sessionId = resolveSessionId({ env: options.env }); let brokerEndpoint = null; - if (!options.disableBroker) { + if (options.disableBroker) { + assertBrokerSessionActive(sessionId); + } else { brokerEndpoint = options.brokerEndpoint ?? options.env?.[BROKER_ENDPOINT_ENV] ?? process.env[BROKER_ENDPOINT_ENV] ?? null; - if (!brokerEndpoint && options.reuseExistingBroker) { - brokerEndpoint = loadBrokerSession(cwd)?.endpoint ?? null; + if (brokerEndpoint) { + // A supplied endpoint bypasses the lifecycle functions that normally + // perform admission under the broker-state lock. + assertBrokerSessionActive(sessionId); } - if (!brokerEndpoint && !options.reuseExistingBroker) { - const brokerSession = await ensureBrokerSession(cwd, { env: options.env }); + if (!brokerEndpoint && sessionId && options.reuseExistingBroker) { + const brokerSession = await reuseBrokerSession(cwd, { + ...options.brokerOptions, + env: options.env, + killProcess: terminateProcessTree + }); + brokerEndpoint = brokerSession?.endpoint ?? null; + } + if (!brokerEndpoint && sessionId && !options.reuseExistingBroker) { + const brokerSession = await ensureBrokerSession(cwd, { + ...options.brokerOptions, + env: options.env, + killProcess: terminateProcessTree + }); brokerEndpoint = brokerSession?.endpoint ?? null; } } @@ -349,6 +371,12 @@ export class CodexAppServerClient { ? new BrokerCodexAppServerClient(cwd, { ...options, brokerEndpoint }) : new SpawnedCodexAppServerClient(cwd, options); await client.initialize(); + try { + assertBrokerSessionActive(sessionId); + } catch (error) { + await client.close(); + throw error; + } return client; } } diff --git a/plugins/codex/scripts/lib/broker-lifecycle.mjs b/plugins/codex/scripts/lib/broker-lifecycle.mjs index ef763819..3388f4ce 100644 --- a/plugins/codex/scripts/lib/broker-lifecycle.mjs +++ b/plugins/codex/scripts/lib/broker-lifecycle.mjs @@ -1,4 +1,5 @@ import fs from "node:fs"; +import { createHash } from "node:crypto"; import net from "node:net"; import os from "node:os"; import path from "node:path"; @@ -6,11 +7,318 @@ import process from "node:process"; import { spawn } from "node:child_process"; import { fileURLToPath } from "node:url"; import { createBrokerEndpoint, parseBrokerEndpoint } from "./broker-endpoint.mjs"; -import { resolveStateDir } from "./state.mjs"; +import { resolveStateDir, resolveStateRoot } 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"; const BROKER_STATE_FILE = "broker.json"; +const SESSION_ID_ENV = "CODEX_COMPANION_SESSION_ID"; +const BROKER_STATE_LOCK_TIMEOUT_MS = 5000; +const BROKER_SHUTDOWN_TIMEOUT_MS = 1000; +const BROKER_LOCK_TIMEOUT_CODE = "EBROKERSTATELOCKTIMEOUT"; +export const BROKER_OWNER_ENDED_CODE = "EBROKEROWNERENDED"; +export const BROKER_CLEANUP_INCOMPLETE_CODE = "EBROKERCLEANUPINCOMPLETE"; +const MAX_BROKER_STATE_BYTES = 64 * 1024; +const ENDED_SESSIONS_DIR = ".ended-sessions"; +const BROKER_LOCK_SESSION_FILE = "session"; + +let brokerLockTokenSeq = 0; + +export function resolveSessionId(options = {}) { + if (options.sessionId) { + return options.sessionId; + } + const env = options.env ?? process.env; + return env[SESSION_ID_ENV] ?? null; +} + +function brokerSessionOwners(session) { + const owners = []; + if (Array.isArray(session?.sessionIds)) { + owners.push(...session.sessionIds); + } + if (session?.sessionId) { + owners.push(session.sessionId); + } + return [...new Set(owners.filter(Boolean))]; +} + +function withBrokerSessionOwner(session, sessionId) { + if (!sessionId) { + return session; + } + const owners = brokerSessionOwners(session); + if (!owners.includes(sessionId)) { + owners.push(sessionId); + } + return { + ...session, + sessionId: owners[0] ?? sessionId, + sessionIds: owners + }; +} + +async function sleep(ms) { + await new Promise((resolve) => setTimeout(resolve, ms)); +} + +function isBrokerLockTimeout(error) { + return error?.code === BROKER_LOCK_TIMEOUT_CODE; +} + +function brokerOwnerEndedError(sessionId) { + return Object.assign( + new Error(`Broker owner session ${sessionId ?? "unknown"} ended while the broker was starting.`), + { code: BROKER_OWNER_ENDED_CODE } + ); +} + +function brokerCleanupIncompleteError(reason, count = 0, cause = null) { + return Object.assign(new Error(`Broker cleanup stopped before completion: ${reason}.`), { + code: BROKER_CLEANUP_INCOMPLETE_CODE, + count, + reason, + ...(cause ? { cause } : {}) + }); +} + +function isLockOwnerAlive(token) { + const pid = Number.parseInt(token?.split("-", 1)[0] ?? "", 10); + if (!Number.isSafeInteger(pid) || pid <= 0) { + return true; + } + try { + process.kill(pid, 0); + return true; + } catch (error) { + return error?.code !== "ESRCH"; + } +} + +function reclaimBarrierPrefix(lockDir) { + return `${path.basename(lockDir)}.reclaim-`; +} + +function recoverBrokerReclaimBarriers(lockDir) { + const parent = path.dirname(lockDir); + const prefix = reclaimBarrierPrefix(lockDir); + let blocked = false; + for (const entry of fs.readdirSync(parent, { withFileTypes: true })) { + if (!entry.isDirectory() || !entry.name.startsWith(prefix)) { + continue; + } + const barrier = path.join(parent, entry.name); + const reclaimerPid = Number.parseInt(entry.name.slice(prefix.length).split("-", 1)[0], 10); + if (isLockOwnerAlive(`${reclaimerPid}-reclaimer`)) { + blocked = true; + continue; + } + const movedLock = path.join(barrier, "lock"); + if (fs.existsSync(movedLock)) { + let ownerToken = null; + try { + ownerToken = fs.readFileSync(path.join(movedLock, "owner"), "utf8"); + } catch { + // An unreadable moved lock is preserved conservatively. + } + if (!ownerToken || isLockOwnerAlive(ownerToken)) { + if (!fs.existsSync(lockDir)) { + try { + fs.renameSync(movedLock, lockDir); + fs.rmSync(barrier, { recursive: true, force: true }); + continue; + } catch { + // Another process restored or published the canonical lock. + } + } + blocked = true; + continue; + } + } + fs.rmSync(barrier, { recursive: true, force: true }); + } + return blocked; +} + +function releaseOwnedBrokerLock(lockDir, token) { + const tokenFile = path.join(lockDir, "owner"); + try { + if (fs.readFileSync(tokenFile, "utf8") === token) { + fs.rmSync(lockDir, { recursive: true, force: true }); + return; + } + } catch { + // A reclaimer may have moved the lock behind a visible barrier. + } + const parent = path.dirname(lockDir); + const prefix = reclaimBarrierPrefix(lockDir); + for (const entry of fs.readdirSync(parent, { withFileTypes: true })) { + if (!entry.isDirectory() || !entry.name.startsWith(prefix)) { + continue; + } + const barrier = path.join(parent, entry.name); + const movedLock = path.join(barrier, "lock"); + try { + if (fs.readFileSync(path.join(movedLock, "owner"), "utf8") === token) { + fs.rmSync(movedLock, { recursive: true, force: true }); + fs.rmSync(barrier, { recursive: true, force: true }); + break; + } + } catch { + // This barrier belongs to another lock generation. + } + } + // Orphan recovery can restore this generation while release is inspecting + // the moved path. Recheck the canonical token before returning. + try { + if (fs.readFileSync(tokenFile, "utf8") === token) { + fs.rmSync(lockDir, { recursive: true, force: true }); + } + } catch { + // This owner no longer has a published lock generation. + } +} + +async function withBrokerStateFileLock(stateFile, fn, options = {}) { + const lockDir = `${stateFile}.lock`; + const tokenFile = path.join(lockDir, "owner"); + const token = `${process.pid}-${brokerLockTokenSeq += 1}`; + const timeoutMs = options.timeoutMs ?? BROKER_STATE_LOCK_TIMEOUT_MS; + const deadline = Date.now() + timeoutMs; + fs.mkdirSync(path.dirname(stateFile), { recursive: true }); + + while (true) { + if (recoverBrokerReclaimBarriers(lockDir)) { + if (Date.now() >= deadline) { + throw Object.assign(new Error(`Timed out waiting for broker state lock: ${stateFile}`), { + code: BROKER_LOCK_TIMEOUT_CODE + }); + } + await sleep(25); + continue; + } + let candidate = null; + try { + if (fs.existsSync(lockDir)) { + throw Object.assign(new Error(`Broker state lock exists: ${stateFile}`), { code: "EEXIST" }); + } + candidate = fs.mkdtempSync(`${lockDir}.candidate-${process.pid}-`); + fs.writeFileSync(path.join(candidate, "owner"), token, { encoding: "utf8", flag: "wx" }); + const ownerSessionId = options.ownerSessionId; + if (ownerSessionId && Buffer.byteLength(ownerSessionId, "utf8") <= 1024) { + fs.writeFileSync(path.join(candidate, BROKER_LOCK_SESSION_FILE), ownerSessionId, { + encoding: "utf8", + flag: "wx" + }); + } + fs.renameSync(candidate, lockDir); + candidate = null; + if (recoverBrokerReclaimBarriers(lockDir)) { + if (fs.readFileSync(tokenFile, "utf8") === token) { + fs.rmSync(lockDir, { recursive: true, force: true }); + } + await sleep(25); + continue; + } + break; + } catch (error) { + if (candidate) { + fs.rmSync(candidate, { recursive: true, force: true }); + } + if (!fs.existsSync(lockDir)) { + if (["EEXIST", "ENOTEMPTY", "EPERM"].includes(error?.code)) { + if (Date.now() >= deadline) { + throw Object.assign(new Error(`Timed out waiting for broker state lock: ${stateFile}`), { + code: BROKER_LOCK_TIMEOUT_CODE + }); + } + await sleep(25); + continue; + } + throw error; + } + let stat = null; + try { + stat = fs.lstatSync(lockDir); + } catch { + // The lock may have disappeared between mkdirSync and lstatSync. Retry + // through the normal deadline/backoff path instead of spinning. + } + if (stat && !stat.isDirectory()) { + let removed = false; + try { + // unlinkSync cannot remove a directory. If another contender replaced + // the invalid path with a real lock after lstatSync, this fails safely + // instead of renaming or deleting that contender's lock. + fs.unlinkSync(lockDir); + removed = true; + } catch { + // Lost the replacement race or cannot remove the invalid path. Retry + // through the normal deadline/backoff path below. + } + if (removed) { + continue; + } + stat = null; + } + if (stat) { + let ownerToken = null; + try { + ownerToken = fs.readFileSync(tokenFile, "utf8"); + } catch { + // Canonical locks are atomically published with an owner token. + // An unreadable token is preserved rather than reclaimed unsafely. + } + if (ownerToken && !isLockOwnerAlive(ownerToken)) { + const barrier = fs.mkdtempSync(`${lockDir}.reclaim-${process.pid}-`); + const claimed = path.join(barrier, "lock"); + let reclaimed = false; + try { + const currentToken = fs.readFileSync(tokenFile, "utf8"); + if (currentToken !== ownerToken) { + continue; + } + fs.renameSync(lockDir, claimed); + if (fs.readFileSync(path.join(claimed, "owner"), "utf8") === ownerToken && + !isLockOwnerAlive(ownerToken)) { + fs.rmSync(claimed, { recursive: true, force: true }); + reclaimed = true; + } else { + while (fs.existsSync(lockDir) && Date.now() < deadline) { + await sleep(25); + } + if (!fs.existsSync(lockDir)) { + fs.renameSync(claimed, lockDir); + } + } + } catch { + // Keep the barrier visible until the moved lock is restored or a + // later process recovers it after this reclaimer exits. + } finally { + if (!fs.existsSync(claimed)) { + fs.rmSync(barrier, { recursive: true, force: true }); + } + } + if (reclaimed) { + continue; + } + } + } + if (Date.now() >= deadline) { + throw Object.assign(new Error(`Timed out waiting for broker state lock: ${stateFile}`), { + code: BROKER_LOCK_TIMEOUT_CODE + }); + } + await sleep(25); + } + } + + try { + return await fn(); + } finally { + releaseOwnedBrokerLock(lockDir, token); + } +} export function createBrokerSessionDir(prefix = "cxc-") { return fs.mkdtempSync(path.join(os.tmpdir(), prefix)); @@ -40,19 +348,47 @@ export async function waitForBrokerEndpoint(endpoint, timeoutMs = 2000) { return false; } -export async function sendBrokerShutdown(endpoint) { +export async function sendBrokerShutdown(endpoint, { timeoutMs = BROKER_SHUTDOWN_TIMEOUT_MS } = {}) { + if (timeoutMs <= 0) { + return; + } await new Promise((resolve) => { - const socket = connectToEndpoint(endpoint); + let settled = false; + let timer = null; + const finish = () => { + if (settled) { + return; + } + settled = true; + if (timer) { + clearTimeout(timer); + } + resolve(); + }; + // Graceful shutdown is best-effort: a corrupt/unsupported endpoint makes + // connectToEndpoint (parseBrokerEndpoint) throw synchronously. Never reject, + // so callers always fall through to the forced process/file teardown. + let socket; + try { + socket = connectToEndpoint(endpoint); + } catch { + finish(); + return; + } + timer = setTimeout(() => { + socket.destroy(); + finish(); + }, timeoutMs); socket.setEncoding("utf8"); socket.on("connect", () => { socket.write(`${JSON.stringify({ id: 1, method: "broker/shutdown", params: {} })}\n`); }); socket.on("data", () => { socket.end(); - resolve(); + finish(); }); - socket.on("error", resolve); - socket.on("close", resolve); + socket.on("error", finish); + socket.on("close", finish); }); } @@ -73,23 +409,185 @@ function resolveBrokerStateFile(cwd) { return path.join(resolveStateDir(cwd), BROKER_STATE_FILE); } -export function loadBrokerSession(cwd) { - const stateFile = resolveBrokerStateFile(cwd); - if (!fs.existsSync(stateFile)) { +function readBoundedUtf8(descriptor, maxBytes) { + const buffer = Buffer.allocUnsafe(maxBytes + 1); + const bytesRead = fs.readSync(descriptor, buffer, 0, buffer.length, 0); + return bytesRead > maxBytes ? null : buffer.toString("utf8", 0, bytesRead); +} + +function readBrokerStateFile(stateFile) { + let descriptor = null; + try { + const before = fs.lstatSync(stateFile); + if (!before.isFile() || before.size > MAX_BROKER_STATE_BYTES) { + return null; + } + const noFollow = fs.constants.O_NOFOLLOW ?? 0; + const nonBlock = fs.constants.O_NONBLOCK ?? 0; + descriptor = fs.openSync(stateFile, fs.constants.O_RDONLY | noFollow | nonBlock); + const opened = fs.fstatSync(descriptor); + if (!opened.isFile() || opened.size > MAX_BROKER_STATE_BYTES || + opened.dev !== before.dev || opened.ino !== before.ino) { + return null; + } + const contents = readBoundedUtf8(descriptor, MAX_BROKER_STATE_BYTES); + return contents == null ? null : JSON.parse(contents); + } catch { return null; + } finally { + if (descriptor != null) { + try { + fs.closeSync(descriptor); + } catch { + // Ignore a descriptor already closed by a concurrent test shim. + } + } } +} +function readBrokerLockSession(lockDir) { + const sessionFile = path.join(lockDir, BROKER_LOCK_SESSION_FILE); + let descriptor = null; try { - return JSON.parse(fs.readFileSync(stateFile, "utf8")); + const before = fs.lstatSync(sessionFile); + if (!before.isFile() || before.size > 1024) { + return null; + } + const noFollow = fs.constants.O_NOFOLLOW ?? 0; + const nonBlock = fs.constants.O_NONBLOCK ?? 0; + descriptor = fs.openSync(sessionFile, fs.constants.O_RDONLY | noFollow | nonBlock); + const opened = fs.fstatSync(descriptor); + if (!opened.isFile() || opened.size > 1024 || + opened.dev !== before.dev || opened.ino !== before.ino) { + return null; + } + return readBoundedUtf8(descriptor, 1024); } catch { return null; + } finally { + if (descriptor != null) { + try { + fs.closeSync(descriptor); + } catch { + // Ignore a descriptor already closed by a concurrent test shim. + } + } + } +} + +export function loadBrokerSession(cwd) { + return readBrokerStateFile(resolveBrokerStateFile(cwd)); +} + +function endedSessionFile(sessionId) { + const digest = createHash("sha256").update(sessionId).digest("hex"); + return path.join(resolveStateRoot(), ENDED_SESSIONS_DIR, digest); +} + +export function markBrokerSessionEnded(sessionId) { + if (!sessionId || Buffer.byteLength(sessionId, "utf8") > 1024) { + return false; + } + const markerFile = endedSessionFile(sessionId); + fs.mkdirSync(path.dirname(markerFile), { recursive: true, mode: 0o700 }); + const candidate = `${markerFile}.tmp-${process.pid}-${brokerLockTokenSeq += 1}`; + try { + // Publish only a complete value. A direct `wx` write makes the final path + // visible before all bytes are present, allowing another process to + // transiently treat an ending session as active. + fs.writeFileSync(candidate, sessionId, { encoding: "utf8", flag: "wx", mode: 0o600 }); + fs.renameSync(candidate, markerFile); + } finally { + fs.rmSync(candidate, { force: true }); } + return true; +} + +export function isBrokerSessionEnded(sessionId) { + if (!sessionId || Buffer.byteLength(sessionId, "utf8") > 1024) { + return false; + } + const markerFile = endedSessionFile(sessionId); + let descriptor = null; + try { + const before = fs.lstatSync(markerFile); + if (!before.isFile() || before.size > 1024) { + throw Object.assign(new Error(`Invalid ended-session marker: ${markerFile}`), { + code: "EINVALIDENDEDSESSIONMARKER" + }); + } + const noFollow = fs.constants.O_NOFOLLOW ?? 0; + const nonBlock = fs.constants.O_NONBLOCK ?? 0; + descriptor = fs.openSync(markerFile, fs.constants.O_RDONLY | noFollow | nonBlock); + const opened = fs.fstatSync(descriptor); + if (!opened.isFile() || opened.size > 1024 || + opened.dev !== before.dev || opened.ino !== before.ino) { + throw Object.assign(new Error(`Changed ended-session marker: ${markerFile}`), { + code: "EINVALIDENDEDSESSIONMARKER" + }); + } + if (readBoundedUtf8(descriptor, 1024) !== sessionId) { + throw Object.assign(new Error(`Mismatched ended-session marker: ${markerFile}`), { + code: "EINVALIDENDEDSESSIONMARKER" + }); + } + return true; + } catch (error) { + if (error?.code === "ENOENT") { + return false; + } + throw error; + } finally { + if (descriptor != null) { + try { + fs.closeSync(descriptor); + } catch { + // Ignore a descriptor already closed by a concurrent test shim. + } + } + } +} + +export function assertBrokerSessionActive(sessionId) { + if (sessionId && isBrokerSessionEnded(sessionId)) { + throw brokerOwnerEndedError(sessionId); + } +} + +function applyEndedBrokerOwners(session) { + const endedOwners = new Set( + brokerSessionOwners(session).filter((owner) => isBrokerSessionEnded(owner)) + ); + if (endedOwners.size === 0) { + return { session, endedOwners }; + } + if (!session) { + return { session, endedOwners }; + } + const owners = brokerSessionOwners(session).filter((owner) => !endedOwners.has(owner)); + return { + session: { + ...session, + sessionId: owners[0] ?? null, + sessionIds: owners + }, + endedOwners + }; +} + +// Write broker.json atomically (temp + rename) so a concurrent reader — e.g. an +// unlocked ownership pre-check in another session's teardown — never observes a +// half-written file and mis-parses it. +function writeBrokerStateFile(stateFile, session) { + const tmp = `${stateFile}.tmp-${process.pid}-${(brokerLockTokenSeq += 1)}`; + fs.writeFileSync(tmp, `${JSON.stringify(session, null, 2)}\n`, "utf8"); + fs.renameSync(tmp, stateFile); } export function saveBrokerSession(cwd, session) { const stateDir = resolveStateDir(cwd); fs.mkdirSync(stateDir, { recursive: true }); - fs.writeFileSync(resolveBrokerStateFile(cwd), `${JSON.stringify(session, null, 2)}\n`, "utf8"); + writeBrokerStateFile(resolveBrokerStateFile(cwd), session); } export function clearBrokerSession(cwd) { @@ -110,24 +608,10 @@ async function isBrokerEndpointReady(endpoint) { } } -export async function ensureBrokerSession(cwd, options = {}) { - const existing = loadBrokerSession(cwd); - if (existing && (await isBrokerEndpointReady(existing.endpoint))) { - return existing; - } - - if (existing) { - teardownBrokerSession({ - endpoint: existing.endpoint ?? null, - pidFile: existing.pidFile ?? null, - logFile: existing.logFile ?? null, - sessionDir: existing.sessionDir ?? null, - pid: existing.pid ?? null, - killProcess: options.killProcess ?? null - }); - clearBrokerSession(cwd); - } - +// Spawn a broker process and wait for it to accept connections. Returns an +// unsaved session record, or null if it never became ready. No locking or +// persistence — the caller owns those. +async function spawnReadyBroker(cwd, options) { const sessionDir = createBrokerSessionDir(); const endpointFactory = options.createBrokerEndpoint ?? createBrokerEndpoint; const endpoint = endpointFactory(sessionDir, options.platform); @@ -159,15 +643,381 @@ export async function ensureBrokerSession(cwd, options = {}) { return null; } - const session = { + return { endpoint, pidFile, logFile, sessionDir, - pid: child.pid ?? null + pid: child.pid ?? null, + sessionId: resolveSessionId(options) }; - saveBrokerSession(cwd, session); - return session; +} + +// Tear down whatever broker is recorded for cwd (best effort) so a fresh one +// can replace it. +function discardBrokerSession(cwd, session, options) { + if (!session) { + return; + } + teardownBrokerSession({ + endpoint: session.endpoint ?? null, + pidFile: session.pidFile ?? null, + logFile: session.logFile ?? null, + sessionDir: session.sessionDir ?? null, + pid: session.pid ?? null, + killProcess: options.killProcess ?? null + }); + clearBrokerSession(cwd); +} + +function rejectEndedBrokerOwner(cwd, pending, sessionId, options) { + if (!sessionId || (!pending.endedOwners.has(sessionId) && !isBrokerSessionEnded(sessionId))) { + return; + } + const remainingOwners = brokerSessionOwners(pending.session); + if (pending.session && remainingOwners.length > 0) { + saveBrokerSession(cwd, pending.session); + } else { + discardBrokerSession(cwd, pending.session, options); + } + throw brokerOwnerEndedError(sessionId); +} + +// Adopt cwd's recorded broker for this session if it is still live, else spawn +// and persist a fresh one. Callers run this inside the state-file lock so two +// racing sessions cannot each leave an orphaned broker with no broker.json. +async function adoptOrSpawnBroker(cwd, options) { + const sessionId = resolveSessionId(options); + const pending = applyEndedBrokerOwners(loadBrokerSession(cwd)); + rejectEndedBrokerOwner(cwd, pending, sessionId, options); + const current = pending.session; + if (current && (await isBrokerEndpointReady(current.endpoint))) { + const pendingAfterProbe = applyEndedBrokerOwners(current); + rejectEndedBrokerOwner(cwd, pendingAfterProbe, sessionId, options); + const active = pendingAfterProbe.session; + const withOwner = withBrokerSessionOwner(active, sessionId); + if (withOwner !== current) { + saveBrokerSession(cwd, withOwner); + } + return withOwner; + } + discardBrokerSession(cwd, current, options); + + const session = await spawnReadyBroker(cwd, options); + if (!session) { + if (isBrokerSessionEnded(sessionId)) { + throw brokerOwnerEndedError(sessionId); + } + return null; + } + const spawnedWithOwner = withBrokerSessionOwner(session, session.sessionId); + const pendingAfterSpawn = applyEndedBrokerOwners(spawnedWithOwner); + const activeOwners = brokerSessionOwners(pendingAfterSpawn.session); + if (activeOwners.length === 0) { + teardownBrokerSession({ + endpoint: session.endpoint, + pidFile: session.pidFile, + logFile: session.logFile, + sessionDir: session.sessionDir, + pid: session.pid, + killProcess: options.killProcess ?? null + }); + throw brokerOwnerEndedError(session.sessionId); + } + saveBrokerSession(cwd, pendingAfterSpawn.session); + return pendingAfterSpawn.session; +} + +export async function ensureBrokerSession(cwd, options = {}) { + const stateFile = resolveBrokerStateFile(cwd); + const sessionId = resolveSessionId(options); + const lockOptions = options.lockTimeoutMs == null ? {} : { timeoutMs: options.lockTimeoutMs }; + + // Fast path: reuse a ready broker. The readiness probe runs outside the lock, + // so the broker may have been torn down (or replaced) before we acquired it; + // trust only the locked re-read, never the pre-lock snapshot. + for (let attempt = 0; attempt < 2; attempt += 1) { + const existing = loadBrokerSession(cwd); + if (!existing || !(await isBrokerEndpointReady(existing.endpoint))) { + break; + } + const reused = await withBrokerStateFileLock(stateFile, () => { + const pending = applyEndedBrokerOwners(loadBrokerSession(cwd)); + rejectEndedBrokerOwner(cwd, pending, sessionId, options); + const current = pending.session; + if (!current || current.endpoint !== existing.endpoint) { + return null; + } + const withOwner = withBrokerSessionOwner(current, sessionId); + if (withOwner !== current) { + saveBrokerSession(cwd, withOwner); + } + return withOwner; + }, { + ...lockOptions, + ownerSessionId: sessionId + }); + if (reused) { + return reused; + } + // State changed while we waited for the lock — re-probe: a replacement + // broker may already be live and reusable. + } + + // Slow path: adopt-or-spawn under the lock so concurrent spawns don't orphan + // brokers. resolveStateDir must exist before we can create the lock dir. + fs.mkdirSync(resolveStateDir(cwd), { recursive: true }); + return withBrokerStateFileLock(stateFile, () => adoptOrSpawnBroker(cwd, options), { + ...lockOptions, + ownerSessionId: sessionId + }); +} + +// Reuse cwd's recorded broker without claiming ownership or spawning a new +// broker. Probe-only callers still have to honor the global ended-session +// tombstone; reading broker.json directly can otherwise let an ended session +// reconnect after its teardown was deferred by lock contention. +export async function reuseBrokerSession(cwd, options = {}) { + const stateFile = resolveBrokerStateFile(cwd); + const sessionId = resolveSessionId(options); + // Always take the state lock, even when broker.json has not been published. + // SessionEnd can intentionally leave only a global tombstone when it races + // before the first broker lock; probe-only reuse must reconcile that status + // before it is allowed to fall back to a direct app-server. + fs.mkdirSync(resolveStateDir(cwd), { recursive: true }); + const lockOptions = options.lockTimeoutMs == null ? {} : { timeoutMs: options.lockTimeoutMs }; + return withBrokerStateFileLock(stateFile, async () => { + const pending = applyEndedBrokerOwners(loadBrokerSession(cwd)); + rejectEndedBrokerOwner(cwd, pending, sessionId, options); + + const current = pending.session; + if (!current || brokerSessionOwners(current).length === 0) { + discardBrokerSession(cwd, current, options); + return null; + } + if (!(await isBrokerEndpointReady(current.endpoint))) { + discardBrokerSession(cwd, current, options); + return null; + } + const pendingAfterProbe = applyEndedBrokerOwners(current); + rejectEndedBrokerOwner(cwd, pendingAfterProbe, sessionId, options); + if (pendingAfterProbe.session !== current) { + saveBrokerSession(cwd, pendingAfterProbe.session); + } + return pendingAfterProbe.session; + }, lockOptions); +} + +export async function teardownBrokerForCwd( + cwd, + sessionId, + { + fallbackSession = null, + killProcess = null, + shutdownTimeoutMs = BROKER_SHUTDOWN_TIMEOUT_MS, + lockTimeoutMs = BROKER_STATE_LOCK_TIMEOUT_MS + } = {} +) { + const stateFile = resolveBrokerStateFile(cwd); + const lockDir = `${stateFile}.lock`; + markBrokerSessionEnded(sessionId); + if (!fs.existsSync(stateFile) && !fs.existsSync(lockDir) && !fallbackSession) { + return false; + } + fs.mkdirSync(resolveStateDir(cwd), { recursive: true }); + + try { + return await withBrokerStateFileLock( + stateFile, + async () => { + const recorded = loadBrokerSession(cwd); + const pending = applyEndedBrokerOwners(recorded); + const current = pending.session; + const owners = brokerSessionOwners(current); + if (current && owners.length > 0) { + if (current !== recorded) { + writeBrokerStateFile(stateFile, current); + } + return false; + } + + const session = recorded ?? fallbackSession; + if (session?.endpoint) { + await sendBrokerShutdown(session.endpoint, { timeoutMs: shutdownTimeoutMs }); + } + teardownBrokerSession({ + endpoint: session?.endpoint ?? null, + pidFile: session?.pidFile ?? null, + logFile: session?.logFile ?? null, + sessionDir: session?.sessionDir ?? null, + pid: session?.pid ?? null, + killProcess + }); + if (fs.existsSync(stateFile)) { + fs.unlinkSync(stateFile); + } + return true; + }, + { timeoutMs: lockTimeoutMs } + ); + } catch (error) { + if (isBrokerLockTimeout(error)) { + throw brokerCleanupIncompleteError("lock-timeout"); + } + throw error; + } +} + +export async function teardownBrokersForSession( + sessionId, + { + killProcess = null, + shutdownTimeoutMs = BROKER_SHUTDOWN_TIMEOUT_MS, + lockTimeoutMs = BROKER_STATE_LOCK_TIMEOUT_MS, + budgetMs = null, + excludeCwd = null + } = {} +) { + if (!sessionId) { + return 0; + } + markBrokerSessionEnded(sessionId); + const stateRoot = resolveStateRoot(); + if (!fs.existsSync(stateRoot)) { + return 0; + } + + const deadline = budgetMs != null ? Date.now() + budgetMs : null; + if (deadline != null && Date.now() >= deadline) { + throw brokerCleanupIncompleteError("deadline"); + } + let count = 0; + let teardownError = null; + let incompleteReason = null; + const excludedStateFile = excludeCwd ? resolveBrokerStateFile(excludeCwd) : null; + const stateDirectory = fs.opendirSync(stateRoot); + try { + while (true) { + if (deadline != null && Date.now() >= deadline) { + incompleteReason = "deadline"; + break; + } + const entry = stateDirectory.readSync(); + if (!entry) { + break; + } + if (deadline != null && Date.now() >= deadline) { + incompleteReason = "deadline"; + break; + } + if (!entry.isDirectory() || entry.isSymbolicLink()) { + continue; + } + const stateFile = path.join(stateRoot, entry.name, BROKER_STATE_FILE); + if (stateFile === excludedStateFile) { + continue; + } + const stateExists = fs.existsSync(stateFile); + const lockDir = `${stateFile}.lock`; + const lockSessionId = readBrokerLockSession(lockDir); + if (!stateExists && lockSessionId !== sessionId) { + continue; + } + + // Ownership pre-check without the lock: never block on a lock held for a + // workspace this session does not own. The owner set only ever grows to + // include our sessionId (reuse) or shrinks when we ourselves remove it, so + // an unlocked read cannot falsely exclude a broker we own. A parse failure + // (e.g. a genuinely corrupt file) is NOT treated as "not ours" — fall + // through to the locked re-read, which is authoritative, rather than + // skipping a broker that might belong to this session. + let preview = null; + if (stateExists) { + try { + preview = readBrokerStateFile(stateFile); + } catch { + preview = null; + } + } + if ( + preview && + !brokerSessionOwners(preview).includes(sessionId) && + lockSessionId !== sessionId + ) { + continue; + } + const remaining = deadline != null ? deadline - Date.now() : lockTimeoutMs; + if (remaining <= 0) { + incompleteReason = "deadline"; + break; + } + const entryLockTimeout = Math.min(lockTimeoutMs, remaining); + + try { + await withBrokerStateFileLock( + stateFile, + async () => { + const recorded = readBrokerStateFile(stateFile); + if (!recorded) { + return; + } + const session = applyEndedBrokerOwners(recorded).session; + const owners = brokerSessionOwners(session); + if (owners.length > 0) { + if (session !== recorded) { + writeBrokerStateFile(stateFile, session); + } + return; + } + + // Bound the graceful-shutdown wait by the remaining scan budget, not + // just the per-RPC default: several unresponsive endpoints could + // otherwise each burn shutdownTimeoutMs and push the whole scan past + // the SessionEnd hook's budget. Out of budget → skip the RPC and let + // teardownBrokerSession terminate the process directly. + if (recorded.endpoint) { + const shutdownWait = + deadline != null ? Math.min(shutdownTimeoutMs, deadline - Date.now()) : shutdownTimeoutMs; + if (shutdownWait > 0) { + await sendBrokerShutdown(recorded.endpoint, { timeoutMs: shutdownWait }); + } + } + teardownBrokerSession({ + endpoint: recorded.endpoint ?? null, + pidFile: recorded.pidFile ?? null, + logFile: recorded.logFile ?? null, + sessionDir: recorded.sessionDir ?? null, + pid: recorded.pid ?? null, + killProcess + }); + if (fs.existsSync(stateFile)) { + fs.unlinkSync(stateFile); + } + count += 1; + }, + { timeoutMs: entryLockTimeout } + ); + } catch (error) { + if (!isBrokerLockTimeout(error)) { + teardownError ??= error; + continue; + } + // This entry's lock remained unavailable within the timeout. Skip it so + // the rest of this session's brokers still get torn down instead of + // aborting the whole scan, but report that cleanup was incomplete. + incompleteReason ??= "lock-timeout"; + } + } + } finally { + stateDirectory.closeSync(); + } + if (incompleteReason) { + throw brokerCleanupIncompleteError(incompleteReason, count, teardownError); + } + if (teardownError) { + throw teardownError; + } + return count; } export function teardownBrokerSession({ endpoint = null, pidFile, logFile, sessionDir = null, pid = null, killProcess = null }) { diff --git a/plugins/codex/scripts/lib/state.mjs b/plugins/codex/scripts/lib/state.mjs index 2da23498..ac260859 100644 --- a/plugins/codex/scripts/lib/state.mjs +++ b/plugins/codex/scripts/lib/state.mjs @@ -11,6 +11,10 @@ const FALLBACK_STATE_ROOT_DIR = path.join(os.tmpdir(), "codex-companion"); const STATE_FILE_NAME = "state.json"; const JOBS_DIR_NAME = "jobs"; const MAX_JOBS = 50; +const STATE_LOCK_TIMEOUT_MS = 5000; +const STATE_LOCK_TIMEOUT_CODE = "ESTATELOCKTIMEOUT"; + +let stateWriteSequence = 0; function nowIso() { return new Date().toISOString(); @@ -26,6 +30,239 @@ function defaultState() { }; } +function waitSynchronously(ms) { + Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, ms); +} + +function isLockOwnerAlive(token) { + const pid = Number.parseInt(token?.split("-", 1)[0] ?? "", 10); + if (!Number.isSafeInteger(pid) || pid <= 0) { + return true; + } + try { + process.kill(pid, 0); + return true; + } catch (error) { + return error?.code !== "ESRCH"; + } +} + +function reclaimBarrierPrefix(lockDir) { + return `${path.basename(lockDir)}.reclaim-`; +} + +function recoverStateReclaimBarriers(lockDir) { + const parent = path.dirname(lockDir); + const prefix = reclaimBarrierPrefix(lockDir); + let blocked = false; + for (const entry of fs.readdirSync(parent, { withFileTypes: true })) { + if (!entry.isDirectory() || !entry.name.startsWith(prefix)) { + continue; + } + const barrier = path.join(parent, entry.name); + const reclaimerPid = Number.parseInt(entry.name.slice(prefix.length).split("-", 1)[0], 10); + if (isLockOwnerAlive(`${reclaimerPid}-reclaimer`)) { + blocked = true; + continue; + } + const movedLock = path.join(barrier, "lock"); + if (fs.existsSync(movedLock)) { + let ownerToken = null; + try { + ownerToken = fs.readFileSync(path.join(movedLock, "owner"), "utf8"); + } catch { + // An unreadable moved lock is preserved conservatively. + } + if (!ownerToken || isLockOwnerAlive(ownerToken)) { + if (!fs.existsSync(lockDir)) { + try { + fs.renameSync(movedLock, lockDir); + fs.rmSync(barrier, { recursive: true, force: true }); + continue; + } catch { + // Another process restored or published the canonical lock. + } + } + blocked = true; + continue; + } + } + fs.rmSync(barrier, { recursive: true, force: true }); + } + return blocked; +} + +function releaseOwnedStateLock(lockDir, token) { + const tokenFile = path.join(lockDir, "owner"); + try { + if (fs.readFileSync(tokenFile, "utf8") === token) { + fs.rmSync(lockDir, { recursive: true, force: true }); + return; + } + } catch { + // A reclaimer may have moved the lock behind a visible barrier. + } + const parent = path.dirname(lockDir); + const prefix = reclaimBarrierPrefix(lockDir); + for (const entry of fs.readdirSync(parent, { withFileTypes: true })) { + if (!entry.isDirectory() || !entry.name.startsWith(prefix)) { + continue; + } + const barrier = path.join(parent, entry.name); + const movedLock = path.join(barrier, "lock"); + try { + if (fs.readFileSync(path.join(movedLock, "owner"), "utf8") === token) { + fs.rmSync(movedLock, { recursive: true, force: true }); + fs.rmSync(barrier, { recursive: true, force: true }); + break; + } + } catch { + // This barrier belongs to another lock generation. + } + } + // Orphan recovery may have restored this generation between the moved-token + // check and removal. Recheck the canonical path before returning so a + // completed owner cannot be stranded as a live-looking lock. + try { + if (fs.readFileSync(tokenFile, "utf8") === token) { + fs.rmSync(lockDir, { recursive: true, force: true }); + } + } catch { + // This owner no longer has a published lock generation. + } +} + +function withStateFileLock(cwd, fn, { timeoutMs = STATE_LOCK_TIMEOUT_MS } = {}) { + const stateFile = resolveStateFile(cwd); + const lockDir = `${stateFile}.lock`; + const tokenFile = path.join(lockDir, "owner"); + const token = `${process.pid}-${stateWriteSequence += 1}`; + const deadline = Date.now() + timeoutMs; + fs.mkdirSync(path.dirname(stateFile), { recursive: true }); + + while (true) { + if (recoverStateReclaimBarriers(lockDir)) { + if (Date.now() >= deadline) { + throw Object.assign(new Error(`Timed out waiting for state lock: ${stateFile}`), { + code: STATE_LOCK_TIMEOUT_CODE + }); + } + waitSynchronously(25); + continue; + } + let candidate = null; + try { + if (fs.existsSync(lockDir)) { + throw Object.assign(new Error(`State lock exists: ${stateFile}`), { code: "EEXIST" }); + } + candidate = fs.mkdtempSync(`${lockDir}.candidate-${process.pid}-`); + fs.writeFileSync(path.join(candidate, "owner"), token, { encoding: "utf8", flag: "wx" }); + // Publishing a non-empty, pre-stamped directory is atomic. A live holder + // can therefore never expose an unowned canonical lock path. + fs.renameSync(candidate, lockDir); + candidate = null; + if (recoverStateReclaimBarriers(lockDir)) { + if (fs.readFileSync(tokenFile, "utf8") === token) { + fs.rmSync(lockDir, { recursive: true, force: true }); + } + waitSynchronously(25); + continue; + } + break; + } catch (error) { + if (candidate) { + fs.rmSync(candidate, { recursive: true, force: true }); + } + if (!fs.existsSync(lockDir)) { + if (["EEXIST", "ENOTEMPTY", "EPERM"].includes(error?.code)) { + if (Date.now() >= deadline) { + throw Object.assign(new Error(`Timed out waiting for state lock: ${stateFile}`), { + code: STATE_LOCK_TIMEOUT_CODE + }); + } + waitSynchronously(25); + continue; + } + throw error; + } + let stat = null; + try { + stat = fs.lstatSync(lockDir); + } catch { + // Retry below; the lock may have disappeared after mkdirSync failed. + } + if (stat && !stat.isDirectory()) { + try { + fs.unlinkSync(lockDir); + continue; + } catch { + stat = null; + } + } + if (stat) { + let ownerToken = null; + try { + ownerToken = fs.readFileSync(tokenFile, "utf8"); + } catch { + // Every canonical lock is published with a token. Treat an unreadable + // token conservatively as live instead of risking overlapping writers. + } + if (ownerToken && !isLockOwnerAlive(ownerToken)) { + const barrier = fs.mkdtempSync(`${lockDir}.reclaim-${process.pid}-`); + const claimed = path.join(barrier, "lock"); + let reclaimed = false; + try { + const currentToken = fs.readFileSync(tokenFile, "utf8"); + if (currentToken !== ownerToken) { + continue; + } + fs.renameSync(lockDir, claimed); + if (fs.readFileSync(path.join(claimed, "owner"), "utf8") === ownerToken && + !isLockOwnerAlive(ownerToken)) { + fs.rmSync(claimed, { recursive: true, force: true }); + reclaimed = true; + } else { + while (fs.existsSync(lockDir) && Date.now() < deadline) { + waitSynchronously(25); + } + if (!fs.existsSync(lockDir)) { + fs.renameSync(claimed, lockDir); + } + } + } catch { + // The barrier stays visible until the moved lock is restored or a + // later process recovers it after this reclaimer exits. + } finally { + if (!fs.existsSync(claimed)) { + fs.rmSync(barrier, { recursive: true, force: true }); + } + } + if (reclaimed) { + continue; + } + } + } + if (Date.now() >= deadline) { + throw Object.assign(new Error(`Timed out waiting for state lock: ${stateFile}`), { + code: STATE_LOCK_TIMEOUT_CODE + }); + } + waitSynchronously(25); + } + } + + try { + return fn(); + } finally { + releaseOwnedStateLock(lockDir, token); + } +} + +export function resolveStateRoot() { + const pluginDataDir = process.env[PLUGIN_DATA_ENV]; + return pluginDataDir ? path.join(pluginDataDir, "state") : FALLBACK_STATE_ROOT_DIR; +} + export function resolveStateDir(cwd) { const workspaceRoot = resolveWorkspaceRoot(cwd); let canonicalWorkspaceRoot = workspaceRoot; @@ -38,8 +275,7 @@ export function resolveStateDir(cwd) { const slugSource = path.basename(workspaceRoot) || "workspace"; const slug = slugSource.replace(/[^a-zA-Z0-9._-]+/g, "-").replace(/^-+|-+$/g, "") || "workspace"; const hash = createHash("sha256").update(canonicalWorkspaceRoot).digest("hex").slice(0, 16); - const pluginDataDir = process.env[PLUGIN_DATA_ENV]; - const stateRoot = pluginDataDir ? path.join(pluginDataDir, "state") : FALLBACK_STATE_ROOT_DIR; + const stateRoot = resolveStateRoot(); return path.join(stateRoot, `${slug}-${hash}`); } @@ -89,7 +325,7 @@ function removeFileIfExists(filePath) { } } -export function saveState(cwd, state) { +function saveStateUnlocked(cwd, state) { const previousJobs = loadState(cwd).jobs; ensureStateDir(cwd); const nextJobs = pruneJobs(state.jobs ?? []); @@ -111,14 +347,44 @@ export function saveState(cwd, state) { removeFileIfExists(job.logFile); } - fs.writeFileSync(resolveStateFile(cwd), `${JSON.stringify(nextState, null, 2)}\n`, "utf8"); + const stateFile = resolveStateFile(cwd); + const temporaryStateFile = `${stateFile}.tmp-${process.pid}-${stateWriteSequence += 1}`; + try { + fs.writeFileSync(temporaryStateFile, `${JSON.stringify(nextState, null, 2)}\n`, "utf8"); + fs.renameSync(temporaryStateFile, stateFile); + } finally { + fs.rmSync(temporaryStateFile, { force: true }); + } return nextState; } +export function saveState(cwd, state) { + return withStateFileLock(cwd, () => saveStateUnlocked(cwd, state)); +} + export function updateState(cwd, mutate) { - const state = loadState(cwd); - mutate(state); - return saveState(cwd, state); + return withStateFileLock(cwd, () => { + const state = loadState(cwd); + mutate(state); + return saveStateUnlocked(cwd, state); + }); +} + +export function removeSessionJobs(cwd, sessionId, options = {}) { + const { beforeRemove = () => {}, ...lockOptions } = options; + return withStateFileLock(cwd, () => { + const state = loadState(cwd); + const removedJobs = state.jobs.filter((job) => job.sessionId === sessionId); + if (removedJobs.length === 0) { + return []; + } + beforeRemove(removedJobs); + saveStateUnlocked(cwd, { + ...state, + jobs: state.jobs.filter((job) => job.sessionId !== sessionId) + }); + return removedJobs; + }, lockOptions); } export function generateJobId(prefix = "job") { diff --git a/plugins/codex/scripts/session-lifecycle-hook.mjs b/plugins/codex/scripts/session-lifecycle-hook.mjs index 778571e6..a3c116c2 100644 --- a/plugins/codex/scripts/session-lifecycle-hook.mjs +++ b/plugins/codex/scripts/session-lifecycle-hook.mjs @@ -1,24 +1,39 @@ #!/usr/bin/env node import fs from "node:fs"; +import path from "node:path"; import process from "node:process"; +import { fileURLToPath } from "node:url"; import { terminateProcessTree } from "./lib/process.mjs"; import { BROKER_ENDPOINT_ENV } from "./lib/app-server.mjs"; import { - clearBrokerSession, + BROKER_CLEANUP_INCOMPLETE_CODE, LOG_FILE_ENV, - loadBrokerSession, PID_FILE_ENV, - sendBrokerShutdown, - teardownBrokerSession + markBrokerSessionEnded, + teardownBrokerForCwd, + teardownBrokersForSession } from "./lib/broker-lifecycle.mjs"; -import { loadState, resolveStateFile, saveState } from "./lib/state.mjs"; +import { removeSessionJobs, resolveStateFile, resolveStateRoot } from "./lib/state.mjs"; import { TRANSCRIPT_PATH_ENV } from "./lib/claude-session-transfer.mjs"; import { resolveWorkspaceRoot } from "./lib/workspace.mjs"; export const SESSION_ID_ENV = "CODEX_COMPANION_SESSION_ID"; const PLUGIN_DATA_ENV = "CLAUDE_PLUGIN_DATA"; +// hooks.json gives the SessionEnd hook a 5s budget. Share a shorter deadline +// across cross-workspace job discovery/cleanup and session-keyed broker +// teardown so the cwd fallback retains time before the hook is killed. +const SESSION_END_CLEANUP_BUDGET_MS = 3000; +const SESSION_END_LOCK_TIMEOUT_MS = 750; +const MAX_SESSION_JOB_STATE_BYTES = 1024 * 1024; + +function cleanupIncompleteError(reason) { + return Object.assign(new Error(`Session cleanup stopped before scanning all job state: ${reason}.`), { + code: BROKER_CLEANUP_INCOMPLETE_CODE, + reason + }); +} function readHookInput() { const raw = fs.readFileSync(0, "utf8").trim(); @@ -39,39 +54,158 @@ function appendEnvVar(name, value) { fs.appendFileSync(process.env.CLAUDE_ENV_FILE, `export ${name}=${shellEscape(value)}\n`, "utf8"); } -function cleanupSessionJobs(cwd, sessionId) { - if (!cwd || !sessionId) { - return; - } - - const workspaceRoot = resolveWorkspaceRoot(cwd); +function cleanupWorkspaceSessionJobs(workspaceRoot, sessionId, deadline) { const stateFile = resolveStateFile(workspaceRoot); if (!fs.existsSync(stateFile)) { - return; + return { brokerTeardownSafe: true, error: null }; + } + if (!readSessionJobsFromStateFile(stateFile).complete) { + return { brokerTeardownSafe: false, error: null }; } - const state = loadState(workspaceRoot); - const removedJobs = state.jobs.filter((job) => job.sessionId === sessionId); - if (removedJobs.length === 0) { - return; + let reachedTerminationPhase = false; + try { + removeSessionJobs( + workspaceRoot, + sessionId, + { + timeoutMs: Math.max(0, deadline - Date.now()), + beforeRemove(jobs) { + reachedTerminationPhase = true; + for (const job of jobs) { + const stillRunning = job.status === "queued" || job.status === "running"; + if (!stillRunning) { + continue; + } + try { + terminateProcessTree(job.pid ?? Number.NaN); + } catch { + // Ignore teardown failures during session shutdown. + } + } + } + } + ); + return { brokerTeardownSafe: true, error: null }; + } catch (error) { + return { brokerTeardownSafe: reachedTerminationPhase, error }; } +} - for (const job of removedJobs) { - const stillRunning = job.status === "queued" || job.status === "running"; - if (!stillRunning) { - continue; +function readSessionJobsFromStateFile(stateFile) { + let descriptor = null; + try { + const before = fs.lstatSync(stateFile); + if (!before.isFile() || before.size > MAX_SESSION_JOB_STATE_BYTES) { + return { jobs: [], complete: false }; } - try { - terminateProcessTree(job.pid ?? Number.NaN); - } catch { - // Ignore teardown failures during session shutdown. + const noFollow = fs.constants.O_NOFOLLOW ?? 0; + const nonBlock = fs.constants.O_NONBLOCK ?? 0; + descriptor = fs.openSync(stateFile, fs.constants.O_RDONLY | noFollow | nonBlock); + const opened = fs.fstatSync(descriptor); + if (!opened.isFile() || opened.size > MAX_SESSION_JOB_STATE_BYTES || + opened.dev !== before.dev || opened.ino !== before.ino) { + return { jobs: [], complete: false }; + } + const state = JSON.parse(fs.readFileSync(descriptor, "utf8")); + return { jobs: Array.isArray(state.jobs) ? state.jobs : [], complete: true }; + } catch (error) { + return { jobs: [], complete: error?.code === "ENOENT" }; + } finally { + if (descriptor != null) { + fs.closeSync(descriptor); } } +} - saveState(workspaceRoot, { - ...state, - jobs: state.jobs.filter((job) => job.sessionId !== sessionId) - }); +function findSessionJobWorkspaces(cwd, sessionId, deadline) { + const workspaceRoots = new Set([resolveWorkspaceRoot(cwd)]); + const stateRoot = resolveStateRoot(); + if (!fs.existsSync(stateRoot)) { + return { workspaceRoots, complete: true }; + } + + const stateDirectory = fs.opendirSync(stateRoot); + let complete = true; + let exhausted = false; + try { + while (Date.now() < deadline) { + const entry = stateDirectory.readSync(); + if (!entry) { + exhausted = true; + break; + } + if (entry.isSymbolicLink()) { + complete = false; + continue; + } + if (!entry.isDirectory()) { + continue; + } + const stateFile = path.join(stateRoot, entry.name, "state.json"); + const state = readSessionJobsFromStateFile(stateFile); + complete &&= state.complete; + for (const job of state.jobs) { + if (job.sessionId === sessionId && typeof job.workspaceRoot === "string" && job.workspaceRoot) { + workspaceRoots.add(job.workspaceRoot); + } + } + } + if (!exhausted) { + complete = false; + } + } finally { + stateDirectory.closeSync(); + } + return { workspaceRoots, complete }; +} + +function cleanupSessionJobs(cwd, sessionId, deadline) { + if (!cwd || !sessionId) { + return { discoveryComplete: true, cwdBrokerTeardownSafe: true, error: null }; + } + + let cleanupError = null; + let jobCleanupComplete = true; + let cwdBrokerTeardownSafe = true; + let workspaceIndex = 0; + const discovery = findSessionJobWorkspaces(cwd, sessionId, deadline); + if (!discovery.complete) { + cleanupError = cleanupIncompleteError("job-discovery-incomplete"); + } + for (const workspaceRoot of discovery.workspaceRoots) { + // Always clean the hook cwd first. Bound additional workspace cleanup by + // the shared SessionEnd deadline so broker cleanup and the cwd fallback + // retain time inside the hook's 5-second limit. + if (workspaceIndex > 0 && Date.now() >= deadline) { + jobCleanupComplete = false; + cleanupError ??= cleanupIncompleteError("job-cleanup-deadline"); + break; + } + workspaceIndex += 1; + try { + const cleanup = cleanupWorkspaceSessionJobs(workspaceRoot, sessionId, deadline); + cleanupError ??= cleanup.error; + if (!cleanup.brokerTeardownSafe) { + jobCleanupComplete = false; + cleanupError ??= cleanupIncompleteError("job-state-incomplete"); + if (workspaceIndex === 1) { + cwdBrokerTeardownSafe = false; + } + } + } catch (error) { + cleanupError ??= error; + jobCleanupComplete = false; + if (workspaceIndex === 1) { + cwdBrokerTeardownSafe = false; + } + } + } + return { + discoveryComplete: discovery.complete && jobCleanupComplete, + cwdBrokerTeardownSafe, + error: cleanupError + }; } function handleSessionStart(input) { @@ -80,37 +214,62 @@ function handleSessionStart(input) { appendEnvVar(PLUGIN_DATA_ENV, process.env[PLUGIN_DATA_ENV]); } -async function handleSessionEnd(input) { +export async function handleSessionEnd(input) { const cwd = input.cwd || process.cwd(); - const brokerSession = - loadBrokerSession(cwd) ?? - (process.env[BROKER_ENDPOINT_ENV] - ? { - endpoint: process.env[BROKER_ENDPOINT_ENV], - pidFile: process.env[PID_FILE_ENV] ?? null, - logFile: process.env[LOG_FILE_ENV] ?? null - } - : 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, input.session_id || process.env[SESSION_ID_ENV]); - teardownBrokerSession({ - endpoint: brokerEndpoint, - pidFile, - logFile, - sessionDir, - pid, - killProcess: terminateProcessTree - }); - clearBrokerSession(cwd); + const sessionId = input.session_id || process.env[SESSION_ID_ENV]; + const cleanupDeadline = Date.now() + SESSION_END_CLEANUP_BUDGET_MS; + let cleanupError = null; + try { + markBrokerSessionEnded(sessionId); + } catch (error) { + cleanupError = error; + } + let jobDiscoveryComplete = true; + let cwdBrokerTeardownSafe = true; + try { + const cleanup = cleanupSessionJobs(cwd, sessionId, cleanupDeadline); + cleanupError ??= cleanup.error; + jobDiscoveryComplete = cleanup.discoveryComplete; + cwdBrokerTeardownSafe = cleanup.cwdBrokerTeardownSafe; + } catch (error) { + cleanupError ??= error; + jobDiscoveryComplete = false; + cwdBrokerTeardownSafe = false; + } + + if (sessionId && jobDiscoveryComplete) { + try { + await teardownBrokersForSession(sessionId, { + killProcess: terminateProcessTree, + lockTimeoutMs: SESSION_END_LOCK_TIMEOUT_MS, + budgetMs: Math.max(0, cleanupDeadline - Date.now()), + excludeCwd: cwd + }); + } catch (error) { + cleanupError ??= error; + } + } + + if (cwdBrokerTeardownSafe) { + try { + await teardownBrokerForCwd(cwd, sessionId, { + fallbackSession: process.env[BROKER_ENDPOINT_ENV] + ? { + endpoint: process.env[BROKER_ENDPOINT_ENV], + pidFile: process.env[PID_FILE_ENV] ?? null, + logFile: process.env[LOG_FILE_ENV] ?? null + } + : null, + killProcess: terminateProcessTree, + lockTimeoutMs: SESSION_END_LOCK_TIMEOUT_MS + }); + } catch (error) { + cleanupError ??= error; + } + } + if (cleanupError) { + throw cleanupError; + } } async function main() { @@ -127,7 +286,26 @@ async function main() { } } -main().catch((error) => { - process.stderr.write(`${error instanceof Error ? error.message : String(error)}\n`); - process.exit(1); -}); +function isExecutedDirectly() { + if (!process.argv[1]) { + return false; + } + + try { + return fs.realpathSync(fileURLToPath(import.meta.url)) === + fs.realpathSync(path.resolve(process.argv[1])); + } catch { + // argv[1] may not resolve to a real path (deleted file, loader indirection); + // importing must never crash at module load. + return false; + } +} + +// Only run main() when executed directly (not when imported by tests). +const isMain = isExecutedDirectly(); +if (isMain) { + main().catch((error) => { + process.stderr.write(`${error instanceof Error ? error.message : String(error)}\n`); + process.exit(1); + }); +} diff --git a/tests/broker-lifecycle.test.mjs b/tests/broker-lifecycle.test.mjs new file mode 100644 index 00000000..7e16490f --- /dev/null +++ b/tests/broker-lifecycle.test.mjs @@ -0,0 +1,1964 @@ +import fs from "node:fs"; +import net from "node:net"; +import os from "node:os"; +import path from "node:path"; +import { spawn } from "node:child_process"; +import { pathToFileURL } from "node:url"; +import test from "node:test"; +import assert from "node:assert/strict"; + +import { makeTempDir } from "./helpers.mjs"; +import { createBrokerEndpoint, parseBrokerEndpoint } from "../plugins/codex/scripts/lib/broker-endpoint.mjs"; +import { + loadState, + resolveJobFile, + resolveJobLogFile, + resolveStateDir, + resolveStateFile, + resolveStateRoot, + saveState, + writeJobFile +} from "../plugins/codex/scripts/lib/state.mjs"; +import { + BROKER_CLEANUP_INCOMPLETE_CODE, + BROKER_OWNER_ENDED_CODE, + ensureBrokerSession, + isBrokerSessionEnded, + loadBrokerSession, + resolveSessionId, + saveBrokerSession, + sendBrokerShutdown, + teardownBrokerForCwd, + teardownBrokersForSession +} from "../plugins/codex/scripts/lib/broker-lifecycle.mjs"; +import { CodexAppServerClient } from "../plugins/codex/scripts/lib/app-server.mjs"; +import { handleSessionEnd } from "../plugins/codex/scripts/session-lifecycle-hook.mjs"; + +function sleep(ms) { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +function observeLockAttempt(lockDir) { + const originalMkdirSync = fs.mkdirSync; + const originalExistsSync = fs.existsSync; + let notify; + const attempted = new Promise((resolve) => { + notify = resolve; + }); + fs.mkdirSync = (dir, ...args) => { + if (dir === lockDir) { + notify(); + } + return originalMkdirSync.call(fs, dir, ...args); + }; + fs.existsSync = (target) => { + if (target === lockDir) { + notify(); + } + return originalExistsSync.call(fs, target); + }; + return { + attempted, + restore() { + fs.mkdirSync = originalMkdirSync; + fs.existsSync = originalExistsSync; + } + }; +} + +// Minimal stand-in for app-server-broker.mjs: honors the spawn contract +// (serve --endpoint --cwd --pid-file) enough for waitForBrokerEndpoint. +const FAKE_BROKER_SCRIPT = `import fs from "node:fs"; +import net from "node:net"; + +if (process.env.TEST_BROKER_SPAWN_MARKER) { + fs.writeFileSync(process.env.TEST_BROKER_SPAWN_MARKER, "spawned", "utf8"); +} + +const args = process.argv.slice(2); +const get = (name) => args[args.indexOf(name) + 1]; +const sockPath = get("--endpoint").replace(/^(?:unix|pipe):/, ""); +const server = net.createServer((socket) => socket.end()); +setTimeout(() => { + server.listen(sockPath, () => { + fs.writeFileSync(get("--pid-file"), String(process.pid), "utf8"); + }); +}, Number(process.env.TEST_BROKER_LISTEN_DELAY_MS || 0)); +`; + +function writeFakeBrokerScript() { + const scriptPath = path.join(makeTempDir(), "fake-broker.mjs"); + fs.writeFileSync(scriptPath, FAKE_BROKER_SCRIPT, "utf8"); + return scriptPath; +} + +const stateRootForTest = resolveStateRoot; + +test("resolveStateRoot uses CLAUDE_PLUGIN_DATA/state when set", () => { + const pluginData = makeTempDir(); + const prev = process.env.CLAUDE_PLUGIN_DATA; + process.env.CLAUDE_PLUGIN_DATA = pluginData; + try { + assert.equal(resolveStateRoot(), path.join(pluginData, "state")); + } finally { + if (prev == null) delete process.env.CLAUDE_PLUGIN_DATA; + else process.env.CLAUDE_PLUGIN_DATA = prev; + } +}); + +test("resolveStateRoot falls back to a tmp dir when CLAUDE_PLUGIN_DATA is unset", () => { + const prev = process.env.CLAUDE_PLUGIN_DATA; + delete process.env.CLAUDE_PLUGIN_DATA; + try { + assert.equal(resolveStateRoot(), path.join(os.tmpdir(), "codex-companion")); + } finally { + if (prev != null) process.env.CLAUDE_PLUGIN_DATA = prev; + } +}); + +test("resolveSessionId prefers explicit option, then env, then null", () => { + assert.equal(resolveSessionId({ sessionId: "explicit" }), "explicit"); + assert.equal(resolveSessionId({ env: { CODEX_COMPANION_SESSION_ID: "from-env" } }), "from-env"); +}); + +test("resolveSessionId reads process.env when no option/env given", () => { + const prev = process.env.CODEX_COMPANION_SESSION_ID; + process.env.CODEX_COMPANION_SESSION_ID = "proc-env"; + try { + assert.equal(resolveSessionId({}), "proc-env"); + } finally { + if (prev == null) delete process.env.CODEX_COMPANION_SESSION_ID; + else process.env.CODEX_COMPANION_SESSION_ID = prev; + } +}); + +test("resolveSessionId returns null when nothing is set", () => { + const prev = process.env.CODEX_COMPANION_SESSION_ID; + delete process.env.CODEX_COMPANION_SESSION_ID; + try { + assert.equal(resolveSessionId({ env: {} }), null); + } finally { + if (prev != null) process.env.CODEX_COMPANION_SESSION_ID = prev; + } +}); + +function writeBrokerJson(stateRoot, dirName, session) { + const dir = path.join(stateRoot, dirName); + fs.mkdirSync(dir, { recursive: true }); + const file = path.join(dir, "broker.json"); + fs.writeFileSync(file, JSON.stringify(session), "utf8"); + return file; +} + +function withPluginData(fn) { + const pluginData = makeTempDir(); + const prev = process.env.CLAUDE_PLUGIN_DATA; + process.env.CLAUDE_PLUGIN_DATA = pluginData; + return Promise.resolve(fn(pluginData)).finally(() => { + if (prev == null) delete process.env.CLAUDE_PLUGIN_DATA; + else process.env.CLAUDE_PLUGIN_DATA = prev; + }); +} + +async function withReadyBroker(fn) { + const sessionDir = makeTempDir(); + const endpoint = createBrokerEndpoint(sessionDir); + const target = parseBrokerEndpoint(endpoint); + const requests = []; + const server = net.createServer((socket) => { + socket.setEncoding("utf8"); + socket.on("data", (chunk) => { + requests.push(chunk); + socket.write(`${JSON.stringify({ id: 1, result: {} })}\n`); + socket.end(); + }); + }); + + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(target.path, () => { + server.off("error", reject); + resolve(); + }); + }); + + try { + return await fn({ endpoint, requests, sessionDir }); + } finally { + await new Promise((resolve) => server.close(resolve)); + } +} + +async function withHangingBroker(fn) { + const sessionDir = makeTempDir(); + const endpoint = createBrokerEndpoint(sessionDir); + const target = parseBrokerEndpoint(endpoint); + const requests = []; + const sockets = []; + const server = net.createServer((socket) => { + sockets.push(socket); + socket.setEncoding("utf8"); + socket.on("data", (chunk) => { + requests.push(chunk); + }); + }); + + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(target.path, () => { + server.off("error", reject); + resolve(); + }); + }); + + const destroySockets = () => { + for (const socket of sockets) { + socket.destroy(); + } + }; + + try { + return await fn({ endpoint, requests, sessionDir, destroySockets }); + } finally { + destroySockets(); + await new Promise((resolve) => server.close(resolve)); + } +} + +function waitForChild(child) { + return new Promise((resolve, reject) => { + child.once("error", reject); + child.once("exit", (code, signal) => { + resolve({ code, signal }); + }); + }); +} + +test("teardownBrokersForSession tears down a broker registered for a different cwd", async () => { + await withPluginData(async () => { + const stateRoot = stateRootForTest(); + const sessionDir = makeTempDir(); + const pidFile = path.join(sessionDir, "broker.pid"); + const logFile = path.join(sessionDir, "broker.log"); + fs.writeFileSync(pidFile, "12345\n"); + fs.writeFileSync(logFile, ""); + const brokerJson = writeBrokerJson(stateRoot, "worktree-deadbeefdeadbeef", { + endpoint: "unix:/tmp/codex-test-nonexistent.sock", + pidFile, logFile, sessionDir, pid: 12345, sessionId: "S" + }); + + const killed = []; + const count = await teardownBrokersForSession("S", { killProcess: (pid) => killed.push(pid) }); + + assert.equal(count, 1); + assert.deepEqual(killed, [12345]); + assert.equal(fs.existsSync(brokerJson), false); + }); +}); + +test("teardownBrokersForSession ignores broker.json without sessionId (legacy)", async () => { + await withPluginData(async () => { + const stateRoot = stateRootForTest(); + const brokerJson = writeBrokerJson(stateRoot, "legacy-2222222222222222", { + endpoint: "unix:/tmp/codex-test-nonexistent3.sock", pid: null + }); + const count = await teardownBrokersForSession("S", { killProcess: () => {} }); + assert.equal(count, 0); + assert.equal(fs.existsSync(brokerJson), true); + }); +}); + +test("teardownBrokersForSession is a no-op for empty sessionId", async () => { + await withPluginData(async () => { + const count = await teardownBrokersForSession("", { killProcess: () => {} }); + assert.equal(count, 0); + }); +}); + +test("reusing a ready broker transfers cleanup ownership to the later session", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, requests, sessionDir }) => { + const cwd = makeTempDir(); + const pidFile = path.join(sessionDir, "broker.pid"); + const logFile = path.join(sessionDir, "broker.log"); + fs.writeFileSync(pidFile, "12345\n"); + fs.writeFileSync(logFile, ""); + saveBrokerSession(cwd, { + endpoint, + pidFile, + logFile, + sessionDir, + pid: 12345, + sessionId: "A" + }); + + const reused = await ensureBrokerSession(cwd, { env: { CODEX_COMPANION_SESSION_ID: "B" } }); + assert.equal(reused.endpoint, endpoint); + assert.deepEqual(loadBrokerSession(cwd).sessionIds, ["A", "B"]); + + const killed = []; + assert.equal(await teardownBrokersForSession("A", { killProcess: (pid) => killed.push(pid) }), 0); + assert.deepEqual(killed, []); + assert.equal(loadBrokerSession(cwd).sessionId, "B"); + assert.deepEqual(loadBrokerSession(cwd).sessionIds, ["B"]); + assert.equal(requests.length, 0); + + assert.equal(await teardownBrokersForSession("B", { killProcess: (pid) => killed.push(pid) }), 1); + assert.deepEqual(killed, [12345]); + assert.equal(loadBrokerSession(cwd), null); + assert.equal(requests.length, 1); + }); + }); +}); + +test("handleSessionEnd removes only the ending owner from a shared cwd broker", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, requests, sessionDir }) => { + const cwd = makeTempDir(); + const pidFile = path.join(sessionDir, "broker.pid"); + const logFile = path.join(sessionDir, "broker.log"); + fs.writeFileSync(pidFile, "12345\n"); + fs.writeFileSync(logFile, ""); + saveBrokerSession(cwd, { + endpoint, + pidFile, + logFile, + sessionDir, + pid: 12345, + sessionId: "A", + sessionIds: ["A", "B"] + }); + + await handleSessionEnd({ cwd, session_id: "A" }); + + assert.equal(loadBrokerSession(cwd).sessionId, "B"); + assert.deepEqual(loadBrokerSession(cwd).sessionIds, ["B"]); + assert.equal(requests.length, 0); + }); + }); +}); + +test("concurrent SessionEnd hooks tear down a shared broker after the last owner exits", async () => { + await withPluginData(async (pluginData) => { + await withReadyBroker(async ({ endpoint, requests, sessionDir }) => { + const cwd = makeTempDir(); + const markerDir = makeTempDir(); + const pidFile = path.join(sessionDir, "broker.pid"); + const logFile = path.join(sessionDir, "broker.log"); + fs.writeFileSync(pidFile, "12345\n"); + fs.writeFileSync(logFile, ""); + const brokerJson = writeBrokerJson(stateRootForTest(), "worktree-race-deadbeef", { + endpoint, + pidFile, + logFile, + sessionDir, + pid: 12345, + sessionId: "A", + sessionIds: ["A", "B"] + }); + + const moduleUrl = pathToFileURL(path.resolve("plugins/codex/scripts/lib/broker-lifecycle.mjs")).href; + const script = ` + import fs from "node:fs"; + import path from "node:path"; + + const stateFile = process.env.TEST_BROKER_STATE_FILE; + const markerDir = process.env.TEST_MARKER_DIR; + const sessionId = process.env.TEST_SESSION_ID; + const otherSessionId = sessionId === "A" ? "B" : "A"; + const lockDir = stateFile + ".lock"; + const originalMkdirSync = fs.mkdirSync.bind(fs); + let reachedBarrier = false; + + fs.mkdirSync = (dir, ...args) => { + if (dir === lockDir && !reachedBarrier) { + reachedBarrier = true; + fs.writeFileSync(path.join(markerDir, sessionId + ".ready"), "", "utf8"); + const otherReady = path.join(markerDir, otherSessionId + ".ready"); + const deadline = Date.now() + 2000; + while (!fs.existsSync(otherReady) && Date.now() < deadline) {} + if (!fs.existsSync(otherReady)) { + throw new Error("other SessionEnd did not reach the lock barrier"); + } + } + return originalMkdirSync(dir, ...args); + }; + + const { teardownBrokersForSession } = await import(process.env.TEST_BROKER_MODULE_URL); + await teardownBrokersForSession(sessionId, { killProcess: () => {} }); + `; + + const makeChild = (sessionId) => + spawn(process.execPath, ["--input-type=module", "-e", script], { + cwd: path.resolve("."), + env: { + ...process.env, + CLAUDE_PLUGIN_DATA: pluginData, + TEST_BROKER_STATE_FILE: brokerJson, + TEST_BROKER_MODULE_URL: moduleUrl, + TEST_MARKER_DIR: markerDir, + TEST_SESSION_ID: sessionId + }, + stdio: ["ignore", "pipe", "pipe"] + }); + + const childA = makeChild("A"); + const childB = makeChild("B"); + const [resultA, resultB] = await Promise.all([waitForChild(childA), waitForChild(childB)]); + + assert.deepEqual([resultA.code, resultB.code], [0, 0]); + assert.equal(fs.existsSync(brokerJson), false); + assert.equal(requests.length, 1); + }); + }); +}); + +test("handleSessionEnd still tears down session brokers when job cleanup fails", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + const workspaceStateDir = resolveStateDir(cwd); + saveBrokerSession(cwd, { + endpoint: "unix:/tmp/codex-test-nonexistent-cleanup.sock", + pidFile: null, + logFile: null, + sessionDir: null, + pid: null, + sessionId: "S", + sessionIds: ["S"] + }); + const badLogFile = path.join(workspaceStateDir, "bad.log"); + fs.mkdirSync(badLogFile); + const originalKill = process.kill; + const killedPids = []; + process.kill = (pid, signal) => { + assert.deepEqual(loadState(cwd).jobs.map((job) => job.id), ["running"]); + assert.equal(fs.existsSync(badLogFile), true); + killedPids.push({ pid, signal }); + return true; + }; + const stateFile = path.join(workspaceStateDir, "state.json"); + fs.writeFileSync( + stateFile, + `${JSON.stringify({ + version: 1, + config: { stopReviewGate: false }, + jobs: [{ id: "running", status: "running", sessionId: "S", pid: 12345, logFile: badLogFile }] + }, null, 2)}\n`, + "utf8" + ); + + try { + await assert.rejects(() => handleSessionEnd({ cwd, session_id: "S" }), { code: "EISDIR" }); + } finally { + process.kill = originalKill; + } + assert.deepEqual(killedPids, [{ pid: -12345, signal: "SIGTERM" }]); + assert.deepEqual(loadState(cwd).jobs.map((job) => job.id), ["running"]); + assert.equal(fs.existsSync(path.join(workspaceStateDir, "broker.json")), false); + }); +}); + +test("handleSessionEnd preserves session brokers when job state locking fails", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + const otherWorkspace = makeTempDir(); + const stateFile = resolveStateFile(cwd); + const stateLock = `${stateFile}.lock`; + saveState(cwd, { + jobs: [{ + id: "locked-job", + status: "running", + sessionId: "S", + workspaceRoot: cwd, + pid: 999999999 + }] + }); + saveBrokerSession(cwd, { + endpoint: "unix:/tmp/codex-test-locked-job-cwd.sock", + sessionId: "S", + sessionIds: ["S"] + }); + saveBrokerSession(otherWorkspace, { + endpoint: "unix:/tmp/codex-test-locked-job-other.sock", + sessionId: "S", + sessionIds: ["S"] + }); + + const originalRenameSync = fs.renameSync; + fs.renameSync = (source, destination) => { + if (destination === stateLock) { + throw Object.assign(new Error("state lock denied"), { code: "EACCES" }); + } + return originalRenameSync.call(fs, source, destination); + }; + + try { + await assert.rejects(() => handleSessionEnd({ cwd, session_id: "S" }), { code: "EACCES" }); + } finally { + fs.renameSync = originalRenameSync; + } + + assert.deepEqual(loadState(cwd).jobs.map((job) => job.id), ["locked-job"]); + assert.notEqual(loadBrokerSession(cwd), null); + assert.notEqual(loadBrokerSession(otherWorkspace), null); + }); +}); + +test("handleSessionEnd reports incomplete cleanup when cwd job state is unreadable", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + const otherWorkspace = makeTempDir(); + const stateFile = resolveStateFile(cwd); + fs.mkdirSync(path.dirname(stateFile), { recursive: true }); + fs.writeFileSync(stateFile, "{not-json", "utf8"); + saveBrokerSession(cwd, { + endpoint: "unix:/tmp/codex-test-unreadable-job-cwd.sock", + sessionId: "S", + sessionIds: ["S"] + }); + saveBrokerSession(otherWorkspace, { + endpoint: "unix:/tmp/codex-test-unreadable-job-other.sock", + sessionId: "S", + sessionIds: ["S"] + }); + + await assert.rejects( + () => handleSessionEnd({ cwd, session_id: "S" }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE, reason: "job-discovery-incomplete" } + ); + + assert.equal(fs.readFileSync(stateFile, "utf8"), "{not-json"); + assert.notEqual(loadBrokerSession(cwd), null); + assert.notEqual(loadBrokerSession(otherWorkspace), null); + }); +}); + +test("handleSessionEnd cleans session jobs from a different --cwd workspace before broker teardown", async () => { + await withPluginData(async () => { + const hookCwd = makeTempDir(); + const jobWorkspace = makeTempDir(); + const jobId = "cross-cwd-job"; + const jobLog = resolveJobLogFile(jobWorkspace, jobId); + const job = { + id: jobId, + status: "running", + sessionId: "S", + workspaceRoot: jobWorkspace, + pid: 999999999, + logFile: jobLog + }; + fs.writeFileSync(jobLog, "running\n", "utf8"); + writeJobFile(jobWorkspace, jobId, job); + saveState(jobWorkspace, { jobs: [job] }); + saveBrokerSession(jobWorkspace, { + endpoint: "unix:/tmp/codex-test-cross-cwd.sock", + pidFile: null, + logFile: null, + sessionDir: null, + pid: null, + sessionId: "S", + sessionIds: ["S"] + }); + + await handleSessionEnd({ cwd: hookCwd, session_id: "S" }); + + assert.deepEqual(loadState(jobWorkspace).jobs, []); + assert.equal(fs.existsSync(resolveJobFile(jobWorkspace, jobId)), false); + assert.equal(fs.existsSync(jobLog), false); + assert.equal(loadBrokerSession(jobWorkspace), null); + }); +}); + +test("handleSessionEnd continues cross-cwd cleanup after one workspace cleanup fails", async () => { + await withPluginData(async () => { + const hookCwd = makeTempDir(); + const otherWorkspace = makeTempDir(); + const badLog = path.join(resolveStateDir(hookCwd), "bad.log"); + fs.mkdirSync(badLog, { recursive: true }); + saveState(hookCwd, { + jobs: [{ + id: "bad-job", + status: "completed", + sessionId: "S", + workspaceRoot: hookCwd, + logFile: badLog + }] + }); + + const otherJobId = "other-job"; + const otherLog = resolveJobLogFile(otherWorkspace, otherJobId); + const otherJob = { + id: otherJobId, + status: "running", + sessionId: "S", + workspaceRoot: otherWorkspace, + pid: 999999999, + logFile: otherLog + }; + fs.writeFileSync(otherLog, "running\n", "utf8"); + writeJobFile(otherWorkspace, otherJobId, otherJob); + saveState(otherWorkspace, { jobs: [otherJob] }); + saveBrokerSession(otherWorkspace, { + endpoint: "unix:/tmp/codex-test-other-workspace.sock", + sessionId: "S", + sessionIds: ["S"] + }); + + await assert.rejects(() => handleSessionEnd({ cwd: hookCwd, session_id: "S" }), { code: "EISDIR" }); + + assert.deepEqual(loadState(otherWorkspace).jobs, []); + assert.equal(fs.existsSync(resolveJobFile(otherWorkspace, otherJobId)), false); + assert.equal(fs.existsSync(otherLog), false); + assert.equal(loadBrokerSession(otherWorkspace), null); + }); +}); + +test("handleSessionEnd reports incomplete cleanup when the deadline skips a discovered workspace", async () => { + await withPluginData(async () => { + const hookCwd = makeTempDir(); + const otherWorkspace = makeTempDir(); + saveState(hookCwd, { + jobs: [{ + id: "hook-job", + status: "completed", + sessionId: "S", + workspaceRoot: hookCwd + }] + }); + saveState(otherWorkspace, { + jobs: [{ + id: "other-job", + status: "completed", + sessionId: "S", + workspaceRoot: otherWorkspace + }] + }); + saveBrokerSession(otherWorkspace, { + endpoint: "unix:/tmp/codex-test-job-cleanup-deadline.sock", + sessionId: "S", + sessionIds: ["S"] + }); + + const hookStateFile = resolveStateFile(hookCwd); + const originalRenameSync = fs.renameSync; + const originalNow = Date.now; + let now = 0; + Date.now = () => now; + fs.renameSync = (source, destination) => { + const result = originalRenameSync.call(fs, source, destination); + if (destination === hookStateFile) { + now = 3000; + } + return result; + }; + + try { + await assert.rejects( + () => handleSessionEnd({ cwd: hookCwd, session_id: "S" }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE, reason: "job-cleanup-deadline" } + ); + } finally { + Date.now = originalNow; + fs.renameSync = originalRenameSync; + } + + assert.deepEqual(loadState(otherWorkspace).jobs.map((job) => job.id), ["other-job"]); + assert.notEqual(loadBrokerSession(otherWorkspace), null); + }); +}); + +test("handleSessionEnd reports incomplete cleanup when job-state discovery is incomplete", async () => { + await withPluginData(async () => { + const hookCwd = makeTempDir(); + const jobWorkspace = makeTempDir(); + const stateDir = resolveStateDir(jobWorkspace); + fs.mkdirSync(stateDir, { recursive: true }); + fs.writeFileSync( + path.join(stateDir, "state.json"), + JSON.stringify({ + jobs: [{ + id: "oversized-job", + status: "running", + sessionId: "S", + workspaceRoot: jobWorkspace, + padding: "x".repeat(1024 * 1024) + }] + }), + "utf8" + ); + saveBrokerSession(jobWorkspace, { + endpoint: "unix:/tmp/codex-test-oversized-state.sock", + sessionId: "S", + sessionIds: ["S"] + }); + + await assert.rejects( + () => handleSessionEnd({ cwd: hookCwd, session_id: "S" }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE, reason: "job-discovery-incomplete" } + ); + + assert.equal(loadState(jobWorkspace).jobs.length, 1); + assert.notEqual(loadBrokerSession(jobWorkspace), null); + assert.equal(isBrokerSessionEnded("S"), true); + }); +}); + +test("handleSessionEnd does not cap a complete job-state discovery scan", async () => { + await withPluginData(async () => { + const hookCwd = makeTempDir(); + const stateRoot = stateRootForTest(); + for (let index = 0; index < 1001; index += 1) { + fs.mkdirSync(path.join(stateRoot, `empty-${String(index).padStart(4, "0")}`), { recursive: true }); + } + + await handleSessionEnd({ cwd: hookCwd, session_id: "S" }); + }); +}); + +test("handleSessionEnd gives broker teardown only the remaining shared cleanup budget", async () => { + await withPluginData(async () => { + const hookCwd = makeTempDir(); + const brokerWorkspace = makeTempDir(); + saveBrokerSession(brokerWorkspace, { + endpoint: "unix:/tmp/codex-test-shared-deadline.sock", + sessionId: "S", + sessionIds: ["S"] + }); + + const probe = fs.opendirSync(stateRootForTest()); + const dirPrototype = Object.getPrototypeOf(probe); + probe.closeSync(); + const originalReadSync = dirPrototype.readSync; + const originalNow = Date.now; + let now = 0; + Date.now = () => now; + dirPrototype.readSync = function (...args) { + const entry = originalReadSync.call(this, ...args); + if (!entry) { + now = 3000; + } + return entry; + }; + + try { + await assert.rejects( + () => handleSessionEnd({ cwd: hookCwd, session_id: "S" }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE } + ); + } finally { + Date.now = originalNow; + dirPrototype.readSync = originalReadSync; + } + + assert.notEqual(loadBrokerSession(brokerWorkspace), null); + }); +}); + +test("teardownBrokersForSession times out unresponsive broker shutdown requests", async () => { + await withPluginData(async () => { + await withHangingBroker(async ({ endpoint, requests, sessionDir, destroySockets }) => { + const stateRoot = stateRootForTest(); + const brokerJson = writeBrokerJson(stateRoot, "hanging-shutdown-deadbeef", { + endpoint, + pidFile: null, + logFile: null, + sessionDir, + pid: null, + sessionId: "S", + sessionIds: ["S"] + }); + + let completed = false; + const teardown = teardownBrokersForSession("S", { killProcess: () => {}, shutdownTimeoutMs: 50 }) + .then(() => { + completed = true; + }); + try { + await new Promise((resolve) => setTimeout(resolve, 150)); + assert.equal(completed, true); + } finally { + destroySockets(); + await teardown; + } + assert.equal(fs.existsSync(brokerJson), false); + assert.equal(requests.length, 1); + }); + }); +}); + +test("teardownBrokersForSession reports a same-session locked entry after tearing down the rest", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, requests, sessionDir }) => { + const stateRoot = stateRootForTest(); + + // Entry 1: owned by this session but its lock is held by a concurrent hook + // that never released it. + const lockedCwd = makeTempDir(); + saveBrokerSession(lockedCwd, { + endpoint, pidFile: null, logFile: null, sessionDir, pid: null, + sessionId: "S", sessionIds: ["S"] + }); + const lockedJson = path.join(resolveStateDir(lockedCwd), "broker.json"); + fs.mkdirSync(`${lockedJson}.lock`); + + // Entry 2: another record owned by the same session, later in the scan. + const cleanableJson = writeBrokerJson(stateRoot, "worktree-cleanablebbbb", { + endpoint: "invalid:endpoint", pidFile: null, logFile: null, + sessionDir: null, pid: null, sessionId: "S" + }); + + try { + await assert.rejects( + () => teardownBrokersForSession("S", { killProcess: () => {}, lockTimeoutMs: 60 }), + (error) => { + assert.equal(error.code, BROKER_CLEANUP_INCOMPLETE_CODE); + assert.equal(error.reason, "lock-timeout"); + assert.equal(error.count, 1); + return true; + } + ); + // The locked entry is skipped, the reachable one is still cleaned. + assert.equal(fs.existsSync(cleanableJson), false); + assert.equal(fs.existsSync(lockedJson), true); + assert.equal(requests.length, 0); + } finally { + fs.rmSync(`${lockedJson}.lock`, { recursive: true, force: true }); + } + + // The timed-out entry keeps an end marker. A later reuse applies it under + // the lock, so B becomes the sole owner and can perform final teardown. + const reused = await ensureBrokerSession(lockedCwd, { + env: { CODEX_COMPANION_SESSION_ID: "B" } + }); + assert.deepEqual(reused.sessionIds, ["B"]); + assert.equal(await teardownBrokersForSession("B", { killProcess: () => {} }), 1); + assert.equal(fs.existsSync(lockedJson), false); + assert.equal(requests.length, 1); + }); + }); +}); + +test("ensureBrokerSession cannot re-own a broker after its session end marker is recorded", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, requests, sessionDir }) => { + const cwd = makeTempDir(); + saveBrokerSession(cwd, { + endpoint, pidFile: null, logFile: null, sessionDir, pid: null, + sessionId: "S", sessionIds: ["S"] + }); + const stateFile = path.join(resolveStateDir(cwd), "broker.json"); + const lockDir = `${stateFile}.lock`; + fs.mkdirSync(lockDir); + + try { + await assert.rejects( + () => teardownBrokerForCwd(cwd, "S", { killProcess: () => {}, lockTimeoutMs: 50 }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE, reason: "lock-timeout" } + ); + } finally { + fs.rmSync(lockDir, { recursive: true, force: true }); + } + + await assert.rejects( + () => ensureBrokerSession(cwd, { env: { CODEX_COMPANION_SESSION_ID: "S" } }), + { code: BROKER_OWNER_ENDED_CODE } + ); + assert.equal(loadBrokerSession(cwd), null); + assert.equal(requests.length, 0); + }); + }); +}); + +test("reuseExistingBroker cannot reconnect after its session end marker is recorded", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, requests, sessionDir }) => { + const cwd = makeTempDir(); + saveBrokerSession(cwd, { + endpoint, pidFile: null, logFile: null, sessionDir, pid: null, + sessionId: "S", sessionIds: ["S"] + }); + const stateFile = path.join(resolveStateDir(cwd), "broker.json"); + const lockDir = `${stateFile}.lock`; + fs.mkdirSync(lockDir); + + try { + await assert.rejects( + () => teardownBrokerForCwd(cwd, "S", { killProcess: () => {}, lockTimeoutMs: 50 }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE, reason: "lock-timeout" } + ); + } finally { + fs.rmSync(lockDir, { recursive: true, force: true }); + } + + await assert.rejects( + () => CodexAppServerClient.connect(cwd, { + env: { CODEX_COMPANION_SESSION_ID: "S" }, + reuseExistingBroker: true + }), + { code: BROKER_OWNER_ENDED_CODE } + ); + assert.equal(loadBrokerSession(cwd), null); + assert.equal(requests.length, 0); + }); + }); +}); + +test("reuseExistingBroker rejects a marker-only ended session before direct fallback", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + + assert.equal( + await teardownBrokerForCwd(cwd, "S", { killProcess: () => {} }), + false + ); + assert.equal(loadBrokerSession(cwd), null); + + await assert.rejects( + () => CodexAppServerClient.connect(cwd, { + env: { CODEX_COMPANION_SESSION_ID: "S" }, + reuseExistingBroker: true + }), + { code: BROKER_OWNER_ENDED_CODE } + ); + assert.equal(loadBrokerSession(cwd), null); + }); +}); + +test("teardownBrokersForSession replaces an invalid non-directory lock path", async () => { + await withPluginData(async () => { + const stateRoot = stateRootForTest(); + const brokerJson = writeBrokerJson(stateRoot, "worktree-invalidlockaaa", { + endpoint: "unix:/tmp/codex-test-invalid-lock.sock", + sessionId: "S" + }); + const lockPath = `${brokerJson}.lock`; + fs.writeFileSync(lockPath, "not a directory", "utf8"); + + const count = await teardownBrokersForSession("S", { killProcess: () => {}, lockTimeoutMs: 60 }); + + assert.equal(count, 1); + assert.equal(fs.existsSync(brokerJson), false); + assert.equal(fs.existsSync(lockPath), false); + }); +}); + +test("teardownBrokersForSession never deletes a valid lock that replaces an invalid path", async () => { + await withPluginData(async () => { + const stateRoot = stateRootForTest(); + const brokerJson = writeBrokerJson(stateRoot, "worktree-lockswapaaaa", { + endpoint: "unix:/tmp/codex-test-lock-swap.sock", + sessionId: "S" + }); + const lockPath = `${brokerJson}.lock`; + fs.writeFileSync(lockPath, "invalid", "utf8"); + const old = new Date(Date.now() - 120000); + fs.utimesSync(lockPath, old, old); + const originalLstatSync = fs.lstatSync.bind(fs); + let swapped = false; + fs.lstatSync = (file, ...args) => { + const stat = originalLstatSync(file, ...args); + if (file === lockPath && !swapped) { + swapped = true; + fs.unlinkSync(lockPath); + fs.mkdirSync(lockPath); + } + return stat; + }; + + try { + await assert.rejects( + () => teardownBrokersForSession("S", { killProcess: () => {}, lockTimeoutMs: 60 }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE, reason: "lock-timeout" } + ); + } finally { + fs.lstatSync = originalLstatSync; + fs.rmSync(lockPath, { recursive: true, force: true }); + } + + assert.equal(fs.existsSync(brokerJson), true); + }); +}); + +test("teardownBrokersForSession never waits on a lock for another session's workspace", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, requests, sessionDir }) => { + const stateRoot = stateRootForTest(); + + // An unrelated workspace with a held lock, ordered before ours. + const otherJson = writeBrokerJson(stateRoot, "worktree-0000otheraaaa", { + endpoint: "unix:/tmp/codex-test-other.sock", + pidFile: null, logFile: null, sessionDir: null, pid: null, sessionId: "OTHER" + }); + fs.mkdirSync(`${otherJson}.lock`); + + const ourJson = writeBrokerJson(stateRoot, "worktree-9999oursbbbbb", { + endpoint, pidFile: null, logFile: null, sessionDir, pid: null, sessionId: "S" + }); + + try { + const startedAt = Date.now(); + // A large per-entry lock timeout would stall for seconds if we waited on + // the unrelated lock; the ownership pre-check must skip it outright. + const count = await teardownBrokersForSession("S", { killProcess: () => {}, lockTimeoutMs: 5000 }); + const elapsed = Date.now() - startedAt; + + assert.equal(count, 1); + assert.equal(fs.existsSync(ourJson), false); + assert.equal(fs.existsSync(otherJson), true); + assert.ok(elapsed < 1000, `expected fast teardown, took ${elapsed}ms`); + assert.equal(requests.length, 1); + } finally { + fs.rmSync(`${otherJson}.lock`, { recursive: true, force: true }); + } + }); + }); +}); + +test("teardownBrokersForSession reclaims a stale lock and still tears the broker down", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, requests, sessionDir }) => { + const stateRoot = stateRootForTest(); + const brokerJson = writeBrokerJson(stateRoot, "worktree-stalelock1234", { + endpoint, pidFile: null, logFile: null, sessionDir, pid: null, sessionId: "S" + }); + + // A lock left behind by a crashed holder: present but aged well past the + // stale threshold, so acquisition must reclaim it instead of timing out. + const lockDir = `${brokerJson}.lock`; + fs.mkdirSync(lockDir); + fs.writeFileSync(path.join(lockDir, "owner"), "2147483647-crashed", "utf8"); + const old = new Date(Date.now() - 120000); + fs.utimesSync(lockDir, old, old); + + const startedAt = Date.now(); + const count = await teardownBrokersForSession("S", { killProcess: () => {}, lockTimeoutMs: 2000 }); + const elapsed = Date.now() - startedAt; + + assert.equal(count, 1); + assert.equal(fs.existsSync(brokerJson), false); + assert.equal(requests.length, 1); + assert.ok(elapsed < 1000, `stale lock should be reclaimed promptly, took ${elapsed}ms`); + }); + }); +}); + +test("teardownBrokersForSession immediately reclaims a fresh dead-owner lock", async () => { + await withPluginData(async () => { + const stateRoot = stateRootForTest(); + const brokerJson = writeBrokerJson(stateRoot, "worktree-freshdeadlock", { + endpoint: "unix:/tmp/codex-test-fresh-dead-lock.sock", + pidFile: null, + logFile: null, + sessionDir: null, + pid: null, + sessionId: "S" + }); + const lockDir = `${brokerJson}.lock`; + fs.mkdirSync(lockDir); + fs.writeFileSync(path.join(lockDir, "owner"), "2147483647-crashed", "utf8"); + + const startedAt = Date.now(); + const count = await teardownBrokersForSession("S", { killProcess: () => {}, lockTimeoutMs: 200 }); + + assert.equal(count, 1); + assert.equal(fs.existsSync(brokerJson), false); + assert.equal(fs.existsSync(lockDir), false); + assert.ok(Date.now() - startedAt < 1000); + }); +}); + +test("teardownBrokersForSession reports an incomplete scan once its time budget is exhausted", async () => { + await withPluginData(async () => { + const stateRoot = stateRootForTest(); + // Two same-session entries, both with held locks. With a tiny budget the + // scan must return promptly rather than spending lockTimeoutMs on each. + const a = writeBrokerJson(stateRoot, "worktree-budgetaaaaaaaa", { + endpoint: "unix:/tmp/codex-test-b1.sock", + pidFile: null, logFile: null, sessionDir: null, pid: null, sessionId: "S" + }); + const b = writeBrokerJson(stateRoot, "worktree-budgetbbbbbbbb", { + endpoint: "unix:/tmp/codex-test-b2.sock", + pidFile: null, logFile: null, sessionDir: null, pid: null, sessionId: "S" + }); + fs.mkdirSync(`${a}.lock`); + fs.mkdirSync(`${b}.lock`); + + try { + const startedAt = Date.now(); + await assert.rejects( + () => teardownBrokersForSession("S", { killProcess: () => {}, lockTimeoutMs: 5000, budgetMs: 200 }), + (error) => { + assert.equal(error.code, BROKER_CLEANUP_INCOMPLETE_CODE); + assert.equal(error.reason, "deadline"); + return true; + } + ); + const elapsed = Date.now() - startedAt; + assert.ok(elapsed < 1500, `expected budget-bounded scan, took ${elapsed}ms`); + } finally { + fs.rmSync(`${a}.lock`, { recursive: true, force: true }); + fs.rmSync(`${b}.lock`, { recursive: true, force: true }); + } + }); +}); + +test("teardownBrokersForSession streams the state root instead of reading it all at once", async () => { + await withPluginData(async () => { + const stateRoot = stateRootForTest(); + const brokerJson = writeBrokerJson(stateRoot, "worktree-streamedscan", { + endpoint: "unix:/tmp/codex-test-streamed-scan.sock", + pidFile: null, + logFile: null, + sessionDir: null, + pid: null, + sessionId: "S" + }); + const originalReaddirSync = fs.readdirSync; + fs.readdirSync = (target, ...args) => { + if (target === stateRoot) { + throw new Error("state root must be streamed"); + } + return originalReaddirSync.call(fs, target, ...args); + }; + + try { + assert.equal(await teardownBrokersForSession("S", { killProcess: () => {} }), 1); + } finally { + fs.readdirSync = originalReaddirSync; + } + + assert.equal(fs.existsSync(brokerJson), false); + }); +}); + +test("teardownBrokersForSession does not silently cap a complete state-root scan", async () => { + await withPluginData(async () => { + const stateRoot = stateRootForTest(); + fs.mkdirSync(stateRoot, { recursive: true }); + for (let index = 0; index < 1001; index += 1) { + fs.mkdirSync(path.join(stateRoot, `empty-${String(index).padStart(4, "0")}`)); + } + + const probe = fs.opendirSync(stateRoot); + const dirPrototype = Object.getPrototypeOf(probe); + probe.closeSync(); + const originalReadSync = dirPrototype.readSync; + let reads = 0; + dirPrototype.readSync = function (...args) { + reads += 1; + return originalReadSync.call(this, ...args); + }; + + try { + assert.equal(await teardownBrokersForSession("S", { killProcess: () => {} }), 0); + } finally { + dirPrototype.readSync = originalReadSync; + } + assert.ok(reads > 1000, `expected an uncapped scan, observed ${reads} reads`); + }); +}); + +test("ensureBrokerSession spawns and persists a fresh broker when the recorded one is dead", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + saveBrokerSession(cwd, { + endpoint: "unix:/tmp/codex-test-dead-spawn.sock", + pidFile: null, + logFile: null, + sessionDir: null, + pid: null, + sessionId: "A", + sessionIds: ["A"] + }); + + const session = await ensureBrokerSession(cwd, { + env: { CODEX_COMPANION_SESSION_ID: "B" }, + scriptPath: writeFakeBrokerScript() + }); + try { + assert.ok(session); + assert.notEqual(session.endpoint, "unix:/tmp/codex-test-dead-spawn.sock"); + assert.deepEqual(session.sessionIds, ["B"]); + // The freshly spawned broker is the one persisted (no orphan record). + assert.equal(loadBrokerSession(cwd).endpoint, session.endpoint); + assert.equal(await new Promise((resolve) => { + const socket = net.createConnection({ path: parseBrokerEndpoint(session.endpoint).path }); + socket.on("connect", () => { socket.end(); resolve(true); }); + socket.on("error", () => resolve(false)); + }), true); + } finally { + if (session?.pid) { + try { + process.kill(session.pid); + } catch { + // Already gone. + } + } + } + }); +}); + +test("ensureBrokerSession never spawns without the state lock after acquisition times out", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + const spawnMarker = path.join(makeTempDir(), "spawned"); + const stateFile = path.join(resolveStateDir(cwd), "broker.json"); + fs.mkdirSync(path.dirname(stateFile), { recursive: true }); + fs.mkdirSync(`${stateFile}.lock`); + + try { + await assert.rejects( + () => ensureBrokerSession(cwd, { + scriptPath: writeFakeBrokerScript(), + lockTimeoutMs: 60, + env: { ...process.env, TEST_BROKER_SPAWN_MARKER: spawnMarker } + }), + { code: "EBROKERSTATELOCKTIMEOUT" } + ); + assert.equal(loadBrokerSession(cwd), null); + assert.equal(fs.existsSync(spawnMarker), false); + } finally { + fs.rmSync(`${stateFile}.lock`, { recursive: true, force: true }); + } + }); +}); + +test("ensureBrokerSession does not persist a broker whose owner ended during spawn", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + const spawnMarker = path.join(makeTempDir(), "spawned"); + const stateFile = path.join(resolveStateDir(cwd), "broker.json"); + const killed = []; + const pending = ensureBrokerSession(cwd, { + scriptPath: writeFakeBrokerScript(), + env: { + ...process.env, + CODEX_COMPANION_SESSION_ID: "S", + TEST_BROKER_SPAWN_MARKER: spawnMarker, + TEST_BROKER_LISTEN_DELAY_MS: "200" + }, + killProcess(pid) { + killed.push(pid); + process.kill(pid); + } + }); + + const deadline = Date.now() + 2000; + while (!fs.existsSync(spawnMarker) && Date.now() < deadline) { + await sleep(10); + } + assert.equal(fs.existsSync(spawnMarker), true); + assert.equal(fs.existsSync(stateFile), false); + assert.equal(fs.existsSync(`${stateFile}.lock`), true); + + await assert.rejects( + () => teardownBrokerForCwd(cwd, "S", { killProcess: () => {}, lockTimeoutMs: 50 }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE, reason: "lock-timeout" } + ); + await assert.rejects(pending, { code: BROKER_OWNER_ENDED_CODE }); + assert.equal(loadBrokerSession(cwd), null); + assert.equal(fs.existsSync(`${stateFile}.lock`), false); + assert.equal(killed.length, 1); + }); +}); + +test("mismatched-cwd SessionEnd marks an in-flight broker spawn before broker.json exists", async () => { + await withPluginData(async () => { + const brokerCwd = makeTempDir(); + const hookCwd = makeTempDir(); + const spawnMarker = path.join(makeTempDir(), "spawned"); + const stateFile = path.join(resolveStateDir(brokerCwd), "broker.json"); + const killed = []; + const pending = ensureBrokerSession(brokerCwd, { + scriptPath: writeFakeBrokerScript(), + env: { + ...process.env, + CODEX_COMPANION_SESSION_ID: "S", + TEST_BROKER_SPAWN_MARKER: spawnMarker, + TEST_BROKER_LISTEN_DELAY_MS: "1000" + }, + killProcess(pid) { + killed.push(pid); + process.kill(pid); + } + }); + + const deadline = Date.now() + 2000; + while (!fs.existsSync(spawnMarker) && Date.now() < deadline) { + await sleep(10); + } + assert.equal(fs.existsSync(spawnMarker), true); + assert.equal(fs.existsSync(stateFile), false); + assert.equal(fs.existsSync(`${stateFile}.lock`), true); + + await assert.rejects( + () => handleSessionEnd({ cwd: hookCwd, session_id: "S" }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE, reason: "lock-timeout" } + ); + await assert.rejects(pending, { code: BROKER_OWNER_ENDED_CODE }); + assert.equal(loadBrokerSession(brokerCwd), null); + assert.equal(fs.existsSync(`${stateFile}.lock`), false); + assert.equal(killed.length, 1); + }); +}); + +test("session-wide teardown marks a session joining a ready broker under its lock", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, requests, sessionDir }) => { + const cwd = makeTempDir(); + saveBrokerSession(cwd, { + endpoint, + pidFile: null, + logFile: null, + sessionDir, + pid: null, + sessionId: "A", + sessionIds: ["A"] + }); + const stateFile = path.join(resolveStateDir(cwd), "broker.json"); + const lockDir = `${stateFile}.lock`; + fs.mkdirSync(lockDir); + fs.writeFileSync(path.join(lockDir, "owner"), `${process.pid}-ready-reuse`, "utf8"); + fs.writeFileSync(path.join(lockDir, "session"), "S", "utf8"); + + try { + await assert.rejects( + () => teardownBrokersForSession("S", { + killProcess: () => {}, + lockTimeoutMs: 50 + }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE, reason: "lock-timeout" } + ); + } finally { + fs.rmSync(lockDir, { recursive: true, force: true }); + } + + await assert.rejects( + () => ensureBrokerSession(cwd, { + env: { CODEX_COMPANION_SESSION_ID: "S" } + }), + { code: BROKER_OWNER_ENDED_CODE } + ); + assert.deepEqual(loadBrokerSession(cwd).sessionIds, ["A"]); + assert.equal(requests.length, 0); + }); + }); +}); + +test("CodexAppServerClient does not fall back to a direct app-server after its broker owner ends", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + const spawnMarker = path.join(makeTempDir(), "broker-spawned"); + const directMarker = path.join(makeTempDir(), "direct-spawned"); + const fakeBin = makeTempDir(); + const fakeCodex = path.join(fakeBin, "codex"); + fs.writeFileSync( + fakeCodex, + `#!/bin/sh\nprintf spawned > "${directMarker}"\nexit 1\n`, + { encoding: "utf8", mode: 0o755 } + ); + + const env = { + ...process.env, + PATH: `${fakeBin}${path.delimiter}${process.env.PATH ?? ""}`, + CODEX_COMPANION_SESSION_ID: "S", + TEST_BROKER_SPAWN_MARKER: spawnMarker, + TEST_BROKER_LISTEN_DELAY_MS: "200" + }; + const pending = CodexAppServerClient.connect(cwd, { + env, + brokerOptions: { scriptPath: writeFakeBrokerScript() } + }); + + const deadline = Date.now() + 2000; + while (!fs.existsSync(spawnMarker) && Date.now() < deadline) { + await sleep(10); + } + assert.equal(fs.existsSync(spawnMarker), true); + await assert.rejects( + () => teardownBrokerForCwd(cwd, "S", { killProcess: () => {}, lockTimeoutMs: 50 }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE, reason: "lock-timeout" } + ); + + await assert.rejects(pending, { code: BROKER_OWNER_ENDED_CODE }); + assert.equal(fs.existsSync(directMarker), false); + assert.equal(loadBrokerSession(cwd), null); + }); +}); + +test("CodexAppServerClient rejects an ended session before using a supplied broker endpoint", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + await teardownBrokerForCwd(cwd, "S", { killProcess: () => {} }); + + await assert.rejects( + () => CodexAppServerClient.connect(cwd, { + brokerEndpoint: createBrokerEndpoint(path.join(makeTempDir(), "broker.sock")), + env: { CODEX_COMPANION_SESSION_ID: "S" } + }), + { code: BROKER_OWNER_ENDED_CODE } + ); + }); +}); + +test("CodexAppServerClient rejects a supplied endpoint when its session ends during initialize", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + const sessionDir = makeTempDir(); + const endpoint = createBrokerEndpoint(sessionDir); + const target = parseBrokerEndpoint(endpoint); + let releaseInitialize; + const initializeReceived = new Promise((resolve) => { + releaseInitialize = resolve; + }); + let socket = null; + const server = net.createServer((connected) => { + socket = connected; + connected.once("data", () => releaseInitialize()); + }); + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(target.path, () => { + server.off("error", reject); + resolve(); + }); + }); + + try { + const pending = CodexAppServerClient.connect(cwd, { + brokerEndpoint: endpoint, + env: { CODEX_COMPANION_SESSION_ID: "S" } + }); + await initializeReceived; + await teardownBrokerForCwd(cwd, "S", { killProcess: () => {} }); + socket.write(`${JSON.stringify({ id: 1, result: {} })}\n`); + + await assert.rejects(pending, { code: BROKER_OWNER_ENDED_CODE }); + } finally { + socket?.destroy(); + await new Promise((resolve) => server.close(resolve)); + } + }); +}); + +test("CodexAppServerClient rejects an ended session before direct fallback", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + const directMarker = path.join(makeTempDir(), "direct-spawned"); + const fakeBin = makeTempDir(); + const fakeCodex = path.join(fakeBin, "codex"); + fs.writeFileSync( + fakeCodex, + `#!/bin/sh\nprintf spawned > "${directMarker}"\nexit 1\n`, + { encoding: "utf8", mode: 0o755 } + ); + await teardownBrokerForCwd(cwd, "S", { killProcess: () => {} }); + + await assert.rejects( + () => CodexAppServerClient.connect(cwd, { + disableBroker: true, + env: { + ...process.env, + PATH: `${fakeBin}${path.delimiter}${process.env.PATH ?? ""}`, + CODEX_COMPANION_SESSION_ID: "S" + } + }), + { code: BROKER_OWNER_ENDED_CODE } + ); + assert.equal(fs.existsSync(directMarker), false); + }); +}); + +test("CodexAppServerClient does not fall back after an ended owner's broker spawn times out", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + const stateFile = path.join(resolveStateDir(cwd), "broker.json"); + const spawnMarker = path.join(makeTempDir(), "broker-spawned"); + const directMarker = path.join(makeTempDir(), "direct-spawned"); + const fakeBin = makeTempDir(); + const fakeCodex = path.join(fakeBin, "codex"); + fs.writeFileSync( + fakeCodex, + `#!/bin/sh\nprintf spawned > "${directMarker}"\nexit 1\n`, + { encoding: "utf8", mode: 0o755 } + ); + + const pending = CodexAppServerClient.connect(cwd, { + env: { + ...process.env, + PATH: `${fakeBin}${path.delimiter}${process.env.PATH ?? ""}`, + CODEX_COMPANION_SESSION_ID: "S", + TEST_BROKER_SPAWN_MARKER: spawnMarker, + TEST_BROKER_LISTEN_DELAY_MS: "1000" + }, + brokerOptions: { scriptPath: writeFakeBrokerScript(), timeoutMs: 2000 } + }); + const rejected = assert.rejects(pending, { code: BROKER_OWNER_ENDED_CODE }); + + const deadline = Date.now() + 2000; + while ((!fs.existsSync(spawnMarker) || !fs.existsSync(`${stateFile}.lock`)) && Date.now() < deadline) { + await sleep(10); + } + assert.equal(fs.existsSync(spawnMarker), true); + assert.equal(fs.existsSync(`${stateFile}.lock`), true); + await assert.rejects( + () => teardownBrokerForCwd(cwd, "S", { killProcess: () => {}, lockTimeoutMs: 50 }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE, reason: "lock-timeout" } + ); + + await rejected; + assert.equal(fs.existsSync(directMarker), false); + assert.equal(loadBrokerSession(cwd), null); + }); +}); + +test("ensureBrokerSession does not resurrect a broker torn down while waiting for the state lock", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, sessionDir }) => { + const cwd = makeTempDir(); + saveBrokerSession(cwd, { + endpoint, + pidFile: null, + logFile: null, + sessionDir, + pid: null, + sessionId: "A", + sessionIds: ["A"] + }); + const stateFile = path.join(resolveStateDir(cwd), "broker.json"); + const lockDir = `${stateFile}.lock`; + fs.mkdirSync(lockDir); + const observer = observeLockAttempt(lockDir); + + const pending = ensureBrokerSession(cwd, { + env: { CODEX_COMPANION_SESSION_ID: "B" }, + scriptPath: writeFakeBrokerScript() + }); + + // While B is parked on the lock (readiness probe already passed), A's + // SessionEnd shuts the broker down and removes broker.json. + await observer.attempted; + observer.restore(); + fs.unlinkSync(stateFile); + fs.rmSync(parseBrokerEndpoint(endpoint).path, { force: true }); + fs.rmdirSync(lockDir); + + const session = await pending; + try { + assert.ok(session); + assert.notEqual(session.endpoint, endpoint); + assert.deepEqual(session.sessionIds, ["B"]); + assert.equal(loadBrokerSession(cwd).endpoint, session.endpoint); + } finally { + if (session?.pid) { + try { + process.kill(session.pid); + } catch { + // Already gone. + } + } + } + }); + }); +}); + +test("ensureBrokerSession reuses a live replacement broker after losing the lock race", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint: staleEndpoint, sessionDir: staleDir }) => { + await withReadyBroker(async ({ endpoint: freshEndpoint, sessionDir: freshDir }) => { + const cwd = makeTempDir(); + saveBrokerSession(cwd, { + endpoint: staleEndpoint, + pidFile: null, + logFile: null, + sessionDir: staleDir, + pid: null, + sessionId: "A", + sessionIds: ["A"] + }); + const stateFile = path.join(resolveStateDir(cwd), "broker.json"); + const lockDir = `${stateFile}.lock`; + fs.mkdirSync(lockDir); + const observer = observeLockAttempt(lockDir); + + const pending = ensureBrokerSession(cwd, { env: { CODEX_COMPANION_SESSION_ID: "B" } }); + + // While B waits, the stale broker is replaced by a different live one. + await observer.attempted; + observer.restore(); + saveBrokerSession(cwd, { + endpoint: freshEndpoint, + pidFile: null, + logFile: null, + sessionDir: freshDir, + pid: null, + sessionId: "C", + sessionIds: ["C"] + }); + fs.rmdirSync(lockDir); + + const session = await pending; + assert.equal(session.endpoint, freshEndpoint); + assert.deepEqual(session.sessionIds, ["C", "B"]); + assert.deepEqual(loadBrokerSession(cwd).sessionIds, ["C", "B"]); + }); + }); + }); +}); + +test("teardownBrokersForSession tolerates an unparseable broker.json and keeps cleaning later brokers", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, requests, sessionDir }) => { + const stateRoot = stateRootForTest(); + + // A corrupt/half-written broker.json: the unlocked ownership pre-check + // must not treat a parse failure as "not ours" and abort or skip the + // rest of the scan. + const corruptDir = path.join(stateRoot, "worktree-badjson00000"); + fs.mkdirSync(corruptDir, { recursive: true }); + fs.writeFileSync(path.join(corruptDir, "broker.json"), "{ not valid json", "utf8"); + + const goodJson = writeBrokerJson(stateRoot, "worktree-goodjson11111", { + endpoint, pidFile: null, logFile: null, sessionDir, pid: null, sessionId: "S" + }); + + const count = await teardownBrokersForSession("S", { killProcess: () => {}, lockTimeoutMs: 200 }); + + assert.equal(count, 1); + assert.equal(fs.existsSync(goodJson), false); + assert.equal(requests.length, 1); + }); + }); +}); + +test("teardownBrokersForSession skips symlinked and oversized broker state", async () => { + await withPluginData(async () => { + const stateRoot = stateRootForTest(); + fs.mkdirSync(stateRoot, { recursive: true }); + const outside = makeTempDir(); + fs.writeFileSync(path.join(outside, "broker.json"), JSON.stringify({ sessionId: "S", pid: 12345 })); + fs.symlinkSync(outside, path.join(stateRoot, "symlinked-workspace")); + + const linkedFileDir = path.join(stateRoot, "linked-file-workspace"); + fs.mkdirSync(linkedFileDir); + fs.symlinkSync(path.join(outside, "broker.json"), path.join(linkedFileDir, "broker.json")); + + const oversizedDir = path.join(stateRoot, "oversized-workspace"); + fs.mkdirSync(oversizedDir); + fs.writeFileSync(path.join(oversizedDir, "broker.json"), "x".repeat(64 * 1024 + 1)); + + const killed = []; + const count = await teardownBrokersForSession("S", { killProcess: (pid) => killed.push(pid) }); + + assert.equal(count, 0); + assert.deepEqual(killed, []); + assert.equal(fs.existsSync(path.join(outside, "broker.json")), true); + }); +}); + +test("sendBrokerShutdown returns immediately for a non-positive timeout", async () => { + await withHangingBroker(async ({ endpoint, requests }) => { + await sendBrokerShutdown(endpoint, { timeoutMs: 0 }); + assert.equal(requests.length, 0); + }); +}); + +test("teardownBrokerForCwd records a global ended session before the first broker lock exists", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + const tornDown = await teardownBrokerForCwd(cwd, "S"); + + assert.equal(tornDown, false); + assert.equal(isBrokerSessionEnded("S"), true); + assert.equal(fs.existsSync(resolveStateDir(cwd)), false); + await assert.rejects( + () => ensureBrokerSession(cwd, { env: { CODEX_COMPANION_SESSION_ID: "S" } }), + { code: BROKER_OWNER_ENDED_CODE } + ); + assert.equal(loadBrokerSession(cwd), null); + }); +}); + +test("saveBrokerSession writes broker.json atomically", async () => { + await withPluginData(async () => { + const cwd = makeTempDir(); + const dir = resolveStateDir(cwd); + const stateFile = path.join(dir, "broker.json"); + const operations = []; + const originalWriteFileSync = fs.writeFileSync.bind(fs); + const originalRenameSync = fs.renameSync.bind(fs); + fs.writeFileSync = (file, data, ...args) => { + if (String(file).startsWith(`${stateFile}.tmp-`) || file === stateFile) { + operations.push(["write", String(file)]); + } + return originalWriteFileSync(file, data, ...args); + }; + fs.renameSync = (from, to) => { + if (to === stateFile) { + operations.push(["rename", String(from), String(to)]); + } + return originalRenameSync(from, to); + }; + try { + saveBrokerSession(cwd, { endpoint: "unix:/tmp/x.sock", sessionId: "S", sessionIds: ["S"] }); + } finally { + fs.writeFileSync = originalWriteFileSync; + fs.renameSync = originalRenameSync; + } + + assert.equal(operations[0][0], "write"); + assert.match(operations[0][1], /broker\.json\.tmp-/); + assert.deepEqual(operations[1], ["rename", operations[0][1], stateFile]); + const leftovers = fs.readdirSync(dir).filter((name) => name.includes("broker.json.tmp")); + assert.deepEqual(leftovers, []); + assert.equal(loadBrokerSession(cwd).sessionId, "S"); + }); +}); + +test("teardownBrokersForSession tolerates a corrupt endpoint and keeps cleaning later brokers", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, requests, sessionDir }) => { + const stateRoot = stateRootForTest(); + + // A stale record with an unsupported endpoint: sendBrokerShutdown would + // throw synchronously on it. It must be torn down best-effort, not abort + // the scan. + const corruptPid = 999999999; // non-existent; killProcess is mocked below + const corruptJson = writeBrokerJson(stateRoot, "worktree-corrupt00000", { + endpoint: "garbage-endpoint", pidFile: null, logFile: null, + sessionDir: null, pid: corruptPid, sessionId: "S" + }); + + // A healthy broker owned by the same session, later in the scan. + const goodJson = writeBrokerJson(stateRoot, "worktree-goodendpoint1", { + endpoint, pidFile: null, logFile: null, sessionDir, pid: null, sessionId: "S" + }); + + const killed = []; + const count = await teardownBrokersForSession("S", { killProcess: (pid) => killed.push(pid) }); + + assert.equal(count, 2); + assert.equal(fs.existsSync(corruptJson), false); + assert.equal(fs.existsSync(goodJson), false); + assert.ok(killed.includes(corruptPid)); + assert.equal(requests.length, 1); + }); + }); +}); + +test("teardownBrokersForSession continues after one broker cleanup fails", async () => { + await withPluginData(async () => { + const stateRoot = stateRootForTest(); + const firstJson = writeBrokerJson(stateRoot, "worktree-0000failureaaa", { + endpoint: "unix:/tmp/codex-test-cleanup-failure.sock", + sessionId: "S" + }); + const secondJson = writeBrokerJson(stateRoot, "worktree-9999successbbb", { + endpoint: "unix:/tmp/codex-test-cleanup-success.sock", + sessionId: "S" + }); + const originalUnlinkSync = fs.unlinkSync.bind(fs); + fs.unlinkSync = (file) => { + if (file === firstJson) { + throw Object.assign(new Error("simulated unlink failure"), { code: "EACCES" }); + } + return originalUnlinkSync(file); + }; + + try { + await assert.rejects( + () => teardownBrokersForSession("S", { killProcess: () => {} }), + { code: "EACCES" } + ); + } finally { + fs.unlinkSync = originalUnlinkSync; + } + + assert.equal(fs.existsSync(firstJson), true); + assert.equal(fs.existsSync(secondJson), false); + }); +}); + +test("teardownBrokersForSession does not mutate brokers when the global end marker cannot be published", async () => { + await withPluginData(async () => { + const stateRoot = stateRootForTest(); + const firstJson = writeBrokerJson(stateRoot, "worktree-0000first", { + endpoint: "unix:/tmp/codex-test-marker-failure.sock", + sessionId: "S" + }); + const secondJson = writeBrokerJson(stateRoot, "worktree-9999second", { + endpoint: "unix:/tmp/codex-test-marker-success.sock", + sessionId: "S" + }); + const originalWriteFileSync = fs.writeFileSync.bind(fs); + fs.writeFileSync = (file, ...args) => { + if (path.dirname(String(file)) === path.join(stateRoot, ".ended-sessions")) { + throw Object.assign(new Error("simulated marker write failure"), { code: "EACCES" }); + } + return originalWriteFileSync(file, ...args); + }; + + try { + await assert.rejects( + () => teardownBrokersForSession("S", { killProcess: () => {} }), + { code: "EACCES" } + ); + } finally { + fs.writeFileSync = originalWriteFileSync; + } + + assert.equal(fs.existsSync(firstJson), true); + assert.equal(fs.existsSync(secondJson), true); + }); +}); + +test("an unrelated broker reuse cannot consume another session's ended status", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, sessionDir }) => { + const cwd = makeTempDir(); + saveBrokerSession(cwd, { + endpoint, pidFile: null, logFile: null, sessionDir, pid: null, + sessionId: "A", sessionIds: ["A"] + }); + + assert.equal(await teardownBrokerForCwd(makeTempDir(), "S"), false); + assert.equal(isBrokerSessionEnded("S"), true); + + const reused = await ensureBrokerSession(cwd, { + env: { CODEX_COMPANION_SESSION_ID: "C" } + }); + assert.deepEqual(reused.sessionIds, ["A", "C"]); + assert.equal(isBrokerSessionEnded("S"), true); + + await assert.rejects( + () => ensureBrokerSession(cwd, { env: { CODEX_COMPANION_SESSION_ID: "S" } }), + { code: BROKER_OWNER_ENDED_CODE } + ); + assert.deepEqual(loadBrokerSession(cwd).sessionIds, ["A", "C"]); + }); + }); +}); + +test("handleSessionEnd skips the cwd fallback while broker state remains locked", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, requests, sessionDir }) => { + const cwd = makeTempDir(); + saveBrokerSession(cwd, { + endpoint, pidFile: null, logFile: null, sessionDir, pid: null, + sessionId: "S", sessionIds: ["S"] + }); + // Hold the broker's lock through both teardown attempts. The cwd fallback + // must not use its stale unlocked snapshot to tear the broker down. + const lockDir = `${path.join(resolveStateDir(cwd), "broker.json")}.lock`; + fs.mkdirSync(lockDir); + + try { + await assert.rejects( + () => handleSessionEnd({ cwd, session_id: "S" }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE, reason: "lock-timeout" } + ); + assert.notEqual(loadBrokerSession(cwd), null); + assert.equal(requests.length, 0); + } finally { + fs.rmSync(lockDir, { recursive: true, force: true }); + } + + const reused = await ensureBrokerSession(cwd, { + env: { CODEX_COMPANION_SESSION_ID: "B" } + }); + assert.deepEqual(reused.sessionIds, ["B"]); + await handleSessionEnd({ cwd, session_id: "B" }); + assert.equal(loadBrokerSession(cwd), null); + assert.equal(requests.length, 1); + }); + }); +}); + +test("teardownBrokerForCwd rechecks owners under the lock before fallback teardown", async () => { + await withPluginData(async () => { + await withReadyBroker(async ({ endpoint, requests, sessionDir }) => { + const cwd = makeTempDir(); + saveBrokerSession(cwd, { + endpoint, pidFile: null, logFile: null, sessionDir, pid: null, + sessionId: "A", sessionIds: ["A"] + }); + const stateFile = path.join(resolveStateDir(cwd), "broker.json"); + const lockDir = `${stateFile}.lock`; + fs.mkdirSync(lockDir); + const observer = observeLockAttempt(lockDir); + + const pending = teardownBrokerForCwd(cwd, "A", { + killProcess: () => {}, + lockTimeoutMs: 500 + }); + await observer.attempted; + observer.restore(); + saveBrokerSession(cwd, { + endpoint, pidFile: null, logFile: null, sessionDir, pid: null, + sessionId: "A", sessionIds: ["A", "B"] + }); + fs.rmdirSync(lockDir); + + const tornDown = await pending; + + assert.equal(tornDown, false); + assert.deepEqual(loadBrokerSession(cwd).sessionIds, ["B"]); + assert.equal(requests.length, 0); + }); + }); +}); + +test("teardownBrokersForSession caps shutdown waits to the remaining budget", async () => { + await withPluginData(async () => { + await withHangingBroker(async ({ endpoint, requests, sessionDir, destroySockets }) => { + const stateRoot = stateRootForTest(); + const first = writeBrokerJson(stateRoot, "worktree-hang1aaaaaaaaa", { + endpoint, pidFile: null, logFile: null, sessionDir, pid: null, + sessionId: "S", sessionIds: ["S"] + }); + const second = writeBrokerJson(stateRoot, "worktree-hang2bbbbbbbbb", { + endpoint, pidFile: null, logFile: null, sessionDir: null, pid: null, + sessionId: "S", sessionIds: ["S"] + }); + + try { + const startedAt = Date.now(); + // Endpoints accept but never reply; a per-RPC 1s wait on each would blow + // the budget. The scan must honor budgetMs across shutdown waits. + await assert.rejects( + () => teardownBrokersForSession("S", { + killProcess: () => {}, + shutdownTimeoutMs: 1000, + budgetMs: 300 + }), + { code: BROKER_CLEANUP_INCOMPLETE_CODE } + ); + const elapsed = Date.now() - startedAt; + assert.ok(elapsed < 900, `expected budget-capped shutdown, took ${elapsed}ms`); + assert.equal(Number(fs.existsSync(first)) + Number(fs.existsSync(second)), 1); + assert.equal(requests.length, 1); + } finally { + destroySockets(); + } + }); + }); +}); + +test("handleSessionEnd tears down broker even when cwd mismatches (regression #380)", async () => { + await withPluginData(async () => { + const stateRoot = stateRootForTest(); + const sessionDir = makeTempDir(); + const pidFile = path.join(sessionDir, "broker.pid"); + fs.writeFileSync(pidFile, "999999999\n"); // non-existent pid, harmless to signal + const brokerJson = writeBrokerJson(stateRoot, "worktree-33333333deadbeef", { + endpoint: "unix:/tmp/codex-test-nonexistent4.sock", + pidFile, logFile: null, sessionDir, pid: 999999999, sessionId: "S" + }); + + // cwd is a DIFFERENT path than the broker's workspace — the cwd-based path + // would miss; the session-based path must still tear it down. + await handleSessionEnd({ cwd: makeTempDir(), session_id: "S" }); + + assert.equal(fs.existsSync(brokerJson), false); + }); +}); diff --git a/tests/runtime.test.mjs b/tests/runtime.test.mjs index 8f276835..b3c4f6be 100644 --- a/tests/runtime.test.mjs +++ b/tests/runtime.test.mjs @@ -7,7 +7,13 @@ 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 { + loadBrokerSession, + saveBrokerSession, + teardownBrokerForCwd, + teardownBrokersForSession +} from "../plugins/codex/scripts/lib/broker-lifecycle.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)), ".."); @@ -15,6 +21,37 @@ const PLUGIN_ROOT = path.join(ROOT, "plugins", "codex"); const SCRIPT = path.join(PLUGIN_ROOT, "scripts", "codex-companion.mjs"); const STOP_HOOK = path.join(PLUGIN_ROOT, "scripts", "stop-review-gate-hook.mjs"); const SESSION_HOOK = path.join(PLUGIN_ROOT, "scripts", "session-lifecycle-hook.mjs"); +const TEST_PLUGIN_DATA = makeTempDir("codex-plugin-test-data-"); + +// Runtime tests exercise real SessionEnd and broker teardown paths. Keep their +// state out of the user's fallback /tmp/codex-companion directory so a test +// session can never discover or clean a live plugin session. +process.env.CLAUDE_PLUGIN_DATA = TEST_PLUGIN_DATA; +test.after(async () => { + const stateRoot = path.join(TEST_PLUGIN_DATA, "state"); + const sessionIds = new Set(); + if (fs.existsSync(stateRoot)) { + for (const entry of fs.readdirSync(stateRoot, { withFileTypes: true })) { + if (!entry.isDirectory()) { + continue; + } + try { + const session = JSON.parse(fs.readFileSync(path.join(stateRoot, entry.name, "broker.json"), "utf8")); + for (const sessionId of session.sessionIds ?? [session.sessionId]) { + if (sessionId) { + sessionIds.add(sessionId); + } + } + } catch { + // No broker state for this workspace. + } + } + } + for (const sessionId of sessionIds) { + await teardownBrokersForSession(sessionId, { killProcess: terminateProcessTree }); + } + fs.rmSync(TEST_PLUGIN_DATA, { recursive: true, force: true }); +}); async function waitFor(predicate, { timeoutMs = 5000, intervalMs = 50 } = {}) { const start = Date.now(); @@ -28,6 +65,21 @@ async function waitFor(predicate, { timeoutMs = 5000, intervalMs = 50 } = {}) { throw new Error("Timed out waiting for condition."); } +function buildSessionEnv(binDir, sessionId) { + return { + ...buildEnv(binDir), + CODEX_COMPANION_SESSION_ID: sessionId + }; +} + +function registerBrokerCleanup(t, cwd, sessionId = null) { + t.after(async () => { + await teardownBrokerForCwd(cwd, sessionId, { + killProcess: terminateProcessTree + }); + }); +} + test("setup reports ready when fake codex is installed and authenticated", () => { const binDir = makeTempDir(); installFakeCodex(binDir); @@ -698,6 +750,38 @@ test("session start hook exports the Claude session id, transcript path, and plu ); }); +test("session start hook runs when invoked through a symlinked plugin root", () => { + const repo = makeTempDir(); + const linkParent = makeTempDir(); + const pluginLink = path.join(linkParent, "codex-link"); + fs.symlinkSync(PLUGIN_ROOT, pluginLink, process.platform === "win32" ? "junction" : "dir"); + + const envFile = path.join(makeTempDir(), "claude-env.sh"); + fs.writeFileSync(envFile, "", "utf8"); + const pluginDataDir = makeTempDir(); + const symlinkedSessionHook = path.join(pluginLink, "scripts", "session-lifecycle-hook.mjs"); + + const result = run("node", [symlinkedSessionHook, "SessionStart"], { + cwd: repo, + env: { + ...process.env, + CLAUDE_ENV_FILE: envFile, + CLAUDE_PLUGIN_DATA: pluginDataDir + }, + input: JSON.stringify({ + hook_event_name: "SessionStart", + session_id: "sess-symlink", + cwd: repo + }) + }); + + assert.equal(result.status, 0, result.stderr); + assert.equal( + fs.readFileSync(envFile, "utf8"), + `export CODEX_COMPANION_SESSION_ID='sess-symlink'\nexport CLAUDE_PLUGIN_DATA='${pluginDataDir}'\n` + ); +}); + test("write task output focuses on the Codex result without generic follow-up hints", () => { const repo = makeTempDir(); const binDir = makeTempDir(); @@ -890,7 +974,7 @@ test("task can finish after subagent work even if the parent turn/completed even assert.equal(result.stdout, "Handled the requested task.\nTask prompt accepted.\n"); }); -test("task using the shared broker still completes when Codex spawns subagents", () => { +test("task using the shared broker still completes when Codex spawns subagents", (t) => { const repo = makeTempDir(); const binDir = makeTempDir(); installFakeCodex(binDir, "with-subagent"); @@ -900,16 +984,16 @@ test("task using the shared broker still completes when Codex spawns subagents", run("git", ["commit", "-m", "init"], { cwd: repo }); fs.writeFileSync(path.join(repo, "README.md"), "hello again\n"); - const env = buildEnv(binDir); + const sessionId = "sess-shared-subagents"; + const env = buildSessionEnv(binDir, sessionId); + registerBrokerCleanup(t, repo, sessionId); const review = run("node", [SCRIPT, "review"], { cwd: repo, env }); assert.equal(review.status, 0, review.stderr); - if (!loadBrokerSession(repo)) { - return; - } + assert.ok(loadBrokerSession(repo)); const result = run("node", [SCRIPT, "task", "challenge the current design"], { cwd: repo, @@ -1737,7 +1821,7 @@ test("cancel with a job id can still target an active job from another Claude se assert.equal(state.jobs[0].status, "cancelled"); }); -test("cancel sends turn interrupt to the shared app-server before killing a brokered task", async () => { +test("cancel sends turn interrupt to the shared app-server before killing a brokered task", async (t) => { const repo = makeTempDir(); const binDir = makeTempDir(); const fakeStatePath = path.join(binDir, "fake-codex-state.json"); @@ -1747,7 +1831,9 @@ test("cancel sends turn interrupt to the shared app-server before killing a brok run("git", ["add", "README.md"], { cwd: repo }); run("git", ["commit", "-m", "init"], { cwd: repo }); - const env = buildEnv(binDir); + const sessionId = "sess-cancel-shared-turn"; + const env = buildSessionEnv(binDir, sessionId); + registerBrokerCleanup(t, repo, sessionId); const launched = run("node", [SCRIPT, "task", "--background", "--json", "investigate the flaky worker timeout"], { cwd: repo, env @@ -1790,15 +1876,6 @@ test("cancel sends turn interrupt to the shared app-server before killing a brok turnId: runningJob.turnId }); - const cleanup = run("node", [SESSION_HOOK, "SessionEnd"], { - cwd: repo, - env, - input: JSON.stringify({ - hook_event_name: "SessionEnd", - cwd: repo - }) - }); - assert.equal(cleanup.status, 0, cleanup.stderr); }); test("session end fully cleans up jobs for the ending session", async (t) => { @@ -2116,7 +2193,7 @@ test("stop hook runs the actual task when auth status looks stale", () => { assert.match(payload.reason, /Missing empty-state guard/i); }); -test("commands lazily start and reuse one shared app-server after first use", async () => { +test("commands without a session owner use a direct app-server", (t) => { const repo = makeTempDir(); const binDir = makeTempDir(); const fakeStatePath = path.join(binDir, "fake-codex-state.json"); @@ -2129,6 +2206,34 @@ test("commands lazily start and reuse one shared app-server after first use", as fs.writeFileSync(path.join(repo, "README.md"), "hello again\n"); const env = buildEnv(binDir); + registerBrokerCleanup(t, repo); + + const review = run("node", [SCRIPT, "review"], { + cwd: repo, + env + }); + assert.equal(review.status, 0, review.stderr); + assert.equal(loadBrokerSession(repo), null); + + const fakeState = JSON.parse(fs.readFileSync(fakeStatePath, "utf8")); + assert.equal(fakeState.appServerStarts, 1); +}); + +test("commands lazily start and reuse one shared app-server after first use", async (t) => { + const repo = makeTempDir(); + const binDir = makeTempDir(); + const fakeStatePath = path.join(binDir, "fake-codex-state.json"); + + 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 }); + fs.writeFileSync(path.join(repo, "README.md"), "hello again\n"); + + const sessionId = "sess-lazy-shared-runtime"; + const env = buildSessionEnv(binDir, sessionId); + registerBrokerCleanup(t, repo, sessionId); const review = run("node", [SCRIPT, "review"], { cwd: repo, @@ -2137,9 +2242,7 @@ test("commands lazily start and reuse one shared app-server after first use", as assert.equal(review.status, 0, review.stderr); const brokerSession = loadBrokerSession(repo); - if (!brokerSession) { - return; - } + assert.ok(brokerSession); const adversarial = run("node", [SCRIPT, "adversarial-review"], { cwd: repo, @@ -2150,18 +2253,9 @@ test("commands lazily start and reuse one shared app-server after first use", as const fakeState = JSON.parse(fs.readFileSync(fakeStatePath, "utf8")); assert.equal(fakeState.appServerStarts, 1); - const cleanup = run("node", [SESSION_HOOK, "SessionEnd"], { - cwd: repo, - env, - input: JSON.stringify({ - hook_event_name: "SessionEnd", - cwd: repo - }) - }); - assert.equal(cleanup.status, 0, cleanup.stderr); }); -test("setup reuses an existing shared app-server without starting another one", () => { +test("setup reuses an existing shared app-server without starting another one", (t) => { const repo = makeTempDir(); const binDir = makeTempDir(); const fakeStatePath = path.join(binDir, "fake-codex-state.json"); @@ -2173,7 +2267,9 @@ test("setup reuses an existing shared app-server without starting another one", run("git", ["commit", "-m", "init"], { cwd: repo }); fs.writeFileSync(path.join(repo, "README.md"), "hello again\n"); - const env = buildEnv(binDir); + const sessionId = "sess-setup-shared-runtime"; + const env = buildSessionEnv(binDir, sessionId); + registerBrokerCleanup(t, repo, sessionId); const review = run("node", [SCRIPT, "review"], { cwd: repo, @@ -2182,9 +2278,7 @@ test("setup reuses an existing shared app-server without starting another one", assert.equal(review.status, 0, review.stderr); const brokerSession = loadBrokerSession(repo); - if (!brokerSession) { - return; - } + assert.ok(brokerSession); const setup = run("node", [SCRIPT, "setup", "--json"], { cwd: repo, @@ -2195,18 +2289,9 @@ test("setup reuses an existing shared app-server without starting another one", const fakeState = JSON.parse(fs.readFileSync(fakeStatePath, "utf8")); assert.equal(fakeState.appServerStarts, 1); - const cleanup = run("node", [SESSION_HOOK, "SessionEnd"], { - cwd: repo, - env, - input: JSON.stringify({ - hook_event_name: "SessionEnd", - cwd: repo - }) - }); - assert.equal(cleanup.status, 0, cleanup.stderr); }); -test("status reports shared session runtime when a lazy broker is active", () => { +test("status reports shared session runtime when a lazy broker is active", (t) => { const repo = makeTempDir(); const binDir = makeTempDir(); installFakeCodex(binDir); @@ -2216,19 +2301,21 @@ test("status reports shared session runtime when a lazy broker is active", () => run("git", ["commit", "-m", "init"], { cwd: repo }); fs.writeFileSync(path.join(repo, "README.md"), "hello again\n"); + const sessionId = "sess-status-shared-runtime"; + const env = buildSessionEnv(binDir, sessionId); + registerBrokerCleanup(t, repo, sessionId); + const review = run("node", [SCRIPT, "review"], { cwd: repo, - env: buildEnv(binDir) + env }); assert.equal(review.status, 0, review.stderr); - if (!loadBrokerSession(repo)) { - return; - } + assert.ok(loadBrokerSession(repo)); const result = run("node", [SCRIPT, "status"], { cwd: repo, - env: buildEnv(binDir) + env }); assert.equal(result.status, 0, result.stderr); diff --git a/tests/state.test.mjs b/tests/state.test.mjs index 0f8f57ce..45979e84 100644 --- a/tests/state.test.mjs +++ b/tests/state.test.mjs @@ -5,7 +5,15 @@ import test from "node:test"; import assert from "node:assert/strict"; import { makeTempDir } from "./helpers.mjs"; -import { resolveJobFile, resolveJobLogFile, resolveStateDir, resolveStateFile, saveState } from "../plugins/codex/scripts/lib/state.mjs"; +import { + loadState, + removeSessionJobs, + resolveJobFile, + resolveJobLogFile, + resolveStateDir, + resolveStateFile, + saveState +} from "../plugins/codex/scripts/lib/state.mjs"; test("resolveStateDir uses a temp-backed per-workspace directory", () => { const workspace = makeTempDir(); @@ -103,3 +111,124 @@ test("saveState prunes dropped job artifacts when indexed jobs exceed the cap", .sort() ); }); + +test("removeSessionJobs preserves a job added before acquiring the state lock", () => { + const workspace = makeTempDir(); + const stateFile = resolveStateFile(workspace); + const jobA = { id: "job-a", status: "running", sessionId: "A" }; + const jobB = { + id: "job-b", + status: "running", + sessionId: "B", + jobFile: resolveJobFile(workspace, "job-b"), + logFile: resolveJobLogFile(workspace, "job-b") + }; + saveState(workspace, { jobs: [jobA] }); + fs.writeFileSync(jobB.jobFile, JSON.stringify(jobB), "utf8"); + fs.writeFileSync(jobB.logFile, "running\n", "utf8"); + + const lockDir = `${stateFile}.lock`; + const originalRenameSync = fs.renameSync; + let injected = false; + fs.renameSync = (source, destination) => { + if (destination === lockDir && !injected) { + injected = true; + const state = JSON.parse(fs.readFileSync(stateFile, "utf8")); + fs.writeFileSync(stateFile, `${JSON.stringify({ ...state, jobs: [...state.jobs, jobB] }, null, 2)}\n`, "utf8"); + } + return originalRenameSync.call(fs, source, destination); + }; + + let removed; + try { + removed = removeSessionJobs(workspace, "A"); + } finally { + fs.renameSync = originalRenameSync; + } + + assert.deepEqual(removed.map((job) => job.id), ["job-a"]); + assert.deepEqual(loadState(workspace).jobs.map((job) => job.id), ["job-b"]); + assert.equal(fs.existsSync(jobB.jobFile), true); + assert.equal(fs.existsSync(jobB.logFile), true); +}); + +test("saveState does not release a state lock that has a replacement owner", () => { + const workspace = makeTempDir(); + const stateFile = resolveStateFile(workspace); + const lockDir = `${stateFile}.lock`; + const tokenFile = path.join(lockDir, "owner"); + const originalRenameSync = fs.renameSync; + let replaced = false; + fs.renameSync = (source, destination) => { + const result = originalRenameSync.call(fs, source, destination); + if (destination === stateFile && !replaced) { + replaced = true; + fs.writeFileSync(tokenFile, "replacement-owner", "utf8"); + } + return result; + }; + + try { + saveState(workspace, { jobs: [] }); + } finally { + fs.renameSync = originalRenameSync; + } + + assert.equal(replaced, true); + assert.equal(fs.readFileSync(tokenFile, "utf8"), "replacement-owner"); + fs.rmSync(lockDir, { recursive: true, force: true }); +}); + +test("saveState immediately reclaims a fresh lock whose owner process exited", () => { + const workspace = makeTempDir(); + const stateFile = resolveStateFile(workspace); + const lockDir = `${stateFile}.lock`; + fs.mkdirSync(lockDir, { recursive: true }); + fs.writeFileSync(path.join(lockDir, "owner"), "2147483647-crashed", "utf8"); + + const startedAt = Date.now(); + saveState(workspace, { jobs: [] }); + + assert.ok(Date.now() - startedAt < 1000); + assert.equal(fs.existsSync(lockDir), false); + assert.deepEqual(loadState(workspace).jobs, []); +}); + +test("saveState retries when a contended lock disappears before inspection", () => { + const workspace = makeTempDir(); + const stateFile = resolveStateFile(workspace); + const lockDir = `${stateFile}.lock`; + fs.mkdirSync(lockDir, { recursive: true }); + fs.writeFileSync(path.join(lockDir, "owner"), `${process.pid}-holder`, "utf8"); + const originalExistsSync = fs.existsSync; + let released = false; + fs.existsSync = (target) => { + if (target === lockDir && !released) { + released = true; + fs.rmSync(lockDir, { recursive: true, force: true }); + return true; + } + return originalExistsSync.call(fs, target); + }; + + try { + saveState(workspace, { jobs: [] }); + } finally { + fs.existsSync = originalExistsSync; + } + + assert.equal(released, true); + assert.deepEqual(loadState(workspace).jobs, []); +}); + +test("saveState recovers an orphaned stale-reclaim barrier", () => { + const workspace = makeTempDir(); + const stateFile = resolveStateFile(workspace); + fs.mkdirSync(path.dirname(stateFile), { recursive: true }); + const barrier = fs.mkdtempSync(`${stateFile}.lock.reclaim-2147483647-`); + + saveState(workspace, { jobs: [] }); + + assert.equal(fs.existsSync(barrier), false); + assert.deepEqual(loadState(workspace).jobs, []); +});