dotfiles/home/.pi/agent/extensions/approval-gate.ts

/**
 * Approval Gate Extension
 *
 * Prompts before mutating/dangerous tool calls execute.
 *
 * For `edit` calls, renders a rich diff preview (with surrounding file
 * context) by reusing pi's built-in edit tool renderer inside a custom
 * approval modal. For `bash`/`write`, falls back to a text confirm dialog.
 *
 * Config: ~/.pi/agent/approval-gate.json
 * {
 *   "bashAllowPatterns": ["^pwd$"]
 * }
 */

import { existsSync, readFileSync } from "node:fs";
import { join } from "node:path";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import {
	createEditToolDefinition,
	getAgentDir,
	type EditToolInput,
} from "@earendil-works/pi-coding-agent";
import {
	Box,
	type Component,
	Key,
	matchesKey,
	Spacer,
	Text,
} from "@earendil-works/pi-tui";
import { DynamicBorder } from "@earendil-works/pi-coding-agent";

interface ApprovalGateConfig {
	bashAllowPatterns?: unknown;
}

interface LoadedConfig {
	bashAllowPatterns: RegExp[];
	warnings: string[];
}

const BLOCKED_REASON = "Blocked by approval-gate";
const STATUS_KEY = "approval-gate";
const CONFIG_FILE = "approval-gate.json";

function truncate(value: string, maxLength = 1200): string {
	if (value.length <= maxLength) return value;
	return `${value.slice(0, maxLength)}\n… truncated ${value.length - maxLength} character(s)`;
}

function formatBytes(value: string): string {
	const bytes = Buffer.byteLength(value, "utf8");
	if (bytes < 1024) return `${bytes} B`;
	if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)} KB`;
	return `${(bytes / (1024 * 1024)).toFixed(1)} MB`;
}

function loadConfig(): LoadedConfig {
	const configPath = join(getAgentDir(), CONFIG_FILE);
	const warnings: string[] = [];

	if (!existsSync(configPath)) {
		return { bashAllowPatterns: [], warnings };
	}

	let parsed: ApprovalGateConfig;
	try {
		parsed = JSON.parse(readFileSync(configPath, "utf8")) as ApprovalGateConfig;
	} catch (error) {
		warnings.push(`Failed to parse ${configPath}: ${error instanceof Error ? error.message : String(error)}`);
		return { bashAllowPatterns: [], warnings };
	}

	if (parsed.bashAllowPatterns === undefined) {
		return { bashAllowPatterns: [], warnings };
	}

	if (!Array.isArray(parsed.bashAllowPatterns)) {
		warnings.push(`${configPath}: bashAllowPatterns must be an array of regex pattern strings`);
		return { bashAllowPatterns: [], warnings };
	}

	const bashAllowPatterns: RegExp[] = [];
	for (const pattern of parsed.bashAllowPatterns) {
		if (typeof pattern !== "string") {
			warnings.push(`${configPath}: skipped non-string bash allow pattern`);
			continue;
		}

		try {
			bashAllowPatterns.push(new RegExp(pattern));
		} catch (error) {
			warnings.push(
				`${configPath}: skipped invalid regex ${JSON.stringify(pattern)}: ${
					error instanceof Error ? error.message : String(error)
				}`,
			);
		}
	}

	return { bashAllowPatterns, warnings };
}

function updateStatus(ctx: ExtensionContext, enabled: boolean): void {
	if (!ctx.hasUI) return;
	ctx.ui.setStatus(STATUS_KEY, `approvals: ${enabled ? "✓" : "✗"}`);
}

function isAllowedBashCommand(command: string, patterns: RegExp[]): boolean {
	const trimmed = command.trim();
	return patterns.some((pattern) => pattern.test(trimmed));
}

function summarizeToolCall(toolName: string, input: Record<string, unknown>): string {
	if (toolName === "bash") {
		const command = typeof input.command === "string" ? input.command : JSON.stringify(input, null, 2);
		return `Command:\n\n${truncate(command)}`;
	}

	if (toolName === "edit") {
		const path = typeof input.path === "string" ? input.path : "unknown";
		const edits = Array.isArray(input.edits) ? input.edits : [];
		const editSummaries = edits.slice(0, 3).map((edit, index) => {
			if (!edit || typeof edit !== "object") return `Edit ${index + 1}: <unrecognized edit block>`;
			const block = edit as { oldText?: unknown; newText?: unknown };
			const oldText = typeof block.oldText === "string" ? block.oldText : "<missing oldText>";
			const newText = typeof block.newText === "string" ? block.newText : "<missing newText>";
			return [
				`Edit ${index + 1}:`,
				`  old (${formatBytes(oldText)}): ${truncate(oldText, 500)}`,
				`  new (${formatBytes(newText)}): ${truncate(newText, 500)}`,
			].join("\n");
		});

		const omitted = edits.length > editSummaries.length ? `\n\n… omitted ${edits.length - editSummaries.length} edit block(s)` : "";
		return [`Path: ${path}`, `Edit blocks: ${edits.length}`, "", ...editSummaries].join("\n") + omitted;
	}

	if (toolName === "write") {
		const path = typeof input.path === "string" ? input.path : "unknown";
		const content = typeof input.content === "string" ? input.content : undefined;
		if (content === undefined) {
			return `Path: ${path}\nContent size: unknown`;
		}

		return `Path: ${path}\nContent size: ${formatBytes(content)}\n\nContent preview:\n${truncate(content)}`;
	}

	return truncate(JSON.stringify(input, null, 2));
}

interface EditRenderState {
	callComponent?: Box & {
		preview?: unknown;
		previewArgsKey?: string;
		previewPending?: boolean;
		settledError?: boolean;
	};
}

interface EditRenderContext {
	args: EditToolInput;
	toolCallId: string;
	invalidate: () => void;
	lastComponent: Component | undefined;
	state: EditRenderState;
	cwd: string;
	executionStarted: boolean;
	argsComplete: boolean;
	isPartial: boolean;
	expanded: boolean;
	showImages: boolean;
	isError: boolean;
}

/**
 * Custom approval component for `edit` tool calls.
 *
 * Reuses pi's built-in edit tool `renderCall` to render the same contextual
 * diff preview the user would see in the normal tool-execution row, including
 * surrounding file context and colored intra-line changes. The preview is
 * computed asynchronously (reads the target file), so we re-render on
 * invalidate until it settles.
 */
class EditApprovalComponent implements Component {
	private readonly tui: { requestRender: () => void };
	private readonly theme: { fg: (color: string, text: string) => string; bold: (text: string) => string };
	private readonly editDef: ReturnType<typeof createEditToolDefinition>;
	private readonly input: EditToolInput;
	private readonly cwd: string;
	private readonly toolCallId: string;
	private readonly state: EditRenderState = {};
	private lastComponent: Component | undefined;
	private readonly header: Text;
	private readonly footer: Text;
	private readonly topBorder: DynamicBorder;
	private readonly bottomBorder: DynamicBorder;
	private settled = false;

	public onApprove?: () => void;
	public onReject?: () => void;

	constructor(opts: {
		tui: { requestRender: () => void };
		theme: { fg: (color: string, text: string) => string; bold: (text: string) => string };
		cwd: string;
		input: EditToolInput;
		toolCallId: string;
	}) {
		this.tui = opts.tui;
		this.theme = opts.theme;
		this.cwd = opts.cwd;
		this.input = opts.input;
		this.toolCallId = opts.toolCallId;
		this.editDef = createEditToolDefinition(opts.cwd);

		const editCount = Array.isArray(this.input.edits) ? this.input.edits.length : 0;
		const headerText = `${this.theme.fg("toolTitle", this.theme.bold("edit"))} ${this.theme.fg("accent", this.input.path)}${this.theme.fg("dim", ` · ${editCount} block(s) · approve?`)}`;
		this.header = new Text(headerText, 1, 0);
		this.footer = new Text(this.theme.fg("dim", "y / enter approve · n / esc reject"), 1, 0);
		const borderFn = (s: string) => this.theme.fg("accent", s);
		this.topBorder = new DynamicBorder(borderFn);
		this.bottomBorder = new DynamicBorder(borderFn);
	}

	private buildContext(): EditRenderContext {
		return {
			args: this.input,
			toolCallId: this.toolCallId,
			invalidate: () => this.tui.requestRender(),
			lastComponent: this.lastComponent,
			state: this.state,
			cwd: this.cwd,
			executionStarted: false,
			argsComplete: true,
			isPartial: false,
			expanded: true,
			showImages: false,
			isError: false,
		};
	}

	private renderEditPreview(width: number): string[] {
		const renderCall = this.editDef.renderCall;
		if (!renderCall) return [this.theme.fg("warning", "(diff preview unavailable)")];
		const component = renderCall(this.input, this.theme as never, this.buildContext() as never);
		this.lastComponent = component;
		return component.render(width);
	}

	render(width: number): string[] {
		const lines: string[] = [];
		lines.push(...this.topBorder.render(width));
		lines.push(...this.header.render(width));
		lines.push(...this.renderEditPreview(width));
		lines.push(...new Spacer(1).render(width));
		lines.push(...this.footer.render(width));
		lines.push(...this.bottomBorder.render(width));
		return lines;
	}

	handleInput(data: string): void {
		if (this.settled) return;
		if (matchesKey(data, Key.enter) || data === "y" || data === "Y") {
			this.settled = true;
			this.onApprove?.();
		} else if (matchesKey(data, Key.escape) || data === "n" || data === "N") {
			this.settled = true;
			this.onReject?.();
		}
	}

	invalidate(): void {
		this.header.invalidate();
		this.footer.invalidate();
		this.topBorder.invalidate();
		this.bottomBorder.invalidate();
		this.lastComponent?.invalidate?.();
	}
}

function isEditInput(value: Record<string, unknown>): value is EditToolInput {
	return (
		typeof value.path === "string" &&
		Array.isArray(value.edits) &&
		value.edits.length > 0 &&
		value.edits.every(
			(e) => e && typeof e === "object" && typeof (e as { oldText?: unknown }).oldText === "string" && typeof (e as { newText?: unknown }).newText === "string",
		)
	);
}

async function approveEditCall(input: Record<string, unknown>, ctx: ExtensionContext, toolCallId: string): Promise<boolean> {
	if (ctx.mode !== "tui" || !isEditInput(input)) {
		// Non-TUI or malformed input: fall back to text summary confirm.
		const message = `${summarizeToolCall("edit", input)}\n\nAllow this tool call?`;
		return ctx.ui.confirm("Approve edit?", message);
	}

	return ctx.ui.custom<boolean>((tui, theme, _keybindings, done) => {
		const component = new EditApprovalComponent({
			tui: tui as { requestRender: () => void },
			theme: theme as { fg: (color: string, text: string) => string; bold: (text: string) => string },
			cwd: ctx.cwd,
			input,
			toolCallId,
		});
		component.onApprove = () => done(true);
		component.onReject = () => done(false);
		return component;
	});
}

export default function approvalGateExtension(pi: ExtensionAPI) {
	let enabled = true;
	const config = loadConfig();

	pi.registerCommand("approval-gate", {
		description: "Turn mutating/dangerous tool approval prompts on or off",
		handler: async (args, ctx) => {
			const action = args.trim().toLowerCase();

			if (action === "on") {
				enabled = true;
				updateStatus(ctx, enabled);
				ctx.ui.notify("approval-gate enabled", "info");
				return;
			}

			if (action === "off") {
				enabled = false;
				updateStatus(ctx, enabled);
				ctx.ui.notify("approval-gate disabled", "info");
				return;
			}

			if (action === "" || action === "status") {
				updateStatus(ctx, enabled);
				ctx.ui.notify(`approval-gate is ${enabled ? "enabled" : "disabled"}`, "info");
				return;
			}

			ctx.ui.notify("Usage: /approval-gate [on|off|status]", "warning");
		},
	});

	pi.on("session_start", async (_event, ctx) => {
		updateStatus(ctx, enabled);

		for (const warning of config.warnings) {
			ctx.ui.notify(`approval-gate: ${warning}`, "warning");
		}
	});

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

		const input = event.input as Record<string, unknown>;
		const isBash = event.toolName === "bash";
		const isEdit = event.toolName === "edit";
		const isWrite = event.toolName === "write";

		if (!isBash && !isEdit && !isWrite) return undefined;

		if (isBash) {
			const command = typeof input.command === "string" ? input.command : "";
			if (command && isAllowedBashCommand(command, config.bashAllowPatterns)) {
				return undefined;
			}
		}

		if (!ctx.hasUI) {
			return { block: true, reason: `${BLOCKED_REASON} (no UI available for confirmation)` };
		}

		let approved: boolean;
		if (isEdit) {
			approved = await approveEditCall(input, ctx, event.toolCallId);
		} else {
			approved = await ctx.ui.confirm(
				`Approve ${event.toolName}?`,
				`${summarizeToolCall(event.toolName, input)}\n\nAllow this tool call?`,
			);
		}

		if (!approved) {
			return { block: true, reason: BLOCKED_REASON };
		}

		return undefined;
	});
}