import assert from "node:assert/strict"; import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import test from "node:test"; import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent"; import toolSearchExtension, { platformToolPolicy } from "../extensions/index.ts"; interface ToolSearchResult { details: { loadedGroup?: string; evictedGroups: string[]; activeGroups: string[]; catalogSource: string; }; usage?: unknown; } interface RegisteredTool { execute( id: string, params: { group?: string; query?: string }, signal: AbortSignal | undefined, onUpdate: undefined, ctx: ExtensionContext, ): Promise; } function sourceTool(name: string, description: string) { return { name, description, parameters: { type: "object", properties: {} }, sourceInfo: { type: "extension", path: `/test/${name}.ts` }, }; } function testContext(overrides: Partial = {}): ExtensionContext { return { ui: { setStatus: () => {} }, model: undefined, modelRegistry: { hasConfiguredAuth: () => false }, ...overrides, } as unknown as ExtensionContext; } test("loads whole groups and evicts the least-recently-used group", async () => { const agentDir = await mkdtemp(join(tmpdir(), "my-pi-tool-search-extension-")); const previousAgentDir = process.env.PI_CODING_AGENT_DIR; process.env.PI_CODING_AGENT_DIR = agentDir; try { await writeFile( join(agentDir, "settings.json"), JSON.stringify({ toolSearch: { alwaysEnabled: ["codegraph_explore", "lsp_diagnostics"], showToolSearchFooterStatus: false, maxActiveGroups: 2, maxToolsPerGroup: 4, maxDynamicTools: 8, groupOverrides: {}, }, }), ); const handlers = new Map unknown>(); const registered = new Map(); const activeCalls: string[][] = []; const sourceTools = [ sourceTool("read", "Read a file"), sourceTool("bash", "Run a command"), sourceTool("codegraph_explore", "Explore code relationships"), sourceTool("lsp_diagnostics", "Read diagnostics"), sourceTool("alpha_one", "First alpha capability with all important details"), sourceTool("alpha_two", "Second alpha capability"), sourceTool("beta_one", "Beta capability"), sourceTool("gamma_one", "Gamma capability"), ]; const api = { getAllTools: () => sourceTools, getActiveTools: () => activeCalls.at(-1) ?? sourceTools.map((tool) => tool.name), registerTool: (definition: unknown) => { const tool = definition as RegisteredTool & { name: string }; registered.set(tool.name, tool); }, registerCommand: () => {}, setActiveTools: (names: string[]) => activeCalls.push([...names]), on: (event: string, handler: (...args: any[]) => unknown) => handlers.set(event, handler), } as unknown as ExtensionAPI; toolSearchExtension(api); const ctx = testContext(); handlers.get("session_start")?.({}, ctx); assert.ok(activeCalls.at(-1)?.includes("tool_search")); assert.ok(activeCalls.at(-1)?.includes("codegraph_explore")); assert.ok(!activeCalls.at(-1)?.includes("alpha_one")); const tool = registered.get("tool_search"); const alpha = await tool?.execute("call-1", { group: "alpha" }, undefined, undefined, ctx); assert.equal(alpha?.details.loadedGroup, "alpha"); assert.ok(activeCalls.at(-1)?.includes("alpha_one")); assert.ok(activeCalls.at(-1)?.includes("alpha_two")); assert.ok(activeCalls.at(-2)?.every((name) => activeCalls.at(-1)?.includes(name))); await tool?.execute("call-2", { group: "beta" }, undefined, undefined, ctx); handlers.get("tool_execution_start")?.({ toolName: "alpha_one" }, ctx); const gamma = await tool?.execute("call-3", { group: "gamma" }, undefined, undefined, ctx); assert.deepEqual(gamma?.details.evictedGroups, ["beta"]); assert.deepEqual(gamma?.details.activeGroups, ["alpha", "gamma"]); assert.ok(activeCalls.at(-1)?.includes("alpha_one")); assert.ok(activeCalls.at(-1)?.includes("gamma_one")); assert.ok(!activeCalls.at(-1)?.includes("beta_one")); } finally { if (previousAgentDir === undefined) delete process.env.PI_CODING_AGENT_DIR; else process.env.PI_CODING_AGENT_DIR = previousAgentDir; await rm(agentDir, { recursive: true, force: true }); } }); test("builds and caches a validated catalog with the current model on first search", async () => { const agentDir = await mkdtemp(join(tmpdir(), "my-pi-tool-search-model-")); const previousAgentDir = process.env.PI_CODING_AGENT_DIR; process.env.PI_CODING_AGENT_DIR = agentDir; try { const handlers = new Map unknown>(); const registered = new Map(); const prompts: string[] = []; const sourceTools = [ sourceTool("read", "Read a file"), sourceTool("web_search", "Search the public web with filters and return sources"), sourceTool("web_fetch", "Fetch the complete content of a selected web page"), ]; const api = { getAllTools: () => sourceTools, registerTool: (definition: unknown) => { const tool = definition as RegisteredTool & { name: string }; registered.set(tool.name, tool); }, registerCommand: () => {}, setActiveTools: () => {}, on: (event: string, handler: (...args: any[]) => unknown) => handlers.set(event, handler), } as unknown as ExtensionAPI; toolSearchExtension(api); const generated = { groups: [ { id: "web-research", title: "Web research", summary: "Search and fetch current web sources.", useWhen: ["Current external facts are needed"], avoidWhen: ["The answer is entirely local"], tools: ["web_search", "web_fetch"], }, ], tools: [ { name: "web_search", summary: "Find web sources.", useWhen: ["Discover sources"], avoidWhen: [], keywords: ["search", "ๆœ็ดข"], primaryGroup: "web-research" }, { name: "web_fetch", summary: "Read a selected source.", useWhen: ["A URL is known"], avoidWhen: [], keywords: ["fetch", "ๆŠ“ๅ–"], primaryGroup: "web-research" }, ], }; const usage = { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 2, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } }; const ctx = testContext({ model: { provider: "test", id: "catalog-model" } as ExtensionContext["model"], modelRegistry: { hasConfiguredAuth: () => true, complete: async (_model: unknown, context: { messages: Array<{ content: Array<{ text: string }> }> }) => { prompts.push(context.messages[0]?.content[0]?.text ?? ""); return { content: [{ type: "text", text: JSON.stringify(generated) }], usage }; }, } as unknown as ExtensionContext["modelRegistry"], }); handlers.get("session_start")?.({}, ctx); const result = await registered.get("tool_search")?.execute( "call-1", { group: "web-research" }, undefined, undefined, ctx, ); assert.equal(result?.details.loadedGroup, "web-research"); assert.equal(result?.details.catalogSource, "model"); assert.deepEqual(result?.usage, usage); assert.match(prompts[0] ?? "", /Search the public web with filters and return sources/); const cache = JSON.parse(await readFile(join(agentDir, "tool-search", "catalog-v1.json"), "utf8")); assert.equal(cache.catalog.groups[0].id, "web-research"); } finally { if (previousAgentDir === undefined) delete process.env.PI_CODING_AGENT_DIR; else process.env.PI_CODING_AGENT_DIR = previousAgentDir; await rm(agentDir, { recursive: true, force: true }); } }); test("uses the precomputed bundle groups without calling a model", async () => { const agentDir = await mkdtemp(join(tmpdir(), "my-pi-tool-search-bundle-")); const previousAgentDir = process.env.PI_CODING_AGENT_DIR; process.env.PI_CODING_AGENT_DIR = agentDir; try { const handlers = new Map unknown>(); const registered = new Map(); let modelCalls = 0; const sourceTools = [ sourceTool("read", "Read a file"), sourceTool("tavily_web_search", "Search broadly for current web sources"), sourceTool("tavily_web_fetch", "Fetch a selected Tavily source"), ]; const api = { getAllTools: () => sourceTools, registerTool: (definition: unknown) => { const tool = definition as RegisteredTool & { name: string }; registered.set(tool.name, tool); }, registerCommand: () => {}, setActiveTools: () => {}, on: (event: string, handler: (...args: any[]) => unknown) => handlers.set(event, handler), } as unknown as ExtensionAPI; toolSearchExtension(api); const ctx = testContext({ model: { provider: "test", id: "catalog-model" } as ExtensionContext["model"], modelRegistry: { hasConfiguredAuth: () => true, complete: async () => { modelCalls += 1; throw new Error("the bundle catalog should not invoke a model"); }, } as unknown as ExtensionContext["modelRegistry"], }); handlers.get("session_start")?.({}, ctx); const result = await registered.get("tool_search")?.execute( "call-1", { group: "web-tavily" }, undefined, undefined, ctx, ); assert.equal(result?.details.loadedGroup, "web-tavily"); assert.equal(result?.details.catalogSource, "bundle"); assert.equal(modelCalls, 0); await assert.rejects(readFile(join(agentDir, "tool-search", "catalog-v1.json"), "utf8"), /ENOENT/); } finally { if (previousAgentDir === undefined) delete process.env.PI_CODING_AGENT_DIR; else process.env.PI_CODING_AGENT_DIR = previousAgentDir; await rm(agentDir, { recursive: true, force: true }); } }); test("keeps PowerShell with Bash only on Windows", () => { const definitions = [ sourceTool("bash", "Run Bash commands"), sourceTool("powershell", "Run PowerShell commands"), sourceTool("read", "Read a file"), ]; const nonWindows = platformToolPolicy(definitions, "darwin"); assert.deepEqual(nonWindows.tools.map((tool) => tool.name), ["bash", "read"]); assert.ok(nonWindows.coreTools.includes("bash")); assert.ok(!nonWindows.coreTools.includes("powershell")); const windows = platformToolPolicy(definitions, "win32"); assert.deepEqual(windows.tools.map((tool) => tool.name), ["bash", "powershell", "read"]); assert.ok(windows.coreTools.includes("bash")); assert.ok(windows.coreTools.includes("powershell")); }); test("loads an exact known group without waiting for unknown-tool model enrichment", async () => { const agentDir = await mkdtemp(join(tmpdir(), "my-pi-tool-search-fast-path-")); const previousAgentDir = process.env.PI_CODING_AGENT_DIR; process.env.PI_CODING_AGENT_DIR = agentDir; try { const handlers = new Map unknown>(); const registered = new Map(); let modelCalls = 0; const sourceTools = [ sourceTool("read", "Read a file"), sourceTool("tavily_web_search", "Search broadly for current web sources"), sourceTool("tavily_web_fetch", "Fetch a selected Tavily source"), sourceTool("unrecognized_extra", "An unrecognized third-party capability"), ]; const api = { getAllTools: () => sourceTools, registerTool: (definition: unknown) => { const tool = definition as RegisteredTool & { name: string }; registered.set(tool.name, tool); }, registerCommand: () => {}, setActiveTools: () => {}, on: (event: string, handler: (...args: any[]) => unknown) => handlers.set(event, handler), } as unknown as ExtensionAPI; toolSearchExtension(api); const ctx = testContext({ model: { provider: "test", id: "catalog-model" } as ExtensionContext["model"], modelRegistry: { hasConfiguredAuth: () => true, complete: async () => { modelCalls += 1; throw new Error("an exact known group must not wait for model enrichment"); }, } as unknown as ExtensionContext["modelRegistry"], }); handlers.get("session_start")?.({}, ctx); const result = await registered.get("tool_search")?.execute( "call-1", { group: "web-tavily" }, undefined, undefined, ctx, ); assert.equal(result?.details.loadedGroup, "web-tavily"); assert.equal(result?.details.catalogSource, "hybrid"); assert.equal(modelCalls, 0); } finally { if (previousAgentDir === undefined) delete process.env.PI_CODING_AGENT_DIR; else process.env.PI_CODING_AGENT_DIR = previousAgentDir; await rm(agentDir, { recursive: true, force: true }); } });