fix(server): Pi discovers optional T3 tools on demand

apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.test.ts

@@ -9,6 +9,308 @@ type RequestHook = (
   ctx: { model: { provider: string } },
 ) => Record<string, unknown> | undefined;
 
+interface RegisteredTool {
+  readonly name: string;
+  readonly description: string;
+  readonly parameters: unknown;
+  readonly exposure?: string;
+  readonly promptSnippet?: string;
+  readonly promptGuidelines?: ReadonlyArray<string>;
+  readonly execute: (
+    id: string,
+    args: Record<string, unknown>,
+    signal?: AbortSignal,
+  ) => Promise<{
+    readonly content: ReadonlyArray<{ readonly type: string; readonly text: string }>;
+  }>;
+}
+
+type AgentStartHook = (
+  event: { systemPrompt: string },
+  ctx: { ui: { notify: (message: string, severity: string) => void } },
+) => Promise<{ systemPrompt: string }>;
+
+async function loadMcpBridge(
+  options: {
+    readonly modern?: boolean;
+    readonly toolSearchAvailable?: boolean;
+    readonly toolSearchDisabled?: boolean;
+    readonly allowsTool?: (name: string) => boolean;
+  } = {},
+) {
+  const handlers = new Map<string, AgentStartHook>();
+  const tools: RegisteredTool[] = [];
+  const requests: Array<{ readonly method: string; readonly params?: unknown }> = [];
+  let activeTools = ["read"];
+  const transports: Array<{
+    readonly url: string;
+    readonly authorization: string;
+    readonly signal: AbortSignal | undefined;
+  }> = [];
+  const servers: Array<{ readonly name: string; readonly config: Record<string, unknown> }> = [];
+  const catalog = [
+    { name: "orchestrator_capabilities", description: "Discover available providers and models." },
+    { name: "delegate_task", description: "Delegate work to another agent." },
+    { name: "task_status", description: "Check delegated work." },
+    { name: "preview_snapshot", description: "Inspect the collaborative browser." },
+  ].map((tool) => ({ ...tool, inputSchema: { type: "object", properties: {} } }));
+  const source = NodeModule.stripTypeScriptTypes(
+    PI_T3_MCP_EXTENSION_SOURCE.replace('import { Type } from "typebox";', "").replace(
+      "export default async function",
+      "async function",
+    ),
+  );
+  await NodeVM.runInNewContext(`${source}\nt3McpExtension(pi)`, {
+    process: {
+      env: { T3_MCP_URL: "http://fixture.invalid/mcp", T3_MCP_BEARER_TOKEN: "fixture-token" },
+    },
+    AbortSignal,
+    Type: { Unsafe: (schema: unknown) => schema },
+    fetch: async (
+      url: string,
+      options: { body: string; headers: Record<string, string>; signal?: AbortSignal },
+    ) => {
+      transports.push({
+        url,
+        authorization: options.headers.authorization!,
+        signal: options.signal,
+      });
+      const request = JSON.parse(options.body) as { id: number; method: string; params?: unknown };
+      requests.push(request);
+      const result =
+        request.method === "tools/list"
+          ? { tools: catalog }
+          : request.method === "tools/call"
+            ? { content: [{ type: "text", text: "browser snapshot" }] }
+            : {};
+      return new Response(JSON.stringify({ jsonrpc: "2.0", id: request.id, result }), {
+        headers: { "content-type": "application/json" },
+      });
+    },
+    pi: {
+      on: (name: string, handler: AgentStartHook) => handlers.set(name, handler),
+      registerTool: (tool: RegisteredTool) => {
+        if (options.allowsTool && !options.allowsTool(tool.name)) return;
+        const index = tools.findIndex((current) => current.name === tool.name);
+        if (index === -1) tools.push(tool);
+        else tools[index] = tool;
+      },
+      getActiveTools: () => activeTools,
+      setActiveTools: (names: string[]) => {
+        activeTools = names.filter((name) => options.allowsTool?.(name) ?? true);
+      },
+      getAllTools: () =>
+        options.toolSearchAvailable &&
+        !options.toolSearchDisabled &&
+        (options.allowsTool?.("tool_search") ?? true)
+          ? [{ name: "tool_search", sourceInfo: { path: "builtin:tool-search" } }]
+          : [],
+      ...(options.modern
+        ? {
+            registerMcpServer: (name: string, config: Record<string, unknown>) =>
+              servers.push({ name, config }),
+            unregisterMcpServer: () => servers.splice(0),
+          }
+        : {}),
+    },
+  });
+  return {
+    handlers,
+    tools,
+    requests,
+    servers,
+    transports,
+    getActiveTools: () => activeTools,
+    restoreActiveTools: (names: string[]) => {
+      activeTools = names;
+    },
+  };
+}
+
+describe("Pi MCP tool exposure", () => {
+  it("keeps orchestration direct and optional bridge tools discoverable on modern Pi", async () => {
+    const bridge = await loadMcpBridge({ modern: true, toolSearchAvailable: true });
+    await bridge.handlers.get("session_start")!(
+      { systemPrompt: "" },
+      { ui: { notify: () => undefined } },
+    );
+    const start = bridge.handlers.get("before_agent_start");
+    assert.isDefined(start);
+    const prompt = await start!(
+      { systemPrompt: "Pi system prompt" },
+      { ui: { notify: () => undefined } },
+    );
+    assert.include(prompt.systemPrompt, "orchestrator_capabilities");
+    assert.equal(bridge.servers.length, 0);
+    assert.equal(bridge.tools.length, 8);
+    assert.deepEqual(bridge.getActiveTools(), ["read", "tool_search"]);
+    assert.deepEqual(
+      bridge.tools
+        .filter((tool) => tool.exposure !== "hidden")
+        .map((tool) => [tool.name, tool.exposure]),
+      [
+        ["mcp__t3-code__orchestrator_capabilities", "direct"],
+        ["mcp__t3-code__delegate_task", "direct"],
+        ["mcp__t3-code__task_status", "direct"],
+        ["mcp__t3-code__preview_snapshot", "deferred"],
+      ],
+    );
+    assert.deepEqual(
+      bridge.tools.filter((tool) => tool.exposure === "hidden").map((tool) => tool.name),
+      [
+        "mcp__t3_code__orchestrator_capabilities",
+        "mcp__t3_code__delegate_task",
+        "mcp__t3_code__task_status",
+        "mcp__t3_code__preview_snapshot",
+      ],
+    );
+    const result = await bridge.tools
+      .find((tool) => tool.name === "mcp__t3-code__preview_snapshot")!
+      .execute("call-1", { depth: 2 });
+    assert.equal(result.content[0]?.text, "browser snapshot");
+    assert.equal(bridge.requests.at(-1)?.method, "tools/call");
+  });
+
+  it.each(["legacy Pi", "disabled tool search"])(
+    "keeps tool execution available with %s",
+    async (mode) => {
+      const bridge = await loadMcpBridge({
+        modern: mode !== "legacy Pi",
+        toolSearchAvailable: mode === "disabled tool search",
+        toolSearchDisabled: mode === "disabled tool search",
+      });
+      if (mode !== "legacy Pi") {
+        const start = bridge.handlers.get("session_start");
+        await start!({ systemPrompt: "Pi system prompt" }, { ui: { notify: () => undefined } });
+      }
+      assert.equal(bridge.tools.filter((tool) => tool.exposure !== "hidden").length, 4);
+      assert.isTrue(
+        bridge.tools
+          .filter((tool) => tool.exposure !== "hidden")
+          .every((tool) => tool.exposure === undefined || tool.exposure === "direct"),
+      );
+      const tool = bridge.tools.find((tool) => tool.name === "mcp__t3-code__preview_snapshot");
+      assert.isDefined(tool);
+      const controller = new AbortController();
+      const result = await tool!.execute("call-1", { depth: 2 }, controller.signal);
+      assert.equal(result.content[0]?.text, "browser snapshot");
+      assert.strictEqual(bridge.transports.at(-1)?.signal, controller.signal);
+      assert.equal(bridge.transports.at(-1)?.url, "http://fixture.invalid/mcp");
+      assert.equal(bridge.transports.at(-1)?.authorization, "Bearer fixture-token");
+      assert.deepEqual(JSON.parse(JSON.stringify(bridge.requests.at(-1))), {
+        jsonrpc: "2.0",
+        id: 3,
+        method: "tools/call",
+        params: { name: "preview_snapshot", arguments: { depth: 2 } },
+      });
+      assert.isUndefined(tool?.promptSnippet);
+      assert.isUndefined(tool?.promptGuidelines);
+    },
+  );
+
+  it("preserves legacy wildcard tool selection", async () => {
+    const bridge = await loadMcpBridge({
+      modern: true,
+      toolSearchAvailable: true,
+      allowsTool: (name) =>
+        name === "read" || name === "tool_search" || name.startsWith("mcp__t3-code__"),
+    });
+    await bridge.handlers.get("session_start")!(
+      { systemPrompt: "" },
+      { ui: { notify: () => undefined } },
+    );
+    assert.equal(bridge.tools.length, 4);
+    assert.equal(bridge.tools.filter((tool) => tool.exposure === "direct").length, 3);
+    const tool = bridge.tools.find((tool) => tool.name === "mcp__t3-code__preview_snapshot");
+    assert.isDefined(tool);
+    assert.equal((await tool!.execute("selected", {})).content[0]?.text, "browser snapshot");
+  });
+
+  it.each([false, true])(
+    "reconciles tree loadouts while honoring search exclusion: %s",
+    async (excludeSearch) => {
+      const bridge = await loadMcpBridge({
+        modern: true,
+        toolSearchAvailable: true,
+        allowsTool: (name) =>
+          name !== "mcp__t3-code__delegate_task" && (!excludeSearch || name !== "tool_search"),
+      });
+      await bridge.handlers.get("session_start")!(
+        { systemPrompt: "" },
+        { ui: { notify: () => undefined } },
+      );
+      bridge.restoreActiveTools([
+        "read",
+        "mcp__t3-code__task_status",
+        "mcp__t3-code__preview_snapshot",
+      ]);
+      const tree = bridge.handlers.get("session_tree");
+      assert.isDefined(tree);
+      await tree!({ systemPrompt: "" }, { ui: { notify: () => undefined } });
+      assert.deepEqual(bridge.getActiveTools(), [
+        "read",
+        "mcp__t3-code__task_status",
+        "mcp__t3-code__preview_snapshot",
+        ...(!excludeSearch ? ["tool_search"] : []),
+      ]);
+      assert.isFalse(
+        bridge.tools.some(
+          (tool) => tool.exposure !== "hidden" && tool.name.endsWith("__delegate_task"),
+        ),
+      );
+    },
+  );
+});
+
+describe("Pi tool discovery permissions", () => {
+  it("allows discovery without confirmation and still gates the discovered tool", async () => {
+    type ToolCallHook = (
+      event: { toolName: string; input: unknown },
+      ctx: { ui: { confirm: (title: string, detail: string) => Promise<boolean> } },
+    ) => Promise<{ block: true; reason: string } | undefined>;
+    let toolCall: ToolCallHook | undefined;
+    let searchPath = "builtin:tool-search";
+    const source = NodeModule.stripTypeScriptTypes(
+      PI_T3_MCP_EXTENSION_SOURCE.replace('import { Type } from "typebox";', "").replace(
+        "export default async function",
+        "async function",
+      ),
+    );
+    await NodeVM.runInNewContext(`${source}\nt3McpExtension(pi)`, {
+      process: { env: { T3_PI_RUNTIME_MODE: "approval-required" } },
+      pi: {
+        on: (name: string, handler: ToolCallHook) => {
+          if (name === "tool_call") toolCall = handler;
+        },
+        getAllTools: () => [{ name: "tool_search", sourceInfo: { path: searchPath } }],
+      },
+    });
+    assert.isDefined(toolCall);
+    const confirmations: string[] = [];
+    const ctx = {
+      ui: {
+        confirm: async (title: string) => {
+          confirmations.push(title);
+          return false;
+        },
+      },
+    };
+    assert.isUndefined(
+      await toolCall!({ toolName: "tool_search", input: { query: "preview_snapshot" } }, ctx),
+    );
+    assert.equal(confirmations.length, 0);
+    const result = await toolCall!({ toolName: "mcp__t3-code__preview_snapshot", input: {} }, ctx);
+    assert.equal(result?.block, true);
+    assert.deepEqual(confirmations, ["Allow mcp__t3-code__preview_snapshot?"]);
+
+    // An extension that replaces the search builtin is not known to be read-only.
+    searchPath = "/extensions/custom-search.ts";
+    const replaced = await toolCall!({ toolName: "tool_search", input: {} }, ctx);
+    assert.equal(replaced?.block, true);
+    assert.equal(confirmations.at(-1), "Allow tool_search?");
+  });
+});
+
 async function loadRequestHook(): Promise<RequestHook> {
   const handlers = new Map<string, RequestHook>();
   // Execute the shipped extension with MCP disabled; this path needs no Typebox.