Pi Review Extension

review.ts

import { getLanguageFromPath, highlightCode, type ExtensionAPI, type ExtensionContext } from "@earendil-works/pi-coding-agent";
import { readFile } from "node:fs/promises";
import { resolve } from "node:path";

const REVIEW_TOOLS = new Set(["bash", "write", "edit"]);
const CUSTOM_TYPE = "review-state";
const PREVIEW_LIMIT = 6000;
const BLOCK_LIMIT = 1200;
const MAX_DIFF_LINES = 160;
const MAX_DIFF_INPUT_LINES = 400;
const DIFF_CONTEXT_LINES = 3;

type ReviewToolName = "bash" | "write" | "edit";
type ToolInput = Record<string, unknown>;
type ReviewTheme = ExtensionContext["ui"]["theme"];
type StoredEntry = {
	type?: string;
	customType?: string;
	data?: { enabled?: unknown };
};
type RenderedLine = { text: string; highlighted?: string };
type LinePart = { kind: "context" | "remove" | "add" | "skip"; text: string; highlighted?: string };
type EditBlock = { oldText: string; newText: string };

const truncate = (text: string, limit: number) => {
	if (text.length <= limit) return text;
	return `${text.slice(0, limit)}\n… truncated ${text.length - limit} character(s)`;
};

const textField = (input: ToolInput, key: string) => {
	const value = input[key];
	return typeof value === "string" ? value : "";
};

const lineCount = (text: string) => text === "" ? 0 : text.split("\n").length;

const formatBlock = (title: string, value: string, limit = 900) => `${title}:\n${truncate(value || "(empty)", limit)}`;

const formatBashPreview = (theme: ReviewTheme, command: string) => [
	theme.fg("text", "Review bash tool call"),
	"",
	theme.fg("accent", truncate(command || "(empty)", PREVIEW_LIMIT)),
].join("\n");

const splitLines = (text: string) => {
	if (text === "") return [];
	const lines = text.split("\n");
	if (lines[lines.length - 1] === "") lines.pop();
	return lines;
};

const reviewLanguageFromPath = (path: string) => {
	const language = getLanguageFromPath(path);
	if (language || !path.endsWith(".tmpl")) return language;
	return getLanguageFromPath(path.slice(0, -".tmpl".length));
};

const renderDiffLines = (text: string, path: string): RenderedLine[] => {
	const lines = splitLines(text).slice(0, MAX_DIFF_INPUT_LINES);
	if (lines.length === 0) return [];

	const language = path ? reviewLanguageFromPath(path) : undefined;
	if (!language) return lines.map((line) => ({ text: line }));

	const highlighted = highlightCode(lines.join("\n"), language);
	return lines.map((line, index) => ({ text: line, highlighted: highlighted[index] ?? line }));
};

const highlightedBody = (part: Pick<LinePart, "highlighted">) => part.highlighted !== undefined
	? part.highlighted || " "
	: undefined;

const addedText = (theme: ReviewTheme, part: Pick<LinePart, "text" | "highlighted">) =>
	theme.bg("toolSuccessBg", highlightedBody(part) ?? theme.fg("toolDiffAdded", part.text || " "));
const removedText = (theme: ReviewTheme, part: Pick<LinePart, "text" | "highlighted">) =>
	theme.bg("toolErrorBg", highlightedBody(part) ?? theme.fg("toolDiffRemoved", part.text || " "));
const contextText = (theme: ReviewTheme, part: Pick<LinePart, "text" | "highlighted">) =>
	highlightedBody(part) ?? theme.fg("toolDiffContext", part.text);

const commonPrefixLength = (left: string, right: string) => {
	let index = 0;
	while (index < left.length && index < right.length && left[index] === right[index]) index++;
	return index;
};

const commonSuffixLength = (left: string, right: string, prefixLength: number) => {
	let index = 0;
	while (
		index < left.length - prefixLength &&
		index < right.length - prefixLength &&
		left[left.length - 1 - index] === right[right.length - 1 - index]
	) index++;
	return index;
};

const renderChangedPair = (theme: ReviewTheme, removed: LinePart, added: LinePart) => {
	if (removed.highlighted !== undefined || added.highlighted !== undefined) {
		return [
			`${theme.fg("toolDiffRemoved", "- ")}${removedText(theme, removed)}`,
			`${theme.fg("toolDiffAdded", "+ ")}${addedText(theme, added)}`,
		];
	}

	const prefixLength = commonPrefixLength(removed.text, added.text);
	const suffixLength = commonSuffixLength(removed.text, added.text, prefixLength);
	const removedPrefix = removed.text.slice(0, prefixLength);
	const addedPrefix = added.text.slice(0, prefixLength);
	const removedMiddle = removed.text.slice(prefixLength, removed.text.length - suffixLength);
	const addedMiddle = added.text.slice(prefixLength, added.text.length - suffixLength);
	const removedSuffix = suffixLength ? removed.text.slice(removed.text.length - suffixLength) : "";
	const addedSuffix = suffixLength ? added.text.slice(added.text.length - suffixLength) : "";

	return [
		`${theme.fg("toolDiffRemoved", "- ")}${contextText(theme, { text: removedPrefix })}${removedText(theme, { text: removedMiddle })}${contextText(theme, { text: removedSuffix })}`,
		`${theme.fg("toolDiffAdded", "+ ")}${contextText(theme, { text: addedPrefix })}${addedText(theme, { text: addedMiddle })}${contextText(theme, { text: addedSuffix })}`,
	];
};

const lineDiff = (oldText: string, newText: string, path: string): LinePart[] => {
	const oldLines = renderDiffLines(oldText, path);
	const newLines = renderDiffLines(newText, path);
	const table = Array.from({ length: oldLines.length + 1 }, () => Array<number>(newLines.length + 1).fill(0));

	for (let oldIndex = oldLines.length - 1; oldIndex >= 0; oldIndex--) {
		for (let newIndex = newLines.length - 1; newIndex >= 0; newIndex--) {
			table[oldIndex][newIndex] = oldLines[oldIndex].text === newLines[newIndex].text
				? table[oldIndex + 1][newIndex + 1] + 1
				: Math.max(table[oldIndex + 1][newIndex], table[oldIndex][newIndex + 1]);
		}
	}

	const parts: LinePart[] = [];
	let oldIndex = 0;
	let newIndex = 0;
	while (oldIndex < oldLines.length && newIndex < newLines.length) {
		if (oldLines[oldIndex].text === newLines[newIndex].text) {
			parts.push({ kind: "context", ...oldLines[oldIndex] });
			oldIndex++;
			newIndex++;
		} else if (table[oldIndex + 1][newIndex] >= table[oldIndex][newIndex + 1]) {
			parts.push({ kind: "remove", ...oldLines[oldIndex++] });
		} else {
			parts.push({ kind: "add", ...newLines[newIndex++] });
		}
	}
	while (oldIndex < oldLines.length) parts.push({ kind: "remove", ...oldLines[oldIndex++] });
	while (newIndex < newLines.length) parts.push({ kind: "add", ...newLines[newIndex++] });
	return parts;
};

const compactDiffParts = (parts: LinePart[]) => {
	const changedIndexes = parts.flatMap((part, index) => part.kind === "context" ? [] : [index]);
	if (changedIndexes.length === 0) return parts;

	const visibleIndexes = new Set<number>();
	for (const changedIndex of changedIndexes) {
		const start = Math.max(0, changedIndex - DIFF_CONTEXT_LINES);
		const end = Math.min(parts.length - 1, changedIndex + DIFF_CONTEXT_LINES);
		for (let index = start; index <= end; index++) visibleIndexes.add(index);
	}

	const compacted: LinePart[] = [];
	let hiddenCount = 0;
	for (let index = 0; index < parts.length; index++) {
		const part = parts[index];
		if (visibleIndexes.has(index) || part.kind !== "context") {
			if (hiddenCount > 0) {
				compacted.push({ kind: "skip", text: `… ${hiddenCount} unchanged line(s) hidden` });
				hiddenCount = 0;
			}
			compacted.push(part);
		} else {
			hiddenCount++;
		}
	}
	if (hiddenCount > 0) compacted.push({ kind: "skip", text: `… ${hiddenCount} unchanged line(s) hidden` });
	return compacted;
};

const renderDiff = (theme: ReviewTheme, path: string, oldText: string, newText: string) => {
	const oldLineCount = splitLines(oldText).length;
	const newLineCount = splitLines(newText).length;
	const inputTruncated = oldLineCount > MAX_DIFF_INPUT_LINES || newLineCount > MAX_DIFF_INPUT_LINES;
	const parts = compactDiffParts(lineDiff(oldText, newText, path));
	const lines: string[] = [];
	let index = 0;

	while (index < parts.length && lines.length < MAX_DIFF_LINES) {
		const part = parts[index];
		if (part.kind === "context") {
			lines.push(`${theme.fg("dim", "  ")}${contextText(theme, part)}`);
			index++;
			continue;
		}
		if (part.kind === "skip") {
			lines.push(theme.fg("dim", part.text));
			index++;
			continue;
		}

		const removed: LinePart[] = [];
		const added: LinePart[] = [];
		while (parts[index]?.kind === "remove") removed.push(parts[index++]);
		while (parts[index]?.kind === "add") added.push(parts[index++]);

		const pairCount = Math.min(removed.length, added.length);
		for (let pairIndex = 0; pairIndex < pairCount && lines.length < MAX_DIFF_LINES; pairIndex++) {
			lines.push(...renderChangedPair(theme, removed[pairIndex], added[pairIndex]));
		}
		for (const removedLine of removed.slice(pairCount)) {
			if (lines.length >= MAX_DIFF_LINES) break;
			lines.push(`${theme.fg("toolDiffRemoved", "- ")}${removedText(theme, removedLine)}`);
		}
		for (const addedLine of added.slice(pairCount)) {
			if (lines.length >= MAX_DIFF_LINES) break;
			lines.push(`${theme.fg("toolDiffAdded", "+ ")}${addedText(theme, addedLine)}`);
		}
	}

	if (index < parts.length) lines.push(theme.fg("dim", `… diff truncated after ${MAX_DIFF_LINES} line(s)`));
	if (inputTruncated) lines.push(theme.fg("dim", `… compared first ${MAX_DIFF_INPUT_LINES} line(s) of each side`));
	return lines.join("\n") || theme.fg("dim", "No textual changes detected.");
};

const readExistingFile = async (ctx: ExtensionContext, path: string) => {
	try {
		return await readFile(resolve(ctx.cwd, path), "utf8");
	} catch {
		return undefined;
	}
};

const editBlocks = (input: ToolInput): EditBlock[] => {
	const edits = Array.isArray(input.edits) ? input.edits : [];
	return edits.flatMap((edit) => {
		if (!edit || typeof edit !== "object") return [];
		const block = edit as ToolInput;
		return [{ oldText: textField(block, "oldText"), newText: textField(block, "newText") }];
	});
};

const applyEditBlocks = (content: string, edits: EditBlock[]) => {
	let next = content;
	for (const [index, edit] of edits.entries()) {
		if (!edit.oldText) return { ok: false as const, error: `Edit ${index + 1} is missing oldText.` };
		const first = next.indexOf(edit.oldText);
		if (first === -1) return { ok: false as const, error: `Edit ${index + 1} oldText was not found in the current file.` };
		if (next.indexOf(edit.oldText, first + edit.oldText.length) !== -1) {
			return { ok: false as const, error: `Edit ${index + 1} oldText is not unique in the current file.` };
		}
		next = `${next.slice(0, first)}${edit.newText}${next.slice(first + edit.oldText.length)}`;
	}
	return { ok: true as const, content: next };
};

const formatRawEditBlocks = (input: ToolInput) => {
	const edits = editBlocks(input);
	if (edits.length === 0) return ["No edit blocks found in tool input."];

	return edits.flatMap((edit, index) => [
		`Edit ${index + 1}:`,
		formatBlock("oldText", edit.oldText),
		formatBlock("newText", edit.newText),
	]);
};

const buildPreview = async (ctx: ExtensionContext, toolName: ReviewToolName, input: ToolInput) => {
	if (toolName === "bash") {
		return formatBashPreview(ctx.ui.theme, textField(input, "command"));
	}

	const path = textField(input, "path");
	const existing = path ? await readExistingFile(ctx, path) : undefined;
	const before = existing ?? "";
	if (toolName === "edit") {
		const applied = applyEditBlocks(before, editBlocks(input));
		if (!applied.ok) {
			return truncate([
				"Review edit tool call",
				`Path: ${path || "(missing path)"}`,
				ctx.ui.theme.fg("warning", `Preview warning: ${applied.error}`),
				"",
				...formatRawEditBlocks(input),
			].join("\n"), PREVIEW_LIMIT);
		}

		return truncate([
			"Review edit tool call",
			`Path: ${path || "(missing path)"}`,
			`Current file: ${lineCount(before)} line(s), ${before.length} character(s)`,
			"",
			renderDiff(ctx.ui.theme, path, before, applied.content),
		].join("\n"), PREVIEW_LIMIT);
	}

	const content = textField(input, "content");
	const existingSummary = existing === undefined
		? "Existing file: not found or unreadable; proposed content is shown as additions"
		: `Existing file: ${lineCount(before)} line(s), ${before.length} character(s)`;

	return truncate([
		"Review write tool call",
		`Path: ${path || "(missing path)"}`,
		existingSummary,
		`Proposed content: ${lineCount(content)} line(s), ${content.length} character(s)`,
		"",
		renderDiff(ctx.ui.theme, path, before, content),
	].join("\n"), PREVIEW_LIMIT);
};

const buildDeniedReason = (toolName: ReviewToolName, input: ToolInput, feedback: string | undefined) => {
	const path = toolName === "bash" ? "" : textField(input, "path");
	const target = path ? ` for ${path}` : "";
	const note = feedback?.trim() || "User denied the tool call and did not provide additional feedback.";
	return truncate(`Review denied ${toolName}${target}.\nUser feedback: ${note}`, BLOCK_LIMIT);
};

const readStoredEnabled = (entries: StoredEntry[]) => {
	let enabled = false;
	for (const entry of entries) {
		if (entry.type !== "custom" || entry.customType !== CUSTOM_TYPE) continue;
		if (typeof entry.data?.enabled === "boolean") enabled = entry.data.enabled;
	}
	return enabled;
};

export default function (pi: ExtensionAPI) {
	let enabled = false;

	const persistEnabled = () => {
		pi.appendEntry(CUSTOM_TYPE, { enabled });
	};

	const setStatus = (ctx: ExtensionContext) => {
		const value = enabled ? "on" : "off";
		ctx.ui.setStatus("review", ctx.ui.theme.fg("dim", "review ") + ctx.ui.theme.fg("accent", value));
	};

	pi.on("session_start", async (_event, ctx) => {
		enabled = readStoredEnabled(ctx.sessionManager.getEntries() as StoredEntry[]);
		setStatus(ctx);
	});

	pi.on("tool_call", async (event, ctx) => {
		if (!enabled || !REVIEW_TOOLS.has(event.toolName)) return undefined;

		const toolName = event.toolName as ReviewToolName;
		const input = event.input as ToolInput;
		if (!ctx.hasUI) {
			return { block: true, reason: `Review mode blocked ${toolName}: interactive approval UI is unavailable.` };
		}

		const preview = await buildPreview(ctx, toolName, input);
		const approvalHint = "Press Enter with empty feedback to approve. Type feedback to deny.";
		const prompt = toolName === "bash"
			? `${preview}\n\n${ctx.ui.theme.fg("text", approvalHint)}`
			: `${preview}\n\n${approvalHint}`;
		const feedback = await ctx.ui.input(prompt, "Empty approves; feedback denies");
		if (feedback !== undefined && feedback.trim() === "") return undefined;

		return { block: true, reason: buildDeniedReason(toolName, input, feedback) };
	});

	pi.registerCommand("review", {
		description: "Toggle approval review for bash, write, and edit tool calls (on, off, toggle)",
		handler: async (args, ctx) => {
			const action = args.trim().toLowerCase();
			if (action === "on") enabled = true;
			else if (action === "off") enabled = false;
			else if (action === "" || action === "toggle") enabled = !enabled;
			else {
				ctx.ui.notify("Usage: /review [on|off|toggle]", "warning");
				return;
			}

			persistEnabled();
			setStatus(ctx);
			ctx.ui.notify(`Review mode ${enabled ? "on" : "off"} for this session`, "info");
		},
	});
}