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, } from "./src/index.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 .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) { 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 .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) { 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 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) { 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 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) { 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 }, }; } }, }); // 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 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, }, }; }); }