my-pi/tests/tool-routing.test.ts

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

import toolRoutingExtension, {
	appendToolRouting,
	buildToolRoutingSection,
	getSystemPromptDumpPath,
} from "../extensions/tool-routing.ts";

test("search routing narrows files before requesting matching line numbers", () => {
	const section = buildToolRoutingSection(["find", "grep", "read"]);

	assert.match(section, /Use find to identify likely files or directories/);
	assert.match(section, /Use grep or multi_grep only inside the narrowed path or file/);
	assert.match(section, /obtain matching line numbers/);
	assert.match(section, /narrow the path, glob, or pattern instead of increasing the limit/);
	assert.match(section, /Use read with offset\/limit only for the exact matching region/);
});

test("SSH connection routing requires an explicit named host and concrete task", () => {
	const section = buildToolRoutingSection(["ssh_connect"]);
	assert.match(section, /only when the user explicitly names an imported host/);
	assert.match(section, /separate step and wait for success before calling other ssh_\* tools/);
	assert.match(section, /never infer or substitute a different host/);
});

test("remote cwd routing separates persistent workspace changes from shell commands", () => {
	const section = buildToolRoutingSection(["ssh_cd", "ssh_bash"]);
	assert.match(section, /call it as a separate step and wait for success/);
	assert.match(section, /relative ssh_read\/ssh_write\/ssh_edit\/ssh_find\/ssh_grep/);
	assert.match(section, /intentionally temporary, command-local directory change/);
});

test("remote search routing uses structured bounded SSH tools", () => {
	const section = buildToolRoutingSection(["ssh_find", "ssh_grep"]);
	assert.match(section, /use ssh_find to narrow remote file paths before ssh_grep/);
	assert.match(section, /Do not run find, fd, grep, or rg through ssh_bash/);
	assert.match(section, /when a result says truncated, narrow the path or pattern/);
});

test("context routing analyzes large files without a full read", () => {
	const section = buildToolRoutingSection(["ctx_execute", "ctx_execute_file"]);

	assert.match(section, /Prefer ctx_execute_file over read when deriving an answer from a large file/);
	assert.match(section, /ctx_execute_file is confined to the current project root/);
	assert.match(section, /host permission approval does not bypass this Context Mode boundary/);
});

test("Chrome snapshot routing keeps raw page data out of model context", () => {
	const section = buildToolRoutingSection(["chrome_snapshot", "ctx_execute_file"]);

	assert.match(section, /immediately analyze the returned workspace-relative capture path with ctx_execute_file/);
	assert.match(section, /never use read or cat for the raw JSON/);
	assert.match(section, /take a fresh snapshot after any interaction/);
});

test("hashline routing uses small reads and fresh anchors for edits", () => {
	const section = buildToolRoutingSection(["read", "edit"]);

	assert.match(section, /smallest useful offset\/limit range/);
	assert.match(section, /obtain fresh LINE#HASH anchors with a small read/);
	assert.match(section, /Reuse fresh anchors returned by successful edits/);
});

test("LSP routing covers navigation and post-edit diagnostics", () => {
	const section = buildToolRoutingSection(["lsp_definition", "lsp_references", "lsp_diagnostics"]);

	assert.match(section, /Use LSP for symbol definitions, references/);
	assert.match(section, /post-edit diagnostics/);
});

test("online search routing selects providers by task", () => {
	const tavily = buildToolRoutingSection(["tavily_web_search", "tavily_web_fetch"]);
	assert.match(tavily, /broad discovery when keywords are uncertain/);
	assert.match(tavily, /news\/recent\/trending topics/);
	assert.doesNotMatch(tavily, /official or API documentation/);

	const exa = buildToolRoutingSection(["exa_web_search", "exa_web_fetch", "exa_web_search_advanced"]);
	assert.match(exa, /official or API documentation/);
	assert.match(exa, /explicitly bound the result count and returned text length/);
	assert.doesNotMatch(exa, /Chinese-language pages/);

	const keenable = buildToolRoutingSection(["keenable_search", "keenable_fetch"]);
	assert.match(keenable, /Chinese-language pages, policies and government notices/);
	assert.match(keenable, /site-restricted search/);
	assert.match(keenable, /publication-date filtering/);
});

test("routing includes only guidance for active optional tools", () => {
	const section = buildToolRoutingSection(["read"]);

	assert.doesNotMatch(section, /codegraph_explore first/);
	assert.doesNotMatch(section, /For literal search/);
	assert.doesNotMatch(section, /Use Context Mode/);
	assert.doesNotMatch(section, /ctx_execute_file is confined/);
	assert.doesNotMatch(section, /fresh LINE#HASH anchors/);
	assert.doesNotMatch(section, /Choose one online search provider/);
	assert.doesNotMatch(section, /broad discovery when keywords are uncertain/);
	assert.match(section, /Use read only when exact source is needed/);
});

test("routing is appended once", () => {
	const once = appendToolRouting("base prompt", ["find", "grep"]);
	const twice = appendToolRouting(once, ["find", "grep"]);

	assert.equal(twice, once);
	assert.equal((once.match(/<!-- my-pi-tool-routing -->/g) ?? []).length, 1);
});

test("dump path stays under the current project", () => {
	assert.equal(
		getSystemPromptDumpPath("/workspace/project"),
		"/workspace/project/.pi-debug/effective-system-prompt.md",
	);
});

test("dump-system-prompt writes the last routed prompt seen by the extension", async () => {
	let beforeAgentStart: ((event: any) => { systemPrompt: string }) | undefined;
	let command: { handler: (args: string, ctx: any) => Promise<void> } | undefined;
	const pi = {
		on(name: string, handler: typeof beforeAgentStart) {
			if (name === "before_agent_start") beforeAgentStart = handler;
		},
		registerCommand(_name: string, definition: typeof command) {
			command = definition;
		},
	} as unknown as ExtensionAPI;
	toolRoutingExtension(pi);
	assert.ok(beforeAgentStart);
	assert.ok(command);
	beforeAgentStart({
		systemPrompt: "prompt already modified by an earlier extension",
		systemPromptOptions: { selectedTools: ["find", "grep"] },
	});

	const cwd = await mkdtemp(join(tmpdir(), "my-pi-tool-routing-"));
	try {
		let notification = "";
		await command.handler("", {
			cwd,
			getSystemPrompt: () => "base prompt",
			getSystemPromptOptions: () => ({ selectedTools: ["find", "grep"] }),
			ui: {
				notify(message: string) {
					notification = message;
				},
			},
		});

		const outputPath = getSystemPromptDumpPath(cwd);
		const dumped = await readFile(outputPath, "utf8");
		assert.match(dumped, /^prompt already modified by an earlier extension/);
		assert.match(dumped, /<!-- my-pi-tool-routing -->/);
		assert.match(notification, new RegExp(outputPath.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")));
	} finally {
		await rm(cwd, { recursive: true, force: true });
	}
});