From d60fab8a76d9c169f945341f0ee3bde21903bb55 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20M=C3=BCller?= Date: Fri, 28 Aug 2026 17:40:01 +0200 Subject: [PATCH] fix(coding-agent): repair corrupt REPL protocol frames (#1839) * refactor(coding-agent): unify kernel shutdown * docs(coding-agent): carry over shutdown rationale comments * fix(coding-agent): repair corrupt REPL protocol * fix(coding-agent): cancel the final kernel snapshot before disposal cleanup * fix(coding-agent): block new executions from splicing ahead of the final dispose snapshot flushSnapshotForDispose captured the execution queue while the kernel state was still 'running', so a concurrent execute() (which has no per-cell timeout) could enqueue after the capture; captureSnapshot then queued behind it and its timeout never started, hanging dispose() and shutdown({snapshot:true}) forever. Claim the flush before capturing the queue: non-internal requests are rejected while the final snapshot is flushing, keeping every wait on the dispose path bounded. * fix(coding-agent): close protocol-repair races and bound the repair - Gate queued requests behind an in-flight repair: a request whose queue slot opens during a repair releases the slot and requeues after the repair, so queued executes/snapshots can no longer run on the empty replacement kernel ahead of restoreState (and overwrite its result or clobber the last good snapshot). - Race the repair wait in execute() with opts.signal so an aborted call returns its aborted result promptly instead of pending through repair. - Bound the repair's restore with REPAIR_RESTORE_TIMEOUT_MS so a stalled replacement cannot wedge protocolRepairPromise (and every later execute) forever; a timed-out restore discards the replacement and returns to clean idle. - Stop respawn-looping when the replacement kernel itself emits a corrupt frame during repair: discard it and settle at clean idle instead of superseding into an unbounded chain of repairs. * fix(coding-agent): re-bootstrap repaired kernels and gate busy-waiting cells - A protocol-repaired kernel restored only the user namespace, so live handles (rlm, bash, skills) were missing and the next cell failed with NameErrors. KernelManagerOptions gains bootstrapCode (the provisioner passes its buildRlmBootstrapCode output); the repair re-runs it after restore, bounded like the restore, and discards the replacement to clean idle when it fails. - The repair requeue gate only ran right after taking the queue slot, so a cell busy-waiting on an interrupted execution when corruption arrived could proceed onto the unrestored replacement while blocking the repair's restore behind its slot. The gate now sits after the busy-wait (still exempting the repair's own requests), covering both windows. - Fold the repeated kill-to-clean-idle teardown into killChildToIdle and rename REPAIR_RESTORE_TIMEOUT_MS to REPAIR_STEP_TIMEOUT_MS since it now bounds restore and bootstrap. * fix(coding-agent): join concurrent teardowns into one final snapshot flush Concurrent dispose() and shutdown({ snapshot: true }) (e.g. a session dispose racing the SIGINT handler) each ran flushSnapshotForDispose: the first finisher's finally cleared flushingSnapshotForDispose while the second was still flushing, dropping the execution guard early and enqueueing a duplicate final snapshot behind the first. Memoize the in-flight flush so every teardown joins the same bounded flush. * fix(coding-agent): close protocol repair gaps for schema-invalid frames, startup chunks, and rebootstrap - Fail the protocol on JSON object frames that violate the frame schema (unknown event kind; done/host_request without a string id) instead of silently dropping them, which left the active request unsettled forever. The protocol version handshake is exact, so unknown kinds are corruption, not forward compatibility. - A corrupt frame sharing one stdout chunk with ready no longer lets the startup mark the corrupt child running: failProtocolFrame records the error while starting and doStart fails the start. - A fresh kernel started lazily after a discarded repair re-runs the configured runtime bootstrap before serving user code, so rlm/bash/skills handles are never silently missing. * test(coding-agent): reconcile chain-collapse suites and cover teardown racing the lazy re-bootstrap - Point the graceful-shutdown flush spies at captureSnapshot: the merged final flush (kernel-snapshot-dispose-timeout) snapshots through the bounded captureSnapshot path, not snapshotState, so the old spies let real snapshot requests hit the mocked wire and leaked fake timers. - The superseded repair's booting replacement is now hard-killed without a protocol shutdown (single-kernel-shutdown only drains running kernels), so it may die before journaling its spawn; assert spawn-count stability after teardown instead of an exact boot count. - New composition test: a teardown racing an in-flight lazy re-bootstrap stays bounded, spawns nothing new, and settles the raced cell. * fix(coding-agent): re-check the final-flush guard after the re-bootstrap wait - An external execute admitted before the final dispose flush could park on the lazy re-bootstrap, then take its queue slot after the flush captured pendingExecutions - splicing between the captured queue and the final snapshot and unbounding the teardown. Re-check the flush guard right after the ensure await; unit test proves the re-check is load-bearing. - Deflake the teardown-vs-re-bootstrap composition test: the fake runtime journals its spawn before ready, so poll isRunning before the teardown to race the re-bootstrap deterministically instead of the boot. * fix(coding-agent): close bot-review gaps in protocol repair and teardown seams - Terminal frames (done/host_request) with an EMPTY string id now fail the protocol like id-less ones: the runtime mints uuid hex ids and echoes the host's uuids, so an empty id is corruption that previously slipped past the validator into the silent-drop hang. - An execution aborted while waiting on the lazy re-bootstrap settles promptly: the wait races the abort signal (same pattern as waitForProtocolRepair) and the aborted request returns before parking on the queue slot behind the still-running bootstrap. - A failed kernel startup awaits its own teardown (final snapshot flush included) before surfacing the error, so a replacement provisioner gated on this dispose can no longer race the still-flushing kernel over the same snapshot files. - restart() during the final dispose flush is rejected instead of taking a queue slot and joining the owning shutdown, which could deadlock the flush's snapshot behind the held slot. * test(coding-agent): cover empty-string host_request ids in the frame validator Kills the surviving mutant that relaxed the host_request half of the non-empty-id requirement back to typeof-only: an empty-id host_request now has a corruption test mirroring the empty-id done one (protocol fails and the kernel is repaired instead of the awaiting cell hanging). * fix(coding-agent): restore the saved namespace on lazily reprovisioned kernels After a discarded repair whose failure never implicated the snapshot (the repair start died, or only its bootstrap failed), the lazily started fresh kernel re-ran the bootstrap but silently abandoned the user's saved namespace forever. Track restore separately: killChildToIdle marks both pending; the two snapshot-culprit discard paths (clean restore failure, replacement corrupted mid-repair) clear the restore flag so a poisoned snapshot is never retried; the lazy path restores one-shot best-effort (ordinary startup semantics) before bootstrapping, and a teardown's final snapshot never triggers reprovisioning. Restore-attempt journaling in the fake runtime pins the no-retry and no-double-restore semantics. * fix(coding-agent): protect the saved snapshot from unrestored kernels; phase-aware corruption attribution - The final dispose flush skips its snapshot while pendingRestore is set: a kernel that never restored the saved namespace must not overwrite it (the on-disk snapshot is strictly fresher). One interleaving-independent condition instead of joining the reprovision promise, which would reintroduce the 30-60s teardown stall the teardown guard prevents. The lazy one-shot clear moves after the corruption check so the flag also covers a flush racing an in-flight lazy restore. - The repair give-up path attributes corruption by phase: performRestore clears pendingRestore on success, so the flag still being set means the corruption struck at or before the restore (snapshot stays the prime suspect, loop-safe); corruption strictly after a successful restore keeps the namespace restorable at the cost of at most one bounded restore per later attempt. --- .../kernel-snapshot-dispose-timeout.md | 1 + .../.changes/protocol-frame-repair.md | 1 + .../.changes/single-kernel-shutdown.md | 1 + .../src/core/kernel/repl-manager.ts | 506 +++++++++++--- .../coding-agent/src/core/kernel/shared.ts | 19 +- .../coding-agent/src/core/tools/ipython.ts | 13 +- .../test/agent-session-recursion.test.ts | 2 +- .../test/ipython-bootstrap.test.ts | 4 +- .../test/ipython-provisioner.test.ts | 58 +- .../test/kernel-agent-message-skill.test.ts | 2 +- .../test/repl-kernel-abort.test.ts | 262 +++++++- .../test/repl-kernel-execute.test.ts | 4 +- .../test/repl-kernel-parent-watchdog.test.ts | 12 +- .../repl-kernel-protocol-corruption.test.ts | 635 ++++++++++++++++++ .../test/repl-kernel-shutdown.test.ts | 90 ++- .../test/repl-kernel-startup.test.ts | 8 +- .../test/repl-kernel-state-roundtrip.test.ts | 14 +- 17 files changed, 1480 insertions(+), 152 deletions(-) create mode 100644 packages/coding-agent/.changes/kernel-snapshot-dispose-timeout.md create mode 100644 packages/coding-agent/.changes/protocol-frame-repair.md create mode 100644 packages/coding-agent/.changes/single-kernel-shutdown.md create mode 100644 packages/coding-agent/test/repl-kernel-protocol-corruption.test.ts 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); -- 2.54.0