my-pi/extensions/prompt-guard.ts

import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

import { loadToolCatalog, type ToolCatalog } from "./system-prompt-config.ts";

type AnyMessage = { role?: string; content?: unknown } & Record<string, unknown>;

type AnyHandler = (event: never, ctx: never) => unknown;

export interface PromptGuardOptions {
	/**
	 * Post-process the messages a wrapped `context` handler leaves behind.
	 * Receives the message list snapshot from before the handler ran and the
	 * list after; returns the list the handler's result should report.
	 */
	contextFilter?: (before: AnyMessage[], after: AnyMessage[]) => AnyMessage[];
}

/**
 * Proxy an `ExtensionAPI` for third-party extensions so they cannot bypass the
 * bundle's single system-prompt owner (`extensions/system-prompt.ts`):
 *
 * - `registerTool` applies the tool's `description` override from
 *   `config/prompt/tools.jsonc` when the catalog carries one (lazily loaded;
 *   a catalog error keeps the upstream description).
 * - `before_agent_start` handlers keep every returned field except
 *   `systemPrompt`, which is dropped.
 * - `context` handlers are post-processed through `opts.contextFilter`.
 * - Everything else passes through unchanged.
 */
export function guardExtensionApi(pi: ExtensionAPI, opts: PromptGuardOptions = {}): ExtensionAPI {
	let catalogCache: ToolCatalog | null | undefined;
	const catalog = (): ToolCatalog | null => {
		if (catalogCache === undefined) {
			try {
				catalogCache = loadToolCatalog();
			} catch {
				catalogCache = null;
			}
		}
		return catalogCache;
	};

	const registerTool: ExtensionAPI["registerTool"] = (tool) => {
		const entry = catalog()?.tools[tool.name];
		if (entry !== undefined && !("inherit" in entry) && typeof entry.description === "string") {
			pi.registerTool({ ...tool, description: entry.description });
		} else {
			pi.registerTool(tool);
		}
	};

	const on: ExtensionAPI["on"] = ((event: string, handler: AnyHandler): unknown => {
		let wrapped = handler;
		if (event === "before_agent_start") {
			wrapped = (async (eventPayload: unknown, ctx: unknown) => {
				const result = await handler(eventPayload as never, ctx as never);
				if (result !== null && typeof result === "object" && "systemPrompt" in result) {
					const rest = { ...(result as Record<string, unknown>) };
					delete rest.systemPrompt;
					return rest;
				}
				return result;
			}) as AnyHandler;
		} else if (event === "context" && opts.contextFilter) {
			const filter = opts.contextFilter;
			wrapped = (async (eventPayload: { messages?: AnyMessage[] }, ctx: unknown) => {
				const messages = Array.isArray(eventPayload.messages) ? eventPayload.messages : [];
				const before = messages.slice();
				const result = await handler(eventPayload as never, ctx as never);
				const resultMessages =
					result !== null &&
					typeof result === "object" &&
					Array.isArray((result as { messages?: unknown }).messages)
						? ((result as { messages: AnyMessage[] }).messages)
						: messages;
				const filtered = filter(before, resultMessages);
				if (filtered === resultMessages) {
					return result;
				}
				return {
					...(result !== null && typeof result === "object" ? (result as Record<string, unknown>) : {}),
					messages: filtered,
				};
			}) as AnyHandler;
		}
		// `on` is overloaded per event name; drop to a loosely typed call and
		// forward whatever the real registration returns (e.g. a disposer).
		return (pi.on as (name: string, h: AnyHandler) => unknown)(event, wrapped);
	}) as ExtensionAPI["on"];

	return new Proxy(pi, {
		get(target, prop, receiver) {
			if (prop === "registerTool") return registerTool;
			if (prop === "on") return on;
			const value = Reflect.get(target, prop, receiver);
			return typeof value === "function" ? value.bind(target) : value;
		},
	});
}