Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/mockup/index.ts

Raw
import { fileURLToPath } from "node:url";
import {
	defineTool,
	type ExtensionAPI,
	type ExtensionContext,
	keyHint,
} from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";
import { MOCKUP_DESCRIPTION, mockupParameters } from "./contract.js";
import { closeDebug, dbg, span } from "./src/debug.ts";
import {
	parseBooleanSetting,
	resolveSetting,
	type SettingDeclaration,
} from "./src/pi-ext-settings.ts";

const MOCKUP_SKILL_PATH = fileURLToPath(
	new URL("./guidance/mockup-html/SKILL.md", import.meta.url),
);

// Precedence: env, trusted project, user, true.
const OPEN_BROWSER_SETTING: SettingDeclaration<boolean> = {
	key: "mockup.openBrowser",
	// PI_MOCKUP_NO_OPEN is a negation: env "1" means do not open.
	parse: (raw, source) => {
		const value = parseBooleanSetting(raw, source);
		return source === "env" && value !== undefined ? !value : value;
	},
	default: true,
	env: "PI_MOCKUP_NO_OPEN",
};

type MockupRuntime = ReturnType<
	typeof import("./mockup-tool.js")["createMockupRuntime"]
>;

type ToolOutcome =
	| "accept"
	| "iterate"
	| "cancelled"
	| "previewed"
	| "completed";

function toolOutcome(result: { details?: unknown }): ToolOutcome {
	const details = result.details;
	if (!details || typeof details !== "object") return "completed";
	const candidate = details as { action?: unknown; kind?: unknown };
	if (candidate.action === "accept" || candidate.action === "iterate")
		return candidate.action;
	if (candidate.kind === "cancelled") return "cancelled";
	if (candidate.kind === "preview") return "previewed";
	return "completed";
}

function savedContextCapsule(ctx: ExtensionContext): string | undefined {
	const branch = ctx.sessionManager.getBranch();
	for (let index = branch.length - 1; index >= 0; index--) {
		const entry = branch[index] as unknown as Record<string, unknown>;
		if (entry.type !== "message") continue;
		const message = entry.message;
		if (!message || typeof message !== "object") continue;
		const result = message as Record<string, unknown>;
		if (result.role !== "toolResult" || result.toolName !== "mockup") continue;
		const details = result.details;
		if (!details || typeof details !== "object") return undefined;
		const saved = details as Record<string, unknown>;
		return saved.sessionId === ctx.sessionManager.getSessionId() &&
			typeof saved.contextCapsule === "string"
			? saved.contextCapsule
			: undefined;
	}
	return undefined;
}

export default function registerMockup(pi: ExtensionAPI): void {
	// Keep browser, server, preview, and cache modules outside Pi's startup
	// import graph; tool execution loads them on demand.
	let runtime: MockupRuntime | undefined;
	let runtimePromise: Promise<MockupRuntime> | undefined;

	const loadRuntime = (): Promise<MockupRuntime> => {
		if (runtime) {
			dbg?.("runtime.cache", { state: "loaded" });
			return Promise.resolve(runtime);
		}
		if (runtimePromise) {
			dbg?.("runtime.cache", { state: "pending" });
			return runtimePromise;
		}
		const finish = span?.("runtime.load");
		runtimePromise = import("./mockup-tool.js").then(
			({ createMockupRuntime }) => {
				runtime = createMockupRuntime({
					openBrowser: (ctx) => {
						const setting = resolveSetting(pi, ctx, OPEN_BROWSER_SETTING);
						if (setting.ok) return setting.value;
						ctx.ui.notify(`${setting.error}; browser not opened`, "warning");
						return false;
					},
				});
				return runtime;
			},
		);
		void runtimePromise.then(
			() => finish?.(),
			() => finish?.("error", { kind: "import" }),
		);
		return runtimePromise;
	};

	const tool = defineTool({
		name: "mockup",
		label: "mockup",
		description: MOCKUP_DESCRIPTION,
		promptSnippet:
			"Preview HTML mockups in Pi or review them interactively in a browser",
		parameters: mockupParameters,
		renderShell: "self",
		renderCall(args, theme, context) {
			if (runtime?.tool.renderCall)
				return runtime.tool.renderCall(args, theme, context);
			const text =
				(context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
			const options = Array.isArray(args.options) ? args.options : [];
			const optionCount = options.length;
			const optionWord = `option${optionCount === 1 ? "" : "s"}`;
			const htmlCharacters = options.reduce(
				(total, option) =>
					total + (typeof option.html === "string" ? option.html.length : 0),
				0,
			);
			const phase =
				context.isPartial && !context.argsComplete
					? ` · generating HTML${htmlCharacters > 0 ? ` · ${htmlCharacters} chars` : ""}`
					: context.isPartial && !context.executionStarted
						? " · waiting to run"
						: "";
			let output = `${theme.fg("toolTitle", theme.bold("mockup"))} ${theme.fg("accent", args.title ?? "untitled")} ${theme.fg("toolOutput", `· ${optionCount} ${optionWord}${args.multi ? " · multi-select" : ""}${phase}`)}`;
			if (!context.expanded) {
				output += theme.fg(
					"dim",
					` (${keyHint("app.tools.expand", "to expand")})`,
				);
			} else {
				if (args.question)
					output += `\n${theme.fg("toolOutput", args.question)}`;
				for (const option of options) {
					const description = option.description
						? theme.fg("dim", ` · ${option.description}`)
						: "";
					output += `\n${theme.fg("accent", String(option.id))} ${theme.fg("toolOutput", option.label)}${description}`;
				}
			}
			text.setText(output);
			return text;
		},
		renderResult(result, options, theme, context) {
			if (runtime?.tool.renderResult)
				return runtime.tool.renderResult(result, options, theme, context);
			return new Text(
				theme.fg(options.isPartial ? "accent" : "muted", "Preparing mockup…"),
				0,
				0,
			);
		},
		async execute(toolCallId, params, signal, onUpdate, ctx) {
			const end = span?.("tool.execute");
			try {
				const active = await loadRuntime();
				active.remember(ctx);
				const result = await active.tool.execute(
					toolCallId,
					params,
					signal,
					onUpdate,
					ctx,
				);
				const outcome = toolOutcome(result);
				dbg?.("tool.outcome", { outcome });
				end?.("finish", { outcome });
				return result;
			} catch (error) {
				dbg?.("tool.outcome", { outcome: "failed" });
				end?.("error", { outcome: "failed" });
				throw error;
			}
		},
	});

	pi.on("resources_discover", () => ({ skillPaths: [MOCKUP_SKILL_PATH] }));
	pi.on("session_start", (_event, ctx) => {
		dbg?.("session.start", { mode: ctx.mode });
		runtime?.remember(ctx);
	});
	pi.on("session_shutdown", async (_event, ctx) => {
		dbg?.("session.shutdown");
		try {
			await runtime?.shutdown(ctx);
		} finally {
			closeDebug();
		}
	});
	pi.on("before_agent_start", (event, ctx) => {
		runtime?.agentWorking(ctx);
		const capsule = runtime?.contextCapsule(ctx) ?? savedContextCapsule(ctx);
		if (capsule) event.systemPromptOptions.sections.mockup_context = capsule;
	});
	pi.on("agent_settled", (_event, ctx) => runtime?.agentSettled(ctx));

	pi.registerTool(tool);
}