Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/the-system/index.ts

Raw
import path from "node:path";
import {
	type ExtensionAPI,
	type ExtensionContext,
	isToolCallEventType,
} from "@earendil-works/pi-coding-agent";
import { loadStartupNotice, loadSystem, loadSystemCommand } from "./loaders.js";
import { closeDebug, dbg, span } from "./src/debug.ts";
import { isStaleContextError } from "./src/pi-ext-stale-context.ts";
import { systemArgumentCompletions } from "./system-command-meta.js";

type BoardRuntime = ReturnType<
	typeof import("./board-runtime.js")["createBoardRuntime"]
>;
type BundleSlotTool = ReturnType<
	typeof import("./bundle-slot-tool.js")["createBundleSlotTool"]
>;

interface OperationalRuntime {
	board: BoardRuntime;
	bundleSlot: BundleSlotTool;
}

const OPERATIONAL_TOOL_NAMES = new Set([
	"system_bundle_slot",
	"system_board_propose",
	"system_board_update",
]);

export default function registerTheSystem(pi: ExtensionAPI): void {
	let generation = 0;
	let warningNotified = false;
	let bundleSlotRunActive = false;
	let bundleSlotUsed = false;
	let boardProposalRunActive = false;
	let boardUpdateRunActive = false;
	let bundleMigrationRunActive = false;
	// Scoped tools stay inactive until /system starts governed work; defer their
	// runtime graph so routine Pi startup does not transpile unused tool code.
	let operational: OperationalRuntime | undefined;
	let operationalPromise: Promise<OperationalRuntime> | undefined;
	const loadOperationalRuntime = (): Promise<OperationalRuntime> => {
		if (operational) return Promise.resolve(operational);
		if (operationalPromise) return operationalPromise;
		operationalPromise = Promise.all([
			import("./board-runtime.js"),
			import("./bundle-slot-tool.js"),
		]).then(([{ createBoardRuntime }, { createBundleSlotTool }]) => {
			const runtime = {
				board: createBoardRuntime(),
				bundleSlot: createBundleSlotTool(() => {
					bundleSlotUsed = true;
				}),
			};
			pi.registerTool(runtime.bundleSlot);
			pi.registerTool(runtime.board.tool);
			pi.registerTool(runtime.board.updateTool);
			pi.setActiveTools(
				pi.getActiveTools().filter((name) => !OPERATIONAL_TOOL_NAMES.has(name)),
			);
			operational = runtime;
			return runtime;
		});
		return operationalPromise;
	};
	const activateBundleSlot = () => {
		const bundleSlot = operational?.bundleSlot;
		if (!bundleSlot) throw new Error("The System tools are not loaded");
		const active = pi.getActiveTools();
		if (!active.includes(bundleSlot.name))
			pi.setActiveTools([...active, bundleSlot.name]);
		bundleSlotRunActive = true;
		bundleSlotUsed = false;
	};
	const deactivateBundleSlot = () => {
		const bundleSlot = operational?.bundleSlot;
		const active = pi.getActiveTools();
		if (bundleSlot && active.includes(bundleSlot.name))
			pi.setActiveTools(active.filter((name) => name !== bundleSlot.name));
		bundleSlotRunActive = false;
		bundleSlotUsed = false;
	};
	const activateBoardProposal = async (cwd: string) => {
		const board = operational?.board;
		if (!board) throw new Error("The System tools are not loaded");
		await board.beginProposal(cwd);
		const active = pi.getActiveTools();
		if (!active.includes(board.tool.name))
			pi.setActiveTools([...active, board.tool.name]);
		boardProposalRunActive = true;
	};
	const deactivateBoardProposal = () => {
		const board = operational?.board;
		const active = pi.getActiveTools();
		if (board && active.includes(board.tool.name))
			pi.setActiveTools(active.filter((name) => name !== board.tool.name));
		board?.endProposal();
		boardProposalRunActive = false;
	};
	const activateBoardUpdate = async (cwd: string) => {
		const board = operational?.board;
		if (!board) throw new Error("The System tools are not loaded");
		await board.beginUpdate(cwd);
		const active = pi.getActiveTools();
		if (!active.includes(board.updateTool.name))
			pi.setActiveTools([...active, board.updateTool.name]);
		boardUpdateRunActive = true;
	};
	const deactivateBoardUpdate = () => {
		const board = operational?.board;
		const active = pi.getActiveTools();
		if (board && active.includes(board.updateTool.name))
			pi.setActiveTools(
				active.filter((name) => name !== board.updateTool.name),
			);
		board?.endUpdate();
		boardUpdateRunActive = false;
	};
	const activateBundleMigration = () => {
		bundleMigrationRunActive = true;
	};
	const beginLifecycle = () => {
		if (bundleSlotRunActive) deactivateBundleSlot();
		if (boardProposalRunActive) deactivateBoardProposal();
		if (boardUpdateRunActive) deactivateBoardUpdate();
	};
	const notifyFailure = (
		ctx: ExtensionContext,
		label: string,
		error: unknown,
	) => {
		if (ctx.hasUI)
			ctx.ui.notify(
				`${label}: ${error instanceof Error ? error.message : String(error)}`,
				"warning",
			);
	};

	const syncRuntimeIndicator = async (
		ctx: ExtensionContext,
		ticket: number,
	) => {
		if (ticket !== generation) return;
		try {
			const { startupRuntimeState } = await loadStartupNotice();
			if (ticket !== generation) return;
			const runtime = await startupRuntimeState(ctx.cwd);
			if (ticket !== generation) return;
			if (ctx.hasUI) {
				ctx.ui.setStatus(
					"the-system",
					runtime.active
						? ctx.ui.theme.fg(
								"accent",
								` ${runtime.prefix ?? "?"}${runtime.phaseLabel ? ` · ${runtime.phaseLabel}` : ""}`,
							)
						: undefined,
				);
				if (runtime.notice && !warningNotified) {
					warningNotified = true;
					ctx.ui.notify(runtime.notice, "warning");
				}
			}
		} catch (error) {
			if (!isStaleContextError(error)) throw error;
		}
	};

	pi.on("resources_discover", async (event) => {
		const root = path.join(import.meta.dirname, "skills");
		const skillPaths = [path.join(root, "vertical-slices")];
		const { startupRuntimeState } = await loadStartupNotice();
		const runtime = await startupRuntimeState(event.cwd);
		if (runtime.active && !runtime.notice)
			skillPaths.push(
				...["the-system-gfm", "the-system-tex", "the-system-mermaid"].map(
					(name) => path.join(root, name),
				),
			);
		return { skillPaths };
	});
	pi.on("session_start", (_event, ctx) => {
		dbg?.("session.start");
		deactivateBundleSlot();
		deactivateBoardProposal();
		deactivateBoardUpdate();
		bundleMigrationRunActive = false;
		const ticket = ++generation;
		warningNotified = false;
		const timer = setTimeout(() => {
			if (ticket === generation)
				void syncRuntimeIndicator(ctx, ticket).catch((error) =>
					notifyFailure(ctx, "The System startup check failed", error),
				);
		}, 0);
		timer.unref?.();
	});
	pi.on("session_shutdown", () => {
		dbg?.("session.shutdown");
		closeDebug();
		generation++;
		warningNotified = false;
		bundleSlotRunActive = false;
		bundleSlotUsed = false;
		boardProposalRunActive = false;
		boardUpdateRunActive = false;
		bundleMigrationRunActive = false;
		operational?.board.shutdown();
	});
	pi.on("before_agent_start", async (event, ctx) => {
		void syncRuntimeIndicator(ctx, generation).catch((error) =>
			notifyFailure(ctx, "The System startup check failed", error),
		);
		const { composeSystemSection } = await loadSystem();
		const section = composeSystemSection(ctx.cwd);
		if (section) event.systemPromptOptions.sections.the_system = section;
	});
	pi.on("agent_settled", (_event, ctx) => {
		void syncRuntimeIndicator(ctx, generation).catch((error) =>
			notifyFailure(ctx, "The System status refresh failed", error),
		);
		if (bundleSlotRunActive && bundleSlotUsed) deactivateBundleSlot();
		if (boardProposalRunActive && operational?.board.proposalComplete())
			deactivateBoardProposal();
		if (boardUpdateRunActive && operational?.board.updateComplete())
			deactivateBoardUpdate();
		bundleMigrationRunActive = false;
	});
	pi.on("tool_call", async (event, ctx) => {
		if (
			!isToolCallEventType("write", event) &&
			!isToolCallEventType("edit", event)
		)
			return;
		const { preflightFileMutation } = await loadSystem();
		const ui = {
			hasUI: ctx.hasUI,
			select: (title: string, items: string[]) => ctx.ui.select(title, items),
		};
		const options = { bundleMigration: bundleMigrationRunActive };
		const result = isToolCallEventType("write", event)
			? await preflightFileMutation(ctx.cwd, "write", event.input, ui, options)
			: await preflightFileMutation(ctx.cwd, "edit", event.input, ui, options);
		return result ? { block: true, reason: result.reason } : undefined;
	});

	pi.registerCommand("system", {
		description: "The System project governance",
		getArgumentCompletions: systemArgumentCompletions,
		handler: async (args, ctx) => {
			const finish = span?.("command.handle");
			try {
				const [{ createSystemCommand }, runtime] = await Promise.all([
					loadSystemCommand(),
					loadOperationalRuntime(),
				]);
				await createSystemCommand(pi, {
					beginLifecycle,
					activateBundleTool: activateBundleSlot,
					activateBoardProposal,
					activateBoardUpdate,
					activateBundleMigration,
					openBoard: runtime.board.open,
				}).handler(args, ctx);
				finish?.();
			} catch (error) {
				finish?.("error");
				throw error;
			}
		},
	});
}