my-pi/extensions/tool-routing.ts

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

const ROUTING_MARKER = "<!-- my-pi-tool-routing -->";
const DUMP_RELATIVE_PATH = ".pi-debug/effective-system-prompt.md";

type SelectedTools = BuildSystemPromptOptions["selectedTools"];

function hasTool(selectedTools: SelectedTools, name: string): boolean {
	return selectedTools === undefined || selectedTools.includes(name);
}

export function buildToolRoutingSection(selectedTools: SelectedTools): string {
	const rules: string[] = [];

	if (hasTool(selectedTools, "codegraph_explore")) {
		rules.push(
			"- Use codegraph_explore first for architecture, symbol relationships, code flows, and symbols you are about to change. Treat source returned by it as already read; do not repeat the same exploration with read or search unless information is missing.",
		);
	}

	if (hasTool(selectedTools, "find") || hasTool(selectedTools, "grep") || hasTool(selectedTools, "multi_grep")) {
		rules.push(
			"- For literal search, narrow in stages instead of requesting a repository-wide dump:",
			"  1. Use find to identify likely files or directories. Constrain path/glob and keep the result limit small.",
			"  2. Use grep or multi_grep only inside the narrowed path or file to obtain matching line numbers. Start with a specific or literal pattern, a small context window, and a bounded limit.",
			"  3. Use read with offset/limit only for the exact matching region when full source is still needed.",
			"- Do not run a broad search expected to return hundreds or thousands of files or matches. If results truncate or hit a limit, narrow the path, glob, or pattern instead of increasing the limit or dumping all results.",
		);
	}

	if (hasTool(selectedTools, "ctx_execute") || hasTool(selectedTools, "ctx_execute_file")) {
		rules.push(
			"- Use Context Mode for logs, test/build output, generated data, large-file analysis, and any command whose output may be large or unpredictable. Print only the derived answer needed for the task.",
		);
	}

	if (hasTool(selectedTools, "read")) {
		rules.push(
			"- Use read directly when exact text is needed for editing or when a file is small. Do not read a complete large file merely to copy, compare, hash, count, or summarize it.",
		);
	}

	rules.push(
		"- Stop searching once the current evidence is sufficient. Avoid repeating the same lookup through CodeGraph, find, grep, and read without a concrete information gap.",
	);

	return `${ROUTING_MARKER}\n## Tool Routing\n\n${rules.join("\n")}`;
}

export function appendToolRouting(systemPrompt: string, selectedTools: SelectedTools): string {
	if (systemPrompt.includes(ROUTING_MARKER)) {
		return systemPrompt;
	}
	return `${systemPrompt}\n\n${buildToolRoutingSection(selectedTools)}`;
}

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

export default function toolRoutingExtension(pi: ExtensionAPI): void {
	let lastEffectivePrompt: string | undefined;

	pi.on("before_agent_start", (event) => {
		lastEffectivePrompt = appendToolRouting(event.systemPrompt, event.systemPromptOptions.selectedTools);
		return { systemPrompt: lastEffectivePrompt };
	});

	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 selectedTools = ctx.getSystemPromptOptions().selectedTools;
			const effectivePrompt = lastEffectivePrompt ?? appendToolRouting(ctx.getSystemPrompt(), selectedTools);

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