Developer
5f1c107667
The before_agent_start handler now scans the active session context via ctx.sessionManager.buildSessionContext() for an existing pi-project-map-hint custom message and skips injection when one is already present in the current branch. This prevents duplicate visible hints in advisory/pre-init modes and duplicate hidden hints in strong/strict modes. The hint is automatically re-injected after compaction or /tree navigation removes it from the active path. Also removes the unused hooks/on-prompt.ts prompt-text injector.
537 lines
15 KiB
TypeScript
537 lines
15 KiB
TypeScript
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
|
import { Type } from "typebox";
|
|
import { readFileSync, readdirSync, statSync } from "fs";
|
|
import { join, relative } from "path";
|
|
import {
|
|
initProject,
|
|
patchFile,
|
|
validateMaps,
|
|
reinitPath,
|
|
retrieveContext,
|
|
buildPreInitHint,
|
|
buildAdvisoryReminder,
|
|
buildStrictBypassGuard,
|
|
modeAllowsPreInitHint,
|
|
modeAllowsInjection,
|
|
evaluateStrictBypass,
|
|
discoverContextWindow,
|
|
buildInjectionPayload,
|
|
shouldReinjectForEvent,
|
|
getRootPairMtimes,
|
|
rootPairChanged,
|
|
} from "./src/index.js";
|
|
import { loadConfig } from "./src/config.js";
|
|
import { createLLMClient } from "./src/llm/llm-client.js";
|
|
import { LLMError } from "./src/llm/llm-error.js";
|
|
|
|
/**
|
|
* Get the LLM client for Pi runtime.
|
|
*
|
|
* When running inside Pi, we ALWAYS use Pi's native LLM via
|
|
* @mariozechner/pi-ai's complete() function. This respects the
|
|
* user's /model selection and /login auth. No external fetch().
|
|
*/
|
|
function getPiLLMClient(ctx: any) {
|
|
return createLLMClient("pi", { extensionContext: ctx });
|
|
}
|
|
|
|
function findPiMapFiles(cwd: string): string[] {
|
|
const results: string[] = [];
|
|
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));
|
|
}
|
|
}
|
|
try {
|
|
statSync(join(dir, ".pi-map.md"));
|
|
results.push(relative(cwd, join(dir, ".pi-map.md")));
|
|
} catch {
|
|
// no map in this dir
|
|
}
|
|
}
|
|
walk(cwd);
|
|
return results;
|
|
}
|
|
|
|
function isDirty(content: string): boolean {
|
|
return content.includes("## dirty") && !content.includes("## dirty\n-");
|
|
}
|
|
|
|
function renderProgressBar(
|
|
completed: number,
|
|
total: number,
|
|
currentFile?: string,
|
|
width = 20,
|
|
): string {
|
|
const pct = total > 0 ? completed / total : 0;
|
|
const filled = Math.round(width * pct);
|
|
const bar = "█".repeat(filled) + "░".repeat(width - filled);
|
|
const file = currentFile ? ` → ${currentFile}` : "";
|
|
return `[${bar}] ${completed}/${total}${file}`;
|
|
}
|
|
|
|
const HINT_CUSTOM_TYPE = "pi-project-map-hint";
|
|
|
|
function hintAlreadyInContext(ctx: any): boolean {
|
|
const manager = ctx?.sessionManager;
|
|
if (!manager || typeof manager.buildSessionContext !== "function") {
|
|
return false;
|
|
}
|
|
|
|
const { messages } = manager.buildSessionContext();
|
|
if (!Array.isArray(messages)) return false;
|
|
|
|
return messages.some(
|
|
(m: any) =>
|
|
m &&
|
|
m.role === "custom" &&
|
|
m.customType === HINT_CUSTOM_TYPE,
|
|
);
|
|
}
|
|
|
|
export default function (pi: ExtensionAPI) {
|
|
let lastRootPairMtimes: import("./src/index.js").RootPairMtimes = {};
|
|
pi.registerTool({
|
|
name: "project_map_init",
|
|
label: "Project Map Init",
|
|
description:
|
|
"Generate paired .pi-map.md and .pi-map.index.md analysis files for the entire project or a subdirectory",
|
|
promptSnippet:
|
|
"Initialize paired project map/index artifacts for codebase understanding",
|
|
promptGuidelines: [
|
|
"Use project_map_init when starting work on a new project or after significant restructuring",
|
|
"Run project_map_init when .pi-map.md / .pi-map.index.md files are missing or severely outdated",
|
|
],
|
|
parameters: Type.Object({
|
|
path: Type.Optional(
|
|
Type.String({
|
|
description: "Project root path (default: current directory)",
|
|
}),
|
|
),
|
|
}),
|
|
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
|
|
try {
|
|
const targetPath = params.path || ctx.cwd;
|
|
const client = getPiLLMClient(ctx);
|
|
await initProject(targetPath, {
|
|
verbose: false,
|
|
llmClient: client,
|
|
cacheDir: ctx.cwd,
|
|
onProgress: (info) => {
|
|
const bar = renderProgressBar(
|
|
info.completed,
|
|
info.total,
|
|
info.currentFile,
|
|
);
|
|
_onUpdate?.({
|
|
content: [{ type: "text", text: bar }],
|
|
details: {
|
|
progress:
|
|
info.total > 0
|
|
? Math.round((info.completed / info.total) * 100)
|
|
: 0,
|
|
file: info.currentFile,
|
|
dir: info.dir,
|
|
},
|
|
});
|
|
},
|
|
});
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Generated .pi-map.md files for ${targetPath}`,
|
|
},
|
|
],
|
|
details: { success: true, cwd: ctx.cwd },
|
|
};
|
|
} catch (err: any) {
|
|
const msg = err instanceof LLMError ? err.message : String(err);
|
|
return {
|
|
content: [{ type: "text", text: `Error: ${msg}` }],
|
|
details: { success: false, error: msg },
|
|
};
|
|
}
|
|
},
|
|
});
|
|
|
|
pi.registerTool({
|
|
name: "project_map_patch",
|
|
label: "Project Map Patch",
|
|
description:
|
|
"Update the paired .pi-map.md / .pi-map.index.md artifacts for the directory containing a changed file",
|
|
promptSnippet:
|
|
"Update paired project map/index artifacts after editing a source file",
|
|
promptGuidelines: [
|
|
"Use project_map_patch immediately after editing any source file",
|
|
"Pass the absolute or relative path of the modified file",
|
|
],
|
|
parameters: Type.Object({
|
|
file_path: Type.String({
|
|
description: "Path to the modified file",
|
|
}),
|
|
}),
|
|
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
|
|
try {
|
|
const client = getPiLLMClient(ctx);
|
|
await patchFile(params.file_path, client, ctx.cwd);
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Patched map for ${params.file_path}`,
|
|
},
|
|
],
|
|
details: { success: true },
|
|
};
|
|
} catch (err: any) {
|
|
const msg = err instanceof LLMError ? err.message : String(err);
|
|
return {
|
|
content: [{ type: "text", text: `Error: ${msg}` }],
|
|
details: { success: false, error: msg },
|
|
};
|
|
}
|
|
},
|
|
});
|
|
|
|
pi.registerTool({
|
|
name: "project_map_validate",
|
|
label: "Project Map Validate",
|
|
description:
|
|
"Check all .pi-map.md / .pi-map.index.md files for staleness and discrepancies",
|
|
promptSnippet: "Validate paired project map/index artifacts for accuracy",
|
|
promptGuidelines: [
|
|
"Use project_map_validate before making architectural decisions if you suspect stale data",
|
|
"Use project_map_validate to detect files that were deleted or added outside the agent",
|
|
],
|
|
parameters: Type.Object({
|
|
path: Type.Optional(
|
|
Type.String({
|
|
description: "Project root path (default: current directory)",
|
|
}),
|
|
),
|
|
}),
|
|
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
|
|
try {
|
|
const targetPath = params.path || ctx.cwd;
|
|
const result = await validateMaps(targetPath, {
|
|
fix: false,
|
|
verbose: false,
|
|
});
|
|
const text = result.clean
|
|
? "All .pi-map.md files are clean."
|
|
: `Found ${result.discrepancies.length} discrepancies:\n` +
|
|
result.discrepancies
|
|
.map((d) => ` [${d.type}] ${d.path}: ${d.message}`)
|
|
.join("\n");
|
|
return {
|
|
content: [{ type: "text", text }],
|
|
details: { success: true, clean: result.clean },
|
|
};
|
|
} catch (err: any) {
|
|
const msg = err instanceof LLMError ? err.message : String(err);
|
|
return {
|
|
content: [{ type: "text", text: `Error: ${msg}` }],
|
|
details: { success: false, error: msg },
|
|
};
|
|
}
|
|
},
|
|
});
|
|
|
|
pi.registerTool({
|
|
name: "project_map_reinit",
|
|
label: "Project Map Reinit",
|
|
description:
|
|
"Force full regeneration of all .pi-map.md / .pi-map.index.md artifacts",
|
|
promptSnippet:
|
|
"Force full regeneration of paired project map/index artifacts",
|
|
promptGuidelines: [
|
|
"Use project_map_reinit when validation shows widespread staleness",
|
|
"Use project_map_reinit after pulling major changes from version control",
|
|
],
|
|
parameters: Type.Object({
|
|
path: Type.Optional(
|
|
Type.String({
|
|
description: "Path to regenerate (default: entire project)",
|
|
}),
|
|
),
|
|
}),
|
|
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
|
|
try {
|
|
const targetPath = params.path || ctx.cwd;
|
|
const client = getPiLLMClient(ctx);
|
|
await reinitPath(targetPath, {
|
|
verbose: false,
|
|
llmClient: client,
|
|
cacheDir: ctx.cwd,
|
|
onProgress: (info) => {
|
|
const bar = renderProgressBar(
|
|
info.completed,
|
|
info.total,
|
|
info.currentFile,
|
|
);
|
|
_onUpdate?.({
|
|
content: [{ type: "text", text: bar }],
|
|
details: {
|
|
progress:
|
|
info.total > 0
|
|
? Math.round((info.completed / info.total) * 100)
|
|
: 0,
|
|
file: info.currentFile,
|
|
dir: info.dir,
|
|
},
|
|
});
|
|
},
|
|
});
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Regenerated maps for ${targetPath}`,
|
|
},
|
|
],
|
|
details: { success: true },
|
|
};
|
|
} catch (err: any) {
|
|
const msg = err instanceof LLMError ? err.message : String(err);
|
|
return {
|
|
content: [{ type: "text", text: `Error: ${msg}` }],
|
|
details: { success: false, error: msg },
|
|
};
|
|
}
|
|
},
|
|
});
|
|
|
|
pi.registerTool({
|
|
name: "project_map_context",
|
|
label: "Project Map Context",
|
|
description:
|
|
"Retrieve a compact markdown context bundle for a natural-language query using paired project map/index metadata",
|
|
promptSnippet:
|
|
"Get relevant project context for a task without reading every source file",
|
|
promptGuidelines: [
|
|
"Use project_map_context when you need to understand a task area before diving into source",
|
|
"Pass a concise query describing the feature, bug, or area you want to explore",
|
|
"Always read the suggested indexes first, then maps, then verify from source",
|
|
],
|
|
parameters: Type.Object({
|
|
query: Type.String({
|
|
description:
|
|
"Natural-language query describing the task or area to explore",
|
|
}),
|
|
}),
|
|
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
|
|
try {
|
|
const bundle = retrieveContext(params.query, ctx.cwd);
|
|
return {
|
|
content: [{ type: "text", text: bundle }],
|
|
details: { success: true },
|
|
};
|
|
} catch (err: any) {
|
|
const msg = err instanceof LLMError ? err.message : String(err);
|
|
return {
|
|
content: [{ type: "text", text: `Error: ${msg}` }],
|
|
details: { success: false, error: msg },
|
|
};
|
|
}
|
|
},
|
|
});
|
|
|
|
// Auto-load .pi-map.md files on session start
|
|
pi.on("session_start", async (_event, ctx) => {
|
|
const mapFiles = findPiMapFiles(ctx.cwd);
|
|
if (mapFiles.length === 0) return;
|
|
|
|
const dirtyFiles = mapFiles.filter((f) => {
|
|
try {
|
|
const content = readFileSync(join(ctx.cwd, f), "utf8");
|
|
return isDirty(content);
|
|
} catch {
|
|
return false;
|
|
}
|
|
});
|
|
|
|
if (dirtyFiles.length > 0) {
|
|
ctx.ui.notify(
|
|
`pi-project-map: ${dirtyFiles.length} dirty packages detected. Run project_map_validate or project_map_reinit.`,
|
|
"warning",
|
|
);
|
|
}
|
|
});
|
|
|
|
// Inject maintenance instructions before agent starts, but only once
|
|
// within the current branch of context. Re-inject after compaction or
|
|
// tree navigation removes the hint from the active path.
|
|
pi.on("before_agent_start", async (_event, _ctx) => {
|
|
const config = loadConfig(_ctx.cwd);
|
|
const mapFiles = findPiMapFiles(_ctx.cwd);
|
|
|
|
// 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 {};
|
|
}
|
|
if (hintAlreadyInContext(_ctx)) return {};
|
|
return {
|
|
message: {
|
|
customType: HINT_CUSTOM_TYPE,
|
|
content: buildPreInitHint(),
|
|
display: true,
|
|
},
|
|
};
|
|
}
|
|
|
|
// Slice 4: advisory mode shows a visible lightweight reminder after init.
|
|
// No root-pair preload, no per-turn reinjection.
|
|
if (config.promptInjectionMode === "advisory") {
|
|
if (hintAlreadyInContext(_ctx)) return {};
|
|
return {
|
|
message: {
|
|
customType: HINT_CUSTOM_TYPE,
|
|
content: buildAdvisoryReminder(),
|
|
display: true,
|
|
},
|
|
};
|
|
}
|
|
|
|
// Maps exist but mode does not permit automatic artifact injection.
|
|
if (!modeAllowsInjection(config.promptInjectionMode)) {
|
|
return {};
|
|
}
|
|
|
|
// Slice 3b: detect root-pair artifact changes
|
|
const currentMtimes = getRootPairMtimes(_ctx.cwd);
|
|
const hasPrevious =
|
|
lastRootPairMtimes.mapMtime !== undefined ||
|
|
lastRootPairMtimes.indexMtime !== undefined;
|
|
const artifactChanged =
|
|
hasPrevious && rootPairChanged(currentMtimes, lastRootPairMtimes);
|
|
lastRootPairMtimes = currentMtimes;
|
|
|
|
// Slice 3a/3b: avoid redundant reinjection by scanning outgoing context
|
|
const eventType = artifactChanged ? "artifact_change" : "agent_start";
|
|
const decision = shouldReinjectForEvent(
|
|
{
|
|
messages: _event?.messages,
|
|
type: eventType,
|
|
},
|
|
config.promptInjectionMode,
|
|
);
|
|
if (!decision.needed) {
|
|
return {};
|
|
}
|
|
|
|
if (hintAlreadyInContext(_ctx)) return {};
|
|
|
|
// Slice 2: post-init root-pair preload + budgeted expansion
|
|
const contextWindow = discoverContextWindow(_ctx);
|
|
const payload = buildInjectionPayload(_ctx.cwd, config, contextWindow);
|
|
return {
|
|
message: {
|
|
customType: HINT_CUSTOM_TYPE,
|
|
content: payload.content,
|
|
display: payload.display,
|
|
},
|
|
};
|
|
});
|
|
|
|
// Per-turn context scanning for reinjection in strong/strict modes
|
|
pi.on("context", async (event: any, ctx: any) => {
|
|
const config = loadConfig(ctx.cwd);
|
|
const mapFiles = findPiMapFiles(ctx.cwd);
|
|
|
|
if (mapFiles.length === 0) {
|
|
return {};
|
|
}
|
|
|
|
if (!modeAllowsInjection(config.promptInjectionMode)) {
|
|
return {};
|
|
}
|
|
|
|
// 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 ||
|
|
lastRootPairMtimes.indexMtime !== undefined;
|
|
const artifactChanged =
|
|
hasPrevious && rootPairChanged(currentMtimes, lastRootPairMtimes);
|
|
lastRootPairMtimes = currentMtimes;
|
|
|
|
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: event?.type,
|
|
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,
|
|
},
|
|
};
|
|
});
|
|
}
|