Implement layered maps and context retrieval

tests/retrieve.test.ts

@@ -0,0 +1,244 @@
+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");
+	});
+});