@zihanw/pi-forge 0.2.0 → 0.3.1

package/src/lifecycle.ts

@@ -0,0 +1,171 @@
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import {
+	compileMessages,
+	compileSystemPrompt,
+	createPromptVariableStore,
+	getLatestUserMessage,
+	markSessionVariablesClean,
+	resetTurnVariables,
+} from "./compiler.ts";
+import { applyFinalizeRegexRulesToMessage } from "./regex.ts";
+import { STATE_ENTRY_TYPE, VARIABLE_ENTRY_TYPE, type PiForgeRuntimeState } from "./runtime-state.ts";
+import type { PromptStackDiagnostic, PromptVariableStore, PromptVariableValue } from "./types.ts";
+
+export interface LifecycleDeps {
+	reloadStacks(ctx: ExtensionContext, preferredId?: string): void;
+	refreshWebEditorHost(ctx: ExtensionContext): void;
+	notifyActivePreset(ctx: ExtensionContext, detail: string): void;
+	syncActiveToolPolicy(ctx?: ExtensionContext): void;
+	activeId(): string | undefined;
+	persistActiveSelection(id: string): void;
+	recordCompileDiagnostics(ctx: ExtensionContext, diagnostics: PromptStackDiagnostic[]): void;
+}
+
+export function registerLifecycleHandlers(pi: ExtensionAPI, state: PiForgeRuntimeState, deps: LifecycleDeps): void {
+	pi.on("session_start", async (event, ctx) => {
+		restoreBranchScopedRuntime(ctx, state, deps);
+		deps.refreshWebEditorHost(ctx);
+		deps.notifyActivePreset(ctx, "after session " + event.reason);
+	});
+
+	pi.on("session_tree", async (_event, ctx) => {
+		restoreBranchScopedRuntime(ctx, state, deps);
+		deps.refreshWebEditorHost(ctx);
+		deps.notifyActivePreset(ctx, "after tree navigation");
+	});
+
+	pi.on("session_compact", async (_event, ctx) => {
+		restoreBranchScopedRuntime(ctx, state, deps);
+		deps.refreshWebEditorHost(ctx);
+		deps.notifyActivePreset(ctx, "after compaction");
+	});
+
+	pi.on("turn_start", async () => {
+		deps.syncActiveToolPolicy();
+		const id = deps.activeId();
+		if (id) deps.persistActiveSelection(id);
+	});
+
+	pi.on("before_agent_start", async (event, ctx) => {
+		state.currentSystemPromptOptions = event.systemPromptOptions;
+		state.currentLatestUserMessage = event.prompt;
+		state.currentVariableStore = createPromptVariableStore(state.sessionVariables);
+		resetTurnVariables(state.currentVariableStore);
+		state.contextRewritePending = true;
+
+		if (!state.active) return;
+
+		const result = compileSystemPrompt(
+			state.active.stack,
+			{ options: event.systemPromptOptions, ctx, latestUserMessage: event.prompt, now: new Date(), variables: state.currentVariableStore },
+			event.systemPrompt,
+		);
+		deps.recordCompileDiagnostics(ctx, result.diagnostics);
+		persistVariablesIfDirty(pi, state, state.currentVariableStore);
+
+		return { systemPrompt: result.systemPrompt };
+	});
+
+	pi.on("context", async (event, ctx) => {
+		if (!state.active || !state.currentSystemPromptOptions || !state.contextRewritePending) return;
+
+		// Rewrite the message layout only for the first provider request of a user-submitted prompt.
+		// Tool-result follow-up turns must receive Pi's natural context; otherwise post-history
+		// prompt blocks such as COT / {{lastUserMessage}} are re-appended after every tool call
+		// and the model restarts its planning instead of continuing from the tool result.
+		state.contextRewritePending = false;
+
+		if (!state.currentVariableStore) state.currentVariableStore = createPromptVariableStore(state.sessionVariables);
+		const latestUserMessage = getLatestUserMessage(event.messages) ?? state.currentLatestUserMessage;
+		const result = compileMessages(
+			state.active.stack,
+			{ options: state.currentSystemPromptOptions, ctx, latestUserMessage, now: new Date(), variables: state.currentVariableStore },
+			event.messages,
+		);
+		deps.recordCompileDiagnostics(ctx, [...state.latestCompileDiagnostics, ...result.diagnostics]);
+		persistVariablesIfDirty(pi, state, state.currentVariableStore);
+		return { messages: result.messages };
+	});
+
+	pi.on("message_end", async (event, ctx) => {
+		if (!state.active) return;
+		const diagnostics: PromptStackDiagnostic[] = [];
+		const message = applyFinalizeRegexRulesToMessage(state.active.stack, event.message, diagnostics);
+		if (diagnostics.length > 0) deps.recordCompileDiagnostics(ctx, [...state.latestCompileDiagnostics, ...diagnostics]);
+		if (!message) return;
+		return { message };
+	});
+
+	pi.on("agent_end", async () => {
+		persistVariablesIfDirty(pi, state, state.currentVariableStore);
+		state.currentSystemPromptOptions = undefined;
+		state.currentLatestUserMessage = undefined;
+		state.currentVariableStore = undefined;
+		state.contextRewritePending = false;
+	});
+}
+
+function restoreBranchScopedRuntime(ctx: ExtensionContext, state: PiForgeRuntimeState, deps: LifecycleDeps): void {
+	state.sessionVariables = getRestoredVariables(ctx);
+	state.currentVariableStore = undefined;
+	const restoredActiveId = getRestoredActiveId(ctx);
+	state.lastPersistedActiveId = restoredActiveId;
+	deps.reloadStacks(ctx, restoredActiveId);
+}
+
+function getCurrentBranchEntries(ctx: ExtensionContext): unknown[] {
+	const leafId = ctx.sessionManager.getLeafId();
+	if (leafId === null) return [];
+	const sessionManager = ctx.sessionManager as {
+		getBranch?: (fromId?: string) => unknown[];
+		getEntries: () => unknown[];
+	};
+	return sessionManager.getBranch ? sessionManager.getBranch(leafId ?? undefined) : sessionManager.getEntries();
+}
+
+function getRestoredActiveId(ctx: ExtensionContext): string | undefined {
+	const entries = getCurrentBranchEntries(ctx);
+	for (let i = entries.length - 1; i >= 0; i--) {
+		const entry = entries[i] as { type?: string; customType?: string; data?: { activeStackId?: unknown } };
+		if (entry.type === "custom" && entry.customType === STATE_ENTRY_TYPE) {
+			return typeof entry.data?.activeStackId === "string" ? entry.data.activeStackId : undefined;
+		}
+	}
+	return undefined;
+}
+
+function getRestoredVariables(ctx: ExtensionContext): Record<string, PromptVariableValue> {
+	const entries = getCurrentBranchEntries(ctx);
+	for (let i = entries.length - 1; i >= 0; i--) {
+		const entry = entries[i] as { type?: string; customType?: string; data?: { variables?: unknown } };
+		if (entry.type !== "custom" || entry.customType !== VARIABLE_ENTRY_TYPE) continue;
+		if (!entry.data || typeof entry.data.variables !== "object" || Array.isArray(entry.data.variables)) return {};
+		return normalizeVariableRecord(entry.data.variables as Record<string, unknown>);
+	}
+	return {};
+}
+
+function persistVariablesIfDirty(pi: ExtensionAPI, state: PiForgeRuntimeState, store: PromptVariableStore | undefined): void {
+	if (!store?.sessionDirty) return;
+	state.sessionVariables = { ...store.session };
+	pi.appendEntry(VARIABLE_ENTRY_TYPE, { variables: state.sessionVariables });
+	markSessionVariablesClean(store);
+}
+
+function normalizeVariableRecord(value: Record<string, unknown>): Record<string, PromptVariableValue> {
+	const result: Record<string, PromptVariableValue> = {};
+	for (const [key, raw] of Object.entries(value)) {
+		if (isPromptVariableValue(raw)) result[key] = raw;
+	}
+	return result;
+}
+
+function isPromptVariableValue(value: unknown): value is PromptVariableValue {
+	if (value === null) return true;
+	const type = typeof value;
+	if (type === "string" || type === "boolean") return true;
+	if (type === "number") return Number.isFinite(value);
+	if (Array.isArray(value)) return value.every(isPromptVariableValue);
+	if (!value || typeof value !== "object") return false;
+	return Object.values(value as Record<string, unknown>).every(isPromptVariableValue);
+}

package/src/loader.ts

@@ -1,38 +1,51 @@
 import { existsSync, readdirSync, readFileSync } from "node:fs";
 import { basename, join } from "node:path";
+import { validateRegexConfig } from "./regex.ts";
+import { promptStackReadDirs } from "./storage.ts";
 import type {
 	LoadedPromptStack,
+	PromptResourcePolicy,
 	PromptStack,
 	PromptStackDiagnostic,
 	PromptStackItem,
 	PromptStackRole,
-	PromptStateConfig,
-	PromptStateDefinition,
-	PromptStateScope,
-	PromptStateValue,
 } from "./types.ts";
 import { SUPPORTED_SLOTS } from "./types.ts";
 
 const VALID_ROLES = new Set<PromptStackRole>(["system", "user", "assistant", "custom"]);
+const VALID_CHAT_HISTORY_TOOL_MODES = new Set(["keep", "drop"]);
 
-export function promptStacksDir(cwd: string): string {
-	return join(cwd, ".pi", "prompt-stacks");
-}
+export { isInsidePromptStackStorage, legacyPromptStacksDir, promptStackPath, promptStackReadDirs, promptStacksDir } from "./storage.ts";
 
 export function loadPromptStacks(cwd: string): LoadedPromptStack[] {
-	const dir = promptStacksDir(cwd);
-	if (!existsSync(dir)) return [];
+	const files = promptStackFiles(cwd);
+	const loaded = files.map(loadPromptStackFile);
+	annotateDuplicateStackIds(loaded);
+	return loaded;
+}
 
-	let entries: string[];
-	try {
-		entries = readdirSync(dir).filter((name) => name.endsWith(".json"));
-	} catch {
-		return [];
+function promptStackFiles(cwd: string): string[] {
+	const files: string[] = [];
+	const shadowedNames = new Set<string>();
+
+	for (const dir of promptStackReadDirs(cwd)) {
+		if (!existsSync(dir)) continue;
+
+		let entries: string[];
+		try {
+			entries = readdirSync(dir).filter((name) => name.endsWith(".json"));
+		} catch {
+			continue;
+		}
+
+		for (const name of entries.sort()) {
+			if (shadowedNames.has(name)) continue;
+			shadowedNames.add(name);
+			files.push(join(dir, name));
+		}
 	}
 
-	const loaded = entries.sort().map((name) => loadPromptStackFile(join(dir, name)));
-	annotateDuplicateStackIds(loaded);
-	return loaded;
+	return files;
 }
 
 function annotateDuplicateStackIds(stacks: LoadedPromptStack[]): void {
@@ -129,7 +142,7 @@ function normalizeStack(raw: unknown, filePath: string, diagnostics: PromptStack
 
 	const obj = raw as Record<string, unknown>;
 	const id = typeof obj.id === "string" && obj.id.trim() ? obj.id.trim() : basename(filePath, ".json");
-	const schemaVersion = obj.schemaVersion === 1 ? 1 : 1;
+	const schemaVersion = 1;
 	if (obj.schemaVersion !== 1) {
 		diagnostics.push({ level: "warning", message: "Missing or unsupported schemaVersion; assuming 1." });
 	}
@@ -139,6 +152,13 @@ function normalizeStack(raw: unknown, filePath: string, diagnostics: PromptStack
 		diagnostics.push({ level: "error", message: "Prompt stack must contain an items array." });
 	}
 
+	if (isPlainObject(obj.state) && Object.keys(obj.state).length > 0) {
+		diagnostics.push({
+			level: "info",
+			message: "state is no longer supported and was ignored; use stack.variables and template variable macros instead.",
+		});
+	}
+
 	return {
 		schemaVersion,
 		type: obj.type === "pi-forge.prompt-stack" ? "pi-forge.prompt-stack" : undefined,
@@ -149,8 +169,10 @@ function normalizeStack(raw: unknown, filePath: string, diagnostics: PromptStack
 		mode: obj.mode === "append" || obj.mode === "prepend" || obj.mode === "replace" ? obj.mode : undefined,
 		defaults: isPlainObject(obj.defaults) ? (obj.defaults as PromptStack["defaults"]) : undefined,
 		context: isPlainObject(obj.context) ? (obj.context as PromptStack["context"]) : undefined,
+		tools: normalizeResourcePolicy(obj.tools, "tools", diagnostics),
+		skills: normalizeResourcePolicy(obj.skills, "skills", diagnostics),
 		variables: normalizeStringRecord(obj.variables),
-		state: normalizeStateConfig(obj.state, diagnostics),
+		regex: normalizeRegexConfig(obj.regex, diagnostics),
 		items,
 		import: isPlainObject(obj.import) ? (obj.import as Record<string, unknown>) : undefined,
 	};
@@ -213,6 +235,7 @@ export function validatePromptStack(stack: PromptStack): PromptStackDiagnostic[]
 				diagnostics.push({ level: "warning", message: `Unsupported slot: ${item.slot}`, itemId: item.id });
 			}
 			if (item.enabled !== false && item.slot === "chat-history") chatHistoryCount++;
+			if (item.slot === "chat-history") validateChatHistoryOptions(item, diagnostics);
 		}
 
 		if (item.kind === "block" && item.role !== "system" && item.enabled !== false && !item.role) {
@@ -232,15 +255,64 @@ export function validatePromptStack(stack: PromptStack): PromptStackDiagnostic[]
 		});
 	}
 
+	diagnostics.push(...validateResourcePolicy(stack.tools, "tools"));
+	diagnostics.push(...validateResourcePolicy(stack.skills, "skills"));
+	if ((stack.mode === "append" || stack.mode === "prepend") && hasPolicyEntries(stack.skills)) {
+		diagnostics.push({
+			level: "warning",
+			message: "skills policy only filters pi-forge skills slots. Use mode \"replace\" if you need the base Pi prompt to omit filtered skills.",
+		});
+	}
+
+	diagnostics.push(...validateRegexConfig(stack.regex));
+
 	return diagnostics;
 }
 
+function validateChatHistoryOptions(
+	item: Extract<PromptStackItem, { kind: "slot" }>,
+	diagnostics: PromptStackDiagnostic[],
+): void {
+	const options = item.options;
+	if (!options) return;
+
+	for (const key of ["includeLastUserMessage", "stripAssistantThinking", "includeSummaries"]) {
+		const value = options[key];
+		if (value !== undefined && typeof value !== "boolean") {
+			diagnostics.push({ level: "warning", message: `chat-history option ${key} should be a boolean.`, itemId: item.id });
+		}
+	}
+
+	if (options.roles !== undefined && !isStringArray(options.roles)) {
+		diagnostics.push({ level: "error", message: "chat-history option roles must be an array of strings.", itemId: item.id });
+	}
+
+	if (options.toolMode !== undefined && (typeof options.toolMode !== "string" || !VALID_CHAT_HISTORY_TOOL_MODES.has(options.toolMode))) {
+		diagnostics.push({ level: "error", message: 'chat-history option toolMode must be "keep" or "drop".', itemId: item.id });
+	}
+
+	for (const key of ["maxMessages", "maxChars"]) {
+		const value = options[key];
+		if (value !== undefined && !isPositiveInteger(value)) {
+			diagnostics.push({ level: "error", message: `chat-history option ${key} must be a positive integer.`, itemId: item.id });
+		}
+	}
+}
+
 function normalizeId(value: unknown, fallback: string): string {
 	if (typeof value === "string" && value.trim()) return value.trim();
 	if (typeof value === "number" && Number.isFinite(value)) return String(value);
 	return fallback;
 }
 
+function isStringArray(value: unknown): value is string[] {
+	return Array.isArray(value) && value.every((item) => typeof item === "string");
+}
+
+function isPositiveInteger(value: unknown): value is number {
+	return typeof value === "number" && Number.isInteger(value) && value > 0;
+}
+
 function isPlainObject(value: unknown): value is Record<string, unknown> {
 	return !!value && typeof value === "object" && !Array.isArray(value);
 }
@@ -254,56 +326,69 @@ function normalizeStringRecord(value: unknown): Record<string, string> | undefin
 	return result;
 }
 
-function normalizeStateConfig(value: unknown, diagnostics: PromptStackDiagnostic[]): PromptStateConfig | undefined {
+function normalizeResourcePolicy(value: unknown, label: string, diagnostics: PromptStackDiagnostic[]): PromptResourcePolicy | undefined {
 	if (value === undefined) return undefined;
 	if (!isPlainObject(value)) {
-		diagnostics.push({ level: "warning", message: "state must be an object when provided." });
+		diagnostics.push({ level: "error", message: `${label} policy must be an object when provided.` });
 		return undefined;
 	}
+	const allow = normalizePolicyPatterns(value.allow, `${label}.allow`, diagnostics);
+	const deny = normalizePolicyPatterns(value.deny, `${label}.deny`, diagnostics);
+	if (allow && deny) {
+		diagnostics.push({ level: "error", message: `${label} policy must use either allow or deny, not both.` });
+		return { allow };
+	}
+	if (allow) return { allow };
+	if (deny) return { deny };
+	return {};
+}
 
-	const definitionsRaw = value.definitions;
-	const definitions: Record<string, PromptStateDefinition> = {};
-
-	if (definitionsRaw !== undefined && !isPlainObject(definitionsRaw)) {
-		diagnostics.push({ level: "warning", message: "state.definitions must be an object when provided." });
+function normalizePolicyPatterns(value: unknown, label: string, diagnostics: PromptStackDiagnostic[]): string[] | undefined {
+	if (value === undefined) return undefined;
+	if (!Array.isArray(value)) {
+		diagnostics.push({ level: "error", message: `${label} must be an array of strings when provided.` });
+		return undefined;
 	}
 
-	if (isPlainObject(definitionsRaw)) {
-		for (const [name, rawDefinition] of Object.entries(definitionsRaw)) {
-			if (!isPlainObject(rawDefinition)) {
-				diagnostics.push({ level: "warning", message: `state definition for ${name} must be an object.` });
-				continue;
-			}
+	const patterns: string[] = [];
+	for (const [index, item] of value.entries()) {
+		if (typeof item !== "string" || !item.trim()) {
+			diagnostics.push({ level: "error", message: `${label}[${index}] must be a non-empty string.` });
+			continue;
+		}
+		patterns.push(item.trim());
+	}
+	return patterns.length > 0 ? patterns : undefined;
+}
 
-			const definition: PromptStateDefinition = {};
-			if (typeof rawDefinition.type === "string") definition.type = rawDefinition.type;
-			if (isStateScope(rawDefinition.scope)) definition.scope = rawDefinition.scope;
-			if (typeof rawDefinition.description === "string") definition.description = rawDefinition.description;
-			if (typeof rawDefinition.agentWritable === "boolean") definition.agentWritable = rawDefinition.agentWritable;
-			if (typeof rawDefinition.userWritable === "boolean") definition.userWritable = rawDefinition.userWritable;
-			if (isPromptStateValue(rawDefinition.default)) definition.default = rawDefinition.default;
-			else if (rawDefinition.default !== undefined) {
-				diagnostics.push({ level: "warning", message: `state definition default for ${name} is not a JSON-compatible value.` });
+function validateResourcePolicy(policy: PromptResourcePolicy | undefined, label: string): PromptStackDiagnostic[] {
+	const diagnostics: PromptStackDiagnostic[] = [];
+	if ((policy?.allow?.length ?? 0) > 0 && (policy?.deny?.length ?? 0) > 0) {
+		diagnostics.push({ level: "error", message: `${label} policy must use either allow or deny, not both.` });
+	}
+	for (const key of ["allow", "deny"] as const) {
+		const values = policy?.[key];
+		if (!values) continue;
+		const seen = new Set<string>();
+		for (const value of values) {
+			if (seen.has(value)) {
+				diagnostics.push({ level: "warning", message: `Duplicate ${label}.${key} pattern: ${value}.` });
 			}
-			definitions[name] = definition;
+			seen.add(value);
 		}
 	}
-
-	return {
-		schemaVersion: value.schemaVersion === 1 ? 1 : undefined,
-		definitions: Object.keys(definitions).length > 0 ? definitions : undefined,
-	};
+	return diagnostics;
 }
 
-function isStateScope(value: unknown): value is PromptStateScope {
-	return value === "static" || value === "session" || value === "turn";
+function hasPolicyEntries(policy: PromptResourcePolicy | undefined): boolean {
+	return !!policy && ((policy.allow?.length ?? 0) > 0 || (policy.deny?.length ?? 0) > 0);
 }
 
-function isPromptStateValue(value: unknown): value is PromptStateValue {
-	if (value === null) return true;
-	const type = typeof value;
-	if (type === "string" || type === "number" || type === "boolean") return Number.isFinite(value as number) || type !== "number";
-	if (Array.isArray(value)) return value.every(isPromptStateValue);
-	if (!isPlainObject(value)) return false;
-	return Object.values(value).every(isPromptStateValue);
+function normalizeRegexConfig(value: unknown, diagnostics: PromptStackDiagnostic[]): PromptStack["regex"] | undefined {
+	if (value === undefined) return undefined;
+	if (!isPlainObject(value)) {
+		diagnostics.push({ level: "warning", message: "regex must be an object when provided." });
+		return undefined;
+	}
+	return value as PromptStack["regex"];
 }