pi-ext

import type {
	ExtensionAPI,
	ExtensionContext,
	SessionCompactEvent,
} from "@earendil-works/pi-coding-agent";
import { closeDebug, dbg } from "./src/debug.ts";

const CONTINUATION_TYPE = "continuity-resume";
const GOAL_STATE_TYPE = "goal-state";
const SPECIALIZED_CONTINUATION_TOOLS: Set<string> = new Set(["update_goal"]);

type PendingContinuation = {
	sessionId: string;
	timer: ReturnType<typeof setTimeout>;
};

export function buildContinuationPrompt(
	sessionFile: string | undefined,
	compactionEntryId: string,
): string {
	const sessionSource =
		sessionFile === undefined
			? "This session is ephemeral, so no persisted history is available."
			: [
					`Persisted session JSONL: ${JSON.stringify(sessionFile)}.`,
					"If intent or the next action remains unclear, inspect only the active branch ending at the compaction entry by following parentId links.",
					"Do not treat append order as branch order or launch a nested Pi process.",
				].join(" ");

	return [
		"Compaction completed. Recover the user's current intent from the compaction summary and current worktree.",
		sessionSource,
		`Compaction entry ID: ${JSON.stringify(compactionEntryId)}.`,
		"The worktree is authoritative for file state; session history is authoritative for user intent and constraints.",
		"If actionable work remains, continue it. If the intent is already satisfied, or all remaining work is blocked, state that briefly and stop.",
		"Do not rerun completed checks, retry known blockers, invent follow-up work, or act merely to demonstrate continuation.",
	].join("\n\n");
}

function isRecord(value: unknown): value is Record<string, unknown> {
	return value !== null && typeof value === "object" && !Array.isArray(value);
}

function hasActiveGoal(ctx: ExtensionContext): boolean {
	const sessionId = ctx.sessionManager.getSessionId();
	const entries = ctx.sessionManager.getBranch();
	for (let index = entries.length - 1; index >= 0; index--) {
		const entry = entries[index];
		if (
			entry === undefined ||
			entry.type !== "custom" ||
			entry.customType !== GOAL_STATE_TYPE ||
			!isRecord(entry.data)
		)
			continue;
		return (
			entry.data["status"] === "active" &&
			(entry.data["sessionId"] === undefined ||
				entry.data["sessionId"] === sessionId)
		);
	}
	return false;
}

function hasSpecializedContinuation(
	pi: ExtensionAPI,
	ctx: ExtensionContext,
): boolean {
	if (
		hasActiveGoal(ctx) ||
		pi.getActiveTools().some((name) => SPECIALIZED_CONTINUATION_TOOLS.has(name))
	)
		return true;

	let owned = false;
	let querying = true;
	// Replies must arrive during emit; late async replies cannot claim this turn.
	pi.events.emit("continuation:query", {
		sessionId: ctx.sessionManager.getSessionId(),
		respond(state: unknown) {
			if (querying && (state === "working" || state === "done")) owned = true;
		},
	});
	querying = false;
	return owned;
}

export default function continuityExtension(pi: ExtensionAPI): void {
	let pending: PendingContinuation | undefined;

	const cancelPending = (): void => {
		if (!pending) return;
		clearTimeout(pending.timer);
		pending = undefined;
		dbg?.("continuation", { outcome: "cancelled" });
	};

	const sameSession = (ctx: ExtensionContext, sessionId: string): boolean => {
		try {
			return ctx.sessionManager.getSessionId() === sessionId;
		} catch {
			return false;
		}
	};

	const scheduleContinuation = (
		event: SessionCompactEvent,
		ctx: ExtensionContext,
	): void => {
		cancelPending();
		if (event.willRetry) {
			dbg?.("continuation", { outcome: "skipped-retry" });
			return;
		}
		if (ctx.hasPendingMessages()) {
			dbg?.("continuation", { outcome: "skipped-pending" });
			return;
		}
		if (hasSpecializedContinuation(pi, ctx)) {
			dbg?.("continuation", { outcome: "skipped-specialized" });
			return;
		}

		const sessionId = ctx.sessionManager.getSessionId();
		const prompt = buildContinuationPrompt(
			ctx.sessionManager.getSessionFile(),
			event.compactionEntry.id,
		);
		const candidate: PendingContinuation = {
			sessionId,
			timer: setTimeout(() => {
				if (pending !== candidate) return;
				pending = undefined;
				if (
					!sameSession(ctx, sessionId) ||
					!ctx.isIdle() ||
					ctx.hasPendingMessages() ||
					hasSpecializedContinuation(pi, ctx)
				) {
					dbg?.("continuation", { outcome: "skipped-stale" });
					return;
				}
				try {
					pi.sendMessage(
						{
							customType: CONTINUATION_TYPE,
							content: prompt,
							display: false,
							details: {
								compactionEntryId: event.compactionEntry.id,
								reason: event.reason,
							},
						},
						{ triggerTurn: true, deliverAs: "followUp" },
					);
					dbg?.("continuation", { outcome: "sent" });
				} catch (error) {
					dbg?.("continuation", { outcome: "failed" });
					ctx.ui.notify(
						`Continuity could not resume: ${error instanceof Error ? error.message : String(error)}`,
						"error",
					);
				}
			}, 0),
		};
		pending = candidate;
		dbg?.("continuation", { outcome: "scheduled" });
	};

	pi.on("session_start", () => dbg?.("session.start"));
	pi.on("session_compact", scheduleContinuation);
	pi.on("agent_start", cancelPending);
	pi.on("session_shutdown", () => {
		cancelPending();
		dbg?.("session.shutdown");
		closeDebug();
	});
}

export const __test: {
	CONTINUATION_TYPE: typeof CONTINUATION_TYPE;
	GOAL_STATE_TYPE: typeof GOAL_STATE_TYPE;
	SPECIALIZED_CONTINUATION_TOOLS: typeof SPECIALIZED_CONTINUATION_TOOLS;
	hasActiveGoal: typeof hasActiveGoal;
} = {
	CONTINUATION_TYPE,
	GOAL_STATE_TYPE,
	SPECIALIZED_CONTINUATION_TOOLS,
	hasActiveGoal,
};