@wahack/pi-coding-agent 15.11.0

package/src/eval/js/context-manager.ts

@@ -0,0 +1,502 @@
+import { logger, Snowflake, workerHostEntry } from "@oh-my-pi/pi-utils";
+import type { ToolSession } from "../../tools";
+import { ToolAbortError, ToolError } from "../../tools/tool-errors";
+import { callSessionTool, type JsStatusEvent } from "./tool-bridge";
+import { WorkerCore } from "./worker-core";
+// Coding-agent binary/bundle workers route through the CLI entrypoint with a
+// hidden argv mode, so compiled/npm builds only need one JavaScript entry.
+import type {
+	JsDisplayOutput,
+	RunErrorPayload,
+	SessionSnapshot,
+	Transport,
+	WorkerInbound,
+	WorkerOutbound,
+} from "./worker-protocol";
+
+export { rewriteImports, wrapCode } from "./shared/rewrite-imports";
+export type { JsDisplayOutput } from "./worker-protocol";
+
+export interface VmRunState {
+	signal?: AbortSignal;
+	onText?: (chunk: string) => void;
+	onDisplay?: (output: JsDisplayOutput) => void;
+}
+
+interface WorkerHandle {
+	mode: "worker" | "inline";
+	send(msg: WorkerInbound): void;
+	onMessage(handler: (msg: WorkerOutbound) => void): () => void;
+	close(): Promise<boolean>;
+	terminate(): Promise<void>;
+}
+
+interface PendingRun {
+	runId: string;
+	runState: VmRunState;
+	toolSession: ToolSession;
+	resolve(value: { value: unknown }): void;
+	reject(error: Error): void;
+	toolCalls: Map<string, AbortController>;
+	settled: boolean;
+}
+
+interface JsSession {
+	sessionKey: string;
+	worker: WorkerHandle;
+	state: "alive" | "dead";
+	pending: Map<string, PendingRun>;
+}
+
+const sessions = new Map<string, JsSession>();
+const startingSessions = new Map<string, Promise<JsSession>>();
+const resettingSessions = new Map<string, Promise<void>>();
+// Worker startup (module-graph import + WorkerCore construction) is infrastructure
+// cost, not user compute. Floor it independently of Bun's 5s default per-test timeout
+// so a slow cold-start under load isn't aborted mid-init — terminating a still-
+// initializing Bun worker triggers the same kind of terminate-race that motivates
+// avoiding `vm.runInContext` (see shared/indirect-eval.ts), here surfacing as a
+// SIGILL/SIGSEGV. Callers that pass a larger per-cell budget still dominate.
+const WORKER_INIT_TIMEOUT_MS = 15_000;
+const WORKER_CLOSE_TIMEOUT_MS = 1_000;
+
+export async function executeInVmContext(options: {
+	sessionKey: string;
+	sessionId: string;
+	cwd: string;
+	session: ToolSession;
+	localRoots?: Record<string, string>;
+	reset?: boolean;
+	code: string;
+	filename: string;
+	timeoutMs?: number;
+	runState: VmRunState;
+}): Promise<{ value: unknown }> {
+	if (options.reset) {
+		// Coalesce concurrent resets: an existing in-flight reset already
+		// produces a fresh context, so a follow-up `reset: true` cell should
+		// just wait for it rather than failing the user-visible call.
+		const inFlight = resettingSessions.get(options.sessionKey);
+		if (inFlight) await inFlight.catch(() => undefined);
+		else {
+			const resetPromise = resetVmContext(options.sessionKey);
+			resettingSessions.set(
+				options.sessionKey,
+				resetPromise.then(() => undefined),
+			);
+			try {
+				await resetPromise;
+			} finally {
+				resettingSessions.delete(options.sessionKey);
+			}
+		}
+	} else {
+		// Internal coordination: wait for any in-flight reset to settle and
+		// then run on the freshly-rebuilt context.
+		const inFlight = resettingSessions.get(options.sessionKey);
+		if (inFlight) await inFlight.catch(() => undefined);
+	}
+	const session = await acquireSession(
+		options.sessionKey,
+		{ cwd: options.cwd, sessionId: options.sessionId, localRoots: options.localRoots },
+		options.timeoutMs,
+	);
+	return await runOnce(session, options);
+}
+
+export async function resetVmContext(sessionKey: string): Promise<void> {
+	const session = sessions.get(sessionKey) ?? (await startingSessions.get(sessionKey)?.catch(() => undefined));
+	if (!session) return;
+	sessions.delete(sessionKey);
+	await killSession(session, new ToolError("JS context reset"), { force: false });
+}
+
+export async function disposeAllVmContexts(): Promise<void> {
+	const pending = [...startingSessions.values()];
+	startingSessions.clear();
+	const started = await Promise.allSettled(pending);
+	const all = [...sessions.values()];
+	for (const result of started) {
+		if (result.status !== "fulfilled") continue;
+		if (!all.includes(result.value)) all.push(result.value);
+	}
+	sessions.clear();
+	await Promise.all(all.map(session => killSession(session, new ToolError("JS context disposed"), { force: false })));
+}
+
+async function runOnce(
+	session: JsSession,
+	options: {
+		sessionId: string;
+		cwd: string;
+		session: ToolSession;
+		localRoots?: Record<string, string>;
+		code: string;
+		filename: string;
+		runState: VmRunState;
+	},
+): Promise<{ value: unknown }> {
+	const runId = `r-${Snowflake.next()}`;
+	const { promise, resolve, reject } = Promise.withResolvers<{ value: unknown }>();
+	const pending: PendingRun = {
+		runId,
+		runState: options.runState,
+		toolSession: options.session,
+		resolve,
+		reject,
+		toolCalls: new Map(),
+		settled: false,
+	};
+	session.pending.set(runId, pending);
+
+	const onAbort = (): void => {
+		const reason = options.runState.signal?.reason;
+		const abortError = reasonToError(reason, "Execution aborted");
+		// Cancel any in-flight tool calls first.
+		for (const ctrl of pending.toolCalls.values()) ctrl.abort(abortError);
+		// Hard-kill the worker — only way to interrupt synchronous user code.
+		void killSessionFor(session, abortError, { force: true });
+	};
+
+	if (options.runState.signal?.aborted) {
+		queueMicrotask(onAbort);
+	} else {
+		options.runState.signal?.addEventListener("abort", onAbort, { once: true });
+	}
+
+	try {
+		session.worker.send({
+			type: "run",
+			runId,
+			code: options.code,
+			filename: options.filename,
+			snapshot: { cwd: options.cwd, sessionId: options.sessionId, localRoots: options.localRoots },
+		});
+		return await promise;
+	} finally {
+		options.runState.signal?.removeEventListener("abort", onAbort);
+		session.pending.delete(runId);
+	}
+}
+
+async function acquireSession(sessionKey: string, snapshot: SessionSnapshot, timeoutMs?: number): Promise<JsSession> {
+	const existing = sessions.get(sessionKey);
+	if (existing && existing.state === "alive") return existing;
+	const starting = startingSessions.get(sessionKey);
+	if (starting) return await starting;
+
+	const startup = (async (): Promise<JsSession> => {
+		const worker = await spawnJsWorker();
+		const session: JsSession = {
+			sessionKey,
+			worker,
+			state: "alive",
+			pending: new Map(),
+		};
+		const { promise: readyPromise, resolve: resolveReady, reject: rejectReady } = Promise.withResolvers<void>();
+		let resolved = false;
+		const unsubscribe = worker.onMessage(msg => {
+			if (!resolved && msg.type === "ready") {
+				resolved = true;
+				resolveReady();
+				return;
+			}
+			if (!resolved && msg.type === "init-failed") {
+				resolved = true;
+				rejectReady(errorFromPayload(msg.error));
+				return;
+			}
+			handleSessionMessage(session, msg);
+		});
+		try {
+			// Init headroom is the fixed infrastructure floor; the caller's per-cell timeout
+			// dominates when larger so users can grant more by raising `timeout` on a cell.
+			const readyTimeoutMs = Math.max(WORKER_INIT_TIMEOUT_MS, timeoutMs ?? 0);
+			await raceWithTimeout(readyPromise, readyTimeoutMs, "Timed out initializing JS eval worker");
+			worker.send({ type: "init", snapshot });
+			sessions.set(sessionKey, session);
+			return session;
+		} catch (error) {
+			unsubscribe();
+			await worker.terminate().catch(() => undefined);
+			throw error;
+		}
+	})();
+	startingSessions.set(sessionKey, startup);
+	try {
+		return await startup;
+	} finally {
+		if (startingSessions.get(sessionKey) === startup) startingSessions.delete(sessionKey);
+	}
+}
+
+function handleSessionMessage(session: JsSession, msg: WorkerOutbound): void {
+	switch (msg.type) {
+		case "text": {
+			const pending = session.pending.get(msg.runId);
+			pending?.runState.onText?.(msg.chunk);
+			return;
+		}
+		case "display": {
+			const pending = session.pending.get(msg.runId);
+			pending?.runState.onDisplay?.(msg.output);
+			return;
+		}
+		case "tool-call":
+			void handleToolCall(session, msg);
+			return;
+		case "result":
+			settlePending(session, msg);
+			return;
+		case "log":
+			logWorkerMessage(msg);
+			return;
+		case "ready":
+		case "init-failed":
+		case "closed":
+			return;
+	}
+}
+
+async function handleToolCall(session: JsSession, msg: Extract<WorkerOutbound, { type: "tool-call" }>): Promise<void> {
+	const pending = session.pending.get(msg.runId);
+	if (!pending) {
+		safeSend(session, {
+			type: "tool-reply",
+			id: msg.id,
+			reply: { ok: false, error: { message: "Run no longer active" } },
+		});
+		return;
+	}
+	const ctrl = new AbortController();
+	pending.toolCalls.set(msg.id, ctrl);
+	try {
+		const value = await callSessionTool(msg.name, msg.args, {
+			session: pending.toolSession,
+			signal: ctrl.signal,
+			emitStatus: (event: JsStatusEvent) => pending.runState.onDisplay?.({ type: "status", event }),
+		});
+		safeSend(session, { type: "tool-reply", id: msg.id, reply: { ok: true, value } });
+	} catch (error) {
+		safeSend(session, { type: "tool-reply", id: msg.id, reply: { ok: false, error: toErrorPayload(error) } });
+	} finally {
+		pending.toolCalls.delete(msg.id);
+	}
+}
+
+function settlePending(session: JsSession, msg: Extract<WorkerOutbound, { type: "result" }>): void {
+	const pending = session.pending.get(msg.runId);
+	if (!pending || pending.settled) return;
+	pending.settled = true;
+	if (msg.ok) {
+		pending.resolve({ value: undefined });
+		return;
+	}
+	pending.reject(errorFromPayload(msg.error));
+}
+
+async function killSessionFor(session: JsSession, error: Error, options: { force: boolean }): Promise<void> {
+	if (sessions.get(session.sessionKey) === session) {
+		sessions.delete(session.sessionKey);
+	}
+	await killSession(session, error, options);
+}
+
+async function killSession(session: JsSession, error: Error, options: { force: boolean }): Promise<void> {
+	if (session.state === "dead") return;
+	session.state = "dead";
+	for (const pending of session.pending.values()) {
+		if (pending.settled) continue;
+		pending.settled = true;
+		for (const ctrl of pending.toolCalls.values()) ctrl.abort(error);
+		pending.reject(error);
+	}
+	session.pending.clear();
+	if (options.force) {
+		await session.worker.terminate().catch(() => undefined);
+		return;
+	}
+	if (await session.worker.close().catch(() => false)) return;
+	await session.worker.terminate().catch(() => undefined);
+}
+
+function safeSend(session: JsSession, msg: WorkerInbound): void {
+	if (session.state !== "alive") return;
+	try {
+		session.worker.send(msg);
+	} catch (err) {
+		logger.debug("js worker send failed", { error: err instanceof Error ? err.message : String(err) });
+	}
+}
+
+function reasonToError(reason: unknown, fallback: string): Error {
+	if (reason instanceof Error) return reason;
+	if (typeof reason === "string") return new ToolAbortError(reason);
+	return new ToolAbortError(fallback);
+}
+
+function errorFromPayload(payload: RunErrorPayload): Error {
+	if (payload.isAbort) {
+		const err = new ToolAbortError(payload.message || "Execution aborted");
+		if (payload.stack) err.stack = payload.stack;
+		return err;
+	}
+	const ctor = payload.isToolError ? ToolError : Error;
+	const error = new ctor(payload.message);
+	if (payload.name) error.name = payload.name;
+	if (payload.stack) error.stack = payload.stack;
+	return error;
+}
+
+function toErrorPayload(error: unknown): RunErrorPayload {
+	if (error instanceof Error) {
+		return {
+			name: error.name,
+			message: error.message,
+			stack: error.stack,
+			isAbort: error.name === "AbortError" || error.name === "ToolAbortError",
+			isToolError: error instanceof ToolError || error.name === "ToolError",
+		};
+	}
+	return { message: String(error) };
+}
+
+function logWorkerMessage(msg: Extract<WorkerOutbound, { type: "log" }>): void {
+	if (msg.level === "debug") logger.debug(msg.msg, msg.meta);
+	else if (msg.level === "warn") logger.warn(msg.msg, msg.meta);
+	else logger.error(msg.msg, msg.meta);
+}
+
+async function raceWithTimeout<T>(promise: Promise<T>, timeoutMs: number, reason: string): Promise<T> {
+	const timeoutSignal = AbortSignal.timeout(timeoutMs);
+	const { promise: timeoutPromise, reject } = Promise.withResolvers<never>();
+	const onAbort = (): void => reject(new ToolError(reason));
+	timeoutSignal.addEventListener("abort", onAbort, { once: true });
+	try {
+		return await Promise.race([promise, timeoutPromise]);
+	} finally {
+		timeoutSignal.removeEventListener("abort", onAbort);
+	}
+}
+
+async function spawnJsWorker(): Promise<WorkerHandle> {
+	try {
+		const hostEntry = workerHostEntry();
+		const worker = hostEntry
+			? new Worker(hostEntry, { type: "module", argv: ["__omp_js_eval_worker"] })
+			: new Worker(new URL("./worker-entry.ts", import.meta.url).href, { type: "module" });
+		return wrapBunWorker(worker);
+	} catch (err) {
+		logger.warn("Bun Worker spawn failed; using inline JS eval worker (no sync-loop guard)", {
+			error: err instanceof Error ? err.message : String(err),
+		});
+		return spawnInlineWorker();
+	}
+}
+
+function wrapBunWorker(worker: Worker): WorkerHandle {
+	return {
+		mode: "worker",
+		send(msg) {
+			worker.postMessage(msg);
+		},
+		onMessage(handler) {
+			const wrap = (event: MessageEvent): void => handler(event.data as WorkerOutbound);
+			worker.addEventListener("message", wrap);
+			return () => worker.removeEventListener("message", wrap);
+		},
+		async close() {
+			const { promise: closed, resolve } = Promise.withResolvers<boolean>();
+			let settled = false;
+			let sawClosedAck = false;
+			let sawWorkerExit = false;
+			let timeout: NodeJS.Timeout | undefined;
+			let unsubscribe = (): void => {};
+			const finish = (value: boolean): void => {
+				if (settled) return;
+				settled = true;
+				if (timeout) clearTimeout(timeout);
+				unsubscribe();
+				worker.removeEventListener("close", onClose);
+				resolve(value);
+			};
+			const finishIfClosed = (): void => {
+				if (sawClosedAck && sawWorkerExit) finish(true);
+			};
+			const onClose = (): void => {
+				sawWorkerExit = true;
+				finishIfClosed();
+			};
+			unsubscribe = this.onMessage(msg => {
+				if (msg.type !== "closed") return;
+				sawClosedAck = true;
+				finishIfClosed();
+			});
+			worker.addEventListener("close", onClose);
+			timeout = setTimeout(() => finish(false), WORKER_CLOSE_TIMEOUT_MS);
+			worker.postMessage({ type: "close" } satisfies WorkerInbound);
+			return await closed;
+		},
+		async terminate() {
+			worker.terminate();
+		},
+	};
+}
+
+/**
+ * Inline fallback for environments where Bun cannot spawn the worker entry
+ * (e.g. some test runners). Preserves behavior but cannot interrupt synchronous
+ * infinite loops because user code runs on the main thread.
+ */
+function spawnInlineWorker(): WorkerHandle {
+	const hostListeners = new Set<(message: WorkerOutbound) => void>();
+	const workerListeners = new Set<(message: WorkerInbound) => void>();
+	const workerTransport: Transport = {
+		send: msg =>
+			queueMicrotask(() => {
+				for (const listener of hostListeners) listener(msg);
+			}),
+		onMessage: handler => {
+			workerListeners.add(handler);
+			return () => workerListeners.delete(handler);
+		},
+		close: () => {},
+	};
+	new WorkerCore(workerTransport);
+	return {
+		mode: "inline",
+		send: msg =>
+			queueMicrotask(() => {
+				for (const listener of workerListeners) listener(msg);
+			}),
+		onMessage: handler => {
+			hostListeners.add(handler);
+			return () => hostListeners.delete(handler);
+		},
+		async close() {
+			const { promise: closed, resolve } = Promise.withResolvers<boolean>();
+			let settled = false;
+			let timeout: NodeJS.Timeout | undefined;
+			let unsubscribe = (): void => {};
+			const finish = (value: boolean): void => {
+				if (settled) return;
+				settled = true;
+				if (timeout) clearTimeout(timeout);
+				unsubscribe();
+				hostListeners.clear();
+				workerListeners.clear();
+				resolve(value);
+			};
+			unsubscribe = this.onMessage(msg => {
+				if (msg.type === "closed") finish(true);
+			});
+			this.send({ type: "close" });
+			timeout = setTimeout(() => finish(false), WORKER_CLOSE_TIMEOUT_MS);
+			return await closed;
+		},
+		async terminate() {
+			hostListeners.clear();
+			workerListeners.clear();
+		},
+	};
+}

package/src/eval/js/executor.ts

@@ -0,0 +1,173 @@
+import { DEFAULT_MAX_BYTES, OutputSink } from "../../session/streaming-output";
+import type { ToolSession } from "../../tools";
+import { resolveOutputMaxColumns, resolveOutputSinkHeadBytes } from "../../tools/output-meta";
+import { isEvalTimeoutControlEvent } from "../bridge-timeout";
+import { executeInVmContext, type JsDisplayOutput } from "./context-manager";
+import type { JsStatusEvent } from "./shared/types";
+
+export interface JsExecutorOptions {
+	cwd?: string;
+	timeoutMs?: number;
+	deadlineMs?: number;
+	/**
+	 * Runtime-work budget (ms). Used for worker cold-start headroom and
+	 * timeout-annotation text when the caller drives cancellation via the eval
+	 * watchdog `signal` instead of `deadlineMs`/`timeoutMs`. Never arms a timer.
+	 */
+	idleTimeoutMs?: number;
+	onChunk?: (chunk: string) => Promise<void> | void;
+	onStatus?: (event: JsStatusEvent) => void;
+	signal?: AbortSignal;
+	sessionId: string;
+	reset?: boolean;
+	sessionFile?: string;
+	artifactPath?: string;
+	artifactId?: string;
+	session: ToolSession;
+	/** On-disk roots the helpers substitute for internal-URL schemes (e.g. `local://`). */
+	localRoots?: Record<string, string>;
+}
+
+export interface JsResult {
+	output: string;
+	exitCode: number | undefined;
+	cancelled: boolean;
+	truncated: boolean;
+	artifactId?: string;
+	totalLines: number;
+	totalBytes: number;
+	outputLines: number;
+	outputBytes: number;
+	displayOutputs: JsDisplayOutput[];
+}
+
+function getExecutionTimeoutMs(options: Pick<JsExecutorOptions, "deadlineMs" | "timeoutMs">): number | undefined {
+	if (options.deadlineMs !== undefined) {
+		return Math.max(1, options.deadlineMs - Date.now());
+	}
+	return options.timeoutMs;
+}
+
+function isAbortError(error: unknown): boolean {
+	return (
+		(error instanceof DOMException && (error.name === "AbortError" || error.name === "TimeoutError")) ||
+		(error instanceof Error && (error.name === "AbortError" || error.name === "TimeoutError"))
+	);
+}
+
+function isTimeoutReason(reason: unknown): boolean {
+	return (
+		(reason instanceof DOMException && reason.name === "TimeoutError") ||
+		(reason instanceof Error && reason.name === "TimeoutError")
+	);
+}
+
+function formatJsTimeoutAnnotation(timeoutMs: number | undefined): string {
+	// Timeout cancellation force-kills the worker (the only way to interrupt
+	// synchronous user code), which discards the persistent VM state. Say so,
+	// or the model will keep referencing variables that no longer exist.
+	const reset = "The JS worker was force-killed and its VM state was reset; variables from earlier cells are gone.";
+	if (timeoutMs === undefined) return `Command timed out. ${reset}`;
+	const secs = Math.max(1, Math.round(timeoutMs / 1000));
+	return `Command timed out after ${secs} seconds. ${reset}`;
+}
+
+export async function executeJs(code: string, options: JsExecutorOptions): Promise<JsResult> {
+	const displayOutputs: JsDisplayOutput[] = [];
+	const outputSink = new OutputSink({
+		artifactPath: options.artifactPath,
+		artifactId: options.artifactId,
+		spillThreshold: DEFAULT_MAX_BYTES,
+		headBytes: resolveOutputSinkHeadBytes(options.session.settings),
+		maxColumns: resolveOutputMaxColumns(options.session.settings),
+		onChunk: chunk => options.onChunk?.(chunk),
+	});
+	const legacyTimeoutMs = getExecutionTimeoutMs(options);
+	const timeoutSignal =
+		typeof legacyTimeoutMs === "number" && Number.isFinite(legacyTimeoutMs) && legacyTimeoutMs > 0
+			? AbortSignal.timeout(legacyTimeoutMs)
+			: undefined;
+	const signal =
+		options.signal && timeoutSignal
+			? AbortSignal.any([options.signal, timeoutSignal])
+			: (options.signal ?? timeoutSignal);
+	// The eval tool drives cancellation via its own watchdog `signal` and passes
+	// only the runtime-work budget; use it solely as worker cold-start headroom
+	// and never derive a competing fixed timer from it.
+	const acquireBudgetMs = legacyTimeoutMs ?? options.idleTimeoutMs;
+
+	try {
+		await executeInVmContext({
+			sessionKey: options.sessionId,
+			sessionId: options.sessionId,
+			cwd: options.cwd ?? options.session.cwd,
+			session: options.session,
+			localRoots: options.localRoots,
+			reset: options.reset,
+			code,
+			filename: `js-cell-${crypto.randomUUID()}.js`,
+			timeoutMs: acquireBudgetMs,
+			runState: {
+				signal,
+				onText: chunk => outputSink.push(chunk),
+				onDisplay: output => {
+					if (output.type === "status") {
+						// Timeout-control events drive the eval watchdog only; never
+						// store or render them as cell output.
+						options.onStatus?.(output.event);
+						if (isEvalTimeoutControlEvent(output.event)) return;
+					}
+					displayOutputs.push(output);
+				},
+			},
+		});
+		const summary = await outputSink.dump();
+		return {
+			output: summary.output,
+			exitCode: 0,
+			cancelled: false,
+			truncated: summary.truncated,
+			artifactId: summary.artifactId,
+			totalLines: summary.totalLines,
+			totalBytes: summary.totalBytes,
+			outputLines: summary.outputLines,
+			outputBytes: summary.outputBytes,
+			displayOutputs,
+		};
+	} catch (error) {
+		if (signal?.aborted || isAbortError(error)) {
+			const timedOut = Boolean(timeoutSignal?.aborted) || isTimeoutReason(options.signal?.reason);
+			if (timedOut) {
+				outputSink.push(formatJsTimeoutAnnotation(legacyTimeoutMs ?? options.idleTimeoutMs));
+			}
+			const summary = await outputSink.dump();
+			return {
+				output: summary.output,
+				exitCode: undefined,
+				cancelled: true,
+				truncated: summary.truncated,
+				artifactId: summary.artifactId,
+				totalLines: summary.totalLines,
+				totalBytes: summary.totalBytes,
+				outputLines: summary.outputLines,
+				outputBytes: summary.outputBytes,
+				displayOutputs,
+			};
+		}
+		const message = error instanceof Error ? (error.stack ?? error.message) : String(error);
+		outputSink.push(message);
+		const summary = await outputSink.dump();
+		return {
+			output: summary.output,
+			exitCode: 1,
+			cancelled: false,
+			truncated: summary.truncated,
+			artifactId: summary.artifactId,
+			totalLines: summary.totalLines,
+			totalBytes: summary.totalBytes,
+			outputLines: summary.outputLines,
+			outputBytes: summary.outputBytes,
+			displayOutputs,
+		};
+	}
+}