my-pi/extensions/system-prompt.ts

import type { BeforeAgentStartEvent, ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { mkdir, writeFile } from "node:fs/promises";
import { dirname, resolve } from "node:path";

import {
	applyToolCatalog,
	isSubagentCustomPrompt,
	loadRoutingRules,
	loadSectionsConfig,
	loadToolCatalog,
	orderSections,
	readPreamble,
	renderRouting,
	spliceCover,
	type MutablePromptOptions,
	type RoutingConfig,
	type SectionsConfig,
	type ToolCatalog,
} from "./system-prompt-config.ts";

const DUMP_RELATIVE_PATH = ".pi-debug/effective-system-prompt.md";

export function getSystemPromptDumpPath(cwd: string): string {
	return resolve(cwd, DUMP_RELATIVE_PATH);
}

/** The options object as Pi 1.3 normalizes it; `MutablePromptOptions` omits `forceSystemPrompt`/`customPrompt`. */
type RuntimePromptOptions = MutablePromptOptions & {
	forceSystemPrompt?: string;
	customPrompt?: string;
};

/**
 * Injectable seams so the composer can be exercised under `node --test`
 * without touching the real `config/prompt/` files.
 */
export interface SystemPromptComposerDeps {
	loadToolCatalog?: () => ToolCatalog;
	loadRoutingRules?: () => RoutingConfig;
	loadSectionsConfig?: () => SectionsConfig;
	readPreamble?: () => string | undefined;
}

export interface SystemPromptComposer {
	handleBeforeAgentStart: (
		event: BeforeAgentStartEvent,
		ctx: ExtensionContext,
	) => { systemPrompt: string } | undefined;
	getLastComposed: () => string | undefined;
}

/**
 * The bundle's single system-prompt owner. It is the only extension allowed to
 * return `systemPrompt`; every other bundle or third-party extension must
 * contribute structurally through `event.systemPromptOptions`.
 *
 * Per turn it applies `config/prompt/tools.jsonc` (tool snippets/guidelines),
 * renders `config/prompt/routing.jsonc` into a `<tool_routing>` section, orders
 * contributed sections per `config/prompt/sections.jsonc`, then reads the
 * re-rendered prompt and splices in `config/prompt/preamble.md` before
 * returning the final text.
 *
 * If an earlier extension violated the contract by setting
 * `forceSystemPrompt`, the prompt is already frozen: the composer degrades to
 * text splicing (preamble + appended `<tool_routing>` when absent) and warns
 * once. Config load/validation failures warn once and leave Pi's default
 * prompt untouched — the handler never throws.
 */
export function createSystemPromptComposer(deps: SystemPromptComposerDeps = {}): SystemPromptComposer {
	const loadCatalog = deps.loadToolCatalog ?? loadToolCatalog;
	const loadRouting = deps.loadRoutingRules ?? loadRoutingRules;
	const loadSections = deps.loadSectionsConfig ?? loadSectionsConfig;
	const loadPreamble = deps.readPreamble ?? readPreamble;

	let lastComposed: string | undefined;
	let forcedNotified = false;
	let configErrorNotified = false;
	const notifiedMissingSets = new Set<string>();

	function handleBeforeAgentStart(
		event: BeforeAgentStartEvent,
		ctx: ExtensionContext,
	): { systemPrompt: string } | undefined {
		const options = event.systemPromptOptions as unknown as RuntimePromptOptions | undefined;
		const forced = options !== undefined && options.forceSystemPrompt !== undefined;
		if (forced && ctx.hasUI && !forcedNotified) {
			forcedNotified = true;
			ctx.ui.notify(
				"检测到其他扩展直接改写了系统提示词（违反单一入口规则），本轮降级为文本拼接。",
				"warning",
			);
		}

		let catalog: ToolCatalog;
		let routing: RoutingConfig;
		let sections: SectionsConfig;
		let preamble: string | undefined;
		try {
			catalog = loadCatalog();
			routing = loadRouting();
			sections = loadSections();
			preamble = loadPreamble();
		} catch (error) {
			if (ctx.hasUI && !configErrorNotified) {
				configErrorNotified = true;
				const detail = error instanceof Error ? error.message : String(error);
				ctx.ui.notify(`系统提示词配置加载失败：${detail}`, "warning");
			}
			return undefined;
		}

		// A subagent session is identified by its customPrompt carrying the
		// pi-subagents <active_agent/> tag — never by the rendered prompt, which
		// a main session may legitimately contain the literal tag inside.
		const isSubagent = isSubagentCustomPrompt(options?.customPrompt);
		const audience = isSubagent ? "subagent" : "main";

		let routingText = "";
		if (options) {
			const missing = applyToolCatalog(options, catalog);
			if (missing.length > 0 && ctx.hasUI) {
				const key = [...missing].sort().join("\n");
				if (!notifiedMissingSets.has(key)) {
					notifiedMissingSets.add(key);
					ctx.ui.notify(`以下工具未在 config/prompt/tools.jsonc 登记：${missing.join(", ")}`, "warning");
				}
			}
			routingText = renderRouting(routing.rules, options.selectedTools, audience);
			options.sections ??= {};
			if (isSubagent) {
				for (const name of sections.subagentExclude) {
					delete options.sections[name];
				}
			}
			if (routingText.length > 0) {
				options.sections.tool_routing = routingText;
			}
			options.sections = orderSections(options.sections, sections.order);
		} else {
			routingText = renderRouting(routing.rules, undefined, audience);
		}

		let rendered = event.systemPrompt;
		if (routingText.length > 0 && !rendered.includes("<tool_routing>")) {
			// Degrade path: the prompt was frozen by forceSystemPrompt, so the
			// contributed section cannot render — append it as plain text.
			rendered += `\n\n<tool_routing>\n${routingText}\n</tool_routing>`;
		}

		// A non-empty customPrompt (pi-subagents agent prompt, SYSTEM.md) owns
		// the whole preamble — including identities like <active_agent/> that
		// the permission system reads — so the bundle cover must not replace it.
		const ownsPreamble = typeof options?.customPrompt === "string" && options.customPrompt.length > 0;
		const final = preamble && !ownsPreamble ? spliceCover(rendered, preamble) : rendered;
		lastComposed = final;
		return { systemPrompt: final };
	}

	return {
		handleBeforeAgentStart,
		getLastComposed: () => lastComposed,
	};
}

export default function systemPromptExtension(pi: ExtensionAPI): void {
	const composer = createSystemPromptComposer();

	pi.on("before_agent_start", (event, ctx) => composer.handleBeforeAgentStart(event, ctx));

	pi.registerCommand("dump-system-prompt", {
		description: "Write the effective Pi system prompt to .pi-debug/effective-system-prompt.md",
		handler: async (_args, ctx) => {
			const outputPath = getSystemPromptDumpPath(ctx.cwd);
			const effectivePrompt = composer.getLastComposed() ?? ctx.getSystemPrompt();

			await mkdir(dirname(outputPath), { recursive: true });
			await writeFile(outputPath, effectivePrompt, "utf8");
			ctx.ui.notify(`System prompt written to ${outputPath} (${effectivePrompt.length} chars)`, "info");
		},
	});
}