diff --git a/packages/coding-agent/.changes/kernel-snapshot-dispose-timeout.md b/packages/coding-agent/.changes/kernel-snapshot-dispose-timeout.md new file mode 100644 index 000000000..e06391971 --- /dev/null +++ b/packages/coding-agent/.changes/kernel-snapshot-dispose-timeout.md @@ -0,0 +1 @@ +- Fixed graceful Python kernel disposal so timed-out final snapshots are cancelled before teardown. diff --git a/packages/coding-agent/.changes/protocol-frame-repair.md b/packages/coding-agent/.changes/protocol-frame-repair.md new file mode 100644 index 000000000..c85c2545f --- /dev/null +++ b/packages/coding-agent/.changes/protocol-frame-repair.md @@ -0,0 +1 @@ +- Fixed invalid kernel protocol frames hanging requests by rejecting the affected request and replacing the kernel from its latest state snapshot. diff --git a/packages/coding-agent/.changes/single-kernel-shutdown.md b/packages/coding-agent/.changes/single-kernel-shutdown.md new file mode 100644 index 000000000..b0d2e6d7b --- /dev/null +++ b/packages/coding-agent/.changes/single-kernel-shutdown.md @@ -0,0 +1 @@ +- Fixed kernel teardown so session cleanup and signal handling share one bounded graceful shutdown path. diff --git a/packages/coding-agent/src/core/kernel/repl-manager.ts b/packages/coding-agent/src/core/kernel/repl-manager.ts index afd46757e..a78ae3282 100644 --- a/packages/coding-agent/src/core/kernel/repl-manager.ts +++ b/packages/coding-agent/src/core/kernel/repl-manager.ts @@ -17,7 +17,7 @@ import { type ExecuteOptions, type ExecuteResult, errorMessage, - HOST_REQUEST_DISPOSE_TIMEOUT_MS, + HOST_REQUEST_SHUTDOWN_TIMEOUT_MS, installSignalHandlersOnce, isRecord, KERNEL_ABORT_GRACE_MS, @@ -29,6 +29,7 @@ import { type KernelDiffDisplay, type KernelManagerOptions, type KernelSentAgentMessage, + type KernelShutdownOptions, type KernelStartOptions, liveKernels, MAX_ATTACHMENT_DATA_CHARS, @@ -37,7 +38,6 @@ import { parseDiffDisplay, parseSentAgentMessage, raceStartupWithAbort, - SNAPSHOT_DISPOSE_TIMEOUT_MS, SNAPSHOT_EXECUTION_TIMEOUT_MS, } from "./shared.js"; import { @@ -49,6 +49,7 @@ import { const REPL_PROTOCOL_VERSION = 3; const READY_TIMEOUT_MS = 30_000; +const REPAIR_STEP_TIMEOUT_MS = 30_000; // Runtime-minted host-request ids never repeat; the bound only guards a // misbehaving runtime from growing the dedup set forever. const MAX_HANDLED_HOST_REQUEST_IDS = 1024; @@ -86,6 +87,38 @@ interface ActiveExecution { reject: (error: Error) => void; } +// Complete event vocabulary of protocol version 2 (see prime-agent-runtime/src/rlm/repl.md). +// The version handshake is exact, so an unknown kind is corruption, not a newer runtime. +const PROTOCOL_EVENT_KINDS = new Set([ + "ready", + "stdout", + "stderr", + "result", + "display", + "host_request", + "error", + "done", +]); + +/** + * Reason a JSON object still isn't a valid protocol frame, or undefined. + * `done` and `host_request` route strictly by non-empty string id (the runtime + * mints uuid hex ids and echoes the host's uuids); silently dropping an id-less + * one would leave the awaiting request unsettled forever. + */ +function invalidProtocolFrameReason(event: Record): string | undefined { + if (typeof event.event !== "string" || !PROTOCOL_EVENT_KINDS.has(event.event)) { + return "unknown protocol event"; + } + if ( + (event.event === "done" || event.event === "host_request") && + (typeof event.id !== "string" || event.id === "") + ) { + return `${event.event} frame without id`; + } + return undefined; +} + function asStringArray(value: unknown): string[] { return Array.isArray(value) ? value.filter((v): v is string => typeof v === "string") : []; } @@ -103,7 +136,7 @@ function asReasonArray(value: unknown): { name: string; reason: string }[] { export class ReplKernelManager { private readonly options: Pick< KernelManagerOptions, - "python" | "cwd" | "env" | "sessionId" | "hostHandlers" | "pythonSkills" | "snapshot" + "python" | "cwd" | "env" | "sessionId" | "hostHandlers" | "pythonSkills" | "snapshot" | "bootstrapCode" >; private readonly handledHostRequestIds = new Set(); private child?: ChildProcess; @@ -129,10 +162,26 @@ export class ReplKernelManager { private startGeneration = 0; /** Generation whose graceful shutdown() owns the teardown, so the exit handler must not run it. */ private gracefulShutdownGeneration?: number; + private gracefulShutdownPromise?: Promise; /** Memoized so concurrent callers all await the same in-flight startup. */ private startPromise?: Promise; /** Pending debounced auto-snapshot, if one has been scheduled. */ private snapshotTimer?: ReturnType; + /** While the final dispose snapshot is flushing, new external executions are rejected. */ + private flushingSnapshotForDispose = false; + /** In-flight final snapshot flush; concurrent teardowns join it instead of re-flushing. */ + private snapshotFlushForDispose?: Promise; + /** Repairs a child whose dedicated protocol stream emitted an invalid frame. */ + private protocolRepairPromise?: Promise; + private protocolRepairOwner?: { superseded: boolean }; + /** Corruption seen while still "starting" (e.g. ready and garbage in one chunk) fails that start. */ + private startupProtocolError?: Error; + /** A repair discarded its kernel: the next fresh start must re-run the runtime bootstrap. */ + private pendingRebootstrap = false; + /** Restore the saved namespace on that fresh start too (false when the snapshot itself is the declared culprit). */ + private pendingRestore = false; + private rebootstrapPromise?: Promise; + private teardownInFlight = 0; constructor(options: KernelManagerOptions) { this.options = { @@ -143,6 +192,7 @@ export class ReplKernelManager { hostHandlers: options.hostHandlers, pythonSkills: options.pythonSkills, snapshot: options.snapshot, + bootstrapCode: options.bootstrapCode, }; } @@ -213,11 +263,16 @@ export class ReplKernelManager { this.child = child; if (child.pid !== undefined) recordOrphanProcessState(child.pid, true); this.readyDeferred = createDeferred(); + this.startupProtocolError = undefined; this.wireChild(child); try { const protocol = await this.waitForReady(child); if (this.startStale(generation)) throw new Error("Kernel start superseded"); + // Ready and a corrupt frame can share one stdout chunk: ready resolved the + // deferred synchronously before the corruption was parsed, so the rejection + // in failProtocolFrame was a no-op. Never mark such a child running. + if (this.startupProtocolError) throw this.startupProtocolError; if (protocol !== REPL_PROTOCOL_VERSION) { throw new Error( `Kernel runtime speaks protocol ${protocol}, expected ${REPL_PROTOCOL_VERSION}. ` + @@ -257,10 +312,19 @@ export class ReplKernelManager { try { event = JSON.parse(line); } catch { - this.appendKernelDiagnostic(`unparseable protocol line: ${line.slice(0, 200)}`); - continue; + this.failProtocolFrame(child, `unparseable protocol line: ${line.slice(0, 200)}`); + return; + } + if (!isRecord(event)) { + this.failProtocolFrame(child, `non-object protocol line: ${line.slice(0, 200)}`); + return; + } + const invalidReason = invalidProtocolFrameReason(event); + if (invalidReason) { + this.failProtocolFrame(child, `${invalidReason}: ${line.slice(0, 200)}`); + return; } - if (isRecord(event)) this.handleEvent(event); + this.handleEvent(event); } }); @@ -295,6 +359,226 @@ export class ReplKernelManager { }); } + private failProtocolFrame(child: ChildProcess, diagnostic: string): void { + if (this.child !== child) return; + this.appendKernelDiagnostic(diagnostic); + const error = new Error(`Kernel protocol error: ${diagnostic}`); + if (this.state === "starting") this.startupProtocolError = error; + this.readyDeferred?.reject(error); + this.rejectActiveExecution(error); + if (this.teardownInFlight > 0 || this.state !== "running") return; + + if (this.protocolRepairOwner) { + // A repair's own replacement child corrupted: discard it instead of respawn-looping. + this.appendKernelDiagnostic("replacement kernel corrupted during protocol repair; giving up"); + this.protocolRepairOwner.superseded = true; + // performRestore clears pendingRestore, so it still being set means the + // corruption struck at or before the restore phase: the snapshot stays + // the prime suspect (ambiguous attribution, loop-safe — retrying it + // would re-trigger the corruption). Corruption strictly after a + // successful restore never implicates the snapshot; keeping the flag + // costs at most one bounded restore per later attempt. + const snapshotSuspect = this.pendingRestore; + this.killChildToIdle(); + if (snapshotSuspect) this.pendingRestore = false; + return; + } + const owner = { superseded: false }; + this.protocolRepairOwner = owner; + const repair = this.repairProtocolChild(child, owner); + this.protocolRepairPromise = repair; + void repair.then( + () => { + if (this.protocolRepairPromise === repair) this.protocolRepairPromise = undefined; + if (this.protocolRepairOwner === owner) this.protocolRepairOwner = undefined; + }, + (repairError) => { + this.appendKernelDiagnostic(`protocol repair failed: ${errorMessage(repairError)}`); + if (this.protocolRepairPromise === repair) this.protocolRepairPromise = undefined; + if (this.protocolRepairOwner === owner) this.protocolRepairOwner = undefined; + }, + ); + } + + private async repairProtocolChild(child: ChildProcess, owner: { superseded: boolean }): Promise { + if (this.child !== child || this.state === "shutdown") return; + this.killChildToIdle(); + + const start = this.start(); + const generation = this.startGeneration; + try { + await start; + } catch (error) { + this.finishFailedProtocolRepair(owner, error); + return; + } + if (this.startStale(generation) || (this.state as string) !== "running") { + this.finishFailedProtocolRepair(owner); + return; + } + + const restored = await this.performRestore(true); + if (this.startStale(generation) || (this.state as string) !== "running") { + this.finishFailedProtocolRepair(owner); + return; + } + if (this.options.snapshot && restored === null) { + if (owner.superseded || this.protocolRepairOwner !== owner) return; + this.appendKernelDiagnostic("protocol repair restore failed; discarding replacement kernel"); + this.killChildToIdle(); + // The snapshot is the declared culprit; the lazy path must not retry it. + this.pendingRestore = false; + return; + } + + // Restore revives only the user namespace; live handles (rlm, bash, skills) + // come from the runtime bootstrap, so a repaired kernel must re-run it. + if (!this.options.bootstrapCode) return; + const bootstrapped = await this.bootstrapRepairedKernel(this.options.bootstrapCode); + if (this.startStale(generation) || (this.state as string) !== "running") { + this.finishFailedProtocolRepair(owner); + return; + } + if (!bootstrapped) { + if (owner.superseded || this.protocolRepairOwner !== owner) return; + this.appendKernelDiagnostic("protocol repair bootstrap failed; discarding replacement kernel"); + this.killChildToIdle(); + } + } + + /** Bounded bootstrap of a repaired kernel; false when it failed. Never throws. */ + private async bootstrapRepairedKernel(code: string): Promise { + try { + const r = await this.enqueueRequest( + { type: "execute", code }, + code, + { internal: true, protocolRepair: true }, + REPAIR_STEP_TIMEOUT_MS, + ); + if (r.status !== "ok") { + this.appendKernelDiagnostic( + `protocol repair bootstrap ${r.status === "aborted" ? "timed out" : "failed"}: ${r.error?.evalue ?? r.stderr}`, + ); + return false; + } + this.pendingRebootstrap = false; + return true; + } catch (error) { + this.appendKernelDiagnostic(`protocol repair bootstrap error: ${errorMessage(error)}`); + return false; + } + } + + /** + * A fresh kernel started after a discarded repair has none of the runtime + * bootstrap's live handles (rlm, bash, skills) and an empty namespace: + * reprovision (restore, then bootstrap) before any user request. A failed + * re-bootstrap discards the kernel again instead of serving user code on an + * unprovisioned namespace. + */ + private async ensureKernelRebootstrapped(signal?: AbortSignal): Promise { + const code = this.options.bootstrapCode; + const needsRestore = Boolean(this.options.snapshot) && this.pendingRestore; + const needsBootstrap = Boolean(code) && this.pendingRebootstrap; + // An in-flight repair owns its kernel's restore/bootstrap sequence, and + // a teardown's final snapshot must never trigger reprovisioning. + if ( + (!needsRestore && !needsBootstrap) || + this.protocolRepairPromise || + this.teardownInFlight > 0 || + this.state !== "running" + ) { + return; + } + let task = this.rebootstrapPromise; + if (!task) { + const started = this.reprovisionFreshKernel(code); + task = started; + this.rebootstrapPromise = started; + void started.finally(() => { + if (this.rebootstrapPromise === started) this.rebootstrapPromise = undefined; + }); + } + // An aborted request never executes, so it may skip the wait; race the + // signal like waitForProtocolRepair does instead of riding out the + // bootstrap bound after a mid-wait abort. + if (signal) { + if (signal.aborted) return; + let onAbort: () => void = () => {}; + const aborted = new Promise((resolve) => { + onAbort = resolve; + signal.addEventListener("abort", onAbort, { once: true }); + }); + try { + await Promise.race([task, aborted]); + } finally { + signal.removeEventListener("abort", onAbort); + } + if (signal.aborted) return; + } + const ok = await task; // bounded by REPAIR_STEP_TIMEOUT_MS + if (!ok) throw new Error("Kernel bootstrap failed after protocol repair"); + } + + /** Restore (one-shot, best-effort) then bootstrap the lazily started fresh kernel. */ + private async reprovisionFreshKernel(code: string | undefined): Promise { + if (this.options.snapshot && this.pendingRestore) { + await this.performRestore(true); // clears pendingRestore on success + // Corrupted during the restore: the spawned repair owns the kernel now. + if (this.protocolRepairPromise || this.state !== "running") return false; + // One attempt per discard: a clean restore failure falls back to an + // empty namespace (ordinary startup semantics), never a retry loop. + this.pendingRestore = false; + } + if (!code || !this.pendingRebootstrap) return true; + const ok = await this.bootstrapRepairedKernel(code); + if (!ok && this.state === "running") this.killChildToIdle(); + return ok; + } + + /** Kill the current child and settle at clean idle, so the next start spawns fresh. */ + private killChildToIdle(): void { + // The discarded kernel carried the runtime bootstrap and (possibly) the + // restored namespace; a lazily started replacement must reprovision both. + this.pendingRebootstrap = true; + this.pendingRestore = true; + this.state = "shutdown"; + liveKernels.delete(this); + this.cleanupResources("SIGKILL"); + this.state = "idle"; + } + + private finishFailedProtocolRepair(owner: { superseded: boolean }, error?: unknown): void { + if (error) this.appendKernelDiagnostic(`protocol repair start failed: ${errorMessage(error)}`); + if (owner.superseded || this.protocolRepairOwner !== owner) return; + if (this.state === "shutdown") this.state = "idle"; + } + + private supersedeProtocolRepair(): void { + if (this.protocolRepairOwner) this.protocolRepairOwner.superseded = true; + } + + /** Wait until no protocol repair is pending; resolves early when the signal aborts. */ + private async waitForProtocolRepair(signal?: AbortSignal): Promise { + while (this.protocolRepairPromise && !signal?.aborted) { + const repair = this.protocolRepairPromise; + if (!signal) { + await repair; + continue; + } + let onAbort: () => void = () => {}; + const aborted = new Promise((resolve) => { + onAbort = resolve; + signal.addEventListener("abort", onAbort, { once: true }); + }); + try { + await Promise.race([repair, aborted]); + } finally { + signal.removeEventListener("abort", onAbort); + } + } + } + private async waitForReady(child: ChildProcess): Promise { const ready = this.readyDeferred; if (!ready) throw new Error("Kernel ready state is missing"); @@ -433,6 +717,7 @@ export class ReplKernelManager { } async execute(code: string, opts: ExecuteOptions = {}): Promise { + await this.waitForProtocolRepair(opts.signal); const result = await this.enqueueExecute(code, opts); // Refresh the on-disk snapshot after real work so a later resume (or a // crash before graceful shutdown) revives the most recent namespace. @@ -465,6 +750,21 @@ export class ReplKernelManager { if ((this.state as string) === "shutdown") { throw new Error("Kernel has been shut down"); } + if (this.flushingSnapshotForDispose && !opts.internal) { + throw new Error("Kernel is shutting down"); + } + if (!opts.protocolRepair) await this.ensureKernelRebootstrapped(opts.signal); + // Aborted while waiting on the re-bootstrap: settle now instead of parking + // on the queue slot behind the still-running bootstrap. + if (opts.signal?.aborted) { + return { stdout: "", stderr: "", status: "aborted", durationMs: 0 }; + } + // Re-check: a final flush may have started while this request awaited the + // lazy re-bootstrap; admitting it now would splice it between the flush's + // captured queue and the final snapshot, unbounding the teardown. + if (this.flushingSnapshotForDispose && !opts.internal) { + throw new Error("Kernel is shutting down"); + } const prev = this.executionQueue; let resolveNext: () => void = () => {}; @@ -483,6 +783,13 @@ export class ReplKernelManager { if ((this.state as string) === "shutdown") { throw new Error("Kernel has been shut down"); } + // A repair started while this request was queued or busy-waiting: release + // the slot so the repair's own restore can run, then requeue behind it. + if (this.protocolRepairPromise && !opts.protocolRepair) { + resolveNext(); + await this.waitForProtocolRepair(opts.signal); + return this.enqueueRequest(requestFields, code, opts, executionTimeoutMs); + } if (executionTimeoutMs === undefined) { return await this.executeInner(requestFields, code, opts, started); } @@ -889,56 +1196,82 @@ export class ReplKernelManager { } if (result === "timeout") { this.appendKernelDiagnostic( - `timed out waiting ${timeoutMs}ms for ${tasks.length} host request task(s) during dispose`, + `timed out waiting ${timeoutMs}ms for ${tasks.length} host request task(s) during shutdown`, ); } } - /** Resolves true when this call performed the cleanup (false: a concurrent teardown won). */ - async shutdown(opts: { snapshot?: boolean } = {}): Promise { + /** Resolves true when this call performed the cleanup (false: a concurrent teardown won; a joiner's options are ignored - the first caller's policy wins). */ + async shutdown(opts: KernelShutdownOptions = {}): Promise { + const inFlightShutdown = this.gracefulShutdownPromise; + if (inFlightShutdown) { + await inFlightShutdown; + return false; + } + + this.teardownInFlight++; + this.supersedeProtocolRepair(); + const operation = this.performShutdown(opts); + this.gracefulShutdownPromise = operation; + try { + return await operation; + } finally { + this.teardownInFlight--; + if (this.gracefulShutdownPromise === operation) this.gracefulShutdownPromise = undefined; + } + } + + private async performShutdown(opts: KernelShutdownOptions): Promise { if (this.state === "shutdown") { liveKernels.delete(this); + if (this.gracefulShutdownGeneration === this.startGeneration) return false; this.cleanupResources(); return true; } // Captured before any await: teardowns and newer starts bump the counter. const generation = this.startGeneration; - // Best-effort final flush (bounded) before teardown — used by signal handlers - // so a SIGINT/SIGTERM exit doesn't lose work the debounced snapshot hasn't saved. if (opts.snapshot) { await this.flushSnapshotForDispose(); - if (this.startStale(generation)) return false; // superseded mid-flush: the newer owner already cleaned this kernel + if (this.startStale(generation)) return false; } + // Protocol shutdown first: the runtime closes MCP servers and kills live bash() process groups a bare hard-kill would leak. + const protocolShutdownAvailable = this.state === "running"; this.state = "shutdown"; liveKernels.delete(this); - // Claim the teardown: our own child's exit handler must not run cleanupResources - // (which bumps the generation and would misread this call as superseded). A - // concurrent kill()/dispose() still bumps the generation and supersedes us. this.gracefulShutdownGeneration = generation; let shutdownTimer: ReturnType | undefined; let doneWaiterId: string | undefined; let performedCleanup = false; - const shutdownDeadline = new Promise((_resolve, reject) => { - shutdownTimer = globalThis.setTimeout( - () => reject(new Error(`Kernel did not shut down within ${KERNEL_SHUTDOWN_TIMEOUT_MS}ms`)), - KERNEL_SHUTDOWN_TIMEOUT_MS, - ); - shutdownTimer.unref?.(); - }); try { - if (this.child?.stdin && !this.child.stdin.destroyed) { + if (opts.drainHostRequests) { + const inFlightHostRequests = [...this.inFlightHostRequests]; + if (inFlightHostRequests.length > 0) { + await this.waitForHostRequestsToSettle(inFlightHostRequests, HOST_REQUEST_SHUTDOWN_TIMEOUT_MS); + } + } + if ( + protocolShutdownAvailable && + !this.startStale(generation) && + this.child?.stdin && + !this.child.stdin.destroyed + ) { const requestId = uuid(); doneWaiterId = requestId; const doneReply = new Promise((resolve) => { this.pendingDoneWaiters.set(requestId, resolve); }); + const shutdownDeadline = new Promise((_resolve, reject) => { + shutdownTimer = globalThis.setTimeout( + () => reject(new Error(`Kernel did not shut down within ${KERNEL_SHUTDOWN_TIMEOUT_MS}ms`)), + KERNEL_SHUTDOWN_TIMEOUT_MS, + ); + shutdownTimer.unref?.(); + }); const send = this.writeLine({ type: "shutdown", id: requestId }); send.catch(() => undefined); - // A kernel that exits without delivering its shutdown done must not stall the deadline. const kernelExit = this.waitForKernelExit(); const gracefulReply = Promise.all([send, doneReply]); - // Abandoned by the race, a late send failure must not reject unhandled. gracefulReply.catch(() => undefined); await Promise.race([gracefulReply, kernelExit, shutdownDeadline]); await Promise.race([kernelExit, shutdownDeadline]); @@ -951,8 +1284,6 @@ export class ReplKernelManager { if (shutdownTimer) globalThis.clearTimeout(shutdownTimer); if (doneWaiterId) this.pendingDoneWaiters.delete(doneWaiterId); if (this.gracefulShutdownGeneration === generation) this.gracefulShutdownGeneration = undefined; - // A superseded shutdown must not tear down the newer start's kernel. Ownership is decided - // here, before cleanupResources bumps the generation and would misread this call as superseded. if (!this.startStale(generation)) { this.cleanupResources(); performedCleanup = true; @@ -963,6 +1294,12 @@ export class ReplKernelManager { } async restart(): Promise { + // A final dispose flush owns the queue tail. Taking a slot now and joining + // the in-flight shutdown would deadlock: the flush's snapshot waits on our + // slot while we wait on the flush's shutdown. + if (this.flushingSnapshotForDispose) { + throw new Error("Kernel is shutting down"); + } const prev = this.executionQueue; let resolveNext: () => void = () => {}; this.executionQueue = new Promise((r) => { @@ -982,6 +1319,7 @@ export class ReplKernelManager { } async kill(): Promise { + this.supersedeProtocolRepair(); this.state = "shutdown"; liveKernels.delete(this); this.cleanupResources("SIGKILL"); @@ -1045,14 +1383,27 @@ export class ReplKernelManager { * (rlm, skills) over anything restored. Never throws. */ async restoreState(): Promise { + return this.performRestore(false); + } + + /** Repair restores bypass the repair gate and are bounded so a stalled kernel cannot wedge it. */ + private async performRestore(protocolRepair: boolean): Promise { const cfg = this.options.snapshot; if (!cfg) return null; try { - const r = await this.enqueueRequest({ type: "restore", path: cfg.path }, "", { internal: true }); + const r = await this.enqueueRequest( + { type: "restore", path: cfg.path }, + "", + { internal: true, protocolRepair }, + protocolRepair ? REPAIR_STEP_TIMEOUT_MS : undefined, + ); if (r.status !== "ok" || !r.doneFields) { - this.appendKernelDiagnostic(`state restore failed: ${r.error?.evalue ?? r.stderr}`); + this.appendKernelDiagnostic( + `state restore ${r.status === "aborted" ? "timed out" : "failed"}: ${r.error?.evalue ?? r.stderr}`, + ); return null; } + this.pendingRestore = false; return { restored: asStringArray(r.doneFields.restored), failed: asReasonArray(r.doneFields.failed), @@ -1100,83 +1451,46 @@ export class ReplKernelManager { } } - /** Best-effort final snapshot before a graceful dispose, bounded by a timeout. */ - private async flushSnapshotForDispose(): Promise { - if (!this.options.snapshot || !this.isRunning) return; - let timeout: ReturnType | undefined; - const guard = new Promise((resolve) => { - timeout = globalThis.setTimeout(resolve, SNAPSHOT_DISPOSE_TIMEOUT_MS); - if (timeout && typeof timeout === "object" && "unref" in timeout) timeout.unref(); + private flushSnapshotForDispose(): Promise { + // Concurrent teardowns (dispose vs a signal-handler shutdown) join one flush: + // a second flusher would clear the execution guard while the first is still + // snapshotting and enqueue a duplicate final snapshot behind it. + this.snapshotFlushForDispose ??= this.runSnapshotFlushForDispose().finally(() => { + this.snapshotFlushForDispose = undefined; }); - try { - await Promise.race([this.snapshotState().then(() => undefined), guard]); - } finally { - if (timeout) clearTimeout(timeout); - } + return this.snapshotFlushForDispose; } - /** Graceful cleanup. Waits briefly for in-flight host request handlers before killing the child. */ - dispose(): Promise { - return (async () => { - // Captured before any await: teardowns and newer starts bump the counter. - const generation = this.startGeneration; - // Final namespace flush while the kernel is still live (session end / reload). - await this.flushSnapshotForDispose(); - if (this.startStale(generation)) return; // superseded mid-flush: the newer owner already cleaned this kernel - this.state = "shutdown"; - liveKernels.delete(this); - // Claim the teardown so the child's exit handler does not run - // cleanupResources mid-dispose (same contract as shutdown()). - this.gracefulShutdownGeneration = generation; - const inFlightHostRequests = [...this.inFlightHostRequests]; - try { - if (inFlightHostRequests.length > 0) { - await this.waitForHostRequestsToSettle(inFlightHostRequests, HOST_REQUEST_DISPOSE_TIMEOUT_MS); - } - if (!this.startStale(generation)) { - // Bounded protocol shutdown first: the runtime's shutdown branch - // closes MCP servers and kills live bash() process groups, which - // a bare hard-kill would leak until the orphan reaper runs. - await this.requestProtocolShutdown(KERNEL_SHUTDOWN_TIMEOUT_MS); - } - } finally { - if (this.gracefulShutdownGeneration === generation) this.gracefulShutdownGeneration = undefined; - if (!this.startStale(generation)) this.cleanupResources(); // else: superseded, the newer owner already cleaned - } - })(); - } - - /** Best-effort bounded protocol shutdown; the caller's hard kill remains the backstop. */ - private async requestProtocolShutdown(timeoutMs: number): Promise { - const stdin = this.child?.stdin; - if (!stdin || stdin.destroyed) return; - const requestId = uuid(); - let timer: ReturnType | undefined; + private async runSnapshotFlushForDispose(): Promise { + if (!this.options.snapshot || !this.isRunning) return; + // A kernel that never restored the saved namespace must not overwrite it: + // the on-disk snapshot is strictly fresher than this namespace. + if (this.pendingRestore) return; + // Block new external executions so none can splice ahead of the final snapshot and stall dispose. + this.flushingSnapshotForDispose = true; try { - const doneReply = new Promise((resolve) => { - this.pendingDoneWaiters.set(requestId, resolve); - }); - const send = this.writeLine({ type: "shutdown", id: requestId }); - send.catch(() => undefined); - const kernelExit = this.waitForKernelExit(); - const deadline = new Promise((resolve) => { - timer = globalThis.setTimeout(resolve, timeoutMs); - timer.unref?.(); - }); - const gracefulReply = Promise.all([send, doneReply]).then(() => undefined); - gracefulReply.catch(() => undefined); - await Promise.race([gracefulReply, kernelExit, deadline]); - await Promise.race([kernelExit, deadline]); - } catch { - // Best-effort: cleanupResources still hard-kills the child. + const pendingExecutions = this.executionQueue; + if (this.activeExecution) void this.interrupt().catch(() => undefined); + let timeout: ReturnType | undefined; + const queueSettled = await Promise.race([ + pendingExecutions.then(() => true), + new Promise((resolve) => { + timeout = globalThis.setTimeout(() => resolve(false), SNAPSHOT_EXECUTION_TIMEOUT_MS); + timeout.unref?.(); + }), + ]); + if (timeout) globalThis.clearTimeout(timeout); + if (!queueSettled) return; + await this.captureSnapshot({ executionTimeoutMs: SNAPSHOT_EXECUTION_TIMEOUT_MS }); } finally { - if (timer) globalThis.clearTimeout(timer); - this.pendingDoneWaiters.delete(requestId); + // Reset: a superseding start() can revive this kernel for new work. + this.flushingSnapshotForDispose = false; } } /** Synchronous best-effort cleanup. Safe to call from `process.on('exit')`. */ disposeSync(): void { + this.supersedeProtocolRepair(); this.state = "shutdown"; liveKernels.delete(this); this.cleanupResources(); diff --git a/packages/coding-agent/src/core/kernel/shared.ts b/packages/coding-agent/src/core/kernel/shared.ts index 287f8a28c..25d43b917 100644 --- a/packages/coding-agent/src/core/kernel/shared.ts +++ b/packages/coding-agent/src/core/kernel/shared.ts @@ -3,12 +3,9 @@ import type { KernelBootstrapProgressHandler, KernelPythonSkill } from "./bootst import type { RestoreResult, SnapshotResult } from "./state-snapshot.js"; export const DEFAULT_MAX_OUTPUT_CHARS = 65536; -export const HOST_REQUEST_DISPOSE_TIMEOUT_MS = 5000; +export const HOST_REQUEST_SHUTDOWN_TIMEOUT_MS = 5000; export const KERNEL_SHUTDOWN_TIMEOUT_MS = 5000; export const DEFAULT_SNAPSHOT_DEBOUNCE_MS = 1500; -// Cap how long a graceful dispose waits on the final snapshot; the debounced -// on-disk copy is the fallback if this is exceeded. -export const SNAPSHOT_DISPOSE_TIMEOUT_MS = 5000; export const SNAPSHOT_EXECUTION_TIMEOUT_MS = 5000; export const KERNEL_ABORT_GRACE_MS = 1000; export const KERNEL_BUSY_REUSE_WAIT_MS = 5000; @@ -57,6 +54,8 @@ export interface KernelManagerOptions { pythonSkills?: readonly KernelPythonSkill[]; /** Persist/revive the user namespace across kernel restarts and session resume. */ snapshot?: KernelSnapshotConfig; + /** Runtime bootstrap re-run on a protocol-repaired kernel so live handles (rlm, bash, skills) exist again. */ + bootstrapCode?: string; } export interface KernelStartOptions { @@ -73,6 +72,8 @@ export interface ExecuteOptions { maxOutputChars?: number; /** Synthetic host cell (snapshot/restore/list); excluded from lastCellCode attribution. */ internal?: boolean; + /** The protocol repair's own restore; exempt from waiting on the repair it belongs to. */ + protocolRepair?: boolean; } /** MIME tag the `edit` skill emits diff payloads under. */ @@ -268,16 +269,20 @@ export function createDeferred(): Deferred { return { promise, resolve, reject }; } +export interface KernelShutdownOptions { + snapshot?: boolean; + drainHostRequests?: boolean; +} + /** Public surface every kernel client exposes to the provisioner and session layer. */ export interface KernelClient { readonly ownerSessionId: string | undefined; readonly isRunning: boolean; start(options?: KernelStartOptions): Promise; execute(code: string, opts?: ExecuteOptions): Promise; - shutdown(opts?: { snapshot?: boolean }): Promise; + shutdown(opts?: KernelShutdownOptions): Promise; restart(): Promise; kill(): Promise; - dispose(): Promise; disposeSync(): void; snapshotState(): Promise; pruneOversizedVariables(): Promise; @@ -293,7 +298,7 @@ let signalHandlersInstalled = false; registerSessionResourceCleanup((sessionId) => { for (const k of liveKernels) { if (!sessionId || k.ownerSessionId === sessionId) { - void k.dispose(); + void k.shutdown({ snapshot: true, drainHostRequests: true }); } } }); diff --git a/packages/coding-agent/src/core/tools/ipython.ts b/packages/coding-agent/src/core/tools/ipython.ts index bedf72d0e..485c236cc 100644 --- a/packages/coding-agent/src/core/tools/ipython.ts +++ b/packages/coding-agent/src/core/tools/ipython.ts @@ -362,7 +362,7 @@ export class IpythonKernelProvisioner { if (!pending) return; try { const m = await pending; - await m.dispose(); + await m.shutdown({ snapshot: true, drainHostRequests: true }); } catch { // a failed startup already cleaned up after itself } @@ -454,6 +454,7 @@ export class IpythonKernelProvisioner { // without bash, where the runtime's teaching error fires instead). const shellPath = resolveKernelBashShell(this.options?.shellPath); const commandPrefix = this.options?.commandPrefix; + const bootstrapCode = buildRlmBootstrapCode(this.options?.pythonSkills); const m = new ReplKernelManager({ python: this.options?.python, cwd: this.cwd, @@ -470,6 +471,7 @@ export class IpythonKernelProvisioner { snapshot: snapshotDir ? { path: snapshotPathIn(snapshotDir), manifestPath: manifestPathIn(snapshotDir) } : undefined, + bootstrapCode, }); let pendingRestore: RestoreResult | undefined; try { @@ -500,7 +502,7 @@ export class IpythonKernelProvisioner { } } this.emitStartupProgress("Preparing Python runtime..."); - const bootstrap = await m.execute(buildRlmBootstrapCode(this.options?.pythonSkills), { + const bootstrap = await m.execute(bootstrapCode, { signal: startupSignal, }); if (bootstrap.status !== "ok") { @@ -508,8 +510,11 @@ export class IpythonKernelProvisioner { throw new Error(`Failed to initialize rlm runtime in the Python kernel:\n${details}`); } } catch (error) { - // Never leak the kernel process if startup fails after spawn. - void m.dispose(); + // Never leak the kernel process if startup fails after spawn — and never + // surface the failure before the teardown (final snapshot flush included) + // finished, or a replacement provisioner gated on this dispose could + // race the still-flushing kernel over the same snapshot files. + await m.shutdown({ snapshot: true, drainHostRequests: true }).catch(() => undefined); throw error; } // Only tell the model what was revived once the kernel is actually usable — diff --git a/packages/coding-agent/test/agent-session-recursion.test.ts b/packages/coding-agent/test/agent-session-recursion.test.ts index 2c373352b..793801106 100644 --- a/packages/coding-agent/test/agent-session-recursion.test.ts +++ b/packages/coding-agent/test/agent-session-recursion.test.ts @@ -3906,7 +3906,7 @@ print(_result.name) expect(finished.status).toBe("ok"); expect(finished.stdout.trim()).toBe("detached-worker"); } finally { - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }); }); diff --git a/packages/coding-agent/test/ipython-bootstrap.test.ts b/packages/coding-agent/test/ipython-bootstrap.test.ts index 1b8dadb2e..27957ac42 100644 --- a/packages/coding-agent/test/ipython-bootstrap.test.ts +++ b/packages/coding-agent/test/ipython-bootstrap.test.ts @@ -93,7 +93,7 @@ describeIfKernel("RLM bootstrap (real kernel)", () => { expect(envResult.status).toBe("ok"); expect(envResult.stdout.trim()).toBe("1"); } finally { - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }, 60_000); @@ -135,7 +135,7 @@ describeIfKernel("RLM bootstrap (real kernel)", () => { expect(second.diffs?.[0]?.path).toBe(realpathSync(join(secondDir, "same.txt"))); expect(first.diffs?.[0]?.path).not.toBe(second.diffs?.[0]?.path); } finally { - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }, 60_000); }); diff --git a/packages/coding-agent/test/ipython-provisioner.test.ts b/packages/coding-agent/test/ipython-provisioner.test.ts index 9935cdce9..dd3a56afb 100644 --- a/packages/coding-agent/test/ipython-provisioner.test.ts +++ b/packages/coding-agent/test/ipython-provisioner.test.ts @@ -65,6 +65,45 @@ function createBusyKernelContext( return { ctx, setWorkingMessage }; } +function writeFakeReplRuntime(markerPath: string): string { + const python = join(tempDir, "python-repl"); + writeFileSync( + python, + `#!/usr/bin/env node +const fs = require("node:fs"); +const readline = require("node:readline"); +const emit = (event) => process.stdout.write(JSON.stringify(event) + "\\n"); +emit({ event: "ready", protocol: 2, python: process.version }); +const input = readline.createInterface({ input: process.stdin }); +input.on("line", (line) => { + const request = JSON.parse(line); + if (request.type === "restore") { + emit({ event: "done", id: request.id, status: "ok", restored: [], failed: [] }); + return; + } + if (request.type === "snapshot") { + setTimeout(() => { + fs.writeFileSync(${JSON.stringify(markerPath)}, "1"); + emit({ event: "done", id: request.id, status: "ok", saved: [], skipped: [], bytes: 0 }); + }, 150); + return; + } + if (request.type === "execute") { + emit({ event: "error", id: request.id, ename: "RuntimeError", evalue: "bootstrap refused", traceback: [] }); + emit({ event: "done", id: request.id, status: "error" }); + return; + } + if (request.type === "shutdown") { + emit({ event: "done", id: request.id, status: "ok" }); + process.exit(0); + } +}); +`, + ); + chmodSync(python, 0o755); + return python; +} + describe("IpythonKernelProvisioner", () => { beforeEach(() => { tempDir = mkdtempSync(join(tmpdir(), "prime-agent-provisioner-")); @@ -77,6 +116,23 @@ describe("IpythonKernelProvisioner", () => { } }); + it("does not surface a failed startup before the kernel's final snapshot flush finished", async () => { + const marker = join(tempDir, "snapshot-flushed"); + const snapshotDir = join(tempDir, "snapshots"); + mkdirSync(snapshotDir, { recursive: true }); + const python = writeFakeReplRuntime(marker); + const provisioner = new IpythonKernelProvisioner(tempDir, { python, snapshotDir }); + try { + await expect(provisioner.ensure()).rejects.toThrow(/Failed to initialize rlm runtime/); + // The failed kernel's teardown (final snapshot flush included) completed + // before the failure surfaced, so a replacement provisioner gated on this + // one cannot race the still-flushing kernel over the same snapshot files. + expect(existsSync(marker)).toBe(true); + } finally { + await provisioner.dispose(); + } + }); + it("memoizes concurrent ensure() calls into one startup", async () => { const { python, countRuns } = writeFakePython(); const provisioner = new IpythonKernelProvisioner(tempDir, { python }); @@ -352,7 +408,7 @@ describe("ReplKernelManager session cleanup during startup", () => { await expect(startup).rejects.toThrow(/Kernel exited before ready|disposed during startup/); expect(manager.isRunning).toBe(false); } finally { - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }); }); diff --git a/packages/coding-agent/test/kernel-agent-message-skill.test.ts b/packages/coding-agent/test/kernel-agent-message-skill.test.ts index f7cd4ac3b..f56b7218c 100644 --- a/packages/coding-agent/test/kernel-agent-message-skill.test.ts +++ b/packages/coding-agent/test/kernel-agent-message-skill.test.ts @@ -277,6 +277,6 @@ background_send = asyncio.create_task(send_later())`, expect(host.lateSentAgentMessageHandlers.size).toBe(256); expect(host.lateSentAgentMessageHandlers.has("request-0")).toBe(false); expect(host.lateSentAgentMessageHandlers.has("request-299")).toBe(true); - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); }); }); diff --git a/packages/coding-agent/test/repl-kernel-abort.test.ts b/packages/coding-agent/test/repl-kernel-abort.test.ts index 6187e4fa9..b3283a419 100644 --- a/packages/coding-agent/test/repl-kernel-abort.test.ts +++ b/packages/coding-agent/test/repl-kernel-abort.test.ts @@ -216,7 +216,7 @@ describe("ReplKernelManager abort handling", () => { manager.disposeSync(); }); - it("starts the snapshot timeout after earlier kernel work finishes", async () => { + it("cancels a hung final snapshot execution before teardown", async () => { vi.useFakeTimers(); const manager = new ReplKernelManager({ cwd: process.cwd(), @@ -240,33 +240,273 @@ describe("ReplKernelManager abort handling", () => { ); }), ); + const cleanupResources = vi.fn(); Object.assign( manager as unknown as { state: "running"; executionQueue: Promise; executeInner: typeof executeInner; start: () => Promise; + cleanupResources: () => void; }, - { state: "running", executionQueue: previousExecution, executeInner, start: async () => {} }, + { state: "running", executionQueue: previousExecution, executeInner, start: async () => {}, cleanupResources }, ); - const snapshot = ( - manager as unknown as { - captureSnapshot: (options?: { executionTimeoutMs?: number }) => Promise; - } - ).captureSnapshot({ executionTimeoutMs: 5000 }); - await vi.advanceTimersByTimeAsync(5000); + const disposal = manager.shutdown({ snapshot: true, drainHostRequests: true }); expect(executeInner).not.toHaveBeenCalled(); - releaseQueue(); await waitForCalls(executeInner, 1); const signal = executeInner.mock.calls[0]?.[2].signal; expect(signal?.aborted).toBe(false); await vi.advanceTimersByTimeAsync(4999); expect(signal?.aborted).toBe(false); + expect(cleanupResources).not.toHaveBeenCalled(); await vi.advanceTimersByTimeAsync(1); + expect(signal?.aborted).toBe(true); - await expect(snapshot).resolves.toBeNull(); + await expect(disposal).resolves.toBe(true); + expect(cleanupResources).toHaveBeenCalledOnce(); + }); + + it("tears down when the final snapshot is blocked behind a hung execution", async () => { + vi.useFakeTimers(); + const manager = new ReplKernelManager({ + cwd: process.cwd(), + snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" }, + }); + const executeInner = vi.fn(); + const interrupt = vi.fn(async () => {}); + const cleanupResources = vi.fn(); + Object.assign( + manager as unknown as { + state: "running"; + executionQueue: Promise; + activeExecution: object; + executeInner: typeof executeInner; + interrupt: typeof interrupt; + cleanupResources: () => void; + }, + { + state: "running", + executionQueue: new Promise(() => {}), + activeExecution: {}, + executeInner, + interrupt, + cleanupResources, + }, + ); + + const disposal = manager.shutdown({ snapshot: true, drainHostRequests: true }); + expect(interrupt).toHaveBeenCalledOnce(); + await vi.advanceTimersByTimeAsync(4999); + expect(cleanupResources).not.toHaveBeenCalled(); + await vi.advanceTimersByTimeAsync(1); + + await expect(disposal).resolves.toBe(true); + expect(executeInner).not.toHaveBeenCalled(); + expect(cleanupResources).toHaveBeenCalledOnce(); + }); + + it("rejects executions enqueued during the final snapshot flush so dispose stays bounded", async () => { + vi.useFakeTimers(); + const manager = new ReplKernelManager({ + cwd: process.cwd(), + snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" }, + }); + let releaseQueue: () => void = () => {}; + const previousExecution = new Promise((resolve) => { + releaseQueue = resolve; + }); + const executeInner = vi.fn( + async ( + _requestFields: Record, + _code: string, + opts: { signal?: AbortSignal }, + ): Promise<{ stdout: string; stderr: string; status: "aborted"; durationMs: number }> => + await new Promise((resolve) => { + opts.signal?.addEventListener( + "abort", + () => resolve({ stdout: "", stderr: "", status: "aborted", durationMs: 5000 }), + { once: true }, + ); + }), + ); + const cleanupResources = vi.fn(); + Object.assign( + manager as unknown as { + state: "running"; + executionQueue: Promise; + executeInner: typeof executeInner; + start: () => Promise; + cleanupResources: () => void; + }, + { state: "running", executionQueue: previousExecution, executeInner, start: async () => {}, cleanupResources }, + ); + + const disposal = manager.shutdown({ snapshot: true, drainHostRequests: true }); + // A cell arriving mid-flush must not splice ahead of the final snapshot. + await expect(manager.execute("1 + 1")).rejects.toThrow("Kernel is shutting down"); + releaseQueue(); + await waitForCalls(executeInner, 1); + expect(executeInner.mock.calls[0]?.[0].type).toBe("snapshot"); + await vi.advanceTimersByTimeAsync(5000); + + await expect(disposal).resolves.toBe(true); + expect(executeInner).toHaveBeenCalledOnce(); + expect(cleanupResources).toHaveBeenCalledOnce(); + }); + + it("rejects an execution whose re-bootstrap wait crossed the start of the final flush", async () => { + const manager = new ReplKernelManager({ + cwd: process.cwd(), + snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" }, + bootstrapCode: "boot-code", + }); + let releaseBootstrap: () => void = () => {}; + const bootstrapBlocked = new Promise((resolve) => { + releaseBootstrap = resolve; + }); + let releaseSnapshot: () => void = () => {}; + const snapshotBlocked = new Promise((resolve) => { + releaseSnapshot = resolve; + }); + const calls: string[] = []; + const ok = { stdout: "", stderr: "", status: "ok" as const, durationMs: 0 }; + const executeInner = vi.fn(async (requestFields: Record & { type: string }, code: string) => { + if (requestFields.type === "snapshot") { + calls.push("snapshot"); + await snapshotBlocked; + return { ...ok, doneFields: { saved: [], skipped: [], bytes: 0 } }; + } + if (code === "boot-code") { + calls.push("bootstrap"); + await bootstrapBlocked; + return ok; + } + calls.push("cell"); + return ok; + }); + const cleanupResources = vi.fn(); + Object.assign( + manager as unknown as { + state: "running"; + pendingRebootstrap: boolean; + executeInner: typeof executeInner; + start: () => Promise; + cleanupResources: () => void; + }, + { state: "running", pendingRebootstrap: true, executeInner, start: async () => {}, cleanupResources }, + ); + + // Admitted before the flush: passes the first guard, then parks on the + // in-flight lazy re-bootstrap. + const cell = manager.execute("user-cell"); + cell.catch(() => undefined); + await waitForCalls(executeInner, 1); + expect(calls).toEqual(["bootstrap"]); + + // The teardown's final flush starts while the cell is still parked. + const teardown = manager.shutdown({ snapshot: true }); + await new Promise((resolve) => globalThis.setTimeout(resolve, 0)); + releaseBootstrap(); + + // Un-bootstrapped no more, the cell must NOT splice between the flush's + // captured queue and the final snapshot; the guard re-check rejects it. + await expect(cell).rejects.toThrow("Kernel is shutting down"); + releaseSnapshot(); + await expect(teardown).resolves.toBe(true); + expect(calls).toEqual(["bootstrap", "snapshot"]); + }); + + it("settles an aborted execution promptly while the lazy re-bootstrap is still running", async () => { + const manager = new ReplKernelManager({ cwd: process.cwd(), bootstrapCode: "boot-code" }); + let releaseBootstrap: () => void = () => {}; + const bootstrapBlocked = new Promise((resolve) => { + releaseBootstrap = resolve; + }); + const ok = { stdout: "", stderr: "", status: "ok" as const, durationMs: 0 }; + const executeInner = vi.fn(async (_requestFields: Record & { type: string }, code: string) => { + if (code === "boot-code") { + await bootstrapBlocked; + } + return ok; + }); + Object.assign( + manager as unknown as { + state: "running"; + pendingRebootstrap: boolean; + executeInner: typeof executeInner; + start: () => Promise; + }, + { state: "running", pendingRebootstrap: true, executeInner, start: async () => {} }, + ); + + const controller = new AbortController(); + const cell = manager.execute("user-cell", { signal: controller.signal }); + await waitForCalls(executeInner, 1); + + // Aborted mid-wait: the cell must not ride out the bootstrap bound. + controller.abort(); + await expect(cell).resolves.toMatchObject({ status: "aborted" }); + releaseBootstrap(); + }); + + it("rejects restart() during the final snapshot flush instead of deadlocking the teardown", async () => { + const manager = new ReplKernelManager({ + cwd: process.cwd(), + snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" }, + }); + Object.assign( + manager as unknown as { state: "running"; flushingSnapshotForDispose: boolean; start: () => Promise }, + { + state: "running", + flushingSnapshotForDispose: true, + start: async () => {}, + }, + ); + + // Taking a queue slot here and joining the owning shutdown would leave the + // flush's snapshot parked behind the slot forever (mutual wait). + await expect(manager.restart()).rejects.toThrow("Kernel is shutting down"); + }); + + it("joins concurrent teardowns into a single final snapshot flush", async () => { + const manager = new ReplKernelManager({ + cwd: process.cwd(), + snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" }, + }); + const executeInner = vi.fn( + async ( + _requestFields: Record & { type: string }, + ): Promise<{ stdout: string; stderr: string; status: "ok"; durationMs: number }> => ({ + stdout: "", + stderr: "", + status: "ok", + durationMs: 0, + }), + ); + const cleanupResources = vi.fn(); + Object.assign( + manager as unknown as { + state: "running"; + executionQueue: Promise; + executeInner: typeof executeInner; + start: () => Promise; + cleanupResources: () => void; + }, + { state: "running", executionQueue: Promise.resolve(), executeInner, start: async () => {}, cleanupResources }, + ); + + // A session dispose racing a signal-handler shutdown must share one final + // flush instead of queueing a second snapshot behind the first. + await Promise.all([ + manager.shutdown({ snapshot: true, drainHostRequests: true }), + manager.shutdown({ snapshot: true }), + ]); + + const snapshotRequests = executeInner.mock.calls.filter((call) => call[0].type === "snapshot"); + expect(snapshotRequests).toHaveLength(1); + expect(cleanupResources).toHaveBeenCalled(); }); it("routes null-id and stale-id stream events into backgroundOutput, not stdout", async () => { @@ -328,7 +568,7 @@ describe("ReplKernelManager abort handling", () => { stdin: { destroyed: false, destroy: () => undefined }, }; - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); const types = writeLine.mock.calls.map((call) => (call[0] as { type?: string }).type); expect(types).toContain("shutdown"); expect(killSignals).toContain("SIGTERM"); diff --git a/packages/coding-agent/test/repl-kernel-execute.test.ts b/packages/coding-agent/test/repl-kernel-execute.test.ts index d8c4023e5..c653bce98 100644 --- a/packages/coding-agent/test/repl-kernel-execute.test.ts +++ b/packages/coding-agent/test/repl-kernel-execute.test.ts @@ -37,7 +37,7 @@ describeIf("ReplKernelManager execute (real runtime)", () => { }); afterEach(async () => { - await manager?.dispose(); + await manager?.shutdown({ snapshot: true, drainHostRequests: true }); manager = undefined; if (dir) { rmSync(dir, { recursive: true, force: true }); @@ -128,7 +128,7 @@ describeIf("ReplKernelManager execute (real runtime)", () => { expect(r.status).toBe("ok"); const pid = Number(r.result); expect(Number.isInteger(pid)).toBe(true); - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); let alive = true; for (let i = 0; i < 100; i++) { try { diff --git a/packages/coding-agent/test/repl-kernel-parent-watchdog.test.ts b/packages/coding-agent/test/repl-kernel-parent-watchdog.test.ts index f7f375ad6..5924881fc 100644 --- a/packages/coding-agent/test/repl-kernel-parent-watchdog.test.ts +++ b/packages/coding-agent/test/repl-kernel-parent-watchdog.test.ts @@ -65,7 +65,7 @@ describe("repl kernel parent watchdog", () => { await expect(manager.execute("x")).rejects.toThrow(/Kernel exited before ready/); } finally { errorSpy.mockRestore(); - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); } expect(readFileSync(envDump, "utf8")).toMatch(new RegExp(`^PRIME_AGENT_KERNEL_OWNER_PID=${process.pid}$`, "m")); @@ -95,7 +95,7 @@ describe("repl kernel parent watchdog", () => { const records = existsSync(journalPath) ? readJournalRecords(journalPath) : []; expect(records.some((r) => r.pid === 999999 && !r.active)).toBe(false); } finally { - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }); @@ -115,7 +115,7 @@ describe("repl kernel parent watchdog", () => { const records = readJournalRecords(journalPath); expect(records.some((r) => r.pid === 999999 && !r.active)).toBe(true); } finally { - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }); @@ -160,7 +160,7 @@ describe("repl kernel parent watchdog", () => { } finally { internals.child = undefined; internals.state = "idle"; - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }); @@ -173,7 +173,7 @@ describe("repl kernel parent watchdog", () => { shutdown(): Promise; kill(): Promise; }; - internals.state = "starting"; + internals.state = "running"; // A live child handle keeps waitForKernelExit parked so the send actually blocks the shutdown. const child = Object.assign(new EventEmitter(), { exitCode: null, @@ -267,7 +267,7 @@ describe("repl kernel parent watchdog", () => { internals.child = undefined; internals.startPromise = undefined; internals.state = "idle"; - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }); }); diff --git a/packages/coding-agent/test/repl-kernel-protocol-corruption.test.ts b/packages/coding-agent/test/repl-kernel-protocol-corruption.test.ts new file mode 100644 index 000000000..93bdfa062 --- /dev/null +++ b/packages/coding-agent/test/repl-kernel-protocol-corruption.test.ts @@ -0,0 +1,635 @@ +import { chmodSync, existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { ReplKernelManager } from "../src/core/kernel/index.js"; + +let tempDir = ""; + +function writeFakeRuntime(path: string): void { + writeFileSync( + path, + `#!/usr/bin/env node +const fs = require("node:fs"); +const readline = require("node:readline"); +const countPath = process.env.FAKE_REPL_SPAWN_COUNT; +const count = fs.existsSync(countPath) ? Number(fs.readFileSync(countPath, "utf8")) + 1 : 1; +fs.writeFileSync(countPath, String(count)); +let state = {}; +let boot = "missing"; +setInterval(() => { + if (fs.existsSync(process.env.FAKE_REPL_EMIT_GARBAGE)) { + fs.unlinkSync(process.env.FAKE_REPL_EMIT_GARBAGE); + process.stdout.write("BROKEN-SPONTANEOUS\\n"); + } +}, 25); +const emit = (event) => process.stdout.write(JSON.stringify(event) + "\\n"); +if (count > 1 && fs.existsSync(process.env.FAKE_REPL_DIE_ON_BOOT)) { + fs.writeFileSync(process.env.FAKE_REPL_DIE_ON_BOOT + "-died", "1"); + process.exit(1); +} +if (fs.existsSync(process.env.FAKE_REPL_CORRUPT_BOOT)) { + process.stdout.write("BROKEN-BOOT\\n"); +} else if (fs.existsSync(process.env.FAKE_REPL_READY_WITH_GARBAGE)) { + process.stdout.write( + JSON.stringify({ event: "ready", protocol: 2, python: process.version }) + "\\nBROKEN-WITH-READY\\n", + ); +} else if (!(count > 1 && fs.existsSync(process.env.FAKE_REPL_DELAY_READY))) { + emit({ event: "ready", protocol: 2, python: process.version }); +} +const input = readline.createInterface({ input: process.stdin }); +input.on("line", (line) => { + const request = JSON.parse(line); + if (request.type === "execute") { + if (request.code === "hang") { + emit({ event: "stdout", id: request.id, text: "hanging" }); + return; + } + if (request.code === "bootstrap" && fs.existsSync(process.env.FAKE_REPL_CORRUPT_BOOTSTRAP)) { + fs.unlinkSync(process.env.FAKE_REPL_CORRUPT_BOOTSTRAP); + process.stdout.write("BROKEN-BOOTSTRAP\\n"); + return; + } + if (request.code === "bootstrap" && fs.existsSync(process.env.FAKE_REPL_FAIL_BOOTSTRAP)) { + emit({ event: "error", id: request.id, ename: "RuntimeError", evalue: "bootstrap refused", traceback: [] }); + emit({ event: "done", id: request.id, status: "error" }); + return; + } + if (request.code === "bootstrap") boot = "live"; + if (request.code === "read-boot") emit({ event: "stdout", id: request.id, text: boot }); + if (request.code === "seed") state.value = "persisted"; + if (request.code === "read") emit({ event: "stdout", id: request.id, text: state.value || "fresh" }); + if (request.code === "corrupt-active") { + process.stdout.write("BROKEN-" + "x".repeat(400) + "\\n"); + return; + } + if (request.code === "corrupt-idle") { + process.stdout.write(JSON.stringify({ event: "done", id: request.id, status: "ok" }) + "\\n42\\n"); + return; + } + if (request.code === "slow-bootstrap") { + setTimeout(() => { + boot = "live"; + emit({ event: "done", id: request.id, status: "ok" }); + }, 250); + return; + } + if (request.code === "corrupt-unknown-kind") { + emit({ event: "bogus", id: request.id }); + return; + } + if (request.code === "corrupt-idless-done") { + emit({ event: "done", status: "ok" }); + return; + } + if (request.code === "corrupt-empty-id-done") { + emit({ event: "done", id: "", status: "ok" }); + return; + } + if (request.code === "corrupt-empty-id-host-request") { + emit({ event: "host_request", id: "", data: { type: "noop" } }); + return; + } + emit({ event: "done", id: request.id, status: "ok" }); + return; + } + if (request.type === "snapshot") { + if (fs.existsSync(process.env.FAKE_REPL_CORRUPT_SNAPSHOT)) { + process.stdout.write("BROKEN-SNAPSHOT\\n"); + return; + } + fs.writeFileSync(request.path, JSON.stringify(state)); + fs.writeFileSync(request.manifest_path, "{}"); + emit({ event: "done", id: request.id, status: "ok", saved: Object.keys(state), skipped: [], bytes: 1 }); + return; + } + if (request.type === "restore") { + fs.appendFileSync(process.env.FAKE_REPL_RESTORE_LOG, "r"); + if (fs.existsSync(process.env.FAKE_REPL_CORRUPT_RESTORE)) { + process.stdout.write("BROKEN-RESTORE\\n"); + return; + } + if (fs.existsSync(process.env.FAKE_REPL_FAIL_RESTORE)) { + emit({ event: "done", id: request.id, status: "error", reason: "restore refused" }); + return; + } + state = fs.existsSync(request.path) ? JSON.parse(fs.readFileSync(request.path, "utf8")) : {}; + emit({ event: "done", id: request.id, status: "ok", restored: Object.keys(state), failed: [] }); + return; + } + if (request.type === "shutdown") { + emit({ event: "done", id: request.id, status: "ok" }); + process.exit(0); + } +}); +`, + ); + chmodSync(path, 0o755); +} + +function spawnCount(path: string): number { + return existsSync(path) ? Number(readFileSync(path, "utf8")) : 0; +} + +function restoreCount(path: string): number { + return existsSync(path) ? readFileSync(path, "utf8").length : 0; +} + +describe("ReplKernelManager corrupt protocol repair", () => { + beforeEach(() => { + tempDir = mkdtempSync(join(tmpdir(), "prime-agent-repl-corrupt-")); + }); + + afterEach(() => { + if (tempDir) { + rmSync(tempDir, { recursive: true, force: true }); + tempDir = ""; + } + }); + + function newManager(options: { snapshot?: boolean; debounceMs?: number; bootstrapCode?: string } = {}): { + manager: ReplKernelManager; + bootCorruptionPath: string; + bootstrapCorruptionPath: string; + bootstrapFailurePath: string; + countPath: string; + delayReadyPath: string; + dieOnBootPath: string; + emitGarbagePath: string; + readyGarbagePath: string; + restoreLogPath: string; + restoreCorruptionPath: string; + restoreFailurePath: string; + snapshotCorruptionPath: string; + snapshotPath: string; + } { + const python = join(tempDir, "python"); + const bootCorruptionPath = join(tempDir, "corrupt-boot"); + const bootstrapCorruptionPath = join(tempDir, "corrupt-bootstrap"); + const bootstrapFailurePath = join(tempDir, "fail-bootstrap"); + const countPath = join(tempDir, "spawn-count"); + const delayReadyPath = join(tempDir, "delay-ready"); + const dieOnBootPath = join(tempDir, "die-on-boot"); + const emitGarbagePath = join(tempDir, "emit-garbage"); + const readyGarbagePath = join(tempDir, "ready-with-garbage"); + const restoreCorruptionPath = join(tempDir, "corrupt-restore"); + const restoreLogPath = join(tempDir, "restore-log"); + const restoreFailurePath = join(tempDir, "fail-restore"); + const snapshotCorruptionPath = join(tempDir, "corrupt-snapshot"); + const snapshotPath = join(tempDir, "state.json"); + writeFakeRuntime(python); + return { + manager: new ReplKernelManager({ + python, + cwd: tempDir, + env: { + FAKE_REPL_CORRUPT_BOOT: bootCorruptionPath, + FAKE_REPL_CORRUPT_BOOTSTRAP: bootstrapCorruptionPath, + FAKE_REPL_DIE_ON_BOOT: dieOnBootPath, + FAKE_REPL_FAIL_BOOTSTRAP: bootstrapFailurePath, + FAKE_REPL_CORRUPT_RESTORE: restoreCorruptionPath, + FAKE_REPL_CORRUPT_SNAPSHOT: snapshotCorruptionPath, + FAKE_REPL_DELAY_READY: delayReadyPath, + FAKE_REPL_EMIT_GARBAGE: emitGarbagePath, + FAKE_REPL_FAIL_RESTORE: restoreFailurePath, + FAKE_REPL_READY_WITH_GARBAGE: readyGarbagePath, + FAKE_REPL_RESTORE_LOG: restoreLogPath, + FAKE_REPL_SPAWN_COUNT: countPath, + }, + snapshot: options.snapshot + ? { + path: snapshotPath, + manifestPath: join(tempDir, "manifest.json"), + debounceMs: options.debounceMs ?? 1, + } + : undefined, + bootstrapCode: options.bootstrapCode, + }), + bootCorruptionPath, + bootstrapCorruptionPath, + bootstrapFailurePath, + countPath, + delayReadyPath, + dieOnBootPath, + emitGarbagePath, + readyGarbagePath, + restoreLogPath, + restoreCorruptionPath, + restoreFailurePath, + snapshotCorruptionPath, + snapshotPath, + }; + } + + it("rejects an active request, replaces the child, and restores the last snapshot", async () => { + const { manager, countPath, restoreLogPath, snapshotPath } = newManager({ snapshot: true }); + try { + await manager.execute("seed"); + await expect.poll(() => existsSync(snapshotPath)).toBe(true); + + const corrupt = manager.execute("corrupt-active"); + await expect(corrupt).rejects.toThrow(/Kernel protocol error: unparseable protocol line: BROKEN-/); + await expect(corrupt).rejects.not.toThrow("x".repeat(200)); + + const result = await manager.execute("read"); + expect(result).toMatchObject({ status: "ok", stdout: "persisted" }); + expect(spawnCount(countPath)).toBe(2); + // The repair's successful restore is the only one; the lazy path must not re-restore. + expect(restoreCount(restoreLogPath)).toBe(1); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("runs a queued execute only after the repair restores state", async () => { + const { manager, countPath, snapshotPath } = newManager({ snapshot: true }); + try { + await manager.execute("seed"); + await expect.poll(() => existsSync(snapshotPath)).toBe(true); + + const corrupt = manager.execute("corrupt-active"); + const queued = manager.execute("read"); + await expect(corrupt).rejects.toThrow("Kernel protocol error"); + await expect(queued).resolves.toMatchObject({ status: "ok", stdout: "persisted" }); + expect(spawnCount(countPath)).toBe(2); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("returns an aborted result promptly while the repair is still starting up", async () => { + const { manager, delayReadyPath } = newManager(); + try { + await manager.start(); + writeFileSync(delayReadyPath, "1"); + await expect(manager.execute("corrupt-active")).rejects.toThrow("unparseable protocol line"); + + const controller = new AbortController(); + const pending = manager.execute("read", { signal: controller.signal }); + controller.abort(); + await expect(pending).resolves.toMatchObject({ status: "aborted" }); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("gives up instead of respawn-looping when the replacement corrupts during restore", async () => { + const { manager, countPath, restoreCorruptionPath, restoreLogPath, snapshotPath } = newManager({ + snapshot: true, + }); + try { + await manager.execute("seed"); + await expect.poll(() => existsSync(snapshotPath)).toBe(true); + writeFileSync(restoreCorruptionPath, "1"); + + await expect(manager.execute("corrupt-active")).rejects.toThrow("Kernel protocol error"); + // The single replacement is discarded and no further children spawn. + await new Promise((r) => setTimeout(r, 400)); + expect(spawnCount(countPath)).toBe(2); + // The next execute waits out the abandoned repair and starts a fresh kernel. + await expect(manager.execute("read")).resolves.toMatchObject({ status: "ok", stdout: "fresh" }); + expect(spawnCount(countPath)).toBe(3); + // The snapshot was the declared culprit: exactly one (failed) restore, no lazy retry. + expect(restoreCount(restoreLogPath)).toBe(1); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("discards the replacement kernel when the repair restore fails", async () => { + const { manager, countPath, restoreFailurePath, restoreLogPath, snapshotPath } = newManager({ snapshot: true }); + try { + await manager.execute("seed"); + await expect.poll(() => existsSync(snapshotPath)).toBe(true); + writeFileSync(restoreFailurePath, "1"); + + await expect(manager.execute("corrupt-active")).rejects.toThrow("Kernel protocol error"); + await expect(manager.execute("read")).resolves.toMatchObject({ status: "ok", stdout: "fresh" }); + expect(spawnCount(countPath)).toBe(3); + // The snapshot was the declared culprit: exactly one (failed) restore, no lazy retry. + expect(restoreCount(restoreLogPath)).toBe(1); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("re-runs the runtime bootstrap on the repaired kernel", async () => { + const { manager, countPath, snapshotPath } = newManager({ snapshot: true, bootstrapCode: "bootstrap" }); + try { + await manager.execute("bootstrap"); + await manager.execute("seed"); + await expect.poll(() => existsSync(snapshotPath)).toBe(true); + + await expect(manager.execute("corrupt-active")).rejects.toThrow("Kernel protocol error"); + await expect(manager.execute("read-boot")).resolves.toMatchObject({ status: "ok", stdout: "live" }); + await expect(manager.execute("read")).resolves.toMatchObject({ status: "ok", stdout: "persisted" }); + expect(spawnCount(countPath)).toBe(2); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("requeues a cell that was busy-waiting on an interrupted cell when corruption arrived", async () => { + const { manager, countPath, emitGarbagePath, snapshotPath } = newManager({ snapshot: true }); + try { + await manager.execute("seed"); + await expect.poll(() => existsSync(snapshotPath)).toBe(true); + + // An aborted hung cell stays active, so the next cell busy-waits on it. + const controller = new AbortController(); + let onHanging: () => void = () => {}; + const hanging = new Promise((r) => { + onHanging = r; + }); + const hung = manager.execute("hang", { signal: controller.signal, onStream: () => onHanging() }); + await hanging; + controller.abort(); + await expect(hung).resolves.toMatchObject({ status: "aborted" }); + + const queued = manager.execute("read"); + await new Promise((r) => setTimeout(r, 100)); + writeFileSync(emitGarbagePath, "1"); + await expect(queued).resolves.toMatchObject({ status: "ok", stdout: "persisted" }); + expect(spawnCount(countPath)).toBe(2); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("repairs an object frame with an unknown event kind instead of hanging the request", async () => { + const { manager, countPath } = newManager(); + try { + await expect(manager.execute("corrupt-unknown-kind")).rejects.toThrow( + /Kernel protocol error: unknown protocol event/, + ); + await expect(manager.execute("read")).resolves.toMatchObject({ status: "ok", stdout: "fresh" }); + expect(spawnCount(countPath)).toBe(2); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("repairs a done frame without an id instead of hanging the request", async () => { + const { manager, countPath } = newManager(); + try { + await expect(manager.execute("corrupt-idless-done")).rejects.toThrow( + /Kernel protocol error: done frame without id/, + ); + await expect(manager.execute("read")).resolves.toMatchObject({ status: "ok", stdout: "fresh" }); + expect(spawnCount(countPath)).toBe(2); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("repairs a done frame with an empty-string id instead of hanging the request", async () => { + const { manager, countPath } = newManager(); + try { + await expect(manager.execute("corrupt-empty-id-done")).rejects.toThrow( + /Kernel protocol error: done frame without id/, + ); + await expect(manager.execute("read")).resolves.toMatchObject({ status: "ok", stdout: "fresh" }); + expect(spawnCount(countPath)).toBe(2); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("repairs a host_request frame with an empty-string id instead of hanging the request", async () => { + const { manager, countPath } = newManager(); + try { + await expect(manager.execute("corrupt-empty-id-host-request")).rejects.toThrow( + /Kernel protocol error: host_request frame without id/, + ); + await expect(manager.execute("read")).resolves.toMatchObject({ status: "ok", stdout: "fresh" }); + expect(spawnCount(countPath)).toBe(2); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("fails startup when ready and a corrupt frame arrive in one chunk", async () => { + const { manager, countPath, readyGarbagePath } = newManager(); + try { + writeFileSync(readyGarbagePath, "1"); + await expect(manager.start()).rejects.toThrow("unparseable protocol line: BROKEN-WITH-READY"); + expect(manager.isRunning).toBe(false); + expect(spawnCount(countPath)).toBe(1); + + rmSync(readyGarbagePath); + await expect(manager.start()).resolves.toBeUndefined(); + expect(manager.isRunning).toBe(true); + expect(spawnCount(countPath)).toBe(2); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("restores the saved namespace on the fresh kernel when only the repair bootstrap failed", async () => { + const { manager, bootstrapFailurePath, countPath, snapshotPath } = newManager({ + snapshot: true, + bootstrapCode: "bootstrap", + }); + try { + await manager.execute("bootstrap"); + await manager.execute("seed"); + await expect.poll(() => existsSync(snapshotPath)).toBe(true); + writeFileSync(bootstrapFailurePath, "1"); + + // The repair restores kernel 2 fine but its bootstrap fails, and the + // lazy attempt on kernel 3 fails the same way: the snapshot was never + // the culprit, so neither discard may poison it. + await expect(manager.execute("corrupt-active")).rejects.toThrow("Kernel protocol error"); + await expect(manager.execute("read")).rejects.toThrow("Kernel bootstrap failed after protocol repair"); + + rmSync(bootstrapFailurePath); + await expect(manager.execute("read")).resolves.toMatchObject({ status: "ok", stdout: "persisted" }); + await expect(manager.execute("read-boot")).resolves.toMatchObject({ status: "ok", stdout: "live" }); + expect(spawnCount(countPath)).toBe(4); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("restores the saved namespace on the fresh kernel when the repair start failed", async () => { + const { manager, countPath, dieOnBootPath, snapshotPath } = newManager({ snapshot: true }); + try { + await manager.execute("seed"); + await expect.poll(() => existsSync(snapshotPath)).toBe(true); + writeFileSync(dieOnBootPath, "1"); + + // The replacement dies before ready, so the repair start fails without + // ever implicating the snapshot. + await expect(manager.execute("corrupt-active")).rejects.toThrow("Kernel protocol error"); + await expect.poll(() => existsSync(`${dieOnBootPath}-died`)).toBe(true); + rmSync(dieOnBootPath); + + await expect(manager.execute("read")).resolves.toMatchObject({ status: "ok", stdout: "persisted" }); + expect(spawnCount(countPath)).toBe(3); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("keeps the namespace restorable when corruption strikes after the repair's restore succeeded", async () => { + const { manager, bootstrapCorruptionPath, countPath, restoreLogPath, snapshotPath } = newManager({ + snapshot: true, + bootstrapCode: "bootstrap", + }); + try { + await manager.execute("bootstrap"); + await manager.execute("seed"); + await expect.poll(() => existsSync(snapshotPath)).toBe(true); + writeFileSync(bootstrapCorruptionPath, "1"); + + // Kernel 2 restores cleanly, then corrupts (one-shot) during the + // repair's bootstrap: the snapshot was never implicated. + await expect(manager.execute("corrupt-active")).rejects.toThrow("Kernel protocol error"); + + await expect(manager.execute("read")).resolves.toMatchObject({ status: "ok", stdout: "persisted" }); + await expect(manager.execute("read-boot")).resolves.toMatchObject({ status: "ok", stdout: "live" }); + expect(spawnCount(countPath)).toBe(3); + expect(restoreCount(restoreLogPath)).toBe(2); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("final flush never overwrites the saved snapshot with an unrestored namespace", async () => { + const { manager, dieOnBootPath, snapshotPath } = newManager({ snapshot: true }); + try { + await manager.execute("seed"); + await expect.poll(() => existsSync(snapshotPath)).toBe(true); + writeFileSync(dieOnBootPath, "1"); + + // Repair-start failure leaves the namespace pending restore. + await expect(manager.execute("corrupt-active")).rejects.toThrow("Kernel protocol error"); + await expect.poll(() => existsSync(`${dieOnBootPath}-died`)).toBe(true); + rmSync(dieOnBootPath); + + // restart() resurrects a running kernel without ever reprovisioning it; + // the teardown's final flush must not snapshot that empty namespace + // over the strictly fresher saved one. + await manager.restart(); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + expect(JSON.parse(readFileSync(snapshotPath, "utf8"))).toEqual({ value: "persisted" }); + }); + + it("re-runs the bootstrap on the fresh kernel after a discarded repair", async () => { + const { manager, countPath, restoreFailurePath, snapshotPath } = newManager({ + snapshot: true, + bootstrapCode: "bootstrap", + }); + try { + await manager.execute("bootstrap"); + await manager.execute("seed"); + await expect.poll(() => existsSync(snapshotPath)).toBe(true); + writeFileSync(restoreFailurePath, "1"); + + // The repair's restore fails, so the replacement kernel is discarded. + await expect(manager.execute("corrupt-active")).rejects.toThrow("Kernel protocol error"); + // The lazily started fresh kernel must carry the runtime bootstrap again. + await expect(manager.execute("read-boot")).resolves.toMatchObject({ status: "ok", stdout: "live" }); + expect(spawnCount(countPath)).toBe(3); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("stays bounded when a teardown races the lazy re-bootstrap of a fresh kernel", async () => { + const { manager, countPath, restoreFailurePath, snapshotPath } = newManager({ + snapshot: true, + bootstrapCode: "slow-bootstrap", + }); + await manager.execute("slow-bootstrap"); + await manager.execute("seed"); + await expect.poll(() => existsSync(snapshotPath)).toBe(true); + writeFileSync(restoreFailurePath, "1"); + await expect(manager.execute("corrupt-active")).rejects.toThrow("Kernel protocol error"); + + // The fresh kernel's re-bootstrap (internal + protocolRepair) is in flight + // when the teardown starts: it must pass the final-flush execution guard, + // so the flush's queue wait settles and the shutdown stays bounded. + const pending = manager.execute("read-boot"); + pending.catch(() => undefined); + await expect.poll(() => spawnCount(countPath)).toBe(3); + // The fake journals its spawn before ready: wait for "running" so the + // teardown deterministically races the re-bootstrap, not the boot. + await expect.poll(() => manager.isRunning).toBe(true); + await expect(manager.shutdown({ snapshot: true, drainHostRequests: true })).resolves.toBe(true); + + expect(manager.isRunning).toBe(false); + expect(spawnCount(countPath)).toBe(3); + // The raced cell must settle either way — served before the teardown or + // rejected by it — never hang. + const outcome = await pending.then( + (r) => `${r.status}:${r.stdout}`, + (e) => String(e), + ); + expect(outcome).toMatch(/ok:live|shut down|shutting down/); + }); + + it("repairs a non-object frame received while idle", async () => { + const { manager, countPath } = newManager(); + try { + await expect(manager.execute("corrupt-idle")).resolves.toMatchObject({ status: "ok" }); + await expect(manager.execute("read")).resolves.toMatchObject({ status: "ok", stdout: "fresh" }); + expect(spawnCount(countPath)).toBe(2); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("does not respawn repeatedly when startup emits a corrupt frame", async () => { + const { manager, bootCorruptionPath, countPath } = newManager(); + try { + writeFileSync(bootCorruptionPath, "1"); + await expect(manager.start()).rejects.toThrow("unparseable protocol line: BROKEN-BOOT"); + expect(spawnCount(countPath)).toBe(1); + + rmSync(bootCorruptionPath); + await expect(manager.start()).resolves.toBeUndefined(); + expect(spawnCount(countPath)).toBe(2); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); + + it("does not repair corruption during the dispose snapshot flush", async () => { + const { manager, countPath, snapshotCorruptionPath } = newManager({ + snapshot: true, + debounceMs: 60_000, + }); + await manager.start(); + writeFileSync(snapshotCorruptionPath, "1"); + + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + + expect(manager.isRunning).toBe(false); + expect(spawnCount(countPath)).toBe(1); + }); + + it("stands down when shutdown supersedes the repair", async () => { + const { manager, countPath, delayReadyPath } = newManager(); + try { + await manager.start(); + writeFileSync(delayReadyPath, "1"); + await expect(manager.execute("corrupt-active")).rejects.toThrow("unparseable protocol line"); + + await expect(manager.shutdown()).resolves.toBe(true); + expect(manager.isRunning).toBe(false); + // The superseded repair's booting replacement is hard-killed (it never + // reached "running", so no graceful protocol shutdown); it may die + // before it journals its spawn. Nothing new may spawn afterwards. + const spawnsAfterShutdown = spawnCount(countPath); + expect(spawnsAfterShutdown).toBeLessThanOrEqual(2); + await expect(manager.execute("read")).rejects.toThrow("Kernel has been shut down"); + await new Promise((r) => setTimeout(r, 200)); + expect(spawnCount(countPath)).toBe(spawnsAfterShutdown); + } finally { + await manager.shutdown({ snapshot: true, drainHostRequests: true }); + } + }); +}); diff --git a/packages/coding-agent/test/repl-kernel-shutdown.test.ts b/packages/coding-agent/test/repl-kernel-shutdown.test.ts index 4b93b1104..f39bf469f 100644 --- a/packages/coding-agent/test/repl-kernel-shutdown.test.ts +++ b/packages/coding-agent/test/repl-kernel-shutdown.test.ts @@ -11,6 +11,7 @@ type ShutdownInternals = { wireChild: (child: ShutdownInternals["child"]) => void; pendingDoneWaiters: Map void>; inFlightHostRequests: Set>; + kernelStderr: string; child: EventEmitter & { exitCode: number | null; signalCode: NodeJS.Signals | null; @@ -25,11 +26,16 @@ type ShutdownInternals = { function configuredManager( onSend: (request: Record, internals: ShutdownInternals) => void | Promise, hostHandlers?: HostRequestHandlers, + withSnapshot = false, ): { manager: ReplKernelManager; internals: ShutdownInternals; } { - const manager = new ReplKernelManager({ cwd: process.cwd(), hostHandlers }); + const manager = new ReplKernelManager({ + cwd: process.cwd(), + hostHandlers, + snapshot: withSnapshot ? { path: "/tmp/shutdown-test.dill", manifestPath: "/tmp/shutdown-test.json" } : undefined, + }); const internals = manager as unknown as ShutdownInternals; const child = Object.assign(new EventEmitter(), { exitCode: null, @@ -51,13 +57,21 @@ function configuredManager( } describe("ReplKernelManager graceful shutdown", () => { - it("bounds a stuck stdin write with the aggregate shutdown deadline", async () => { + it.each([ + ["optional policy", {}], + ["snapshot-and-drain policy", { snapshot: true, drainHostRequests: true }], + ] as const)("uses the same diagnostic deadline for the %s", async (_label, options) => { vi.useFakeTimers(); try { - const { manager, internals } = configuredManager(() => new Promise(() => {})); - const shutdown = manager.shutdown(); + const { manager, internals } = configuredManager(() => new Promise(() => {}), undefined, true); + // The final flush snapshots through the private captureSnapshot (bounded, prune-free). + vi.spyOn(manager as unknown as { captureSnapshot: () => Promise }, "captureSnapshot").mockResolvedValue( + null, + ); + const shutdown = manager.shutdown(options); await vi.advanceTimersByTimeAsync(5_000); - await shutdown; + await expect(shutdown).resolves.toBe(true); + expect(internals.kernelStderr).toContain("Kernel did not shut down within 5000ms"); expect(internals.child).toBeUndefined(); } finally { vi.useRealTimers(); @@ -104,15 +118,20 @@ describe("ReplKernelManager graceful shutdown", () => { } }); - it("settles an in-flight host request before dispose tears the kernel down", async () => { + it("flushes a snapshot and drains host requests when session teardown selects both policies", async () => { let releaseHandler: (() => void) | undefined; const handlerBlocked = new Promise((resolve) => { releaseHandler = resolve; }); const sentReplies: Record[] = []; const { manager, internals } = configuredManager( - async (request) => { + async (request, state) => { if (request.type === "host_reply") sentReplies.push(request); + if (request.type === "shutdown") { + state.handleEvent({ event: "done", id: request.id, status: "ok" }); + state.child.exitCode = 0; + state.child.emit("exit", 0, null); + } }, { "test.slow": async () => { @@ -120,14 +139,20 @@ describe("ReplKernelManager graceful shutdown", () => { return { answer: 42, status: "weird" }; }, }, + true, ); + const snapshot = vi + .spyOn(manager as unknown as { captureSnapshot: () => Promise }, "captureSnapshot") + .mockResolvedValue(null); internals.handleEvent({ event: "host_request", id: "hr-1", data: { type: "test.slow" } }); - expect(internals.inFlightHostRequests.size).toBe(1); + const shutdown = manager.shutdown({ snapshot: true, drainHostRequests: true }); + await new Promise((resolve) => globalThis.setTimeout(resolve, 0)); - const dispose = manager.dispose(); + expect(snapshot).toHaveBeenCalledOnce(); + expect(internals.writeLine).not.toHaveBeenCalled(); releaseHandler?.(); - await dispose; + await shutdown; expect(internals.inFlightHostRequests.size).toBe(0); expect(sentReplies).toEqual([ @@ -158,13 +183,18 @@ describe("ReplKernelManager graceful shutdown", () => { return { answer: 42 }; }, }, + true, ); + const snapshot = vi + .spyOn(manager as unknown as { captureSnapshot: () => Promise }, "captureSnapshot") + .mockResolvedValue(null); internals.handleEvent({ event: "host_request", id: "hr-1", data: { type: "test.slow" } }); const tracked = [...internals.inFlightHostRequests]; expect(tracked).toHaveLength(1); await manager.shutdown(); + expect(snapshot).not.toHaveBeenCalled(); expect(internals.child).toBeUndefined(); // The reply write fails against the torn-down child; the task must still settle. @@ -173,6 +203,46 @@ describe("ReplKernelManager graceful shutdown", () => { expect(internals.inFlightHostRequests.size).toBe(0); }); + it("shares one wire exchange across concurrent graceful shutdown calls", async () => { + let exitKernel: (() => void) | undefined; + const { manager, internals } = configuredManager((request, state) => { + state.handleEvent({ event: "done", id: request.id, status: "ok" }); + exitKernel = () => { + state.child.exitCode = 0; + state.child.emit("exit", 0, null); + }; + }); + + const owner = manager.shutdown(); + await new Promise((resolve) => globalThis.setTimeout(resolve, 0)); + const concurrent = manager.shutdown(); + exitKernel?.(); + + await expect(owner).resolves.toBe(true); + await expect(concurrent).resolves.toBe(false); + expect(internals.writeLine).toHaveBeenCalledOnce(); + }); + + it("stands down when a hard teardown supersedes the snapshot flush", async () => { + let releaseSnapshot: (() => void) | undefined; + const snapshotBlocked = new Promise((resolve) => { + releaseSnapshot = () => resolve(null); + }); + const { manager, internals } = configuredManager(() => undefined, undefined, true); + vi.spyOn(manager as unknown as { captureSnapshot: () => Promise }, "captureSnapshot").mockReturnValue( + snapshotBlocked, + ); + const child = internals.child; + + const shutdown = manager.shutdown({ snapshot: true }); + await new Promise((resolve) => globalThis.setTimeout(resolve, 0)); + await manager.kill(); + releaseSnapshot?.(); + + await expect(shutdown).resolves.toBe(false); + expect(child.kill).toHaveBeenCalledOnce(); + }); + it("restart after a graceful shutdown starts the kernel again", async () => { const { manager, internals } = configuredManager((request, state) => { if (request.type !== "shutdown") return; diff --git a/packages/coding-agent/test/repl-kernel-startup.test.ts b/packages/coding-agent/test/repl-kernel-startup.test.ts index d41789ac1..219453f72 100644 --- a/packages/coding-agent/test/repl-kernel-startup.test.ts +++ b/packages/coding-agent/test/repl-kernel-startup.test.ts @@ -35,7 +35,7 @@ describe("ReplKernelManager startup", () => { ); } finally { errorSpy.mockRestore(); - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }); @@ -52,7 +52,7 @@ describe("ReplKernelManager startup", () => { await expect(manager.execute("print(1)")).rejects.toThrow(/speaks protocol 1, expected 3/); } finally { errorSpy.mockRestore(); - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }); @@ -66,7 +66,7 @@ describe("ReplKernelManager startup", () => { await expect(manager.start()).rejects.toThrow(/ENOENT/); } finally { errorSpy.mockRestore(); - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }); @@ -87,7 +87,7 @@ describe("ReplKernelManager startup", () => { } finally { vi.useRealTimers(); errorSpy.mockRestore(); - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); } }); }); diff --git a/packages/coding-agent/test/repl-kernel-state-roundtrip.test.ts b/packages/coding-agent/test/repl-kernel-state-roundtrip.test.ts index 130ea1364..e08ece79a 100644 --- a/packages/coding-agent/test/repl-kernel-state-roundtrip.test.ts +++ b/packages/coding-agent/test/repl-kernel-state-roundtrip.test.ts @@ -60,7 +60,7 @@ describeIfKernel("repl kernel state snapshot round-trip (real runtime)", { tags: expect(existsSync(snapshotPath)).toBe(true); expect(existsSync(manifestPath)).toBe(true); } finally { - await writer.dispose(); + await writer.shutdown({ snapshot: true, drainHostRequests: true }); } const reader = newManager(); @@ -72,7 +72,7 @@ describeIfKernel("repl kernel state snapshot round-trip (real runtime)", { tags: const echo = await reader.execute("print(x, double(x), sum(df))"); expect(echo.stdout.trim()).toBe("42 84 6"); } finally { - await reader.dispose(); + await reader.shutdown({ snapshot: true, drainHostRequests: true }); } }, 60_000); @@ -87,7 +87,7 @@ describeIfKernel("repl kernel state snapshot round-trip (real runtime)", { tags: const restore = await manager.restoreState(); expect(restore).toEqual({ restored: [], failed: [], path: join(freshDir, "missing.dill") }); } finally { - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); rmSync(freshDir, { recursive: true, force: true }); } }, 60_000); @@ -125,7 +125,7 @@ describeIfKernel("repl kernel state snapshot round-trip (real runtime)", { tags: const echo = await manager.execute("print(kept_number + 1, kept_text, 'In' in dir(), 'Out' in dir())"); expect(echo.stdout.trim()).toBe("42 hello False False"); } finally { - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); rmSync(ipythonArtifactDir, { recursive: true, force: true }); } }, 60_000); @@ -141,7 +141,7 @@ describeIfKernel("repl kernel state snapshot round-trip (real runtime)", { tags: expect(names).not.toContain("_hidden"); expect(names).not.toContain("rlm"); } finally { - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); rmSync(listDir, { recursive: true, force: true }); } }, 60_000); @@ -170,7 +170,7 @@ describeIfKernel("repl kernel state snapshot round-trip (real runtime)", { tags: expect(remaining).toContain("small_text"); expect(remaining).not.toContain("large_text"); } finally { - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); rmSync(boundedDir, { recursive: true, force: true }); } }, 60_000); @@ -187,7 +187,7 @@ describeIfKernel("repl kernel state snapshot round-trip (real runtime)", { tags: await manager.execute("auto_var = 'persisted'"); await expect.poll(() => existsSync(autoPath), { timeout: 10_000 }).toBe(true); } finally { - await manager.dispose(); + await manager.shutdown({ snapshot: true, drainHostRequests: true }); rmSync(autoDir, { recursive: true, force: true }); } }, 60_000);