mirror of
https://bitbucket.org/siakitem/my-pi.git
synced 2026-08-28 08:35:57 +00:00
383 lines
15 KiB
TypeScript
383 lines
15 KiB
TypeScript
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;
|
|
loadedGroups: string[];
|
|
evictedGroups: string[];
|
|
activeGroups: string[];
|
|
catalogSource: string;
|
|
};
|
|
usage?: unknown;
|
|
}
|
|
|
|
interface RegisteredTool {
|
|
execute(
|
|
id: string,
|
|
params: { group?: string; groups?: string[]; query?: string },
|
|
signal: AbortSignal | undefined,
|
|
onUpdate: undefined,
|
|
ctx: ExtensionContext,
|
|
): Promise<ToolSearchResult>;
|
|
}
|
|
|
|
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> = {}): 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<string, (...args: any[]) => unknown>();
|
|
const registered = new Map<string, RegisteredTool>();
|
|
const activeCalls: string[][] = [];
|
|
const commands = new Map<string, (...args: any[]) => unknown>();
|
|
const notifications: string[] = [];
|
|
let failActivation = false;
|
|
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: (name: string, definition: { handler: (...args: any[]) => unknown }) => commands.set(name, definition.handler),
|
|
setActiveTools: (names: string[]) => {
|
|
if (failActivation) throw new Error("simulated setActiveTools failure");
|
|
activeCalls.push([...names]);
|
|
},
|
|
on: (event: string, handler: (...args: any[]) => unknown) => handlers.set(event, handler),
|
|
} as unknown as ExtensionAPI;
|
|
|
|
toolSearchExtension(api);
|
|
const ctx = testContext({
|
|
hasUI: true,
|
|
ui: {
|
|
setStatus: () => {},
|
|
notify: (message: string) => notifications.push(message),
|
|
} as unknown as ExtensionContext["ui"],
|
|
});
|
|
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"));
|
|
|
|
const batch = await tool?.execute("call-4", { groups: ["beta", "gamma"] }, undefined, undefined, ctx);
|
|
assert.deepEqual(batch?.details.loadedGroups, ["beta", "gamma"]);
|
|
assert.deepEqual(batch?.details.evictedGroups, ["alpha"]);
|
|
assert.deepEqual(batch?.details.activeGroups, ["gamma", "beta"]);
|
|
|
|
const rejected = await tool?.execute(
|
|
"call-5",
|
|
{ groups: ["alpha", "beta", "gamma"] },
|
|
undefined,
|
|
undefined,
|
|
ctx,
|
|
);
|
|
assert.deepEqual(rejected?.details.loadedGroups, []);
|
|
assert.deepEqual(rejected?.details.activeGroups, ["gamma", "beta"]);
|
|
|
|
await commands.get("tool-search-status")?.("", ctx);
|
|
assert.match(notifications.at(-1) ?? "", /Dynamic groups: 2\/2/);
|
|
assert.match(notifications.at(-1) ?? "", /LRU oldest → newest: beta → gamma/);
|
|
assert.match(notifications.at(-1) ?? "", /Last eviction: alpha \(group cap 2\)/);
|
|
|
|
failActivation = true;
|
|
await assert.rejects(
|
|
tool!.execute("call-6", { groups: ["alpha", "gamma"] }, undefined, undefined, ctx),
|
|
/simulated setActiveTools failure/,
|
|
);
|
|
failActivation = false;
|
|
await commands.get("tool-search-status")?.("", ctx);
|
|
assert.match(notifications.at(-1) ?? "", /LRU oldest → newest: beta → gamma/);
|
|
assert.match(notifications.at(-1) ?? "", /Last eviction: alpha \(group cap 2\)/);
|
|
} 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<string, (...args: any[]) => unknown>();
|
|
const registered = new Map<string, RegisteredTool>();
|
|
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<string, (...args: any[]) => unknown>();
|
|
const registered = new Map<string, RegisteredTool>();
|
|
const activeCalls: string[][] = [];
|
|
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("ctx_execute", "Run commands in a context sandbox"),
|
|
sourceTool("ctx_execute_file", "Analyze a file in a context sandbox"),
|
|
sourceTool("ctx_batch_execute", "Run a bounded command batch"),
|
|
];
|
|
const api = {
|
|
getAllTools: () => sourceTools,
|
|
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({
|
|
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);
|
|
assert.ok(["ctx_execute", "ctx_execute_file", "ctx_batch_execute"].every((name) => activeCalls.at(-1)?.includes(name)));
|
|
|
|
const ambiguous = await registered.get("tool_search")?.execute(
|
|
"call-1",
|
|
{ query: "web search" },
|
|
undefined,
|
|
undefined,
|
|
ctx,
|
|
);
|
|
assert.equal(ambiguous?.details.loadedGroup, undefined);
|
|
assert.deepEqual(ambiguous?.details.activeGroups, []);
|
|
|
|
const result = await registered.get("tool_search")?.execute(
|
|
"call-2",
|
|
{ query: "tavily_web_search" },
|
|
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<string, (...args: any[]) => unknown>();
|
|
const registered = new Map<string, RegisteredTool>();
|
|
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 });
|
|
}
|
|
});
|