import { describe, it, expect, beforeEach, afterEach } from "vitest"; import { mkdtempSync, writeFileSync, mkdirSync, rmSync } from "fs"; import { join } from "path"; import { tmpdir } from "os"; import { retrieveContext } from "../src/retrieve.js"; describe("retrieve", () => { let dir: string; beforeEach(() => { dir = mkdtempSync(join(tmpdir(), "pi-ret-")); }); afterEach(() => { rmSync(dir, { recursive: true }); }); function writeMap( relDir: string, role: string, files: { name: string; purpose: string; exports?: string[] }[], tags?: string[], symbols?: string[], ) { const d = join(dir, relDir); if (!d.startsWith(dir)) throw new Error("invalid path"); mkdirSync(d, { recursive: true }); const fileLines = files .map((f) => { const exp = f.exports?.length ? ` | exp: ${f.exports.join(", ")}` : ""; return `- ${f.name} | ${f.purpose}${exp}`; }) .join("\n"); const tagLine = tags?.length ? tags.join(", ") : "-"; const symLines = symbols?.length ? symbols.map((s) => `- ${s}`).join("\n") : "-"; writeFileSync( join(d, ".pi-map.md"), `# ${relDir} dir: ${relDir} index: ${relDir}/.pi-map.index.md ## role ${role} ## files ${fileLines} ## arch Test arch ## tags ${tagLine} ## symbols ${symLines} ## workflows - ## dirty - `, ); writeFileSync( join(d, ".pi-map.index.md"), `# ${relDir} (index) dir: ${relDir} ## role ${role} ## parent - ## children - ## files ${files.map((f) => `- ${f.name}`).join("\n")} ## links index: ${relDir}/.pi-map.index.md map: ${relDir}/.pi-map.md ## workflows - ## dirty - `, ); } it("returns empty bundle when no maps exist", () => { const bundle = retrieveContext("auth", dir); expect(bundle).toContain("# Context bundle: auth"); expect(bundle).toContain("No relevant directories found"); }); it("ranks directories by query relevance", () => { writeMap( "src/auth", "Auth layer: JWT issuance, validation, refresh.", [{ name: "tokens.ts", purpose: "JWT gen/val", exports: ["issueToken"] }], ["auth", "jwt"], ["issueToken"], ); writeMap( "src/utils", "Shared utilities and helpers.", [{ name: "helpers.ts", purpose: "Helpers" }], ["utils"], ); const bundle = retrieveContext("jwt validation", dir); expect(bundle).toContain("src/auth/.pi-map.index.md"); expect(bundle).toContain("src/auth/.pi-map.md"); }); it("limits results to top 3 by default", () => { for (let i = 0; i < 5; i++) { writeMap( `src/pkg${i}`, `Package ${i} logic`, [{ name: `file${i}.ts`, purpose: `Feature ${i}` }], [`pkg${i}`], ); } // Query matching all const bundle = retrieveContext("pkg", dir); const indexMatches = bundle.match(/\.pi-map\.index\.md/g) || []; expect(indexMatches.length).toBeLessThanOrEqual(3); }); it("includes likely files from matched directories", () => { writeMap( "src/auth", "Auth layer", [ { name: "tokens.ts", purpose: "JWT tokens", exports: ["issueToken"] }, { name: "middleware.ts", purpose: "Auth middleware", exports: ["requireAuth"], }, ], ["auth"], ["issueToken", "requireAuth"], ); const bundle = retrieveContext("auth middleware", dir); expect(bundle).toContain("src/auth/tokens.ts"); expect(bundle).toContain("src/auth/middleware.ts"); }); it("includes relevant symbols when present", () => { writeMap( "src/auth", "Auth layer", [{ name: "tokens.ts", purpose: "JWT tokens" }], ["auth"], ["issueToken", "verifyToken"], ); const bundle = retrieveContext("auth", dir); expect(bundle).toContain("## relevant symbols"); expect(bundle).toContain("issueToken"); expect(bundle).toContain("verifyToken"); }); it("omits symbol section when no symbols exist", () => { writeMap( "src/utils", "Utilities", [{ name: "helpers.ts", purpose: "Helpers" }], ["utils"], [], ); const bundle = retrieveContext("utils", dir); expect(bundle).not.toContain("## relevant symbols"); }); it("includes instructions in every bundle", () => { writeMap( "src/cli", "CLI entrypoint", [{ name: "cli.ts", purpose: "CLI" }], ["cli"], ); const bundle = retrieveContext("cli", dir); expect(bundle).toContain("## instructions"); expect(bundle).toContain("Read the indexes first"); }); it("uses stable section order", () => { writeMap( "src/auth", "Auth layer", [{ name: "tokens.ts", purpose: "JWT tokens" }], ["auth"], ["issueToken"], ); const bundle = retrieveContext("auth", dir); const queryIdx = bundle.indexOf("## query"); const indexIdx = bundle.indexOf("## relevant indexes"); const mapIdx = bundle.indexOf("## relevant maps"); const fileIdx = bundle.indexOf("## likely files"); const symIdx = bundle.indexOf("## relevant symbols"); const instIdx = bundle.indexOf("## instructions"); expect(queryIdx).toBeGreaterThan(-1); expect(indexIdx).toBeGreaterThan(queryIdx); expect(mapIdx).toBeGreaterThan(indexIdx); expect(fileIdx).toBeGreaterThan(mapIdx); expect(symIdx).toBeGreaterThan(fileIdx); expect(instIdx).toBeGreaterThan(symIdx); }); it("scores file names and purposes", () => { writeMap( "src/auth", "Auth layer", [ { name: "validateToken.ts", purpose: "Token validation logic", exports: ["validateToken"], }, ], [], [], ); const bundle = retrieveContext("validate token", dir); expect(bundle).toContain("src/auth/.pi-map.index.md"); }); });