416 lines
14 KiB
TypeScript
416 lines
14 KiB
TypeScript
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 () => ({
|
|
ok: true,
|
|
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 5 tools", () => {
|
|
expect(Object.keys(registeredTools)).toHaveLength(5);
|
|
expect(registeredTools).toHaveProperty("project_map_init");
|
|
expect(registeredTools).toHaveProperty("project_map_patch");
|
|
expect(registeredTools).toHaveProperty("project_map_validate");
|
|
expect(registeredTools).toHaveProperty("project_map_reinit");
|
|
expect(registeredTools).toHaveProperty("project_map_context");
|
|
});
|
|
|
|
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 and index exist and are up to date", async () => {
|
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
|
writeFileSync(
|
|
join(dir, ".pi-map.md"),
|
|
"# .\ndir: .\n\nindex: ./.pi-map.index.md\n\n## Project Map Protocol\n\n1. Read this protocol and the root `.pi-map.index.md` first.\n\nTrust boundary: index routes, map orients, source decides.\n\n## role\nTest\n## files\n-\n## arch\nTest arch\n## tags\n-\n## symbols\n-\n## workflows\n-\n## dirty\n-\n",
|
|
);
|
|
writeFileSync(
|
|
join(dir, ".pi-map.index.md"),
|
|
"# . (index)\ndir: .\n\n## Project Map Protocol\n\n1. Read this protocol and the root `.pi-map.index.md` first.\n\nTrust boundary: index routes, map orients, source decides.\n\n## role\nTest\n## parent\n-\n## children\n-\n## files\n-\n## links\nindex: ./.pi-map.index.md\nmap: ./.pi-map.md\n## workflows\n-\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("project_map_context tool", () => {
|
|
it("returns a context bundle for a query", async () => {
|
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
|
// Create a minimal map + index
|
|
writeFileSync(
|
|
join(dir, ".pi-map.md"),
|
|
"# .\ndir: .\n\nindex: ./.pi-map.index.md\n\n## Project Map Protocol\n\nTrust boundary: index routes, map orients, source decides.\n\n## role\nTest project\n## files\n- test.ts | Test file\n## arch\nTest\n## tags\ntest\n## symbols\n- main\n## workflows\n-\n## dirty\n-\n",
|
|
);
|
|
writeFileSync(
|
|
join(dir, ".pi-map.index.md"),
|
|
"# . (index)\ndir: .\n\n## role\nTest project\n## parent\n-\n## children\n-\n## files\n- test.ts\n## links\nindex: ./.pi-map.index.md\nmap: ./.pi-map.md\n## workflows\n-\n## dirty\n-\n",
|
|
);
|
|
mockCtx.cwd = dir;
|
|
|
|
const tool = registeredTools.project_map_context;
|
|
const result = await tool.execute(
|
|
"tool-1",
|
|
{ query: "test" },
|
|
null,
|
|
null,
|
|
mockCtx,
|
|
);
|
|
expect(result.details.success).toBe(true);
|
|
expect(result.content[0].text).toContain("# Context bundle: test");
|
|
expect(result.content[0].text).toContain("## relevant indexes");
|
|
expect(result.content[0].text).toContain("## instructions");
|
|
});
|
|
|
|
it("returns no-results bundle when no maps match", async () => {
|
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
|
mockCtx.cwd = dir;
|
|
|
|
const tool = registeredTools.project_map_context;
|
|
const result = await tool.execute(
|
|
"tool-1",
|
|
{ query: "nonexistent" },
|
|
null,
|
|
null,
|
|
mockCtx,
|
|
);
|
|
expect(result.details.success).toBe(true);
|
|
expect(result.content[0].text).toContain("No relevant directories found");
|
|
});
|
|
});
|
|
|
|
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 hidden hint when maps exist and mode is strong", async () => {
|
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
|
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTest\n");
|
|
writeFileSync(
|
|
join(dir, ".pi-project-map.json"),
|
|
JSON.stringify({ promptInjectionMode: "strong" }),
|
|
);
|
|
mockCtx.cwd = dir;
|
|
|
|
const handler = registeredEvents.before_agent_start;
|
|
const result = await handler(null, mockCtx);
|
|
|
|
expect(result).toHaveProperty("message");
|
|
expect(result.message.content).toContain("root `.pi-map.index.md`");
|
|
expect(result.message.content).toContain("project_map_patch");
|
|
expect(result.message.content).toContain("project_map_validate");
|
|
expect(result.message.display).toBe(false);
|
|
});
|
|
|
|
it("injects canonical root-pair block with real file content when maps exist and mode is strong", async () => {
|
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
|
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nRootMap\n");
|
|
writeFileSync(
|
|
join(dir, ".pi-map.index.md"),
|
|
"# . (index)\n## role\nRootIndex\n",
|
|
);
|
|
writeFileSync(
|
|
join(dir, ".pi-project-map.json"),
|
|
JSON.stringify({ promptInjectionMode: "strong" }),
|
|
);
|
|
mockCtx.cwd = dir;
|
|
|
|
const handler = registeredEvents.before_agent_start;
|
|
const result = await handler(null, mockCtx);
|
|
|
|
expect(result).toHaveProperty("message");
|
|
expect(result.message.display).toBe(false);
|
|
expect(result.message.content).toContain(
|
|
"<!-- PI_MAP_ROOT_PAIR_START -->",
|
|
);
|
|
expect(result.message.content).toContain("<!-- PI_MAP_ROOT_PAIR_END -->");
|
|
expect(result.message.content).toContain("RootMap");
|
|
expect(result.message.content).toContain("RootIndex");
|
|
expect(result.message.content).toContain("Trust boundary:");
|
|
});
|
|
|
|
it("returns empty object when no maps exist and mode is off", async () => {
|
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
|
writeFileSync(
|
|
join(dir, ".pi-project-map.json"),
|
|
JSON.stringify({ promptInjectionMode: "off" }),
|
|
);
|
|
mockCtx.cwd = dir;
|
|
|
|
const handler = registeredEvents.before_agent_start;
|
|
const result = await handler(null, mockCtx);
|
|
|
|
expect(result).toEqual({});
|
|
});
|
|
|
|
it("shows visible pre-init hint when no maps exist and mode is strong", async () => {
|
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
|
writeFileSync(
|
|
join(dir, ".pi-project-map.json"),
|
|
JSON.stringify({ promptInjectionMode: "strong" }),
|
|
);
|
|
mockCtx.cwd = dir;
|
|
|
|
const handler = registeredEvents.before_agent_start;
|
|
const result = await handler(null, mockCtx);
|
|
|
|
expect(result).toHaveProperty("message");
|
|
expect(result.message.display).toBe(true);
|
|
expect(result.message.content).toContain("project_map_init");
|
|
expect(result.message.content).toContain("📋");
|
|
});
|
|
|
|
it("shows visible pre-init hint when no maps exist and mode is advisory", async () => {
|
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
|
writeFileSync(
|
|
join(dir, ".pi-project-map.json"),
|
|
JSON.stringify({ promptInjectionMode: "advisory" }),
|
|
);
|
|
mockCtx.cwd = dir;
|
|
|
|
const handler = registeredEvents.before_agent_start;
|
|
const result = await handler(null, mockCtx);
|
|
|
|
expect(result).toHaveProperty("message");
|
|
expect(result.message.display).toBe(true);
|
|
expect(result.message.content).toContain("project_map_init");
|
|
});
|
|
|
|
it("returns empty object when maps exist and mode is advisory", async () => {
|
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
|
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTest\n");
|
|
writeFileSync(
|
|
join(dir, ".pi-project-map.json"),
|
|
JSON.stringify({ promptInjectionMode: "advisory" }),
|
|
);
|
|
mockCtx.cwd = dir;
|
|
|
|
const handler = registeredEvents.before_agent_start;
|
|
const result = await handler(null, mockCtx);
|
|
|
|
expect(result).toEqual({});
|
|
});
|
|
|
|
it("shows visible pre-init hint when no maps exist and mode is strict", async () => {
|
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
|
writeFileSync(
|
|
join(dir, ".pi-project-map.json"),
|
|
JSON.stringify({ promptInjectionMode: "strict" }),
|
|
);
|
|
mockCtx.cwd = dir;
|
|
|
|
const handler = registeredEvents.before_agent_start;
|
|
const result = await handler(null, mockCtx);
|
|
|
|
expect(result).toHaveProperty("message");
|
|
expect(result.message.display).toBe(true);
|
|
expect(result.message.content).toContain("project_map_init");
|
|
});
|
|
|
|
it("injects hidden hint when maps exist and mode is strict", async () => {
|
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
|
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTest\n");
|
|
writeFileSync(
|
|
join(dir, ".pi-project-map.json"),
|
|
JSON.stringify({ promptInjectionMode: "strict" }),
|
|
);
|
|
mockCtx.cwd = dir;
|
|
|
|
const handler = registeredEvents.before_agent_start;
|
|
const result = await handler(null, mockCtx);
|
|
|
|
expect(result).toHaveProperty("message");
|
|
expect(result.message.display).toBe(false);
|
|
expect(result.message.content).toContain("root `.pi-map.index.md`");
|
|
});
|
|
});
|
|
});
|