oh-my-pi/packages/coding-agent/test/system-prompt-templates.test.ts

import { afterEach, describe, expect, test, vi } from "bun:test";
import * as fs from "node:fs/promises";
import * as os from "node:os";
import * as path from "node:path";
import { type AgentTool, INTENT_FIELD } from "@oh-my-pi/pi-agent-core";
import { buildSystemPrompt, buildSystemPromptToolMetadata } from "@oh-my-pi/pi-coding-agent/system-prompt";
import { prompt } from "@oh-my-pi/pi-utils";
import Handlebars from "handlebars";
import * as z from "zod/v4";

const baseGitContext = {
	isRepo: true,
	currentBranch: "feature/tests",
	mainBranch: "main",
	status: "M packages/coding-agent/src/prompts/system/custom-system-prompt.md",
	commits: "abc123 Fix tests",
};

const systemPromptsDir = path.resolve(import.meta.dir, "../src/prompts/system");

const baseRenderContext: prompt.TemplateContext = {
	TASK_TOOL_NAME: "task",
	ARGUMENTS: "alpha beta",
	agent: "You are a delegated worker",
	agentsMdSearch: { files: [] },
	appendPrompt: "Appendix instructions",
	arguments: "alpha beta",
	base: "Base system prompt",
	content: "Rule content",
	context: "Background context",
	contextFile: "/tmp/context.md",
	contextFiles: [{ path: "/tmp/context/a.md", content: "Alpha context" }],
	customPrompt: "Custom prompt body",
	cwd: "/tmp/pi-issue-147",
	date: "2026-02-24",
	dateTime: "2026-02-24T12:00:00Z",
	editToolName: "edit",
	environment: [{ label: "OS", value: "Darwin" }],
	finalPlanFilePath: "local://PLAN_FINAL.md",
	git: baseGitContext,
	intentField: INTENT_FIELD,
	intentTracing: true,
	iterative: true,
	maxRetries: 3,
	modifiedFiles: ["packages/coding-agent/src/config/prompt-templates.ts"],
	name: "rs-no-unwrap",
	path: "packages/coding-agent/src/config/prompt-templates.ts",
	planContent: "1. Read code\n2. Add tests",
	planExists: true,
	planFilePath: "local://PLAN.md",
	readFiles: ["packages/coding-agent/src/prompts/system/custom-system-prompt.md"],
	repeatToolDescriptions: true,
	reentry: false,
	request: "Create an agent to review prompt templates",
	retryCount: 1,
	rules: [{ name: "rs-no-unwrap", description: "Avoid unwrap", globs: ["**/*.rs"] }],
	skills: [{ name: "system-prompts", description: "Prompt design skill" }],
	systemPromptCustomization: "System customization",
	toolInfo: [{ name: "read", label: "Read", description: "Reads files" }],
	toolRefs: {
		read: "read",
		search: "search",
		find: "find",
		edit: "edit",
		task: "task",
		web_search: "web_search",
		todo_write: "todo_write",
		inspect_image: "inspect_image",
		search_tool_bm25: "search_tool_bm25",
		lsp: "lsp",
		ast_grep: "ast_grep",
		ast_edit: "ast_edit",
		grep: "grep",
		write: "write",
	},
	tools: ["read", "search", "find", "edit", "task", "web_search", "todo_write"],
	worktree: "/tmp/pi-issue-147",
	writeToolName: "write",
};

async function loadSystemPromptTemplates(): Promise<Map<string, string>> {
	const templates = new Map<string, string>();
	const glob = new Bun.Glob("*.md");

	for await (const fileName of glob.scan({ cwd: systemPromptsDir, onlyFiles: true })) {
		const templatePath = path.join(systemPromptsDir, fileName);
		templates.set(fileName, await Bun.file(templatePath).text());
	}

	return templates;
}

function countOccurrences(text: string, needle: string): number {
	if (!needle) return 0;
	return text.split(needle).length - 1;
}

async function withTempDir(run: (dir: string) => Promise<void>): Promise<void> {
	const dir = await fs.mkdtemp(path.join(os.tmpdir(), "pi-system-prompt-"));
	try {
		await run(dir);
	} finally {
		await fs.rm(dir, { recursive: true, force: true });
	}
}

describe("system Handlebars prompt templates", () => {
	afterEach(() => {
		vi.restoreAllMocks();
	});

	test("parses and compiles every system template", async () => {
		const templates = await loadSystemPromptTemplates();
		expect(templates.size).toBeGreaterThan(0);

		for (const [fileName, template] of templates) {
			expect(() => Handlebars.parse(template), `Failed parsing ${fileName}`).not.toThrow();
			expect(() => Handlebars.compile(template), `Failed compiling ${fileName}`).not.toThrow();
		}
	});

	test("custom-system-prompt renders project section for context and git combinations", async () => {
		const templatePath = path.join(systemPromptsDir, "custom-system-prompt.md");
		const template = await Bun.file(templatePath).text();

		const both = prompt.render(template, {
			...baseRenderContext,
			contextFiles: [{ path: "a.txt", content: "A" }],
			git: { ...baseGitContext, isRepo: true },
		});
		expect(both).toContain("<project>");
		expect(both).toContain("## Context");
		expect(both).toContain("## Version Control");

		const contextOnly = prompt.render(template, {
			...baseRenderContext,
			contextFiles: [{ path: "a.txt", content: "A" }],
			git: { isRepo: false },
		});
		expect(contextOnly).toContain("<project>");
		expect(contextOnly).toContain("## Context");
		expect(contextOnly).not.toContain("## Version Control");

		const gitOnly = prompt.render(template, {
			...baseRenderContext,
			contextFiles: [],
			git: {
				isRepo: true,
				currentBranch: "feature/tests",
				mainBranch: "main",
				status: "clean",
				commits: "abc123 test commit",
			},
		});
		expect(gitOnly).toContain("<project>");
		expect(gitOnly).not.toContain("## Context");
		expect(gitOnly).toContain("## Version Control");

		const neither = prompt.render(template, {
			...baseRenderContext,
			contextFiles: [],
			git: { isRepo: false },
		});
		expect(neither).not.toContain("<project>");
		expect(neither).not.toContain("## Context");
		expect(neither).not.toContain("## Version Control");
	});

	test("subagent system owns shared context while user prompt only owns assignment", async () => {
		const systemTemplate = await Bun.file(path.join(systemPromptsDir, "subagent-system-prompt.md")).text();
		const userTemplate = await Bun.file(path.join(systemPromptsDir, "subagent-user-prompt.md")).text();

		const subagentSystem = prompt.render(systemTemplate, {
			...baseRenderContext,
			context: "Shared task background",
			agent: "You are a task agent.",
		});
		const subagentUser = prompt.render(userTemplate, {
			...baseRenderContext,
			context: "Shared task background",
			assignment: "Do the task.",
		});

		expect(subagentSystem).toContain("[CONTEXT]\nShared task background\n[/CONTEXT]");
		expect(subagentSystem).toContain("[ROLE]");
		expect(subagentUser).toContain("Complete the assignment below, thoroughly:");
		expect(subagentUser).toContain("Do the task.");
		expect(subagentUser).not.toContain("[CONTEXT]");
		expect(subagentUser).not.toContain("Shared task background");
	});
	test("system-prompt renders MCP discovery hint when enabled", async () => {
		const templatePath = path.join(systemPromptsDir, "system-prompt.md");
		const template = await Bun.file(templatePath).text();

		const rendered = prompt.render(template, {
			...baseRenderContext,
			mcpDiscoveryMode: true,
			hasMCPDiscoveryServers: true,
			mcpDiscoveryServerSummaries: ["github (2 tools)", "slack (1 tool)"],
		});

		expect(rendered).toContain("## Discovery");
		expect(rendered).toContain("Discoverable MCP servers in this session: github (2 tools), slack (1 tool).");
		expect(rendered).not.toContain("Example discoverable MCP tools:");
		expect(rendered).toContain("call `search_tool_bm25` before concluding no such tool exists");
	});

	test("buildSystemPrompt keeps system and project as separate ordered blocks with date context in project", async () => {
		await withTempDir(async dir => {
			const { systemPrompt } = await buildSystemPrompt({
				cwd: dir,
				contextFiles: [],
				skills: [],
				rules: [],
				toolNames: ["read"],
				workspaceTree: {
					rootPath: dir,
					rendered: ".\n  - src/        1m",
					truncated: false,
					totalLines: 2,
					agentsMdFiles: [],
				},
			});

			expect(systemPrompt).toHaveLength(2);
			expect(systemPrompt[0]).toContain("[CONTRACT]");
			expect(systemPrompt[0]).not.toContain("current working directory");
			expect(systemPrompt[1]).toContain("<workstation>");
			expect(systemPrompt[1]).toContain("<workspace-tree>");
			expect(systemPrompt[1]).toContain("Today is ");
			expect(systemPrompt[1]).toContain(`current working directory is '${dir}'.`);
			expect(systemPrompt[1].indexOf("</workspace-tree>")).toBeLessThan(systemPrompt[1].indexOf("Today is "));
		});
	});
	test("buildSystemPrompt renders workspace tree after directory context in project prompt", async () => {
		await withTempDir(async dir => {
			const { systemPrompt } = await buildSystemPrompt({
				cwd: dir,
				contextFiles: [],
				skills: [],
				rules: [],
				toolNames: ["read"],
				workspaceTree: {
					rootPath: dir,
					rendered: ".\n  - src/        1m",
					truncated: true,
					totalLines: 2,
					agentsMdFiles: ["packages/coding-agent/AGENTS.md"],
				},
			});

			const projectPrompt = systemPrompt[1] ?? "";

			expect(projectPrompt).toContain("<workspace-tree>");
			expect(projectPrompt).toContain("Working directory layout (sorted by mtime, recent first; depth ≤ 3):");
			expect(projectPrompt).toContain("(some entries elided to keep the tree short");
			expect(projectPrompt.indexOf("</dir-context>")).toBeLessThan(projectPrompt.indexOf("<workspace-tree>"));
		});
	});

	test("buildSystemPrompt deduplicates always-apply rules already present in SYSTEM.md", async () => {
		const duplicateRule = ["Use static imports.", "", "Do not use dynamic loading."].join("\n");
		const distinctRule = "Validate inputs at boundaries.";

		await withTempDir(async dir => {
			const configDir = path.join(dir, ".agent");
			await fs.mkdir(configDir, { recursive: true });
			await fs.writeFile(
				path.join(configDir, "SYSTEM.md"),
				["Project instructions", "", duplicateRule, "", "Trailing note"].join("\n"),
			);

			const { systemPrompt } = await buildSystemPrompt({
				cwd: dir,
				contextFiles: [],
				skills: [],
				rules: [],
				toolNames: ["read"],
				customPrompt: "Custom prompt body",
				alwaysApplyRules: [
					{ name: "no-dynamic-loading", content: duplicateRule, path: "/tmp/no-dynamic-loading.md" },
					{ name: "validate-boundaries", content: distinctRule, path: "/tmp/validate-boundaries.md" },
				],
			});

			const prompt = systemPrompt.join("\n\n");

			expect(countOccurrences(prompt, "Use static imports.")).toBe(1);
			expect(countOccurrences(prompt, "Do not use dynamic loading.")).toBe(1);
			expect(countOccurrences(prompt, distinctRule)).toBe(1);
		});
	});

	test("buildSystemPrompt deduplicates always-apply rules already present in customPrompt", async () => {
		const duplicateRule = ["Keep functions small.", "", "Extract shared helpers on the second use."].join("\n");
		const distinctRule = "Surface failures explicitly to callers.";

		const { systemPrompt } = await buildSystemPrompt({
			cwd: os.tmpdir(),
			contextFiles: [],
			skills: [],
			rules: [],
			toolNames: ["read"],
			customPrompt: ["Custom guidance", "", duplicateRule, "", "More custom guidance"].join("\n"),
			alwaysApplyRules: [
				{ name: "small-functions", content: duplicateRule, path: "/tmp/small-functions.md" },
				{ name: "truthful-failures", content: distinctRule, path: "/tmp/truthful-failures.md" },
			],
		});

		const prompt = systemPrompt.join("\n\n");

		expect(countOccurrences(prompt, "Keep functions small.")).toBe(1);
		expect(countOccurrences(prompt, "Extract shared helpers on the second use.")).toBe(1);
		expect(countOccurrences(prompt, distinctRule)).toBe(1);
	});

	test("buildSystemPromptToolMetadata captures custom wire names", () => {
		const editTool = {
			name: "edit",
			label: "Edit",
			description: "Edits files",
			parameters: z.object({}),
			customWireName: "apply_patch",
			execute: async () => ({ content: [] }),
		} satisfies AgentTool;

		const metadata = buildSystemPromptToolMetadata(new Map([["edit", editTool]]));

		expect(metadata.get("edit")?.wireName).toBe("apply_patch");
	});

	test("buildSystemPrompt references overridden tool wire names", async () => {
		const { systemPrompt } = await buildSystemPrompt({
			cwd: os.tmpdir(),
			contextFiles: [],
			skills: [],
			rules: [],
			toolNames: ["read", "search", "find", "edit", "lsp", "bash", "eval"],
			tools: new Map([
				["read", { label: "Read", description: "Reads files" }],
				["search", { label: "Search", description: "Searches files" }],
				["find", { label: "Find", description: "Finds files" }],
				["edit", { label: "Edit", description: "Edits files", wireName: "apply_patch" }],
				["lsp", { label: "LSP", description: "Queries language servers" }],
				["bash", { label: "Bash", description: "Runs shell commands" }],
				["eval", { label: "Eval", description: "Runs eval cells" }],
			]),
		});

		const promptText = systemPrompt.join("\n\n");

		expect(promptText).toContain("Edit: `apply_patch`");
		expect(promptText).toContain("surgical text edits → `apply_patch`");
		expect(promptText).not.toContain("Edit: `edit`");
	});

	test("buildSystemPrompt omits CPU info when os.cpus fails", async () => {
		vi.spyOn(os, "cpus").mockImplementation(() => {
			throw new Error("os.cpus() failed");
		});

		const { systemPrompt } = await buildSystemPrompt({
			cwd: os.tmpdir(),
			contextFiles: [],
			skills: [],
			rules: [],
			toolNames: ["read"],
		});

		const projectPrompt = systemPrompt[1] ?? "";

		const workstation = /<workstation>\n(?<content>[\s\S]*?)\n<\/workstation>/u.exec(projectPrompt)?.groups?.content;
		expect(workstation).toContain("OS:");
		expect(workstation).not.toContain("CPU:");
	});
});