pi-map/tests/retrieve.test.ts

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