my-pi/pi-tool-search/test/extension.test.ts

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, { createToolSearchExtension, platformToolPolicy } from "../extensions/index.ts";

interface ToolSearchResult {
  details: {
    loadedGroup?: string;
    loadedGroups: string[];
    evictedGroups: string[];
    activeGroups: string[];
    catalogSource: string;
  };
  usage?: unknown;
}

interface RegisteredTool {
  execute(
    id: string,
    params: { group?: string; groups?: 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;
  const configPath = join(agentDir, "tool-search-config.json");
  process.env.PI_CODING_AGENT_DIR = agentDir;

  try {
    await writeFile(
      configPath,
      JSON.stringify({
        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 commands = new Map<string, (...args: any[]) => unknown>();
    const notifications: string[] = [];
    let failActivation = false;
    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: (name: string, definition: { handler: (...args: any[]) => unknown }) => commands.set(name, definition.handler),
      setActiveTools: (names: string[]) => {
        if (failActivation) throw new Error("simulated setActiveTools failure");
        activeCalls.push([...names]);
      },
      on: (event: string, handler: (...args: any[]) => unknown) => handlers.set(event, handler),
    } as unknown as ExtensionAPI;

    createToolSearchExtension({ configPath })(api);
    const ctx = testContext({
      hasUI: true,
      ui: {
        setStatus: () => {},
        notify: (message: string) => notifications.push(message),
      } as unknown as ExtensionContext["ui"],
    });
    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"));

    const batch = await tool?.execute("call-4", { groups: ["beta", "gamma"] }, undefined, undefined, ctx);
    assert.deepEqual(batch?.details.loadedGroups, ["beta", "gamma"]);
    assert.deepEqual(batch?.details.evictedGroups, ["alpha"]);
    assert.deepEqual(batch?.details.activeGroups, ["gamma", "beta"]);

    const rejected = await tool?.execute(
      "call-5",
      { groups: ["alpha", "beta", "gamma"] },
      undefined,
      undefined,
      ctx,
    );
    assert.deepEqual(rejected?.details.loadedGroups, []);
    assert.deepEqual(rejected?.details.activeGroups, ["gamma", "beta"]);

    await commands.get("tool-search-status")?.("", ctx);
    assert.match(notifications.at(-1) ?? "", /Dynamic groups: 2\/2/);
    assert.match(notifications.at(-1) ?? "", /LRU oldest → newest: beta → gamma/);
    assert.match(notifications.at(-1) ?? "", /Last eviction: alpha \(group cap 2\)/);

    failActivation = true;
    await assert.rejects(
      tool!.execute("call-6", { groups: ["alpha", "gamma"] }, undefined, undefined, ctx),
      /simulated setActiveTools failure/,
    );
    failActivation = false;
    await commands.get("tool-search-status")?.("", ctx);
    assert.match(notifications.at(-1) ?? "", /LRU oldest → newest: beta → gamma/);
    assert.match(notifications.at(-1) ?? "", /Last eviction: alpha \(group cap 2\)/);
  } 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("research_search", "Search the public web with filters and return sources"),
      sourceTool("research_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: ["research_search", "research_fetch"],
        },
      ],
      tools: [
        { name: "research_search", summary: "Find web sources.", useWhen: ["Discover sources"], avoidWhen: [], keywords: ["search", "搜索"], primaryGroup: "web-research" },
        { name: "research_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, "my-pi", "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>();
    const activeCalls: string[][] = [];
    let modelCalls = 0;
    const sourceTools = [
      sourceTool("read", "Read a file"),
      sourceTool("web_search", "Search broadly for current web sources"),
      sourceTool("web_fetch", "Fetch a selected web source"),
      sourceTool("context_tree_query", "Inspect the context tree ledger"),
      sourceTool("overflow_query", "Inspect a stored oversized tool result"),
      sourceTool("memory_search", "Search durable memories"),
    ];
    const api = {
      getAllTools: () => sourceTools,
      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({
      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);
    assert.ok(["context_tree_query", "overflow_query", "memory_search"].every((name) => activeCalls.at(-1)?.includes(name)));

    const ambiguous = await registered.get("tool_search")?.execute(
      "call-1",
      { query: "web search" },
      undefined,
      undefined,
      ctx,
    );
    assert.equal(ambiguous?.details.loadedGroup, undefined);
    assert.deepEqual(ambiguous?.details.activeGroups, []);

    const result = await registered.get("tool_search")?.execute(
      "call-2",
      { query: "web_search" },
      undefined,
      undefined,
      ctx,
    );
    assert.equal(result?.details.loadedGroup, "web-search");
    assert.equal(result?.details.catalogSource, "bundle");
    assert.equal(modelCalls, 0);
    await assert.rejects(readFile(join(agentDir, "my-pi", "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 });
  }
});

test("keeps PowerShell with Bash only on Windows", () => {
  const definitions = [
    sourceTool("bash", "Run Bash commands"),
    sourceTool("powershell", "Run PowerShell commands"),
    sourceTool("read", "Read a file"),
  ];

  const nonWindows = platformToolPolicy(definitions, "darwin");
  assert.deepEqual(nonWindows.tools.map((tool) => tool.name), ["bash", "read"]);
  assert.ok(nonWindows.coreTools.includes("bash"));
  assert.ok(!nonWindows.coreTools.includes("powershell"));

  const windows = platformToolPolicy(definitions, "win32");
  assert.deepEqual(windows.tools.map((tool) => tool.name), ["bash", "powershell", "read"]);
  assert.ok(windows.coreTools.includes("bash"));
  assert.ok(windows.coreTools.includes("powershell"));
});

test("drops hidden and script-only tools but keeps reachable exposures", () => {
  const definitions = [
    sourceTool("read", "Read a file"),
    { ...sourceTool("hidden_helper", "Internal helper"), exposure: "hidden" },
    { ...sourceTool("script_only", "Only callable from scripts"), exposure: "codemode" },
    { ...sourceTool("codemode", "Batch several tool calls in one script"), exposure: "model-only" },
    { ...sourceTool("lazy_tool", "Activated on demand"), exposure: "deferred" },
  ];

  const { tools } = platformToolPolicy(definitions, "darwin");
  assert.deepEqual(tools.map((tool) => tool.name), ["read", "codemode", "lazy_tool"]);
});

test("pins a model-only codemode tool instead of hiding it", async () => {
  const agentDir = await mkdtemp(join(tmpdir(), "my-pi-tool-search-codemode-"));
  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 activeCalls: string[][] = [];
    const sourceTools = [
      sourceTool("read", "Read a file"),
      { ...sourceTool("codemode", "Batch several tool calls in one script"), exposure: "model-only" },
    ];
    const api = {
      getAllTools: () => sourceTools,
      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;
    // The bundle config pins "codemode" in alwaysEnabled; exposure "model-only"
    // stays reachable, so it must land in the active set, not a dynamic group.
    toolSearchExtension(api);
    handlers.get("session_start")?.({}, testContext());

    assert.ok(activeCalls.at(-1)?.includes("codemode"));
    assert.ok(activeCalls.at(-1)?.includes("tool_search"));
    assert.ok(activeCalls.at(-1)?.includes("read"));
  } 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("loads an exact known group without waiting for unknown-tool model enrichment", async () => {
  const agentDir = await mkdtemp(join(tmpdir(), "my-pi-tool-search-fast-path-"));
  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("web_search", "Search broadly for current web sources"),
      sourceTool("web_fetch", "Fetch a selected web source"),
      sourceTool("unrecognized_extra", "An unrecognized third-party capability"),
    ];
    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("an exact known group must not wait for model enrichment");
        },
      } as unknown as ExtensionContext["modelRegistry"],
    });
    handlers.get("session_start")?.({}, ctx);

    const result = await registered.get("tool_search")?.execute(
      "call-1",
      { group: "web-search" },
      undefined,
      undefined,
      ctx,
    );
    assert.equal(result?.details.loadedGroup, "web-search");
    assert.equal(result?.details.catalogSource, "hybrid");
    assert.equal(modelCalls, 0);
  } 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 });
  }
});