feat: vendor grouped tool search

This commit is contained in:
云服务部-叶林立
2026-08-19 16:43:36 +08:00
parent 779a845e18
commit 3507f85363
22 changed files with 2236 additions and 6 deletions
+40
View File
@@ -0,0 +1,40 @@
import assert from "node:assert/strict";
import test from "node:test";
import type { ToolInfo } from "@earendil-works/pi-coding-agent";
import { BUNDLE_GROUP_DEFINITIONS, createBundleCatalog } from "../extensions/bundle-groups.ts";
import { buildManifestHash } from "../extensions/catalog.ts";
function tool(name: string): ToolInfo {
return {
name,
description: `Full description for ${name}`,
parameters: { type: "object", properties: {} } as ToolInfo["parameters"],
sourceInfo: { source: "extension", scope: "user", path: `/test/${name}.ts` } as unknown as ToolInfo["sourceInfo"],
};
}
const constraints = { maxToolsPerGroup: 8, groupOverrides: {} };
test("bundle seed names are unique and cover the declared my-pi tools", () => {
const names = BUNDLE_GROUP_DEFINITIONS.flatMap((group) => group.tools);
assert.equal(new Set(names).size, names.length);
const tools = names.map(tool);
const hash = buildManifestHash(tools, constraints);
const result = createBundleCatalog(tools, hash, constraints);
assert.deepEqual(result.unknownTools, []);
assert.equal(result.coveredNames.size, names.length);
assert.equal(result.catalog.tools.length, names.length);
assert.equal(result.catalog.tools.find((card) => card.name === "ctx_purge")?.primaryGroup, "context-administration");
assert.equal(result.catalog.tools.find((card) => card.name === "tavily_web_fetch")?.primaryGroup, "web-tavily");
});
test("user overrides take priority and unknown tools retain deterministic fallback groups", () => {
const tools = [tool("tavily_web_search"), tool("custom_analyze")];
const configured = { maxToolsPerGroup: 8, groupOverrides: { preferred: ["tavily_web_search"] } };
const hash = buildManifestHash(tools, configured);
const result = createBundleCatalog(tools, hash, configured);
assert.deepEqual(result.unknownTools.map((item) => item.name), ["custom_analyze"]);
assert.equal(result.catalog.tools.find((card) => card.name === "tavily_web_search")?.primaryGroup, "preferred");
assert.equal(result.catalog.tools.find((card) => card.name === "custom_analyze")?.primaryGroup, "custom");
});
+82
View File
@@ -0,0 +1,82 @@
import assert from "node:assert/strict";
import { mkdtemp, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import test from "node:test";
import type { ToolInfo } from "@earendil-works/pi-coding-agent";
import {
buildManifestHash,
createFallbackCatalog,
parseGeneratedCatalog,
rankGroups,
readCachedCatalog,
writeCachedCatalog,
} from "../extensions/catalog.ts";
function tool(name: string, description: string): ToolInfo {
return {
name,
description,
parameters: { type: "object", properties: { query: { type: "string" } } } as ToolInfo["parameters"],
sourceInfo: { source: "extension", scope: "user", path: `/test/${name}.ts` } as unknown as ToolInfo["sourceInfo"],
};
}
const constraints = { maxToolsPerGroup: 2, groupOverrides: {} };
test("fallback grouping chunks large prefixes and applies explicit overrides", () => {
const tools = [tool("ctx_one", "First"), tool("ctx_two", "Second"), tool("ctx_three", "Third")];
const configured = { maxToolsPerGroup: 2, groupOverrides: { preferred: ["ctx_three"] } };
const hash = buildManifestHash(tools, configured);
const catalog = createFallbackCatalog(tools, hash, configured);
assert.deepEqual(catalog.groups.find((group) => group.id === "preferred")?.tools, ["ctx_three"]);
assert.ok(catalog.groups.every((group) => group.tools.length <= 2));
assert.equal(catalog.tools.find((card) => card.name === "ctx_three")?.primaryGroup, "preferred");
});
test("validates exact generated assignments and ranks generated metadata", () => {
const tools = [tool("web_search", "Find sources"), tool("web_fetch", "Read a source")];
const hash = buildManifestHash(tools, constraints);
const generated = {
groups: [
{
id: "web-research",
title: "Web research",
summary: "Find and read current sources.",
useWhen: ["需要网页搜索"],
avoidWhen: [],
tools: ["web_search", "web_fetch"],
},
],
tools: [
{ name: "web_search", summary: "Find sources", useWhen: [], avoidWhen: [], keywords: ["搜索"], primaryGroup: "web-research" },
{ name: "web_fetch", summary: "Read sources", useWhen: [], avoidWhen: [], keywords: ["抓取"], primaryGroup: "web-research" },
],
};
const catalog = parseGeneratedCatalog(JSON.stringify(generated), tools, hash, constraints, "test/model");
assert.equal(rankGroups(catalog, "帮我搜索网页")[0]?.group.id, "web-research");
assert.throws(
() => parseGeneratedCatalog(JSON.stringify({ ...generated, tools: generated.tools.slice(0, 1) }), tools, hash, constraints, "test/model"),
/omitted/,
);
});
test("writes a private cache and rejects a stale manifest hash", async () => {
const directory = await mkdtemp(join(tmpdir(), "tool-search-catalog-"));
try {
const tools = [tool("web_search", "Find sources")];
const hash = buildManifestHash(tools, constraints);
const generated = {
groups: [{ id: "web", title: "Web", summary: "Find sources", useWhen: [], avoidWhen: [], tools: ["web_search"] }],
tools: [{ name: "web_search", summary: "Find sources", useWhen: [], avoidWhen: [], keywords: [], primaryGroup: "web" }],
};
const catalog = parseGeneratedCatalog(JSON.stringify(generated), tools, hash, constraints, "test/model");
const path = join(directory, "nested", "catalog.json");
writeCachedCatalog(path, catalog);
assert.equal(readCachedCatalog(path, tools, hash, constraints)?.groups[0]?.id, "web");
assert.equal(readCachedCatalog(path, tools, "stale", constraints), undefined);
} finally {
await rm(directory, { recursive: true, force: true });
}
});
+105
View File
@@ -0,0 +1,105 @@
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 { BUNDLE_TOOL_SEARCH_DEFAULTS, ensureToolSearchDefaults, readToolSearchConfig } from "../extensions/config.ts";
async function withAgentDir(run: (agentDir: string) => Promise<void>): Promise<void> {
const agentDir = await mkdtemp(join(tmpdir(), "my-pi-tool-search-"));
try {
await run(agentDir);
} finally {
await rm(agentDir, { recursive: true, force: true });
}
}
test("writes bundle defaults when toolSearch is absent", async () => {
await withAgentDir(async (agentDir) => {
assert.equal(ensureToolSearchDefaults(agentDir), "updated");
assert.deepEqual(JSON.parse(await readFile(join(agentDir, "settings.json"), "utf8")), {
toolSearch: {
alwaysEnabled: [...BUNDLE_TOOL_SEARCH_DEFAULTS.alwaysEnabled],
showToolSearchFooterStatus: false,
maxActiveGroups: 3,
maxToolsPerGroup: 8,
maxDynamicTools: 20,
groupOverrides: {},
},
});
});
});
test("fills missing defaults and preserves explicit settings", async () => {
await withAgentDir(async (agentDir) => {
const path = join(agentDir, "settings.json");
await writeFile(path, JSON.stringify({ theme: "dark", toolSearch: { alwaysEnabled: ["multi_grep"], maxActiveGroups: 2 } }), "utf8");
assert.equal(ensureToolSearchDefaults(agentDir), "updated");
assert.deepEqual(JSON.parse(await readFile(path, "utf8")), {
theme: "dark",
toolSearch: {
alwaysEnabled: ["multi_grep"],
maxActiveGroups: 2,
showToolSearchFooterStatus: false,
maxToolsPerGroup: 8,
maxDynamicTools: 20,
groupOverrides: {},
},
});
});
});
test("normalizes invalid runtime values without overwriting the file", async () => {
await withAgentDir(async (agentDir) => {
const path = join(agentDir, "settings.json");
const original = JSON.stringify({
toolSearch: {
alwaysEnabled: ["one", "one", 3],
showToolSearchFooterStatus: "no",
maxActiveGroups: 0,
maxToolsPerGroup: 4,
maxDynamicTools: -1,
groupOverrides: { web: ["search", 2], empty: [] },
},
});
await writeFile(path, original, "utf8");
assert.deepEqual(readToolSearchConfig(agentDir), {
alwaysEnabled: ["one"],
showToolSearchFooterStatus: false,
maxActiveGroups: 3,
maxToolsPerGroup: 4,
maxDynamicTools: 20,
groupOverrides: { web: ["search"] },
});
assert.equal(await readFile(path, "utf8"), original);
});
});
test("preserves explicit complete configuration", async () => {
await withAgentDir(async (agentDir) => {
const path = join(agentDir, "settings.json");
const original = `${JSON.stringify({
toolSearch: {
alwaysEnabled: [],
showToolSearchFooterStatus: true,
maxActiveGroups: 1,
maxToolsPerGroup: 2,
maxDynamicTools: 2,
groupOverrides: { custom: ["one"] },
},
}, null, 2)}\n`;
await writeFile(path, original, "utf8");
assert.equal(ensureToolSearchDefaults(agentDir), "unchanged");
assert.equal(await readFile(path, "utf8"), original);
});
});
test("does not overwrite malformed settings", async () => {
await withAgentDir(async (agentDir) => {
const path = join(agentDir, "settings.json");
await writeFile(path, "{not-json", "utf8");
assert.equal(ensureToolSearchDefaults(agentDir), "skipped-invalid");
assert.equal(await readFile(path, "utf8"), "{not-json");
});
});
+251
View File
@@ -0,0 +1,251 @@
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 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<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 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<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>();
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 });
}
});