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

1230 lines
38 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";
import {
ROOT_PAIR_START_MARKER,
TRUST_BOUNDARY_TEXT,
} from "../src/prompt-injection.js";
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");
expect(registeredEvents).toHaveProperty("context");
});
});
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("shows visible advisory reminder 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).toHaveProperty("message");
expect(result.message.display).toBe(true);
expect(result.message.content).toContain("advisory mode active");
expect(result.message.content).not.toContain(ROOT_PAIR_START_MARKER);
});
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("defaults to strong mode and injects hidden root pair when no config file exists", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nDefaultMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nDefaultIndex\n",
);
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("DefaultMap");
expect(result.message.content).toContain("DefaultIndex");
});
it("returns empty object when maps exist and mode is off", 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: "off" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.before_agent_start;
const result = await handler(null, mockCtx);
expect(result).toEqual({});
});
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`");
});
it("skips reinjection when outgoing context already contains marker", 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(
{
messages: [
{ content: "<!-- PI_MAP_ROOT_PAIR_START --> already injected" },
],
type: "agent_start",
},
mockCtx,
);
expect(result).toEqual({});
});
it("reinjects when root-pair artifact changes since last check", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nOriginal\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strong" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.before_agent_start;
// First call establishes baseline mtimes and injects
await handler(null, mockCtx);
// Modify root map on disk
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nUpdated\n");
// Second call should detect the change and reinject
const result = await handler(null, mockCtx);
expect(result).toHaveProperty("message");
expect(result.message.content).toContain("Updated");
expect(result.message.content).toContain("TestIndex");
});
});
describe("context event", () => {
it("registers a context handler", () => {
expect(registeredEvents).toHaveProperty("context");
expect(typeof registeredEvents.context).toBe("function");
});
it("injects on relevant turn when marker is absent", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strong" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const result = await handler(
{ messages: [{ content: "edit this file" }], type: "edit_intent" },
mockCtx,
);
expect(result).toHaveProperty("message");
expect(result.message.content).toContain(
"<!-- PI_MAP_ROOT_PAIR_START -->",
);
expect(result.message.content).toContain("TestMap");
expect(result.message.content).toContain("TestIndex");
});
it("skips injection when marker is already present in messages", 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.context;
const result = await handler(
{
messages: [{ content: "<!-- PI_MAP_ROOT_PAIR_START -->" }],
type: "edit_intent",
},
mockCtx,
);
expect(result).toEqual({});
});
it("skips injection when marker is present in provider payload", 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.context;
const result = await handler(
{
messages: [{ content: "hi" }],
payload: { body: "<!-- PI_MAP_ROOT_PAIR_START -->" },
type: "edit_intent",
},
mockCtx,
);
expect(result).toEqual({});
});
it("skips injection for irrelevant event types", 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.context;
const result = await handler(
{ messages: [{ content: "hi" }], type: "user_chat" },
mockCtx,
);
expect(result).toEqual({});
});
it("skips injection when 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.context;
const result = await handler(
{ messages: [{ content: "edit this" }], type: "edit_intent" },
mockCtx,
);
expect(result).toEqual({});
});
it("reinjects on compaction event when marker is absent", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strong" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const result = await handler(
{ messages: [{ content: "compacting" }], type: "compaction" },
mockCtx,
);
expect(result).toHaveProperty("message");
expect(result.message.content).toContain(
"<!-- PI_MAP_ROOT_PAIR_START -->",
);
expect(result.message.content).toContain("TestMap");
expect(result.message.content).toContain("TestIndex");
});
it("skips reinjection on compaction when marker is present", 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.context;
const result = await handler(
{
messages: [{ content: "<!-- PI_MAP_ROOT_PAIR_START -->" }],
type: "compaction",
},
mockCtx,
);
expect(result).toEqual({});
});
it("reinjects on artifact_change event even when marker is present", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strong" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const result = await handler(
{
messages: [{ content: "<!-- PI_MAP_ROOT_PAIR_START -->" }],
type: "artifact_change",
},
mockCtx,
);
expect(result).toHaveProperty("message");
expect(result.message.content).toContain("TestMap");
expect(result.message.content).toContain("TestIndex");
});
it("reinjects when root-pair artifact changes since last check", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nOriginal\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strong" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
// Initialize mtimes with an irrelevant turn (no injection)
await handler(
{ messages: [{ content: "hi" }], type: "user_chat" },
mockCtx,
);
// Modify root map on disk
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nUpdated\n");
// Next call should detect the change and reinject even on an irrelevant turn
const result = await handler(
{ messages: [{ content: "hi" }], type: "user_chat" },
mockCtx,
);
expect(result).toHaveProperty("message");
expect(result.message.content).toContain("Updated");
expect(result.message.content).toContain("TestIndex");
});
it("sequence: irrelevant turn skips, first relevant turn reinjects", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strong" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const skip = await handler(
{ messages: [{ content: "hi" }], type: "user_chat" },
mockCtx,
);
expect(skip).toEqual({});
const inject = await handler(
{ messages: [{ content: "edit this file" }], type: "edit_intent" },
mockCtx,
);
expect(inject).toHaveProperty("message");
expect(inject.message.content).toContain("PI_MAP_ROOT_PAIR_START");
});
it("sequence: relevant turn reinjects, subsequent relevant turn skips when marker present", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strong" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const first = await handler(
{ messages: [{ content: "edit this file" }], type: "edit_intent" },
mockCtx,
);
expect(first).toHaveProperty("message");
expect(first.message.content).toContain("PI_MAP_ROOT_PAIR_START");
const second = await handler(
{
messages: [
{
content: `<!-- PI_MAP_ROOT_PAIR_START --> already injected`,
},
],
type: "edit_intent",
},
mockCtx,
);
expect(second).toEqual({});
});
it("sequence: edit_intent then compaction reinjects only when marker absent", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strong" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const edit = await handler(
{ messages: [{ content: "edit this file" }], type: "edit_intent" },
mockCtx,
);
expect(edit).toHaveProperty("message");
expect(edit.message.content).toContain("PI_MAP_ROOT_PAIR_START");
const compaction = await handler(
{
messages: [
{
content: `<!-- PI_MAP_ROOT_PAIR_START --> already injected`,
},
],
type: "compaction",
},
mockCtx,
);
expect(compaction).toEqual({});
});
it("sequence: architecture_sensitive after irrelevant turn reinjects once", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strong" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const skip = await handler(
{ messages: [{ content: "hi" }], type: "user_chat" },
mockCtx,
);
expect(skip).toEqual({});
const first = await handler(
{
messages: [{ content: "discuss the architecture" }],
type: "architecture_sensitive",
},
mockCtx,
);
expect(first).toHaveProperty("message");
expect(first.message.content).toContain("PI_MAP_ROOT_PAIR_START");
const second = await handler(
{
messages: [
{
content: `<!-- PI_MAP_ROOT_PAIR_START --> already injected`,
},
],
type: "architecture_sensitive",
},
mockCtx,
);
expect(second).toEqual({});
});
it("sequence: artifact_change invalidation triggers reinjection even after skipped turns", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nOriginal\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strong" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
// Establish baseline mtimes with an irrelevant turn
await handler(
{ messages: [{ content: "hi" }], type: "user_chat" },
mockCtx,
);
// Relevant turn that finds marker present skips reinjection
const skip = await handler(
{
messages: [
{
content: `<!-- PI_MAP_ROOT_PAIR_START --> injected`,
},
],
type: "edit_intent",
},
mockCtx,
);
expect(skip).toEqual({});
// Modify root artifacts on disk
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nUpdated\n");
// Even an irrelevant turn must reinject because artifacts changed
const result = await handler(
{ messages: [{ content: "hi" }], type: "user_chat" },
mockCtx,
);
expect(result).toHaveProperty("message");
expect(result.message.content).toContain("Updated");
expect(result.message.content).toContain("TestIndex");
});
it("sequence: strict mode shows bypass guard on sensitive turns when protocol path is missing", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strict" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const first = await handler(
{ messages: [{ content: "edit this file" }], type: "edit_intent" },
mockCtx,
);
expect(first).toHaveProperty("message");
expect(first.message.display).toBe(true);
expect(first.message.content).toContain("Strict project-map guard");
expect(first.message.content).toContain("protocol path missing");
const second = await handler(
{
messages: [
{
content: `${ROOT_PAIR_START_MARKER} ${TRUST_BOUNDARY_TEXT} already injected`,
},
],
type: "architecture_sensitive",
},
mockCtx,
);
expect(second).toEqual({});
const third = await handler(
{ messages: [{ content: "compacting" }], type: "compaction" },
mockCtx,
);
expect(third).toHaveProperty("message");
expect(third.message.content).toContain("PI_MAP_ROOT_PAIR_START");
});
it("sequence: strict mode behaves like strong when protocol path is present", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strict" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const first = await handler(
{
messages: [
{
content: `${ROOT_PAIR_START_MARKER} ${TRUST_BOUNDARY_TEXT}`,
},
],
type: "edit_intent",
},
mockCtx,
);
expect(first).toEqual({});
const second = await handler(
{
messages: [
{
content: `${ROOT_PAIR_START_MARKER} ${TRUST_BOUNDARY_TEXT} still present`,
},
],
type: "compaction",
},
mockCtx,
);
expect(second).toEqual({});
});
it("sequence: mixed relevant and irrelevant turns preserves deterministic behavior", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strong" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const turns = [
{ type: "user_chat", expectInject: false },
{ type: "edit_intent", expectInject: true },
{ type: "user_chat", expectInject: false },
{ type: "architecture_sensitive", expectInject: false }, // marker present from prior injection
{ type: "compaction", expectInject: false }, // marker still present
];
let markerPrefix = "";
for (const turn of turns) {
const result = await handler(
{
messages: [{ content: `${markerPrefix}${turn.type}` }],
type: turn.type,
},
mockCtx,
);
if (turn.expectInject) {
expect(result).toHaveProperty("message");
expect(result.message.content).toContain("PI_MAP_ROOT_PAIR_START");
markerPrefix = "<!-- PI_MAP_ROOT_PAIR_START --> ";
} else {
expect(result).toEqual({});
}
}
});
it("sequence: strict mode bypass guard is cleared by inline marker on the next sensitive turn", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strict" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const guard = await handler(
{ messages: [{ content: "edit this file" }], type: "edit_intent" },
mockCtx,
);
expect(guard).toHaveProperty("message");
expect(guard.message.display).toBe(true);
expect(guard.message.content).toContain("Strict project-map guard");
const cleared = await handler(
{
messages: [
{
content: "[PI_MAP_BYPASS: emergency patch] edit this file",
},
],
type: "edit_intent",
},
mockCtx,
);
expect(cleared).toHaveProperty("message");
expect(cleared.message.display).toBe(false);
expect(cleared.message.content).toContain("PI_MAP_ROOT_PAIR_START");
const later = await handler(
{ messages: [{ content: "edit another file" }], type: "edit_intent" },
mockCtx,
);
expect(later).toHaveProperty("message");
expect(later.message.display).toBe(true);
expect(later.message.content).toContain("Strict project-map guard");
});
it("sequence: strict mode forces reinjection on artifact_change even when protocol path is missing", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nOriginalMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strict" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
// Establish baseline mtimes on an irrelevant turn.
await handler(
{ messages: [{ content: "hi" }], type: "user_chat" },
mockCtx,
);
// A sensitive turn with no protocol path should be guarded.
const guard = await handler(
{ messages: [{ content: "edit this file" }], type: "edit_intent" },
mockCtx,
);
expect(guard).toHaveProperty("message");
expect(guard.message.content).toContain("Strict project-map guard");
// Modify root artifacts outside the agent.
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nUpdatedMap\n");
// Artifact invalidation must force reinjection, not re-guard.
const result = await handler(
{ messages: [{ content: "hi" }], type: "user_chat" },
mockCtx,
);
expect(result).toHaveProperty("message");
expect(result.message.display).toBe(false);
expect(result.message.content).toContain("UpdatedMap");
expect(result.message.content).toContain("PI_MAP_ROOT_PAIR_START");
});
it("defaults to strong mode and reinjects on edit_intent when no config file exists", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nDefaultMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nDefaultIndex\n",
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const result = await handler(
{ messages: [{ content: "edit this file" }], type: "edit_intent" },
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("DefaultMap");
expect(result.message.content).toContain("DefaultIndex");
});
it("off mode skips context reinjection even on edit_intent", 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: "off" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const result = await handler(
{ messages: [{ content: "edit this file" }], type: "edit_intent" },
mockCtx,
);
expect(result).toEqual({});
});
it("strict mode guards heuristic sensitive action in a generic user_chat turn", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n");
writeFileSync(
join(dir, ".pi-map.index.md"),
"# . (index)\n## role\nTestIndex\n",
);
writeFileSync(
join(dir, ".pi-project-map.json"),
JSON.stringify({ promptInjectionMode: "strict" }),
);
mockCtx.cwd = dir;
const handler = registeredEvents.context;
const result = await handler(
{ messages: [{ content: "refactor the module" }], type: "user_chat" },
mockCtx,
);
expect(result).toHaveProperty("message");
expect(result.message.display).toBe(true);
expect(result.message.content).toContain("Strict project-map guard");
expect(result.message.content).toContain("protocol path missing");
expect(result.message.content).not.toContain(ROOT_PAIR_START_MARKER);
});
it("strict mode does not guard non-sensitive user_chat turns", 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.context;
const result = await handler(
{ messages: [{ content: "what is the weather?" }], type: "user_chat" },
mockCtx,
);
expect(result).toEqual({});
});
});
});