feat(tool-search): expand dynamic group activation

2026-08-28 08:35:57 +00:00 · 2026-08-28 12:05:11 +08:00
parent bf874455db
commit 2dcec0207c
15 changed files with 897 additions and 119 deletions
@@ -18,6 +18,7 @@ import {
  parseGeneratedCatalog,
  rankGroups,
  readCachedCatalog,
+  selectConfidentGroup,
  writeCachedCatalog,
  type CatalogConstraints,
  type GroupCard,
@@ -32,6 +33,7 @@ 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";
@@ -102,7 +104,7 @@ function groupDescription(catalog: ToolCatalog, source: CatalogSource): string {
    })
    .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.",
+    "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)",
@@ -131,6 +133,7 @@ export default function toolSearchExtension(pi: ExtensionAPI): void {
  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 {
@@ -171,12 +174,32 @@ export default function toolSearchExtension(pi: ExtensionAPI): void {
    return catalog.groups.find((group) => group.id === id);
  }

-  function dynamicToolCount(): number {
-    return unique([...activeGroups.keys()].flatMap((id) => groupById(id)?.tools ?? [])).length;
+  function dynamicToolCount(groupIds: Iterable<string> = activeGroups.keys()): number {
+    return unique([...groupIds].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 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 {
@@ -189,7 +212,7 @@ export default function toolSearchExtension(pi: ExtensionAPI): void {
    ctx.ui.setStatus(
      "tool-search",
      config.showToolSearchFooterStatus
-        ? `${activeToolCount} / ${tools.length + 1} tools · ${activeGroups.size} / ${config.maxActiveGroups} groups`
+        ? `${activeToolCount} / ${tools.length + pinnedTools.size + 1} tools · ${activeGroups.size} / ${config.maxActiveGroups} groups`
        : undefined,
    );
  }
@@ -199,71 +222,147 @@ export default function toolSearchExtension(pi: ExtensionAPI): void {
      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`,
+      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 requestedId = params.group?.trim().toLowerCase();
-        let selected = requestedId ? groupById(requestedId) : undefined;
-        const generation: CatalogGenerationResult = selected
-          ? {}
-          : await ensureModelCatalog(ctx, signal);
+        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);

-        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;
+        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];
+        }

-        if (!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(
-            params.group ? `Unknown group: ${params.group}` : "No group was activated. Provide an exact group id from the catalog.",
+            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: { loadedGroup: undefined, evictedGroups: [], activeGroups: [...activeGroups.keys()], candidates: candidates.map(({ group, score }) => ({ id: group.id, score })), catalogSource },
+            details: emptyDetails(candidates.map(({ group, score }) => ({ id: group.id, score }))),
            ...(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}.`);
+        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: { loadedGroup: undefined, evictedGroups: [], activeGroups: [...activeGroups.keys()], candidates: [], catalogSource },
+            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[] = [];
-        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);
-          }
+        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);
        }
-        activeGroups.set(selected.id, ++clock);
-        applyActiveTools();
-        registerToolSearch();
-        updateStatus(ctx);

-        lines.push(`Loaded group: ${selected.id} (${groupTools.join(", ")})`);
+        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: selected.id,
+            loadedGroup: selectedIds.at(-1),
+            loadedGroups: selectedIds,
            evictedGroups,
            activeGroups: [...activeGroups.keys()],
            candidates: ranked.map(({ group, score }) => ({ id: group.id, score })),
@@ -379,17 +478,23 @@ export default function toolSearchExtension(pi: ExtensionAPI): void {
    if (catalogChanged) {
      activeGroups.clear();
      clock = 0;
+      lastEviction = "none";
      attemptedGenerationHash = undefined;
      installBaseCatalog(true);
    }

-    let capacityChanged = false;
+    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);
-      capacityChanged = true;
+      refreshEvictions.push(evicted);
    }
+    const capacityChanged = refreshEvictions.length > 0;
+    if (capacityChanged) lastEviction = `${refreshEvictions.join(", ")} (${[...refreshReasons].join("; ")}; settings refresh)`;
    if (catalogChanged || policyChanged || capacityChanged) {
      registerToolSearch();
      applyActiveTools();
@@ -397,6 +502,13 @@ export default function toolSearchExtension(pi: ExtensionAPI): void {
    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) => {
@@ -407,6 +519,7 @@ export default function toolSearchExtension(pi: ExtensionAPI): void {
      }
      activeGroups.clear();
      attemptedGenerationHash = undefined;
+      lastEviction = "none";
      installBaseCatalog(false);
      registerToolSearch();
      applyActiveTools();