pi-map/pi-extension.ts

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "typebox";
import { readFileSync, readdirSync, statSync } from "fs";
import { join, relative } from "path";
import {
	initProject,
	patchFile,
	validateMaps,
	reinitPath,
	retrieveContext,
	buildPreInitHint,
	buildAdvisoryReminder,
	buildStrictBypassGuard,
	modeAllowsPreInitHint,
	modeAllowsInjection,
	evaluateStrictBypass,
	discoverContextWindow,
	buildInjectionPayload,
	shouldReinjectForEvent,
	getRootPairMtimes,
	rootPairChanged,
} from "./src/index.js";
import { loadConfig } from "./src/config.js";
import { createLLMClient } from "./src/llm/llm-client.js";
import { LLMError } from "./src/llm/llm-error.js";

/**
 * Get the LLM client for Pi runtime.
 *
 * When running inside Pi, we ALWAYS use Pi's native LLM via
 * @mariozechner/pi-ai's complete() function. This respects the
 * user's /model selection and /login auth. No external fetch().
 */
function getPiLLMClient(ctx: any) {
	return createLLMClient("pi", { extensionContext: ctx });
}

function findPiMapFiles(cwd: string): string[] {
	const results: string[] = [];
	function walk(dir: string) {
		let entries: import("fs").Dirent[];
		try {
			entries = readdirSync(dir, { withFileTypes: true });
		} catch {
			return;
		}
		for (const entry of entries) {
			if (
				entry.isDirectory() &&
				!entry.name.startsWith(".") &&
				entry.name !== "node_modules"
			) {
				walk(join(dir, entry.name));
			}
		}
		try {
			statSync(join(dir, ".pi-map.md"));
			results.push(relative(cwd, join(dir, ".pi-map.md")));
		} catch {
			// no map in this dir
		}
	}
	walk(cwd);
	return results;
}

function isDirty(content: string): boolean {
	return content.includes("## dirty") && !content.includes("## dirty\n-");
}

function renderProgressBar(
	completed: number,
	total: number,
	currentFile?: string,
	width = 20,
): string {
	const pct = total > 0 ? completed / total : 0;
	const filled = Math.round(width * pct);
	const bar = "█".repeat(filled) + "░".repeat(width - filled);
	const file = currentFile ? ` → ${currentFile}` : "";
	return `[${bar}] ${completed}/${total}${file}`;
}

const HINT_CUSTOM_TYPE = "pi-project-map-hint";

function hintAlreadyInContext(ctx: any): boolean {
	const manager = ctx?.sessionManager;
	if (!manager || typeof manager.buildSessionContext !== "function") {
		return false;
	}

	const { messages } = manager.buildSessionContext();
	if (!Array.isArray(messages)) return false;

	return messages.some(
		(m: any) =>
			m &&
			m.role === "custom" &&
			m.customType === HINT_CUSTOM_TYPE,
	);
}

export default function (pi: ExtensionAPI) {
	let lastRootPairMtimes: import("./src/index.js").RootPairMtimes = {};
	pi.registerTool({
		name: "project_map_init",
		label: "Project Map Init",
		description:
			"Generate paired .pi-map.md and .pi-map.index.md analysis files for the entire project or a subdirectory",
		promptSnippet:
			"Initialize paired project map/index artifacts for codebase understanding",
		promptGuidelines: [
			"Use project_map_init when starting work on a new project or after significant restructuring",
			"Run project_map_init when .pi-map.md / .pi-map.index.md files are missing or severely outdated",
		],
		parameters: Type.Object({
			path: Type.Optional(
				Type.String({
					description: "Project root path (default: current directory)",
				}),
			),
		}),
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			try {
				const targetPath = params.path || ctx.cwd;
				const client = getPiLLMClient(ctx);
				await initProject(targetPath, {
					verbose: false,
					llmClient: client,
					cacheDir: ctx.cwd,
					onProgress: (info) => {
						const bar = renderProgressBar(
							info.completed,
							info.total,
							info.currentFile,
						);
						_onUpdate?.({
							content: [{ type: "text", text: bar }],
							details: {
								progress:
									info.total > 0
										? Math.round((info.completed / info.total) * 100)
										: 0,
								file: info.currentFile,
								dir: info.dir,
							},
						});
					},
				});
				return {
					content: [
						{
							type: "text",
							text: `Generated .pi-map.md files for ${targetPath}`,
						},
					],
					details: { success: true, cwd: ctx.cwd },
				};
			} catch (err: any) {
				const msg = err instanceof LLMError ? err.message : String(err);
				return {
					content: [{ type: "text", text: `Error: ${msg}` }],
					details: { success: false, error: msg },
				};
			}
		},
	});

	pi.registerTool({
		name: "project_map_patch",
		label: "Project Map Patch",
		description:
			"Update the paired .pi-map.md / .pi-map.index.md artifacts for the directory containing a changed file",
		promptSnippet:
			"Update paired project map/index artifacts after editing a source file",
		promptGuidelines: [
			"Use project_map_patch immediately after editing any source file",
			"Pass the absolute or relative path of the modified file",
		],
		parameters: Type.Object({
			file_path: Type.String({
				description: "Path to the modified file",
			}),
		}),
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			try {
				const client = getPiLLMClient(ctx);
				await patchFile(params.file_path, client, ctx.cwd);
				return {
					content: [
						{
							type: "text",
							text: `Patched map for ${params.file_path}`,
						},
					],
					details: { success: true },
				};
			} catch (err: any) {
				const msg = err instanceof LLMError ? err.message : String(err);
				return {
					content: [{ type: "text", text: `Error: ${msg}` }],
					details: { success: false, error: msg },
				};
			}
		},
	});

	pi.registerTool({
		name: "project_map_validate",
		label: "Project Map Validate",
		description:
			"Check all .pi-map.md / .pi-map.index.md files for staleness and discrepancies. Optionally repair them.",
		promptSnippet: "Validate paired project map/index artifacts for accuracy",
		promptGuidelines: [
			"Use project_map_validate before making architectural decisions if you suspect stale data",
			"Use project_map_validate to detect files that were deleted or added outside the agent",
			"Set fix=true to repair localized discrepancies without running a full project_map_reinit",
		],
		parameters: Type.Object({
			path: Type.Optional(
				Type.String({
					description: "Project root path (default: current directory)",
				}),
			),
			fix: Type.Optional(
				Type.Boolean({
					description:
						"Repair discrepancies automatically (default: false). Requires an LLM client.",
					default: false,
				}),
			),
		}),
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			try {
				const targetPath = params.path || ctx.cwd;
				const client = params.fix ? getPiLLMClient(ctx) : undefined;
				const result = await validateMaps(targetPath, {
					fix: params.fix ?? false,
					verbose: false,
					llmClient: client,
					cacheDir: ctx.cwd,
				});
				const text = result.clean
					? "All .pi-map.md files are clean."
					: `Found ${result.discrepancies.length} discrepancies:\n` +
						result.discrepancies
							.map((d) => `  [${d.type}] ${d.path}: ${d.message}`)
							.join("\n");
				return {
					content: [{ type: "text", text }],
					details: { success: true, clean: result.clean },
				};
			} catch (err: any) {
				const msg = err instanceof LLMError ? err.message : String(err);
				return {
					content: [{ type: "text", text: `Error: ${msg}` }],
					details: { success: false, error: msg },
				};
			}
		},
	});

	pi.registerTool({
		name: "project_map_reinit",
		label: "Project Map Reinit",
		description:
			"Regenerate .pi-map.md / .pi-map.index.md artifacts for a subtree plus its ancestors, falling back to full regeneration only when the subtree covers more than the configured percentage of project files (default: 10%)",
		promptSnippet:
			"Regenerate paired project map/index artifacts for a subtree or the whole project",
		promptGuidelines: [
			"Use project_map_reinit only after project_map_patch and project_map_validate --fix cannot resolve the staleness",
			"For localized changes, prefer project_map_patch <changed-file> or project_map_validate with fix=true",
			"Use project_map_reinit for widespread structural damage (e.g. broken links across many directories) or after large merges",
			"When reinit falls back to full regeneration, it is because the target subtree covers more than the configured reinitFullThresholdPercent of project files",
		],
		parameters: Type.Object({
			path: Type.Optional(
				Type.String({
					description: "Path to regenerate (default: entire project)",
				}),
			),
		}),
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			try {
				const targetPath = params.path || ctx.cwd;
				const client = getPiLLMClient(ctx);
				await reinitPath(targetPath, {
					verbose: false,
					llmClient: client,
					cacheDir: ctx.cwd,
					onProgress: (info) => {
						const bar = renderProgressBar(
							info.completed,
							info.total,
							info.currentFile,
						);
						_onUpdate?.({
							content: [{ type: "text", text: bar }],
							details: {
								progress:
									info.total > 0
										? Math.round((info.completed / info.total) * 100)
										: 0,
								file: info.currentFile,
								dir: info.dir,
							},
						});
					},
				});
				return {
					content: [
						{
							type: "text",
							text: `Regenerated maps for ${targetPath}`,
						},
					],
					details: { success: true },
				};
			} catch (err: any) {
				const msg = err instanceof LLMError ? err.message : String(err);
				return {
					content: [{ type: "text", text: `Error: ${msg}` }],
					details: { success: false, error: msg },
				};
			}
		},
	});

	pi.registerTool({
		name: "project_map_context",
		label: "Project Map Context",
		description:
			"Retrieve a compact markdown context bundle for a natural-language query using paired project map/index metadata",
		promptSnippet:
			"Get relevant project context for a task without reading every source file",
		promptGuidelines: [
			"Use project_map_context when you need to understand a task area before diving into source",
			"Pass a concise query describing the feature, bug, or area you want to explore",
			"Always read the suggested indexes first, then maps, then verify from source",
		],
		parameters: Type.Object({
			query: Type.String({
				description:
					"Natural-language query describing the task or area to explore",
			}),
		}),
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			try {
				const bundle = retrieveContext(params.query, ctx.cwd);
				return {
					content: [{ type: "text", text: bundle }],
					details: { success: true },
				};
			} catch (err: any) {
				const msg = err instanceof LLMError ? err.message : String(err);
				return {
					content: [{ type: "text", text: `Error: ${msg}` }],
					details: { success: false, error: msg },
				};
			}
		},
	});

	// Auto-load .pi-map.md files on session start
	pi.on("session_start", async (_event, ctx) => {
		const mapFiles = findPiMapFiles(ctx.cwd);
		if (mapFiles.length === 0) return;

		const dirtyFiles = mapFiles.filter((f) => {
			try {
				const content = readFileSync(join(ctx.cwd, f), "utf8");
				return isDirty(content);
			} catch {
				return false;
			}
		});

		if (dirtyFiles.length > 0) {
			ctx.ui.notify(
				`pi-project-map: ${dirtyFiles.length} dirty package(s) detected. Run project_map_validate first; use project_map_reinit only if staleness is widespread.`,
				"warning",
			);
		}
	});

	// Inject maintenance instructions before agent starts, but only once
	// within the current branch of context. Re-inject after compaction or
	// tree navigation removes the hint from the active path.
	pi.on("before_agent_start", async (_event, _ctx) => {
		const config = loadConfig(_ctx.cwd);
		const mapFiles = findPiMapFiles(_ctx.cwd);

		// Mode is off: no injection at all
		if (config.promptInjectionMode === "off") {
			return {};
		}

		// No maps exist yet: show visible pre-init hint, but only if mode allows it
		if (mapFiles.length === 0) {
			if (!modeAllowsPreInitHint(config.promptInjectionMode)) {
				return {};
			}
			if (hintAlreadyInContext(_ctx)) return {};
			return {
				message: {
					customType: HINT_CUSTOM_TYPE,
					content: buildPreInitHint(),
					display: true,
				},
			};
		}

		// Slice 4: advisory mode shows a visible lightweight reminder after init.
		// No root-pair preload, no per-turn reinjection.
		if (config.promptInjectionMode === "advisory") {
			if (hintAlreadyInContext(_ctx)) return {};
			return {
				message: {
					customType: HINT_CUSTOM_TYPE,
					content: buildAdvisoryReminder(),
					display: true,
				},
			};
		}

		// Maps exist but mode does not permit automatic artifact injection.
		if (!modeAllowsInjection(config.promptInjectionMode)) {
			return {};
		}

		// Slice 3b: detect root-pair artifact changes
		const currentMtimes = getRootPairMtimes(_ctx.cwd);
		const hasPrevious =
			lastRootPairMtimes.mapMtime !== undefined ||
			lastRootPairMtimes.indexMtime !== undefined;
		const artifactChanged =
			hasPrevious && rootPairChanged(currentMtimes, lastRootPairMtimes);
		lastRootPairMtimes = currentMtimes;

		// Slice 3a/3b: avoid redundant reinjection by scanning outgoing context
		const eventType = artifactChanged ? "artifact_change" : "agent_start";
		const decision = shouldReinjectForEvent(
			{
				messages: _event?.messages,
				type: eventType,
			},
			config.promptInjectionMode,
		);
		if (!decision.needed) {
			return {};
		}

		if (hintAlreadyInContext(_ctx)) return {};

		// Slice 2: post-init root-pair preload + budgeted expansion
		const contextWindow = discoverContextWindow(_ctx);
		const payload = buildInjectionPayload(_ctx.cwd, config, contextWindow);
		return {
			message: {
				customType: HINT_CUSTOM_TYPE,
				content: payload.content,
				display: payload.display,
			},
		};
	});

	// Per-turn context scanning for reinjection in strong/strict modes
	pi.on("context", async (event: any, ctx: any) => {
		const config = loadConfig(ctx.cwd);
		const mapFiles = findPiMapFiles(ctx.cwd);

		if (mapFiles.length === 0) {
			return {};
		}

		if (!modeAllowsInjection(config.promptInjectionMode)) {
			return {};
		}

		// Slice 3b: detect root-pair artifact changes early. Invalidation always
		// forces reinjection, even in strict mode, because the context is stale.
		const currentMtimes = getRootPairMtimes(ctx.cwd);
		const hasPrevious =
			lastRootPairMtimes.mapMtime !== undefined ||
			lastRootPairMtimes.indexMtime !== undefined;
		const artifactChanged =
			hasPrevious && rootPairChanged(currentMtimes, lastRootPairMtimes);
		lastRootPairMtimes = currentMtimes;

		if (artifactChanged) {
			const decision = shouldReinjectForEvent(
				{
					messages: event?.messages,
					type: "artifact_change",
					payload: event?.payload,
				},
				config.promptInjectionMode,
			);
			if (!decision.needed) {
				return {};
			}
			const contextWindow = discoverContextWindow(ctx);
			const payload = buildInjectionPayload(ctx.cwd, config, contextWindow);
			return {
				message: {
					customType: "pi-project-map-hint",
					content: payload.content,
					display: payload.display,
				},
			};
		}

		// Slice 4: strict-mode bypass guard for sensitive actions with missing protocol path.
		if (config.promptInjectionMode === "strict") {
			const bypass = evaluateStrictBypass(event, config.promptInjectionMode);
			if (bypass.guard) {
				return {
					message: {
						customType: "pi-project-map-hint",
						content: buildStrictBypassGuard(bypass.reason),
						display: true,
					},
				};
			}
		}

		const decision = shouldReinjectForEvent(
			{
				messages: event?.messages,
				type: event?.type,
				payload: event?.payload,
			},
			config.promptInjectionMode,
		);
		if (!decision.needed) {
			return {};
		}

		const contextWindow = discoverContextWindow(ctx);
		const payload = buildInjectionPayload(ctx.cwd, config, contextWindow);
		return {
			message: {
				customType: "pi-project-map-hint",
				content: payload.content,
				display: payload.display,
			},
		};
	});
}