Files

550 lines
23 KiB
TypeScript

/**
* pi-tool-search — model-built tool groups with bounded dynamic schema loading.
*
* Imported from https://github.com/tuansondinh/pi-tool-search and maintained
* locally from snapshot ddfb23646fd3957b791214de278e23aa393c9b13 (v0.3.6).
*/
import { randomUUID } from "node:crypto";
import { unlinkSync } from "node:fs";
import { join } from "node:path";
import { getAgentDir, type ExtensionAPI, type ExtensionContext, type ToolInfo } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import {
buildCatalogModelInput,
buildManifestHash,
createFallbackCatalog,
parseGeneratedCatalog,
rankGroups,
readCachedCatalog,
selectConfidentGroup,
writeCachedCatalog,
type CatalogConstraints,
type GroupCard,
type ToolCatalog,
} from "./catalog.ts";
import {
catalogPreservesBundleAssignments,
createBundleCatalog,
type BundleCatalogResult,
} from "./bundle-groups.ts";
import { ensureToolSearchDefaults, readToolSearchConfig } from "./config.ts";
const TOOL_SEARCH_NAME = "tool_search";
const POWERSHELL_TOOL_NAME = "powershell";
const MAX_BATCH_GROUPS = 5;
const CORE_TOOLS = ["read", "write", "edit", "bash", "grep", "find"];
type CatalogSource = "bundle" | "cache" | "fallback" | "hybrid" | "model";
type ModelUsage = Awaited<ReturnType<ExtensionContext["modelRegistry"]["complete"]>>["usage"];
type CatalogGenerationResult = { notice?: string; usage?: ModelUsage };
export function platformToolPolicy<T extends { name: string }>(
definitions: readonly T[],
platform: NodeJS.Platform,
): { tools: T[]; coreTools: string[] } {
const tools = platform === "win32"
? [...definitions]
: definitions.filter((tool) => tool.name !== POWERSHELL_TOOL_NAME);
const coreTools = platform === "win32"
? [...CORE_TOOLS, POWERSHELL_TOOL_NAME]
: [...CORE_TOOLS];
return { tools, coreTools };
}
function unique(values: Iterable<string>): string[] {
return [...new Set(values)];
}
function catalogPrompt(toolsJson: string, maxToolsPerGroup: number, fixedGroupsJson: string): string {
return [
"Build a compact retrieval catalog for the provided Pi tools.",
"The catalog is metadata only: never rename tools or invent parameters.",
"Group tools that are commonly needed together in one workflow, but keep providers or security-sensitive administration separate when that improves routing.",
`Every tool must appear exactly once and every group must contain 1-${maxToolsPerGroup} tools.`,
"Tools listed in <fixed-groups> must keep exactly those primaryGroup ids; these are curated bundle assignments.",
"Summaries must preserve capability, important boundaries, and when not to use a tool.",
"Keywords should include common English terms and useful Chinese equivalents when applicable.",
"Return JSON only with this exact shape:",
'{"groups":[{"id":"kebab-case","title":"...","summary":"...","useWhen":["..."],"avoidWhen":["..."],"tools":["exact_tool_name"]}],"tools":[{"name":"exact_tool_name","summary":"...","useWhen":["..."],"avoidWhen":["..."],"keywords":["..."],"primaryGroup":"kebab-case"}]}',
"Do not include markdown fences or explanatory text.",
"",
"<fixed-groups>",
fixedGroupsJson,
"</fixed-groups>",
"",
"<tools>",
toolsJson,
"</tools>",
].join("\n");
}
function responseText(content: Array<{ type: string; text?: string }>): string {
return content
.filter((item): item is { type: string; text: string } => item.type === "text" && typeof item.text === "string")
.map((item) => item.text)
.join("\n");
}
function groupDescription(catalog: ToolCatalog, source: CatalogSource): string {
const sourceLabel =
source === "bundle"
? "precomputed my-pi bundle catalog; no model generation required"
: source === "hybrid"
? "my-pi bundle catalog plus deterministic groups for unrecognized tools; the first search may enrich the unknown tools"
: source === "fallback"
? "deterministic fallback"
: `${source} catalog`;
const groups = catalog.groups
.map((group) => {
const when = group.useWhen.slice(0, 2).join("; ");
const detail = [group.summary, when].filter(Boolean).join(" Use when: ").slice(0, 360);
return ` ${group.id}: ${detail}`;
})
.join("\n");
return [
"Activate one or more complete tool groups for the current task. Prefer exact group ids from this catalog; use query only when no id clearly matches.",
`Catalog source: ${sourceLabel}.`,
"Available groups:",
groups || " (no hidden tool groups)",
].join("\n");
}
export default function toolSearchExtension(pi: ExtensionAPI): void {
const agentDir = getAgentDir();
const cachePath = join(agentDir, "tool-search", "catalog-v1.json");
const defaultResult = ensureToolSearchDefaults(agentDir);
if (defaultResult === "skipped-invalid") {
console.warn("my-pi: skipped pi-tool-search defaults because settings.json or toolSearch is invalid");
}
let config = readToolSearchConfig(agentDir);
let tools: ToolInfo[] = [];
let catalog: ToolCatalog = createFallbackCatalog([], "", {
maxToolsPerGroup: Math.min(config.maxToolsPerGroup, config.maxDynamicTools),
groupOverrides: {},
});
let catalogSource: CatalogSource = "fallback";
let manifestHash = "";
let policySignature = "";
let attemptedGenerationHash: string | undefined;
let clock = 0;
const activeGroups = new Map<string, number>();
const pinnedTools = new Set<string>();
let bundleState: BundleCatalogResult = { catalog, coveredNames: new Set(), unknownTools: [] };
let lastEviction = "none";
function constraints(): CatalogConstraints {
return {
maxToolsPerGroup: Math.min(config.maxToolsPerGroup, config.maxDynamicTools),
groupOverrides: config.groupOverrides,
};
}
function installBaseCatalog(useCache: boolean): void {
bundleState = createBundleCatalog(tools, manifestHash, constraints());
if (bundleState.unknownTools.length === 0) {
catalog = bundleState.catalog;
catalogSource = "bundle";
return;
}
const cached = useCache ? readCachedCatalog(cachePath, tools, manifestHash, constraints()) : undefined;
if (cached && catalogPreservesBundleAssignments(cached, bundleState.catalog, bundleState.coveredNames)) {
catalog = cached;
catalogSource = "cache";
return;
}
catalog = bundleState.catalog;
catalogSource = "hybrid";
}
function fixedBundleGroupsJson(): string {
return JSON.stringify(
bundleState.catalog.groups
.map((group) => ({
id: group.id,
tools: group.tools.filter((name) => bundleState.coveredNames.has(name)),
}))
.filter((group) => group.tools.length > 0),
);
}
function groupById(id: string): GroupCard | undefined {
return catalog.groups.find((group) => group.id === id);
}
function dynamicToolCount(groupIds: Iterable<string> = activeGroups.keys()): number {
return unique([...groupIds].flatMap((id) => groupById(id)?.tools ?? [])).length;
}
function leastRecentlyUsedGroup(
groups: Map<string, number> = activeGroups,
protectedGroups: ReadonlySet<string> = new Set(),
): string | undefined {
return [...groups]
.filter(([id]) => !protectedGroups.has(id))
.sort((left, right) => left[1] - right[1] || left[0].localeCompare(right[0]))[0]?.[0];
}
function currentStatusText(): string {
const lru = [...activeGroups]
.sort((left, right) => left[1] - right[1] || left[0].localeCompare(right[0]))
.map(([id]) => id);
return [
"Tool Search status",
`Catalog: ${catalogSource}`,
`Pinned tools: ${pinnedTools.size}`,
`Dynamic groups: ${activeGroups.size}/${config.maxActiveGroups}`,
`Dynamic tools: ${dynamicToolCount()}/${config.maxDynamicTools}`,
`LRU oldest → newest: ${lru.join(" → ") || "none"}`,
`Last eviction: ${lastEviction}`,
].join("\n");
}
function applyActiveTools(): void {
const groupTools = [...activeGroups.keys()].flatMap((id) => groupById(id)?.tools ?? []);
pi.setActiveTools(unique([TOOL_SEARCH_NAME, ...pinnedTools, ...groupTools]));
}
function updateStatus(ctx: Pick<ExtensionContext, "ui">): void {
const activeToolCount = unique([TOOL_SEARCH_NAME, ...pinnedTools, ...[...activeGroups.keys()].flatMap((id) => groupById(id)?.tools ?? [])]).length;
ctx.ui.setStatus(
"tool-search",
config.showToolSearchFooterStatus
? `${activeToolCount} / ${tools.length + pinnedTools.size + 1} tools · ${activeGroups.size} / ${config.maxActiveGroups} groups`
: undefined,
);
}
function registerToolSearch(): void {
pi.registerTool({
name: TOOL_SEARCH_NAME,
label: "Tool Search",
description: groupDescription(catalog, catalogSource),
promptSnippet: `Activate relevant tool groups on demand; up to ${config.maxActiveGroups} dynamic groups can remain active. Use groups for workflows that need several groups together`,
parameters: Type.Object({
group: Type.Optional(Type.String({ description: "Exact group id from the tool_search catalog" })),
groups: Type.Optional(
Type.Array(Type.String({ description: "Exact group id from the tool_search catalog" }), {
description: `Exact group ids to activate atomically (maximum ${MAX_BATCH_GROUPS})`,
minItems: 1,
maxItems: MAX_BATCH_GROUPS,
}),
),
query: Type.Optional(Type.String({ description: "Natural-language task used to rank groups when an exact id is unclear" })),
}),
async execute(_toolCallId, params, signal, _onUpdate, ctx) {
const singleGroup = typeof params.group === "string" ? params.group.trim() : "";
const batchGroups = Array.isArray(params.groups) ? params.groups : [];
const query = typeof params.query === "string" ? params.query.trim() : "";
const hasSingle = singleGroup.length > 0;
const hasBatch = batchGroups.length > 0;
const hasQuery = query.length > 0;
const modes = Number(hasSingle) + Number(hasBatch) + Number(hasQuery);
const lines: string[] = [];
const emptyDetails = (candidates: Array<{ id: string; score: number }> = []) => ({
loadedGroup: undefined,
loadedGroups: [] as string[],
evictedGroups: [] as string[],
activeGroups: [...activeGroups.keys()],
candidates,
catalogSource,
});
if (modes !== 1) {
lines.push("Provide exactly one of group, groups, or query.");
return { content: [{ type: "text", text: lines.join("\n") }], details: emptyDetails() };
}
if (batchGroups.length > MAX_BATCH_GROUPS) {
lines.push(`groups accepts at most ${MAX_BATCH_GROUPS} group ids.`);
return { content: [{ type: "text", text: lines.join("\n") }], details: emptyDetails() };
}
const requestedIds = unique(
(hasBatch ? batchGroups : hasSingle ? [singleGroup] : [])
.map((id) => id.trim().toLowerCase())
.filter(Boolean),
);
let selectedGroups = requestedIds.map(groupById).filter((group): group is GroupCard => group !== undefined);
const needsCatalog = hasQuery || selectedGroups.length !== requestedIds.length;
const generation: CatalogGenerationResult = needsCatalog ? await ensureModelCatalog(ctx, signal) : {};
if (generation.notice) lines.push(generation.notice);
selectedGroups = requestedIds.map(groupById).filter((group): group is GroupCard => group !== undefined);
const ranked = hasQuery ? rankGroups(catalog, query).slice(0, 3) : [];
if (hasQuery) {
const selected = selectConfidentGroup(ranked, query);
if (selected) selectedGroups = [selected];
}
const missingIds = requestedIds.filter((id) => !groupById(id));
if (selectedGroups.length === 0 || missingIds.length > 0) {
const candidates = ranked.length > 0 ? ranked : catalog.groups.slice(0, 5).map((group) => ({ group, score: 0 }));
lines.push(
missingIds.length > 0
? `Unknown groups: ${missingIds.join(", ")}`
: hasQuery
? "No group was activated because the query match was weak or ambiguous. Use an exact group id from the candidates."
: "No group was activated. Provide an exact group id from the catalog.",
`Candidates: ${candidates.map(({ group }) => group.id).join(", ") || "none"}`,
);
return {
content: [{ type: "text", text: lines.join("\n") }],
details: emptyDetails(candidates.map(({ group, score }) => ({ id: group.id, score }))),
...(generation.usage ? { usage: generation.usage } : {}),
};
}
const selectedIds = selectedGroups.map((group) => group.id);
const selectedTools = unique(selectedGroups.flatMap((group) => group.tools));
if (selectedIds.length > config.maxActiveGroups || selectedTools.length > config.maxDynamicTools) {
lines.push(
`Requested workflow needs ${selectedIds.length} groups and ${selectedTools.length} tools, exceeding the configured limits of ${config.maxActiveGroups} groups and ${config.maxDynamicTools} tools.`,
);
return {
content: [{ type: "text", text: lines.join("\n") }],
details: emptyDetails(),
...(generation.usage ? { usage: generation.usage } : {}),
};
}
const requestedSet = new Set(selectedIds);
const nextActive = new Map(activeGroups);
for (const id of selectedIds) if (!nextActive.has(id)) nextActive.set(id, Number.MAX_SAFE_INTEGER);
const evictedGroups: string[] = [];
const evictionReasons = new Set<string>();
while (nextActive.size > config.maxActiveGroups || dynamicToolCount(nextActive.keys()) > config.maxDynamicTools) {
if (nextActive.size > config.maxActiveGroups) evictionReasons.add(`group cap ${config.maxActiveGroups}`);
if (dynamicToolCount(nextActive.keys()) > config.maxDynamicTools) evictionReasons.add(`tool cap ${config.maxDynamicTools}`);
const evicted = leastRecentlyUsedGroup(nextActive, requestedSet);
if (!evicted) break;
nextActive.delete(evicted);
evictedGroups.push(evicted);
}
const previousActive = new Map(activeGroups);
const previousClock = clock;
const previousLastEviction = lastEviction;
activeGroups.clear();
for (const [id, recency] of nextActive) activeGroups.set(id, recency);
for (const id of selectedIds) activeGroups.set(id, ++clock);
if (evictedGroups.length > 0) {
lastEviction = `${evictedGroups.join(", ")} (${[...evictionReasons].join("; ")})`;
}
try {
applyActiveTools();
registerToolSearch();
updateStatus(ctx);
} catch (error) {
activeGroups.clear();
for (const [id, recency] of previousActive) activeGroups.set(id, recency);
clock = previousClock;
lastEviction = previousLastEviction;
try {
applyActiveTools();
registerToolSearch();
updateStatus(ctx);
} catch {
// Preserve the original host activation failure while retaining rolled-back internal state.
}
throw error;
}
if (selectedGroups.length === 1) {
lines.push(`Loaded group: ${selectedGroups[0]?.id} (${selectedTools.join(", ")})`);
} else {
lines.push(`Loaded groups: ${selectedIds.join(", ")} (${selectedTools.length} unique tools)`);
}
if (evictedGroups.length > 0) lines.push(`Evicted least-recently-used: ${evictedGroups.join(", ")}`);
lines.push(`Active groups: ${[...activeGroups.keys()].join(", ")}`);
return {
content: [{ type: "text", text: lines.join("\n") }],
details: {
loadedGroup: selectedIds.at(-1),
loadedGroups: selectedIds,
evictedGroups,
activeGroups: [...activeGroups.keys()],
candidates: ranked.map(({ group, score }) => ({ id: group.id, score })),
catalogSource,
},
...(generation.usage ? { usage: generation.usage } : {}),
};
},
});
}
async function ensureModelCatalog(
ctx: ExtensionContext,
signal: AbortSignal | undefined,
): Promise<{ notice?: string; usage?: ModelUsage }> {
if (catalogSource !== "hybrid" || attemptedGenerationHash === manifestHash || tools.length === 0) return {};
attemptedGenerationHash = manifestHash;
if (!ctx.model || !ctx.modelRegistry.hasConfiguredAuth(ctx.model)) {
return { notice: "Using the precomputed bundle catalog plus deterministic groups for unrecognized tools because the current model is unavailable or unauthenticated." };
}
let generationUsage: ModelUsage | undefined;
try {
const response = await ctx.modelRegistry.complete(
ctx.model,
{
messages: [
{
role: "user",
content: [
{
type: "text",
text: catalogPrompt(
buildCatalogModelInput(tools),
constraints().maxToolsPerGroup,
fixedBundleGroupsJson(),
),
},
],
timestamp: Date.now(),
},
],
},
{
signal,
reasoningEffort: "low",
cacheRetention: "none",
sessionId: randomUUID(),
maxTokens: 12_000,
},
);
generationUsage = response.usage;
catalog = parseGeneratedCatalog(
responseText(response.content),
tools,
manifestHash,
constraints(),
`${ctx.model.provider}/${ctx.model.id}`,
);
if (!catalogPreservesBundleAssignments(catalog, bundleState.catalog, bundleState.coveredNames)) {
throw new Error("generated catalog changed one or more curated my-pi bundle assignments");
}
catalogSource = "model";
try {
writeCachedCatalog(cachePath, catalog);
} catch (error) {
console.warn(`my-pi: could not cache pi-tool-search catalog: ${error instanceof Error ? error.message : String(error)}`);
}
registerToolSearch();
return {
notice: `Generated and cached ${catalog.groups.length} groups while preserving the precomputed my-pi assignments.`,
usage: generationUsage,
};
} catch (error) {
return {
notice: `Model enrichment failed; using the precomputed bundle catalog plus deterministic unknown-tool groups (${error instanceof Error ? error.message : String(error)}).`,
...(generationUsage ? { usage: generationUsage } : {}),
};
}
}
function refreshState(ctx: Pick<ExtensionContext, "ui">, forceReset: boolean): void {
const nextConfig = readToolSearchConfig(agentDir);
const registeredTools = pi.getAllTools().filter((tool) => tool.name !== TOOL_SEARCH_NAME);
const platformPolicy = platformToolPolicy(registeredTools, process.platform);
const allTools = platformPolicy.tools;
const availableNames = new Set(allTools.map((tool) => tool.name));
const nextPinned = new Set(
[...platformPolicy.coreTools, ...nextConfig.alwaysEnabled].filter((name) => availableNames.has(name)),
);
const hiddenTools = allTools.filter((tool) => !nextPinned.has(tool.name));
const nextConstraints = {
maxToolsPerGroup: Math.min(nextConfig.maxToolsPerGroup, nextConfig.maxDynamicTools),
groupOverrides: nextConfig.groupOverrides,
};
const nextHash = buildManifestHash(hiddenTools, nextConstraints);
const nextPolicySignature = JSON.stringify({
alwaysEnabled: [...nextPinned].sort(),
maxActiveGroups: nextConfig.maxActiveGroups,
maxDynamicTools: nextConfig.maxDynamicTools,
showToolSearchFooterStatus: nextConfig.showToolSearchFooterStatus,
});
const catalogChanged = forceReset || nextHash !== manifestHash;
const policyChanged = nextPolicySignature !== policySignature;
config = nextConfig;
tools = hiddenTools;
pinnedTools.clear();
for (const name of nextPinned) pinnedTools.add(name);
manifestHash = nextHash;
policySignature = nextPolicySignature;
if (catalogChanged) {
activeGroups.clear();
clock = 0;
lastEviction = "none";
attemptedGenerationHash = undefined;
installBaseCatalog(true);
}
const refreshEvictions: string[] = [];
const refreshReasons = new Set<string>();
while (activeGroups.size > config.maxActiveGroups || dynamicToolCount() > config.maxDynamicTools) {
if (activeGroups.size > config.maxActiveGroups) refreshReasons.add(`group cap ${config.maxActiveGroups}`);
if (dynamicToolCount() > config.maxDynamicTools) refreshReasons.add(`tool cap ${config.maxDynamicTools}`);
const evicted = leastRecentlyUsedGroup();
if (!evicted) break;
activeGroups.delete(evicted);
refreshEvictions.push(evicted);
}
const capacityChanged = refreshEvictions.length > 0;
if (capacityChanged) lastEviction = `${refreshEvictions.join(", ")} (${[...refreshReasons].join("; ")}; settings refresh)`;
if (catalogChanged || policyChanged || capacityChanged) {
registerToolSearch();
applyActiveTools();
}
updateStatus(ctx);
}
pi.registerCommand("tool-search-status", {
description: "Show Tool Search catalog, capacity, active-group LRU order, and the latest eviction reason",
handler: async (_args, ctx) => {
if (ctx.hasUI) ctx.ui.notify(currentStatusText(), "info");
},
});
pi.registerCommand("tool-search-rebuild", {
description: "Invalidate the generated tool-group catalog; rebuild lazily on the next tool_search call",
handler: async (_args, ctx) => {
try {
unlinkSync(cachePath);
} catch (error) {
if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error;
}
activeGroups.clear();
attemptedGenerationHash = undefined;
lastEviction = "none";
installBaseCatalog(false);
registerToolSearch();
applyActiveTools();
updateStatus(ctx);
if (ctx.hasUI) {
const message =
bundleState.unknownTools.length === 0
? "Restored the precomputed my-pi tool-group catalog; no model rebuild is needed"
: "Restored my-pi groups; unrecognized tools may be enriched on the next tool_search call";
ctx.ui.notify(message, "info");
}
},
});
pi.on("session_start", (_event, ctx) => {
refreshState(ctx, true);
});
pi.on("turn_start", (_event, ctx) => {
refreshState(ctx, false);
});
pi.on("tool_execution_start", (event) => {
const group = catalog.tools.find((tool) => tool.name === event.toolName)?.primaryGroup;
if (group && activeGroups.has(group)) activeGroups.set(group, ++clock);
});
}