feat(prompt): implement prompt injection slice 2

pi-extension.ts

@@ -11,6 +11,8 @@ import {
 	buildPreInitHint,
 	modeAllowsPreInitHint,
 	modeAllowsInjection,
+	discoverContextWindow,
+	buildInjectionPayload,
 } from "./src/index.js";
 import { loadConfig } from "./src/config.js";
 import { createLLMClient } from "./src/llm/llm-client.js";
@@ -367,17 +369,18 @@ export default function (pi: ExtensionAPI) {
 		}
 
 		// Maps exist but advisory mode does not permit automatic artifact injection yet.
-		// Full root-pair injection with canonical markers comes in later slices.
 		if (!modeAllowsInjection(config.promptInjectionMode)) {
 			return {};
 		}
 
+		// Slice 2: post-init root-pair preload + budgeted expansion
+		const contextWindow = discoverContextWindow(_ctx);
+		const payload = buildInjectionPayload(_ctx.cwd, config, contextWindow);
 		return {
 			message: {
 				customType: "pi-project-map-hint",
-				content:
+				content: payload.content,
-					"📋 Project map active: Start with the root `.pi-map.index.md`, use indexes first for routing, read the local `.pi-map.md` plus source before edits, run `project_map_patch` after source edits, and run `project_map_validate` before freshness-sensitive architectural handoff.",
+				display: payload.display,
-				display: false,
 			},
 		};
 	});

src/index.ts

@@ -31,6 +31,10 @@ export {
 	modeAllowsPreInitHint,
 	modeAllowsInjection,
 	modeRequiresProtocolPath,
+	discoverContextWindow,
+	estimateTokens,
+	findAllArtifactPairs,
+	buildInjectionPayload,
 	ROOT_PAIR_START_MARKER,
 	ROOT_PAIR_END_MARKER,
 	TRUST_BOUNDARY_TEXT,

src/prompt-injection.ts

@@ -1,3 +1,5 @@
+import { readdirSync, statSync, readFileSync } from "fs";
+import { join, relative } from "path";
 import type { SkillConfig, PromptInjectionMode } from "./config.js";
 
 /**
@@ -92,3 +94,179 @@ export function modeAllowsInjection(mode: PromptInjectionMode): boolean {
 export function modeRequiresProtocolPath(mode: PromptInjectionMode): boolean {
 	return mode === "strict";
 }
+
+// ---------------------------------------------------------------------------
+// Context-window discovery (Slice 2)
+// ---------------------------------------------------------------------------
+
+const KNOWN_MODEL_CONTEXT_WINDOWS: Record<string, number> = {
+	"gpt-4o": 128_000,
+	"gpt-4o-mini": 128_000,
+	"gpt-4-turbo": 128_000,
+	"gpt-4": 8_192,
+	"claude-3-5-sonnet": 200_000,
+	"claude-3-opus": 200_000,
+	"kimi-for-coding": 200_000,
+	k2p6: 1_000_000,
+	"kimi-k2-thinking": 256_000,
+};
+
+/**
+ * Discover the active model's context-window size from Pi runtime metadata.
+ * Returns `undefined` when unavailable so callers fall back to the absolute cap.
+ */
+export function discoverContextWindow(ctx: any): number | undefined {
+	const model = ctx?.model;
+	if (model) {
+		if (typeof model.contextWindow === "number" && model.contextWindow > 0) {
+			return model.contextWindow;
+		}
+		if (
+			typeof model.maxContextTokens === "number" &&
+			model.maxContextTokens > 0
+		) {
+			return model.maxContextTokens;
+		}
+		if (typeof model.id === "string") {
+			const known = KNOWN_MODEL_CONTEXT_WINDOWS[model.id];
+			if (known) return known;
+		}
+	}
+	return undefined;
+}
+
+// ---------------------------------------------------------------------------
+// Token estimation (best-effort, deterministic)
+// ---------------------------------------------------------------------------
+
+/**
+ * Rough token estimate from character count.
+ * 1 token ≈ 4 chars for English/prose is a conservative heuristic.
+ */
+export function estimateTokens(text: string): number {
+	return Math.ceil(text.length / 4);
+}
+
+// ---------------------------------------------------------------------------
+// Artifact-pair discovery
+// ---------------------------------------------------------------------------
+
+export interface ArtifactPair {
+	dir: string;
+	mapPath: string;
+	indexPath: string;
+}
+
+/**
+ * Find every directory under `cwd` that contains both `.pi-map.md` and
+ * `.pi-map.index.md`.  Returns shallowest directories first for predictable
+ * structural coverage.
+ */
+export function findAllArtifactPairs(cwd: string): ArtifactPair[] {
+	const results: ArtifactPair[] = [];
+
+	function walk(dir: string) {
+		let entries: import("fs").Dirent[];
+		try {
+			entries = readdirSync(dir, { withFileTypes: true });
+		} catch {
+			return;
+		}
+		for (const entry of entries) {
+			if (
+				entry.isDirectory() &&
+				!entry.name.startsWith(".") &&
+				entry.name !== "node_modules"
+			) {
+				walk(join(dir, entry.name));
+			}
+		}
+		const mapPath = join(dir, ".pi-map.md");
+		const indexPath = join(dir, ".pi-map.index.md");
+		try {
+			statSync(mapPath);
+			statSync(indexPath);
+			results.push({
+				dir: relative(cwd, dir) || ".",
+				mapPath: relative(cwd, mapPath),
+				indexPath: relative(cwd, indexPath),
+			});
+		} catch {
+			// skip dirs without the paired artifacts
+		}
+	}
+
+	walk(cwd);
+
+	// shallow-first: sort by path depth then alphabetically
+	results.sort((a, b) => {
+		const depthA = a.dir.split(/[/\\]/).length;
+		const depthB = b.dir.split(/[/\\]/).length;
+		if (depthA !== depthB) return depthA - depthB;
+		return a.dir.localeCompare(b.dir);
+	});
+
+	return results;
+}
+
+// ---------------------------------------------------------------------------
+// Budgeted expansion
+// ---------------------------------------------------------------------------
+
+/**
+ * Build the full injection payload for a project that already has artifacts.
+ *
+ * Guarantees:
+ * 1. Root pair is always present first.
+ * 2. Additional pairs are appended shallow-first while the budget allows.
+ * 3. A brief maintenance reminder is prepended.
+ */
+export function buildInjectionPayload(
+	cwd: string,
+	config: Pick<SkillConfig, "contextBudgetPercent" | "contextBudgetMaxTokens">,
+	contextWindow: number | undefined,
+): { content: string; display: false } {
+	const budget = computeInjectionBudget(config, contextWindow);
+
+	const pairs = findAllArtifactPairs(cwd);
+	const rootIndex = pairs.findIndex((p) => p.dir === ".");
+	let rootPair: ArtifactPair | undefined;
+	if (rootIndex >= 0) {
+		rootPair = pairs.splice(rootIndex, 1)[0];
+	}
+
+	let usedTokens = 0;
+	const parts: string[] = [];
+
+	// Maintenance reminder (lightweight)
+	const reminder =
+		"📋 Project map active: Start with the root `.pi-map.index.md`, use indexes first for routing, read the local `.pi-map.md` plus source before edits, run `project_map_patch` after source edits, and run `project_map_validate` before freshness-sensitive architectural handoff.";
+	parts.push(reminder);
+	usedTokens += estimateTokens(reminder);
+
+	// Root pair (guaranteed)
+	if (rootPair) {
+		const indexContent = readFileSync(join(cwd, rootPair.indexPath), "utf8");
+		const mapContent = readFileSync(join(cwd, rootPair.mapPath), "utf8");
+		const block = buildRootPairBlock(indexContent, mapContent);
+		parts.push(block);
+		usedTokens += estimateTokens(block);
+	}
+
+	// Budgeted expansion (shallow-first, deterministic)
+	for (const pair of pairs) {
+		const indexContent = readFileSync(join(cwd, pair.indexPath), "utf8");
+		const mapContent = readFileSync(join(cwd, pair.mapPath), "utf8");
+		const pairTokens =
+			estimateTokens(indexContent) + estimateTokens(mapContent);
+		if (usedTokens + pairTokens > budget) {
+			break;
+		}
+		parts.push(
+			`\n## ${pair.dir}\n\n### Index\n${indexContent}\n\n### Map\n${mapContent}`,
+		);
+		usedTokens += pairTokens;
+	}
+
+	return { content: parts.join("\n\n"), display: false };
+}

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");
+		});
+	});
 });