Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/good-job/index.ts

Raw
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import type { AutocompleteItem } from "@earendil-works/pi-tui";
import { boundedError, feedbackTrigger } from "./core.js";
import type { GoodJobRuntime } from "./implementation.js";
import { closeDebug, dbg, span } from "./src/debug.ts";

const SUBCOMMANDS = ["list", "status", "open", "learn"] as const;

export default function goodJobExtension(pi: ExtensionAPI): void {
	let runtime: Promise<GoodJobRuntime> | undefined;
	let generation = 0;
	let sessionReady: Promise<void> = Promise.resolve();
	let startupTimer: ReturnType<typeof setTimeout> | undefined;
	const getRuntime = () =>
		(runtime ??= import("./implementation.js").then(
			({ createGoodJobRuntime }) => createGoodJobRuntime(pi),
		));

	pi.registerCommand("gj", {
		description: "List, inspect, and learn from explicit feedback",
		getArgumentCompletions(prefix): AutocompleteItem[] | null {
			const normalized = prefix.trimStart().toLowerCase();
			if (normalized.includes(" ")) return null;
			const matches = SUBCOMMANDS.filter((command) =>
				command.startsWith(normalized),
			).map((command) => ({ value: command, label: command }));
			return matches.length ? matches : null;
		},
		async handler(args, ctx) {
			await ensureStarted(ctx);
			await (await getRuntime()).runCommand(args, ctx);
		},
	});

	pi.on("input", async (event, ctx) => {
		if (event.source === "extension") return { action: "continue" };
		const feedback = feedbackTrigger(event.text);
		if (!feedback) return { action: "continue" };
		const ui = ctx.hasUI ? ctx.ui : undefined;
		try {
			await ensureStarted(ctx);
			await (await getRuntime()).capture(feedback.kind, feedback.feedback, ctx);
		} catch (error) {
			ui?.notify(`gj could not queue: ${boundedError(error)}`, "error");
		}
		return { action: "continue" };
	});

	const startRuntime = (
		ctx: Parameters<GoodJobRuntime["start"]>[0],
		expectedGeneration = generation,
	) => {
		const currentGeneration = expectedGeneration;
		if (generation !== currentGeneration) return Promise.resolve();
		const end = span?.("runtime.start", { generation: currentGeneration });
		const ui = ctx.hasUI ? ctx.ui : undefined;
		sessionReady = getRuntime()
			.then(async (loaded) => {
				if (generation !== currentGeneration) return false;
				await loaded.start(ctx);
				return true;
			})
			.then((started) =>
				end?.("finish", { status: started ? "ready" : "skipped" }),
			)
			.catch((error) => {
				end?.("error", { type: "startup" });
				if (generation === currentGeneration)
					ui?.notify(`gj startup failed: ${boundedError(error)}`, "error");
			});
		return sessionReady;
	};

	const ensureStarted = (ctx: Parameters<GoodJobRuntime["start"]>[0]) => {
		if (startupTimer) {
			clearTimeout(startupTimer);
			startupTimer = undefined;
			return startRuntime(ctx);
		}
		return sessionReady;
	};

	pi.on("session_start", (_event, ctx) => {
		dbg?.("session.start", { mode: ctx.mode });
		const currentGeneration = ++generation;
		if (startupTimer) clearTimeout(startupTimer);
		startupTimer = setTimeout(() => {
			startupTimer = undefined;
			void startRuntime(ctx, currentGeneration);
		}, 50);
		startupTimer.unref?.();
	});

	pi.on("session_shutdown", async (_event, ctx) => {
		dbg?.("session.shutdown", { mode: ctx.mode });
		try {
			generation++;
			if (startupTimer) {
				clearTimeout(startupTimer);
				startupTimer = undefined;
			}
			await sessionReady;
			const loaded = await runtime?.catch(() => undefined);
			const end = span?.("runtime.stop");
			try {
				if (loaded) await loaded.stop(ctx);
				end?.("finish", { status: loaded ? "stopped" : "idle" });
			} catch (error) {
				end?.("error", { type: "shutdown" });
				throw error;
			}
		} finally {
			closeDebug();
		}
	});
}