pi-map/tests/pi-extension.test.ts

import { describe, it, expect, vi, beforeEach } from "vitest";
import { mkdtempSync, writeFileSync, existsSync, unlinkSync } from "fs";
import { join } from "path";
import { tmpdir, homedir } from "os";

vi.mock("@mariozechner/pi-coding-agent", () => ({
	ExtensionAPI: class {},
}));

vi.mock("typebox", () => ({
	Type: {
		Object: (props: unknown) => props,
		Optional: (prop: unknown) => prop,
		String: (opts: unknown) => ({ type: "string", ...(opts as object) }),
	},
}));

// Mock Pi's AI module so PiLLMClient works in tests
vi.mock("@mariozechner/pi-ai", () => ({
	complete: vi.fn(async (_model: any, _context: any) => ({
		role: "assistant" as const,
		content: [{ type: "text", text: "PURPOSE: Test file\nDEPS: none\nCONCEPTS: testing" }],
	})),
}));

import extension from "../pi-extension.js";

const CACHE_FILE = join(homedir(), ".cache", "pi-project-map", "llm-cache.json");
function clearCache() {
	if (existsSync(CACHE_FILE)) unlinkSync(CACHE_FILE);
}

describe("pi-extension", () => {
	let registeredTools: Record<string, any> = {};
	let registeredEvents: Record<string, any> = {};
	let mockCtx: any;
	let mockNotify: ReturnType<typeof vi.fn>;

	beforeEach(() => {
		clearCache();
		registeredTools = {};
		registeredEvents = {};
		mockNotify = vi.fn();
		mockCtx = {
			cwd: "/home/project",
			model: { provider: "openai", id: "gpt-4o-mini", api: "openai-completions" },
			modelRegistry: {
				getApiKeyAndHeaders: vi.fn(async () => ({ apiKey: "test-key", headers: {} })),
			},
			ui: { notify: mockNotify },
		};

		const mockPi = {
			registerTool: vi.fn((tool: any) => {
				registeredTools[tool.name] = tool;
			}),
			on: vi.fn((event: string, handler: any) => {
				registeredEvents[event] = handler;
			}),
		};

		extension(mockPi);
	});

	describe("tool registration", () => {
		it("registers 4 tools", () => {
			expect(Object.keys(registeredTools)).toHaveLength(4);
			expect(registeredTools).toHaveProperty("project_map_init");
			expect(registeredTools).toHaveProperty("project_map_patch");
			expect(registeredTools).toHaveProperty("project_map_validate");
			expect(registeredTools).toHaveProperty("project_map_reinit");
		});

		it("registers session_start and before_agent_start events", () => {
			expect(registeredEvents).toHaveProperty("session_start");
			expect(registeredEvents).toHaveProperty("before_agent_start");
		});
	});

	describe("project_map_init tool", () => {
		it("generates maps with mocked Pi LLM", async () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
			writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`);
			mockCtx.cwd = dir;

			const tool = registeredTools.project_map_init;
			const result = await tool.execute("tool-1", {}, null, null, mockCtx);
			expect(result.details.success).toBe(true);
			expect(result.content[0].text).toContain("Generated");
		});

		it("returns error when no Pi model configured", async () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
			writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`);
			mockCtx.cwd = dir;
			mockCtx.model = null;

			const tool = registeredTools.project_map_init;
			const result = await tool.execute("tool-1", {}, null, null, mockCtx);
			expect(result.details.success).toBe(false);
			expect(result.content[0].text).toContain("no model configured");
		});
	});

	describe("project_map_patch tool", () => {
		it("patches map with mocked Pi LLM", async () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
			const file = join(dir, "test.ts");
			writeFileSync(file, `export const x = ${Date.now()};`);
			mockCtx.cwd = dir;

			const tool = registeredTools.project_map_patch;
			const result = await tool.execute(
				"tool-1",
				{ file_path: file },
				null,
				null,
				mockCtx,
			);
			expect(result.details.success).toBe(true);
			expect(result.content[0].text).toContain("Patched map");
		});
	});

	describe("project_map_validate tool", () => {
		it("reports clean when map exists and is up to date", async () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
			writeFileSync(
				join(dir, ".pi-map.md"),
				"# .\n## role\nTest\n## files\n## arch\n## dirty\n-\n",
			);
			mockCtx.cwd = dir;

			const tool = registeredTools.project_map_validate;
			const result = await tool.execute("tool-1", {}, null, null, mockCtx);
			expect(result.details.clean).toBe(true);
			expect(result.content[0].text).toContain("clean");
		});

		it("reports missing when files exist but no map", async () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
			writeFileSync(join(dir, "test.ts"), "export const x = 1;");
			mockCtx.cwd = dir;

			const tool = registeredTools.project_map_validate;
			const result = await tool.execute("tool-1", {}, null, null, mockCtx);
			expect(result.details.clean).toBe(false);
			expect(result.content[0].text).toContain("missing");
		});
	});

	describe("project_map_reinit tool", () => {
		it("regenerates maps with mocked Pi LLM", async () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
			writeFileSync(join(dir, "test.ts"), `export const x = ${Date.now()};`);
			mockCtx.cwd = dir;

			const tool = registeredTools.project_map_reinit;
			const result = await tool.execute("tool-1", {}, null, null, mockCtx);
			expect(result.details.success).toBe(true);
			expect(result.content[0].text).toContain("Regenerated");
		});
	});

	describe("session_start event", () => {
		it("notifies when dirty .pi-map.md files exist", async () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
			writeFileSync(
				join(dir, ".pi-map.md"),
				"# .\n## dirty\n2024-01-01: patched\n",
			);
			mockCtx.cwd = dir;

			const handler = registeredEvents.session_start;
			await handler(null, mockCtx);

			expect(mockNotify).toHaveBeenCalledWith(
				expect.stringContaining("1 dirty packages detected"),
				"warning",
			);
		});

		it("does not notify when all clean", async () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
			writeFileSync(join(dir, ".pi-map.md"), "# .\n## dirty\n-\n");
			mockCtx.cwd = dir;

			const handler = registeredEvents.session_start;
			await handler(null, mockCtx);

			expect(mockNotify).not.toHaveBeenCalled();
		});

		it("does nothing when no maps exist", async () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
			mockCtx.cwd = dir;

			const handler = registeredEvents.session_start;
			await handler(null, mockCtx);

			expect(mockNotify).not.toHaveBeenCalled();
		});
	});

	describe("before_agent_start event", () => {
		it("injects hint when .pi-map.md files exist", async () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
			writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTest\n");
			mockCtx.cwd = dir;

			const handler = registeredEvents.before_agent_start;
			const result = await handler(null, mockCtx);

			expect(result).toHaveProperty("message");
			expect(result.message.content).toContain("project_map_patch");
			expect(result.message.display).toBe(false);
		});

		it("returns empty object when no maps exist", async () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
			mockCtx.cwd = dir;

			const handler = registeredEvents.before_agent_start;
			const result = await handler(null, mockCtx);

			expect(result).toEqual({});
		});
	});
});