feat: vendor grouped tool search

pi-tool-search/extensions/index.ts

@@ -0,0 +1,423 @@
+/**
+ * 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,
+  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 CORE_TOOLS = ["read", "write", "edit", "bash", "grep", "find"];
+
+type CatalogSource = "bundle" | "cache" | "fallback" | "hybrid" | "model";
+type ModelUsage = Awaited<ReturnType<ExtensionContext["modelRegistry"]["complete"]>>["usage"];
+
+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 a complete tool group for the current task. Prefer an exact group id 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: [] };
+
+  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(): number {
+    return unique([...activeGroups.keys()].flatMap((id) => groupById(id)?.tools ?? [])).length;
+  }
+
+  function leastRecentlyUsedGroup(): string | undefined {
+    return [...activeGroups].sort((left, right) => left[1] - right[1] || left[0].localeCompare(right[0]))[0]?.[0];
+  }
+
+  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 + 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; at most ${config.maxActiveGroups} dynamic groups remain active`,
+      parameters: Type.Object({
+        group: Type.Optional(Type.String({ description: "Exact group id from the tool_search catalog" })),
+        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 requestedId = params.group?.trim().toLowerCase();
+        const preGenerationGroup = requestedId ? groupById(requestedId) : undefined;
+        const generation = await ensureModelCatalog(ctx, signal);
+        const lines: string[] = [];
+        if (generation.notice) lines.push(generation.notice);
+
+        let selected = requestedId ? groupById(requestedId) : undefined;
+        if (!selected && preGenerationGroup) {
+          const previousNames = new Set(preGenerationGroup.tools);
+          selected = [...catalog.groups]
+            .map((group) => ({ group, overlap: group.tools.filter((name) => previousNames.has(name)).length }))
+            .sort((left, right) => right.overlap - left.overlap || left.group.id.localeCompare(right.group.id))
+            .find((candidate) => candidate.overlap > 0)?.group;
+          if (selected) lines.push(`Mapped initial fallback group ${preGenerationGroup.id} to generated group ${selected.id}.`);
+        }
+        const ranked = params.query ? rankGroups(catalog, params.query).slice(0, 3) : [];
+        if (!selected && params.query && (ranked[0]?.score ?? 0) > 0) selected = ranked[0]?.group;
+
+        if (!selected) {
+          const candidates = ranked.length > 0 ? ranked : catalog.groups.slice(0, 5).map((group) => ({ group, score: 0 }));
+          lines.push(
+            params.group ? `Unknown group: ${params.group}` : "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: { loadedGroup: undefined, evictedGroups: [], activeGroups: [...activeGroups.keys()], candidates: candidates.map(({ group, score }) => ({ id: group.id, score })), catalogSource },
+            ...(generation.usage ? { usage: generation.usage } : {}),
+          };
+        }
+
+        const groupTools = unique(selected.tools);
+        if (groupTools.length > config.maxDynamicTools) {
+          lines.push(`Group ${selected.id} has ${groupTools.length} tools, exceeding maxDynamicTools=${config.maxDynamicTools}.`);
+          return {
+            content: [{ type: "text", text: lines.join("\n") }],
+            details: { loadedGroup: undefined, evictedGroups: [], activeGroups: [...activeGroups.keys()], candidates: [], catalogSource },
+            ...(generation.usage ? { usage: generation.usage } : {}),
+          };
+        }
+
+        const evictedGroups: string[] = [];
+        if (!activeGroups.has(selected.id)) {
+          while (
+            activeGroups.size >= config.maxActiveGroups ||
+            (activeGroups.size > 0 && dynamicToolCount() + groupTools.length > config.maxDynamicTools)
+          ) {
+            const evicted = leastRecentlyUsedGroup();
+            if (!evicted) break;
+            activeGroups.delete(evicted);
+            evictedGroups.push(evicted);
+          }
+        }
+        activeGroups.set(selected.id, ++clock);
+        applyActiveTools();
+        registerToolSearch();
+        updateStatus(ctx);
+
+        lines.push(`Loaded group: ${selected.id} (${groupTools.join(", ")})`);
+        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: selected.id,
+            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 allTools = pi.getAllTools().filter((tool) => tool.name !== TOOL_SEARCH_NAME);
+    const availableNames = new Set(allTools.map((tool) => tool.name));
+    const nextPinned = new Set([...CORE_TOOLS, ...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;
+      attemptedGenerationHash = undefined;
+      installBaseCatalog(true);
+    }
+
+    let capacityChanged = false;
+    while (activeGroups.size > config.maxActiveGroups || dynamicToolCount() > config.maxDynamicTools) {
+      const evicted = leastRecentlyUsedGroup();
+      if (!evicted) break;
+      activeGroups.delete(evicted);
+      capacityChanged = true;
+    }
+    if (catalogChanged || policyChanged || capacityChanged) {
+      registerToolSearch();
+      applyActiveTools();
+    }
+    updateStatus(ctx);
+  }
+
+  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;
+      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);
+  });
+}