feat(prompt): implement prompt injection slice 1
This commit is contained in:
+30
-1
@@ -8,7 +8,11 @@ import {
|
|||||||
validateMaps,
|
validateMaps,
|
||||||
reinitPath,
|
reinitPath,
|
||||||
retrieveContext,
|
retrieveContext,
|
||||||
|
buildPreInitHint,
|
||||||
|
modeAllowsPreInitHint,
|
||||||
|
modeAllowsInjection,
|
||||||
} from "./src/index.js";
|
} from "./src/index.js";
|
||||||
|
import { loadConfig } from "./src/config.js";
|
||||||
import { createLLMClient } from "./src/llm/llm-client.js";
|
import { createLLMClient } from "./src/llm/llm-client.js";
|
||||||
import { LLMError } from "./src/llm/llm-error.js";
|
import { LLMError } from "./src/llm/llm-error.js";
|
||||||
|
|
||||||
@@ -340,8 +344,33 @@ export default function (pi: ExtensionAPI) {
|
|||||||
|
|
||||||
// Inject maintenance instructions before agent starts
|
// Inject maintenance instructions before agent starts
|
||||||
pi.on("before_agent_start", async (_event, _ctx) => {
|
pi.on("before_agent_start", async (_event, _ctx) => {
|
||||||
|
const config = loadConfig(_ctx.cwd);
|
||||||
const mapFiles = findPiMapFiles(_ctx.cwd);
|
const mapFiles = findPiMapFiles(_ctx.cwd);
|
||||||
if (mapFiles.length === 0) return {};
|
|
||||||
|
// Mode is off: no injection at all
|
||||||
|
if (config.promptInjectionMode === "off") {
|
||||||
|
return {};
|
||||||
|
}
|
||||||
|
|
||||||
|
// No maps exist yet: show visible pre-init hint, but only if mode allows it
|
||||||
|
if (mapFiles.length === 0) {
|
||||||
|
if (!modeAllowsPreInitHint(config.promptInjectionMode)) {
|
||||||
|
return {};
|
||||||
|
}
|
||||||
|
return {
|
||||||
|
message: {
|
||||||
|
customType: "pi-project-map-hint",
|
||||||
|
content: buildPreInitHint(),
|
||||||
|
display: true,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
// Maps exist but advisory mode does not permit automatic artifact injection yet.
|
||||||
|
// Full root-pair injection with canonical markers comes in later slices.
|
||||||
|
if (!modeAllowsInjection(config.promptInjectionMode)) {
|
||||||
|
return {};
|
||||||
|
}
|
||||||
|
|
||||||
return {
|
return {
|
||||||
message: {
|
message: {
|
||||||
|
|||||||
@@ -1,6 +1,8 @@
|
|||||||
import { existsSync, readFileSync } from "fs";
|
import { existsSync, readFileSync } from "fs";
|
||||||
import { join } from "path";
|
import { join } from "path";
|
||||||
|
|
||||||
|
export type PromptInjectionMode = "off" | "advisory" | "strong" | "strict";
|
||||||
|
|
||||||
export interface SkillConfig {
|
export interface SkillConfig {
|
||||||
ignorePatterns: string[];
|
ignorePatterns: string[];
|
||||||
smallPackageThreshold: number;
|
smallPackageThreshold: number;
|
||||||
@@ -11,6 +13,9 @@ export interface SkillConfig {
|
|||||||
autoInjectPrompt: boolean;
|
autoInjectPrompt: boolean;
|
||||||
tagCap: number;
|
tagCap: number;
|
||||||
workflowHintCap: number;
|
workflowHintCap: number;
|
||||||
|
promptInjectionMode: PromptInjectionMode;
|
||||||
|
contextBudgetPercent: number;
|
||||||
|
contextBudgetMaxTokens: number;
|
||||||
}
|
}
|
||||||
|
|
||||||
export const DEFAULT_CONFIG: SkillConfig = {
|
export const DEFAULT_CONFIG: SkillConfig = {
|
||||||
@@ -43,6 +48,9 @@ export const DEFAULT_CONFIG: SkillConfig = {
|
|||||||
autoInjectPrompt: true,
|
autoInjectPrompt: true,
|
||||||
tagCap: 8,
|
tagCap: 8,
|
||||||
workflowHintCap: 5,
|
workflowHintCap: 5,
|
||||||
|
promptInjectionMode: "strong",
|
||||||
|
contextBudgetPercent: 15,
|
||||||
|
contextBudgetMaxTokens: 100_000,
|
||||||
};
|
};
|
||||||
|
|
||||||
export function loadConfig(cwd: string = process.cwd()): SkillConfig {
|
export function loadConfig(cwd: string = process.cwd()): SkillConfig {
|
||||||
|
|||||||
@@ -23,3 +23,15 @@ export {
|
|||||||
} from "./directory-model.js";
|
} from "./directory-model.js";
|
||||||
export { populateRoutingMetadata } from "./routing-metadata.js";
|
export { populateRoutingMetadata } from "./routing-metadata.js";
|
||||||
export { retrieveContext } from "./retrieve.js";
|
export { retrieveContext } from "./retrieve.js";
|
||||||
|
export {
|
||||||
|
buildRootPairBlock,
|
||||||
|
hasRootPairMarker,
|
||||||
|
buildPreInitHint,
|
||||||
|
computeInjectionBudget,
|
||||||
|
modeAllowsPreInitHint,
|
||||||
|
modeAllowsInjection,
|
||||||
|
modeRequiresProtocolPath,
|
||||||
|
ROOT_PAIR_START_MARKER,
|
||||||
|
ROOT_PAIR_END_MARKER,
|
||||||
|
TRUST_BOUNDARY_TEXT,
|
||||||
|
} from "./prompt-injection.js";
|
||||||
|
|||||||
@@ -0,0 +1,94 @@
|
|||||||
|
import type { SkillConfig, PromptInjectionMode } from "./config.js";
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Canonical markers for injected root-pair content.
|
||||||
|
* These must be stable across turns and easy to scan in outgoing context.
|
||||||
|
*/
|
||||||
|
export const ROOT_PAIR_START_MARKER = "<!-- PI_MAP_ROOT_PAIR_START -->";
|
||||||
|
export const ROOT_PAIR_END_MARKER = "<!-- PI_MAP_ROOT_PAIR_END -->";
|
||||||
|
export const TRUST_BOUNDARY_TEXT =
|
||||||
|
"Trust boundary: index routes, map orients, source decides.";
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Build a canonical root-pair block wrapping index and map content.
|
||||||
|
*/
|
||||||
|
export function buildRootPairBlock(
|
||||||
|
indexContent: string,
|
||||||
|
mapContent: string,
|
||||||
|
): string {
|
||||||
|
return [
|
||||||
|
ROOT_PAIR_START_MARKER,
|
||||||
|
"## Project Map Protocol",
|
||||||
|
"",
|
||||||
|
"1. Read this protocol and the root `.pi-map.index.md` first.",
|
||||||
|
"",
|
||||||
|
TRUST_BOUNDARY_TEXT,
|
||||||
|
"",
|
||||||
|
"### Root index",
|
||||||
|
indexContent,
|
||||||
|
"",
|
||||||
|
"### Root map",
|
||||||
|
mapContent,
|
||||||
|
ROOT_PAIR_END_MARKER,
|
||||||
|
].join("\n");
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Check whether a message content string contains the canonical root-pair marker.
|
||||||
|
*/
|
||||||
|
export function hasRootPairMarker(content: string): boolean {
|
||||||
|
return content.includes(ROOT_PAIR_START_MARKER);
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Build a lightweight user-visible pre-init startup hint.
|
||||||
|
* No synthetic artifact content is injected — only a prompt to run init.
|
||||||
|
*/
|
||||||
|
export function buildPreInitHint(): string {
|
||||||
|
return [
|
||||||
|
"📋 Project maps not initialized.",
|
||||||
|
"",
|
||||||
|
"Run `project_map_init` to generate paired `.pi-map.md` and `.pi-map.index.md` artifacts for this project.",
|
||||||
|
"After init, the agent will automatically use the root index for routing and the root map for orientation.",
|
||||||
|
].join("\n");
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Compute the effective injection budget in tokens.
|
||||||
|
* Uses the smaller of (percent of context window) and absolute cap.
|
||||||
|
* Falls back to absolute cap if context window is unknown.
|
||||||
|
*/
|
||||||
|
export function computeInjectionBudget(
|
||||||
|
config: Pick<SkillConfig, "contextBudgetPercent" | "contextBudgetMaxTokens">,
|
||||||
|
contextWindow?: number,
|
||||||
|
): number {
|
||||||
|
const absolute = config.contextBudgetMaxTokens;
|
||||||
|
if (contextWindow === undefined || contextWindow <= 0) {
|
||||||
|
return absolute;
|
||||||
|
}
|
||||||
|
const relative = Math.floor(
|
||||||
|
(contextWindow * config.contextBudgetPercent) / 100,
|
||||||
|
);
|
||||||
|
return Math.min(relative, absolute);
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Determine whether the active mode permits pre-init hints.
|
||||||
|
*/
|
||||||
|
export function modeAllowsPreInitHint(mode: PromptInjectionMode): boolean {
|
||||||
|
return mode !== "off";
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Determine whether the active mode permits automatic artifact injection after init.
|
||||||
|
*/
|
||||||
|
export function modeAllowsInjection(mode: PromptInjectionMode): boolean {
|
||||||
|
return mode === "strong" || mode === "strict";
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Resolve whether a given mode requires the protocol path for sensitive actions.
|
||||||
|
*/
|
||||||
|
export function modeRequiresProtocolPath(mode: PromptInjectionMode): boolean {
|
||||||
|
return mode === "strict";
|
||||||
|
}
|
||||||
@@ -271,9 +271,13 @@ describe("pi-extension", () => {
|
|||||||
});
|
});
|
||||||
|
|
||||||
describe("before_agent_start event", () => {
|
describe("before_agent_start event", () => {
|
||||||
it("injects layered protocol hint when project map files exist", async () => {
|
it("injects hidden hint when maps exist and mode is strong", async () => {
|
||||||
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
||||||
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTest\n");
|
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTest\n");
|
||||||
|
writeFileSync(
|
||||||
|
join(dir, ".pi-project-map.json"),
|
||||||
|
JSON.stringify({ promptInjectionMode: "strong" }),
|
||||||
|
);
|
||||||
mockCtx.cwd = dir;
|
mockCtx.cwd = dir;
|
||||||
|
|
||||||
const handler = registeredEvents.before_agent_start;
|
const handler = registeredEvents.before_agent_start;
|
||||||
@@ -286,8 +290,12 @@ describe("pi-extension", () => {
|
|||||||
expect(result.message.display).toBe(false);
|
expect(result.message.display).toBe(false);
|
||||||
});
|
});
|
||||||
|
|
||||||
it("returns empty object when no maps exist", async () => {
|
it("returns empty object when no maps exist and mode is off", async () => {
|
||||||
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
||||||
|
writeFileSync(
|
||||||
|
join(dir, ".pi-project-map.json"),
|
||||||
|
JSON.stringify({ promptInjectionMode: "off" }),
|
||||||
|
);
|
||||||
mockCtx.cwd = dir;
|
mockCtx.cwd = dir;
|
||||||
|
|
||||||
const handler = registeredEvents.before_agent_start;
|
const handler = registeredEvents.before_agent_start;
|
||||||
@@ -295,5 +303,86 @@ describe("pi-extension", () => {
|
|||||||
|
|
||||||
expect(result).toEqual({});
|
expect(result).toEqual({});
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it("shows visible pre-init hint when no maps exist and mode is strong", async () => {
|
||||||
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
||||||
|
writeFileSync(
|
||||||
|
join(dir, ".pi-project-map.json"),
|
||||||
|
JSON.stringify({ promptInjectionMode: "strong" }),
|
||||||
|
);
|
||||||
|
mockCtx.cwd = dir;
|
||||||
|
|
||||||
|
const handler = registeredEvents.before_agent_start;
|
||||||
|
const result = await handler(null, mockCtx);
|
||||||
|
|
||||||
|
expect(result).toHaveProperty("message");
|
||||||
|
expect(result.message.display).toBe(true);
|
||||||
|
expect(result.message.content).toContain("project_map_init");
|
||||||
|
expect(result.message.content).toContain("📋");
|
||||||
|
});
|
||||||
|
|
||||||
|
it("shows visible pre-init hint when no maps exist and mode is advisory", async () => {
|
||||||
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
||||||
|
writeFileSync(
|
||||||
|
join(dir, ".pi-project-map.json"),
|
||||||
|
JSON.stringify({ promptInjectionMode: "advisory" }),
|
||||||
|
);
|
||||||
|
mockCtx.cwd = dir;
|
||||||
|
|
||||||
|
const handler = registeredEvents.before_agent_start;
|
||||||
|
const result = await handler(null, mockCtx);
|
||||||
|
|
||||||
|
expect(result).toHaveProperty("message");
|
||||||
|
expect(result.message.display).toBe(true);
|
||||||
|
expect(result.message.content).toContain("project_map_init");
|
||||||
|
});
|
||||||
|
|
||||||
|
it("returns empty object when maps exist and mode is advisory", async () => {
|
||||||
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
||||||
|
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTest\n");
|
||||||
|
writeFileSync(
|
||||||
|
join(dir, ".pi-project-map.json"),
|
||||||
|
JSON.stringify({ promptInjectionMode: "advisory" }),
|
||||||
|
);
|
||||||
|
mockCtx.cwd = dir;
|
||||||
|
|
||||||
|
const handler = registeredEvents.before_agent_start;
|
||||||
|
const result = await handler(null, mockCtx);
|
||||||
|
|
||||||
|
expect(result).toEqual({});
|
||||||
|
});
|
||||||
|
|
||||||
|
it("shows visible pre-init hint when no maps exist and mode is strict", async () => {
|
||||||
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
||||||
|
writeFileSync(
|
||||||
|
join(dir, ".pi-project-map.json"),
|
||||||
|
JSON.stringify({ promptInjectionMode: "strict" }),
|
||||||
|
);
|
||||||
|
mockCtx.cwd = dir;
|
||||||
|
|
||||||
|
const handler = registeredEvents.before_agent_start;
|
||||||
|
const result = await handler(null, mockCtx);
|
||||||
|
|
||||||
|
expect(result).toHaveProperty("message");
|
||||||
|
expect(result.message.display).toBe(true);
|
||||||
|
expect(result.message.content).toContain("project_map_init");
|
||||||
|
});
|
||||||
|
|
||||||
|
it("injects hidden hint when maps exist and mode is strict", async () => {
|
||||||
|
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
||||||
|
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nTest\n");
|
||||||
|
writeFileSync(
|
||||||
|
join(dir, ".pi-project-map.json"),
|
||||||
|
JSON.stringify({ promptInjectionMode: "strict" }),
|
||||||
|
);
|
||||||
|
mockCtx.cwd = dir;
|
||||||
|
|
||||||
|
const handler = registeredEvents.before_agent_start;
|
||||||
|
const result = await handler(null, mockCtx);
|
||||||
|
|
||||||
|
expect(result).toHaveProperty("message");
|
||||||
|
expect(result.message.display).toBe(false);
|
||||||
|
expect(result.message.content).toContain("root `.pi-map.index.md`");
|
||||||
|
});
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -0,0 +1,100 @@
|
|||||||
|
import { describe, it, expect } from "vitest";
|
||||||
|
import {
|
||||||
|
buildRootPairBlock,
|
||||||
|
hasRootPairMarker,
|
||||||
|
buildPreInitHint,
|
||||||
|
computeInjectionBudget,
|
||||||
|
modeAllowsPreInitHint,
|
||||||
|
modeAllowsInjection,
|
||||||
|
modeRequiresProtocolPath,
|
||||||
|
ROOT_PAIR_START_MARKER,
|
||||||
|
ROOT_PAIR_END_MARKER,
|
||||||
|
TRUST_BOUNDARY_TEXT,
|
||||||
|
} from "../src/prompt-injection.js";
|
||||||
|
|
||||||
|
describe("prompt-injection helpers", () => {
|
||||||
|
describe("canonical markers", () => {
|
||||||
|
it("builds a root-pair block with markers", () => {
|
||||||
|
const block = buildRootPairBlock("index content", "map content");
|
||||||
|
expect(block).toContain(ROOT_PAIR_START_MARKER);
|
||||||
|
expect(block).toContain(ROOT_PAIR_END_MARKER);
|
||||||
|
expect(block).toContain(TRUST_BOUNDARY_TEXT);
|
||||||
|
expect(block).toContain("index content");
|
||||||
|
expect(block).toContain("map content");
|
||||||
|
});
|
||||||
|
|
||||||
|
it("detects root-pair marker in content", () => {
|
||||||
|
const block = buildRootPairBlock("i", "m");
|
||||||
|
expect(hasRootPairMarker(block)).toBe(true);
|
||||||
|
expect(hasRootPairMarker("plain text")).toBe(false);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
describe("pre-init hint", () => {
|
||||||
|
it("returns a visible hint with no synthetic artifacts", () => {
|
||||||
|
const hint = buildPreInitHint();
|
||||||
|
expect(hint).toContain("project_map_init");
|
||||||
|
expect(hint).toContain("📋");
|
||||||
|
expect(hint).not.toContain(ROOT_PAIR_START_MARKER);
|
||||||
|
expect(hint).not.toContain("### Root index");
|
||||||
|
expect(hint).not.toContain("### Root map");
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
describe("budget calculation", () => {
|
||||||
|
it("uses relative percent when smaller than absolute cap", () => {
|
||||||
|
const budget = computeInjectionBudget(
|
||||||
|
{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
|
||||||
|
200_000,
|
||||||
|
);
|
||||||
|
expect(budget).toBe(30_000);
|
||||||
|
});
|
||||||
|
|
||||||
|
it("uses absolute cap when smaller than relative percent", () => {
|
||||||
|
const budget = computeInjectionBudget(
|
||||||
|
{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
|
||||||
|
10_000_000,
|
||||||
|
);
|
||||||
|
expect(budget).toBe(100_000);
|
||||||
|
});
|
||||||
|
|
||||||
|
it("falls back to absolute cap when context window is unknown", () => {
|
||||||
|
const budget = computeInjectionBudget(
|
||||||
|
{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
|
||||||
|
undefined,
|
||||||
|
);
|
||||||
|
expect(budget).toBe(100_000);
|
||||||
|
});
|
||||||
|
|
||||||
|
it("falls back to absolute cap when context window is zero", () => {
|
||||||
|
const budget = computeInjectionBudget(
|
||||||
|
{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
|
||||||
|
0,
|
||||||
|
);
|
||||||
|
expect(budget).toBe(100_000);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
describe("mode helpers", () => {
|
||||||
|
it("allows pre-init hint for all modes except off", () => {
|
||||||
|
expect(modeAllowsPreInitHint("off")).toBe(false);
|
||||||
|
expect(modeAllowsPreInitHint("advisory")).toBe(true);
|
||||||
|
expect(modeAllowsPreInitHint("strong")).toBe(true);
|
||||||
|
expect(modeAllowsPreInitHint("strict")).toBe(true);
|
||||||
|
});
|
||||||
|
|
||||||
|
it("allows injection only for strong and strict", () => {
|
||||||
|
expect(modeAllowsInjection("off")).toBe(false);
|
||||||
|
expect(modeAllowsInjection("advisory")).toBe(false);
|
||||||
|
expect(modeAllowsInjection("strong")).toBe(true);
|
||||||
|
expect(modeAllowsInjection("strict")).toBe(true);
|
||||||
|
});
|
||||||
|
|
||||||
|
it("requires protocol path only for strict", () => {
|
||||||
|
expect(modeRequiresProtocolPath("off")).toBe(false);
|
||||||
|
expect(modeRequiresProtocolPath("advisory")).toBe(false);
|
||||||
|
expect(modeRequiresProtocolPath("strong")).toBe(false);
|
||||||
|
expect(modeRequiresProtocolPath("strict")).toBe(true);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
});
|
||||||
Reference in New Issue
Block a user