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,
	modeAllowsPreInitHint,
	modeAllowsInjection,
	discoverContextWindow,
	buildInjectionPayload,
} 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}`;
}

export default function (pi: ExtensionAPI) {
	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",
		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",
		],
		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 result = await validateMaps(targetPath, {
					fix: false,
					verbose: false,
				});
				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:
			"Force full regeneration of all .pi-map.md / .pi-map.index.md artifacts",
		promptSnippet:
			"Force full regeneration of paired project map/index artifacts",
		promptGuidelines: [
			"Use project_map_reinit when validation shows widespread staleness",
			"Use project_map_reinit after pulling major changes from version control",
		],
		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 packages detected. Run project_map_validate or project_map_reinit.`,
				"warning",
			);
		}
	});

	// Inject maintenance instructions before agent starts
	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 {};
			}
			return {
				message: {
					customType: "pi-project-map-hint",
					content: buildPreInitHint(),
					display: true,
				},
			};
		}

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

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