my-pi/pi-tool-search/extensions/index.ts

/**
 * 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 POWERSHELL_TOOL_NAME = "powershell";
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 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();
        let selected = requestedId ? groupById(requestedId) : undefined;
        const generation: CatalogGenerationResult = selected
          ? {}
          : await ensureModelCatalog(ctx, signal);
        const lines: string[] = [];
        if (generation.notice) lines.push(generation.notice);

        if (!selected && requestedId) selected = groupById(requestedId);
        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 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;
      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);
  });
}