feat(prompt): implement prompt injection slice 3
This commit is contained in:
@@ -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",
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -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: "<!-- PI_MAP_ROOT_PAIR_START --> 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(
|
||||
"<!-- PI_MAP_ROOT_PAIR_START -->",
|
||||
);
|
||||
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: "<!-- PI_MAP_ROOT_PAIR_START -->" }],
|
||||
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: "<!-- PI_MAP_ROOT_PAIR_START -->" },
|
||||
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(
|
||||
"<!-- PI_MAP_ROOT_PAIR_START -->",
|
||||
);
|
||||
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: "<!-- PI_MAP_ROOT_PAIR_START -->" }],
|
||||
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: "<!-- PI_MAP_ROOT_PAIR_START -->" }],
|
||||
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: `<!-- PI_MAP_ROOT_PAIR_START --> 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: `<!-- PI_MAP_ROOT_PAIR_START --> 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: `<!-- PI_MAP_ROOT_PAIR_START --> 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: `<!-- PI_MAP_ROOT_PAIR_START --> 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: `<!-- PI_MAP_ROOT_PAIR_START --> 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 = "<!-- PI_MAP_ROOT_PAIR_START --> ";
|
||||
} else {
|
||||
expect(result).toEqual({});
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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-"));
|
||||
|
||||
Reference in New Issue
Block a user