From 19666c900eea0a96d6bce1998817d89d6f683110 Mon Sep 17 00:00:00 2001 From: Alex Blank Date: Thu, 11 Jun 2026 21:32:38 +0200 Subject: [PATCH] feat(prompt): implement prompt injection slice 3 --- pi-extension.ts | 73 +++++ src/index.ts | 11 + src/merge.ts | 12 +- src/prompt-injection.ts | 193 ++++++++++++ tests/merge.test.ts | 103 ++++++ tests/pi-extension.test.ts | 555 +++++++++++++++++++++++++++++++++ tests/prompt-injection.test.ts | 536 +++++++++++++++++++++++++++++++ 7 files changed, 1479 insertions(+), 4 deletions(-) create mode 100644 tests/merge.test.ts diff --git a/pi-extension.ts b/pi-extension.ts index 436af05..97e59c7 100644 --- a/pi-extension.ts +++ b/pi-extension.ts @@ -13,6 +13,9 @@ import { modeAllowsInjection, discoverContextWindow, buildInjectionPayload, + shouldReinjectForEvent, + getRootPairMtimes, + rootPairChanged, } from "./src/index.js"; import { loadConfig } from "./src/config.js"; import { createLLMClient } from "./src/llm/llm-client.js"; @@ -76,6 +79,7 @@ function renderProgressBar( } export default function (pi: ExtensionAPI) { + let lastRootPairMtimes: import("./src/index.js").RootPairMtimes = {}; pi.registerTool({ name: "project_map_init", label: "Project Map Init", @@ -373,6 +377,28 @@ export default function (pi: ExtensionAPI) { return {}; } + // Slice 3b: detect root-pair artifact changes + const currentMtimes = getRootPairMtimes(_ctx.cwd); + const hasPrevious = + lastRootPairMtimes.mapMtime !== undefined || + lastRootPairMtimes.indexMtime !== undefined; + const artifactChanged = + hasPrevious && rootPairChanged(currentMtimes, lastRootPairMtimes); + lastRootPairMtimes = currentMtimes; + + // Slice 3a/3b: avoid redundant reinjection by scanning outgoing context + const eventType = artifactChanged ? "artifact_change" : "agent_start"; + const decision = shouldReinjectForEvent( + { + messages: _event?.messages, + type: eventType, + }, + config.promptInjectionMode, + ); + if (!decision.needed) { + return {}; + } + // Slice 2: post-init root-pair preload + budgeted expansion const contextWindow = discoverContextWindow(_ctx); const payload = buildInjectionPayload(_ctx.cwd, config, contextWindow); @@ -384,4 +410,51 @@ export default function (pi: ExtensionAPI) { }, }; }); + + // Per-turn context scanning for reinjection in strong/strict modes + pi.on("context", async (event: any, ctx: any) => { + const config = loadConfig(ctx.cwd); + const mapFiles = findPiMapFiles(ctx.cwd); + + if (mapFiles.length === 0) { + return {}; + } + + if (!modeAllowsInjection(config.promptInjectionMode)) { + return {}; + } + + // Slice 3b: detect root-pair artifact changes + const currentMtimes = getRootPairMtimes(ctx.cwd); + const hasPrevious = + lastRootPairMtimes.mapMtime !== undefined || + lastRootPairMtimes.indexMtime !== undefined; + const artifactChanged = + hasPrevious && rootPairChanged(currentMtimes, lastRootPairMtimes); + lastRootPairMtimes = currentMtimes; + + const eventType = artifactChanged ? "artifact_change" : event?.type; + + const decision = shouldReinjectForEvent( + { + messages: event?.messages, + type: eventType, + payload: event?.payload, + }, + config.promptInjectionMode, + ); + if (!decision.needed) { + return {}; + } + + const contextWindow = discoverContextWindow(ctx); + const payload = buildInjectionPayload(ctx.cwd, config, contextWindow); + return { + message: { + customType: "pi-project-map-hint", + content: payload.content, + display: payload.display, + }, + }; + }); } diff --git a/src/index.ts b/src/index.ts index d61f6f0..610a74e 100644 --- a/src/index.ts +++ b/src/index.ts @@ -35,6 +35,17 @@ export { estimateTokens, findAllArtifactPairs, buildInjectionPayload, + outgoingMessagesHaveMarker, + providerPayloadHasMarker, + isRelevantTurnForReinjection, + shouldReinjectForEvent, + detectEditIntent, + detectArchitectureSensitiveReasoning, + getRootPairMtimes, + rootPairChanged, + type RootPairMtimes, + type RelevantTurnType, + type ReinjectDecision, ROOT_PAIR_START_MARKER, ROOT_PAIR_END_MARKER, TRUST_BOUNDARY_TEXT, diff --git a/src/merge.ts b/src/merge.ts index a9949f1..db94ece 100644 --- a/src/merge.ts +++ b/src/merge.ts @@ -48,8 +48,10 @@ export function mergeFileData( for (const cls of ast.classes) { const classExports: string[] = [`class:${cls.name}`]; for (const method of cls.methods) { - const paramStr = method.params.join(", "); - const returnStr = method.returns ? ` → ${method.returns}` : ""; + const paramStr = method.params.join(", ").replace(/\s+/g, " "); + const returnStr = method.returns + ? ` → ${method.returns.replace(/\s+/g, " ")}` + : ""; classExports.push(`method:${method.name}(${paramStr})${returnStr}`); for (const call of method.calls) { classExports.push(`call:${call}`); @@ -65,8 +67,10 @@ export function mergeFileData( // Encode top-level functions if (ast && ast.functions.length > 0) { for (const func of ast.functions) { - const paramStr = func.params.join(", "); - const returnStr = func.returns ? ` → ${func.returns}` : ""; + const paramStr = func.params.join(", ").replace(/\s+/g, " "); + const returnStr = func.returns + ? ` → ${func.returns.replace(/\s+/g, " ")}` + : ""; dedupedExports.push(`func:${func.name}(${paramStr})${returnStr}`); for (const call of func.calls) { dedupedExports.push(`call:${call}`); diff --git a/src/prompt-injection.ts b/src/prompt-injection.ts index 32de23e..657bc9a 100644 --- a/src/prompt-injection.ts +++ b/src/prompt-injection.ts @@ -2,6 +2,199 @@ import { readdirSync, statSync, readFileSync } from "fs"; import { join, relative } from "path"; import type { SkillConfig, PromptInjectionMode } from "./config.js"; +// --------------------------------------------------------------------------- +// Outgoing-context scanning (Slice 3) +// --------------------------------------------------------------------------- + +export interface ReinjectDecision { + needed: boolean; + reason?: string; +} + +/** + * Decide whether reinjection is needed for a given event. + * Checks mode, scans outgoing messages for the canonical marker, + * falls back to payload inspection, and evaluates relevant-turn triggers. + */ +export function shouldReinjectForEvent( + event: { + messages?: Array<{ content?: unknown; text?: string }>; + type?: string; + payload?: unknown; + }, + mode: PromptInjectionMode, +): ReinjectDecision { + if (mode !== "strong" && mode !== "strict") { + return { needed: false, reason: "mode_does_not_require_reinjection" }; + } + + // Slice 3b: root-pair artifact changes invalidate prior injections + if (event.type === "artifact_change") { + return { needed: true, reason: "root_pair_artifact_changed" }; + } + + if (event.messages && outgoingMessagesHaveMarker(event.messages)) { + return { needed: false, reason: "marker_present_in_messages" }; + } + + if (event.payload && providerPayloadHasMarker(event.payload)) { + return { needed: false, reason: "marker_present_in_payload" }; + } + + if (event.type && isRelevantTurnForReinjection(event.type)) { + return { needed: true, reason: "relevant_turn_and_marker_absent" }; + } + + return { needed: false, reason: "not_a_relevant_turn" }; +} + +function extractTextFromMessage(m: { + content?: unknown; + text?: string; +}): string { + if (typeof m.content === "string") return m.content; + if (Array.isArray(m.content)) { + return m.content + .map((block: any) => { + if (typeof block === "string") return block; + if (typeof block.text === "string") return block.text; + if (typeof block.content === "string") return block.content; + return ""; + }) + .join(""); + } + return m.text ?? ""; +} + +/** + * Scan an array of message-like objects for the canonical root-pair marker. + * Normalizes array-based content blocks (e.g., TextContent[]) to strings. + */ +export function outgoingMessagesHaveMarker( + messages: Array<{ content?: unknown; text?: string }>, +): boolean { + return messages.some((m) => hasRootPairMarker(extractTextFromMessage(m))); +} + +/** + * Fallback inspection of a final provider payload for the canonical marker. + * Accepts either a string or an object that will be JSON-stringified. + */ +export function providerPayloadHasMarker(payload: unknown): boolean { + if (!payload) return false; + if (typeof payload === "string") { + return hasRootPairMarker(payload); + } + try { + return hasRootPairMarker(JSON.stringify(payload)); + } catch { + return false; + } +} + +/** + * Relevant-turn types that trigger reinjection in strong mode. + */ +export type RelevantTurnType = + | "agent_start" + | "edit_intent" + | "architecture_sensitive" + | "compaction" + | "artifact_change"; + +/** + * Determine whether a given event type is a relevant reinjection trigger. + */ +export function isRelevantTurnForReinjection(eventType: string): boolean { + const relevant: RelevantTurnType[] = [ + "agent_start", + "edit_intent", + "architecture_sensitive", + "compaction", + "artifact_change", + ]; + return relevant.includes(eventType as RelevantTurnType); +} + +/** + * Heuristic detection of edit intent from message content. + */ +export function detectEditIntent( + messages: Array<{ content?: unknown; text?: string }>, +): boolean { + const combined = messages.map(extractTextFromMessage).join(" "); + const patterns = [ + /\b(edit|modify|update|change|refactor|fix|patch|rewrite|delete|remove|add)\s+(the|a|this|that|these|those|file|code|function|method|class|module|component|line)\b/i, + /\b(write|create|generate)\s+(new|a|the)\s+(file|function|class|module|component)\b/i, + /```[\s\S]*?\b(edit|modify|update|change|refactor|fix|delete|remove)\b/i, + /\bapply\s+(the|a|this|that)\s+(change|edit|patch|fix)\b/i, + ]; + return patterns.some((p) => p.test(combined)); +} + +/** + * Heuristic detection of architecture-sensitive reasoning from message content. + */ +export function detectArchitectureSensitiveReasoning( + messages: Array<{ content?: unknown; text?: string }>, +): boolean { + const combined = messages.map(extractTextFromMessage).join(" "); + const patterns = [ + /\barchitectur(e|al)\b/i, + /\bdesign\s+(decision|pattern|choice|principle|review)\b/i, + /\b(restructure|reorganize|redesign|rearchitect)\b/i, + /\b(system|high.level|macro|structural)\s+(design|architecture|overview)\b/i, + /\bdependency\s+(injection|graph|cycle|inversion)\b/i, + /\b(api|interface|contract|schema|protocol)\s+(design|change|migration|breaking)\b/i, + /\b(microservice|monolith|modular|layered|hexagonal|clean)\s+arch/i, + /\bdata\s+(model|flow|pipeline|architecture)\b/i, + /\b(scalability|performance|security|maintainability)\s+(concern|tradeoff|decision)\b/i, + ]; + return patterns.some((p) => p.test(combined)); +} + +// --------------------------------------------------------------------------- +// Root-pair artifact change detection +// --------------------------------------------------------------------------- + +export interface RootPairMtimes { + mapMtime?: number; + indexMtime?: number; +} + +/** + * Read current mtimes for the root pair artifacts, if they exist. + */ +export function getRootPairMtimes(cwd: string): RootPairMtimes { + const mapPath = join(cwd, ".pi-map.md"); + const indexPath = join(cwd, ".pi-map.index.md"); + const result: RootPairMtimes = {}; + try { + result.mapMtime = statSync(mapPath).mtimeMs; + } catch { + // ignore + } + try { + result.indexMtime = statSync(indexPath).mtimeMs; + } catch { + // ignore + } + return result; +} + +/** + * Compare current root-pair mtimes against previously recorded ones. + */ +export function rootPairChanged( + current: RootPairMtimes, + previous: RootPairMtimes, +): boolean { + return ( + current.mapMtime !== previous.mapMtime || + current.indexMtime !== previous.indexMtime + ); +} + /** * Canonical markers for injected root-pair content. * These must be stable across turns and easy to scan in outgoing context. diff --git a/tests/merge.test.ts b/tests/merge.test.ts new file mode 100644 index 0000000..ae8668f --- /dev/null +++ b/tests/merge.test.ts @@ -0,0 +1,103 @@ +import { describe, it, expect } from "vitest"; +import { mergeFileData } from "../src/merge.js"; + +describe("mergeFileData", () => { + it("collapses multi-line parameters into single-line func signatures", () => { + const result = mergeFileData( + "api.ts", + { purpose: "API helpers", exports: [], deps: [] }, + { + exports: ["shouldReinjectForEvent"], + deps: [], + classes: [], + functions: [ + { + name: "shouldReinjectForEvent", + params: [ + "event: {\n\t\tmessages?: Array<{ content?: unknown; text?: string }>;\n\t\ttype?: string;\n\t\tpayload?: unknown;\n\t}", + "mode: PromptInjectionMode", + ], + returns: "ReinjectDecision", + calls: ["outgoingMessagesHaveMarker"], + raises: [], + }, + ], + }, + ); + + const funcExport = result.exports.find((e) => + e.startsWith("func:shouldReinjectForEvent"), + ); + expect(funcExport).toBeDefined(); + expect(funcExport).not.toContain("\n"); + expect(funcExport).not.toContain("\t"); + expect(funcExport).toBe( + "func:shouldReinjectForEvent(event: { messages?: Array<{ content?: unknown; text?: string }>; type?: string; payload?: unknown; }, mode: PromptInjectionMode) → ReinjectDecision", + ); + }); + + it("collapses multi-line return types into single-line func signatures", () => { + const result = mergeFileData( + "types.ts", + { purpose: "Types", exports: [], deps: [] }, + { + exports: ["complexReturn"], + deps: [], + classes: [], + functions: [ + { + name: "complexReturn", + params: ["x: number"], + returns: "{\n a: string;\n b: number;\n}", + calls: [], + raises: [], + }, + ], + }, + ); + + const funcExport = result.exports.find((e) => + e.startsWith("func:complexReturn"), + ); + expect(funcExport).toBeDefined(); + expect(funcExport).not.toContain("\n"); + expect(funcExport).toBe( + "func:complexReturn(x: number) → { a: string; b: number; }", + ); + }); + + it("collapses multi-line method parameters into single-line method signatures", () => { + const result = mergeFileData( + "class.ts", + { purpose: "Class file", exports: [], deps: [] }, + { + exports: ["MyClass"], + deps: [], + classes: [ + { + name: "MyClass", + methods: [ + { + name: "doThing", + params: ["opts: {\n a: string;\n b: number;\n}"], + returns: "void", + calls: [], + raises: [], + }, + ], + }, + ], + functions: [], + }, + ); + + const methodExport = result.exports.find((e) => + e.startsWith("method:doThing"), + ); + expect(methodExport).toBeDefined(); + expect(methodExport).not.toContain("\n"); + expect(methodExport).toBe( + "method:doThing(opts: { a: string; b: number; }) → void", + ); + }); +}); diff --git a/tests/pi-extension.test.ts b/tests/pi-extension.test.ts index d0567e3..5864c66 100644 --- a/tests/pi-extension.test.ts +++ b/tests/pi-extension.test.ts @@ -93,6 +93,7 @@ describe("pi-extension", () => { it("registers session_start and before_agent_start events", () => { expect(registeredEvents).toHaveProperty("session_start"); expect(registeredEvents).toHaveProperty("before_agent_start"); + expect(registeredEvents).toHaveProperty("context"); }); }); @@ -411,5 +412,559 @@ describe("pi-extension", () => { expect(result.message.display).toBe(false); expect(result.message.content).toContain("root `.pi-map.index.md`"); }); + + it("skips reinjection when outgoing context already contains marker", 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; + const result = await handler( + { + messages: [ + { content: " already injected" }, + ], + type: "agent_start", + }, + mockCtx, + ); + + expect(result).toEqual({}); + }); + + it("reinjects when root-pair artifact changes since last check", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nOriginal\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nTestIndex\n", + ); + writeFileSync( + join(dir, ".pi-project-map.json"), + JSON.stringify({ promptInjectionMode: "strong" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.before_agent_start; + + // First call establishes baseline mtimes and injects + await handler(null, mockCtx); + + // Modify root map on disk + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nUpdated\n"); + + // Second call should detect the change and reinject + const result = await handler(null, mockCtx); + + expect(result).toHaveProperty("message"); + expect(result.message.content).toContain("Updated"); + expect(result.message.content).toContain("TestIndex"); + }); + }); + + describe("context event", () => { + it("registers a context handler", () => { + expect(registeredEvents).toHaveProperty("context"); + expect(typeof registeredEvents.context).toBe("function"); + }); + + it("injects on relevant turn when marker is absent", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nTestIndex\n", + ); + writeFileSync( + join(dir, ".pi-project-map.json"), + JSON.stringify({ promptInjectionMode: "strong" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.context; + const result = await handler( + { messages: [{ content: "edit this file" }], type: "edit_intent" }, + mockCtx, + ); + + expect(result).toHaveProperty("message"); + expect(result.message.content).toContain( + "", + ); + expect(result.message.content).toContain("TestMap"); + expect(result.message.content).toContain("TestIndex"); + }); + + it("skips injection when marker is already present in messages", 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.context; + const result = await handler( + { + messages: [{ content: "" }], + type: "edit_intent", + }, + mockCtx, + ); + + expect(result).toEqual({}); + }); + + it("skips injection when marker is present in provider payload", 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.context; + const result = await handler( + { + messages: [{ content: "hi" }], + payload: { body: "" }, + type: "edit_intent", + }, + mockCtx, + ); + + expect(result).toEqual({}); + }); + + it("skips injection for irrelevant event types", 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.context; + const result = await handler( + { messages: [{ content: "hi" }], type: "user_chat" }, + mockCtx, + ); + + expect(result).toEqual({}); + }); + + it("skips injection when 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.context; + const result = await handler( + { messages: [{ content: "edit this" }], type: "edit_intent" }, + mockCtx, + ); + + expect(result).toEqual({}); + }); + + it("reinjects on compaction event when marker is absent", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nTestIndex\n", + ); + writeFileSync( + join(dir, ".pi-project-map.json"), + JSON.stringify({ promptInjectionMode: "strong" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.context; + const result = await handler( + { messages: [{ content: "compacting" }], type: "compaction" }, + mockCtx, + ); + + expect(result).toHaveProperty("message"); + expect(result.message.content).toContain( + "", + ); + expect(result.message.content).toContain("TestMap"); + expect(result.message.content).toContain("TestIndex"); + }); + + it("skips reinjection on compaction when marker is present", 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.context; + const result = await handler( + { + messages: [{ content: "" }], + type: "compaction", + }, + mockCtx, + ); + + expect(result).toEqual({}); + }); + + it("reinjects on artifact_change event even when marker is present", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nTestIndex\n", + ); + writeFileSync( + join(dir, ".pi-project-map.json"), + JSON.stringify({ promptInjectionMode: "strong" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.context; + const result = await handler( + { + messages: [{ content: "" }], + type: "artifact_change", + }, + mockCtx, + ); + + expect(result).toHaveProperty("message"); + expect(result.message.content).toContain("TestMap"); + expect(result.message.content).toContain("TestIndex"); + }); + + it("reinjects when root-pair artifact changes since last check", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nOriginal\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nTestIndex\n", + ); + writeFileSync( + join(dir, ".pi-project-map.json"), + JSON.stringify({ promptInjectionMode: "strong" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.context; + + // Initialize mtimes with an irrelevant turn (no injection) + await handler( + { messages: [{ content: "hi" }], type: "user_chat" }, + mockCtx, + ); + + // Modify root map on disk + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nUpdated\n"); + + // Next call should detect the change and reinject even on an irrelevant turn + const result = await handler( + { messages: [{ content: "hi" }], type: "user_chat" }, + mockCtx, + ); + + expect(result).toHaveProperty("message"); + expect(result.message.content).toContain("Updated"); + expect(result.message.content).toContain("TestIndex"); + }); + + it("sequence: irrelevant turn skips, first relevant turn reinjects", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nTestIndex\n", + ); + writeFileSync( + join(dir, ".pi-project-map.json"), + JSON.stringify({ promptInjectionMode: "strong" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.context; + + const skip = await handler( + { messages: [{ content: "hi" }], type: "user_chat" }, + mockCtx, + ); + expect(skip).toEqual({}); + + const inject = await handler( + { messages: [{ content: "edit this file" }], type: "edit_intent" }, + mockCtx, + ); + expect(inject).toHaveProperty("message"); + expect(inject.message.content).toContain("PI_MAP_ROOT_PAIR_START"); + }); + + it("sequence: relevant turn reinjects, subsequent relevant turn skips when marker present", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nTestIndex\n", + ); + writeFileSync( + join(dir, ".pi-project-map.json"), + JSON.stringify({ promptInjectionMode: "strong" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.context; + + const first = await handler( + { messages: [{ content: "edit this file" }], type: "edit_intent" }, + mockCtx, + ); + expect(first).toHaveProperty("message"); + expect(first.message.content).toContain("PI_MAP_ROOT_PAIR_START"); + + const second = await handler( + { + messages: [ + { + content: ` already injected`, + }, + ], + type: "edit_intent", + }, + mockCtx, + ); + expect(second).toEqual({}); + }); + + it("sequence: edit_intent then compaction reinjects only when marker absent", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nTestIndex\n", + ); + writeFileSync( + join(dir, ".pi-project-map.json"), + JSON.stringify({ promptInjectionMode: "strong" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.context; + + const edit = await handler( + { messages: [{ content: "edit this file" }], type: "edit_intent" }, + mockCtx, + ); + expect(edit).toHaveProperty("message"); + expect(edit.message.content).toContain("PI_MAP_ROOT_PAIR_START"); + + const compaction = await handler( + { + messages: [ + { + content: ` already injected`, + }, + ], + type: "compaction", + }, + mockCtx, + ); + expect(compaction).toEqual({}); + }); + + it("sequence: architecture_sensitive after irrelevant turn reinjects once", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nTestIndex\n", + ); + writeFileSync( + join(dir, ".pi-project-map.json"), + JSON.stringify({ promptInjectionMode: "strong" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.context; + + const skip = await handler( + { messages: [{ content: "hi" }], type: "user_chat" }, + mockCtx, + ); + expect(skip).toEqual({}); + + const first = await handler( + { + messages: [{ content: "discuss the architecture" }], + type: "architecture_sensitive", + }, + mockCtx, + ); + expect(first).toHaveProperty("message"); + expect(first.message.content).toContain("PI_MAP_ROOT_PAIR_START"); + + const second = await handler( + { + messages: [ + { + content: ` already injected`, + }, + ], + type: "architecture_sensitive", + }, + mockCtx, + ); + expect(second).toEqual({}); + }); + + it("sequence: artifact_change invalidation triggers reinjection even after skipped turns", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nOriginal\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nTestIndex\n", + ); + writeFileSync( + join(dir, ".pi-project-map.json"), + JSON.stringify({ promptInjectionMode: "strong" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.context; + + // Establish baseline mtimes with an irrelevant turn + await handler( + { messages: [{ content: "hi" }], type: "user_chat" }, + mockCtx, + ); + + // Relevant turn that finds marker present skips reinjection + const skip = await handler( + { + messages: [ + { + content: ` injected`, + }, + ], + type: "edit_intent", + }, + mockCtx, + ); + expect(skip).toEqual({}); + + // Modify root artifacts on disk + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nUpdated\n"); + + // Even an irrelevant turn must reinject because artifacts changed + const result = await handler( + { messages: [{ content: "hi" }], type: "user_chat" }, + mockCtx, + ); + expect(result).toHaveProperty("message"); + expect(result.message.content).toContain("Updated"); + expect(result.message.content).toContain("TestIndex"); + }); + + it("sequence: strict mode reinjects on relevant turns same as strong mode", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nTestIndex\n", + ); + writeFileSync( + join(dir, ".pi-project-map.json"), + JSON.stringify({ promptInjectionMode: "strict" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.context; + + const first = await handler( + { messages: [{ content: "edit this file" }], type: "edit_intent" }, + mockCtx, + ); + expect(first).toHaveProperty("message"); + expect(first.message.content).toContain("PI_MAP_ROOT_PAIR_START"); + + const second = await handler( + { + messages: [ + { + content: ` already injected`, + }, + ], + type: "architecture_sensitive", + }, + mockCtx, + ); + expect(second).toEqual({}); + + const third = await handler( + { messages: [{ content: "compacting" }], type: "compaction" }, + mockCtx, + ); + expect(third).toHaveProperty("message"); + expect(third.message.content).toContain("PI_MAP_ROOT_PAIR_START"); + }); + + it("sequence: mixed relevant and irrelevant turns preserves deterministic behavior", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nTestIndex\n", + ); + writeFileSync( + join(dir, ".pi-project-map.json"), + JSON.stringify({ promptInjectionMode: "strong" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.context; + + const turns = [ + { type: "user_chat", expectInject: false }, + { type: "edit_intent", expectInject: true }, + { type: "user_chat", expectInject: false }, + { type: "architecture_sensitive", expectInject: false }, // marker present from prior injection + { type: "compaction", expectInject: false }, // marker still present + ]; + + let markerPrefix = ""; + for (const turn of turns) { + const result = await handler( + { + messages: [{ content: `${markerPrefix}${turn.type}` }], + type: turn.type, + }, + mockCtx, + ); + if (turn.expectInject) { + expect(result).toHaveProperty("message"); + expect(result.message.content).toContain("PI_MAP_ROOT_PAIR_START"); + markerPrefix = " "; + } else { + expect(result).toEqual({}); + } + } + }); }); }); diff --git a/tests/prompt-injection.test.ts b/tests/prompt-injection.test.ts index c04a807..5960445 100644 --- a/tests/prompt-injection.test.ts +++ b/tests/prompt-injection.test.ts @@ -14,6 +14,12 @@ import { estimateTokens, findAllArtifactPairs, buildInjectionPayload, + outgoingMessagesHaveMarker, + providerPayloadHasMarker, + shouldReinjectForEvent, + isRelevantTurnForReinjection, + detectEditIntent, + detectArchitectureSensitiveReasoning, ROOT_PAIR_START_MARKER, ROOT_PAIR_END_MARKER, TRUST_BOUNDARY_TEXT, @@ -172,6 +178,536 @@ describe("prompt-injection helpers", () => { }); }); + describe("outgoing message marker scan", () => { + it("detects marker in string content", () => { + const messages = [{ content: `hello ${ROOT_PAIR_START_MARKER} world` }]; + expect(outgoingMessagesHaveMarker(messages)).toBe(true); + }); + + it("detects marker in text field", () => { + const messages = [{ text: `before ${ROOT_PAIR_START_MARKER} after` }]; + expect(outgoingMessagesHaveMarker(messages)).toBe(true); + }); + + it("detects marker in array content blocks", () => { + const messages = [ + { + content: [ + { type: "text", text: "prefix" }, + { type: "text", text: ROOT_PAIR_START_MARKER }, + ], + }, + ]; + expect(outgoingMessagesHaveMarker(messages)).toBe(true); + }); + + it("returns false when no marker is present", () => { + const messages = [{ content: "plain text" }, { text: "more text" }]; + expect(outgoingMessagesHaveMarker(messages)).toBe(false); + }); + + it("returns false for empty messages", () => { + expect(outgoingMessagesHaveMarker([])).toBe(false); + }); + }); + + describe("provider payload marker fallback", () => { + it("detects marker in string payload", () => { + expect( + providerPayloadHasMarker(`foo ${ROOT_PAIR_START_MARKER} bar`), + ).toBe(true); + }); + + it("detects marker in JSON-stringified object", () => { + expect( + providerPayloadHasMarker({ body: { text: ROOT_PAIR_START_MARKER } }), + ).toBe(true); + }); + + it("returns false for plain string without marker", () => { + expect(providerPayloadHasMarker("plain payload")).toBe(false); + }); + + it("returns false for null payload", () => { + expect(providerPayloadHasMarker(null)).toBe(false); + }); + + it("returns false for undefined payload", () => { + expect(providerPayloadHasMarker(undefined)).toBe(false); + }); + }); + + describe("reinjection decision", () => { + it("rejects reinjection for off mode", () => { + const decision = shouldReinjectForEvent( + { messages: [{ content: "hi" }], type: "agent_start" }, + "off", + ); + expect(decision.needed).toBe(false); + expect(decision.reason).toBe("mode_does_not_require_reinjection"); + }); + + it("rejects reinjection for advisory mode", () => { + const decision = shouldReinjectForEvent( + { messages: [{ content: "hi" }], type: "agent_start" }, + "advisory", + ); + expect(decision.needed).toBe(false); + expect(decision.reason).toBe("mode_does_not_require_reinjection"); + }); + + it("rejects reinjection when marker is present in messages", () => { + const decision = shouldReinjectForEvent( + { + messages: [{ content: ROOT_PAIR_START_MARKER }], + type: "agent_start", + }, + "strong", + ); + expect(decision.needed).toBe(false); + expect(decision.reason).toBe("marker_present_in_messages"); + }); + + it("rejects reinjection when marker is present in payload", () => { + const decision = shouldReinjectForEvent( + { + messages: [{ content: "hi" }], + payload: ROOT_PAIR_START_MARKER, + type: "agent_start", + }, + "strong", + ); + expect(decision.needed).toBe(false); + expect(decision.reason).toBe("marker_present_in_payload"); + }); + + it("allows reinjection for agent_start in strong mode when marker absent", () => { + const decision = shouldReinjectForEvent( + { messages: [{ content: "hi" }], type: "agent_start" }, + "strong", + ); + expect(decision.needed).toBe(true); + expect(decision.reason).toBe("relevant_turn_and_marker_absent"); + }); + + it("allows reinjection for edit_intent in strict mode when marker absent", () => { + const decision = shouldReinjectForEvent( + { messages: [{ content: "edit this file" }], type: "edit_intent" }, + "strict", + ); + expect(decision.needed).toBe(true); + expect(decision.reason).toBe("relevant_turn_and_marker_absent"); + }); + + it("allows reinjection for architecture_sensitive in strong mode", () => { + const decision = shouldReinjectForEvent( + { + messages: [{ content: "discuss architecture" }], + type: "architecture_sensitive", + }, + "strong", + ); + expect(decision.needed).toBe(true); + expect(decision.reason).toBe("relevant_turn_and_marker_absent"); + }); + + it("rejects reinjection for irrelevant event types", () => { + const decision = shouldReinjectForEvent( + { messages: [{ content: "hi" }], type: "user_chat" }, + "strong", + ); + expect(decision.needed).toBe(false); + expect(decision.reason).toBe("not_a_relevant_turn"); + }); + + it("treats compaction as a relevant turn", () => { + expect(isRelevantTurnForReinjection("compaction")).toBe(true); + }); + + it("treats artifact_change as a relevant turn", () => { + expect(isRelevantTurnForReinjection("artifact_change")).toBe(true); + }); + + it("allows reinjection for compaction in strong mode when marker absent", () => { + const decision = shouldReinjectForEvent( + { messages: [{ content: "hi" }], type: "compaction" }, + "strong", + ); + expect(decision.needed).toBe(true); + expect(decision.reason).toBe("relevant_turn_and_marker_absent"); + }); + + it("rejects reinjection for compaction when marker is present", () => { + const decision = shouldReinjectForEvent( + { + messages: [{ content: ROOT_PAIR_START_MARKER }], + type: "compaction", + }, + "strong", + ); + expect(decision.needed).toBe(false); + expect(decision.reason).toBe("marker_present_in_messages"); + }); + + it("forces reinjection for artifact_change even when marker is present", () => { + const decision = shouldReinjectForEvent( + { + messages: [{ content: ROOT_PAIR_START_MARKER }], + type: "artifact_change", + }, + "strong", + ); + expect(decision.needed).toBe(true); + expect(decision.reason).toBe("root_pair_artifact_changed"); + }); + + it("rejects reinjection for all relevant turn types in off mode", () => { + const types = [ + "agent_start", + "edit_intent", + "architecture_sensitive", + "compaction", + "artifact_change", + ]; + for (const type of types) { + const decision = shouldReinjectForEvent( + { messages: [{ content: "hi" }], type }, + "off", + ); + expect(decision.needed).toBe(false); + expect(decision.reason).toBe("mode_does_not_require_reinjection"); + } + }); + + it("rejects reinjection for all relevant turn types in advisory mode", () => { + const types = [ + "agent_start", + "edit_intent", + "architecture_sensitive", + "compaction", + "artifact_change", + ]; + for (const type of types) { + const decision = shouldReinjectForEvent( + { messages: [{ content: "hi" }], type }, + "advisory", + ); + expect(decision.needed).toBe(false); + expect(decision.reason).toBe("mode_does_not_require_reinjection"); + } + }); + + it("allows reinjection for strict mode on relevant turns", () => { + const decision = shouldReinjectForEvent( + { messages: [{ content: "hi" }], type: "edit_intent" }, + "strict", + ); + expect(decision.needed).toBe(true); + expect(decision.reason).toBe("relevant_turn_and_marker_absent"); + }); + + it("rejects reinjection in strict mode when marker is present", () => { + const decision = shouldReinjectForEvent( + { + messages: [{ content: ROOT_PAIR_START_MARKER }], + type: "architecture_sensitive", + }, + "strict", + ); + expect(decision.needed).toBe(false); + expect(decision.reason).toBe("marker_present_in_messages"); + }); + + it("forces reinjection for artifact_change in strict mode even with marker", () => { + const decision = shouldReinjectForEvent( + { + messages: [{ content: ROOT_PAIR_START_MARKER }], + type: "artifact_change", + }, + "strict", + ); + expect(decision.needed).toBe(true); + expect(decision.reason).toBe("root_pair_artifact_changed"); + }); + + it("rejects reinjection for unknown event types even in strong mode", () => { + const decision = shouldReinjectForEvent( + { messages: [{ content: "hi" }], type: "custom_event" }, + "strong", + ); + expect(decision.needed).toBe(false); + expect(decision.reason).toBe("not_a_relevant_turn"); + }); + + it("prefers message marker over payload marker when both present", () => { + const decision = shouldReinjectForEvent( + { + messages: [{ content: ROOT_PAIR_START_MARKER }], + payload: "no marker here", + type: "edit_intent", + }, + "strong", + ); + expect(decision.needed).toBe(false); + expect(decision.reason).toBe("marker_present_in_messages"); + }); + + it("falls back to payload when messages are empty", () => { + const decision = shouldReinjectForEvent( + { + messages: [], + payload: ROOT_PAIR_START_MARKER, + type: "edit_intent", + }, + "strong", + ); + expect(decision.needed).toBe(false); + expect(decision.reason).toBe("marker_present_in_payload"); + }); + }); + + describe("edit intent detection", () => { + it("detects edit-the-file patterns", () => { + expect(detectEditIntent([{ content: "edit the file" }])).toBe(true); + expect(detectEditIntent([{ content: "modify this code" }])).toBe(true); + expect(detectEditIntent([{ content: "update the function" }])).toBe(true); + }); + + it("detects change-the-class patterns", () => { + expect(detectEditIntent([{ content: "change the class" }])).toBe(true); + expect(detectEditIntent([{ content: "refactor the module" }])).toBe(true); + expect(detectEditIntent([{ content: "fix the bug" }])).toBe(true); + }); + + it("detects write-new patterns", () => { + expect(detectEditIntent([{ content: "write a file" }])).toBe(true); + expect(detectEditIntent([{ content: "create a component" }])).toBe(true); + expect(detectEditIntent([{ content: "generate a function" }])).toBe(true); + expect(detectEditIntent([{ content: "write new file" }])).toBe(true); + }); + + it("detects code-block edit intent", () => { + expect( + detectEditIntent([ + { + content: "```ts\nedit the function\n```", + }, + ]), + ).toBe(true); + }); + + it("detects apply-change patterns", () => { + expect( + detectEditIntent([{ content: "apply the change to fix this" }]), + ).toBe(true); + expect(detectEditIntent([{ content: "apply this patch" }])).toBe(true); + }); + + it("returns false for non-edit content", () => { + expect(detectEditIntent([{ content: "what is the weather?" }])).toBe( + false, + ); + expect(detectEditIntent([{ content: "explain how this works" }])).toBe( + false, + ); + expect(detectEditIntent([{ content: "hello world" }])).toBe(false); + }); + + it("handles array content blocks", () => { + expect( + detectEditIntent([ + { + content: [ + { type: "text", text: "please " }, + { type: "text", text: "edit the file" }, + ], + }, + ]), + ).toBe(true); + }); + + it("returns false for empty messages", () => { + expect(detectEditIntent([])).toBe(false); + }); + + it("detects delete/remove patterns", () => { + expect(detectEditIntent([{ content: "delete the old method" }])).toBe( + true, + ); + expect(detectEditIntent([{ content: "remove the unused code" }])).toBe( + true, + ); + }); + + it("detects rewrite/patch patterns", () => { + expect(detectEditIntent([{ content: "rewrite the component" }])).toBe( + true, + ); + expect(detectEditIntent([{ content: "patch the file" }])).toBe(true); + }); + }); + + describe("architecture-sensitive reasoning detection", () => { + it("detects architecture decision patterns", () => { + expect( + detectArchitectureSensitiveReasoning([ + { content: "this is an architecture decision" }, + ]), + ).toBe(true); + expect( + detectArchitectureSensitiveReasoning([ + { content: "the architectural pattern we chose" }, + ]), + ).toBe(true); + }); + + it("detects design pattern patterns", () => { + expect( + detectArchitectureSensitiveReasoning([ + { content: "design pattern review" }, + ]), + ).toBe(true); + expect( + detectArchitectureSensitiveReasoning([ + { content: "design choice analysis" }, + ]), + ).toBe(true); + }); + + it("detects restructure patterns", () => { + expect( + detectArchitectureSensitiveReasoning([ + { content: "restructure the codebase" }, + ]), + ).toBe(true); + expect( + detectArchitectureSensitiveReasoning([ + { content: "reorganize the modules" }, + ]), + ).toBe(true); + expect( + detectArchitectureSensitiveReasoning([ + { content: "redesign the API layer" }, + ]), + ).toBe(true); + expect( + detectArchitectureSensitiveReasoning([ + { content: "rearchitect the system" }, + ]), + ).toBe(true); + }); + + it("detects system design patterns", () => { + expect( + detectArchitectureSensitiveReasoning([ + { content: "system design overview" }, + ]), + ).toBe(true); + expect( + detectArchitectureSensitiveReasoning([ + { content: "high-level architecture" }, + ]), + ).toBe(true); + }); + + it("detects dependency patterns", () => { + expect( + detectArchitectureSensitiveReasoning([ + { content: "dependency injection strategy" }, + ]), + ).toBe(true); + expect( + detectArchitectureSensitiveReasoning([ + { content: "dependency graph analysis" }, + ]), + ).toBe(true); + }); + + it("detects API/interface design patterns", () => { + expect( + detectArchitectureSensitiveReasoning([ + { content: "API design migration" }, + ]), + ).toBe(true); + expect( + detectArchitectureSensitiveReasoning([ + { content: "interface contract change" }, + ]), + ).toBe(true); + }); + + it("detects architecture style patterns", () => { + expect( + detectArchitectureSensitiveReasoning([ + { content: "microservice architecture" }, + ]), + ).toBe(true); + expect( + detectArchitectureSensitiveReasoning([ + { content: "clean architecture principles" }, + ]), + ).toBe(true); + }); + + it("detects data and quality attribute patterns", () => { + expect( + detectArchitectureSensitiveReasoning([ + { content: "data model design" }, + ]), + ).toBe(true); + expect( + detectArchitectureSensitiveReasoning([ + { content: "scalability tradeoff decision" }, + ]), + ).toBe(true); + expect( + detectArchitectureSensitiveReasoning([ + { content: "security concern analysis" }, + ]), + ).toBe(true); + expect( + detectArchitectureSensitiveReasoning([ + { content: "performance concern review" }, + ]), + ).toBe(true); + }); + + it("returns false for non-architecture content", () => { + expect( + detectArchitectureSensitiveReasoning([ + { content: "what is the weather?" }, + ]), + ).toBe(false); + expect( + detectArchitectureSensitiveReasoning([ + { content: "fix the typo in readme" }, + ]), + ).toBe(false); + expect( + detectArchitectureSensitiveReasoning([ + { content: "add a console log" }, + ]), + ).toBe(false); + }); + + it("handles array content blocks", () => { + expect( + detectArchitectureSensitiveReasoning([ + { + content: [ + { type: "text", text: "let's discuss " }, + { type: "text", text: "the architecture" }, + ], + }, + ]), + ).toBe(true); + }); + + it("returns false for empty messages", () => { + expect(detectArchitectureSensitiveReasoning([])).toBe(false); + }); + }); + describe("injection payload", () => { it("always includes the root pair", () => { const dir = mkdtempSync(join(tmpdir(), "pi-payload-test-"));