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("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("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(//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 } | 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, //); assert.match(notification, new RegExp(outputPath.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"))); } finally { await rm(cwd, { recursive: true, force: true }); } });