diff --git a/pi-extension.ts b/pi-extension.ts index 0824539..b783170 100644 --- a/pi-extension.ts +++ b/pi-extension.ts @@ -8,7 +8,11 @@ import { validateMaps, reinitPath, retrieveContext, + buildPreInitHint, + modeAllowsPreInitHint, + modeAllowsInjection, } from "./src/index.js"; +import { loadConfig } from "./src/config.js"; import { createLLMClient } from "./src/llm/llm-client.js"; import { LLMError } from "./src/llm/llm-error.js"; @@ -340,8 +344,33 @@ export default function (pi: ExtensionAPI) { // Inject maintenance instructions before agent starts pi.on("before_agent_start", async (_event, _ctx) => { + const config = loadConfig(_ctx.cwd); const mapFiles = findPiMapFiles(_ctx.cwd); - if (mapFiles.length === 0) return {}; + + // Mode is off: no injection at all + if (config.promptInjectionMode === "off") { + return {}; + } + + // No maps exist yet: show visible pre-init hint, but only if mode allows it + if (mapFiles.length === 0) { + if (!modeAllowsPreInitHint(config.promptInjectionMode)) { + return {}; + } + return { + message: { + customType: "pi-project-map-hint", + content: buildPreInitHint(), + display: true, + }, + }; + } + + // 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 {}; + } return { message: { diff --git a/src/config.ts b/src/config.ts index 6926d0e..931df96 100644 --- a/src/config.ts +++ b/src/config.ts @@ -1,6 +1,8 @@ import { existsSync, readFileSync } from "fs"; import { join } from "path"; +export type PromptInjectionMode = "off" | "advisory" | "strong" | "strict"; + export interface SkillConfig { ignorePatterns: string[]; smallPackageThreshold: number; @@ -11,6 +13,9 @@ export interface SkillConfig { autoInjectPrompt: boolean; tagCap: number; workflowHintCap: number; + promptInjectionMode: PromptInjectionMode; + contextBudgetPercent: number; + contextBudgetMaxTokens: number; } export const DEFAULT_CONFIG: SkillConfig = { @@ -43,6 +48,9 @@ export const DEFAULT_CONFIG: SkillConfig = { autoInjectPrompt: true, tagCap: 8, workflowHintCap: 5, + promptInjectionMode: "strong", + contextBudgetPercent: 15, + contextBudgetMaxTokens: 100_000, }; export function loadConfig(cwd: string = process.cwd()): SkillConfig { diff --git a/src/index.ts b/src/index.ts index 24ee729..37facc6 100644 --- a/src/index.ts +++ b/src/index.ts @@ -23,3 +23,15 @@ export { } from "./directory-model.js"; export { populateRoutingMetadata } from "./routing-metadata.js"; export { retrieveContext } from "./retrieve.js"; +export { + buildRootPairBlock, + hasRootPairMarker, + buildPreInitHint, + computeInjectionBudget, + modeAllowsPreInitHint, + modeAllowsInjection, + modeRequiresProtocolPath, + ROOT_PAIR_START_MARKER, + ROOT_PAIR_END_MARKER, + TRUST_BOUNDARY_TEXT, +} from "./prompt-injection.js"; diff --git a/src/prompt-injection.ts b/src/prompt-injection.ts new file mode 100644 index 0000000..7d89754 --- /dev/null +++ b/src/prompt-injection.ts @@ -0,0 +1,94 @@ +import type { SkillConfig, PromptInjectionMode } from "./config.js"; + +/** + * Canonical markers for injected root-pair content. + * These must be stable across turns and easy to scan in outgoing context. + */ +export const ROOT_PAIR_START_MARKER = ""; +export const ROOT_PAIR_END_MARKER = ""; +export const TRUST_BOUNDARY_TEXT = + "Trust boundary: index routes, map orients, source decides."; + +/** + * Build a canonical root-pair block wrapping index and map content. + */ +export function buildRootPairBlock( + indexContent: string, + mapContent: string, +): string { + return [ + ROOT_PAIR_START_MARKER, + "## Project Map Protocol", + "", + "1. Read this protocol and the root `.pi-map.index.md` first.", + "", + TRUST_BOUNDARY_TEXT, + "", + "### Root index", + indexContent, + "", + "### Root map", + mapContent, + ROOT_PAIR_END_MARKER, + ].join("\n"); +} + +/** + * Check whether a message content string contains the canonical root-pair marker. + */ +export function hasRootPairMarker(content: string): boolean { + return content.includes(ROOT_PAIR_START_MARKER); +} + +/** + * Build a lightweight user-visible pre-init startup hint. + * No synthetic artifact content is injected — only a prompt to run init. + */ +export function buildPreInitHint(): string { + return [ + "📋 Project maps not initialized.", + "", + "Run `project_map_init` to generate paired `.pi-map.md` and `.pi-map.index.md` artifacts for this project.", + "After init, the agent will automatically use the root index for routing and the root map for orientation.", + ].join("\n"); +} + +/** + * Compute the effective injection budget in tokens. + * Uses the smaller of (percent of context window) and absolute cap. + * Falls back to absolute cap if context window is unknown. + */ +export function computeInjectionBudget( + config: Pick, + contextWindow?: number, +): number { + const absolute = config.contextBudgetMaxTokens; + if (contextWindow === undefined || contextWindow <= 0) { + return absolute; + } + const relative = Math.floor( + (contextWindow * config.contextBudgetPercent) / 100, + ); + return Math.min(relative, absolute); +} + +/** + * Determine whether the active mode permits pre-init hints. + */ +export function modeAllowsPreInitHint(mode: PromptInjectionMode): boolean { + return mode !== "off"; +} + +/** + * Determine whether the active mode permits automatic artifact injection after init. + */ +export function modeAllowsInjection(mode: PromptInjectionMode): boolean { + return mode === "strong" || mode === "strict"; +} + +/** + * Resolve whether a given mode requires the protocol path for sensitive actions. + */ +export function modeRequiresProtocolPath(mode: PromptInjectionMode): boolean { + return mode === "strict"; +} diff --git a/tests/pi-extension.test.ts b/tests/pi-extension.test.ts index bea8aad..6a9d229 100644 --- a/tests/pi-extension.test.ts +++ b/tests/pi-extension.test.ts @@ -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`"); + }); }); }); diff --git a/tests/prompt-injection.test.ts b/tests/prompt-injection.test.ts new file mode 100644 index 0000000..bc3d174 --- /dev/null +++ b/tests/prompt-injection.test.ts @@ -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); + }); + }); +});