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);
}

function hasAnyTool(selectedTools: SelectedTools, names: string[]): boolean {
	return names.some((name) => hasTool(selectedTools, 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, "lsp_definition") ||
		hasTool(selectedTools, "lsp_references") ||
		hasTool(selectedTools, "lsp_diagnostics")
	) {
		rules.push(
			"- Use LSP for symbol definitions, references, hover/type information, and post-edit diagnostics when the language server supports the file.",
		);
	}

	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.",
		);
	}

	const tavilyTools = ["tavily_web_search", "tavily_web_fetch"];
	const exaTools = ["exa_web_search", "exa_web_fetch", "exa_web_search_advanced"];
	const keenableTools = ["keenable_search", "keenable_fetch"];
	if (hasAnyTool(selectedTools, [...tavilyTools, ...exaTools, ...keenableTools])) {
		rules.push(
			"- Choose one online search provider based on the task instead of querying Tavily, Exa, and Keenable by default; cross-check with another provider only when source quality or uncertainty warrants it.",
		);
	}
	if (hasAnyTool(selectedTools, tavilyTools)) {
		rules.push(
			"- Use Tavily tools for broad discovery when keywords are uncertain, for news/recent/trending topics, or when many candidate sources are useful. Tavily results can mix official pages, blogs, forums, and secondary sources, so keep max_results modest, avoid raw content unless needed, and fetch only selected pages when depth matters.",
		);
	}
	if (hasAnyTool(selectedTools, exaTools)) {
		rules.push(
			"- Use Exa tools when source precision matters: official or API documentation, release/version information, academic papers, research material, and high-quality fact confirmation. Use exa_web_fetch to continue reading selected pages; for exa_web_search_advanced, explicitly bound the result count and returned text length to prevent oversized responses.",
		);
	}
	if (hasAnyTool(selectedTools, keenableTools)) {
		rules.push(
			"- Use Keenable tools for Chinese-language pages, policies and government notices, site-restricted search, and publication-date filtering. Prefer its compact search results and fetch only selected pages to keep returned content controlled; KEENABLE_API_KEY is optional but raises rate limits.",
		);
	}

	if (hasTool(selectedTools, "ctx_execute") || hasTool(selectedTools, "ctx_execute_file")) {
		rules.push(
			"- Use Context Mode for logs, test/build output, generated data, large-file exploration/analysis/summarization, and any command whose output may be large or unpredictable. Prefer ctx_execute_file over read when deriving an answer from a large file inside the workspace, and print only the derived answer needed for the task.",
		);
	}

	if (hasTool(selectedTools, "ctx_execute_file")) {
		rules.push(
			"- ctx_execute_file is confined to the current project root. Do not call it for absolute paths outside the workspace, ../ traversal, or symlinks that resolve outside the workspace; host permission approval does not bypass this Context Mode boundary. Use another explicitly authorized tool for those files.",
		);
	}

	if (hasTool(selectedTools, "read")) {
		rules.push(
			"- Use read only when exact source is needed or the file is small. Before editing, request the smallest useful offset/limit range; do not read a complete large file merely to copy, compare, hash, count, explore, or summarize it.",
		);
	}

	if (hasTool(selectedTools, "read") && hasTool(selectedTools, "edit")) {
		rules.push(
			"- For existing-file changes, obtain fresh LINE#HASH anchors with a small read, then use anchored edit operations. Reuse fresh anchors returned by successful edits for chained changes; re-read only when anchors are stale or the next target region was not shown.",
		);
	}

	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");
		},
	});
}