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 }, }; }); }