feat(prompt): implement prompt injection slice 1

2026-06-11 17:01:10 +02:00
parent c6064f8d94
commit 56560d9d56
6 changed files with 335 additions and 3 deletions
@@ -271,9 +271,13 @@ describe("pi-extension", () => {
 	});

 	describe("before_agent_start event", () => {
-		it("injects layered protocol hint when project map files exist", async () => {
+		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;
@@ -286,8 +290,12 @@ describe("pi-extension", () => {
 			expect(result.message.display).toBe(false);
 		});

-		it("returns empty object when no maps exist", async () => {
+		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;
@@ -295,5 +303,86 @@ describe("pi-extension", () => {

 			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`");
+		});
 	});
 });
@@ -0,0 +1,100 @@
+import { describe, it, expect } from "vitest";
+import {
+	buildRootPairBlock,
+	hasRootPairMarker,
+	buildPreInitHint,
+	computeInjectionBudget,
+	modeAllowsPreInitHint,
+	modeAllowsInjection,
+	modeRequiresProtocolPath,
+	ROOT_PAIR_START_MARKER,
+	ROOT_PAIR_END_MARKER,
+	TRUST_BOUNDARY_TEXT,
+} from "../src/prompt-injection.js";
+
+describe("prompt-injection helpers", () => {
+	describe("canonical markers", () => {
+		it("builds a root-pair block with markers", () => {
+			const block = buildRootPairBlock("index content", "map content");
+			expect(block).toContain(ROOT_PAIR_START_MARKER);
+			expect(block).toContain(ROOT_PAIR_END_MARKER);
+			expect(block).toContain(TRUST_BOUNDARY_TEXT);
+			expect(block).toContain("index content");
+			expect(block).toContain("map content");
+		});
+
+		it("detects root-pair marker in content", () => {
+			const block = buildRootPairBlock("i", "m");
+			expect(hasRootPairMarker(block)).toBe(true);
+			expect(hasRootPairMarker("plain text")).toBe(false);
+		});
+	});
+
+	describe("pre-init hint", () => {
+		it("returns a visible hint with no synthetic artifacts", () => {
+			const hint = buildPreInitHint();
+			expect(hint).toContain("project_map_init");
+			expect(hint).toContain("📋");
+			expect(hint).not.toContain(ROOT_PAIR_START_MARKER);
+			expect(hint).not.toContain("### Root index");
+			expect(hint).not.toContain("### Root map");
+		});
+	});
+
+	describe("budget calculation", () => {
+		it("uses relative percent when smaller than absolute cap", () => {
+			const budget = computeInjectionBudget(
+				{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
+				200_000,
+			);
+			expect(budget).toBe(30_000);
+		});
+
+		it("uses absolute cap when smaller than relative percent", () => {
+			const budget = computeInjectionBudget(
+				{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
+				10_000_000,
+			);
+			expect(budget).toBe(100_000);
+		});
+
+		it("falls back to absolute cap when context window is unknown", () => {
+			const budget = computeInjectionBudget(
+				{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
+				undefined,
+			);
+			expect(budget).toBe(100_000);
+		});
+
+		it("falls back to absolute cap when context window is zero", () => {
+			const budget = computeInjectionBudget(
+				{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
+				0,
+			);
+			expect(budget).toBe(100_000);
+		});
+	});
+
+	describe("mode helpers", () => {
+		it("allows pre-init hint for all modes except off", () => {
+			expect(modeAllowsPreInitHint("off")).toBe(false);
+			expect(modeAllowsPreInitHint("advisory")).toBe(true);
+			expect(modeAllowsPreInitHint("strong")).toBe(true);
+			expect(modeAllowsPreInitHint("strict")).toBe(true);
+		});
+
+		it("allows injection only for strong and strict", () => {
+			expect(modeAllowsInjection("off")).toBe(false);
+			expect(modeAllowsInjection("advisory")).toBe(false);
+			expect(modeAllowsInjection("strong")).toBe(true);
+			expect(modeAllowsInjection("strict")).toBe(true);
+		});
+
+		it("requires protocol path only for strict", () => {
+			expect(modeRequiresProtocolPath("off")).toBe(false);
+			expect(modeRequiresProtocolPath("advisory")).toBe(false);
+			expect(modeRequiresProtocolPath("strong")).toBe(false);
+			expect(modeRequiresProtocolPath("strict")).toBe(true);
+		});
+	});
+});