pi-map/pi-extension.ts

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "typebox";
import { execSync } from "child_process";
import { readFileSync } from "fs";
import { join } from "path";

function runCommand(
	command: string,
	args: string[],
	cwd: string,
): { stdout: string; stderr: string; success: boolean } {
	try {
		const result = execSync(`npx project-map ${command} ${args.join(" ")}`, {
			cwd,
			encoding: "utf8",
			stdio: ["pipe", "pipe", "pipe"],
		});
		return { stdout: result, stderr: "", success: true };
	} catch (error: any) {
		return {
			stdout: error.stdout || "",
			stderr: error.stderr || error.message,
			success: false,
		};
	}
}

function findPiMapFiles(cwd: string): string[] {
	try {
		const result = execSync('find . -name ".pi-map.md" -type f', {
			cwd,
			encoding: "utf8",
		});
		return result.trim().split("\n").filter(Boolean);
	} catch {
		return [];
	}
}

export default function (pi: ExtensionAPI) {
	// Register custom tools
	pi.registerTool({
		name: "project_map_init",
		label: "Project Map Init",
		description:
			"Generate .pi-map.md analysis files for the entire project or a subdirectory",
		promptSnippet:
			"Initialize project analysis files 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 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) {
			const targetPath = params.path || ctx.cwd;
			const result = runCommand(
				"init",
				targetPath === ctx.cwd ? [] : [targetPath],
				ctx.cwd,
			);
			return {
				content: [{ type: "text", text: result.stdout || result.stderr }],
				details: { success: result.success, cwd: ctx.cwd },
			};
		},
	});

	pi.registerTool({
		name: "project_map_patch",
		label: "Project Map Patch",
		description:
			"Update .pi-map.md for the directory containing a changed file",
		promptSnippet: "Update project analysis 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) {
			const result = runCommand("patch", [params.file_path], ctx.cwd);
			return {
				content: [{ type: "text", text: result.stdout || result.stderr }],
				details: { success: result.success },
			};
		},
	});

	pi.registerTool({
		name: "project_map_validate",
		label: "Project Map Validate",
		description: "Check all .pi-map.md files for staleness and discrepancies",
		promptSnippet: "Validate project analysis files 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) {
			const targetPath = params.path || ctx.cwd;
			const result = runCommand(
				"validate",
				targetPath === ctx.cwd ? [] : [targetPath],
				ctx.cwd,
			);
			return {
				content: [{ type: "text", text: result.stdout || result.stderr }],
				details: {
					success: result.success,
					clean: result.stdout.includes("clean"),
				},
			};
		},
	});

	pi.registerTool({
		name: "project_map_reinit",
		label: "Project Map Reinit",
		description: "Force full regeneration of all .pi-map.md files",
		promptSnippet: "Force full regeneration of project analysis files",
		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) {
			const args = params.path ? [params.path] : [];
			const result = runCommand("reinit", args, ctx.cwd);
			return {
				content: [{ type: "text", text: result.stdout || result.stderr }],
				details: { success: result.success },
			};
		},
	});

	// 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;

		// Check for dirty markers
		const dirtyFiles = mapFiles.filter((f) => {
			try {
				const content = readFileSync(join(ctx.cwd, f), "utf8");
				return content.includes("## dirty") && !content.includes("## dirty\n-");
			} 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) => {
		// Only inject if .pi-map.md files exist
		const mapFiles = findPiMapFiles(_ctx.cwd);
		if (mapFiles.length === 0) return {};

		return {
			message: {
				customType: "pi-project-map-hint",
				content:
					"📋 Project map active: If you modify any source file, run `project_map_patch` with the file path. If you suspect staleness, run `project_map_validate`.",
				display: false, // Don't show in UI, only in LLM context
			},
		};
	});
}