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

	if (hasTool(selectedTools, "ssh_connect")) {
		rules.push(
			"- Use ssh_connect only when the user explicitly names an imported host as part of a concrete remote task. Call it as a separate step and wait for success before calling other ssh_* tools; never infer or substitute a different host.",
		);
	}

	if (hasTool(selectedTools, "ssh_cd")) {
		rules.push(
			"- Treat ssh_cd as a reviewed persistent workspace transition: call it as a separate step and wait for success before issuing ssh_bash or relative ssh_read/ssh_write/ssh_edit/ssh_find/ssh_grep calls that depend on the new remote cwd. Use cd inside ssh_bash only for an intentionally temporary, command-local directory change.",
		);
	}

	if (hasTool(selectedTools, "ssh_find") || hasTool(selectedTools, "ssh_grep")) {
		rules.push(
			"- For remote SSH searches, use ssh_find to narrow remote file paths before ssh_grep searches file contents. Keep path and limit bounded; when a result says truncated, narrow the path or pattern instead of increasing the limit. Do not run find, fd, grep, or rg through ssh_bash.",
		);
	}

	if (hasAnyTool(selectedTools, ["web_search", "web_fetch"])) {
		rules.push(
			"- Use web_search for current, real-world, or citable web information; provider selection is handled internally. Keep queries focused and use site/published_* filters only when they materially narrow results. Then read selected pages with web_fetch instead of repeating similar searches.",
		);
	}

	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, "chrome_snapshot") && hasTool(selectedTools, "ctx_execute_file")) {
		rules.push(
			"- After chrome_snapshot or a Chrome action with includeSnapshot=true, immediately analyze the returned workspace-relative capture path with ctx_execute_file; never use read or cat for the raw JSON. Print only the uids, text, or state needed for the next action, and take a fresh snapshot after any interaction because the previous capture may be stale.",
		);
	}

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