pi-ext

// ultra entry point: keep startup limited to command, result rendering, and the auto-enable check.

import { fileURLToPath } from "node:url";
import type {
	ExtensionAPI,
	ExtensionCommandContext,
	ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { getMarkdownTheme } from "@earendil-works/pi-coding-agent";
import { type AutocompleteItem, Markdown } from "@earendil-works/pi-tui";
import type { WorkflowResult } from "./engine.ts";
import { formatUltraResultMarkdown } from "./result.ts";
import {
	effectiveModelTiers,
	fanoutToolPolicyPrompt,
	loadUltraSettings,
	modelTierPrompt,
} from "./settings.ts";
import { closeDebug, dbg, span } from "./src/debug.ts";

const ULTRA_AUTHORING_SKILL_PATH = fileURLToPath(
	new URL("./prompts/ultra-authoring/SKILL.md", import.meta.url),
);

interface UltraRuntime {
	activateWorkflowTool(): void;
	subagentFlagSection(ctx: ExtensionContext): Promise<string | undefined>;
	getArgumentCompletions(prefix: string): AutocompleteItem[] | null;
	handleCommand(
		argString: string,
		ctx: ExtensionCommandContext,
	): Promise<
		| "completed"
		| "incomplete"
		| "aborted"
		| "disabled"
		| "invalid"
		| "failed"
		| "skipped"
	>;
}

export default function ultraExtension(pi: ExtensionAPI) {
	pi.on("resources_discover", (_event, ctx) => {
		const settings = loadUltraSettings(pi, ctx);
		return settings.enabled
			? { skillPaths: [ULTRA_AUTHORING_SKILL_PATH] }
			: undefined;
	});

	pi.registerMessageRenderer<WorkflowResult>(
		"ultra-result",
		(message, options, theme) => {
			if (!message.details)
				return new Markdown(
					"ultra result missing details",
					0,
					0,
					getMarkdownTheme(),
				);
			return new Markdown(
				formatUltraResultMarkdown(message.details, options.expanded),
				0,
				0,
				getMarkdownTheme(),
				{ color: (text) => theme.fg("customMessageText", text) },
			);
		},
	);

	let runtimePromise: Promise<UltraRuntime> | undefined;
	const runtime = () =>
		(runtimePromise ??= import("./implementation.ts").then(
			({ createUltraRuntime }) => createUltraRuntime(pi),
		));

	pi.on("session_start", async (_event, ctx) => {
		dbg?.("session.start");
		const settings = loadUltraSettings(pi, ctx);
		if (settings.enabled && settings.autoEnable)
			(await runtime()).activateWorkflowTool();
	});

	pi.on("session_shutdown", () => {
		dbg?.("session.shutdown");
		closeDebug();
	});

	pi.on("before_agent_start", async (event, ctx) => {
		if (!pi.getActiveTools().includes("run_workflow")) return;
		const settings = loadUltraSettings(pi, ctx);
		if (!settings.enabled) return;
		event.systemPromptOptions.sections.ultra_model_tier_policy =
			modelTierPrompt(effectiveModelTiers(settings, ctx.model?.provider));
		event.systemPromptOptions.sections.ultra_fanout_tool_policy =
			fanoutToolPolicyPrompt(settings.fanoutToolAllowlist);
		// run_workflow is active, so the runtime is already loaded.
		const flags = await (await runtime()).subagentFlagSection(ctx);
		if (flags) event.systemPromptOptions.sections.ultra_subagent_flags = flags;
	});

	pi.registerCommand("ultra", {
		description:
			"Enable ultra, execute workflow instructions, or run a saved workflow.",
		getArgumentCompletions: async (prefix) =>
			(await runtime()).getArgumentCompletions(prefix),
		handler: async (argString, ctx) => {
			const finish = span?.("command.handle");
			try {
				const outcome = await (await runtime()).handleCommand(argString, ctx);
				finish?.("finish", { outcome });
			} catch (error) {
				finish?.("error", { outcome: "failed" });
				throw error;
			}
		},
	});
}