diff --git a/pi-extension.ts b/pi-extension.ts index b783170..436af05 100644 --- a/pi-extension.ts +++ b/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: - "📋 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: false, + content: payload.content, + display: payload.display, }, }; }); diff --git a/src/index.ts b/src/index.ts index 37facc6..d61f6f0 100644 --- a/src/index.ts +++ b/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, diff --git a/src/prompt-injection.ts b/src/prompt-injection.ts index 7d89754..32de23e 100644 --- a/src/prompt-injection.ts +++ b/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 = { + "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, + 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 }; +} diff --git a/tests/pi-extension.test.ts b/tests/pi-extension.test.ts index 6a9d229..d0567e3 100644 --- a/tests/pi-extension.test.ts +++ b/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( + "", + ); + expect(result.message.content).toContain(""); + 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( diff --git a/tests/prompt-injection.test.ts b/tests/prompt-injection.test.ts index bc3d174..c04a807 100644 --- a/tests/prompt-injection.test.ts +++ b/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"); + }); + }); });