pi-repl-py 0.6.13 → 0.7.0

package/src/engine/index.ts

@@ -1,48 +1,46 @@
-// --- EngineManager: the host half of pi-repl's evaluator, driving a real ipykernel over ---
-// --- ZMTP directly (no guest.py middleman). Owns venv resolution, spawn, queue,   ---
-// --- snapshots, abort grace, and teardown — the wire lives in kernel.ts.         ---
+// --- EngineManager: venv resolution, spawn, queue, snapshots, abort grace, teardown; the wire lives in kernel.ts ---
 
 import { existsSync, mkdirSync, readdirSync, readFileSync, renameSync, rmSync, statSync, writeFileSync } from "node:fs";
 import { homedir } from "node:os";
 import { dirname, join } from "node:path";
-import { KernelClient, type SnapshotEntry } from "./kernel.js";
+import { type HelperLoadResult, KernelClient, type SnapshotEntry } from "./kernel.js";
+
+export type { HelperLoadResult } from "./kernel.js";
 
 function installVenvPython(): string {
 	return join(homedir(), ".pi", "agent", "pi-repl", "venv", "bin", "python3");
 }
 
-/** Prefer a venv with ipykernel; else $PYTHON or python3. */
 function resolvePythonPath(_cwd: string | undefined): string {
-	// Only ever use the install venv: a project or repo `.venv` may lack ipykernel and
-	// shadow the good environment, killing the kernel. No auto-picking.
+	// --- only the install venv: a repo `.venv` may lack ipykernel and would shadow the good one ---
 	const installVenv = installVenvPython();
 	if (existsSync(installVenv)) return installVenv;
 	return process.env.PYTHON ?? "python3";
 }
 
 const DEFAULT_MAX_OUTPUT_CHARS = 46080;
-/** Per-line cap: one genuinely oversized line must not own the channel budget, while legitimately long
- * REPL output (JSON, reprs, errors) still fits under the cap in one piece. Generous enough that only
- * pathological giant lines are trimmed, unlike pi's grep where the line cap keeps matches terse. */
+/** Per-line cap: one giant line must not own the channel budget while long JSON/reprs/errors still pass whole. */
 export const MAX_OUTPUT_LINE_CHARS = 4096;
 const ABORT_GRACE_MS = 20_000;
 const DEFAULT_SNAPSHOT_DEBOUNCE_MS = 1500;
-/** Total snapshot size cap (base64 payload). Per-entry entries are capped at the same
- * bound; larger bindings are reported as skipped names. Mirrors the pi-codex scheme. */
+const DEFAULT_SNAPSHOT_PERIOD_MS = 120_000;
+/** Guard: a periodic refresh stands down for already-heavy namespaces — they re-arm on name churn anyway. */
+export const FORCED_SNAPSHOT_MAX_BYTES = 8 * 1024 * 1024;
+/** Snapshot size cap, also per-entry; oversized bindings are reported as skipped names. */
 const DEFAULT_SNAPSHOT_MAX_BYTES = 128 * 1024 * 1024;
+const RESTORE_QUIET_MS = 250;
+/** Restore-cell deadline: a poisoned pickle would wedge the kernel's single queue forever — kill and mark skipped. */
+const DEFAULT_RESTORE_DEADLINE_MS = 90_000;
 
 interface EngineExecuteError {
-	/** Error class name, e.g. "TypeError". */
 	name: string;
 	message: string;
-	/** Stack trace, split into lines. */
 	stack: string[];
 }
 
 export interface ExecuteResult {
 	stdout: string;
 	stderr: string;
-	/** Rendered value of the cell's final expression, when it has one. */
 	result?: string;
 	status: "ok" | "error" | "aborted";
 	error?: EngineExecuteError;
@@ -53,13 +51,11 @@ export interface ExecuteOptions {
 	/** Aborting cancels the cell via kernel interrupt; the namespace is preserved. */
 	signal?: AbortSignal;
 	onStream?: (chunk: string, name: "stdout" | "stderr") => void;
-	/** Cap stdout / stderr / result at this many characters. Default 45K. */
 	maxOutputChars?: number;
 }
 
 export interface SnapshotResult {
 	path: string;
-	/** Top-level names successfully serialized. */
 	saved: string[];
 	/** Names that could not be serialized, with reasons. */
 	failed: { name: string; reason: string }[];
@@ -74,14 +70,18 @@ export interface RestoreResult {
 export interface EngineOptions {
 	cwd?: string;
 	env?: Record<string, string>;
-	/** Persist/revive the namespace across engine restarts. */
 	snapshot?: {
 		path: string;
-		/** Debounce for the auto-snapshot after each ok cell. Default 1500 ms. */
 		debounceMs?: number;
 		/** Total base64 payload cap; also the per-entry cap. Oversized entries are skipped with a reason. Default 128 MiB. */
 		maxBytes?: number;
+		/** Force a refresh when the last persisted snapshot is older than this, even if no name changed. 0 disables. Default 2 min. */
+		periodMs?: number;
 	};
+	/** Do not revive the snapshot on this engine (used after a wedged restore was detected once). */
+	skipRestore?: boolean;
+	/** True when this engine's snapshot was inherited from a /fork'd parent session. */
+	forkInherited?: boolean;
 }
 
 // --- process-wide cleanup: a child does not die with its parent, so SIGKILL live kernels on exit ---
@@ -102,7 +102,6 @@ function truncateWithMarker(text: string, maxChars: number, wasTruncated: boolea
 	return `${text.slice(0, maxChars)}\n[... output truncated at ${maxChars} chars ...]`;
 }
 
-/** Cap each individual line, so one giant line cannot own the whole channel budget (like grep's line cap). */
 export function capLinesForContext(text: string): { text: string; trimmed: boolean } {
 	const lines = text.split("\n");
 	let trimmed = false;
@@ -116,9 +115,7 @@ export function capLinesForContext(text: string): { text: string; trimmed: boole
 
 const DEFAULT_KEEP_SNAPSHOTS = 25;
 
-/** Scan the state root for per-session snapshot dirs and delete all but the newest `keep`,
- * so a long-lived machine does not accumulate one directory per session forever. The
- * current session's dir is exempt; a snapshot dir without a usable manifest is ignored. */
+/** Keep the newest `keep` snapshot dirs; the live dir is exempt and manifest-less dirs are ignored. */
 export function pruneSnapshotDirs(stateRoot: string, keep: number = DEFAULT_KEEP_SNAPSHOTS, currentDir?: string): void {
 	const entries: { dir: string; mtimeMs: number }[] = [];
 	try {
@@ -141,12 +138,7 @@ export function pruneSnapshotDirs(stateRoot: string, keep: number = DEFAULT_KEEP
 	}
 }
 
-// --- Orphaned-snapshot sweep: snapshot dirs are keyed by conversation file basename, so when
-// --- an owning conversation is deleted (pi removes the .jsonl), its directory becomes dead
-// --- weight. This drops any state dir whose conversation file exists in NONE of the project
-// --- session roots, so deleting a conversation deletes its snapshots with it. Safety rules:
-// --- only dirs that look like ours (contain a namespace.snapshot manifest) are touched, and
-// --- the live session plus the no-session "ephemeral" fallback dir are always exempt. ---
+// --- orphan sweep: a state dir whose conversation file exists in no project root dies with it; only manifest dirs are touched, live + ephemeral exempt ---
 export function pruneOrphanedSnapshotDirs(
 	stateRoot: string,
 	sessionsRoot: string | undefined,
@@ -158,7 +150,11 @@ export function pruneOrphanedSnapshotDirs(
 		for (const proj of readdirSync(sessionsRoot, { withFileTypes: true })) {
 			if (!proj.isDirectory()) continue;
 			for (const f of readdirSync(join(sessionsRoot, proj.name))) {
-				if (f.endsWith(".jsonl")) liveNames.add(f.slice(0, -".jsonl".length));
+				if (!f.endsWith(".jsonl")) continue;
+				const name = f.slice(0, -".jsonl".length);
+				// --- both dir formats (legacy bare-name and slug-keyed) are live while their conversation lives ---
+				liveNames.add(name);
+				liveNames.add(`${proj.name}__${name}`);
 			}
 		}
 	} catch {
@@ -186,28 +182,56 @@ export class EngineManager {
 	private startPromise?: Promise<void>;
 	private executionQueue: Promise<unknown> = Promise.resolve();
 	private snapshotTimer?: ReturnType<typeof setTimeout>;
-	/** User cells currently running on the kernel; the debounced snapshot never cuts in front of one. */
+	/** In-flight user cells; the debounced snapshot never cuts in front of one. */
 	private inFlightCells = 0;
-	/** Last-seen top-level namespace names; snapshots are gated on this set changing. */
+	/** Last-seen namespace names; the snapshot is gated on this set changing. */
 	private lastNamespaceNames?: string[];
 	private pythonPath?: string;
+	private restoredKernel?: KernelClient;
+	/** A wedged revive marks the engine: later kernels boot without restoring. */
+	private restoreSkipped: boolean;
+	private restoreTimer?: ReturnType<typeof setTimeout>;
+	private restoreResolve?: (result: RestoreResult | null) => void;
+	private restorePromise?: Promise<RestoreResult | null>;
+	private restoreSettledResult?: RestoreResult | null;
+	private helperReport: readonly HelperLoadResult[] | null = null;
+	/** Whether the current boot's report has been handed out (once per boot). */
+	private helperReportTaken = true;
+	private readonly forkInherited: boolean;
+	/** When the last snapshot was persisted; 0 = never. Drives the periodic refresh. */
+	private lastPersistedAt = 0;
+	/** Payload bytes of the last persisted snapshot; the periodic refresh stands down above FORCED_SNAPSHOT_MAX_BYTES. */
+	private lastSnapshotBytes = 0;
 
 	constructor(options: EngineOptions = {}) {
 		this.options = options;
+		this.forkInherited = options.forkInherited ?? false;
+		this.restoreSkipped = options.skipRestore ?? false;
+		// no snapshot capability: recovery is trivially "nothing to revive"
+		if (!options.snapshot) this.settleRestore(null);
 	}
 
 	get isRunning(): boolean {
 		return this.state === "running" && (this.kernel?.isRunning ?? false);
 	}
 
-	// -- state can change to "shutdown" from kill()/dispose() at any time; read it
-	// through a method so TS doesn't narrow the union and flag a false "no overlap" --
+	/** True when this engine is a fork that inherited its parent's namespace (drives the fork toast). */
+	get inheritedFromFork(): boolean {
+		return this.forkInherited;
+	}
+
+	/** Current boot's helper verdicts, handed out once per boot (first cell of a session/rebuild); null when nothing to announce. */
+	takeHelperReport(): readonly HelperLoadResult[] | null {
+		if (this.helperReportTaken) return null;
+		this.helperReportTaken = true;
+		return this.helperReport;
+	}
+
+	// --- state can flip to shutdown at any time; read it via a method so TS can't narrow the union away ---
 	private isShutdown(): boolean {
 		return this.state === "shutdown";
 	}
 
-	//lifecycle
-
 	async start(): Promise<void> {
 		if (this.state === "shutdown") throw new Error("Engine has been shut down");
 		if (!this.startPromise) {
@@ -234,14 +258,17 @@ export class EngineManager {
 				env: this.options.env,
 				timeoutMs,
 			});
-			// --- an unexpected kernel death must not survive the next execute: drop the dying
-			// --- instance and clear the boot cache so start() rebuilds it on the next cell. ---
+			// --- a fresh boot's helper verdicts are announced once, on the first cell after it ---
+			this.helperReport = this.kernel.helperReport;
+			this.helperReportTaken = false;
+			// --- drop a dead kernel so the next execute rebuilds; never resume a zombie ---
 			const current = this.kernel;
 			current.setOnUnexpectedExit(() => {
 				if (this.kernel !== current) return;
 				this.kernel = undefined;
 				this.startPromise = undefined;
 				this.lastNamespaceNames = undefined;
+				this.helperReport = null;
 			});
 		} catch (error) {
 			if (this.state === "starting") this.state = "idle";
@@ -255,6 +282,8 @@ export class EngineManager {
 			throw new Error("Engine has been shut down");
 		}
 		this.state = "running";
+		// --- recovery runs in the first quiet gap — never ahead of a user cell, never on the first call's critical path ---
+		this.maybeScheduleRestore();
 	}
 
 	/** Abrupt teardown: SIGKILL the kernel; safe from process.on("exit"). */
@@ -270,7 +299,6 @@ export class EngineManager {
 		this.killSync();
 	}
 
-	/** Graceful cleanup: flush a final snapshot, then terminate the kernel. */
 	async dispose(): Promise<void> {
 		if (this.state === "running") {
 			await this.snapshotState().catch(() => null);
@@ -296,12 +324,19 @@ export class EngineManager {
 				throw new Error("Engine has been shut down");
 			}
 			await this.start();
-			// --- the kernel may have died after the boot promise resolved but before the async
-			// --- exit event surfaced it; drop the zombie and rebuild so the next cell runs. ---
+			// --- the kernel may have died after boot resolved but before its exit event; drop the zombie and rebuild ---
 			if (this.kernel && !this.kernel.isRunning) {
 				this.kernel = undefined;
 				this.startPromise = undefined;
 				await this.start();
+				// --- a mid-session rebuild revives the snapshot BEFORE the triggering cell (startup recovery is background); a wedged revive kills the kernel and the retry skips the restore ---
+				await this.restoreWithReap().catch(() => null);
+				// --- read health via the getter: TS narrowed this.kernel away, but start() may have replaced it ---
+				if (!this.isRunning) {
+					this.kernel = undefined;
+					this.startPromise = undefined;
+					await this.start();
+				}
 			}
 			if (this.isShutdown()) {
 				throw new Error("Engine has been shut down");
@@ -333,8 +368,7 @@ export class EngineManager {
 					onStream: opts.onStream,
 					maxOutputChars: maxChars,
 				});
-				// --- names gate runs off the critical path so the next execute's kernel
-				// --- request enqueues before the list-names hop, not behind it ---
+				// --- names gate runs off the critical path so the next cell enqueues before it ---
 				if (r.status === "ok") setImmediate(() => void this.scheduleSnapshotIfChanged());
 				const status: ExecuteResult["status"] = opts.signal?.aborted ? "aborted" : r.status;
 				// Channel cap (truncateWithMarker), then per-line cap; both append a marker so truncation is explicit.
@@ -376,45 +410,157 @@ export class EngineManager {
 			// --- an incomplete snapshot must not overwrite the last good file ---
 			if (reply.complete === false) return null;
 			mkdirSync(dirname(config.path), { recursive: true });
-			// --- write to a temp file then rename so a crash mid-write can never corrupt
-			// --- the last good snapshot (the restore side parses or returns null) ---
+			// --- atomic write: temp file + rename, so a crash can't corrupt the last good snapshot ---
 			const tmp = `${config.path}.tmp`;
-			writeFileSync(tmp, JSON.stringify({ version: 2, entries: reply.entries, failed: reply.failed }));
+			writeFileSync(tmp, JSON.stringify({ version: 3, entries: reply.entries, failed: reply.failed }));
 			renameSync(tmp, config.path);
+			this.lastPersistedAt = Date.now();
+			this.lastSnapshotBytes = reply.entries.reduce((n, e) => n + e.payload.length, 0);
+			await this.advanceSnapshotGate();
 			return { path: config.path, saved: reply.entries.map((e) => e.name), failed: reply.failed };
 		} catch {
 			return null;
 		}
 	}
 
+	/** On success, advance the name-diff gate to today's names — a failed write never blocks the retry. */
+	private async advanceSnapshotGate(): Promise<void> {
+		const names = await this.listNamespaceNames();
+		if (names !== null && names.length > 0) this.lastNamespaceNames = [...names].sort();
+	}
+
+	/** Restore outcome (never rejects); the background quiet-gap job — this promise is the lifecycle's only announce hook. */
+	restoreResult(): Promise<RestoreResult | null> {
+		if (this.restoreSettledResult !== undefined) return Promise.resolve(this.restoreSettledResult);
+		if (!this.restorePromise) {
+			this.restorePromise = new Promise<RestoreResult | null>((resolve) => {
+				this.restoreResolve = resolve;
+			});
+		}
+		return this.restorePromise;
+	}
+
+	/** True when the restore was deliberately skipped (prior wedge); lets the lifecycle say exactly why. */
+	restoreWasSkipped(): boolean {
+		return this.restoreSkipped;
+	}
+
+	private settleRestore(result: RestoreResult | null): void {
+		if (this.restoreSettledResult !== undefined) return;
+		this.restoreSettledResult = result;
+		this.restoreResolve?.(result);
+		this.restoreResolve = undefined;
+	}
+
+	/** Background revive: fires in the first quiet gap (never ahead of a user cell); mid-session rebuilds force it synchronously. */
+	private maybeScheduleRestore(): void {
+		const config = this.options.snapshot;
+		if (!config) return;
+		if (this.restoreSkipped) {
+			this.settleRestore(null);
+			return;
+		}
+		if (this.kernel && this.kernel === this.restoredKernel) return; // this kernel already revived
+		if (!existsSync(config.path)) {
+			this.settleRestore(null);
+			return;
+		}
+		if (this.restoreTimer) return; // already armed
+		const arm = () => {
+			this.restoreTimer = setTimeout(() => {
+				this.restoreTimer = undefined;
+				// --- quiet-gap rule: a pickling restore must not queue ahead of the user's next cell ---
+				if (this.inFlightCells > 0 || !this.kernel?.isRunning) {
+					arm();
+					return;
+				}
+				void this.runRestore();
+			}, RESTORE_QUIET_MS);
+			this.restoreTimer.unref?.();
+		};
+		arm();
+	}
+
+	/** Restore-cell watchdog: an unpickling that never returns would wedge the single queue forever — kill, skip, and rebuild honestly. */
+	private async restoreWithReap(): Promise<RestoreResult | null> {
+		const config = this.options.snapshot;
+		if (!config || !this.kernel || this.restoreSkipped) {
+			this.settleRestore(null);
+			return null;
+		}
+		if (this.kernel === this.restoredKernel) return this.restoreResult();
+		const deadlineMs =
+			Number(process.env.PI_REPL_BOOT_TIMEOUT_MS ?? this.options.env?.PI_REPL_BOOT_TIMEOUT_MS ?? 0) ||
+			DEFAULT_RESTORE_DEADLINE_MS;
+		const reaper = setTimeout(() => {
+			this.restoreSkipped = true;
+			this.settleRestore(null);
+			this.kernel?.kill();
+		}, deadlineMs);
+		reaper.unref?.();
+		try {
+			return await this.restoreState(false).catch(() => null);
+		} finally {
+			clearTimeout(reaper);
+		}
+	}
+
+	private async runRestore(): Promise<void> {
+		await this.restoreWithReap();
+	}
+
+	/** Restore, idempotent per kernel: a second call shares the in-flight outcome. */
 	async restoreState(skip = false): Promise<RestoreResult | null> {
-		// --- start unconditionally: the boot deadline in the lifecycle bounds this call, so
-		// --- booting eagerly here (even with nothing to revive) is what makes a wedged boot
-		// --- detectable instead of deferring the wedge to the first cell. ---
+		// --- start unconditionally: direct callers may not have started, and a wedged boot is only visible mid-attempt ---
 		await this.start();
-		if (skip) return null;
+		if (skip) {
+			this.settleRestore(null);
+			return null;
+		}
 		const config = this.options.snapshot;
-		if (!config) return null;
-		if (!existsSync(config.path)) return null;
+		const kernel = this.kernel;
+		if (!config || !kernel || this.restoreSkipped) {
+			this.settleRestore(null);
+			return null;
+		}
+		if (kernel === this.restoredKernel) {
+			return this.restoreResult();
+		}
+		// claim the kernel now so the quiet-gap scheduler cannot start a second restore cell
+		this.restoredKernel = kernel;
+		if (!existsSync(config.path)) {
+			this.settleRestore(null);
+			return null;
+		}
 		try {
 			const payload = JSON.parse(readFileSync(config.path, "utf8")) as {
 				version?: number;
 				entries?: SnapshotEntry[];
 				vars?: Record<string, string>;
+				failed?: { name: string; reason: string }[];
 			};
-			// --- version 1 files (pre-source-capture) are still restorable: their vars are plain pickles ---
+			// --- v1 files (pre-source-capture) restore via plain pickles; v3 value entries are zlib-compressed ---
 			const entries: SnapshotEntry[] =
-				payload.version === 2
+				payload.version !== undefined && payload.version >= 2
 					? (payload.entries ?? [])
 					: Object.entries(payload.vars ?? {}).map(([name, b64]) => ({ name, kind: "value", payload: b64 }));
-			const reply = await this.kernel!.restore(entries);
-			return { path: config.path, restored: reply.restored, failed: reply.failed };
+			const reply = await kernel.restore(entries, payload.version === 3);
+			// --- merge save-time skips (oversized bindings) into the result so the resume notice names every loss ---
+			const failed = [...(payload.failed ?? [])];
+			const seen = new Set(failed.map((f) => f.name));
+			for (const f of reply.failed) {
+				if (!seen.has(f.name)) failed.push(f);
+				seen.add(f.name);
+			}
+			const result: RestoreResult = { path: config.path, restored: reply.restored, failed };
+			this.settleRestore(result);
+			return result;
 		} catch {
+			this.settleRestore(null);
 			return null;
 		}
 	}
 
-	/** The conversation's state dir exists, so this engine is a resume, not a first run. */
 	hasSnapshotHistory(): boolean {
 		const config = this.options.snapshot;
 		return config ? existsSync(dirname(config.path)) : false;
@@ -429,8 +575,7 @@ export class EngineManager {
 		}
 	}
 
-	/** Snapshot only if the set of top-level names changed since the last snapshot. Names-only
-	 * comparison is cheap (no pickling); a cell that reuses existing state skips the heavy dump. */
+	/** Snapshot on name change, or when the last persisted snapshot went stale (periodMs) — same-name mutations would otherwise never re-arm; a failed write leaves the gate in place. */
 	private async scheduleSnapshotIfChanged(): Promise<void> {
 		const config = this.options.snapshot;
 		if (!config) return;
@@ -438,8 +583,14 @@ export class EngineManager {
 		if (names === null || names.length === 0) return;
 		const key = [...names].sort().join(",");
 		const prev = this.lastNamespaceNames ? [...this.lastNamespaceNames].sort().join(",") : undefined;
-		if (prev !== undefined && prev === key) return; // nothing changed
-		this.lastNamespaceNames = [...names].sort();
+		const changed = prev === undefined || prev !== key;
+		const periodMs = config.periodMs ?? DEFAULT_SNAPSHOT_PERIOD_MS;
+		const stale =
+			periodMs > 0 &&
+			this.lastPersistedAt > 0 &&
+			Date.now() - this.lastPersistedAt >= periodMs &&
+			this.lastSnapshotBytes <= FORCED_SNAPSHOT_MAX_BYTES;
+		if (!changed && !stale) return;
 		this.scheduleSnapshot();
 	}
 
@@ -451,9 +602,7 @@ export class EngineManager {
 		const fire = () => {
 			this.snapshotTimer = undefined;
 			if (this.inFlightCells > 0) {
-				// --- a pickling cell would wait ahead of the user's next request on the
-				// --- kernel's single queue; the snapshot only lands in a real quiet gap,
-				// --- so re-arm the full quiet window and let activity settle instead ---
+				// --- pickling must not queue ahead of the user's next request; re-arm the quiet window until the kernel is idle ---
 				this.snapshotTimer = setTimeout(fire, quiet);
 				this.snapshotTimer.unref?.();
 				return;