@oh-my-pi/pi-coding-agent 17.2.4 → 17.2.5

package/src/tools/computer/supervisor.ts

@@ -1,25 +1,38 @@
-import type { DesktopAction, DesktopCapabilities, DesktopCapture, DesktopSessionOptions } from "@oh-my-pi/pi-natives";
+import type { DesktopCapabilities } from "@oh-my-pi/pi-natives";
 import { withTimeout } from "@oh-my-pi/pi-utils/async";
 import * as logger from "@oh-my-pi/pi-utils/logger";
+import { Snowflake } from "@oh-my-pi/pi-utils/snowflake";
 import { workerHostEntry } from "@oh-my-pi/pi-utils/worker-host";
+import type { ToolSession } from "../index";
 import { ToolAbortError, ToolError } from "../tool-errors";
 import {
 	COMPUTER_WORKER_ARG,
-	type ComputerWorkerError,
+	type ComputerRunOk,
+	type ComputerSessionSnapshot,
 	type ComputerWorkerInbound,
 	type ComputerWorkerOutbound,
+	type RunErrorPayload,
 } from "./protocol";
 
 const START_TIMEOUT_MS = 10_000;
 const CLOSE_TIMEOUT_MS = 1_500;
+const GRACE_MS = 750;
 const SMOKE_TIMEOUT_MS = 5_000;
+const RESTART_MESSAGE = "computer worker restarted; captures and ax refs were reset";
 
+/** Runs desktop scripts and owns their persistent worker session. */
 export interface ComputerController {
-	readonly capabilities: DesktopCapabilities | undefined;
-	execute(actions: DesktopAction[], signal?: AbortSignal): Promise<DesktopCapture>;
+	run(
+		code: string,
+		timeoutMs: number,
+		snapshot: ComputerSessionSnapshot,
+		signal?: AbortSignal,
+	): Promise<ComputerRunOk>;
+	capabilities(): Promise<DesktopCapabilities | undefined>;
 	close(): Promise<void>;
 }
 
+/** Minimal Bun worker lifecycle surface used by the supervisor. */
 export interface ComputerWorkerHandle {
 	send(message: ComputerWorkerInbound): void;
 	onMessage(handler: (message: ComputerWorkerOutbound) => void): () => void;
@@ -27,6 +40,7 @@ export interface ComputerWorkerHandle {
 	terminate(): Promise<void>;
 }
 
+/** Startup and shutdown deadlines for a computer worker. */
 export interface ComputerSupervisorTimeouts {
 	startMs: number;
 	closeMs: number;
@@ -36,13 +50,22 @@ const DEFAULT_TIMEOUTS: ComputerSupervisorTimeouts = {
 	startMs: START_TIMEOUT_MS,
 	closeMs: CLOSE_TIMEOUT_MS,
 };
+
+/** Dispatches a tool call requested from desktop JavaScript. */
+export type ComputerSessionToolCaller = (
+	name: string,
+	args: unknown,
+	options: { session: ToolSession; signal?: AbortSignal; emitStatus?: () => void },
+) => Promise<unknown>;
+
+/** Creates an isolated computer worker handle. */
 export type ComputerWorkerFactory = () => ComputerWorkerHandle;
 
-function workerError(error: ComputerWorkerError): Error {
-	const result = new ToolError(error.message);
-	result.name = error.name;
-	if (error.stack) result.stack = error.stack;
-	return result;
+interface PendingRun {
+	resolve(value: ComputerRunOk): void;
+	reject(error: unknown): void;
+	signal?: AbortSignal;
+	toolCalls: Map<string, AbortController>;
 }
 
 function wrapWorker(worker: Worker): ComputerWorkerHandle {
@@ -56,10 +79,19 @@ function wrapWorker(worker: Worker): ComputerWorkerHandle {
 			return () => worker.removeEventListener("message", listener);
 		},
 		onError(handler) {
-			const listener = (event: ErrorEvent): void =>
+			const onError = (event: ErrorEvent): void =>
 				handler(event.error instanceof Error ? event.error : new Error(event.message));
-			worker.addEventListener("error", listener);
-			return () => worker.removeEventListener("error", listener);
+			const onMessageError = (event: MessageEvent): void =>
+				handler(new Error(`Computer worker message error: ${String(event.data)}`));
+			const onClose = (): void => handler(new Error("Computer worker exited"));
+			worker.addEventListener("error", onError);
+			worker.addEventListener("messageerror", onMessageError);
+			worker.addEventListener("close", onClose);
+			return () => {
+				worker.removeEventListener("error", onError);
+				worker.removeEventListener("messageerror", onMessageError);
+				worker.removeEventListener("close", onClose);
+			};
 		},
 		async terminate() {
 			worker.terminate();
@@ -67,6 +99,7 @@ function wrapWorker(worker: Worker): ComputerWorkerHandle {
 	};
 }
 
+/** Spawns the computer worker through the active CLI host when available. */
 export function spawnComputerWorker(): ComputerWorkerHandle {
 	const hostEntry = workerHostEntry();
 	const worker = hostEntry
@@ -75,133 +108,116 @@ export function spawnComputerWorker(): ComputerWorkerHandle {
 	return wrapWorker(worker);
 }
 
-interface PendingRequest {
-	resolve(capture: DesktopCapture): void;
-	reject(error: unknown): void;
+function errorFromPayload(payload: RunErrorPayload): Error {
+	const error = payload.isAbort
+		? new ToolAbortError(payload.message)
+		: payload.isToolError
+			? new ToolError(payload.message)
+			: new Error(payload.message);
+	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 { name: "Error", message: String(error), isAbort: false, isToolError: false };
 }
 
+/** Supervises one lazy, crash-isolated computer worker per agent session. */
 export class ComputerSupervisor implements ComputerController {
+	readonly #session: ToolSession;
+	readonly #createWorker: ComputerWorkerFactory;
+	readonly #callSessionTool: ComputerSessionToolCaller;
+	readonly #timeouts: ComputerSupervisorTimeouts;
 	#worker?: ComputerWorkerHandle;
 	#startPromise?: Promise<void>;
-	#capabilities?: DesktopCapabilities;
-	#pending = new Map<string, PendingRequest>();
+	#startReject?: (error: unknown) => void;
+	#startResolve?: () => void;
+	#latestCapabilities?: DesktopCapabilities;
+	#pending = new Map<string, PendingRun>();
 	#nextId = 0;
-	#serial: Promise<void> = Promise.resolve();
 	#closed = false;
 	#unsubscribeMessage?: () => void;
 	#unsubscribeError?: () => void;
 
 	constructor(
-		private readonly options: DesktopSessionOptions,
-		private readonly createWorker: ComputerWorkerFactory = spawnComputerWorker,
-		private readonly timeouts: ComputerSupervisorTimeouts = DEFAULT_TIMEOUTS,
-	) {}
-
-	get capabilities(): DesktopCapabilities | undefined {
-		return this.#capabilities;
+		session: ToolSession,
+		createWorker: ComputerWorkerFactory = spawnComputerWorker,
+		timeouts: ComputerSupervisorTimeouts = DEFAULT_TIMEOUTS,
+		callSessionTool: ComputerSessionToolCaller = async () => {
+			throw new ToolError("Computer session tool bridge is unavailable");
+		},
+	) {
+		this.#session = session;
+		this.#createWorker = createWorker;
+		this.#timeouts = timeouts;
+		this.#callSessionTool = callSessionTool;
 	}
 
-	execute(actions: DesktopAction[], signal?: AbortSignal): Promise<DesktopCapture> {
-		if (this.#closed) return Promise.reject(new ToolError("Computer session is closed"));
-		const run = (): Promise<DesktopCapture> => this.#execute(actions, signal);
-		const result = this.#serial.then(run, run);
-		this.#serial = result.then(
-			() => undefined,
-			() => undefined,
-		);
-		return result;
+	async capabilities(): Promise<DesktopCapabilities | undefined> {
+		return this.#latestCapabilities;
 	}
 
-	async #execute(actions: DesktopAction[], signal?: AbortSignal): Promise<DesktopCapture> {
+	async run(
+		code: string,
+		timeoutMs: number,
+		snapshot: ComputerSessionSnapshot,
+		signal?: AbortSignal,
+	): Promise<ComputerRunOk> {
 		if (this.#closed) throw new ToolError("Computer session is closed");
 		if (signal?.aborted) throw new ToolAbortError();
 		await this.#start();
 		if (signal?.aborted) throw new ToolAbortError();
+
 		const id = `computer-${++this.#nextId}`;
-		const request = Promise.withResolvers<DesktopCapture>();
-		this.#pending.set(id, request);
-		this.#worker!.send({ type: "execute", id, actions });
-		if (!signal) return request.promise;
+		const { promise, resolve, reject } = Promise.withResolvers<ComputerRunOk>();
+		const pending: PendingRun = { resolve, reject, signal, toolCalls: new Map() };
+		this.#pending.set(id, pending);
+		const abort = (): void => {
+			this.#safeSend({ type: "abort", id });
+			for (const controller of pending.toolCalls.values()) controller.abort(signal?.reason);
+		};
+		if (signal?.aborted) abort();
+		else signal?.addEventListener("abort", abort, { once: true });
 
-		const aborted = Promise.withResolvers<never>();
-		const onAbort = (): void => aborted.reject(new ToolAbortError());
-		signal.addEventListener("abort", onAbort, { once: true });
 		try {
-			return await Promise.race([request.promise, aborted.promise]);
-		} catch (error) {
-			if (error instanceof ToolAbortError) await this.#terminate(error);
-			throw error;
+			this.#worker?.send({ type: "run", id, code, timeoutMs, session: snapshot });
+			return await this.#raceWithGrace(promise, timeoutMs);
 		} finally {
-			signal.removeEventListener("abort", onAbort);
+			signal?.removeEventListener("abort", abort);
+			this.#pending.delete(id);
 		}
 	}
 
 	#start(): Promise<void> {
 		if (this.#startPromise) return this.#startPromise;
-		const ready = Promise.withResolvers<void>();
+		const started = Promise.withResolvers<void>();
+		this.#startReject = started.reject;
+		this.#startResolve = started.resolve;
 		try {
-			logger.debug("Starting native computer worker", {
-				backend: this.options.backend,
-				display: this.options.display,
-				maxWidth: this.options.maxWidth,
-				maxHeight: this.options.maxHeight,
-			});
-			const worker = this.createWorker();
+			logger.debug("Starting computer worker");
+			const worker = this.#createWorker();
 			this.#worker = worker;
-			this.#unsubscribeMessage = worker.onMessage(message => {
-				if (message.type === "ready") {
-					this.#capabilities = message.capabilities;
-					logger.debug("Native computer worker ready", {
-						backend: message.capabilities.backend,
-						capturePermission: message.capabilities.capturePermission,
-						inputPermission: message.capabilities.inputPermission,
-						displayCount: message.capabilities.displayCount,
-					});
-					ready.resolve();
-					return;
-				}
-				if (message.type === "result") {
-					this.#capabilities = message.capabilities;
-					const pending = this.#pending.get(message.id);
-					this.#pending.delete(message.id);
-					pending?.resolve(message.capture);
-					logger.debug("Native computer capture completed", {
-						requestId: message.id,
-						backend: message.capture.backend,
-						capturePermission: message.capture.capturePermission,
-						inputPermission: message.capture.inputPermission,
-						width: message.capture.width,
-						height: message.capture.height,
-					});
-					return;
-				}
-				if (message.type === "error") {
-					const error = workerError(message.error);
-					logger.warn("Native computer worker request failed", {
-						requestId: message.id,
-						errorName: message.error.name,
-					});
-					if (message.id) {
-						const pending = this.#pending.get(message.id);
-						this.#pending.delete(message.id);
-						pending?.reject(error);
-					} else {
-						ready.reject(error);
-					}
-				}
-			});
+			this.#unsubscribeMessage = worker.onMessage(message => this.#handleMessage(message));
 			this.#unsubscribeError = worker.onError(error => {
-				ready.reject(error);
-				void this.#terminate(error);
+				void this.#workerFailed(error);
 			});
-			worker.send({ type: "init", options: this.options });
 		} catch (error) {
-			ready.reject(error);
+			started.reject(error);
 		}
 		this.#startPromise = withTimeout(
-			ready.promise,
-			this.timeouts.startMs,
-			"Timed out starting native computer worker",
+			started.promise,
+			this.#timeouts.startMs,
+			"Timed out starting computer worker",
 		).catch(async error => {
 			await this.#terminate(error);
 			throw error;
@@ -209,21 +225,110 @@ export class ComputerSupervisor implements ComputerController {
 		return this.#startPromise;
 	}
 
+	#handleMessage(message: ComputerWorkerOutbound): void {
+		if (message.type === "ready") {
+			this.#startResolve?.();
+			this.#startResolve = undefined;
+			this.#startReject = undefined;
+			return;
+		}
+		if (message.type === "result") {
+			const pending = this.#pending.get(message.id);
+			if (!pending) return;
+			this.#pending.delete(message.id);
+			if (message.ok) {
+				this.#latestCapabilities = message.payload.capabilities;
+				pending.resolve(message.payload);
+			} else {
+				pending.reject(errorFromPayload(message.error));
+			}
+			return;
+		}
+		if (message.type === "tool-call") {
+			void this.#dispatchToolCall(message);
+		}
+	}
+
+	async #dispatchToolCall(message: Extract<ComputerWorkerOutbound, { type: "tool-call" }>): Promise<void> {
+		const pending = this.#pending.get(message.runId);
+		if (!pending) {
+			this.#safeSend({
+				type: "tool-reply",
+				id: message.id,
+				reply: {
+					ok: false,
+					error: { name: "ToolError", message: "No active run for tool call", isToolError: true, isAbort: false },
+				},
+			});
+			return;
+		}
+		const controller = new AbortController();
+		pending.toolCalls.set(message.id, controller);
+		const onParentAbort = (): void => controller.abort(pending.signal?.reason);
+		if (pending.signal?.aborted) onParentAbort();
+		else pending.signal?.addEventListener("abort", onParentAbort, { once: true });
+		try {
+			const value = await this.#callSessionTool(message.name, message.args, {
+				session: this.#session,
+				signal: controller.signal,
+				emitStatus: () => {},
+			});
+			this.#safeSend({ type: "tool-reply", id: message.id, reply: { ok: true, value } });
+		} catch (error) {
+			this.#safeSend({ type: "tool-reply", id: message.id, reply: { ok: false, error: toErrorPayload(error) } });
+		} finally {
+			pending.toolCalls.delete(message.id);
+			pending.signal?.removeEventListener("abort", onParentAbort);
+		}
+	}
+
+	#safeSend(message: ComputerWorkerInbound): void {
+		try {
+			this.#worker?.send(message);
+		} catch (error) {
+			logger.debug("Computer worker send failed", {
+				error: error instanceof Error ? error.message : String(error),
+			});
+		}
+	}
+
+	async #raceWithGrace(promise: Promise<ComputerRunOk>, timeoutMs: number): Promise<ComputerRunOk> {
+		const timeoutSignal = AbortSignal.timeout(timeoutMs + GRACE_MS);
+		const timeout = Promise.withResolvers<never>();
+		const onTimeout = (): void => timeout.reject(new ToolError(RESTART_MESSAGE));
+		timeoutSignal.addEventListener("abort", onTimeout, { once: true });
+		try {
+			return await Promise.race([promise, timeout.promise]);
+		} catch (error) {
+			if (error instanceof ToolError && error.message === RESTART_MESSAGE) await this.#terminate(error);
+			throw error;
+		} finally {
+			timeoutSignal.removeEventListener("abort", onTimeout);
+		}
+	}
+
+	async #workerFailed(error: Error): Promise<void> {
+		logger.warn("Computer worker failed", { error: error.message });
+		await this.#terminate(error);
+	}
+
 	async #terminate(reason: unknown): Promise<void> {
-		logger.debug("Terminating native computer worker", {
-			reason: reason instanceof Error ? reason.name : typeof reason,
-		});
 		const worker = this.#worker;
 		this.#worker = undefined;
 		this.#startPromise = undefined;
-		this.#capabilities = undefined;
+		this.#startReject?.(reason);
+		this.#startReject = undefined;
+		this.#startResolve = undefined;
 		this.#unsubscribeMessage?.();
 		this.#unsubscribeMessage = undefined;
 		this.#unsubscribeError?.();
 		this.#unsubscribeError = undefined;
-		for (const pending of this.#pending.values()) pending.reject(reason);
+		for (const pending of this.#pending.values()) {
+			for (const controller of pending.toolCalls.values()) controller.abort(reason);
+			pending.reject(reason);
+		}
 		this.#pending.clear();
-		await worker?.terminate();
+		await worker?.terminate().catch(() => undefined);
 	}
 
 	async close(): Promise<void> {
@@ -237,7 +342,7 @@ export class ComputerSupervisor implements ComputerController {
 		});
 		try {
 			worker.send({ type: "close" });
-			await withTimeout(closed.promise, this.timeouts.closeMs, "Timed out closing native computer worker");
+			await withTimeout(closed.promise, this.#timeouts.closeMs, "Timed out closing computer worker");
 		} catch {
 			// Forced termination below is the bounded close fallback.
 		} finally {
@@ -249,6 +354,7 @@ export class ComputerSupervisor implements ComputerController {
 
 const ownedSupervisors = new Map<string, Set<ComputerController>>();
 
+/** Registers a controller for owner-scoped session cleanup. */
 export function registerComputerController(ownerId: string | undefined, controller: ComputerController): () => void {
 	if (!ownerId) return () => {};
 	const controllers = ownedSupervisors.get(ownerId) ?? new Set<ComputerController>();
@@ -260,6 +366,7 @@ export function registerComputerController(ownerId: string | undefined, controll
 	};
 }
 
+/** Closes every computer session owned by an agent session. */
 export async function releaseComputerSessionsForOwner(ownerId: string | undefined): Promise<void> {
 	if (!ownerId) return;
 	const controllers = ownedSupervisors.get(ownerId);
@@ -268,24 +375,20 @@ export async function releaseComputerSessionsForOwner(ownerId: string | undefine
 	await Promise.allSettled(Array.from(controllers, controller => controller.close()));
 }
 
+/** Verifies computer worker startup, messaging, and bounded shutdown. */
 export async function smokeTestComputerWorker(
 	timeoutMs = SMOKE_TIMEOUT_MS,
 	createWorker: ComputerWorkerFactory = spawnComputerWorker,
 ): Promise<void> {
 	const worker = createWorker();
-	const exchange = async (
-		message: ComputerWorkerInbound,
-		expected: ComputerWorkerOutbound["type"],
-		failureMessage: string,
-	): Promise<void> => {
+	const waitFor = async (expected: ComputerWorkerOutbound["type"], failureMessage: string): Promise<void> => {
 		const response = Promise.withResolvers<void>();
 		const unsubscribeMessage = worker.onMessage(received => {
 			if (received.type === expected) response.resolve();
-			else if (received.type === "error") response.reject(workerError(received.error));
+			else if (received.type === "result" && !received.ok) response.reject(errorFromPayload(received.error));
 		});
 		const unsubscribeError = worker.onError(error => response.reject(error));
 		try {
-			worker.send(message);
 			await withTimeout(response.promise, timeoutMs, failureMessage);
 		} finally {
 			unsubscribeMessage();
@@ -294,12 +397,12 @@ export async function smokeTestComputerWorker(
 	};
 
 	try {
-		await exchange(
-			{ type: "init", options: { backend: "auto", display: "all", maxWidth: 1920, maxHeight: 1200 } },
-			"ready",
-			"Computer worker smoke initialization timed out",
-		);
-		await exchange({ type: "close" }, "closed", "Computer worker smoke close timed out");
+		const pong = waitFor("pong", "Computer worker smoke ping timed out");
+		worker.send({ type: "ping", id: `computer-smoke-${Snowflake.next()}` });
+		await pong;
+		const closed = waitFor("closed", "Computer worker smoke close timed out");
+		worker.send({ type: "close" });
+		await closed;
 	} finally {
 		await worker.terminate();
 	}

package/src/tools/computer/worker-entry.ts

@@ -3,9 +3,12 @@ import { consumeWorkerInbox, isWorkerHostSelector } from "@oh-my-pi/pi-utils/wor
 import type { ComputerWorkerInbound, ComputerWorkerTransport } from "./protocol";
 import { ComputerWorkerCore } from "./worker";
 
-export function startComputerWorker(): void {
-	if (!parentPort) throw new Error("computer-worker-entry: missing parentPort");
+let started = false;
 
+/** Starts the computer worker once when running inside a Bun worker thread. */
+export function startComputerWorker(): void {
+	if (started || !parentPort) return;
+	started = true;
 	const port = parentPort;
 	const inbox = consumeWorkerInbox();
 	const transport: ComputerWorkerTransport = {
@@ -26,9 +29,11 @@ export function startComputerWorker(): void {
 	new ComputerWorkerCore(transport);
 }
 
-// Bun workers report `import.meta.main === false`. The source fallback still
-// enters this file directly, while packaged CLI workers carry the selector and
-// start the named entry only after installing its inbox.
-if (!Bun.isMainThread && !process.argv.some(isWorkerHostSelector) && import.meta.path === Bun.main) {
+// Direct-source fallback: loaded as a worker's entry module outside a CLI
+// host there is no selector argv, so start immediately. When any CLI-host
+// worker re-enters cli.ts (which imports this module statically), the
+// selector guard defers to the host's dispatch — an unguarded auto-start
+// would hijack the message port of every other worker kind.
+if (!Bun.argv.some(isWorkerHostSelector)) {
 	startComputerWorker();
 }