alex/pi-map
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(prompt): implement prompt injection slice 4

Browse Source
This commit is contained in:
Alex
2026-06-11 22:54:51 +02:00
parent 19666c900e
commit 58e8bd31d3
5 changed files with 965 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/pi-extension.test.ts
+264 -5
View File
@@ -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: `<!-- PI_MAP_ROOT_PAIR_START --> 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({});
});
});
});