my-pi/extensions/system-prompt-config.ts

import { readFileSync, statSync } from "node:fs";
import { fileURLToPath } from "node:url";

/**
 * Pure helpers for the bundle's single system-prompt owner
 * (`extensions/system-prompt.ts`). No Pi runtime imports so `node --test` can
 * exercise everything here directly.
 */

/** Entry in `config/prompt/tools.jsonc`. */
export type ToolCatalogEntry =
	| { inherit: true }
	| {
			snippet?: string | null;
			guidelines?: string[];
			description?: string;
			descriptionAppend?: string;
	  };

export interface ToolCatalog {
	version: 1;
	tools: Record<string, ToolCatalogEntry>;
}

export interface RoutingRule {
	whenAll?: string[];
	whenAny?: string[];
	/** Which sessions the rule renders for; absent means both. */
	scope?: "main" | "subagent";
	text: string;
}

export interface RoutingConfig {
	version: 1;
	rules: RoutingRule[];
}

export interface SectionsConfig {
	version: 1;
	order: string[];
	/** Contributed sections removed from subagent (customPrompt) sessions. */
	subagentExclude: string[];
}

/**
 * The subset of Pi's (normalized) systemPromptOptions this module reads and
 * mutates. Declared structurally so the installed host version does not leak
 * into the config helper's types.
 */
export interface MutablePromptOptions {
	selectedTools?: string[];
	toolSnippets?: Record<string, string>;
	toolGuidelines?: Record<string, string[]>;
	promptGuidelines?: string[];
	sections?: Record<string, string>;
}

const PROMPT_CONFIG_DIR = new URL("../config/prompt/", import.meta.url);

export function resolvePreamblePath(): string {
	return fileURLToPath(new URL("preamble.md", PROMPT_CONFIG_DIR));
}

export function resolveToolsPath(): string {
	return fileURLToPath(new URL("tools.jsonc", PROMPT_CONFIG_DIR));
}

export function resolveRoutingPath(): string {
	return fileURLToPath(new URL("routing.jsonc", PROMPT_CONFIG_DIR));
}

export function resolveSectionsPath(): string {
	return fileURLToPath(new URL("sections.jsonc", PROMPT_CONFIG_DIR));
}

export function resolveSubagentPreamblePath(): string {
	return fileURLToPath(new URL("subagent/preamble.md", PROMPT_CONFIG_DIR));
}

export function resolveSubagentAgentsDir(): string {
	return fileURLToPath(new URL("subagent/agents/", PROMPT_CONFIG_DIR));
}

/**
 * Strip `//` and `/* ... *\/` comments while respecting string literals, so a
 * `//` inside a quoted value (for example `"https://…"`) is never removed.
 */
export function stripJsoncComments(source: string): string {
	let out = "";
	let inString = false;
	for (let i = 0; i < source.length; i++) {
		const ch = source[i];
		if (inString) {
			out += ch;
			if (ch === "\\") {
				if (i + 1 < source.length) out += source[i + 1];
				i += 1;
				continue;
			}
			if (ch === '"') inString = false;
			continue;
		}
		if (ch === '"') {
			inString = true;
			out += ch;
			continue;
		}
		if (ch === "/" && source[i + 1] === "/") {
			while (i < source.length && source[i] !== "\n") i++;
			continue;
		}
		if (ch === "/" && source[i + 1] === "*") {
			i += 2;
			while (i < source.length && !(source[i] === "*" && source[i + 1] === "/")) i++;
			i += 2;
			continue;
		}
		out += ch;
	}
	return out;
}

interface FileCacheEntry {
	mtimeMs: number;
	size: number;
	value: unknown;
}

const fileCache = new Map<string, FileCacheEntry>();

/**
 * Read, strip, parse and validate a JSONC config file, cached by mtime+size so
 * a `config/prompt/` edit takes effect on the next turn without re-parsing
 * unchanged files every turn.
 */
function loadJsonc<T>(path: string, validate: (raw: unknown) => T): T {
	const stat = statSync(path);
	const cached = fileCache.get(path);
	if (cached && cached.mtimeMs === stat.mtimeMs && cached.size === stat.size) {
		return cached.value as T;
	}
	const raw = JSON.parse(stripJsoncComments(readFileSync(path, "utf8")));
	const value = validate(raw);
	fileCache.set(path, { mtimeMs: stat.mtimeMs, size: stat.size, value });
	return value;
}

function fail(path: string, message: string): never {
	throw new Error(`config/prompt/${path}: ${message}`);
}

function isPlainObject(value: unknown): value is Record<string, unknown> {
	return typeof value === "object" && value !== null && !Array.isArray(value);
}

function isStringArray(value: unknown): value is string[] {
	return Array.isArray(value) && value.every((entry) => typeof entry === "string");
}

const TOOL_ENTRY_KEYS = new Set(["inherit", "snippet", "guidelines", "description", "descriptionAppend"]);

function validateToolEntry(name: string, raw: unknown): ToolCatalogEntry {
	if (!isPlainObject(raw)) {
		fail("tools.jsonc", `tool "${name}" must be an object`);
	}
	for (const key of Object.keys(raw)) {
		if (!TOOL_ENTRY_KEYS.has(key)) {
			fail("tools.jsonc", `tool "${name}" has unknown key "${key}"`);
		}
	}
	if ("inherit" in raw) {
		if (raw.inherit !== true || Object.keys(raw).length !== 1) {
			fail("tools.jsonc", `tool "${name}": "inherit" must be the only key and equal true`);
		}
		return { inherit: true };
	}
	const hasField =
		"snippet" in raw || "guidelines" in raw || "description" in raw || "descriptionAppend" in raw;
	if (!hasField) {
		fail("tools.jsonc", `tool "${name}" needs "inherit": true or at least one of snippet/guidelines/description/descriptionAppend`);
	}
	if ("snippet" in raw && raw.snippet !== null && typeof raw.snippet !== "string") {
		fail("tools.jsonc", `tool "${name}".snippet must be a string or null`);
	}
	if ("guidelines" in raw && !isStringArray(raw.guidelines)) {
		fail("tools.jsonc", `tool "${name}".guidelines must be an array of strings`);
	}
	if ("description" in raw && typeof raw.description !== "string") {
		fail("tools.jsonc", `tool "${name}".description must be a string`);
	}
	if ("descriptionAppend" in raw && typeof raw.descriptionAppend !== "string") {
		fail("tools.jsonc", `tool "${name}".descriptionAppend must be a string`);
	}
	return raw as ToolCatalogEntry;
}

export function validateToolCatalog(raw: unknown): ToolCatalog {
	if (!isPlainObject(raw) || raw.version !== 1 || !isPlainObject(raw.tools)) {
		fail("tools.jsonc", 'expected { "version": 1, "tools": { ... } }');
	}
	const tools: Record<string, ToolCatalogEntry> = {};
	for (const [name, entry] of Object.entries(raw.tools)) {
		tools[name] = validateToolEntry(name, entry);
	}
	return { version: 1, tools };
}

const ROUTING_RULE_KEYS = new Set(["whenAll", "whenAny", "scope", "text"]);

function validateRoutingRule(index: number, raw: unknown): RoutingRule {
	if (!isPlainObject(raw)) {
		fail("routing.jsonc", `rule #${index} must be an object`);
	}
	for (const key of Object.keys(raw)) {
		if (!ROUTING_RULE_KEYS.has(key)) {
			fail("routing.jsonc", `rule #${index} has unknown key "${key}"`);
		}
	}
	if (typeof raw.text !== "string") {
		fail("routing.jsonc", `rule #${index}.text must be a string`);
	}
	if ("whenAll" in raw && !isStringArray(raw.whenAll)) {
		fail("routing.jsonc", `rule #${index}.whenAll must be an array of strings`);
	}
	if ("whenAny" in raw && !isStringArray(raw.whenAny)) {
		fail("routing.jsonc", `rule #${index}.whenAny must be an array of strings`);
	}
	if ("scope" in raw && raw.scope !== "main" && raw.scope !== "subagent") {
		fail("routing.jsonc", `rule #${index}.scope must be "main" or "subagent"`);
	}
	return raw as RoutingRule;
}

export function validateRoutingConfig(raw: unknown): RoutingConfig {
	if (!isPlainObject(raw) || raw.version !== 1 || !Array.isArray(raw.rules)) {
		fail("routing.jsonc", 'expected { "version": 1, "rules": [ ... ] }');
	}
	return {
		version: 1,
		rules: raw.rules.map((rule, index) => validateRoutingRule(index, rule)),
	};
}

const SECTION_NAME_PATTERN = /^[a-z][a-z0-9_-]*$/;

export function validateSectionsConfig(raw: unknown): SectionsConfig {
	if (!isPlainObject(raw) || raw.version !== 1 || !isStringArray(raw.order)) {
		fail("sections.jsonc", 'expected { "version": 1, "order": [ ... ] }');
	}
	for (const name of raw.order) {
		if (!SECTION_NAME_PATTERN.test(name) || name === "preamble") {
			fail("sections.jsonc", `"${name}" is not a valid section name`);
		}
	}
	if ("subagentExclude" in raw && !isStringArray(raw.subagentExclude)) {
		fail("sections.jsonc", `"subagentExclude" must be an array of section names`);
	}
	for (const name of raw.subagentExclude ?? []) {
		if (!SECTION_NAME_PATTERN.test(name) || name === "preamble") {
			fail("sections.jsonc", `"subagentExclude" entry "${name}" is not a valid section name`);
		}
	}
	return {
		version: 1,
		order: raw.order,
		subagentExclude: isStringArray(raw.subagentExclude) ? raw.subagentExclude : [],
	};
}

export function loadToolCatalog(): ToolCatalog {
	return loadJsonc(resolveToolsPath(), validateToolCatalog);
}

export function loadRoutingRules(): RoutingConfig {
	return loadJsonc(resolveRoutingPath(), validateRoutingConfig);
}

export function loadSectionsConfig(): SectionsConfig {
	return loadJsonc(resolveSectionsPath(), validateSectionsConfig);
}

/**
 * Read the bundle-managed prompt preamble.
 *
 * Returns `undefined` when the file is absent, unreadable, or blank so the
 * caller leaves Pi's own default preamble untouched instead of forcing an
 * empty one.
 */
export function readPreamble(): string | undefined {
	let raw: string;
	try {
		raw = readFileSync(resolvePreamblePath(), "utf8");
	} catch {
		return undefined;
	}
	const trimmed = raw.replace(/\s+$/, "");
	return trimmed.length > 0 ? trimmed : undefined;
}

/**
 * Read the bundle's subagent base preamble (`config/prompt/subagent/preamble.md`),
 * re-read on each call so the file is never cached between spawns.
 * Blank or missing → `undefined`, which keeps pi-subagents' upstream composition.
 */
export function readSubagentPreamble(): string | undefined {
	let raw: string;
	try {
		raw = readFileSync(resolveSubagentPreamblePath(), "utf8");
	} catch {
		return undefined;
	}
	const trimmed = raw.replace(/\s+$/, "");
	return trimmed.length > 0 ? trimmed : undefined;
}

/**
 * The `<active_agent name="..."/>` tag pi-subagents stamps into every child
 * prompt — tested ONLY against `systemPromptOptions.customPrompt`, never the
 * rendered prompt: a main session may contain the literal tag via project
 * docs and must still be treated as main.
 */
const ACTIVE_AGENT_TAG = /<active_agent\s+name=["'][^"']+["'][^>]*>/i;

export function isSubagentCustomPrompt(customPrompt: unknown): boolean {
	return typeof customPrompt === "string" && ACTIVE_AGENT_TAG.test(customPrompt);
}

/**
 * Apply the tool catalog to the mutable prompt options: a `snippet` string
 * replaces the tool's `<tools>` line (`null` hides it), `guidelines` replaces
 * the tool's `<rules>` lines; `{inherit:true}` entries are left untouched.
 * `description` and `descriptionAppend` are not applied here — they are
 * registration-time concerns handled by `extensions/prompt-guard.ts`.
 *
 * Returns the selected tool names that have no catalog entry at all.
 */
export function applyToolCatalog(options: MutablePromptOptions, catalog: ToolCatalog): string[] {
	const snippets = (options.toolSnippets ??= {});
	const guidelines = (options.toolGuidelines ??= {});
	const missing: string[] = [];
	for (const name of options.selectedTools ?? []) {
		const entry = catalog.tools[name];
		if (entry === undefined) {
			missing.push(name);
			continue;
		}
		if ("inherit" in entry) continue;
		if (entry.snippet !== undefined) {
			if (entry.snippet === null) {
				delete snippets[name];
			} else {
				snippets[name] = entry.snippet;
			}
		}
		if (entry.guidelines !== undefined) {
			guidelines[name] = [...entry.guidelines];
		}
	}
	return missing;
}

/**
 * Render the matching routing rules as `- <text>` lines joined by newlines.
 * A rule applies when every `whenAll` tool is active and at least one
 * `whenAny` tool is active (a rule with neither always applies). An absent
 * `selectedTools` means every tool counts as active, matching the previous
 * tool-routing behavior.
 *
 * `audience` filters `scope`: rules without a scope render for both "main"
 * and "subagent" sessions, scoped rules only for their own audience.
 */
export function renderRouting(
	rules: RoutingRule[],
	selectedTools: string[] | undefined,
	audience: "main" | "subagent" = "main",
): string {
	const isActive = (name: string) => selectedTools === undefined || selectedTools.includes(name);
	const lines: string[] = [];
	for (const rule of rules) {
		if (rule.scope !== undefined && rule.scope !== audience) continue;
		if (rule.whenAll !== undefined && !rule.whenAll.every(isActive)) continue;
		if (rule.whenAny !== undefined && !rule.whenAny.some(isActive)) continue;
		lines.push(`- ${rule.text}`);
	}
	return lines.join("\n");
}

/**
 * Rebuild the sections map in configured order; contributed sections not
 * listed keep their original insertion order after the listed ones.
 */
export function orderSections(
	sections: Record<string, string>,
	order: readonly string[],
): Record<string, string> {
	const ordered: Record<string, string> = {};
	for (const name of order) {
		if (Object.hasOwn(sections, name)) ordered[name] = sections[name];
	}
	for (const [name, value] of Object.entries(sections)) {
		if (!Object.hasOwn(ordered, name)) ordered[name] = value;
	}
	return ordered;
}

/**
 * Matches the newline-delimited boundary before the first top-level prompt
 * section. Nested markup such as `<project_instructions path="...">` carries
 * attributes and therefore never matches a bare tag line, so the first match
 * is the end of the preamble.
 */
const FIRST_SECTION_BOUNDARY = /\n<[a-z][a-z0-9_-]*>\n/;

/**
 * Replace the rendered prompt's preamble with `preamble`, keeping everything
 * after the first top-level section byte-for-byte identical.
 *
 * Pi builds the preamble at the top of the rendered prompt, so swapping that
 * prefix preserves section order: options-based replacement (`customPrompt`)
 * would drop `tools`/`rules`/`docs` and rebuild them at the end. A rendered
 * prompt without any top-level section degenerates to the preamble alone.
 */
export function spliceCover(rendered: string, preamble: string): string {
	const trimmed = preamble.replace(/\s+$/, "");
	const boundary = rendered.search(FIRST_SECTION_BOUNDARY);
	if (boundary < 0) {
		return trimmed;
	}
	return trimmed + rendered.slice(boundary);
}