pi-repl-py 0.6.14 → 0.7.0

package/src/engine/kernel.ts

@@ -1,5 +1,3 @@
-// --- KernelClient: one ipykernel subprocess driven directly over ZMTP (no guest middleman). ---
-
 import { type ChildProcess, spawn } from "node:child_process";
 import { randomUUID } from "node:crypto";
 import { existsSync, readdirSync, readFileSync, rmSync } from "node:fs";
@@ -9,6 +7,7 @@ import { resolveHelperDirs } from "./helpers-locate.js";
 import {
 	type ConnectionFile,
 	executeRequest,
+	isTrustedMessage,
 	JupyterSession,
 	NAMES_MIME,
 	type ParsedMessage,
@@ -36,20 +35,18 @@ export interface CellResult {
 	result?: string;
 	error?: { name: string; message: string; stack: string[] };
 	status: "ok" | "error" | "aborted";
-	/** Per-channel output was capped; the host adds a truncation marker. */
 	truncated?: { stdout: boolean; stderr: boolean };
 }
 
 export interface CellOptions {
 	signal?: AbortSignal;
 	onStream?: (chunk: string, name: "stdout" | "stderr") => void;
-	/** Cap per-channel output accumulation. Default 1 MiB (the old guest cap). */
 	maxOutputChars?: number;
 }
 
 export interface SnapshotEntry {
 	name: string;
-	/** "value" pickles the object; "def" re-executes captured source (functions and classes). */
+	/** "value" = zlib-compressed pickle (v3 files; v2 are plain); "def" re-executes captured source (functions and classes). */
 	kind: "value" | "def";
 	payload: string;
 }
@@ -60,12 +57,15 @@ export interface SnapshotReply {
 	complete: boolean;
 }
 
-// --- boot preload: exec each helper; ls()/help() are gone, discovery is globals() ---
+export interface HelperLoadResult {
+	name: string;
+	ok: boolean;
+	/** First error line, e.g. "NameError: name 'x' is not defined"; undefined when ok. */
+	error?: string;
+}
 
 /** Read the helpers dir (same skip rules as the extension's prompt loader). */
-/** A directory for the kernel to start in; if the requested cwd is gone, fall back to the
- * evaluator's own cwd rather than letting spawn() die with ENOENT. A deleted project dir is
- * a real resume case (pi guards it too) — the kernel must still come up. */
+/** Kernel cwd falls back to the host cwd when the requested dir is gone (a deleted project is a real resume case). */
 function resolveCwd(requested?: string): string {
 	if (requested && existsSync(requested)) return requested;
 	return process.cwd();
@@ -97,14 +97,8 @@ function buildSkipList(helperNames: string[]): string {
 
 function snapshotCode(helperNames: string[], maxBytes: number): string {
 	const skip = buildSkipList(helperNames);
-	// --- functions and classes defined in cells cannot be pickled by reference, so their
-	// --- source is captured instead and re-executed on restore. getsource works for
-	// --- functions because the code object carries the cell's filename in linecache; for
-	// --- classes inspect's module-file lookup misses, so a class is captured by locating
-	// --- its header from a member method's co_firstlineno and dedent-scanning the block.
-	// --- fallback pickles the value and reports it if that also fails; per-entry and total
-	// --- byte caps mirror the pi-codex scheme: oversized bindings become skipped names.
-	return `import pickle as _pk, base64 as _b64, json as _js, inspect as _in, linecache as _lc
+	// --- defs/classes can't be pickled by reference: capture source and re-exec on restore; oversized entries are skipped by name ---
+	return `import pickle as _pk, base64 as _b64, json as _js, zlib as _zl, inspect as _in, linecache as _lc
 def _repl_class_source(_c):
     _m = getattr(_c, '__init__', None)
     if _m is None or not _in.isfunction(_m):
@@ -174,7 +168,7 @@ for _k, _v in list(globals().items()):
         __repl_kind = 'value'
     try:
         if __repl_p is None:
-            __repl_p = _b64.b64encode(_pk.dumps(_v)).decode()
+            __repl_p = _b64.b64encode(_zl.compress(_pk.dumps(_v), 1)).decode()
         __repl_b = len(__repl_p)
         if __repl_b > __repl_max:
             __repl_f.append({'name': _k, 'reason': 'exceeds per-entry snapshot cap'})
@@ -188,15 +182,14 @@ for _k, _v in list(globals().items()):
 get_ipython().display_pub.publish({${JSON.stringify(SNAPSHOT_MIME)}: _js.dumps({'version': 2, 'entries': __repl_e, 'failed': __repl_f})})`;
 }
 
-function restoreCode(entries: SnapshotEntry[]): string {
+function restoreCode(entries: SnapshotEntry[], compressedValues: boolean): string {
 	const per = entries
 		.map(({ name, kind, payload }) => {
 			const n = JSON.stringify(name);
 			const body =
 				kind === "def"
 					? // re-execute captured source and register it in linecache under the code
-						// object's filename so a later snapshot can capture it as source again;
-						// exec also binds the name the source defines.
+						// --- exec also registers the source in linecache so a later snapshot can capture it again ---
 						`__repl_src = _b64.b64decode(${JSON.stringify(payload)}).decode()
     exec(__repl_src, globals())
     __repl_obj = globals().get(${n})
@@ -207,7 +200,9 @@ function restoreCode(entries: SnapshotEntry[]): string {
             __repl_fname = getattr(getattr(__repl_init, '__code__', None), 'co_filename', None)
         if __repl_fname:
             _lc.cache[__repl_fname] = (len(__repl_src.splitlines()), None, __repl_src.splitlines(True), __repl_fname)`
-					: `globals()[${n}] = _pk.loads(_b64.b64decode(${JSON.stringify(payload)}))`;
+					: compressedValues
+						? `globals()[${n}] = _pk.loads(_zl.decompress(_b64.b64decode(${JSON.stringify(payload)})))`
+						: `globals()[${n}] = _pk.loads(_b64.b64decode(${JSON.stringify(payload)}))`;
 			return `try:
     ${body}
     __repl_r['restored'].append(${n})
@@ -215,7 +210,7 @@ except Exception as _e:
     __repl_r['failed'].append({'name': ${n}, 'reason': str(_e)})`;
 		})
 		.join("\n");
-	return `import pickle as _pk, base64 as _b64, json as _js, linecache as _lc
+	return `import pickle as _pk, base64 as _b64, json as _js, zlib as _zl, linecache as _lc
 __repl_r = {'restored': [], 'failed': []}
 ${per}
 get_ipython().display_pub.publish({${JSON.stringify(RESTORE_MIME)}: _js.dumps(__repl_r)})`;
@@ -241,7 +236,6 @@ interface ActiveCell {
 	result?: string;
 	error?: { name: string; message: string; stack: string[] };
 	status: CellResult["status"];
-	/** Private-MIME payloads published by this cell (snapshot/restore/names). */
 	payloads: Record<string, string>;
 	lastActivity: number;
 	timedOut: boolean;
@@ -252,11 +246,8 @@ interface ActiveCell {
 	resolve(result: CellResult & { payloads: Record<string, string> }): void;
 	reject(error: Error): void;
 	settled: boolean;
-	/** The shell execute_reply arrived (carries the authoritative status). */
 	replySeen: boolean;
-	/** The matching iopub status idle arrived (published after all output). */
 	idleSeen: boolean;
-	/** The execute_reply content, held until both halves are seen. */
 	reply?: ParsedMessage;
 }
 
@@ -267,6 +258,7 @@ export class KernelClient {
 	private iopub?: ZmtpSocket;
 	private readonly session: JupyterSession;
 	private readonly helperSources: { name: string; source: string }[];
+	private helperBootReport: HelperLoadResult[] = [];
 	private readonly timeoutMs: number;
 	private activeCell?: ActiveCell;
 	private connectionFilePath?: string;
@@ -274,6 +266,10 @@ export class KernelClient {
 	/** Serializes all kernel ops: one execute at a time, snapshots between cells. */
 	private queue: Promise<unknown> = Promise.resolve();
 	private _onUnexpectedExit?: () => void;
+	get helperReport(): readonly HelperLoadResult[] {
+		return this.helperBootReport;
+	}
+
 	/** Engine hook: an unexpected kernel death (not a deliberate kill) should drop the instance. */
 	setOnUnexpectedExit(fn: () => void): void {
 		this._onUnexpectedExit = fn;
@@ -282,7 +278,7 @@ export class KernelClient {
 	private silenceKillTimer?: ReturnType<typeof setTimeout>;
 	private pendingReplies = new Map<
 		string,
-		{ resolve(m: ParsedMessage): void; timer?: ReturnType<typeof setTimeout> }
+		{ resolve(m: ParsedMessage): void; timer?: ReturnType<typeof setTimeout>; expectedType: string }
 	>();
 
 	private constructor(conn: ConnectionFile, opts: KernelOptions) {
@@ -293,7 +289,6 @@ export class KernelClient {
 		this.timeoutMs = opts.timeoutMs ?? 0;
 	}
 
-	/** Spawn ipykernel, connect all channels, and wait until it answers. */
 	static async start(pythonPath: string, opts: KernelOptions = {}): Promise<KernelClient> {
 		const connPath = join(tmpdir(), `pi-repl-kernel-${randomUUID()}.json`);
 		const child = spawn(pythonPath, ["-m", "ipykernel", "-f", connPath, "--no-stdout"], {
@@ -325,7 +320,6 @@ export class KernelClient {
 				conn = readConnectionFile(connPath);
 				break;
 			} catch {
-				// --- not present yet, or mid-write: retry ---
 				await new Promise((resolve) => setTimeout(resolve, 25));
 			}
 		}
@@ -334,8 +328,7 @@ export class KernelClient {
 		kc.child = child;
 		kc.connectionFilePath = connPath;
 		child.on("exit", () => {
-			// --- a dead kernel settles the running cell; the engine rebuilds. clear child/ready
-			// --- so isRunning reflects death and the engine never resumes a zombie process. ---
+			// --- clear child/ready on death so isRunning reflects it and the engine never resumes a zombie ---
 			kc.settleActive(new Error("kernel process exited"));
 			kc.child = undefined;
 			kc.ready = false;
@@ -351,7 +344,7 @@ export class KernelClient {
 		try {
 			await kc.connectChannels(conn);
 			await kc.probeReady();
-			await kc.preload();
+			kc.helperBootReport = await kc.preload();
 			kc.ready = true;
 		} catch (error) {
 			kc.kill();
@@ -384,16 +377,28 @@ export class KernelClient {
 		return this.waitForReply(msgId, KERNEL_READY_TIMEOUT_MS, "kernel_info_reply").then(() => {});
 	}
 
-	/** Exec every helper file into the kernel namespace. No custom intrinsics. */
-	private preload(): Promise<void> {
-		let code = "";
-		for (const h of this.helperSources) code += `\n${h.source}\n`;
-		return this.executeCell(code, { maxOutputChars: DEFAULT_MAX_OUTPUT_CHARS }).then(() => {});
+	/** One cell per helper: a broken helper (syntax error or top-level raise) must cost itself alone, not abort the rest. */
+	private async preload(): Promise<HelperLoadResult[]> {
+		const report: HelperLoadResult[] = [];
+		for (const h of this.helperSources) {
+			const res = await this.executeCell(h.source, { maxOutputChars: DEFAULT_MAX_OUTPUT_CHARS });
+			if (res.status === "ok") {
+				report.push({ name: h.name, ok: true });
+			} else {
+				const err = res.error;
+				report.push({
+					name: h.name,
+					ok: false,
+					error: err ? `${err.name}: ${err.message}` : "failed to load",
+				});
+			}
+		}
+		return report;
 	}
 
 	private onShellMessage(frames: Buffer[]): void {
 		const msg = this.session.parseMessage(frames);
-		if (!msg) return;
+		if (!isTrustedMessage(msg)) return;
 		const active = this.activeCell;
 		if (msg.msg_type === "execute_reply" && active && msg.parent.msg_id === active.msgId) {
 			// --- the shell reply races the iopub stream: record it, settle only after idle ---
@@ -407,14 +412,14 @@ export class KernelClient {
 
 	private onControlMessage(frames: Buffer[]): void {
 		const msg = this.session.parseMessage(frames);
-		if (!msg) return;
+		if (!isTrustedMessage(msg)) return;
 		// interrupt_reply / shutdown_reply — nothing awaits them; keep draining.
 		this.resolveReply(msg);
 	}
 
 	private onIopubMessage(frames: Buffer[]): void {
 		const msg = this.session.parseMessage(frames);
-		if (!msg) return;
+		if (!isTrustedMessage(msg)) return;
 		const active = this.activeCell;
 		if (!active || msg.parent.msg_id !== active.msgId) return;
 		active.lastActivity = Date.now();
@@ -490,13 +495,15 @@ export class KernelClient {
 				reject(new Error(`kernel did not answer ${expectedType} in time`));
 			}, timeoutMs);
 			timer.unref?.();
-			this.pendingReplies.set(msgId, { resolve, timer });
+			this.pendingReplies.set(msgId, { resolve, timer, expectedType });
 		});
 	}
 
 	private resolveReply(msg: ParsedMessage): void {
 		const pending = this.pendingReplies.get(msg.parent.msg_id as string);
 		if (!pending) return;
+		// --- any reply echoes the parent id; only the awaited type settles the wait, the timer stays armed for it ---
+		if (msg.msg_type !== pending.expectedType) return;
 		this.pendingReplies.delete(msg.parent.msg_id as string);
 		if (pending.timer) clearTimeout(pending.timer);
 		pending.resolve(msg);
@@ -614,8 +621,7 @@ export class KernelClient {
 				if (quiet >= this.timeoutMs && !active.settled) {
 					active.timedOut = true;
 					this.interrupt();
-					// --- the interrupt is a real KeyboardInterrupt, but a cell that swallows/ignores
-					// --- it never replies; escalate to a kill so the queue is freed, mirroring index.ts. ---
+					// --- a cell that swallows the interrupt never replies; escalate to a kill so the queue frees ---
 					this.silenceKillTimer ??= setTimeout(() => {
 						if (!active.settled) this.kill();
 					}, SILENCE_KILL_GRACE_MS);
@@ -670,10 +676,13 @@ export class KernelClient {
 		});
 	}
 
-	restore(entries: SnapshotEntry[]): Promise<{ restored: string[]; failed: { name: string; reason: string }[] }> {
+	restore(
+		entries: SnapshotEntry[],
+		compressedValues = false,
+	): Promise<{ restored: string[]; failed: { name: string; reason: string }[] }> {
 		if (entries.length === 0) return Promise.resolve({ restored: [], failed: [] });
 		return this.enqueue(async () => {
-			const res = await this.executeCellNow(restoreCode(entries), { maxOutputChars: 8_000_000 });
+			const res = await this.executeCellNow(restoreCode(entries, compressedValues), { maxOutputChars: 8_000_000 });
 			const payload = res.payloads[RESTORE_MIME];
 			if (payload === undefined) return { restored: [], failed: [] };
 			try {

package/src/engine/session.ts

@@ -1,12 +1,8 @@
-// --- Jupyter messaging over ZMTP: [identities] <IDS|MSG> [sig, h, p, m, c] ---
-// ids are empty for a client's own channels (kernel ROUTER strips them);
-// sig = hex(HMAC-SHA256(key, h||p||m||c)) over the exact bytes; key from the
-// connection file. Checked against jupyter_client's session.py.
+// --- Jupyter over ZMTP: [<IDS|MSG>] sig h p m c; sig = hex(HMAC-SHA256(key, h||p||m||c)); ids are empty for client channels ---
 
 import { createHmac, randomUUID } from "node:crypto";
 import { readFileSync } from "node:fs";
 
-/** The kernel's connection file: ip/ports/key, written by ipykernel at boot. */
 export interface ConnectionFile {
 	ip: string;
 	transport: "tcp" | "ipc";
@@ -37,7 +33,6 @@ export interface JupyterHeader {
 	version: string;
 }
 
-/** A parsed inbound message: JSON parts + whether the signature verified. */
 export interface ParsedMessage {
 	msg_id: string;
 	msg_type: string;
@@ -52,7 +47,6 @@ function pack(obj: unknown): Buffer {
 	return Buffer.from(JSON.stringify(obj));
 }
 
-/** One client-side session: mints ids, signs and frames outbound messages. */
 export class JupyterSession {
 	readonly sessionId: string;
 	readonly username: string;
@@ -77,7 +71,6 @@ export class JupyterSession {
 		return Buffer.from(hmac.digest("hex"), "ascii");
 	}
 
-	/** Build the wire frames for an outbound message: [DELIM, sig, h, p, m, c]. */
 	buildFrames(
 		msgType: string,
 		content: Record<string, unknown>,
@@ -100,7 +93,6 @@ export class JupyterSession {
 		return [DELIM, signature, h, p, m, c];
 	}
 
-	/** Parse an inbound multipart message (identities stripped by ZMTP); null if malformed. */
 	parseMessage(frames: Buffer[]): ParsedMessage | null {
 		// --- indexOf uses ===; frames are distinct Buffers, so match by value ---
 		const delimIdx = frames.findIndex((f) => f.equals(DELIM));
@@ -123,12 +115,7 @@ export class JupyterSession {
 }
 
 export function executeRequest(code: string, silent: boolean): Record<string, unknown> {
-	// --- store_history is always false: IPython retains every last-expression result in its
-	// --- In/Out history, and that retention is NOT reclaimable from user cells (deleting Out
-	// --- entries and _/__/___ from user_ns leaves the objects alive). With history off, cells
-	// --- stop feeding that growth entirely. The contract does not depend on In/Out: results
-	// --- are published over iopub via the display hook (single-mode execution, unaffected by
-	// --- store_history) and returned in the cell's transcript. ---
+	// --- store_history off: IPython's In/Out pins every result and can't be reclaimed from cells (62MB → 400+MB); the display hook still publishes results ---
 	return {
 		code,
 		silent,
@@ -144,7 +131,11 @@ export const SNAPSHOT_MIME = "application/vnd.pi-repl.snapshot+json";
 export const RESTORE_MIME = "application/vnd.pi-repl.restore+json";
 export const NAMES_MIME = "application/vnd.pi-repl.names+json";
 
-/** Read a private-MIME payload out of an execute_result/display_data content. */
+/** Route-gate: drop malformed and unsigned traffic — anything that fails the kernel's HMAC is not the kernel. */
+export function isTrustedMessage(msg: ParsedMessage | null): msg is ParsedMessage {
+	return msg !== null && msg.signatureOk;
+}
+
 export function readPayload(content: Record<string, unknown>, mime: string): string | null {
 	const data = content.data;
 	if (data && typeof data === "object") {

package/src/engine/zmtp.ts

@@ -1,12 +1,10 @@
-// --- ZMTP 3.0 wire protocol, by hand (bun can't load libzmq's bindings). ---
-// DEALER for shell/control, SUB for iopub; greeting 0xff..0x7f + READY each side.
+// --- ZMTP 3.0 by hand (bun can't load libzmq's bindings); DEALER shell/control, SUB iopub; greeting 0xff..0x7f + READY each side ---
 
 import { connect, type Socket } from "node:net";
 
 const GREETING_SIGNATURE = Buffer.from([0xff, 0, 0, 0, 0, 0, 0, 0, 0x01, 0x7f]);
 const NULL_MECHANISM = Buffer.concat([Buffer.from("NULL"), Buffer.alloc(16)]);
 
-/** The client (non-server) half of the 64-byte ZMTP 3.0 greeting. */
 function buildGreeting(): Buffer {
 	return Buffer.concat([
 		GREETING_SIGNATURE,
@@ -21,7 +19,6 @@ const FRAME_MORE = 0x01;
 const FRAME_LONG = 0x02;
 const GREETING_LENGTH = 64;
 
-/** Serialize one ZMTP frame: flags byte, short/8-byte length, body. */
 export function encodeFrame(body: Uint8Array, more: boolean): Buffer {
 	const flags = more ? FRAME_MORE : 0;
 	if (body.length <= 255) {
@@ -44,7 +41,6 @@ export class ZmtpFrameParser {
 	private total = 0;
 	private current: Buffer[] = []; // frames of the in-progress message
 
-	/** @returns one or more complete messages consumed from `chunk`. */
 	feed(chunk: Uint8Array): Buffer[][] {
 		if (chunk.length > 0) {
 			this.chunks.push(Buffer.from(chunk));
@@ -71,7 +67,6 @@ export class ZmtpFrameParser {
 		return messages;
 	}
 
-	/** The first `n` bytes across the chunk list, without consuming them. */
 	private peekBytes(n: number): Buffer {
 		if (this.chunks[0].length >= n) return this.chunks[0].subarray(0, n);
 		const parts: Buffer[] = [];
@@ -85,7 +80,6 @@ export class ZmtpFrameParser {
 		return Buffer.concat(parts);
 	}
 
-	/** Consume `n` bytes from the front of the chunk list. */
 	private take(n: number): Buffer {
 		const first = this.chunks[0];
 		if (first.length >= n) {
@@ -119,7 +113,6 @@ export class ZmtpFrameParser {
 	}
 }
 
-/** Socket-type string carried in the READY metadata (ZMTP "Socket-Type"). */
 export type ZmtpSocketType = "DEALER" | "SUB";
 
 interface ReadReady {
@@ -127,7 +120,6 @@ interface ReadReady {
 	reject(error: Error): void;
 }
 
-/** One ZMTP client connection: TCP socket + greeting/READY handshake → complete multipart messages. */
 export class ZmtpSocket {
 	private socket?: Socket;
 	private parser = new ZmtpFrameParser();
@@ -140,7 +132,6 @@ export class ZmtpSocket {
 		this.socket = socket;
 	}
 
-	/** Connect to `host:port` and complete the ZMTP handshake for `socketType`. */
 	static connect(opts: { host: string; port: number; socketType: ZmtpSocketType }): Promise<ZmtpSocket> {
 		const socket = connect({ host: opts.host, port: opts.port });
 		const z = new ZmtpSocket(socket);
@@ -160,7 +151,6 @@ export class ZmtpSocket {
 						);
 						return;
 					}
-					// --- greeting done: announce our socket type, then await the peer's ---
 					z.send([buildReadyMetadata(opts.socketType)]);
 				}
 			}
@@ -214,14 +204,12 @@ export class ZmtpSocket {
 		this.onMessage?.(message);
 	}
 
-	/** Send a multipart message (DEALER) or a single subscription frame (SUB). */
 	send(frames: Uint8Array[]): void {
 		for (let i = 0; i < frames.length; i++) {
 			this.socket?.write(encodeFrame(frames[i], i < frames.length - 1));
 		}
 	}
 
-	/** SUB only: subscribe to a topic prefix (empty = all traffic). */
 	subscribe(topic: Uint8Array): void {
 		this.send([Buffer.concat([Buffer.from([0x01]), topic])]);
 	}

package/src/extension/helpers.ts

@@ -1,4 +1,4 @@
-/** Loads helpers from project then global dirs; `helper_description` surfaces verbatim (no signature parsing). */
+/** helper_description surfaces verbatim — there is no signature parsing. */
 
 import { existsSync, readdirSync, readFileSync } from "node:fs";
 import { homedir } from "node:os";
@@ -12,13 +12,12 @@ interface HelperEntry {
 	description: string; // full helper_description body, "" if absent
 }
 
-/** Extract `helper_description` verbatim; no signature parsing. */
 function parseDescription(source: string): string {
 	const m = source.match(/helper_description\s*=\s*("""|''')([\s\S]*?)\1/);
 	return m ? m[2].trim() : "";
 }
 
-/** Merge entries from ordered dirs; first-seen name wins, so a project helper shadows the global one. */
+/** First-seen name wins, so a project helper shadows a global one. */
 function loadHelperEntries(dirs: string[]): HelperEntry[] {
 	const seen = new Set<string>();
 	const entries: HelperEntry[] = [];
@@ -41,20 +40,32 @@ function loadHelperEntries(dirs: string[]): HelperEntry[] {
 	return entries;
 }
 
-/** The prompt-facing list for ONE dir: verbatim description, or an introspection pointer. */
+/** One helper as the model sees it: verbatim description, or a pointer when the file lacks a triple-quoted helper_description. */
+function formatHelperLine(entry: HelperEntry): string {
+	return entry.description
+		? entry.description.replace(/\n/g, "\n  ")
+		: `${entry.name} (no helper_description — define one as a triple-quoted string in the file; inspect with print(${entry.name}.__doc__))`;
+}
+
 export function buildHelpersMap(dir?: string): string[] {
-	return loadHelperEntries([dir ?? DEFAULT_HELPERS_DIR]).map((t) =>
-		t.description
-			? t.description.replace(/\n/g, "\n  ")
-			: `${t.name} (no description, inspect it with print(${t.name}.__doc__))`,
-	);
+	return loadHelperEntries([dir ?? DEFAULT_HELPERS_DIR]).map(formatHelperLine);
+}
+
+/** Project .pi/helpers first, global fallback, project shadows; honors the kernel's env overrides so the two sides can't diverge. */
+export function buildHelpersMapForCwd(cwd: string, globalDir?: string, helpersDir?: string): string[] {
+	const dirs = helpersDir ? [helpersDir] : resolveHelperDirs(cwd, globalDir);
+	return loadHelperEntries(dirs).map(formatHelperLine);
 }
 
-/** The prompt-facing list at a cwd: project .pi/helpers first (up to the git root), global fallback, project shadows. */
-export function buildHelpersMapForCwd(cwd: string, globalDir?: string): string[] {
-	return loadHelperEntries(resolveHelperDirs(cwd, globalDir)).map((t) =>
-		t.description
-			? t.description.replace(/\n/g, "\n  ")
-			: `${t.name} (no description, inspect it with print(${t.name}.__doc__))`,
-	);
+/** System-prompt block for this session's helpers; undefined when there are none. */
+export function buildHelpersPromptSection(cwd: string): string | undefined {
+	const map = buildHelpersMapForCwd(cwd, process.env.PI_HELPERS_GLOBAL_DIR, process.env.PI_HELPERS_DIR);
+	if (map.length === 0) return undefined;
+	const lines = [
+		"Preloaded helpers, use them as any loaded function or variable:",
+		...map.map((line) => `  - ${line.replace(/\n/g, "\n    ")}`),
+		"",
+		"If output begins with <repl_helpers_failed>, those helpers failed to load at boot — do not call the listed names.",
+	];
+	return lines.join("\n");
 }

package/src/extension/preview/candidates.ts

@@ -1,10 +1,6 @@
-// --- candidates: a Python-flavored file detector plus generic line scoring ---
-
 import { descriptor } from "./descriptor.js";
 import type { Candidate } from "./types.js";
 
-// --- Python file effects: verb + literal path, in the idioms this evaluator runs ---
-
 const PY_CHAINED_METHOD: ReadonlyArray<[string, string, number]> = [
 	["write_text", "write", 95],
 	["write_bytes", "write", 95],

package/src/extension/preview/descriptor.ts

@@ -1,4 +1,3 @@
-// --- descriptor: collapse the raw line into one readable, safe, width-capped string ---
 const DESCRIPTOR_MAX_WIDTH = 64;
 
 function collapseWhitespace(text: string): string {

package/src/extension/preview/types.ts

@@ -1,10 +1,8 @@
-/** A scored preview candidate for one cell in the transcript header. */
 export interface Candidate {
 	text: string;
 	score: number;
 }
 
-/** The one-line semantic summary a collapsed cell shows. */
 export interface CellPreview {
 	text: string;
 }

package/src/extension/prompt.ts

@@ -1,22 +1,18 @@
-// --- execute tool: the model-facing contract, shaped exactly like pi's built-in tools ---
-// description = rich short behavior; promptSnippet = one-liner; guidelines = flat bullets.
-
 export const executeToolDescription =
 	"Execute Python in a Jupyter notebook. Your workspace is one notebook where every cell runs in a " +
 	"shared Python environment: variables, functions, imports, classes, and data defined in one cell " +
-	"stay available to later cells for the life of the notebook. Cells return their last expression " +
-	"(auto-displayed) plus stdout/stderr; output is truncated to 45K with a marker.";
+	"stay available and reuseable for later cells for the life of the notebook. " +
+	"Output is truncated to 45K with an explicit marker ([... output truncated at 46080 chars ...]).";
 
-export const executePromptSnippet = "Run Python cells in a Jupyter notebook (read, write, run, search, and more)";
+export const executePromptSnippet = "Run Python in a Jupyter notebook (read, write, run, search, and more)";
 
-// --- the model-facing guidelines, flat bullets like pi's own tool contributions ---
 export function buildPromptGuidelines(preloaded: string[]): string[] {
 	return [
 		"Write modern idiomatic Python.",
-		"Find, filter, fetch, sample: narrow the output in Python, then print only the window that decides the next step (a head, a shape, a slice), not the whole.",
-		"Make surgical, precise changes over rewrites or whole-file dumps: replace, verify, read the file back before trusting it.",
-		"Prefer to reuse existing variables, functions, imports, classes, and data from prior cells/namespaces over recomputing.",
-		"If output begins with <repl_engine_reset>, the runtime rebuilt and the notebook was restored from the last snapshot; reverify surviving names before building on them.",
+		"Find, Filter, Fetch: narrow down the output in Python and only print the exact slice you'll ever need, be scrupulous.",
+		"Make surgical, precise, atomic changes over rewrites or whole-file dumps: read the exact lines FIRST, anchor them, replace, verify, read the file back before trusting it",
+		"AMORTIZE existing variables, functions, imports, classes, and data from prior cells/namespaces over re-deriving, recomputing or reconstructing.",
+		"If output begins with <repl_engine_reset>, the runtime rebuilt and the notebook was restored from a recent snapshot; reverify surviving states before building on them.",
 		...(preloaded.length
 			? [
 					[