pi-ext

import { uuidv7 } from "@earendil-works/pi-ai";
import type {
	ExtensionAPI,
	ExtensionCommandContext,
} from "@earendil-works/pi-coding-agent";
import {
	BorderedLoader,
	convertToLlm,
	serializeConversation,
} from "@earendil-works/pi-coding-agent";
import {
	boundedError,
	buildUserMessage,
	getHandoffMessages,
	responseText,
	SYSTEM_PROMPT,
} from "./core.ts";
import { closeDebug, dbg, span } from "./src/debug.ts";

async function generatePrompt(
	ctx: ExtensionCommandContext,
	conversation: string,
	goal: string,
): Promise<string | null> {
	const model = ctx.model;
	if (!model) return null;
	const end = span?.("generate");
	return ctx.ui.custom<string | null>((tui, theme, _keybindings, done) => {
		const loader = new BorderedLoader(tui, theme, "Generating handoff prompt…");
		loader.onAbort = () => done(null);
		ctx.modelRegistry
			.complete(
				model,
				{
					systemPrompt: SYSTEM_PROMPT,
					messages: [buildUserMessage(conversation, goal)],
				},
				{ signal: loader.signal, cacheRetention: "none", sessionId: uuidv7() },
			)
			.then((response) => {
				if (response.stopReason === "aborted") {
					end?.("finish", { status: "aborted" });
					done(null);
					return;
				}
				end?.("finish", { status: "generated" });
				done(responseText(response));
			})
			.catch((error) => {
				end?.("error", { status: "failed" });
				ctx.ui.notify(
					`handoff generation failed: ${boundedError(error)}`,
					"error",
				);
				done(null);
			});
		return loader;
	});
}

export default function handoffExtension(pi: ExtensionAPI): void {
	pi.registerCommand("handoff", {
		description: "Transfer context to a new focused session",
		async handler(args, ctx) {
			dbg?.("handoff.start", { mode: ctx.mode });
			if (ctx.mode !== "tui") {
				dbg?.("handoff.outcome", { status: "rejected", reason: "mode" });
				ctx.ui.notify("handoff requires interactive mode", "error");
				return;
			}
			if (!ctx.model) {
				dbg?.("handoff.outcome", { status: "rejected", reason: "model" });
				ctx.ui.notify("No model selected", "error");
				return;
			}
			const goal = args.trim();
			if (!goal) {
				dbg?.("handoff.outcome", { status: "rejected", reason: "goal" });
				ctx.ui.notify("Usage: /handoff <goal for new thread>", "error");
				return;
			}

			const messages = getHandoffMessages(ctx.sessionManager.getBranch());
			if (messages.length === 0) {
				dbg?.("handoff.outcome", { status: "rejected", reason: "empty" });
				ctx.ui.notify("No conversation to hand off", "error");
				return;
			}
			const conversation = serializeConversation(convertToLlm(messages));
			const parentSession = ctx.sessionManager.getSessionFile();

			const generated = await generatePrompt(ctx, conversation, goal);
			if (generated === null) {
				dbg?.("handoff.outcome", {
					status: "cancelled",
					messageCount: messages.length,
				});
				ctx.ui.notify("Cancelled", "info");
				return;
			}

			const prompt = await ctx.ui.editor("Edit handoff prompt", generated);
			if (prompt === undefined) {
				dbg?.("handoff.outcome", { status: "cancelled" });
				ctx.ui.notify("Cancelled", "info");
				return;
			}
			dbg?.("handoff.outcome", { status: "edited" });

			// The original ctx is stale after session replacement; use the replacement context.
			const result = await ctx.newSession({
				...(parentSession === undefined ? {} : { parentSession }),
				withSession: async (replacement) => {
					replacement.ui.setEditorText(prompt);
					replacement.ui.notify("Handoff ready. Submit when ready.", "info");
				},
			});
			dbg?.("handoff.outcome", {
				status: result.cancelled ? "cancelled" : "created",
				messageCount: messages.length,
			});
			if (result.cancelled) ctx.ui.notify("New session cancelled", "info");
		},
	});

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