feat(prompt): implement prompt injection slice 2
This commit is contained in:
+7
-4
@@ -11,6 +11,8 @@ import {
|
||||
buildPreInitHint,
|
||||
modeAllowsPreInitHint,
|
||||
modeAllowsInjection,
|
||||
discoverContextWindow,
|
||||
buildInjectionPayload,
|
||||
} from "./src/index.js";
|
||||
import { loadConfig } from "./src/config.js";
|
||||
import { createLLMClient } from "./src/llm/llm-client.js";
|
||||
@@ -367,17 +369,18 @@ export default function (pi: ExtensionAPI) {
|
||||
}
|
||||
|
||||
// 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 {};
|
||||
}
|
||||
|
||||
// Slice 2: post-init root-pair preload + budgeted expansion
|
||||
const contextWindow = discoverContextWindow(_ctx);
|
||||
const payload = buildInjectionPayload(_ctx.cwd, config, contextWindow);
|
||||
return {
|
||||
message: {
|
||||
customType: "pi-project-map-hint",
|
||||
content:
|
||||
"📋 Project map active: Start with the root `.pi-map.index.md`, use indexes first for routing, read the local `.pi-map.md` plus source before edits, run `project_map_patch` after source edits, and run `project_map_validate` before freshness-sensitive architectural handoff.",
|
||||
display: false,
|
||||
content: payload.content,
|
||||
display: payload.display,
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
@@ -31,6 +31,10 @@ export {
|
||||
modeAllowsPreInitHint,
|
||||
modeAllowsInjection,
|
||||
modeRequiresProtocolPath,
|
||||
discoverContextWindow,
|
||||
estimateTokens,
|
||||
findAllArtifactPairs,
|
||||
buildInjectionPayload,
|
||||
ROOT_PAIR_START_MARKER,
|
||||
ROOT_PAIR_END_MARKER,
|
||||
TRUST_BOUNDARY_TEXT,
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
import { readdirSync, statSync, readFileSync } from "fs";
|
||||
import { join, relative } from "path";
|
||||
import type { SkillConfig, PromptInjectionMode } from "./config.js";
|
||||
|
||||
/**
|
||||
@@ -92,3 +94,179 @@ export function modeAllowsInjection(mode: PromptInjectionMode): boolean {
|
||||
export function modeRequiresProtocolPath(mode: PromptInjectionMode): boolean {
|
||||
return mode === "strict";
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Context-window discovery (Slice 2)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const KNOWN_MODEL_CONTEXT_WINDOWS: Record<string, number> = {
|
||||
"gpt-4o": 128_000,
|
||||
"gpt-4o-mini": 128_000,
|
||||
"gpt-4-turbo": 128_000,
|
||||
"gpt-4": 8_192,
|
||||
"claude-3-5-sonnet": 200_000,
|
||||
"claude-3-opus": 200_000,
|
||||
"kimi-for-coding": 200_000,
|
||||
k2p6: 1_000_000,
|
||||
"kimi-k2-thinking": 256_000,
|
||||
};
|
||||
|
||||
/**
|
||||
* Discover the active model's context-window size from Pi runtime metadata.
|
||||
* Returns `undefined` when unavailable so callers fall back to the absolute cap.
|
||||
*/
|
||||
export function discoverContextWindow(ctx: any): number | undefined {
|
||||
const model = ctx?.model;
|
||||
if (model) {
|
||||
if (typeof model.contextWindow === "number" && model.contextWindow > 0) {
|
||||
return model.contextWindow;
|
||||
}
|
||||
if (
|
||||
typeof model.maxContextTokens === "number" &&
|
||||
model.maxContextTokens > 0
|
||||
) {
|
||||
return model.maxContextTokens;
|
||||
}
|
||||
if (typeof model.id === "string") {
|
||||
const known = KNOWN_MODEL_CONTEXT_WINDOWS[model.id];
|
||||
if (known) return known;
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Token estimation (best-effort, deterministic)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Rough token estimate from character count.
|
||||
* 1 token ≈ 4 chars for English/prose is a conservative heuristic.
|
||||
*/
|
||||
export function estimateTokens(text: string): number {
|
||||
return Math.ceil(text.length / 4);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Artifact-pair discovery
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export interface ArtifactPair {
|
||||
dir: string;
|
||||
mapPath: string;
|
||||
indexPath: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Find every directory under `cwd` that contains both `.pi-map.md` and
|
||||
* `.pi-map.index.md`. Returns shallowest directories first for predictable
|
||||
* structural coverage.
|
||||
*/
|
||||
export function findAllArtifactPairs(cwd: string): ArtifactPair[] {
|
||||
const results: ArtifactPair[] = [];
|
||||
|
||||
function walk(dir: string) {
|
||||
let entries: import("fs").Dirent[];
|
||||
try {
|
||||
entries = readdirSync(dir, { withFileTypes: true });
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
for (const entry of entries) {
|
||||
if (
|
||||
entry.isDirectory() &&
|
||||
!entry.name.startsWith(".") &&
|
||||
entry.name !== "node_modules"
|
||||
) {
|
||||
walk(join(dir, entry.name));
|
||||
}
|
||||
}
|
||||
const mapPath = join(dir, ".pi-map.md");
|
||||
const indexPath = join(dir, ".pi-map.index.md");
|
||||
try {
|
||||
statSync(mapPath);
|
||||
statSync(indexPath);
|
||||
results.push({
|
||||
dir: relative(cwd, dir) || ".",
|
||||
mapPath: relative(cwd, mapPath),
|
||||
indexPath: relative(cwd, indexPath),
|
||||
});
|
||||
} catch {
|
||||
// skip dirs without the paired artifacts
|
||||
}
|
||||
}
|
||||
|
||||
walk(cwd);
|
||||
|
||||
// shallow-first: sort by path depth then alphabetically
|
||||
results.sort((a, b) => {
|
||||
const depthA = a.dir.split(/[/\\]/).length;
|
||||
const depthB = b.dir.split(/[/\\]/).length;
|
||||
if (depthA !== depthB) return depthA - depthB;
|
||||
return a.dir.localeCompare(b.dir);
|
||||
});
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Budgeted expansion
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Build the full injection payload for a project that already has artifacts.
|
||||
*
|
||||
* Guarantees:
|
||||
* 1. Root pair is always present first.
|
||||
* 2. Additional pairs are appended shallow-first while the budget allows.
|
||||
* 3. A brief maintenance reminder is prepended.
|
||||
*/
|
||||
export function buildInjectionPayload(
|
||||
cwd: string,
|
||||
config: Pick<SkillConfig, "contextBudgetPercent" | "contextBudgetMaxTokens">,
|
||||
contextWindow: number | undefined,
|
||||
): { content: string; display: false } {
|
||||
const budget = computeInjectionBudget(config, contextWindow);
|
||||
|
||||
const pairs = findAllArtifactPairs(cwd);
|
||||
const rootIndex = pairs.findIndex((p) => p.dir === ".");
|
||||
let rootPair: ArtifactPair | undefined;
|
||||
if (rootIndex >= 0) {
|
||||
rootPair = pairs.splice(rootIndex, 1)[0];
|
||||
}
|
||||
|
||||
let usedTokens = 0;
|
||||
const parts: string[] = [];
|
||||
|
||||
// Maintenance reminder (lightweight)
|
||||
const reminder =
|
||||
"📋 Project map active: Start with the root `.pi-map.index.md`, use indexes first for routing, read the local `.pi-map.md` plus source before edits, run `project_map_patch` after source edits, and run `project_map_validate` before freshness-sensitive architectural handoff.";
|
||||
parts.push(reminder);
|
||||
usedTokens += estimateTokens(reminder);
|
||||
|
||||
// Root pair (guaranteed)
|
||||
if (rootPair) {
|
||||
const indexContent = readFileSync(join(cwd, rootPair.indexPath), "utf8");
|
||||
const mapContent = readFileSync(join(cwd, rootPair.mapPath), "utf8");
|
||||
const block = buildRootPairBlock(indexContent, mapContent);
|
||||
parts.push(block);
|
||||
usedTokens += estimateTokens(block);
|
||||
}
|
||||
|
||||
// Budgeted expansion (shallow-first, deterministic)
|
||||
for (const pair of pairs) {
|
||||
const indexContent = readFileSync(join(cwd, pair.indexPath), "utf8");
|
||||
const mapContent = readFileSync(join(cwd, pair.mapPath), "utf8");
|
||||
const pairTokens =
|
||||
estimateTokens(indexContent) + estimateTokens(mapContent);
|
||||
if (usedTokens + pairTokens > budget) {
|
||||
break;
|
||||
}
|
||||
parts.push(
|
||||
`\n## ${pair.dir}\n\n### Index\n${indexContent}\n\n### Map\n${mapContent}`,
|
||||
);
|
||||
usedTokens += pairTokens;
|
||||
}
|
||||
|
||||
return { content: parts.join("\n\n"), display: false };
|
||||
}
|
||||
|
||||
@@ -290,6 +290,33 @@ describe("pi-extension", () => {
|
||||
expect(result.message.display).toBe(false);
|
||||
});
|
||||
|
||||
it("injects canonical root-pair block with real file content when maps exist and mode is strong", async () => {
|
||||
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
||||
writeFileSync(join(dir, ".pi-map.md"), "# .\n## role\nRootMap\n");
|
||||
writeFileSync(
|
||||
join(dir, ".pi-map.index.md"),
|
||||
"# . (index)\n## role\nRootIndex\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(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("<!-- PI_MAP_ROOT_PAIR_END -->");
|
||||
expect(result.message.content).toContain("RootMap");
|
||||
expect(result.message.content).toContain("RootIndex");
|
||||
expect(result.message.content).toContain("Trust boundary:");
|
||||
});
|
||||
|
||||
it("returns empty object when no maps exist and mode is off", async () => {
|
||||
const dir = mkdtempSync(join(tmpdir(), "pi-ext-test-"));
|
||||
writeFileSync(
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { mkdtempSync, writeFileSync, mkdirSync } from "fs";
|
||||
import { join } from "path";
|
||||
import { tmpdir } from "os";
|
||||
import {
|
||||
buildRootPairBlock,
|
||||
hasRootPairMarker,
|
||||
@@ -7,6 +10,10 @@ import {
|
||||
modeAllowsPreInitHint,
|
||||
modeAllowsInjection,
|
||||
modeRequiresProtocolPath,
|
||||
discoverContextWindow,
|
||||
estimateTokens,
|
||||
findAllArtifactPairs,
|
||||
buildInjectionPayload,
|
||||
ROOT_PAIR_START_MARKER,
|
||||
ROOT_PAIR_END_MARKER,
|
||||
TRUST_BOUNDARY_TEXT,
|
||||
@@ -97,4 +104,133 @@ describe("prompt-injection helpers", () => {
|
||||
expect(modeRequiresProtocolPath("strict")).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe("context-window discovery", () => {
|
||||
it("reads contextWindow from model metadata", () => {
|
||||
const ctx = { model: { contextWindow: 256_000 } };
|
||||
expect(discoverContextWindow(ctx)).toBe(256_000);
|
||||
});
|
||||
|
||||
it("reads maxContextTokens from model metadata", () => {
|
||||
const ctx = { model: { maxContextTokens: 128_000 } };
|
||||
expect(discoverContextWindow(ctx)).toBe(128_000);
|
||||
});
|
||||
|
||||
it("falls back to known model lookup by id", () => {
|
||||
expect(discoverContextWindow({ model: { id: "gpt-4o-mini" } })).toBe(
|
||||
128_000,
|
||||
);
|
||||
expect(discoverContextWindow({ model: { id: "kimi-for-coding" } })).toBe(
|
||||
200_000,
|
||||
);
|
||||
});
|
||||
|
||||
it("returns undefined for unknown models", () => {
|
||||
expect(
|
||||
discoverContextWindow({ model: { id: "unknown-model" } }),
|
||||
).toBeUndefined();
|
||||
});
|
||||
|
||||
it("returns undefined when model is missing", () => {
|
||||
expect(discoverContextWindow({})).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("token estimation", () => {
|
||||
it("estimates tokens from character count", () => {
|
||||
expect(estimateTokens("")).toBe(0);
|
||||
expect(estimateTokens("abcd")).toBe(1);
|
||||
expect(estimateTokens("abcde")).toBe(2);
|
||||
});
|
||||
});
|
||||
|
||||
describe("artifact-pair discovery", () => {
|
||||
it("finds paired artifacts shallow-first", () => {
|
||||
const dir = mkdtempSync(join(tmpdir(), "pi-pair-test-"));
|
||||
mkdirSync(join(dir, "src", "nested"), { recursive: true });
|
||||
|
||||
writeFileSync(join(dir, ".pi-map.md"), "# root");
|
||||
writeFileSync(join(dir, ".pi-map.index.md"), "# root index");
|
||||
writeFileSync(join(dir, "src", ".pi-map.md"), "# src");
|
||||
writeFileSync(join(dir, "src", ".pi-map.index.md"), "# src index");
|
||||
writeFileSync(join(dir, "src", "nested", ".pi-map.md"), "# nested");
|
||||
writeFileSync(
|
||||
join(dir, "src", "nested", ".pi-map.index.md"),
|
||||
"# nested index",
|
||||
);
|
||||
|
||||
const pairs = findAllArtifactPairs(dir);
|
||||
expect(pairs.map((p) => p.dir)).toEqual([".", "src", "src/nested"]);
|
||||
});
|
||||
|
||||
it("ignores directories missing one of the paired artifacts", () => {
|
||||
const dir = mkdtempSync(join(tmpdir(), "pi-pair-test-"));
|
||||
writeFileSync(join(dir, ".pi-map.md"), "# root");
|
||||
// missing .pi-map.index.md intentionally
|
||||
|
||||
expect(findAllArtifactPairs(dir)).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe("injection payload", () => {
|
||||
it("always includes the root pair", () => {
|
||||
const dir = mkdtempSync(join(tmpdir(), "pi-payload-test-"));
|
||||
writeFileSync(join(dir, ".pi-map.md"), "# root map");
|
||||
writeFileSync(join(dir, ".pi-map.index.md"), "# root index");
|
||||
|
||||
const payload = buildInjectionPayload(
|
||||
dir,
|
||||
{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
|
||||
undefined,
|
||||
);
|
||||
|
||||
expect(payload.content).toContain(ROOT_PAIR_START_MARKER);
|
||||
expect(payload.content).toContain(ROOT_PAIR_END_MARKER);
|
||||
expect(payload.content).toContain("# root map");
|
||||
expect(payload.content).toContain("# root index");
|
||||
expect(payload.display).toBe(false);
|
||||
});
|
||||
|
||||
it("expands additional pairs when budget allows", () => {
|
||||
const dir = mkdtempSync(join(tmpdir(), "pi-payload-test-"));
|
||||
mkdirSync(join(dir, "src"), { recursive: true });
|
||||
writeFileSync(join(dir, ".pi-map.md"), "# root map");
|
||||
writeFileSync(join(dir, ".pi-map.index.md"), "# root index");
|
||||
writeFileSync(join(dir, "src", ".pi-map.md"), "# src map");
|
||||
writeFileSync(join(dir, "src", ".pi-map.index.md"), "# src index");
|
||||
|
||||
const payload = buildInjectionPayload(
|
||||
dir,
|
||||
{ contextBudgetPercent: 100, contextBudgetMaxTokens: 10_000 },
|
||||
10_000,
|
||||
);
|
||||
|
||||
expect(payload.content).toContain("# src map");
|
||||
expect(payload.content).toContain("# src index");
|
||||
});
|
||||
|
||||
it("stops expanding when budget is exhausted", () => {
|
||||
const dir = mkdtempSync(join(tmpdir(), "pi-payload-test-"));
|
||||
mkdirSync(join(dir, "src"), { recursive: true });
|
||||
writeFileSync(join(dir, ".pi-map.md"), "# root map");
|
||||
writeFileSync(join(dir, ".pi-map.index.md"), "# root index");
|
||||
writeFileSync(
|
||||
join(dir, "src", ".pi-map.md"),
|
||||
`# src map ${"x".repeat(400)}`,
|
||||
);
|
||||
writeFileSync(
|
||||
join(dir, "src", ".pi-map.index.md"),
|
||||
`# src index ${"x".repeat(400)}`,
|
||||
);
|
||||
|
||||
const payload = buildInjectionPayload(
|
||||
dir,
|
||||
{ contextBudgetPercent: 100, contextBudgetMaxTokens: 10 },
|
||||
10_000,
|
||||
);
|
||||
|
||||
expect(payload.content).toContain("# root map");
|
||||
expect(payload.content).not.toContain("# src map x");
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user