@oh-my-pi/pi-coding-agent 14.5.12 → 14.5.13

package/src/tools/browser/tab-supervisor.ts

@@ -0,0 +1,429 @@
+import { getPuppeteerDir, logger, Snowflake } from "@oh-my-pi/pi-utils";
+import type { Page, Target } from "puppeteer-core";
+import type { ToolSession } from "../../sdk";
+import { expandPath } from "../path-utils";
+import { ToolAbortError, ToolError } from "../tool-errors";
+import { pickElectronTarget } from "./attach";
+import { type BrowserHandle, type BrowserKindTag, holdBrowser, releaseBrowser } from "./registry";
+import type {
+	ReadyInfo,
+	RunErrorPayload,
+	RunResultOk,
+	SessionSnapshot,
+	Transferable,
+	Transport,
+	WorkerInbound,
+	WorkerInitPayload,
+	WorkerOutbound,
+} from "./tab-protocol";
+import { WorkerCore } from "./tab-worker";
+
+interface WorkerHandle {
+	send(msg: WorkerInbound, transferList?: Transferable[]): void;
+	onMessage(handler: (msg: WorkerOutbound) => void): () => void;
+	terminate(): Promise<void>;
+	readonly mode: "worker" | "inline";
+}
+
+export type DialogPolicy = "accept" | "dismiss";
+
+export interface TabSession {
+	name: string;
+	browser: BrowserHandle;
+	targetId: string;
+	worker: WorkerHandle;
+	state: "alive" | "dead";
+	info: ReadyInfo;
+	pending: Map<string, { resolve: (result: RunResultOk) => void; reject: (error: unknown) => void }>;
+	dialogPolicy?: DialogPolicy;
+	kindTag: BrowserKindTag;
+}
+
+export interface AcquireTabOptions {
+	url?: string;
+	waitUntil?: "load" | "domcontentloaded" | "networkidle0" | "networkidle2";
+	viewport?: { width: number; height: number; deviceScaleFactor?: number };
+	target?: string;
+	signal?: AbortSignal;
+	timeoutMs: number;
+	dialogs?: DialogPolicy;
+}
+
+export interface AcquireTabResult {
+	tab: TabSession;
+	created: boolean;
+}
+
+export interface RunInTabOptions {
+	code: string;
+	timeoutMs: number;
+	signal?: AbortSignal;
+	session: ToolSession;
+}
+
+export interface ReleaseTabOptions {
+	kill?: boolean;
+}
+
+const tabs = new Map<string, TabSession>();
+const GRACE_MS = 750;
+
+export function getTab(name: string): TabSession | undefined {
+	return tabs.get(name);
+}
+
+export function listTabs(): TabSession[] {
+	return [...tabs.values()];
+}
+
+export async function acquireTab(
+	name: string,
+	browser: BrowserHandle,
+	opts: AcquireTabOptions,
+): Promise<AcquireTabResult> {
+	const existing = tabs.get(name);
+	if (existing) {
+		if (existing.browser === browser && existing.state === "alive") {
+			if (opts.dialogs !== undefined && opts.dialogs !== existing.dialogPolicy) {
+				await releaseTab(name, { kill: false });
+			} else {
+				if (opts.url) {
+					await runInTabWithSnapshot(
+						name,
+						{
+							code: `await tab.goto(${JSON.stringify(opts.url)}, { waitUntil: ${JSON.stringify(opts.waitUntil ?? "networkidle2")} });`,
+							timeoutMs: opts.timeoutMs,
+							signal: opts.signal,
+						},
+						{ cwd: process.cwd() },
+					);
+				}
+				return { tab: tabs.get(name)!, created: false };
+			}
+		} else {
+			await releaseTab(name, { kill: false });
+		}
+	}
+
+	const initPayload = await buildInitPayload(browser, opts);
+	const worker = await spawnTabWorker();
+	const { promise, resolve, reject } = Promise.withResolvers<ReadyInfo>();
+	const unlisten = worker.onMessage(msg => {
+		if (msg.type === "ready") resolve(msg.info);
+		else if (msg.type === "init-failed") reject(errorFromPayload(msg.error));
+		else if (msg.type === "log") logWorkerMessage(msg);
+	});
+	let info: ReadyInfo;
+	try {
+		worker.send({ type: "init", payload: initPayload });
+		info = await raceWithTimeout(promise, opts.timeoutMs + GRACE_MS, "Timed out initializing browser tab worker");
+	} catch (error) {
+		unlisten();
+		await worker.terminate().catch(() => undefined);
+		if (browser.refCount === 0) await releaseBrowser(browser, { kill: false });
+		throw error;
+	}
+	unlisten();
+
+	holdBrowser(browser);
+	const tab: TabSession = {
+		name,
+		browser,
+		targetId: info.targetId,
+		worker,
+		state: "alive",
+		info,
+		pending: new Map(),
+		dialogPolicy: opts.dialogs,
+		kindTag: browser.kind.kind,
+	};
+	worker.onMessage(msg => handleTabMessage(tab, msg));
+	tabs.set(name, tab);
+	return { tab, created: true };
+}
+
+export async function runInTab(name: string, opts: RunInTabOptions): Promise<RunResultOk> {
+	return await runInTabWithSnapshot(
+		name,
+		{ code: opts.code, timeoutMs: opts.timeoutMs, signal: opts.signal },
+		{ cwd: opts.session.cwd, browserScreenshotDir: expandBrowserScreenshotDir(opts.session) },
+	);
+}
+
+async function runInTabWithSnapshot(
+	name: string,
+	opts: { code: string; timeoutMs: number; signal?: AbortSignal },
+	snapshot: SessionSnapshot,
+): Promise<RunResultOk> {
+	const tab = tabs.get(name);
+	if (!tab || tab.state === "dead") throw new ToolError(`Tab ${JSON.stringify(name)} is not alive. Reopen it.`);
+	if (tab.pending.size > 0) throw new ToolError(`Tab ${JSON.stringify(name)} is busy`);
+	const id = Snowflake.next();
+	const { promise, resolve, reject } = Promise.withResolvers<RunResultOk>();
+	tab.pending.set(id, { resolve, reject });
+	const abort = (): void => tab.worker.send({ type: "abort", id });
+	if (opts.signal?.aborted) abort();
+	else opts.signal?.addEventListener("abort", abort, { once: true });
+	try {
+		tab.worker.send({ type: "run", id, name, code: opts.code, timeoutMs: opts.timeoutMs, session: snapshot });
+		return await raceWithTimeout(
+			promise,
+			opts.timeoutMs + GRACE_MS,
+			"Browser code execution hung past grace; tab killed",
+			async reason => await forceKillTab(name, reason),
+		);
+	} finally {
+		opts.signal?.removeEventListener("abort", abort);
+		tab.pending.delete(id);
+	}
+}
+
+export async function releaseTab(name: string, opts: ReleaseTabOptions = {}): Promise<boolean> {
+	const tab = tabs.get(name);
+	if (!tab) {
+		logger.debug("releaseTab: unknown tab", { name });
+		return false;
+	}
+	const wasAlive = tab.state === "alive";
+	tab.state = "dead";
+	const closeError = new ToolError(`Tab ${JSON.stringify(name)} was closed`);
+	for (const [id, pending] of tab.pending) {
+		try {
+			tab.worker.send({ type: "abort", id });
+		} catch {}
+		pending.reject(closeError);
+	}
+	tab.pending.clear();
+	let forced = false;
+	if (wasAlive) {
+		try {
+			tab.worker.send({ type: "close" });
+			await waitForClosed(tab);
+		} catch {
+			forced = true;
+		}
+	}
+	await tab.worker.terminate().catch(() => undefined);
+	if (forced && tab.kindTag === "headless") await closeOrphanTarget(tab);
+	await releaseBrowser(tab.browser, { kill: opts.kill ?? false });
+	tabs.delete(name);
+	return true;
+}
+
+export async function releaseAllTabs(opts: ReleaseTabOptions = {}): Promise<number> {
+	const names = [...tabs.keys()];
+	let count = 0;
+	for (const name of names) {
+		if (await releaseTab(name, opts)) count++;
+	}
+	return count;
+}
+
+export async function dropHeadlessTabs(): Promise<void> {
+	const names = [...tabs.values()].filter(tab => tab.kindTag === "headless").map(tab => tab.name);
+	for (const name of names) await releaseTab(name);
+}
+
+async function buildInitPayload(browser: BrowserHandle, opts: AcquireTabOptions): Promise<WorkerInitPayload> {
+	const safeDir = getPuppeteerDir();
+	const browserWSEndpoint = browser.browser.wsEndpoint();
+	if (!browserWSEndpoint) throw new ToolError("Browser websocket endpoint is unavailable");
+	if (browser.kind.kind === "headless") {
+		return {
+			mode: "headless",
+			browserWSEndpoint,
+			safeDir,
+			viewport: opts.viewport,
+			dialogs: opts.dialogs,
+			url: opts.url,
+			waitUntil: opts.waitUntil,
+			timeoutMs: opts.timeoutMs,
+		};
+	}
+	const page = await pickElectronTarget(browser.browser, opts.target);
+	const targetId = await targetIdForPage(page);
+	return {
+		mode: "attach",
+		browserWSEndpoint,
+		safeDir,
+		targetId,
+		dialogs: opts.dialogs,
+	};
+}
+
+function handleTabMessage(tab: TabSession, msg: WorkerOutbound): void {
+	if (msg.type === "result") {
+		const pending = tab.pending.get(msg.id);
+		if (!pending) return;
+		tab.pending.delete(msg.id);
+		if (msg.ok) {
+			pending.resolve(msg.payload);
+			return;
+		}
+		pending.reject(errorFromPayload(msg.error));
+		return;
+	}
+	if (msg.type === "ready") {
+		tab.info = msg.info;
+		return;
+	}
+	if (msg.type === "log") logWorkerMessage(msg);
+}
+
+async function forceKillTab(name: string, reason: string): Promise<void> {
+	const tab = tabs.get(name);
+	if (!tab) return;
+	tab.state = "dead";
+	const error = new ToolError(reason);
+	for (const pending of tab.pending.values()) pending.reject(error);
+	tab.pending.clear();
+	await tab.worker.terminate().catch(() => undefined);
+	if (tab.kindTag === "headless") await closeOrphanTarget(tab);
+	await releaseBrowser(tab.browser, { kill: false });
+	tabs.delete(name);
+}
+
+async function closeOrphanTarget(tab: TabSession): Promise<void> {
+	for (const target of tab.browser.browser.targets()) {
+		if ((await targetIdForTarget(target).catch(() => "")) !== tab.targetId) continue;
+		const page = await target.page().catch(() => null);
+		await page?.close().catch(() => undefined);
+		return;
+	}
+}
+
+async function waitForClosed(tab: TabSession): Promise<void> {
+	const { promise, resolve } = Promise.withResolvers<void>();
+	const unsubscribe = tab.worker.onMessage(msg => {
+		if (msg.type === "closed") resolve();
+	});
+	try {
+		await raceWithTimeout(promise, GRACE_MS, "Timed out closing browser tab worker");
+	} finally {
+		unsubscribe();
+	}
+}
+
+function expandBrowserScreenshotDir(session: ToolSession): string | undefined {
+	const value = session.settings.get("browser.screenshotDir") as string | undefined;
+	return value ? expandPath(value) : undefined;
+}
+
+async function targetIdForPage(page: Page): Promise<string> {
+	return await targetIdForTarget(page.target());
+}
+
+async function targetIdForTarget(target: Target): Promise<string> {
+	const raw = target as unknown as { _targetId?: unknown };
+	if (typeof raw._targetId === "string") return raw._targetId;
+	const session = await target.createCDPSession();
+	try {
+		const info = (await session.send("Target.getTargetInfo")) as { targetInfo?: { targetId?: string } };
+		if (info.targetInfo?.targetId) return info.targetInfo.targetId;
+		throw new ToolError("Target id unavailable from CDP target info");
+	} finally {
+		await session.detach().catch(() => undefined);
+	}
+}
+
+function errorFromPayload(payload: RunErrorPayload): Error {
+	const error = payload.isAbort
+		? new ToolAbortError()
+		: payload.isToolError
+			? new ToolError(payload.message)
+			: new Error(payload.message);
+	error.name = payload.name;
+	if (payload.stack) error.stack = payload.stack;
+	return 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,
+	onTimeout?: (reason: string) => Promise<void>,
+): 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]);
+	} catch (error) {
+		if (error instanceof ToolError && error.message === reason) await onTimeout?.(reason);
+		throw error;
+	} finally {
+		timeoutSignal.removeEventListener("abort", onAbort);
+	}
+}
+
+async function spawnTabWorker(): Promise<WorkerHandle> {
+	try {
+		const url = new URL("./tab-worker-entry.ts", import.meta.url);
+		const worker = new Worker(url.href, { type: "module" });
+		return wrapBunWorker(worker);
+	} catch (err) {
+		logger.warn("Bun Worker spawn failed; using inline tab 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, transferList) {
+			worker.postMessage(msg, { transfer: transferList ?? [] });
+		},
+		onMessage(handler) {
+			const wrap = (event: MessageEvent): void => handler(event.data as WorkerOutbound);
+			worker.addEventListener("message", wrap);
+			return () => worker.removeEventListener("message", wrap);
+		},
+		async terminate() {
+			worker.terminate();
+		},
+	};
+}
+
+/**
+ * Inline fallback for environments where Bun cannot compile or spawn the worker
+ * entry. This preserves normal browser 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 as WorkerOutbound);
+			}),
+		onMessage: handler => {
+			const typed = handler as (message: WorkerInbound) => void;
+			workerListeners.add(typed);
+			return () => workerListeners.delete(typed);
+		},
+		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 terminate() {},
+	};
+}

package/src/tools/browser/tab-worker-entry.ts

@@ -0,0 +1,21 @@
+import { parentPort } from "node:worker_threads";
+import type { Transport, WorkerInbound, WorkerOutbound } from "./tab-protocol";
+import { WorkerCore } from "./tab-worker";
+
+if (!parentPort) throw new Error("tab-worker-entry: missing parentPort");
+
+const transport: Transport = {
+	send(msg, transferList) {
+		parentPort!.postMessage(msg, transferList ?? []);
+	},
+	onMessage(handler) {
+		const wrap = (message: unknown): void => handler(message as WorkerOutbound | WorkerInbound);
+		parentPort!.on("message", wrap);
+		return () => parentPort!.off("message", wrap);
+	},
+	close() {
+		parentPort!.close();
+	},
+};
+
+new WorkerCore(transport);