my-pi/tests/system-prompt.test.ts

import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import assert from "node:assert/strict";
import { readFileSync } from "node:fs";
import { mkdtemp, readFile, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { dirname, join, resolve } from "node:path";
import test from "node:test";
import { fileURLToPath } from "node:url";

import {
	applyToolCatalog,
	loadRoutingRules,
	loadSectionsConfig,
	loadToolCatalog,
	orderSections,
	readPreamble,
	renderRouting,
	resolvePreamblePath,
	spliceCover,
	stripJsoncComments,
	type MutablePromptOptions,
	type RoutingRule,
} from "../extensions/system-prompt-config.ts";
import systemPromptExtension, {
	createSystemPromptComposer,
	getSystemPromptDumpPath,
} from "../extensions/system-prompt.ts";
import { guardExtensionApi } from "../extensions/prompt-guard.ts";
import {
	CONTEXT_MODE_ANCHOR_PREFIX,
	filterContextModeMessages,
} from "../extensions/context-mode-filter.ts";
import { BUNDLE_GROUP_DEFINITIONS } from "../pi-tool-search/extensions/bundle-groups.ts";
import { BUNDLE_TOOL_SEARCH_DEFAULTS } from "../pi-tool-search/extensions/config.ts";

const repositoryRoot = resolve(dirname(fileURLToPath(import.meta.url)), "..");
const promptConfigDir = join(repositoryRoot, "config", "prompt");

const BUNDLE_PREAMBLE =
	"You are an expert coding assistant operating inside pi, a coding agent harness. " +
	"You help users by reading files, executing commands, editing code, and writing new files.";

// ── JSONC parsing ───────────────────────────────────────────────────────────

test("stripJsoncComments removes comments but keeps // inside strings", () => {
	const source = `{
		// line comment
		"url": "https://example.com//path", /* block comment */
		"escaped": "a\\"//not a comment"
	}`;
	assert.deepEqual(JSON.parse(stripJsoncComments(source)), {
		url: "https://example.com//path",
		escaped: 'a"//not a comment',
	});
});

// ── Config files load and validate ──────────────────────────────────────────

test("config/prompt files load with the expected shapes", () => {
	const catalog = loadToolCatalog();
	assert.equal(catalog.version, 1);
	assert.ok(Object.keys(catalog.tools).length > 0);

	const routing = loadRoutingRules();
	assert.equal(routing.version, 1);
	assert.ok(routing.rules.length > 0);
	for (const rule of routing.rules) assert.equal(typeof rule.text, "string");

	const sections = loadSectionsConfig();
	assert.equal(sections.version, 1);
	assert.deepEqual(sections.order, [
		"memory_policy",
		"remote_ssh",
		"chrome_profile_bridge",
		"tool_routing",
	]);
});

test("readPreamble reads the bundle preamble from config/prompt/", () => {
	assert.equal(
		resolvePreamblePath(),
		join(promptConfigDir, "preamble.md"),
	);
	assert.equal(readPreamble(), BUNDLE_PREAMBLE);
});

// ── Catalog completeness ────────────────────────────────────────────────────

const CORE_TOOL_NAMES = [
	"read",
	"bash",
	"edit",
	"write",
	"grep",
	"find",
	"ls",
	"powershell",
	"tool_search",
];

function requiredToolNames(): Set<string> {
	const required = new Set(CORE_TOOL_NAMES);
	for (const group of BUNDLE_GROUP_DEFINITIONS) {
		for (const name of group.tools) required.add(name);
	}
	for (const name of BUNDLE_TOOL_SEARCH_DEFAULTS.alwaysEnabled) required.add(name);
	return required;
}

test("tools.jsonc covers every tool the bundle can expose, and nothing else", () => {
	const required = requiredToolNames();
	const catalog = loadToolCatalog();
	for (const name of required) {
		assert.ok(name in catalog.tools, `tool "${name}" is missing from config/prompt/tools.jsonc`);
	}
	for (const name of Object.keys(catalog.tools)) {
		assert.ok(required.has(name), `tools.jsonc entry "${name}" is not a bundle-exposed tool`);
	}
});

// ── Forbidden phrases ───────────────────────────────────────────────────────

test("prompt config carries no retired guidance phrasing", () => {
	const forbidden = [
		"ls, grep, find",
		"read, grep, or edit",
		"read/grep",
		"replace_text",
		"Read/edit files",
		"ctx_fetch_and_index then ctx_search",
	];
	for (const file of ["preamble.md", "tools.jsonc", "routing.jsonc"]) {
		const source = readFileSync(join(promptConfigDir, file), "utf8");
		for (const phrase of forbidden) {
			assert.ok(!source.includes(phrase), `config/prompt/${file} must not contain "${phrase}"`);
		}
	}
});

// ── applyToolCatalog ────────────────────────────────────────────────────────

test("applyToolCatalog applies snippet/guideline overrides and reports unregistered tools", () => {
	const catalog = {
		version: 1 as const,
		tools: {
			bash: { snippet: "custom bash line" },
			edit: { guidelines: ["rule one", "rule two"] },
			read: { inherit: true as const },
			write: { snippet: null },
		},
	};
	const options: MutablePromptOptions = {
		selectedTools: ["bash", "edit", "read", "write", "mystery_tool"],
		toolSnippets: { bash: "orig bash", write: "orig write", read: "orig read" },
		toolGuidelines: { edit: ["old rule"] },
	};

	const missing = applyToolCatalog(options, catalog);

	assert.deepEqual(missing, ["mystery_tool"]);
	assert.equal(options.toolSnippets?.bash, "custom bash line");
	assert.equal(options.toolSnippets?.read, "orig read", "inherit entries keep upstream text");
	assert.ok(!("write" in (options.toolSnippets ?? {})), "snippet: null hides the tools entry");
	assert.deepEqual(options.toolGuidelines?.edit, ["rule one", "rule two"]);
});

// ── renderRouting ───────────────────────────────────────────────────────────

test("renderRouting honours whenAll, whenAny and unconditional rules", () => {
	const rules: RoutingRule[] = [
		{ whenAll: ["a", "b"], text: "all" },
		{ whenAny: ["x", "y"], text: "any" },
		{ text: "always" },
		{ whenAll: ["a"], whenAny: ["b", "z"], text: "both" },
	];

	assert.equal(renderRouting(rules, ["a", "b", "x"]), "- all\n- any\n- always\n- both");
	assert.equal(renderRouting(rules, ["a"]), "- always", "whenAny must match at least one");
	assert.equal(renderRouting(rules, []), "- always");
	assert.equal(
		renderRouting(rules, undefined),
		"- all\n- any\n- always\n- both",
		"absent selectedTools means every tool is active",
	);
	assert.equal(renderRouting([], ["a"]), "");
});

test("bundled routing rules select the same guidance the retired tool-routing extension emitted", () => {
	const routing = loadRoutingRules();

	const literal = renderRouting(routing.rules, ["find", "grep", "read"]);
	assert.match(literal, /Use find to identify likely files or directories/);
	assert.match(literal, /narrow the path, glob, or pattern instead of increasing the limit/);

	const ssh = renderRouting(routing.rules, ["ssh_connect", "ssh_cd", "ssh_bash"]);
	assert.match(ssh, /only when the user explicitly names an imported host/);
	assert.match(ssh, /command-local directory change/);

	const chrome = renderRouting(routing.rules, ["chrome_snapshot", "ctx_execute_file"]);
	assert.match(chrome, /immediately analyze the returned workspace-relative capture path/);

	const minimal = renderRouting(routing.rules, ["read"]);
	assert.doesNotMatch(minimal, /codegraph_explore first/);
	assert.doesNotMatch(minimal, /For literal search/);
	assert.doesNotMatch(minimal, /ssh_connect only when/);
	assert.match(minimal, /Stop searching once the current evidence is sufficient/);
});

// ── orderSections ───────────────────────────────────────────────────────────

test("orderSections applies the configured order and keeps unlisted sections last", () => {
	const sections = {
		late_unlisted: "L",
		tool_routing: "R",
		memory_policy: "M",
		early_unlisted: "E",
	};
	assert.deepEqual(
		Object.keys(orderSections(sections, ["memory_policy", "tool_routing"])),
		["memory_policy", "tool_routing", "late_unlisted", "early_unlisted"],
	);
	assert.deepEqual(orderSections({}, ["memory_policy"]), {});
});

// ── spliceCover (moved from the retired cover extension) ────────────────────

function renderedPrompt(preamble = "PI DEFAULT PREAMBLE") {
	return [
		preamble,
		"<tools>",
		"- read: reads files",
		"</tools>",
		"<rules>",
		"- Be concise",
		"</rules>",
		"<project_context>",
		'<project_instructions path="/work/AGENTS.md">',
		"Nested content with a </project_context>-looking line is not a tag.",
		"</project_instructions>",
		"</project_context>",
		"<cwd>",
		"/work",
		"</cwd>",
	].join("\n");
}

function tagsOf(text: string): string[] {
	return [...text.matchAll(/^<([a-z][a-z0-9_-]*)>$/gm)].map((match) => match[1]);
}

test("spliceCover replaces only the preamble and preserves the remaining bytes", () => {
	const rendered = renderedPrompt();
	const spliced = spliceCover(rendered, "MY CUSTOM COVER");

	assert.ok(spliced.startsWith("MY CUSTOM COVER\n"));
	const originalTail = rendered.slice(rendered.indexOf("\n<tools>"));
	assert.ok(spliced.endsWith(originalTail));
	assert.deepEqual(tagsOf(spliced), tagsOf(rendered));
	assert.deepEqual(tagsOf(spliced), ["tools", "rules", "project_context", "cwd"]);
});

test("spliceCover tolerates a trailing newline and a whitespace-only preamble", () => {
	const rendered = renderedPrompt();
	assert.equal(spliceCover(rendered, "COVER\n\n"), spliceCover(rendered, "COVER"));
	const blank = spliceCover(rendered, "  \n");
	assert.equal(blank, rendered.slice(rendered.indexOf("\n<tools>")));
});

test("spliceCover degenerates to the preamble when the prompt has no sections", () => {
	assert.equal(spliceCover("just a preamble", "NEW"), "NEW");
	assert.equal(spliceCover("", "NEW"), "NEW");
});

// ── Composer: modern structured path ────────────────────────────────────────

interface FakePromptOptions extends MutablePromptOptions {
	forceSystemPrompt?: string;
}

/** Simplified re-render of Pi's prompt from the mutable options object. */
function renderPiPrompt(options: FakePromptOptions): string {
	if (options.forceSystemPrompt !== undefined) return options.forceSystemPrompt;
	const parts = ["PI DEFAULT PREAMBLE"];
	const snippets = Object.entries(options.toolSnippets ?? {});
	if (snippets.length > 0) {
		parts.push(`<tools>\n${snippets.map(([n, s]) => `- ${n}: ${s}`).join("\n")}\n</tools>`);
	}
	const rules = [
		...(options.promptGuidelines ?? []),
		...Object.entries(options.toolGuidelines ?? {}).flatMap(([name, lines]) =>
			lines.map((line) => `${name}: ${line}`),
		),
	];
	if (rules.length > 0) parts.push(`<rules>\n${rules.join("\n")}\n</rules>`);
	parts.push("<cwd>\n/work\n</cwd>");
	for (const [name, text] of Object.entries(options.sections ?? {})) {
		parts.push(`<${name}>\n${text}\n</${name}>`);
	}
	return parts.join("\n");
}

function makeEvent(options: FakePromptOptions | undefined) {
	return {
		systemPromptOptions: options,
		get systemPrompt() {
			return options ? renderPiPrompt(options) : "PI DEFAULT PREAMBLE\n<tools>\n- x\n</tools>";
		},
	};
}

function makeUiCtx() {
	const notices: string[] = [];
	const ctx = {
		hasUI: true,
		ui: { notify: (message: string) => notices.push(message) },
	};
	return { notices, ctx: ctx as never };
}

test("composer splices the preamble, applies the catalog, and orders sections", () => {
	const composer = createSystemPromptComposer();
	const options: FakePromptOptions = {
		selectedTools: ["read", "bash", "edit", "codegraph_explore", "tool_search"],
		toolSnippets: {
			read: "reads files",
			bash: "runs commands",
			edit: "edits",
			codegraph_explore: "explores",
			tool_search: "searches",
		},
		toolGuidelines: { edit: ["upstream edit rule"] },
		sections: { memory_policy: "policy text" },
	};
	const { notices, ctx } = makeUiCtx();

	const result = composer.handleBeforeAgentStart(makeEvent(options) as never, ctx);
	const output = result?.systemPrompt;
	assert.ok(output);

	assert.ok(
		output.startsWith(BUNDLE_PREAMBLE),
		"the bundle preamble replaces Pi's default preamble",
	);
	assert.match(output, /<tool_routing>\n- Use codegraph_explore first/);
	assert.match(output, /search files with the find\/grep tools, not bash/);
	assert.match(output, /LINE#HASH anchors/);
	assert.ok(
		!output.includes("upstream edit rule"),
		"the edit guidelines entry replaces upstream rules",
	);
	assert.ok(
		output.indexOf("<memory_policy>") < output.indexOf("<tool_routing>"),
		"sections.jsonc orders memory_policy before tool_routing",
	);
	assert.deepEqual(notices, []);
	assert.equal(composer.getLastComposed(), output);
});

test("composer degrades to text splicing when an earlier extension froze the prompt", () => {
	const composer = createSystemPromptComposer();
	const options: FakePromptOptions = {
		selectedTools: ["read", "bash"],
		forceSystemPrompt: "FROZEN PROMPT\n<frozen>\nkept\n</frozen>",
	};
	const { notices, ctx } = makeUiCtx();
	const event = makeEvent(options);

	const first = composer.handleBeforeAgentStart(event as never, ctx);
	composer.handleBeforeAgentStart(event as never, ctx);

	assert.match(notices[0] ?? "", /单一入口规则/);
	assert.equal(notices.length, 1, "the violation warning is reported once per process");
	assert.match(first?.systemPrompt ?? "", /<tool_routing>/);
	assert.match(first?.systemPrompt ?? "", /<frozen>\nkept\n<\/frozen>/);
	assert.ok(
		(first?.systemPrompt ?? "").startsWith(BUNDLE_PREAMBLE),
		"the preamble is still spliced onto the frozen prompt",
	);
});

test("composer falls back to text appends when the host exposes no prompt options", () => {
	const composer = createSystemPromptComposer();
	const { ctx } = makeUiCtx();
	const event = { systemPrompt: "PI DEFAULT PREAMBLE\n<tools>\n- x\n</tools>" };

	const result = composer.handleBeforeAgentStart(event as never, ctx);
	const output = result?.systemPrompt ?? "";
	assert.ok(output.startsWith(BUNDLE_PREAMBLE));
	assert.match(output, /<tool_routing>/);
	assert.match(output, /<tools>\n- x\n<\/tools>/);
});

test("composer reports unregistered tools once per distinct name set", () => {
	const composer = createSystemPromptComposer();
	const { notices, ctx } = makeUiCtx();
	const options = (): FakePromptOptions => ({
		selectedTools: ["read", "ghost_tool"],
		toolSnippets: { read: "r" },
	});

	composer.handleBeforeAgentStart(makeEvent(options()) as never, ctx);
	composer.handleBeforeAgentStart(makeEvent(options()) as never, ctx);
	assert.equal(notices.length, 1);
	assert.match(notices[0], /未在 config\/prompt\/tools\.jsonc 登记：.*ghost_tool/);

	composer.handleBeforeAgentStart(
		makeEvent({ selectedTools: ["other_tool"], toolSnippets: {} }) as never,
		ctx,
	);
	assert.equal(notices.length, 2, "a different missing-tool set notifies again");
});

test("composer warns once and returns nothing when the prompt config fails to load", () => {
	const composer = createSystemPromptComposer({
		loadToolCatalog: () => {
			throw new Error("invalid jsonc");
		},
	});
	const { notices, ctx } = makeUiCtx();
	const options: FakePromptOptions = { selectedTools: ["read"], toolSnippets: { read: "r" } };

	assert.equal(composer.handleBeforeAgentStart(makeEvent(options) as never, ctx), undefined);
	assert.equal(composer.handleBeforeAgentStart(makeEvent(options) as never, ctx), undefined);
	assert.equal(notices.length, 1);
	assert.match(notices[0], /invalid jsonc/);
});

// ── prompt guard ────────────────────────────────────────────────────────────

function makeRecordingPi() {
	const tools: Array<{ name: string; description?: string }> = [];
	const handlers = new Map<string, (event: unknown, ctx: unknown) => unknown>();
	const pi = {
		registerTool(def: { name: string; description?: string }) {
			tools.push(def);
		},
		on(name: string, handler: (event: unknown, ctx: unknown) => unknown) {
			handlers.set(name, handler);
		},
		someOtherApi: () => "passthrough",
	};
	return { pi, tools, handlers };
}

test("prompt guard applies catalog description overrides at registration", () => {
	const { pi, tools } = makeRecordingPi();
	const guarded = guardExtensionApi(pi as unknown as ExtensionAPI);

	guarded.registerTool({ name: "edit", description: "upstream edit" } as never);
	guarded.registerTool({ name: "grep", description: "upstream grep" } as never);

	const override = loadToolCatalog().tools.edit;
	assert.ok(override && !("inherit" in override) && typeof override.description === "string");
	assert.equal(tools[0].description, override.description);
	assert.equal(tools[1].description, "upstream grep", "inherit entries keep the upstream description");
});

test("prompt guard strips a returned systemPrompt but keeps other fields", async () => {
	const { pi, handlers } = makeRecordingPi();
	const guarded = guardExtensionApi(pi as unknown as ExtensionAPI);

	guarded.on("before_agent_start", (async () => ({
		systemPrompt: "forced by third party",
		message: "kept",
	})) as never);

	const result = (await handlers.get("before_agent_start")?.({}, {})) as Record<string, unknown>;
	assert.deepEqual(result, { message: "kept" });
});

test("prompt guard runs context results through the provided contextFilter", async () => {
	const { pi, handlers } = makeRecordingPi();
	const guarded = guardExtensionApi(pi as unknown as ExtensionAPI, {
		contextFilter: (before, after) => after.slice(0, before.length),
	});

	guarded.on("context", (async (event: { messages: Array<{ role: string; content: string }> }) => {
		event.messages.push({ role: "user", content: "injected" });
		return { messages: event.messages };
	}) as never);

	const event = { messages: [{ role: "user", content: "real" }] };
	const result = (await handlers.get("context")?.(event, {})) as {
		messages: Array<{ content: string }>;
	};
	assert.deepEqual(result.messages.map((m) => m.content), ["real"]);
});

test("prompt guard returns the handler result untouched when the filter changes nothing", async () => {
	const { pi, handlers } = makeRecordingPi();
	const guarded = guardExtensionApi(pi as unknown as ExtensionAPI, {
		contextFilter: filterContextModeMessages,
	});
	const expected = { messages: [{ role: "user", content: "plain" }] };

	guarded.on("context", (async () => expected) as never);

	const result = await handlers.get("context")?.({ messages: expected.messages.slice() }, {});
	assert.equal(result, expected, "a no-op filter must preserve result identity");
});

test("prompt guard passes unrelated API members through to the real pi object", () => {
	const { pi } = makeRecordingPi();
	const guarded = guardExtensionApi(pi as unknown as ExtensionAPI);
	assert.equal((guarded as unknown as { someOtherApi: () => string }).someOtherApi(), "passthrough");
});

// ── context-mode filter ─────────────────────────────────────────────────────

test("context-mode filter drops only the routing-anchor part", () => {
	const before = [{ role: "user", content: "real" }];
	const anchor = `${CONTEXT_MODE_ANCHOR_PREFIX} ctx_batch_execute > ctx_execute`;
	const after = [
		...before,
		{ role: "user", content: `${anchor}\n\n<active_memory>\nx\n</active_memory>` },
	];
	const filtered = filterContextModeMessages(before, after);
	assert.equal(filtered.length, 2);
	assert.equal(filtered[1].content, "<active_memory>\nx\n</active_memory>");
});

test("context-mode filter removes the appended message when only the anchor remains", () => {
	const before = [{ role: "user", content: "real" }];
	const after = [
		...before,
		{ role: "user", content: `${CONTEXT_MODE_ANCHOR_PREFIX} only anchor` },
	];
	assert.deepEqual(filterContextModeMessages(before, after), before);
});

test("context-mode filter returns the same reference when nothing changed", () => {
	const before = [{ role: "user", content: "real" }];
	const after = [...before, { role: "user", content: "no anchor here" }];
	assert.equal(filterContextModeMessages(before, after), after);
	assert.equal(filterContextModeMessages(after, after), after);
});

test("context-mode filter ignores appended non-user or non-string messages", () => {
	const before = [{ role: "user", content: "real" }];
	const after = [
		...before,
		{ role: "assistant", content: `${CONTEXT_MODE_ANCHOR_PREFIX} not ours` },
		{ role: "user", content: [{ type: "text", text: "structured" }] },
	];
	assert.equal(filterContextModeMessages(before, after), after);
});

// ── Pinned upstream text ────────────────────────────────────────────────────

test("context-mode upstream still appends the single anchor-bearing user message", async () => {
	const packageJson = JSON.parse(
		await readFile(join(repositoryRoot, "package.json"), "utf8"),
	) as { dependencies: Record<string, string> };
	assert.equal(
		packageJson.dependencies["context-mode"],
		"1.0.169",
		"the context-filter pin requires the exact upstream version",
	);
	const upstream = await readFile(
		join(repositoryRoot, "node_modules/context-mode/build/adapters/pi/extension.js"),
		"utf8",
	);
	assert.ok(upstream.includes('parts.push("context-mode active. Hierarchy:'));
	assert.ok(upstream.includes('_pendingContext = extraParts.join("\\n\\n")'));
	assert.ok(upstream.includes("event.messages.push({"));
	assert.ok(upstream.includes('role: "user"'));
});

// ── /dump-system-prompt ─────────────────────────────────────────────────────

test("dump-system-prompt writes the last composed prompt under .pi-debug", async () => {
	let beforeAgentStart: ((event: unknown, ctx: unknown) => unknown) | undefined;
	let command: { handler: (args: string, ctx: unknown) => Promise<void> } | undefined;
	const pi = {
		on(name: string, handler: typeof beforeAgentStart) {
			if (name === "before_agent_start") beforeAgentStart = handler;
		},
		registerCommand(name: string, definition: typeof command) {
			if (name === "dump-system-prompt") command = definition;
		},
	} as unknown as ExtensionAPI;
	systemPromptExtension(pi);
	assert.ok(beforeAgentStart);
	assert.ok(command);

	const { ctx } = makeUiCtx();
	beforeAgentStart(
		makeEvent({ selectedTools: ["read"], toolSnippets: { read: "r" } }) as never,
		ctx,
	);

	const cwd = await mkdtemp(join(tmpdir(), "my-pi-system-prompt-"));
	try {
		await command.handler("", {
			cwd,
			getSystemPrompt: () => "fallback prompt",
			ui: { notify() {} },
		});
		const dumped = await readFile(getSystemPromptDumpPath(cwd), "utf8");
		assert.ok(dumped.startsWith(BUNDLE_PREAMBLE));
		assert.match(dumped, /<tool_routing>/);
	} finally {
		await rm(cwd, { recursive: true, force: true });
	}
});