From 58e8bd31d3a25e34c01eeb23280f2c678dd53b47 Mon Sep 17 00:00:00 2001 From: Alex Blank Date: Thu, 11 Jun 2026 22:54:51 +0200 Subject: [PATCH] feat(prompt): implement prompt injection slice 4 --- pi-extension.ts | 59 ++++- src/index.ts | 9 + src/prompt-injection.ts | 223 ++++++++++++++++- tests/pi-extension.test.ts | 269 ++++++++++++++++++++- tests/prompt-injection.test.ts | 423 +++++++++++++++++++++++++++++++++ 5 files changed, 965 insertions(+), 18 deletions(-) diff --git a/pi-extension.ts b/pi-extension.ts index 97e59c7..40f4dc3 100644 --- a/pi-extension.ts +++ b/pi-extension.ts @@ -9,8 +9,11 @@ import { reinitPath, retrieveContext, buildPreInitHint, + buildAdvisoryReminder, + buildStrictBypassGuard, modeAllowsPreInitHint, modeAllowsInjection, + evaluateStrictBypass, discoverContextWindow, buildInjectionPayload, shouldReinjectForEvent, @@ -372,7 +375,19 @@ export default function (pi: ExtensionAPI) { }; } - // Maps exist but advisory mode does not permit automatic artifact injection yet. + // Slice 4: advisory mode shows a visible lightweight reminder after init. + // No root-pair preload, no per-turn reinjection. + if (config.promptInjectionMode === "advisory") { + return { + message: { + customType: "pi-project-map-hint", + content: buildAdvisoryReminder(), + display: true, + }, + }; + } + + // Maps exist but mode does not permit automatic artifact injection. if (!modeAllowsInjection(config.promptInjectionMode)) { return {}; } @@ -424,7 +439,8 @@ export default function (pi: ExtensionAPI) { return {}; } - // Slice 3b: detect root-pair artifact changes + // Slice 3b: detect root-pair artifact changes early. Invalidation always + // forces reinjection, even in strict mode, because the context is stale. const currentMtimes = getRootPairMtimes(ctx.cwd); const hasPrevious = lastRootPairMtimes.mapMtime !== undefined || @@ -433,12 +449,47 @@ export default function (pi: ExtensionAPI) { hasPrevious && rootPairChanged(currentMtimes, lastRootPairMtimes); lastRootPairMtimes = currentMtimes; - const eventType = artifactChanged ? "artifact_change" : event?.type; + if (artifactChanged) { + const decision = shouldReinjectForEvent( + { + messages: event?.messages, + type: "artifact_change", + 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, + }, + }; + } + + // Slice 4: strict-mode bypass guard for sensitive actions with missing protocol path. + if (config.promptInjectionMode === "strict") { + const bypass = evaluateStrictBypass(event, config.promptInjectionMode); + if (bypass.guard) { + return { + message: { + customType: "pi-project-map-hint", + content: buildStrictBypassGuard(bypass.reason), + display: true, + }, + }; + } + } const decision = shouldReinjectForEvent( { messages: event?.messages, - type: eventType, + type: event?.type, payload: event?.payload, }, config.promptInjectionMode, diff --git a/src/index.ts b/src/index.ts index 610a74e..3eeada9 100644 --- a/src/index.ts +++ b/src/index.ts @@ -31,6 +31,13 @@ export { modeAllowsPreInitHint, modeAllowsInjection, modeRequiresProtocolPath, + hasProtocolPath, + isSensitiveAction, + messagesHaveBypass, + extractBypassReason, + evaluateStrictBypass, + buildAdvisoryReminder, + buildStrictBypassGuard, discoverContextWindow, estimateTokens, findAllArtifactPairs, @@ -46,7 +53,9 @@ export { type RootPairMtimes, type RelevantTurnType, type ReinjectDecision, + type StrictBypassDecision, ROOT_PAIR_START_MARKER, ROOT_PAIR_END_MARKER, TRUST_BOUNDARY_TEXT, + BYPASS_MARKER_PREFIX, } from "./prompt-injection.js"; diff --git a/src/prompt-injection.ts b/src/prompt-injection.ts index 657bc9a..4c36d03 100644 --- a/src/prompt-injection.ts +++ b/src/prompt-injection.ts @@ -116,13 +116,7 @@ export function isRelevantTurnForReinjection(eventType: string): boolean { 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(" "); +function detectEditIntentFromText(combined: string): boolean { 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, @@ -133,12 +127,19 @@ export function detectEditIntent( } /** - * Heuristic detection of architecture-sensitive reasoning from message content. + * Heuristic detection of edit intent from message content. */ -export function detectArchitectureSensitiveReasoning( +export function detectEditIntent( messages: Array<{ content?: unknown; text?: string }>, ): boolean { + if (!messages || messages.length === 0) return false; const combined = messages.map(extractTextFromMessage).join(" "); + return detectEditIntentFromText(combined); +} + +function detectArchitectureSensitiveReasoningFromText( + combined: string, +): boolean { const patterns = [ /\barchitectur(e|al)\b/i, /\bdesign\s+(decision|pattern|choice|principle|review)\b/i, @@ -153,6 +154,17 @@ export function detectArchitectureSensitiveReasoning( 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 { + if (!messages || messages.length === 0) return false; + const combined = messages.map(extractTextFromMessage).join(" "); + return detectArchitectureSensitiveReasoningFromText(combined); +} + // --------------------------------------------------------------------------- // Root-pair artifact change detection // --------------------------------------------------------------------------- @@ -288,6 +300,199 @@ export function modeRequiresProtocolPath(mode: PromptInjectionMode): boolean { return mode === "strict"; } +// --------------------------------------------------------------------------- +// Protocol-path detection (Slice 4) +// --------------------------------------------------------------------------- + +/** + * Inline bypass marker prefix. Agents may include `[PI_MAP_BYPASS: reason]` in + * a message to proceed past a strict-mode guard. + */ +export const BYPASS_MARKER_PREFIX = "[PI_MAP_BYPASS:"; + +function extractTextFromPayload(payload: unknown): string { + if (!payload) return ""; + if (typeof payload === "string") return payload; + try { + return JSON.stringify(payload); + } catch { + return ""; + } +} + +/** + * Check whether the outgoing context contains the full protocol path: + * canonical root-pair marker + trust-boundary instruction. + */ +export function hasProtocolPath( + messages?: Array<{ content?: unknown; text?: string }>, + payload?: unknown, +): boolean { + const sources: string[] = []; + if (messages) { + for (const m of messages) { + sources.push(extractTextFromMessage(m)); + } + } + const payloadText = extractTextFromPayload(payload); + if (payloadText) { + sources.push(payloadText); + } + const combined = sources.join("\n"); + return ( + combined.includes(ROOT_PAIR_START_MARKER) && + combined.includes(TRUST_BOUNDARY_TEXT) + ); +} + +/** + * Determine whether the current turn is a sensitive action. + * + * Uses explicit event types when available, with heuristic fallback from + * message content and provider payload for runtimes that do not emit + * `edit_intent` or `architecture_sensitive` event types. + */ +export function isSensitiveAction( + eventType?: string, + messages?: Array<{ content?: unknown; text?: string }>, + payload?: unknown, +): boolean { + const explicit: RelevantTurnType[] = [ + "edit_intent", + "architecture_sensitive", + ]; + if (eventType && explicit.includes(eventType as RelevantTurnType)) { + return true; + } + const messageText = + messages && messages.length > 0 + ? messages.map(extractTextFromMessage).join("\n") + : ""; + const payloadText = extractTextFromPayload(payload); + const combined = + messageText || payloadText ? `${messageText}\n${payloadText}`.trim() : ""; + if (!combined) return false; + return ( + detectEditIntentFromText(combined) || + detectArchitectureSensitiveReasoningFromText(combined) + ); +} + +/** + * Check whether any message contains a valid explicit bypass marker. + * Empty or whitespace-only reasons are rejected. + */ +export function messagesHaveBypass( + messages?: Array<{ content?: unknown; text?: string }>, + payload?: unknown, +): boolean { + return extractBypassReason(messages, payload) !== undefined; +} + +/** + * Extract the reason from the first `[PI_MAP_BYPASS: reason]` marker, if any. + * Returns `undefined` when the marker is absent or its reason is empty/whitespace. + */ +export function extractBypassReason( + messages?: Array<{ content?: unknown; text?: string }>, + payload?: unknown, +): string | undefined { + const parts: string[] = []; + if (messages) { + parts.push(messages.map(extractTextFromMessage).join("\n")); + } + const payloadText = extractTextFromPayload(payload); + if (payloadText) { + parts.push(payloadText); + } + const combined = parts.join("\n"); + if (!combined) return undefined; + const escapedPrefix = BYPASS_MARKER_PREFIX.replace( + /[.*+?^${}()|[\]\\]/g, + "\\$&", + ); + const match = new RegExp(`${escapedPrefix}\\s*([^\\]]+)\\]`).exec(combined); + if (!match) return undefined; + const reason = match[1].trim(); + return reason.length > 0 ? reason : undefined; +} + +export interface StrictBypassDecision { + guard: boolean; + reason?: string; + bypassMarker?: string; +} + +/** + * Evaluate whether strict mode should block a sensitive turn with a visible + * bypass guard because the protocol path is missing. + * + * Returns `guard: false` for non-strict modes, non-sensitive turns, turns + * where the protocol path is present, or turns that include an explicit bypass + * marker. + */ +export function evaluateStrictBypass( + event: + | { + messages?: Array<{ content?: unknown; text?: string }>; + type?: string; + payload?: unknown; + } + | undefined + | null, + mode: PromptInjectionMode, +): StrictBypassDecision { + if (mode !== "strict") { + return { guard: false }; + } + if (!event) { + return { guard: false }; + } + if (!isSensitiveAction(event.type, event.messages, event.payload)) { + return { guard: false }; + } + const bypassReason = extractBypassReason(event.messages, event.payload); + if (bypassReason !== undefined) { + return { guard: false, bypassMarker: bypassReason }; + } + if (hasProtocolPath(event.messages, event.payload)) { + return { guard: false }; + } + return { + guard: true, + reason: "protocol path missing for sensitive action", + }; +} + +/** + * Build a lightweight user-visible reminder for advisory mode after init. + * No root-pair content is injected automatically. + */ +export function buildAdvisoryReminder(): string { + return [ + "📋 Project map advisory mode active.", + "", + "Root `.pi-map.index.md` and `.pi-map.md` are available but are not automatically injected. Read them manually when you need routing or orientation context, and remember that source remains the final authority before edits.", + ].join("\n"); +} + +/** + * Build a visible strict-mode bypass guard for sensitive turns where the + * protocol path is missing. + */ +export function buildStrictBypassGuard(reason?: string): string { + return [ + "🛑 Strict project-map guard", + "", + reason || + "A sensitive action was detected without the project-map protocol path.", + "", + "The protocol path requires the root `.pi-map.index.md` / `.pi-map.md` pair plus the trust boundary (`index routes, map orients, source decides`) to be present in context.", + "", + "To proceed, either restore the project-map context or include an explicit bypass marker: `[PI_MAP_BYPASS: ]`.", + ].join("\n"); +} + // --------------------------------------------------------------------------- // Context-window discovery (Slice 2) // --------------------------------------------------------------------------- diff --git a/tests/pi-extension.test.ts b/tests/pi-extension.test.ts index 5864c66..8d8e7f5 100644 --- a/tests/pi-extension.test.ts +++ b/tests/pi-extension.test.ts @@ -2,6 +2,10 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { mkdtempSync, writeFileSync, existsSync, unlinkSync } from "fs"; import { join } from "path"; import { tmpdir, homedir } from "os"; +import { + ROOT_PAIR_START_MARKER, + TRUST_BOUNDARY_TEXT, +} from "../src/prompt-injection.js"; vi.mock("@mariozechner/pi-coding-agent", () => ({ ExtensionAPI: class {}, @@ -365,7 +369,7 @@ describe("pi-extension", () => { expect(result.message.content).toContain("project_map_init"); }); - it("returns empty object when maps exist and mode is advisory", async () => { + it("shows visible advisory reminder 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( @@ -377,7 +381,10 @@ describe("pi-extension", () => { const handler = registeredEvents.before_agent_start; const result = await handler(null, mockCtx); - expect(result).toEqual({}); + expect(result).toHaveProperty("message"); + expect(result.message.display).toBe(true); + expect(result.message.content).toContain("advisory mode active"); + expect(result.message.content).not.toContain(ROOT_PAIR_START_MARKER); }); it("shows visible pre-init hint when no maps exist and mode is strict", async () => { @@ -396,6 +403,40 @@ describe("pi-extension", () => { expect(result.message.content).toContain("project_map_init"); }); + it("defaults to strong mode and injects hidden root pair when no config file exists", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nDefaultMap\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nDefaultIndex\n", + ); + 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("DefaultMap"); + expect(result.message.content).toContain("DefaultIndex"); + }); + + it("returns empty object when maps exist and mode is off", 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: "off" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.before_agent_start; + const result = await handler(null, mockCtx); + + expect(result).toEqual({}); + }); + 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"); @@ -882,7 +923,7 @@ describe("pi-extension", () => { expect(result.message.content).toContain("TestIndex"); }); - it("sequence: strict mode reinjects on relevant turns same as strong mode", async () => { + it("sequence: strict mode shows bypass guard on sensitive turns when protocol path is missing", async () => { const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTestMap\n"); writeFileSync( @@ -902,13 +943,15 @@ describe("pi-extension", () => { mockCtx, ); expect(first).toHaveProperty("message"); - expect(first.message.content).toContain("PI_MAP_ROOT_PAIR_START"); + expect(first.message.display).toBe(true); + expect(first.message.content).toContain("Strict project-map guard"); + expect(first.message.content).toContain("protocol path missing"); const second = await handler( { messages: [ { - content: ` already injected`, + content: `${ROOT_PAIR_START_MARKER} ${TRUST_BOUNDARY_TEXT} already injected`, }, ], type: "architecture_sensitive", @@ -925,6 +968,48 @@ describe("pi-extension", () => { expect(third.message.content).toContain("PI_MAP_ROOT_PAIR_START"); }); + it("sequence: strict mode behaves like strong when protocol path 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: "strict" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.context; + + const first = await handler( + { + messages: [ + { + content: `${ROOT_PAIR_START_MARKER} ${TRUST_BOUNDARY_TEXT}`, + }, + ], + type: "edit_intent", + }, + mockCtx, + ); + expect(first).toEqual({}); + + const second = await handler( + { + messages: [ + { + content: `${ROOT_PAIR_START_MARKER} ${TRUST_BOUNDARY_TEXT} still present`, + }, + ], + type: "compaction", + }, + mockCtx, + ); + expect(second).toEqual({}); + }); + 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"); @@ -966,5 +1051,179 @@ describe("pi-extension", () => { } } }); + + it("sequence: strict mode bypass guard is cleared by inline marker on the next sensitive turn", 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 guard = await handler( + { messages: [{ content: "edit this file" }], type: "edit_intent" }, + mockCtx, + ); + expect(guard).toHaveProperty("message"); + expect(guard.message.display).toBe(true); + expect(guard.message.content).toContain("Strict project-map guard"); + + const cleared = await handler( + { + messages: [ + { + content: "[PI_MAP_BYPASS: emergency patch] edit this file", + }, + ], + type: "edit_intent", + }, + mockCtx, + ); + expect(cleared).toHaveProperty("message"); + expect(cleared.message.display).toBe(false); + expect(cleared.message.content).toContain("PI_MAP_ROOT_PAIR_START"); + + const later = await handler( + { messages: [{ content: "edit another file" }], type: "edit_intent" }, + mockCtx, + ); + expect(later).toHaveProperty("message"); + expect(later.message.display).toBe(true); + expect(later.message.content).toContain("Strict project-map guard"); + }); + + it("sequence: strict mode forces reinjection on artifact_change even when protocol path is missing", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nOriginalMap\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; + + // Establish baseline mtimes on an irrelevant turn. + await handler( + { messages: [{ content: "hi" }], type: "user_chat" }, + mockCtx, + ); + + // A sensitive turn with no protocol path should be guarded. + const guard = await handler( + { messages: [{ content: "edit this file" }], type: "edit_intent" }, + mockCtx, + ); + expect(guard).toHaveProperty("message"); + expect(guard.message.content).toContain("Strict project-map guard"); + + // Modify root artifacts outside the agent. + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nUpdatedMap\n"); + + // Artifact invalidation must force reinjection, not re-guard. + const result = await handler( + { messages: [{ content: "hi" }], type: "user_chat" }, + mockCtx, + ); + expect(result).toHaveProperty("message"); + expect(result.message.display).toBe(false); + expect(result.message.content).toContain("UpdatedMap"); + expect(result.message.content).toContain("PI_MAP_ROOT_PAIR_START"); + }); + + it("defaults to strong mode and reinjects on edit_intent when no config file exists", async () => { + const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-")); + writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nDefaultMap\n"); + writeFileSync( + join(dir, ".pi-map.index.md"), + "# . (index)\n## role\nDefaultIndex\n", + ); + 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.display).toBe(false); + expect(result.message.content).toContain("PI_MAP_ROOT_PAIR_START"); + expect(result.message.content).toContain("DefaultMap"); + expect(result.message.content).toContain("DefaultIndex"); + }); + + it("off mode skips context reinjection even on edit_intent", 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: "off" }), + ); + mockCtx.cwd = dir; + + const handler = registeredEvents.context; + const result = await handler( + { messages: [{ content: "edit this file" }], type: "edit_intent" }, + mockCtx, + ); + + expect(result).toEqual({}); + }); + + it("strict mode guards heuristic sensitive action in a generic user_chat turn", 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 result = await handler( + { messages: [{ content: "refactor the module" }], type: "user_chat" }, + mockCtx, + ); + + expect(result).toHaveProperty("message"); + expect(result.message.display).toBe(true); + expect(result.message.content).toContain("Strict project-map guard"); + expect(result.message.content).toContain("protocol path missing"); + expect(result.message.content).not.toContain(ROOT_PAIR_START_MARKER); + }); + + it("strict mode does not guard non-sensitive user_chat turns", 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.context; + const result = await handler( + { messages: [{ content: "what is the weather?" }], type: "user_chat" }, + mockCtx, + ); + + expect(result).toEqual({}); + }); }); }); diff --git a/tests/prompt-injection.test.ts b/tests/prompt-injection.test.ts index 5960445..013cba2 100644 --- a/tests/prompt-injection.test.ts +++ b/tests/prompt-injection.test.ts @@ -10,6 +10,8 @@ import { modeAllowsPreInitHint, modeAllowsInjection, modeRequiresProtocolPath, + hasProtocolPath, + buildAdvisoryReminder, discoverContextWindow, estimateTokens, findAllArtifactPairs, @@ -20,9 +22,15 @@ import { isRelevantTurnForReinjection, detectEditIntent, detectArchitectureSensitiveReasoning, + isSensitiveAction, + messagesHaveBypass, + extractBypassReason, + evaluateStrictBypass, + buildStrictBypassGuard, ROOT_PAIR_START_MARKER, ROOT_PAIR_END_MARKER, TRUST_BOUNDARY_TEXT, + BYPASS_MARKER_PREFIX, } from "../src/prompt-injection.js"; describe("prompt-injection helpers", () => { @@ -708,6 +716,73 @@ describe("prompt-injection helpers", () => { }); }); + describe("protocol path detection", () => { + it("detects protocol path when messages contain marker and trust boundary", () => { + const messages = [ + { + content: `Start ${ROOT_PAIR_START_MARKER} ${TRUST_BOUNDARY_TEXT} end`, + }, + ]; + expect(hasProtocolPath(messages)).toBe(true); + }); + + it("detects protocol path when payload contains marker and trust boundary", () => { + const payload = { + body: `${ROOT_PAIR_START_MARKER}\n${TRUST_BOUNDARY_TEXT}`, + }; + expect(hasProtocolPath(undefined, payload)).toBe(true); + }); + + it("returns false when only marker is present", () => { + const messages = [{ content: `hello ${ROOT_PAIR_START_MARKER} world` }]; + expect(hasProtocolPath(messages)).toBe(false); + }); + + it("returns false when only trust boundary is present", () => { + const messages = [{ content: `hello ${TRUST_BOUNDARY_TEXT} world` }]; + expect(hasProtocolPath(messages)).toBe(false); + }); + + it("returns false when neither marker nor trust boundary is present", () => { + const messages = [{ content: "plain text" }]; + expect(hasProtocolPath(messages)).toBe(false); + }); + + it("returns false for empty messages and undefined payload", () => { + expect(hasProtocolPath([])).toBe(false); + expect(hasProtocolPath(undefined, undefined)).toBe(false); + }); + + it("handles array content blocks", () => { + const messages = [ + { + content: [ + { type: "text", text: ROOT_PAIR_START_MARKER }, + { type: "text", text: TRUST_BOUNDARY_TEXT }, + ], + }, + ]; + expect(hasProtocolPath(messages)).toBe(true); + }); + + it("requires both elements in the same combined source", () => { + const messages = [{ content: ROOT_PAIR_START_MARKER }]; + const payload = { body: TRUST_BOUNDARY_TEXT }; + expect(hasProtocolPath(messages, payload)).toBe(true); + }); + }); + + describe("advisory reminder", () => { + it("returns a visible lightweight reminder with no synthetic artifacts", () => { + const reminder = buildAdvisoryReminder(); + expect(reminder).toContain("advisory mode active"); + expect(reminder).toContain("📋"); + expect(reminder).not.toContain(ROOT_PAIR_START_MARKER); + expect(reminder).not.toContain("### Root index"); + expect(reminder).not.toContain("### Root map"); + }); + }); + describe("injection payload", () => { it("always includes the root pair", () => { const dir = mkdtempSync(join(tmpdir(), "pi-payload-test-")); @@ -769,4 +844,352 @@ describe("prompt-injection helpers", () => { expect(payload.content).not.toContain("# src map x"); }); }); + + describe("sensitive action detection", () => { + it("detects explicit edit_intent event type", () => { + expect(isSensitiveAction("edit_intent", [])).toBe(true); + }); + + it("detects explicit architecture_sensitive event type", () => { + expect(isSensitiveAction("architecture_sensitive", [])).toBe(true); + }); + + it("falls back to edit-intent heuristic for generic events", () => { + expect( + isSensitiveAction("user_chat", [{ content: "edit the file" }]), + ).toBe(true); + }); + + it("falls back to architecture heuristic for generic events", () => { + expect( + isSensitiveAction("user_chat", [ + { content: "this is an architecture decision" }, + ]), + ).toBe(true); + }); + + it("treats compaction as non-sensitive", () => { + expect(isSensitiveAction("compaction", [{ content: "compacting" }])).toBe( + false, + ); + }); + + it("treats agent_start as non-sensitive", () => { + expect(isSensitiveAction("agent_start", [{ content: "hi" }])).toBe(false); + }); + + it("returns false for irrelevant messages without event type", () => { + expect( + isSensitiveAction(undefined, [{ content: "what is the weather?" }]), + ).toBe(false); + }); + + it("returns false when no event type or messages are provided", () => { + expect(isSensitiveAction(undefined, undefined)).toBe(false); + }); + + it("handles array content blocks", () => { + expect( + isSensitiveAction(undefined, [ + { + content: [ + { type: "text", text: "please " }, + { type: "text", text: "refactor the module" }, + ], + }, + ]), + ).toBe(true); + }); + + it("detects sensitive action from payload when messages are empty", () => { + expect(isSensitiveAction("user_chat", [], "edit the file")).toBe(true); + }); + + it("detects architecture-sensitive action from payload", () => { + expect(isSensitiveAction("user_chat", [], "architecture decision")).toBe( + true, + ); + }); + + it("returns false for irrelevant payload without messages or event type", () => { + expect( + isSensitiveAction(undefined, undefined, "what is the weather?"), + ).toBe(false); + }); + + it("combines messages and payload for heuristic detection", () => { + expect( + isSensitiveAction( + "user_chat", + [{ content: "please consider" }], + "the architecture decision", + ), + ).toBe(true); + }); + }); + + describe("bypass marker parsing", () => { + it("detects a valid bypass marker in messages", () => { + expect( + messagesHaveBypass([{ content: "[PI_MAP_BYPASS: temporary fix]" }]), + ).toBe(true); + }); + + it("detects a valid bypass marker prefix", () => { + expect(BYPASS_MARKER_PREFIX).toBe("[PI_MAP_BYPASS:"); + }); + + it("returns false when no marker is present", () => { + expect(messagesHaveBypass([{ content: "plain text" }])).toBe(false); + }); + + it("extracts the bypass reason", () => { + expect( + extractBypassReason([{ content: "[PI_MAP_BYPASS: temporary fix]" }]), + ).toBe("temporary fix"); + }); + + it("returns undefined when marker is absent", () => { + expect(extractBypassReason([{ content: "plain text" }])).toBeUndefined(); + }); + + it("rejects an empty bypass reason", () => { + expect( + extractBypassReason([{ content: "[PI_MAP_BYPASS:]" }]), + ).toBeUndefined(); + }); + + it("rejects a whitespace-only bypass reason", () => { + expect( + extractBypassReason([{ content: "[PI_MAP_BYPASS: ]" }]), + ).toBeUndefined(); + }); + + it("extracts the first reason when multiple markers are present", () => { + expect( + extractBypassReason([ + { + content: "[PI_MAP_BYPASS: first] [PI_MAP_BYPASS: second]", + }, + ]), + ).toBe("first"); + }); + + it("handles array content blocks", () => { + expect( + extractBypassReason([ + { + content: [ + { type: "text", text: "[PI_MAP_BYPASS: " }, + { type: "text", text: "array reason]" }, + ], + }, + ]), + ).toBe("array reason"); + }); + + it("detects bypass marker in payload string", () => { + expect( + messagesHaveBypass(undefined, "[PI_MAP_BYPASS: payload reason]"), + ).toBe(true); + }); + + it("detects bypass marker in payload object", () => { + expect( + messagesHaveBypass(undefined, { + body: "[PI_MAP_BYPASS: object reason]", + }), + ).toBe(true); + }); + + it("extracts reason from payload when messages have none", () => { + expect( + extractBypassReason(undefined, "[PI_MAP_BYPASS: payload reason]"), + ).toBe("payload reason"); + }); + + it("prefers first marker across messages and payload", () => { + expect( + extractBypassReason( + [{ content: "[PI_MAP_BYPASS: msg reason]" }], + "[PI_MAP_BYPASS: payload reason]", + ), + ).toBe("msg reason"); + }); + + it("extracts reason with colon inside", () => { + expect( + extractBypassReason([{ content: "[PI_MAP_BYPASS: fix: edge case]" }]), + ).toBe("fix: edge case"); + }); + + it("rejects bypass marker with only newlines/spaces", () => { + expect( + extractBypassReason([{ content: "[PI_MAP_BYPASS:\n\t ]" }]), + ).toBeUndefined(); + }); + }); + + describe("strict bypass evaluation", () => { + it("never guards in non-strict modes", () => { + for (const mode of ["off", "advisory", "strong"] as const) { + const decision = evaluateStrictBypass( + { messages: [{ content: "edit the file" }], type: "edit_intent" }, + mode, + ); + expect(decision.guard).toBe(false); + } + }); + + it("does not guard non-sensitive turns in strict mode", () => { + const decision = evaluateStrictBypass( + { messages: [{ content: "hi" }], type: "user_chat" }, + "strict", + ); + expect(decision.guard).toBe(false); + }); + + it("guards sensitive turns when protocol path is missing", () => { + const decision = evaluateStrictBypass( + { messages: [{ content: "edit the file" }], type: "edit_intent" }, + "strict", + ); + expect(decision.guard).toBe(true); + expect(decision.reason).toContain("protocol path missing"); + }); + + it("does not guard sensitive turns when protocol path is present", () => { + const decision = evaluateStrictBypass( + { + messages: [ + { + content: `${ROOT_PAIR_START_MARKER} ${TRUST_BOUNDARY_TEXT}`, + }, + ], + type: "edit_intent", + }, + "strict", + ); + expect(decision.guard).toBe(false); + }); + + it("uses heuristic fallback for generic sensitive events", () => { + const decision = evaluateStrictBypass( + { + messages: [{ content: "refactor the module" }], + type: "user_chat", + }, + "strict", + ); + expect(decision.guard).toBe(true); + }); + + it("clears the guard when a valid inline bypass marker is present", () => { + const decision = evaluateStrictBypass( + { + messages: [ + { + content: "[PI_MAP_BYPASS: emergency patch] edit the file", + }, + ], + type: "edit_intent", + }, + "strict", + ); + expect(decision.guard).toBe(false); + expect(decision.bypassMarker).toBe("emergency patch"); + }); + + it("does not clear the guard for an empty bypass marker", () => { + const decision = evaluateStrictBypass( + { + messages: [{ content: "[PI_MAP_BYPASS:] edit the file" }], + type: "edit_intent", + }, + "strict", + ); + expect(decision.guard).toBe(true); + }); + + it("does not guard when event is null or undefined", () => { + expect(evaluateStrictBypass(undefined, "strict").guard).toBe(false); + expect(evaluateStrictBypass(null, "strict").guard).toBe(false); + }); + + it("does not guard when payload-only protocol path is present", () => { + const decision = evaluateStrictBypass( + { + messages: [{ content: "edit the file" }], + payload: `${ROOT_PAIR_START_MARKER}\n${TRUST_BOUNDARY_TEXT}`, + type: "edit_intent", + }, + "strict", + ); + expect(decision.guard).toBe(false); + }); + + it("clears guard when bypass marker is only in payload", () => { + const decision = evaluateStrictBypass( + { + messages: [{ content: "edit the file" }], + payload: "[PI_MAP_BYPASS: emergency patch]", + type: "edit_intent", + }, + "strict", + ); + expect(decision.guard).toBe(false); + expect(decision.bypassMarker).toBe("emergency patch"); + }); + + it("guards when payload-only bypass marker has empty reason", () => { + const decision = evaluateStrictBypass( + { + messages: [{ content: "edit the file" }], + payload: "[PI_MAP_BYPASS:]", + type: "edit_intent", + }, + "strict", + ); + expect(decision.guard).toBe(true); + }); + + it("uses heuristic fallback for payload-only sensitive events", () => { + const decision = evaluateStrictBypass( + { + messages: [], + payload: "refactor the module", + type: "user_chat", + }, + "strict", + ); + expect(decision.guard).toBe(true); + }); + }); + + describe("strict bypass guard message", () => { + it("includes the provided reason", () => { + const guard = buildStrictBypassGuard("custom reason"); + expect(guard).toContain("custom reason"); + expect(guard).toContain("🛑"); + }); + + it("includes a default reason when none is provided", () => { + const guard = buildStrictBypassGuard(); + expect(guard).toContain("sensitive action"); + expect(guard).toContain("protocol path"); + }); + + it("tells the user how to proceed", () => { + const guard = buildStrictBypassGuard(); + expect(guard).toContain("PI_MAP_BYPASS"); + expect(guard).toContain("restore the project-map context"); + }); + + it("does not contain synthetic root-pair content", () => { + const guard = buildStrictBypassGuard(); + expect(guard).not.toContain(ROOT_PAIR_START_MARKER); + expect(guard).not.toContain("### Root index"); + expect(guard).not.toContain("### Root map"); + }); + }); });