feat(prompt): implement prompt injection slice 2

tests/pi-extension.test.ts

@@ -290,6 +290,33 @@ describe("pi-extension", () => {
 			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(

tests/prompt-injection.test.ts

@@ -1,4 +1,7 @@
 import { describe, it, expect } from "vitest";
+import { mkdtempSync, writeFileSync, mkdirSync } from "fs";
+import { join } from "path";
+import { tmpdir } from "os";
 import {
 	buildRootPairBlock,
 	hasRootPairMarker,
@@ -7,6 +10,10 @@ import {
 	modeAllowsPreInitHint,
 	modeAllowsInjection,
 	modeRequiresProtocolPath,
+	discoverContextWindow,
+	estimateTokens,
+	findAllArtifactPairs,
+	buildInjectionPayload,
 	ROOT_PAIR_START_MARKER,
 	ROOT_PAIR_END_MARKER,
 	TRUST_BOUNDARY_TEXT,
@@ -97,4 +104,133 @@ describe("prompt-injection helpers", () => {
 			expect(modeRequiresProtocolPath("strict")).toBe(true);
 		});
 	});
+
+	describe("context-window discovery", () => {
+		it("reads contextWindow from model metadata", () => {
+			const ctx = { model: { contextWindow: 256_000 } };
+			expect(discoverContextWindow(ctx)).toBe(256_000);
+		});
+
+		it("reads maxContextTokens from model metadata", () => {
+			const ctx = { model: { maxContextTokens: 128_000 } };
+			expect(discoverContextWindow(ctx)).toBe(128_000);
+		});
+
+		it("falls back to known model lookup by id", () => {
+			expect(discoverContextWindow({ model: { id: "gpt-4o-mini" } })).toBe(
+				128_000,
+			);
+			expect(discoverContextWindow({ model: { id: "kimi-for-coding" } })).toBe(
+				200_000,
+			);
+		});
+
+		it("returns undefined for unknown models", () => {
+			expect(
+				discoverContextWindow({ model: { id: "unknown-model" } }),
+			).toBeUndefined();
+		});
+
+		it("returns undefined when model is missing", () => {
+			expect(discoverContextWindow({})).toBeUndefined();
+		});
+	});
+
+	describe("token estimation", () => {
+		it("estimates tokens from character count", () => {
+			expect(estimateTokens("")).toBe(0);
+			expect(estimateTokens("abcd")).toBe(1);
+			expect(estimateTokens("abcde")).toBe(2);
+		});
+	});
+
+	describe("artifact-pair discovery", () => {
+		it("finds paired artifacts shallow-first", () => {
+			const dir = mkdtempSync(join(tmpdir(), "pi-pair-test-"));
+			mkdirSync(join(dir, "src", "nested"), { recursive: true });
+
+			writeFileSync(join(dir, ".pi-map.md"), "# root");
+			writeFileSync(join(dir, ".pi-map.index.md"), "# root index");
+			writeFileSync(join(dir, "src", ".pi-map.md"), "# src");
+			writeFileSync(join(dir, "src", ".pi-map.index.md"), "# src index");
+			writeFileSync(join(dir, "src", "nested", ".pi-map.md"), "# nested");
+			writeFileSync(
+				join(dir, "src", "nested", ".pi-map.index.md"),
+				"# nested index",
+			);
+
+			const pairs = findAllArtifactPairs(dir);
+			expect(pairs.map((p) => p.dir)).toEqual([".", "src", "src/nested"]);
+		});
+
+		it("ignores directories missing one of the paired artifacts", () => {
+			const dir = mkdtempSync(join(tmpdir(), "pi-pair-test-"));
+			writeFileSync(join(dir, ".pi-map.md"), "# root");
+			// missing .pi-map.index.md intentionally
+
+			expect(findAllArtifactPairs(dir)).toHaveLength(0);
+		});
+	});
+
+	describe("injection payload", () => {
+		it("always includes the root pair", () => {
+			const dir = mkdtempSync(join(tmpdir(), "pi-payload-test-"));
+			writeFileSync(join(dir, ".pi-map.md"), "# root map");
+			writeFileSync(join(dir, ".pi-map.index.md"), "# root index");
+
+			const payload = buildInjectionPayload(
+				dir,
+				{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
+				undefined,
+			);
+
+			expect(payload.content).toContain(ROOT_PAIR_START_MARKER);
+			expect(payload.content).toContain(ROOT_PAIR_END_MARKER);
+			expect(payload.content).toContain("# root map");
+			expect(payload.content).toContain("# root index");
+			expect(payload.display).toBe(false);
+		});
+
+		it("expands additional pairs when budget allows", () => {
+			const dir = mkdtempSync(join(tmpdir(), "pi-payload-test-"));
+			mkdirSync(join(dir, "src"), { recursive: true });
+			writeFileSync(join(dir, ".pi-map.md"), "# root map");
+			writeFileSync(join(dir, ".pi-map.index.md"), "# root index");
+			writeFileSync(join(dir, "src", ".pi-map.md"), "# src map");
+			writeFileSync(join(dir, "src", ".pi-map.index.md"), "# src index");
+
+			const payload = buildInjectionPayload(
+				dir,
+				{ contextBudgetPercent: 100, contextBudgetMaxTokens: 10_000 },
+				10_000,
+			);
+
+			expect(payload.content).toContain("# src map");
+			expect(payload.content).toContain("# src index");
+		});
+
+		it("stops expanding when budget is exhausted", () => {
+			const dir = mkdtempSync(join(tmpdir(), "pi-payload-test-"));
+			mkdirSync(join(dir, "src"), { recursive: true });
+			writeFileSync(join(dir, ".pi-map.md"), "# root map");
+			writeFileSync(join(dir, ".pi-map.index.md"), "# root index");
+			writeFileSync(
+				join(dir, "src", ".pi-map.md"),
+				`# src map ${"x".repeat(400)}`,
+			);
+			writeFileSync(
+				join(dir, "src", ".pi-map.index.md"),
+				`# src index ${"x".repeat(400)}`,
+			);
+
+			const payload = buildInjectionPayload(
+				dir,
+				{ contextBudgetPercent: 100, contextBudgetMaxTokens: 10 },
+				10_000,
+			);
+
+			expect(payload.content).toContain("# root map");
+			expect(payload.content).not.toContain("# src map x");
+		});
+	});
 });